Refactor flightsql and execution

Author: matthewmturner

crates/datafusion-app/src/flightsql.rs

@@ -15,26 +15,34 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::sync::Arc;
+
 use arrow_flight::sql::client::FlightSqlServiceClient;
 #[cfg(feature = "flightsql")]
 use base64::engine::{general_purpose::STANDARD, Engine as _};
-use datafusion::sql::parser::DFParser;
+use datafusion::{
+    error::{DataFusionError, Result as DFResult},
+    execution::SendableRecordBatchStream,
+    physical_plan::stream::RecordBatchStreamAdapter,
+    sql::parser::DFParser,
+};
 use log::{error, info, warn};
 
 use color_eyre::eyre::{self, Result};
 use tokio::sync::Mutex;
 use tokio_stream::StreamExt;
 use tonic::{transport::Channel, IntoRequest};
 
-// use crate::config::AppConfig;
 #[cfg(feature = "flightsql")]
 use crate::config::BasicAuth;
 
-use crate::{config::FlightSQLConfig, flightsql_benchmarks::FlightSQLBenchmarkStats};
+use crate::{
+    config::FlightSQLConfig, flightsql_benchmarks::FlightSQLBenchmarkStats, ExecOptions, ExecResult,
+};
 
-pub type FlightSQLClient = Mutex<Option<FlightSqlServiceClient<Channel>>>;
+pub type FlightSQLClient = Arc<Mutex<Option<FlightSqlServiceClient<Channel>>>>;
 
-#[derive(Default)]
+#[derive(Clone, Default)]
 pub struct FlightSQLContext {
     config: FlightSQLConfig,
     flightsql_client: FlightSQLClient,
@@ -44,7 +52,7 @@ impl FlightSQLContext {
     pub fn new(config: FlightSQLConfig) -> Self {
         Self {
             config,
-            flightsql_client: Mutex::new(None),
+            flightsql_client: Arc::new(Mutex::new(None)),
         }
     }
 
@@ -157,4 +165,42 @@ impl FlightSQLContext {
             Err(eyre::eyre!("Only a single statement can be benchmarked"))
         }
     }
+
+    pub async fn execute_sql_with_opts(
+        &self,
+        sql: &str,
+        _opts: ExecOptions,
+    ) -> DFResult<ExecResult> {
+        if let Some(ref mut client) = *self.flightsql_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()));
+            }
+            let endpoint = &flight_info.endpoint[0];
+            if let Some(ticket) = &endpoint.ticket {
+                client
+                    .do_get(ticket.clone().into_request())
+                    .await
+                    .map(|stream| {
+                        if let Some(schema) = stream.schema().cloned() {
+                            let mapped_stream = stream
+                                .map(|res| res.map_err(|e| DataFusionError::External(e.into())));
+                            let batch_stream =
+                                Box::pin(RecordBatchStreamAdapter::new(schema, mapped_stream))
+                                    as SendableRecordBatchStream;
+                            Ok(ExecResult::RecordBatchStream(batch_stream))
+                        } else {
+                            Err(DataFusionError::External(
+                                "Missing schema from stream".into(),
+                            ))
+                        }
+                    })
+                    .map_err(|e| DataFusionError::ArrowError(e, None))?
+            } else {
+                return Err(DataFusionError::External("Missing ticket".into()));
+            }
+        } else {
+            return Err(DataFusionError::External("Missing client".into()));
+        }
+    }
 }

crates/datafusion-app/src/lib.rs

@@ -30,3 +30,21 @@ pub mod stats;
 pub mod wasm;
 
 pub use stats::{collect_plan_io_stats, ExecutionStats};
+
+use datafusion::execution::SendableRecordBatchStream;
+
+pub struct ExecOptions {
+    pub limit: Option<usize>,
+    pub flightsql: bool,
+}
+
+impl ExecOptions {
+    pub fn new(limit: Option<usize>, flightsql: bool) -> Self {
+        Self { limit, flightsql }
+    }
+}
+
+pub enum ExecResult {
+    RecordBatchStream(SendableRecordBatchStream),
+    RecordBatchStreamWithMetrics(()),
+}

crates/datafusion-app/src/local.rs

@@ -27,6 +27,7 @@ use futures::TryFutureExt;
 use log::{debug, error, info};
 
 use crate::config::ExecutionConfig;
+use crate::{ExecOptions, ExecResult};
 use color_eyre::eyre::{self, Result};
 use datafusion::common::Result as DFResult;
 use datafusion::execution::{SendableRecordBatchStream, SessionState};
@@ -405,30 +406,14 @@ impl ExecutionContext {
     pub async fn execute_sql_with_opts(
         &self,
         sql: &str,
-        opts: ExecutionOptions,
-    ) -> DFResult<ExecutionResult> {
+        opts: ExecOptions,
+    ) -> DFResult<ExecResult> {
         let df = self.session_ctx.sql(sql).await?;
         let df = if let Some(limit) = opts.limit {
             df.limit(0, Some(limit))?
         } else {
             df
         };
-        Ok(ExecutionResult::RecordBatchStream(
-            df.execute_stream().await,
-        ))
-    }
-}
-
-pub struct ExecutionOptions {
-    limit: Option<usize>,
-}
-
-impl ExecutionOptions {
-    pub fn new(limit: Option<usize>) -> Self {
-        Self { limit }
+        Ok(ExecResult::RecordBatchStream(df.execute_stream().await?))
     }
 }
-
-pub enum ExecutionResult {
-    RecordBatchStream(DFResult<SendableRecordBatchStream>),
-}

src/execution.rs

@@ -17,48 +17,64 @@
 
 pub use datafusion_app::{collect_plan_io_stats, ExecutionStats};
 
+use color_eyre::Result;
 use datafusion::prelude::*;
 #[cfg(feature = "flightsql")]
 use datafusion_app::flightsql::{FlightSQLClient, FlightSQLContext};
-use datafusion_app::local::ExecutionContext;
+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)]
 pub struct AppExecution {
-    context: ExecutionContext,
+    local: ExecutionContext,
     #[cfg(feature = "flightsql")]
-    flightsql_context: FlightSQLContext,
+    flightsql: FlightSQLContext,
 }
 
 impl AppExecution {
-    pub fn new(context: ExecutionContext) -> Self {
+    pub fn new(local: ExecutionContext) -> Self {
         Self {
-            context,
+            local,
             #[cfg(feature = "flightsql")]
-            flightsql_context: FlightSQLContext::default(),
+            flightsql: FlightSQLContext::default(),
         }
     }
 
     pub fn execution_ctx(&self) -> &ExecutionContext {
-        &self.context
+        &self.local
     }
 
     pub fn session_ctx(&self) -> &SessionContext {
-        self.context.session_ctx()
+        self.local.session_ctx()
     }
 
     #[cfg(feature = "flightsql")]
     pub fn flightsql_client(&self) -> &FlightSQLClient {
-        self.flightsql_context.client()
+        self.flightsql.client()
     }
 
     #[cfg(feature = "flightsql")]
     pub fn flightsql_ctx(&self) -> &FlightSQLContext {
-        &self.flightsql_context
+        &self.flightsql
     }
 
     #[cfg(feature = "flightsql")]
     pub fn with_flightsql_ctx(&mut self, flightsql_ctx: FlightSQLContext) {
-        self.flightsql_context = flightsql_ctx;
+        self.flightsql = flightsql_ctx;
+    }
+
+    pub async fn execute_sql_with_opts(&self, sql: &str, opts: ExecOptions) -> Result<ExecResult> {
+        if cfg!(feature = "flightsql") & opts.flightsql {
+            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())
+        }
     }
 }

src/server/http/mod.rs

@@ -77,8 +77,7 @@ impl HttpApp {
     ) -> Result<Self> {
         info!("Listening to HTTP on {addr}");
         let listener = TcpListener::bind(addr).await.unwrap();
-        let state = execution.execution_ctx().clone();
-        let router = create_router(state, config.http_server);
+        let router = create_router(execution, config.http_server);
 
         let metrics_addr: SocketAddr = metrics_addr.parse()?;
         try_start_metrics_server(metrics_addr)?;

src/server/http/router.rs

@@ -19,35 +19,35 @@ use std::{io::Cursor, time::Duration};
 
 use axum::{
     body::Body,
-    extract::{Path, State},
+    extract::{Json, Path, State},
     response::{IntoResponse, Response},
     routing::{get, post},
     Router,
 };
 use datafusion::{arrow::json::ArrayWriter, execution::SendableRecordBatchStream};
-use datafusion_app::local::{ExecutionContext, ExecutionOptions, ExecutionResult};
+use datafusion_app::{ExecOptions, ExecResult};
 use http::{HeaderValue, StatusCode};
 use log::error;
 use serde::Deserialize;
 use tokio_stream::StreamExt;
 use tower_http::{timeout::TimeoutLayer, trace::TraceLayer};
 use tracing::info;
 
-use crate::config::HttpServerConfig;
+use crate::{config::HttpServerConfig, execution::AppExecution};
 
 #[derive(Clone)]
 struct ExecutionState {
-    execution: ExecutionContext,
+    execution: AppExecution,
     config: HttpServerConfig,
 }
 
 impl ExecutionState {
-    pub fn new(execution: ExecutionContext, config: HttpServerConfig) -> Self {
+    pub fn new(execution: AppExecution, config: HttpServerConfig) -> Self {
         Self { execution, config }
     }
 }
 
-pub fn create_router(execution: ExecutionContext, config: HttpServerConfig) -> Router {
+pub fn create_router(execution: AppExecution, config: HttpServerConfig) -> Router {
     let state = ExecutionState::new(execution, config);
     Router::new()
         .route(
@@ -70,14 +70,22 @@ pub fn create_router(execution: ExecutionContext, config: HttpServerConfig) -> R
         .with_state(state)
 }
 
-async fn post_sql_handler(state: State<ExecutionState>, query: String) -> Response {
-    let opts = ExecutionOptions::new(Some(state.config.result_limit));
-    execute_sql_with_opts(state, query, opts).await
+#[derive(Deserialize)]
+struct PostSqlBody {
+    query: String,
+    #[serde(default)]
+    flightsql: bool,
+}
+
+async fn post_sql_handler(state: State<ExecutionState>, Json(body): Json<PostSqlBody>) -> Response {
+    let opts = ExecOptions::new(Some(state.config.result_limit), body.flightsql);
+    execute_sql_with_opts(state, body.query, opts).await
 }
 
 async fn get_catalog_handler(state: State<ExecutionState>) -> Response {
-    let opts = ExecutionOptions::new(None);
-    execute_sql_with_opts(state, "SHOW TABLES".to_string(), opts).await
+    let opts = ExecOptions::new(None, false);
+    let sql = "SHOW TABLES".to_string();
+    execute_sql_with_opts(state, sql, opts).await
 }
 
 #[derive(Deserialize)]
@@ -97,29 +105,26 @@ async fn get_table_handler(
         table,
     } = params;
     let sql = format!("SELECT * FROM \"{catalog}\".\"{schema}\".\"{table}\"");
-    let opts = ExecutionOptions::new(Some(state.config.result_limit));
+    let opts = ExecOptions::new(Some(state.config.result_limit), false);
     execute_sql_with_opts(state, sql, opts).await
 }
 
+// TODO: Maybe rename to something like `response_for_sql`
 async fn execute_sql_with_opts(
     State(state): State<ExecutionState>,
     sql: String,
-    opts: ExecutionOptions,
+    opts: ExecOptions,
 ) -> Response {
     info!("Executing sql: {sql}");
-    let results = state.execution.execute_sql_with_opts(&sql, opts).await;
-    match results {
-        Ok(ExecutionResult::RecordBatchStream(Ok(batch_stream))) => {
-            batch_stream_to_response(batch_stream).await
-        }
-        Err(e) | Ok(ExecutionResult::RecordBatchStream(Err(e))) => {
-            error!("Error executing SQL: {}", e);
-            (
-                StatusCode::BAD_REQUEST,
-                format!("SQL execution failed: {}", e),
-            )
-                .into_response()
-        }
+    match state.execution.execute_sql_with_opts(&sql, opts).await {
+        Ok(ExecResult::RecordBatchStream(stream)) => batch_stream_to_response(stream).await,
+        Ok(_) => (
+            StatusCode::BAD_REQUEST,
+            "Execution failed: unknown result type".to_string(),
+        )
+            .into_response(),
+
+        Err(e) => (StatusCode::BAD_REQUEST, format!("Execution failed: {}", e)).into_response(),
     }
 }