FlightSQL for HTTP refactoring

Author: matthewmturner

crates/datafusion-app/src/config.rs

@@ -239,7 +239,7 @@ pub struct WasmUdfConfig {
 }
 
 #[cfg(feature = "flightsql")]
-#[derive(Clone, Default)]
+#[derive(Clone, Debug, Default)]
 pub struct FlightSQLConfig {
     pub connection_url: String,
     pub benchmark_iterations: usize,

crates/datafusion-app/src/flightsql.rs

@@ -42,24 +42,25 @@ use crate::{
 
 pub type FlightSQLClient = Arc<Mutex<Option<FlightSqlServiceClient<Channel>>>>;
 
-#[derive(Clone, Default)]
+#[derive(Clone, Debug, Default)]
 pub struct FlightSQLContext {
     config: FlightSQLConfig,
-    flightsql_client: FlightSQLClient,
+    client: FlightSQLClient,
 }
 
 impl FlightSQLContext {
     pub fn new(config: FlightSQLConfig) -> Self {
         Self {
             config,
-            flightsql_client: Arc::new(Mutex::new(None)),
+            client: Arc::new(Mutex::new(None)),
         }
     }
 
     pub fn client(&self) -> &FlightSQLClient {
-        &self.flightsql_client
+        &self.client
     }
 
+    // TODO - Make this part of `new` method
     /// Create FlightSQL client from users FlightSQL config
     pub async fn create_client(&self, cli_host: Option<String>) -> Result<()> {
         let final_url = cli_host.unwrap_or(self.config.connection_url.clone());
@@ -74,6 +75,7 @@ impl FlightSQLContext {
                 //
                 // Although that is for HTTP/1.1 and GRPC uses HTTP/2 - so maybe it has changed.
                 // To be tested later with the Tower auth layers to see what they support.
+                // TODO - Do we need this feature block?
                 #[cfg(feature = "flightsql")]
                 {
                     if let Some(token) = &self.config.auth.bearer_token {
@@ -85,7 +87,7 @@ impl FlightSQLContext {
                         client.set_header("Authorization", format!("Basic {encoded_basic}"))
                     }
                 }
-                let mut guard = self.flightsql_client.lock().await;
+                let mut guard = self.client.lock().await;
                 *guard = Some(client);
                 Ok(())
             }
@@ -111,7 +113,7 @@ impl FlightSQLContext {
         let dialect = datafusion::sql::sqlparser::dialect::GenericDialect {};
         let statements = DFParser::parse_sql_with_dialect(query, &dialect)?;
         if statements.len() == 1 {
-            if let Some(ref mut client) = *self.flightsql_client.lock().await {
+            if let Some(ref mut client) = *self.client.lock().await {
                 for _ in 0..iterations {
                     let mut rows = 0;
                     let start = std::time::Instant::now();
@@ -171,7 +173,7 @@ impl FlightSQLContext {
         sql: &str,
         _opts: ExecOptions,
     ) -> DFResult<ExecResult> {
-        if let Some(ref mut client) = *self.flightsql_client.lock().await {
+        if let Some(ref mut client) = *self.client.lock().await {
             let flight_info = client.execute(sql.to_string(), None).await?;
             if flight_info.endpoint.len() != 1 {
                 return Err(DataFusionError::External("More than one endpoint".into()));

justfile

@@ -8,7 +8,7 @@ install-tools:
 
 # Starts a debug HTTP server
 serve-http:
-    RUST_LOG=info cargo r --features=http -- serve-http   
+    RUST_LOG=info cargo r --features=http,flightsql -- serve-http   
 
 # Starts a debug FlightSQL server
 serve-flight-sql:

src/execution.rs

@@ -25,7 +25,7 @@ use datafusion_app::{local::ExecutionContext, ExecOptions, ExecResult};
 
 /// Provides all core execution functionality for execution queries from either a local
 /// `SessionContext` or a remote `FlightSQL` service
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct AppExecution {
     local: ExecutionContext,
     #[cfg(feature = "flightsql")]
@@ -65,16 +65,19 @@ impl AppExecution {
     }
 
     pub async fn execute_sql_with_opts(&self, sql: &str, opts: ExecOptions) -> Result<ExecResult> {
-        if cfg!(feature = "flightsql") & opts.flightsql {
-            self.flightsql
+        #[cfg(feature = "flightsql")]
+        if opts.flightsql {
+            return self
+                .flightsql
                 .execute_sql_with_opts(sql, opts)
                 .await
-                .map_err(|e| e.into())
-        } else {
-            self.local
-                .execute_sql_with_opts(sql, opts)
-                .await
-                .map_err(|e| e.into())
+                .map_err(|e| e.into());
         }
+
+        // If flightsql is not enabled or `opts.flightsql` is false, fall back to local:
+        self.local
+            .execute_sql_with_opts(sql, opts)
+            .await
+            .map_err(|e| e.into())
     }
 }

src/server/http/mod.rs

@@ -23,7 +23,10 @@ use crate::{args::DftArgs, config::AppConfig, execution::AppExecution};
 use axum::Router;
 use color_eyre::Result;
 use datafusion_app::{
-    config::merge_configs, extensions::DftSessionStateBuilder, local::ExecutionContext,
+    config::{merge_configs, AuthConfig, FlightSQLConfig},
+    extensions::DftSessionStateBuilder,
+    flightsql::FlightSQLContext,
+    local::ExecutionContext,
 };
 use router::create_router;
 use tokio::{net::TcpListener, signal};
@@ -112,7 +115,28 @@ pub async fn try_run(cli: DftArgs, config: AppConfig) -> Result<()> {
     if cli.run_ddl {
         execution_ctx.execute_ddl().await;
     }
-    let app_execution = AppExecution::new(execution_ctx);
+    let mut app_execution = AppExecution::new(execution_ctx);
+    #[cfg(feature = "flightsql")]
+    {
+        info!("Setting up FlightSQLContext");
+        let auth = AuthConfig {
+            basic_auth: config.flightsql_client.auth.basic_auth.clone(),
+            bearer_token: config.flightsql_client.auth.bearer_token.clone(),
+        };
+        let flightsql_cfg = FlightSQLConfig::new(
+            config.flightsql_client.connection_url.clone(),
+            config.flightsql_client.benchmark_iterations,
+            auth,
+        );
+
+        let flightsql_context = FlightSQLContext::new(flightsql_cfg.clone());
+        // TODO - Consider adding flag to allow startup even if FlightSQL initiation fails
+        flightsql_context
+            .create_client(Some(flightsql_cfg.connection_url))
+            .await?;
+        app_execution.with_flightsql_ctx(flightsql_context);
+    }
+    info!("Created AppExecution: {app_execution:?}");
     let app = HttpApp::try_new(
         app_execution,
         config.clone(),

src/server/http/router.rs

@@ -78,6 +78,13 @@ struct PostSqlBody {
 }
 
 async fn post_sql_handler(state: State<ExecutionState>, Json(body): Json<PostSqlBody>) -> Response {
+    if body.flightsql && !cfg!(feature = "flightsql") {
+        return (
+            StatusCode::BAD_REQUEST,
+            "FlightSQL is not enabled on this server",
+        )
+            .into_response();
+    }
     let opts = ExecOptions::new(Some(state.config.result_limit), body.flightsql);
     execute_sql_with_opts(state, body.query, opts).await
 }