Merge branch 'enterprise-security-features'

# Please enter a commit message to explain why this merge is necessary,
# especially if it merges an updated upstream into a topic branch.
#
# Lines starting with '#' will be ignored, and an empty message aborts
# the commit.

Author: iPeluwa

Cargo.lock

@@ -17,4 +17,5 @@ datafusion = { workspace = true, default-features = true, features = ["avro"] }
 tokio = { workspace = true, features = ["full"] }
 datafusion-postgres = { path = "../datafusion-postgres", version = "0.6.1" }
 structopt = { version = "0.3", default-features = false }
+log = "0.4"
 env_logger = "0.11"

datafusion-postgres-cli/Cargo.toml

@@ -8,6 +8,7 @@ use datafusion::execution::options::{
 use datafusion::prelude::{SessionConfig, SessionContext};
 use datafusion_postgres::pg_catalog::setup_pg_catalog;
 use datafusion_postgres::{serve, ServerOptions}; // Assuming the crate name is `datafusion_postgres`
+use log::info;
 use structopt::StructOpt;
 
 #[derive(Debug, StructOpt)]
@@ -129,7 +130,7 @@ async fn setup_session_context(
             .register_csv(table_name, table_path, CsvReadOptions::default())
             .await
             .map_err(|e| format!("Failed to register CSV table '{table_name}': {e}"))?;
-        println!("Loaded {table_path} as table {table_name}");
+        info!("Loaded {table_path} as table {table_name}");
     }
 
     // Register JSON tables
@@ -138,7 +139,7 @@ async fn setup_session_context(
             .register_json(table_name, table_path, NdJsonReadOptions::default())
             .await
             .map_err(|e| format!("Failed to register JSON table '{table_name}': {e}"))?;
-        println!("Loaded {table_path} as table {table_name}");
+        info!("Loaded {table_path} as table {table_name}");
     }
 
     // Register Arrow tables
@@ -151,7 +152,7 @@ async fn setup_session_context(
             .register_arrow(table_name, table_path, ArrowReadOptions::default())
             .await
             .map_err(|e| format!("Failed to register Arrow table '{table_name}': {e}"))?;
-        println!("Loaded {table_path} as table {table_name}");
+        info!("Loaded {table_path} as table {table_name}");
     }
 
     // Register Parquet tables
@@ -164,7 +165,7 @@ async fn setup_session_context(
             .register_parquet(table_name, table_path, ParquetReadOptions::default())
             .await
             .map_err(|e| format!("Failed to register Parquet table '{table_name}': {e}"))?;
-        println!("Loaded {table_path} as table {table_name}");
+        info!("Loaded {table_path} as table {table_name}");
     }
 
     // Register Avro tables
@@ -173,7 +174,7 @@ async fn setup_session_context(
             .register_avro(table_name, table_path, AvroReadOptions::default())
             .await
             .map_err(|e| format!("Failed to register Avro table '{table_name}': {e}"))?;
-        println!("Loaded {table_path} as table {table_name}");
+        info!("Loaded {table_path} as table {table_name}");
     }
 
     // Register pg_catalog

datafusion-postgres-cli/src/main.rs

@@ -2,6 +2,7 @@ use std::collections::HashMap;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use log::warn;
 use pgwire::api::auth::{AuthSource, LoginInfo, Password};
 use pgwire::error::{PgWireError, PgWireResult};
 use tokio::sync::RwLock;
@@ -175,7 +176,8 @@ impl AuthManager {
 
                 // Create predefined roles
                 if let Err(e) = auth_manager_spawn.create_predefined_roles().await {
-                    eprintln!("Failed to create predefined roles: {e:?}");
+                    warn!("Failed to create predefined roles: {e:?}");
+
                 }
             }
         });

datafusion-postgres/src/auth.rs

@@ -9,7 +9,7 @@ use datafusion::prelude::SessionContext;
 
 pub mod auth;
 use getset::{Getters, Setters, WithSetters};
-
+use log::{info, warn};
 use pgwire::tokio::process_socket;
 use rustls_pemfile::{certs, pkcs8_private_keys};
 use rustls_pki_types::{CertificateDer, PrivateKeyDer};
@@ -86,26 +86,26 @@ pub async fn serve(
         if let (Some(cert_path), Some(key_path)) = (&opts.tls_cert_path, &opts.tls_key_path) {
             match setup_tls(cert_path, key_path) {
                 Ok(acceptor) => {
-                    println!("TLS enabled using cert: {cert_path} and key: {key_path}");
+                    info!("TLS enabled using cert: {cert_path} and key: {key_path}");
                     Some(acceptor)
                 }
                 Err(e) => {
-                    eprintln!("Failed to setup TLS: {e}. Running without encryption.");
+                    warn!("Failed to setup TLS: {e}. Running without encryption.");
                     None
                 }
             }
         } else {
-            println!("TLS not configured. Running without encryption.");
+            info!("TLS not configured. Running without encryption.");
             None
         };
 
     // Bind to the specified host and port
     let server_addr = format!("{}:{}", opts.host, opts.port);
     let listener = TcpListener::bind(&server_addr).await?;
     if tls_acceptor.is_some() {
-        println!("Listening on {server_addr} with TLS encryption");
+        info!("Listening on {server_addr} with TLS encryption");
     } else {
-        println!("Listening on {server_addr} (unencrypted)");
+        info!("Listening on {server_addr} (unencrypted)");
     }
 
     // Accept incoming connections
@@ -117,13 +117,13 @@ pub async fn serve(
                 // Connection accepted from {addr} - log appropriately based on your logging strategy
 
                 tokio::spawn(async move {
-                    if let Err(_e) = process_socket(socket, tls_acceptor_ref, factory_ref).await {
-                        // Log error or handle appropriately based on your logging strategy
+                    if let Err(e) = process_socket(socket, tls_acceptor_ref, factory_ref).await {
+                        warn!("Error processing socket: {e}");
                     }
                 });
             }
             Err(e) => {
-                eprintln!("Error accept socket: {e}");
+                warn!("Error accept socket: {e}");
             }
         }
     }