feat: add a serve function

Author: sunng87

datafusion-postgres-cli/src/main.rs

@@ -1,13 +1,9 @@
-use std::sync::Arc;
-
 use datafusion::execution::options::{
     ArrowReadOptions, AvroReadOptions, CsvReadOptions, NdJsonReadOptions, ParquetReadOptions,
 };
 use datafusion::prelude::SessionContext;
-use datafusion_postgres::{DfSessionService, HandlerFactory}; // Assuming the crate name is `datafusion_postgres`
-use pgwire::tokio::process_socket;
+use datafusion_postgres::{serve, ServerOptions}; // Assuming the crate name is `datafusion_postgres`
 use structopt::StructOpt;
-use tokio::net::TcpListener;
 
 #[derive(Debug, StructOpt)]
 #[structopt(
@@ -103,33 +99,14 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         println!("Loaded {} as table {}", table_path, table_name);
     }
 
-    // Get the first catalog name from the session context
-    let catalog_name = session_context
-        .catalog_names() // Fixed: Removed .catalog_list()
-        .first()
-        .cloned();
-
-    // Create the handler factory with the session context and catalog name
-    let factory = Arc::new(HandlerFactory(Arc::new(DfSessionService::new(
-        session_context,
-        catalog_name,
-    ))));
+    let server_options = ServerOptions {
+        host: opts.host.clone(),
+        port: opts.port,
+    };
 
-    // Bind to the specified host and port
-    let server_addr = format!("{}:{}", opts.host, opts.port);
-    let listener = TcpListener::bind(&server_addr).await?;
-    println!("Listening on {}", server_addr);
+    serve(session_context, &server_options)
+        .await
+        .map_err(|e| format!("Failed to run server: {}", e))?;
 
-    // Accept incoming connections
-    loop {
-        let (socket, addr) = listener.accept().await?;
-        let factory_ref = factory.clone();
-        println!("Accepted connection from {}", addr);
-
-        tokio::spawn(async move {
-            if let Err(e) = process_socket(socket, None, factory_ref).await {
-                eprintln!("Error processing socket: {}", e);
-            }
-        });
-    }
+    Ok(())
 }

datafusion-postgres/Cargo.toml

@@ -25,4 +25,4 @@ log = "0.4"
 pgwire = { workspace = true }
 postgres-types = "0.2"
 rust_decimal = { version = "1.37", features = ["db-postgres"] }
-tokio = { version = "1.45", features = ["sync"] }
+tokio = { version = "1.45", features = ["sync", "net"] }

datafusion-postgres/src/lib.rs

@@ -4,3 +4,60 @@ mod handlers;
 mod information_schema;
 
 pub use handlers::{DfSessionService, HandlerFactory, Parser};
+
+use std::sync::Arc;
+
+use datafusion::prelude::SessionContext;
+use pgwire::tokio::process_socket;
+use tokio::net::TcpListener;
+
+pub struct ServerOptions {
+    pub host: String,
+    pub port: u16,
+}
+
+impl Default for ServerOptions {
+    fn default() -> Self {
+        ServerOptions {
+            host: "127.0.0.1".to_string(),
+            port: 5432,
+        }
+    }
+}
+
+/// Serve the Datafusion `SessionContext` with Postgres protocol.
+pub async fn serve(
+    session_context: SessionContext,
+    opts: &ServerOptions,
+) -> Result<(), std::io::Error> {
+    // Get the first catalog name from the session context
+    let catalog_name = session_context
+        .catalog_names() // Fixed: Removed .catalog_list()
+        .first()
+        .cloned();
+
+    // Create the handler factory with the session context and catalog name
+    let factory = Arc::new(HandlerFactory(Arc::new(DfSessionService::new(
+        session_context,
+        catalog_name,
+    ))));
+
+    // Bind to the specified host and port
+    let server_addr = format!("{}:{}", opts.host, opts.port);
+    let listener = TcpListener::bind(&server_addr).await?;
+    println!("Listening on {}", server_addr);
+
+    // Accept incoming connections
+    loop {
+        if let Ok((socket, addr)) = listener.accept().await {
+            let factory_ref = factory.clone();
+            println!("Accepted connection from {}", addr);
+
+            tokio::spawn(async move {
+                if let Err(e) = process_socket(socket, None, factory_ref).await {
+                    eprintln!("Error processing socket: {}", e);
+                }
+            });
+        };
+    }
+}