New crate datafusion-flightsql (spiceai#10201)

* add datafusion-flightsql

* formatting

* clippy

* implement do_put; handle handshake/sessions

Author: Jeadie

Cargo.lock

@@ -35,6 +35,7 @@ members = [
     "crates/runtime-api-types",
     "crates/runtime-auth",
     "crates/runtime-async",
+    "crates/datafusion-flightsql",
     "crates/runtime-datafusion",
     "crates/runtime-datafusion-index",
     "crates/runtime-datafusion-udfs",

Cargo.toml

@@ -0,0 +1,37 @@
+[package]
+name = "datafusion-flightsql"
+version = "0.1.0"
+edition.workspace = true
+rust-version = "1.93.1"
+license = "Apache-2.0"
+description = "Arrow Flight SQL service backed by an arbitrary DataFusion SessionContext"
+
+[dependencies]
+arrow = { workspace = true }
+arrow-flight = { workspace = true, features = ["flight-sql-experimental"] }
+arrow-ipc = { workspace = true }
+arrow-schema = { workspace = true }
+arrow_tools = { path = "../arrow_tools" }
+async-stream = { workspace = true }
+bytes = { workspace = true }
+datafusion = { workspace = true }
+futures = { workspace = true }
+http = { workspace = true }
+moka = { workspace = true, features = ["sync"] }
+postcard = { workspace = true }
+prost = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
+snafu = { workspace = true }
+tokio-stream = { workspace = true }
+tonic = { workspace = true }
+tracing = { workspace = true }
+uuid = { workspace = true, features = ["v7"] }
+util = { path = "../util" }
+
+[dev-dependencies]
+tokio = { workspace = true }
+tracing-subscriber = { workspace = true }
+
+[[example]]
+name = "serve_and_query"
+required-features = []

crates/datafusion-flightsql/Cargo.toml

@@ -0,0 +1,131 @@
+/*
+Copyright 2026 The Spice.ai OSS Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+//! Starts a Flight SQL server backed by an in-memory `DataFusion`
+//! [`SessionContext`] and blocks until Ctrl-C.
+//!
+//! ```text
+//! cargo run -p datafusion-flightsql --example serve_and_query
+//! ```
+//!
+//! Once running, connect with any Arrow Flight SQL client.  For example,
+//! using the Python `flightsql-dbapi` library:
+//!
+//! ```python
+//! from flightsql import FlightSQLClient
+//! client = FlightSQLClient(host="localhost", port=50051, insecure=True)
+//! reader = client.execute("SELECT * FROM employees ORDER BY salary DESC")
+//! print(reader.read_all().to_pandas())
+//! ```
+//!
+//! Or with the `arrow-flight` Rust client:
+//!
+//! ```rust,no_run
+//! use arrow_flight::sql::client::FlightSqlServiceClient;
+//! use futures::TryStreamExt;
+//! use tonic::transport::Channel;
+//!
+//! #[tokio::main]
+//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
+//!     let channel = Channel::from_static("http://localhost:50051").connect().await?;
+//!     let mut client = FlightSqlServiceClient::new(channel);
+//!     client.handshake("", "").await?;
+//!
+//!     let mut stmt = client.prepare("SELECT * FROM employees".to_string(), None).await?;
+//!     let info = stmt.execute().await?;
+//!     for endpoint in info.endpoint {
+//!         if let Some(ticket) = endpoint.ticket {
+//!             let batches: Vec<_> = client.do_get(ticket).await?.try_collect().await?;
+//!             println!("{}", arrow::util::pretty::pretty_format_batches(&batches)?);
+//!         }
+//!     }
+//!     Ok(())
+//! }
+//! ```
+
+use std::sync::Arc;
+
+use arrow::array::{Int32Array, StringArray};
+use arrow::datatypes::{DataType, Field, Schema};
+use arrow::record_batch::RecordBatch;
+use datafusion::prelude::SessionContext;
+use datafusion_flightsql::FlightSqlService;
+use tonic::transport::Server;
+
+const ADDR: &str = "0.0.0.0:50051";
+
+#[tokio::main]
+async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    tracing_subscriber::fmt()
+        .with_env_filter("datafusion_flightsql=debug,info")
+        .init();
+
+    // ── Build the SessionContext ──────────────────────────────────────────────
+
+    let ctx = Arc::new(SessionContext::new());
+
+    let batch = RecordBatch::try_new(
+        Arc::new(Schema::new(vec![
+            Field::new("id", DataType::Int32, false),
+            Field::new("name", DataType::Utf8, false),
+            Field::new("department", DataType::Utf8, false),
+            Field::new("salary", DataType::Int32, false),
+        ])),
+        vec![
+            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])),
+            Arc::new(StringArray::from(vec![
+                "Alice", "Bob", "Carol", "Dave", "Eve",
+            ])),
+            Arc::new(StringArray::from(vec![
+                "Engineering",
+                "Marketing",
+                "Engineering",
+                "HR",
+                "Engineering",
+            ])),
+            Arc::new(Int32Array::from(vec![
+                120_000, 85_000, 110_000, 75_000, 130_000,
+            ])),
+        ],
+    )?;
+
+    ctx.register_batch("employees", batch)?;
+    tracing::info!("Registered 'employees' table");
+
+    // You can register additional tables, views, or UDFs on `ctx` here before
+    // the server starts.  Everything registered on the context is immediately
+    // queryable by connected clients.
+
+    // ── Start the Flight SQL server ───────────────────────────────────────────
+
+    let addr = ADDR.parse()?;
+    tracing::info!("Flight SQL server listening on {addr}  (press Ctrl-C to stop)");
+
+    Server::builder()
+        .add_service(FlightSqlService::new(ctx).into_server())
+        .serve_with_shutdown(addr, shutdown_signal())
+        .await?;
+
+    tracing::info!("Server stopped");
+    Ok(())
+}
+
+async fn shutdown_signal() {
+    tokio::signal::ctrl_c()
+        .await
+        .expect("failed to listen for Ctrl-C");
+    tracing::info!("Received Ctrl-C, shutting down");
+}

crates/datafusion-flightsql/examples/serve_and_query.rs

@@ -0,0 +1,86 @@
+/*
+Copyright 2026 The Spice.ai OSS Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+//! `DoAction` / `ListActions` for Flight SQL prepared statements.
+//!
+//! Supported actions:
+//! - `CreatePreparedStatement`
+//! - `ClosePreparedStatement`
+
+use std::sync::Arc;
+
+use arrow_flight::{
+    Action, ActionType,
+    flight_service_server::FlightService,
+    sql::{self, Any, ProstMessageExt},
+};
+use datafusion::prelude::SessionContext;
+use futures::stream::BoxStream;
+use prost::Message;
+use tonic::{Request, Response, Status};
+
+use crate::{FlightSqlService, flightsql::prepared_statement_query, to_tonic_err};
+
+pub(crate) async fn do_action(
+    ctx: Arc<SessionContext>,
+    request: Request<Action>,
+) -> Result<Response<<FlightSqlService as FlightService>::DoActionStream>, Status> {
+    let action = request.into_inner();
+    tracing::trace!("do_action: type={}", action.r#type);
+
+    let msg: Any = Any::decode(&*action.body).map_err(to_tonic_err)?;
+
+    match sql::Command::try_from(msg).map_err(to_tonic_err)? {
+        sql::Command::ActionCreatePreparedStatementRequest(stmt) => {
+            let result =
+                prepared_statement_query::do_action_create_prepared_statement(ctx, stmt).await?;
+            let output = futures::stream::iter(vec![Ok(arrow_flight::Result {
+                body: result.as_any().encode_to_vec().into(),
+            })]);
+            Ok(Response::new(
+                Box::pin(output) as <FlightSqlService as FlightService>::DoActionStream
+            ))
+        }
+        sql::Command::ActionClosePreparedStatementRequest(handle) => {
+            tracing::trace!("close_prepared_statement: {:?}", handle);
+            // Prepared statement handles are self-contained serialised blobs;
+            // there is nothing server-side to clean up.
+            let output = futures::stream::empty();
+            Ok(Response::new(
+                Box::pin(output) as <FlightSqlService as FlightService>::DoActionStream
+            ))
+        }
+        cmd => Err(Status::invalid_argument(format!(
+            "unsupported action command: {cmd:?}"
+        ))),
+    }
+}
+
+pub(crate) fn list() -> Response<BoxStream<'static, Result<ActionType, Status>>> {
+    let actions = vec![
+        ActionType {
+            r#type: "CreatePreparedStatement".to_string(),
+            description: "Creates a reusable prepared statement resource on the server. \
+                          Returns a handle identifying the prepared statement."
+                .to_string(),
+        },
+        ActionType {
+            r#type: "ClosePreparedStatement".to_string(),
+            description: "Closes a prepared statement resource on the server.".to_string(),
+        },
+    ];
+    Response::new(Box::pin(futures::stream::iter(actions.into_iter().map(Ok))))
+}

crates/datafusion-flightsql/src/actions.rs

@@ -0,0 +1,77 @@
+/*
+Copyright 2026 The Spice.ai OSS Authors
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+     https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+use std::sync::Arc;
+
+use arrow_flight::{
+    Ticket,
+    flight_service_server::FlightService,
+    sql::{Any, Command},
+};
+use datafusion::prelude::SessionContext;
+use prost::Message;
+use tonic::{Request, Response, Status};
+
+use crate::{FlightSqlService, flightsql, to_tonic_err};
+
+pub(crate) async fn handle(
+    ctx: Arc<SessionContext>,
+    request: Request<Ticket>,
+) -> Result<Response<<FlightSqlService as FlightService>::DoGetStream>, Status> {
+    let msg: Any = match Message::decode(&*request.get_ref().ticket) {
+        Ok(msg) => msg,
+        Err(_) => return do_get_simple(ctx, request).await,
+    };
+
+    match Command::try_from(msg).map_err(to_tonic_err)? {
+        Command::CommandStatementQuery(cmd) => flightsql::statement_query::do_get(ctx, cmd).await,
+        Command::CommandPreparedStatementQuery(cmd) => {
+            Box::pin(flightsql::prepared_statement_query::do_get(ctx, cmd)).await
+        }
+        Command::CommandPreparedStatementUpdate(cmd) => {
+            Box::pin(flightsql::prepared_statement_update::do_get(ctx, cmd)).await
+        }
+        Command::CommandGetCatalogs(cmd) => flightsql::get_catalogs::do_get(&ctx, cmd),
+        Command::CommandGetDbSchemas(cmd) => flightsql::get_schemas::do_get(&ctx, cmd),
+        Command::CommandGetTables(cmd) => flightsql::get_tables::do_get(ctx, cmd).await,
+        Command::CommandGetPrimaryKeys(cmd) => Ok(flightsql::get_primary_keys::do_get(&cmd)),
+        Command::CommandGetTableTypes(cmd) => flightsql::get_table_types::do_get(cmd),
+        Command::CommandGetSqlInfo(cmd) => flightsql::get_sql_info::do_get(cmd),
+        Command::CommandGetXdbcTypeInfo(cmd) => flightsql::get_xdbc_type_info::do_get(cmd),
+        Command::CommandStatementUpdate(_) => Err(Status::invalid_argument(
+            "CommandStatementUpdate should be sent via DoPut, not DoGet",
+        )),
+        cmd => Err(Status::unimplemented(format!(
+            "DoGet not implemented for command: {cmd:?}"
+        ))),
+    }
+}
+
+async fn do_get_simple(
+    ctx: Arc<SessionContext>,
+    request: Request<Ticket>,
+) -> Result<Response<<FlightSqlService as FlightService>::DoGetStream>, Status> {
+    let ticket = request.into_inner();
+    let sql = std::str::from_utf8(&ticket.ticket)
+        .map_err(|e| Status::invalid_argument(format!("invalid ticket: {e}")))?;
+
+    tracing::trace!("do_get (plain SQL): {sql}");
+
+    let output = FlightSqlService::sql_to_flight_stream(ctx, sql, None).await?;
+    Ok(Response::new(
+        Box::pin(output) as <FlightSqlService as FlightService>::DoGetStream
+    ))
+}