Quick cleanup

Author: backkem

datafusion-flight-sql-server/Cargo.toml

@@ -27,6 +27,7 @@ once_cell = "1.21"
 prost = "0.13"
 tonic.workspace = true
 async-trait.workspace = true
+tonic-async-interceptor = "0.12.0"
 
 [dev-dependencies]
 tokio.workspace = true

datafusion-flight-sql-server/examples/bearer_auth_flight_sql.rs

@@ -24,21 +24,20 @@
 //! Observe the server console output for messages from the interceptor and session provider,
 //! and the client output for the success/failure of each attempt.
 
-use std::sync::Arc;
 use std::time::Duration;
 
+use arrow_flight::flight_service_server::FlightServiceServer;
 use arrow_flight::sql::client::FlightSqlServiceClient;
 use arrow_flight::sql::CommandGetTables;
 use async_trait::async_trait;
 use datafusion::error::{DataFusionError, Result};
+use datafusion::execution::context::SessionState; // SessionState is not in prelude
 use datafusion::prelude::*; // Covers SessionContext, CsvReadOptions, etc.
-use datafusion::execution::context::{SessionConfig, SessionState}; // SessionState is not in prelude
-use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion_flight_sql_server::service::FlightSqlService;
 use datafusion_flight_sql_server::session::SessionStateProvider;
 use tokio::time::sleep;
 use tonic::transport::{Channel, Endpoint, Server};
-use tonic::{metadata::MetadataValue, Request, Status};
+use tonic::{Request, Status};
 
 // UserData struct remains the same
 #[derive(Clone, Debug)]
@@ -76,24 +75,27 @@ async fn bearer_auth_interceptor(mut req: Request<()>) -> Result<Request<()>, St
 }
 
 // Updated MySessionStateProvider
-#[derive(Debug, Clone)] // Removed Default
+// #[derive(Debug, Clone)] // Removed Default
 pub struct MySessionStateProvider {
     base_context: SessionContext,
 }
 
 impl MySessionStateProvider {
-    async fn try_new() -> Result<Self> { // datafusion::error::Result
+    async fn try_new() -> Result<Self> {
+        // datafusion::error::Result
         let ctx = SessionContext::new();
         // Construct path to test.csv relative to CARGO_MANIFEST_DIR of the datafusion-flight-sql-server crate
         let csv_path = concat!(env!("CARGO_MANIFEST_DIR"), "/examples/test.csv");
-        ctx.register_csv("test", csv_path, CsvReadOptions::new()).await?;
+        ctx.register_csv("test", csv_path, CsvReadOptions::new())
+            .await?;
         Ok(Self { base_context: ctx })
     }
 }
 
 #[async_trait]
 impl SessionStateProvider for MySessionStateProvider {
-    async fn new_context(&self, request: &Request<()>) -> Result<SessionState, Status> { // tonic::Result for Status
+    async fn new_context(&self, request: &Request<()>) -> Result<SessionState, Status> {
+        // tonic::Result for Status
         if let Some(user_data) = request.extensions().get::<UserData>() {
             println!(
                 "Session context for user_id: {}. Cloning base context.",
@@ -115,46 +117,48 @@ impl SessionStateProvider for MySessionStateProvider {
 async fn new_client_with_auth(
     dsn: String,
     token: Option<String>,
-) -> Result<FlightSqlServiceClient<Channel>> { // datafusion::error::Result
+) -> Result<FlightSqlServiceClient<Channel>> {
+    // datafusion::error::Result
     let endpoint = Endpoint::from_shared(dsn.clone())
         .map_err(|e| DataFusionError::External(format!("Invalid DSN {}: {}", dsn, e).into()))?
         .connect_timeout(std::time::Duration::from_secs(10));
 
-    let channel = endpoint.connect().await
-        .map_err(|e| DataFusionError::External(format!("Failed to connect to {}: {}", dsn, e).into()))?;
-
-    let service_client = FlightSqlServiceClient::with_interceptor(
-        channel,
-        move |mut req: Request<()>| {
-            if let Some(token_str) = token.clone() {
-                let bearer_token = format!("Bearer {}", token_str);
-                match MetadataValue::try_from(&bearer_token) {
-                    Ok(metadata_val) => req.metadata_mut().insert("authorization", metadata_val),
-                    Err(_) => return Err(Status::invalid_argument("Invalid token format for metadata")),
-                };
-            }
-            Ok(req)
-        },
-    );
-    Ok(service_client)
-}
+    let channel = endpoint.connect().await.map_err(|e| {
+        DataFusionError::External(format!("Failed to connect to {}: {}", dsn, e).into())
+    })?;
 
-fn status_to_df_error(err: tonic::Status) -> DataFusionError {
-    DataFusionError::External(format!("Tonic status error: {}", err).into())
+    let mut service_client = FlightSqlServiceClient::new(channel);
+    if let Some(token_str) = token.clone() {
+        service_client.set_header("authorization", format!("Bearer {}", token_str));
+    }
+    Ok(service_client)
 }
 
 #[tokio::main]
-async fn main() -> Result<()> { // datafusion::error::Result<()>
+async fn main() -> Result<()> {
+    // datafusion::error::Result<()>
     // Server Setup
     let dsn: String = "0.0.0.0:50051".to_string();
-    let state_provider = Arc::new(MySessionStateProvider::try_new().await?);
-    let base_service = FlightSqlService::new_with_state_provider(state_provider).into_service();
-    let wrapped_service = tonic::service::interceptor_fn(base_service, bearer_auth_interceptor);
-    let addr: std::net::SocketAddr = dsn.parse().map_err(|e| DataFusionError::External(format!("Invalid address format {}: {}", dsn, e).into()))?;
+    let state_provider = Box::new(MySessionStateProvider::try_new().await?);
+    let base_service = FlightSqlService::new_with_provider(state_provider);
+    let svc: FlightServiceServer<FlightSqlService> = FlightServiceServer::new(base_service);
+    let addr: std::net::SocketAddr = dsn.parse().map_err(|e| {
+        DataFusionError::External(format!("Invalid address format {}: {}", dsn, e).into())
+    })?;
 
     tokio::spawn(async move {
-        println!("Bearer Authentication Flight SQL server listening on {}", addr);
-        if let Err(e) = Server::builder().add_service(wrapped_service).serve(addr).await {
+        println!(
+            "Bearer Authentication Flight SQL server listening on {}",
+            addr
+        );
+        if let Err(e) = Server::builder()
+            .layer(tonic_async_interceptor::async_interceptor(
+                bearer_auth_interceptor,
+            ))
+            .add_service(svc)
+            .serve(addr)
+            .await
+        {
             eprintln!("Server error: {}", e);
         }
     });
@@ -169,10 +173,18 @@ async fn main() -> Result<()> { // datafusion::error::Result<()>
     println!("\nAttempting GetTables with valid token (token1)...");
     match new_client_with_auth(client_dsn.clone(), Some("token1".to_string())).await {
         Ok(mut client) => {
-            let request = CommandGetTables { catalog: None, db_schema_filter_pattern: None, table_name_filter_pattern: None, table_types: vec![], include_schema: false };
+            let request = CommandGetTables {
+                catalog: None,
+                db_schema_filter_pattern: None,
+                table_name_filter_pattern: None,
+                table_types: vec![],
+                include_schema: false,
+            };
             match client.get_tables(request).await {
-                Ok(response) => println!("GetTables with token1 SUCCEEDED. Response: {:?}", response.into_inner()),
-                Err(e) => eprintln!("GetTables with token1 FAILED: {}", status_to_df_error(e)),
+                Ok(response) => {
+                    println!("GetTables with token1 SUCCEEDED. Response: {:?}", response)
+                }
+                Err(e) => eprintln!("GetTables with token1 FAILED: {}", e),
             }
         }
         Err(e) => eprintln!("Failed to create client with token1: {}", e),
@@ -182,10 +194,19 @@ async fn main() -> Result<()> { // datafusion::error::Result<()>
     println!("\nAttempting GetTables with invalid token (invalidtoken)...");
     match new_client_with_auth(client_dsn.clone(), Some("invalidtoken".to_string())).await {
         Ok(mut client) => {
-            let request = CommandGetTables { catalog: None, db_schema_filter_pattern: None, table_name_filter_pattern: None, table_types: vec![], include_schema: false };
+            let request = CommandGetTables {
+                catalog: None,
+                db_schema_filter_pattern: None,
+                table_name_filter_pattern: None,
+                table_types: vec![],
+                include_schema: false,
+            };
             match client.get_tables(request).await {
-                Ok(response) => println!("GetTables with invalidtoken SUCCEEDED (unexpected). Response: {:?}", response.into_inner()),
-                Err(e) => eprintln!("GetTables with invalidtoken FAILED (as expected): {}", status_to_df_error(e)),
+                Ok(response) => println!(
+                    "GetTables with invalidtoken SUCCEEDED (unexpected). Response: {:?}",
+                    response
+                ),
+                Err(e) => eprintln!("GetTables with invalidtoken FAILED (as expected): {:?}", e),
             }
         }
         Err(e) => eprintln!("Failed to create client with invalidtoken: {}", e),
@@ -195,10 +216,19 @@ async fn main() -> Result<()> { // datafusion::error::Result<()>
     println!("\nAttempting GetTables with no token...");
     match new_client_with_auth(client_dsn.clone(), None).await {
         Ok(mut client) => {
-            let request = CommandGetTables { catalog: None, db_schema_filter_pattern: None, table_name_filter_pattern: None, table_types: vec![], include_schema: false };
+            let request = CommandGetTables {
+                catalog: None,
+                db_schema_filter_pattern: None,
+                table_name_filter_pattern: None,
+                table_types: vec![],
+                include_schema: false,
+            };
             match client.get_tables(request).await {
-                Ok(response) => println!("GetTables with no token SUCCEEDED (unexpected). Response: {:?}", response.into_inner()),
-                Err(e) => eprintln!("GetTables with no token FAILED (as expected): {}", status_to_df_error(e)),
+                Ok(response) => println!(
+                    "GetTables with no token SUCCEEDED (unexpected). Response: {:?}",
+                    response
+                ),
+                Err(e) => eprintln!("GetTables with no token FAILED (as expected): {:?}", e),
             }
         }
         Err(e) => eprintln!("Failed to create client with no token: {}", e),