skade-flight/src/service.rs
//! The Flight SQL service: SQL in, streamed Arrow out.
//!
//! Protocol shape for a generic driver (Power BI's connector, the Flight SQL
//! ODBC driver, `adbc_driver_flightsql`, DataFusion's client):
//!
//! 1. `Handshake` with Basic → bearer token (replayed on every call).
//! 2. `GetSqlInfo`, `GetXdbcTypeInfo`, `GetCatalogs`, `GetDbSchemas`,
//! `GetTables`, `GetTableTypes` — the introspection surface.
//! 3. `GetFlightInfo(CommandStatementQuery)` — we PLAN the statement, logical
//! and physical, once (no execution) and answer with the result schema and
//! a ticket: a signed, expiring handle bound to the caller (ticket.rs).
//! 4. `DoGet(ticket)` — verify the handle, execute the plan cached at step 3,
//! and stream the partitions through the Flight encoder. Nothing is
//! planned again and nothing is collected.
//!
//! Prepared statements are stateless: the handle carries the statement (SQL +
//! bound parameters) under the same signature, so any replica with the same
//! secret can serve any handle and there is no server-side statement table.
//! `DoPutPreparedStatementQuery` binds parameters by answering with a NEW
//! handle that carries them.
//!
//! The ticket is an `Any`-packed `TicketStatementQuery` (arrow-flight's
//! `do_get` dispatches on `Any.type_url`; a bare encode mis-decodes).
use std::pin::Pin;
use std::sync::{Arc, LazyLock};
use arrow_flight::encode::FlightDataEncoderBuilder;
use arrow_flight::error::FlightError;
use arrow_flight::sql::metadata::{SqlInfoData, XdbcTypeInfoData};
use arrow_flight::sql::server::{FlightSqlService, PeekableFlightDataStream};
use arrow_flight::sql::{
ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
ActionCreatePreparedStatementResult, Any, CommandGetCatalogs, CommandGetCrossReference, CommandGetDbSchemas,
CommandGetExportedKeys, CommandGetImportedKeys, CommandGetPrimaryKeys, CommandGetSqlInfo,
CommandGetTableTypes, CommandGetTables, CommandGetXdbcTypeInfo, CommandPreparedStatementQuery,
CommandStatementQuery, CommandStatementUpdate, DoPutPreparedStatementResult, ProstMessageExt, SqlInfo,
TicketStatementQuery,
};
use arrow_flight::{
FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest, HandshakeResponse, Ticket,
};
use futures::{Stream, TryStreamExt};
use prost::Message;
use tonic::{Request, Response, Status, Streaming};
use crate::auth::Auth;
use crate::engine::{CATALOG, Engine, Statement};
use crate::keys::{FOREIGN_KEYS_SCHEMA, PRIMARY_KEYS_SCHEMA, no_foreign_keys, primary_keys};
use crate::ticket::{Claims, PREPARED_TTL, Signer, TICKET_TTL, principal_of_header};
use crate::sqlinfo::{CATALOGS_FRAMES, SQL_INFO, SQL_INFO_FRAMES, TABLE_TYPES_FRAMES, XDBC_TYPE_INFO, XDBC_TYPE_INFO_FRAMES};
type BoxStream = Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
/// The Flight SQL service over one [`Engine`].
#[derive(Clone)]
pub struct SkadeFlightService {
engine: Arc<Engine>,
auth: Arc<Auth>,
/// Mints and verifies statement handles (tickets, prepared handles).
signer: Signer,
sql_info: &'static SqlInfoData,
type_info: &'static XdbcTypeInfoData,
}
impl SkadeFlightService {
pub fn new(engine: Arc<Engine>, auth: Arc<Auth>) -> Self {
// Force the static answers now, not on the first client's call.
LazyLock::force(&SQL_INFO_FRAMES);
LazyLock::force(&XDBC_TYPE_INFO_FRAMES);
LazyLock::force(&TABLE_TYPES_FRAMES);
LazyLock::force(&CATALOGS_FRAMES);
let signer = Signer::for_secret(auth.issued_token());
Self { engine, auth, signer, sql_info: &SQL_INFO, type_info: &XDBC_TYPE_INFO }
}
/// The signer the handles this service mints are verified with (tests
/// forge and tamper with it).
pub fn signer(&self) -> &Signer {
&self.signer
}
pub fn engine(&self) -> &Arc<Engine> {
&self.engine
}
pub fn auth(&self) -> &Arc<Auth> {
&self.auth
}
/// Plan a statement — logical and physical, once — and describe it:
/// schema + a signed ticket bound to `principal`. No execution.
async fn describe(&self, stmt: Statement, principal: &str) -> Result<FlightInfo, Status> {
self.engine.maybe_refresh().await.map_err(Status::from)?;
let planned = self.engine.describe(&stmt).await;
crate::functional_status("skade-flight/get_flight_info_statement", "planned", planned.is_ok(), &stmt.sql);
let planned = planned.map_err(Status::from)?;
// The IPC schema was encoded once when the plan was cached; this is a
// refcount bump, not a re-encode.
let schema_ipc = planned.schema_ipc.clone();
let handle = self.signer.issue(stmt, principal, TICKET_TTL)?;
let ticket = TicketStatementQuery { statement_handle: handle }.as_any().encode_to_vec();
let mut info = FlightInfo::new()
.with_endpoint(FlightEndpoint::new().with_ticket(Ticket::new(ticket)))
.with_total_records(-1)
.with_total_bytes(-1)
.with_ordered(false);
info.schema = schema_ipc;
Ok(info)
}
/// Execute a described statement and stream it out: one encoder task per
/// partition, frames merged through a channel (see encode.rs). The plan is
/// the one `describe` cached; the schema is read off it, not re-planned.
async fn stream(&self, stmt: &Statement) -> Result<Response<BoxStream>, Status> {
let planned = self.engine.describe(stmt).await.map_err(Status::from)?;
let started = self.engine.execute_planned(&planned).await;
crate::functional_status("skade-flight/do_get_statement", "stream_started", started.is_ok(), &stmt.sql);
let partitions = started.map_err(Status::from)?;
Ok(Response::new(crate::encode::encode_partitions(planned.schema.clone(), partitions)?))
}
/// Verify a handle against the caller it is presented by.
fn claims<T>(&self, handle: &[u8], request: &Request<T>) -> Result<Claims, Status> {
self.signer.verify(handle, &principal_of(request))
}
}
/// The caller's principal, off the request's `authorization` header.
fn principal_of<T>(request: &Request<T>) -> String {
principal_of_header(request.metadata().get("authorization").and_then(|v| v.to_str().ok()))
}
fn metadata_info(
descriptor: FlightDescriptor,
ticket: Ticket,
schema: arrow_schema::SchemaRef,
) -> Result<Response<FlightInfo>, Status> {
let info = FlightInfo::new()
.try_with_schema(schema.as_ref())
.map_err(|e| Status::internal(e.to_string()))?
.with_endpoint(FlightEndpoint::new().with_ticket(ticket))
.with_descriptor(descriptor);
Ok(Response::new(info))
}
fn metadata_stream(
schema: arrow_schema::SchemaRef,
batch: Result<arrow_array::RecordBatch, FlightError>,
) -> BoxStream {
Box::pin(
FlightDataEncoderBuilder::new()
.with_schema(schema)
.build(futures::stream::once(async move { batch }))
.map_err(|e: FlightError| Status::internal(e.to_string())),
)
}
#[tonic::async_trait]
impl FlightSqlService for SkadeFlightService {
type FlightService = SkadeFlightService;
async fn register_sql_info(&self, _id: i32, _result: &SqlInfo) {}
// ── Handshake: Basic (or bearer) in, bearer out ─────────────────────
async fn do_handshake(
&self,
request: Request<Streaming<HandshakeRequest>>,
) -> Result<Response<Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send>>>, Status> {
let ok = self.auth.check_request(&request);
crate::functional_status("skade-flight/handshake", "credential_accepted", ok.is_ok(), "");
ok?;
// Bearer in the header, NO message — the Dremio ODBC driver hangs on
// one. Why, measured, lives with the one copy.
nornir_flight::door::handshake_response(self.auth.issued_token())
}
// ── Statements ──────────────────────────────────────────────────────
async fn get_flight_info_statement(
&self,
query: CommandStatementQuery,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let principal = principal_of(&request);
let info = self.describe(Statement::sql(query.query), &principal).await?;
Ok(Response::new(info.with_descriptor(request.into_inner())))
}
async fn do_get_statement(
&self,
ticket: TicketStatementQuery,
request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
let claims = self.claims(&ticket.statement_handle, &request)?;
self.stream(&claims.statement).await
}
// ── Prepared statements: the handle carries the statement, signed ────
async fn do_action_create_prepared_statement(
&self,
query: ActionCreatePreparedStatementRequest,
request: Request<arrow_flight::Action>,
) -> Result<ActionCreatePreparedStatementResult, Status> {
self.engine.maybe_refresh().await.map_err(Status::from)?;
let (schema, params) = self.engine.prepared_shape(&query.query).await.map_err(Status::from)?;
let ipc_opts = arrow_ipc_options();
let dataset_schema: arrow_flight::IpcMessage = arrow_flight::SchemaAsIpc::new(schema.as_ref(), &ipc_opts)
.try_into()
.map_err(|e: arrow_schema::ArrowError| Status::internal(e.to_string()))?;
// `parameter_schema`: one field per placeholder, in `$1, $2 …` order,
// typed as the planner inferred; a placeholder whose type the plan
// cannot pin (e.g. `SELECT $1`) is advertised as Utf8 rather than
// omitted, so the driver still knows how many to bind.
let param_fields: Vec<arrow_schema::Field> = params
.iter()
.map(|(name, ty)| {
arrow_schema::Field::new(name.clone(), ty.clone().unwrap_or(arrow_schema::DataType::Utf8), true)
})
.collect();
let parameter_schema: bytes::Bytes = if param_fields.is_empty() {
Default::default()
} else {
let ps = arrow_schema::Schema::new(param_fields);
let m: arrow_flight::IpcMessage = arrow_flight::SchemaAsIpc::new(&ps, &ipc_opts)
.try_into()
.map_err(|e: arrow_schema::ArrowError| Status::internal(e.to_string()))?;
m.0
};
let handle = self.signer.issue(Statement::sql(query.query), &principal_of(&request), PREPARED_TTL)?;
Ok(ActionCreatePreparedStatementResult {
prepared_statement_handle: handle,
dataset_schema: dataset_schema.0,
parameter_schema,
})
}
async fn do_action_close_prepared_statement(
&self,
_query: ActionClosePreparedStatementRequest,
_request: Request<arrow_flight::Action>,
) -> Result<(), Status> {
Ok(())
}
async fn get_flight_info_prepared_statement(
&self,
cmd: CommandPreparedStatementQuery,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let claims = self.claims(&cmd.prepared_statement_handle, &request)?;
let principal = claims.principal.clone();
let info = self.describe(claims.statement, &principal).await?;
Ok(Response::new(info.with_descriptor(request.into_inner())))
}
async fn do_get_prepared_statement(
&self,
cmd: CommandPreparedStatementQuery,
request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
let claims = self.claims(&cmd.prepared_statement_handle, &request)?;
self.stream(&claims.statement).await
}
/// **Parameter binding.** The bound values arrive as one Arrow batch (one
/// row, one column per placeholder); the answer is a NEW handle carrying
/// them, which the client uses from here on (arrow-rs, arrow-cpp and the
/// ODBC driver all adopt the handle a DoPut result returns). The ODBC
/// driver also sends an EMPTY DoPut before a parameter-less query: that
/// binds nothing and the handle is unchanged.
async fn do_put_prepared_statement_query(
&self,
cmd: CommandPreparedStatementQuery,
request: Request<PeekableFlightDataStream>,
) -> Result<DoPutPreparedStatementResult, Status> {
let claims = self.claims(&cmd.prepared_statement_handle, &request)?;
let principal = claims.principal.clone();
let stream = arrow_flight::decode::FlightRecordBatchStream::new_from_flight_data(
request.into_inner().map_err(|e| FlightError::Tonic(Box::new(e))),
);
let batches: Vec<arrow_array::RecordBatch> =
stream.try_collect().await.map_err(|e| Status::invalid_argument(format!("bind: {e}")))?;
let mut params = Vec::new();
for b in &batches {
params.extend(crate::ticket::params_from_batch(b)?);
}
if params.is_empty() {
return Ok(DoPutPreparedStatementResult::default());
}
let stmt = Statement { sql: claims.statement.sql, params };
let handle = self.signer.issue(stmt, &principal, PREPARED_TTL)?;
Ok(DoPutPreparedStatementResult { prepared_statement_handle: Some(handle) })
}
/// The update door. A session-view DDL (`CREATE VIEW` / `DROP VIEW`,
/// views.rs) is applied and reports 0 rows; every real write is refused
/// by the engine with the rewrite in the message — INVALID_ARGUMENT, since
/// it is the statement that is out of scope, not the server.
async fn do_put_statement_update(
&self,
ticket: CommandStatementUpdate,
_request: Request<PeekableFlightDataStream>,
) -> Result<i64, Status> {
self.engine.maybe_refresh().await.map_err(Status::from)?;
let applied = self.engine.execute_update(&ticket.query).await;
crate::functional_status("skade-flight/do_put_statement_update", "applied", applied.is_ok(), &ticket.query);
applied.map_err(Status::from)
}
// ── Introspection ───────────────────────────────────────────────────
async fn get_flight_info_catalogs(
&self,
query: CommandGetCatalogs,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(request.into_inner(), ticket, query.into_builder().schema())
}
async fn get_flight_info_schemas(
&self,
query: CommandGetDbSchemas,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(request.into_inner(), ticket, query.into_builder().schema())
}
async fn get_flight_info_tables(
&self,
query: CommandGetTables,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(request.into_inner(), ticket, query.into_builder().schema())
}
async fn get_flight_info_table_types(
&self,
query: CommandGetTableTypes,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(request.into_inner(), ticket, query.into_builder().schema())
}
async fn get_flight_info_sql_info(
&self,
query: CommandGetSqlInfo,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(request.into_inner(), ticket, query.into_builder(self.sql_info).schema())
}
async fn get_flight_info_xdbc_type_info(
&self,
query: CommandGetXdbcTypeInfo,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(request.into_inner(), ticket, query.into_builder(self.type_info).schema())
}
async fn do_get_catalogs(
&self,
_query: CommandGetCatalogs,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
Ok(Response::new(CATALOGS_FRAMES.stream()))
}
async fn do_get_schemas(
&self,
query: CommandGetDbSchemas,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
self.engine.maybe_refresh().await.map_err(Status::from)?;
let mut b = query.into_builder();
for ns in self.engine.schema_names() {
b.append(CATALOG, ns);
}
let schema = b.schema();
Ok(Response::new(metadata_stream(schema, b.build())))
}
async fn do_get_tables(
&self,
query: CommandGetTables,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
self.engine.maybe_refresh().await.map_err(Status::from)?;
let tables = self.engine.tables().await.map_err(Status::from)?;
#[cfg(feature = "testmatrix")]
crate::functional_status("skade-flight/get_tables", "listed", true, &format!("{} table(s)", tables.len()));
let mut b = query.into_builder();
for t in tables {
b.append(CATALOG, &t.schema, &t.name, "TABLE", t.arrow.as_ref())
.map_err(|e| Status::internal(e.to_string()))?;
}
let schema = b.schema();
Ok(Response::new(metadata_stream(schema, b.build())))
}
async fn do_get_table_types(
&self,
_query: CommandGetTableTypes,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
Ok(Response::new(TABLE_TYPES_FRAMES.stream()))
}
async fn do_get_sql_info(
&self,
query: CommandGetSqlInfo,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
// No id filter (the usual driver call) → the batch built once at startup.
if query.info.is_empty() {
return Ok(Response::new(SQL_INFO_FRAMES.stream()));
}
let b = query.into_builder(self.sql_info);
let schema = b.schema();
Ok(Response::new(metadata_stream(schema, b.build())))
}
async fn do_get_xdbc_type_info(
&self,
query: CommandGetXdbcTypeInfo,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
// ODBC `SQLGetTypeInfo(SQL_ALL_TYPES)` sends 0: the whole table.
if query.data_type.is_none_or(|dt| dt == 0) {
return Ok(Response::new(XDBC_TYPE_INFO_FRAMES.stream()));
}
let b = query.into_builder(self.type_info);
let schema = b.schema();
// Same null-vs-empty-list repair as the unfiltered table above.
let batch = b.build().map(crate::sqlinfo::create_params_never_null);
// A type code we do not carry (an ODBC concise/verbose alias such as
// SQL_DATETIME=9) would filter to zero rows, and the Flight encoder
// sends no RecordBatch at all for an empty batch: Power BI's ADBC layer
// then dereferences a null result set. Hand the whole table back
// instead; drivers pick their row by data_type themselves.
if matches!(&batch, Ok(b) if b.num_rows() == 0) {
return Ok(Response::new(XDBC_TYPE_INFO_FRAMES.stream()));
}
Ok(Response::new(metadata_stream(schema, batch)))
}
// ── Keys: identifier fields as PK, no FKs (see keys.rs) ─────────────
async fn get_flight_info_primary_keys(
&self,
query: CommandGetPrimaryKeys,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(request.into_inner(), ticket, PRIMARY_KEYS_SCHEMA.clone())
}
async fn get_flight_info_exported_keys(
&self,
query: CommandGetExportedKeys,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(request.into_inner(), ticket, FOREIGN_KEYS_SCHEMA.clone())
}
async fn get_flight_info_imported_keys(
&self,
query: CommandGetImportedKeys,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(request.into_inner(), ticket, FOREIGN_KEYS_SCHEMA.clone())
}
async fn get_flight_info_cross_reference(
&self,
query: CommandGetCrossReference,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(request.into_inner(), ticket, FOREIGN_KEYS_SCHEMA.clone())
}
async fn do_get_primary_keys(
&self,
query: CommandGetPrimaryKeys,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
let schema = query.db_schema.as_deref().unwrap_or(self.engine.options().default_namespace.as_str());
let cols = self.engine.primary_keys(schema, &query.table).await.map_err(Status::from)?;
#[cfg(feature = "testmatrix")]
crate::functional_status("skade-flight/get_primary_keys", "answered", true, &format!("{}.{}: {} col(s)", schema, query.table, cols.len()));
let batch = primary_keys(CATALOG, schema, &query.table, &cols);
Ok(Response::new(metadata_stream(PRIMARY_KEYS_SCHEMA.clone(), Ok(batch))))
}
async fn do_get_exported_keys(
&self,
_query: CommandGetExportedKeys,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
Ok(Response::new(metadata_stream(FOREIGN_KEYS_SCHEMA.clone(), Ok(no_foreign_keys()))))
}
async fn do_get_imported_keys(
&self,
_query: CommandGetImportedKeys,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
Ok(Response::new(metadata_stream(FOREIGN_KEYS_SCHEMA.clone(), Ok(no_foreign_keys()))))
}
async fn do_get_cross_reference(
&self,
_query: CommandGetCrossReference,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
Ok(Response::new(metadata_stream(FOREIGN_KEYS_SCHEMA.clone(), Ok(no_foreign_keys()))))
}
async fn do_get_fallback(
&self,
_request: Request<Ticket>,
message: Any,
) -> Result<Response<BoxStream>, Status> {
Err(Status::unimplemented(format!("unsupported ticket type {}", message.type_url)))
}
}
fn arrow_ipc_options() -> arrow_ipc::writer::IpcWriteOptions {
arrow_ipc::writer::IpcWriteOptions::default()
}