crates/graphar-flight/src/service.rs
use std::{
collections::HashMap,
pin::Pin,
sync::{Arc, LazyLock, Mutex},
};
use arrow::ipc::writer::IpcWriteOptions;
use arrow_array::RecordBatch;
use arrow_flight::{
FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest, HandshakeResponse,
IpcMessage, SchemaAsIpc, Ticket,
decode::FlightRecordBatchStream,
encode::FlightDataEncoderBuilder,
error::FlightError,
sql::{
ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetDbSchemas,
CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandPreparedStatementQuery,
CommandStatementQuery, CommandStatementUpdate, DoPutPreparedStatementResult,
ProstMessageExt, SqlInfo, TicketStatementQuery,
metadata::{SqlInfoData, SqlInfoDataBuilder},
server::{FlightSqlService, PeekableFlightDataStream},
},
};
use futures::{Stream, TryStreamExt};
use graphar_falkordb::FalkorDbLoader;
use prost::Message;
use tonic::{Request, Response, Status, Streaming};
use crate::{
auth::{AuthConfig, Identity},
error::FlightSqlError,
falkor::FalkorExecutor,
registry::{SchemaRegistry, SqlShape, classify_sql, table_candidates},
};
/// Static `GetSqlInfo` payload: what knut advertises about itself to a
/// generic Flight SQL driver (Power BI's ADBC connector, `adbc_driver_flightsql`).
/// The dialect is Cypher, not SQL, so we report read-only with no DDL/DML SQL.
static SQL_INFO: LazyLock<SqlInfoData> = LazyLock::new(|| {
let mut b = SqlInfoDataBuilder::new();
b.append(
SqlInfo::FlightSqlServerName,
"knut graphar-flight (Cypher over FalkorDB)",
);
b.append(SqlInfo::FlightSqlServerVersion, env!("CARGO_PKG_VERSION"));
// Arrow format version, per Schema.fbs.
b.append(SqlInfo::FlightSqlServerArrowVersion, "1.3");
b.append(SqlInfo::FlightSqlServerReadOnly, true);
// We speak Cypher, not SQL DDL/DML — advertise no SQL grammar.
b.append(SqlInfo::SqlDdlCatalog, false);
b.append(SqlInfo::SqlIdentifierQuoteChar, "\"");
b.build().expect("static SqlInfoData builds")
});
type BoxStream = Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
/// Upper bound on the number of auto-schema `RecordBatch`es held between a
/// `GetFlightInfo` and its matching `DoGet`. Each in-flight ad-hoc query needs
/// at most one slot; the cap exists only so a client that calls `GetFlightInfo`
/// without ever following up with `DoGet` (or disappears mid-flight) can't grow
/// the map without bound. When full, the oldest-inserted entry is dropped.
const SCHEMA_CACHE_CAP: usize = 16;
/// Cache of inferred auto-schema results, keyed by resolved Cypher.
///
/// In auto-schema mode (an unregistered query) `GetFlightInfo` must execute the
/// query to learn its schema. Rather than throw that result away and re-run the
/// identical query in `DoGet` (H9), the whole `RecordBatch` is parked here and
/// the matching `DoGet` removes-and-streams it. Bounded (`SCHEMA_CACHE_CAP`) and
/// shared across cloned services, so memory stays flat and any request thread
/// can serve the follow-up.
#[derive(Clone, Default)]
struct SchemaCache {
inner: Arc<Mutex<HashMap<String, RecordBatch>>>,
}
impl SchemaCache {
/// Park an inferred batch under its resolved Cypher. If the cache is at
/// capacity (and this key is new), one existing entry is evicted first so
/// the map can never exceed [`SCHEMA_CACHE_CAP`].
fn insert(&self, cypher: String, batch: RecordBatch) {
let mut map = self.inner.lock().expect("schema cache mutex");
if map.len() >= SCHEMA_CACHE_CAP && !map.contains_key(&cypher) {
// Drop an arbitrary entry — these are short-lived GetFlightInfo→DoGet
// hand-offs, so any abandoned one is fair game.
if let Some(victim) = map.keys().next().cloned() {
map.remove(&victim);
}
}
map.insert(cypher, batch);
}
/// Take the cached batch for a Cypher, removing it (a DoGet consumes it).
fn take(&self, cypher: &str) -> Option<RecordBatch> {
self.inner
.lock()
.expect("schema cache mutex")
.remove(cypher)
}
}
/// A Flight SQL service that accepts Cypher queries instead of SQL.
///
/// ## Query (DoGet)
/// Send any key registered in [`SchemaRegistry`] as the query string.
/// The server executes it against FalkorDB and streams Arrow `RecordBatch`es.
///
/// ## Insert (DoPut)
/// Set the `CommandStatementUpdate` query to:
/// - `VERTICES:<label>` — bulk-insert vertices from the Arrow stream
/// - `EDGES:<src>:<type>:<dst>` — bulk-insert edges from the Arrow stream
#[derive(Clone)]
pub struct CypherFlightService {
executor: FalkorExecutor,
registry: Arc<SchemaRegistry>,
redis_url: Arc<String>,
graph: Arc<String>,
auth: Arc<AuthConfig>,
/// Auto-schema results parked between `GetFlightInfo` and `DoGet` (H9).
schema_cache: SchemaCache,
/// How many `UNWIND`-window queries a **DoPut** keeps in flight (see
/// [`with_put_concurrency`](CypherFlightService::with_put_concurrency)).
/// Defaults to [`default_put_concurrency`].
put_concurrency: usize,
}
/// The default DoPut in-flight window count: the host's parallelism, clamped to
/// a floor of 1.
///
/// FalkorDB's `MultiplexedConnection` carries many concurrent commands over one
/// link, so the ceiling here is the server's ability to absorb them, not the
/// client's socket count. `available_parallelism()` is the same default the
/// in-process bulk loader uses, so DoPut and in-process ingest are tuned alike.
pub fn default_put_concurrency() -> usize {
std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1)
.max(1)
}
impl CypherFlightService {
pub async fn connect(
redis_url: impl Into<String>,
graph: impl Into<String>,
registry: Arc<SchemaRegistry>,
) -> crate::error::Result<Self> {
let redis_url = Arc::new(redis_url.into());
let graph = Arc::new(graph.into());
let executor = FalkorExecutor::connect(redis_url.as_str(), graph.as_str()).await?;
Ok(Self {
executor,
registry,
redis_url,
graph,
auth: Arc::new(AuthConfig::default()),
schema_cache: SchemaCache::default(),
put_concurrency: default_put_concurrency(),
})
}
/// Attach an authentication config. Drives the handshake; the matching
/// per-call enforcement is installed by [`crate::server::run_server_with`].
pub fn with_auth(mut self, auth: Arc<AuthConfig>) -> Self {
self.auth = auth;
self
}
/// Override how many `UNWIND`-window queries a **DoPut** keeps in flight.
/// `0` is clamped to `1` (a serial dispatch — the pre-fix behaviour, kept
/// reachable so the two paths can be A/B'd and so an operator can pin a
/// shared FalkorDB back down).
pub fn with_put_concurrency(mut self, n: usize) -> Self {
self.put_concurrency = n.max(1);
self
}
/// The configured DoPut in-flight window count.
pub fn put_concurrency(&self) -> usize {
self.put_concurrency
}
async fn new_loader(&self) -> crate::error::Result<FalkorDbLoader> {
FalkorDbLoader::connect(&self.redis_url, self.graph.as_str())
.await
.map_err(FlightSqlError::from)
}
/// Resolve an incoming query string to the Cypher to actually run — or
/// refuse, loudly, when the result would not be what was asked for.
///
/// Order: a registered key or table-name alias wins; otherwise the query is
/// classified ([`classify_sql`]): the one honoured SQL shape,
/// `SELECT * FROM <table>` (what generic ADBC/ODBC drivers emit), is
/// rewritten to the table's Cypher; any other SQL is `InvalidArgument`
/// naming every clause that would have been dropped; anything that is not
/// SQL is taken as Cypher verbatim and served through the auto-string path.
///
/// It used to return `String` and silently serve the whole table for
/// `SELECT a, b FROM people WHERE a > 1` — the projection and the predicate
/// thrown away, a wrong row count, no signal. Power BI folds by default, so
/// that was a dashboard showing wrong numbers with nothing to notice.
fn target_cypher(&self, incoming: &str) -> Result<String, Status> {
if let Some(query) = self.registry.resolve_table(incoming) {
return Ok(query);
}
match classify_sql(incoming) {
SqlShape::NotSql => Ok(incoming.to_string()),
SqlShape::Honoured { table } => table_candidates(&table)
.iter()
.find_map(|name| self.registry.resolve_table(name))
.ok_or_else(|| {
Status::not_found(format!(
"no table named `{table}` is registered on this server; \
GetTables lists the ones that are"
))
}),
SqlShape::Unhonoured { table, dropped } => Err(Status::invalid_argument(format!(
"graphar-flight cannot honour this SQL and will not answer it with something \
else: it would have dropped {} and returned {}. This engine \
executes Cypher; the only SQL it serves is `SELECT * FROM <table>` (a \
table from GetTables). Send the query as Cypher, or read the table whole \
and filter client-side.",
dropped
.iter()
.map(|d| format!("`{d}`"))
.collect::<Vec<_>>()
.join(", "),
if table.is_empty() {
"a different result".to_string()
} else {
format!("ALL of `{table}`")
}
))),
}
}
/// The Arrow schema a query returns: the registered schema when known,
/// else inferred by executing it once (auto-string mode).
///
/// In the auto path the executed `RecordBatch` is **not** thrown away — it
/// is parked in [`Self::schema_cache`] keyed by `cypher`, so the matching
/// `DoGet` can stream it without re-running the query (H9). Registered
/// queries need no execution here and cache nothing.
async fn schema_for(&self, cypher: &str) -> Result<arrow_schema::SchemaRef, Status> {
if let Some(s) = self.registry.get(cypher) {
Ok(s)
} else {
let batch = self
.executor
.clone()
.query_auto(cypher)
.await
.map_err(|e| Status::internal(e.to_string()))?;
let schema = batch.schema();
self.schema_cache.insert(cypher.to_string(), batch);
Ok(schema)
}
}
/// The authorization name for a resolved Cypher query: its friendly
/// registered table name when it has one, else the Cypher string itself.
/// This is the key the [`crate::auth::Authorizer`] policy matches against.
fn authz_name(&self, cypher: &str) -> String {
self.registry
.name_of(cypher)
.unwrap_or_else(|| cypher.to_string())
}
/// Consult the per-query authorizer for a resolved Cypher query, using the
/// caller [`Identity`] the interceptor stashed in the request extensions
/// (absent ⇒ [`Identity::Anonymous`], e.g. an open server with no
/// interceptor). Denial → `permission_denied`. A no-op under the default
/// [`crate::auth::Authorizer::AllowAll`].
fn authorize<T>(&self, request: &Request<T>, cypher: &str) -> Result<(), Status> {
let identity = request
.extensions()
.get::<Identity>()
.cloned()
.unwrap_or(Identity::Anonymous);
self.auth.authorize(&identity, &self.authz_name(cypher))
}
/// Execute a (already resolved) Cypher query against FalkorDB and stream
/// the result as Arrow Flight data. Shared by plain and prepared DoGet.
///
/// For an auto-schema query whose batch was parked by a preceding
/// `GetFlightInfo`, the cached batch is taken and streamed as-is — no second
/// execution (H9). A cache miss (DoGet without a prior GetFlightInfo, or an
/// evicted entry) falls back to executing the query fresh. Registered
/// queries are never cached and always execute here against their typed
/// schema.
async fn execute_cypher_stream(&self, cypher: &str) -> Result<Response<BoxStream>, Status> {
let mut exec = self.executor.clone();
let batch = if let Some(schema) = self.registry.get(cypher) {
exec.query(cypher, &schema)
.await
.map_err(|e| Status::internal(e.to_string()))?
} else if let Some(cached) = self.schema_cache.take(cypher) {
cached
} else {
exec.query_auto(cypher)
.await
.map_err(|e| Status::internal(e.to_string()))?
};
let schema = batch.schema();
let stream: BoxStream = Box::pin(
FlightDataEncoderBuilder::new()
.with_schema(schema)
.build(futures::stream::iter([Ok(batch)]))
.map_err(|e: FlightError| Status::internal(e.to_string())),
);
Ok(Response::new(stream))
}
}
#[tonic::async_trait]
impl FlightSqlService for CypherFlightService {
type FlightService = CypherFlightService;
// ── Required: sql info registry (we don't advertise SQL capabilities) ────
async fn register_sql_info(&self, _id: i32, _result: &SqlInfo) {}
// ── Handshake — Basic login → bearer token ───────────────────────────────
//
// The canonical Flight SQL auth flow for generic Windows clients (Power
// BI's ADBC connector, the Flight SQL ODBC driver): the client opens with
// `authorization: Basic <base64(user:pass)>`, the server validates it and
// returns a bearer token in the `authorization` header — and ONLY there;
// see `handshake_response` for the client that hangs on a payload — which
// the client replays on every later call.
async fn do_handshake(
&self,
request: Request<Streaming<HandshakeRequest>>,
) -> Result<
Response<Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send>>>,
Status,
> {
let header = request
.metadata()
.get("authorization")
.and_then(|v| v.to_str().ok());
// Reuse the exact same acceptance rules as the per-call interceptor.
self.auth.check_header(header)?;
handshake_response(self.auth.issued_token())
}
// ── GetFlightInfo — client calls this first to discover schema + ticket ──
//
// If the query is registered, use the known schema (fast, typed).
// Otherwise fall back to auto-string mode: execute the query once to read
// the column names from FalkorDB's response header, build an all-Utf8
// schema, and encode the query in the ticket for DoGet to re-execute.
async fn get_flight_info_statement(
&self,
query: CommandStatementQuery,
request: Request<arrow_flight::FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
// Rewrite `SELECT * FROM <table>` / aliases to the underlying Cypher,
// then authorize the caller for it before describing its schema.
let cypher = self.target_cypher(&query.query)?;
self.authorize(&request, &cypher)?;
let schema = match self.schema_for(&cypher).await {
Ok(s) => {
crate::functional_status(
"graphar-flight/get_flight_info_statement",
"schema_resolved",
true,
&cypher,
);
s
}
Err(e) => {
crate::functional_status(
"graphar-flight/get_flight_info_statement",
"schema_resolved",
false,
&cypher,
);
return Err(e);
}
};
statement_info(cypher, schema, request.into_inner())
}
// ── DoGet — execute Cypher, stream back Arrow data ───────────────────────
async fn do_get_statement(
&self,
ticket: TicketStatementQuery,
request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
let query = String::from_utf8(ticket.statement_handle.to_vec())
.map_err(|e| Status::invalid_argument(format!("invalid ticket: {e}")))?;
self.authorize(&request, &query)?;
let result = self.execute_cypher_stream(&query).await;
crate::functional_status(
"graphar-flight/do_get_statement",
"cypher_stream",
result.is_ok(),
&query,
);
result
}
// ── DoPut — receive Arrow stream, bulk-insert into FalkorDB ──────────────
async fn do_put_statement_update(
&self,
cmd: CommandStatementUpdate,
request: Request<PeekableFlightDataStream>,
) -> Result<i64, Status> {
let directive = cmd.query.trim().to_string();
let action =
parse_put_directive(&directive).map_err(|e| Status::invalid_argument(e.to_string()))?;
// Decode incoming FlightData stream → RecordBatches.
//
// `into_peekable()`, NOT `into_inner()`: the SQL server peeks the first
// message to route the command, and `into_inner()` documents that "any
// messages observed through peek will be lost". Every standard client —
// Arrow's own `FlightDataEncoderBuilder`, pyarrow's
// `do_put(descriptor, schema)` — puts the descriptor on the SAME message
// as the IPC schema, so dropping it dropped the schema and the decoder
// failed with "Received RecordBatch prior to Schema". Peeking keeps it.
//
// A client that instead sends a bare descriptor-only message first is
// still fine: such a message carries no `data_header`, and the filter
// below skips it rather than feeding the IPC decoder an empty header.
let raw = request.into_inner().into_peekable();
let decoded = raw.try_filter(|d| futures::future::ready(!d.data_header.is_empty()));
let batches = FlightRecordBatchStream::new_from_flight_data(
decoded.map_err(|s| FlightError::Tonic(Box::new(s))),
)
.try_collect::<Vec<_>>()
.await
.map_err(|e| Status::internal(e.to_string()))?;
if batches.is_empty() {
return Ok(0);
}
let loader = self
.new_loader()
.await
.map_err(|e| Status::internal(e.to_string()))?;
// ── the DoPut dispatch, routed to the PARALLEL bulk path ─────────────
//
// This used to call the serial `insert_vertices` / `insert_edges`, which
// await one `UNWIND … CREATE` at a time over the single link: the client
// sits idle for a full round-trip between windows, so a Flight SQL DoPut
// ingested at roughly one core no matter how big the box — while the
// in-process bulk loader had been on the concurrent path for a while.
// Same crate, same renderer, two different speeds for the same work.
//
// `insert_vertices_par` / `insert_edges_par` render the SAME per-window
// Cypher (`render_vertex_queries` / `render_edge_queries` — the very
// functions the serial flush shares) with the SAME `_gar_id` scheme, and
// dispatch up to `put_concurrency` of them at once over cloned
// multiplexed connections. Vertices are `CREATE`d over disjoint `_gar_id`
// ranges and edges `MATCH` endpoints that already exist, so both are
// order-independent: concurrency cannot change the graph, only the time
// it takes to build it.
let concurrency = self.put_concurrency;
let n = match action {
PutAction::Vertices { label } => loader
.insert_vertices_par(&label, &batches, concurrency)
.await
.map_err(|e| Status::internal(e.to_string()))?,
PutAction::Edges {
src,
edge_type,
dst,
} => loader
.insert_edges_par(&src, &edge_type, &dst, &batches, concurrency)
.await
.map_err(|e| Status::internal(e.to_string()))?,
};
crate::functional_status(
"graphar-flight/do_put_statement_update",
"parallel_bulk_insert",
true,
&format!("{directive}: {n} rows, concurrency x{concurrency}"),
);
Ok(n as i64)
}
// ── Metadata surface — lets generic ADBC/ODBC Flight SQL drivers ─────────
// (Power BI's connector, adbc_driver_flightsql) introspect knut.
// One catalog (the graph), one schema (`public`), and one "table" per
// registered query.
async fn get_flight_info_catalogs(
&self,
query: CommandGetCatalogs,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let descriptor = request.into_inner();
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(descriptor, ticket, query.into_builder().schema())
}
async fn get_flight_info_schemas(
&self,
query: CommandGetDbSchemas,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let descriptor = request.into_inner();
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(descriptor, ticket, query.into_builder().schema())
}
async fn get_flight_info_tables(
&self,
query: CommandGetTables,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let descriptor = request.into_inner();
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(descriptor, ticket, query.into_builder().schema())
}
async fn get_flight_info_table_types(
&self,
query: CommandGetTableTypes,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let descriptor = request.into_inner();
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(descriptor, ticket, query.into_builder().schema())
}
async fn get_flight_info_sql_info(
&self,
query: CommandGetSqlInfo,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let descriptor = request.into_inner();
let ticket = Ticket::new(query.as_any().encode_to_vec());
metadata_info(descriptor, ticket, query.into_builder(&SQL_INFO).schema())
}
async fn do_get_catalogs(
&self,
query: CommandGetCatalogs,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
let mut builder = query.into_builder();
builder.append(self.graph.as_str());
let schema = builder.schema();
Ok(Response::new(metadata_stream(schema, builder.build())))
}
async fn do_get_schemas(
&self,
query: CommandGetDbSchemas,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
let mut builder = query.into_builder();
builder.append(self.graph.as_str(), DB_SCHEMA);
let schema = builder.schema();
Ok(Response::new(metadata_stream(schema, builder.build())))
}
async fn do_get_tables(
&self,
query: CommandGetTables,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
let mut builder = query.into_builder();
for (name, table_schema) in self.registry.tables() {
builder
.append(
self.graph.as_str(),
DB_SCHEMA,
&name,
"TABLE",
table_schema.as_ref(),
)
.map_err(|e| Status::internal(e.to_string()))?;
}
let schema = builder.schema();
Ok(Response::new(metadata_stream(schema, builder.build())))
}
async fn do_get_table_types(
&self,
query: CommandGetTableTypes,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
let mut builder = query.into_builder();
builder.append("TABLE");
let schema = builder.schema();
Ok(Response::new(metadata_stream(schema, builder.build())))
}
async fn do_get_sql_info(
&self,
query: CommandGetSqlInfo,
_request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
let builder = query.into_builder(&SQL_INFO);
let schema = builder.schema();
Ok(Response::new(metadata_stream(schema, builder.build())))
}
// ── Prepared statements ──────────────────────────────────────────────────
// Stateless: the handle *is* the resolved Cypher, so no server-side
// statement table is needed and any node can serve any handle.
async fn do_action_create_prepared_statement(
&self,
query: ActionCreatePreparedStatementRequest,
request: Request<arrow_flight::Action>,
) -> Result<ActionCreatePreparedStatementResult, Status> {
let cypher = self.target_cypher(&query.query)?;
self.authorize(&request, &cypher)?;
let schema = self.schema_for(&cypher).await?;
let IpcMessage(schema_bytes) =
SchemaAsIpc::new(schema.as_ref(), &IpcWriteOptions::default())
.try_into()
.map_err(|e: arrow::error::ArrowError| Status::internal(e.to_string()))?;
Ok(ActionCreatePreparedStatementResult {
prepared_statement_handle: cypher.into_bytes().into(),
dataset_schema: schema_bytes,
parameter_schema: Default::default(), // no bind parameters
})
}
async fn do_action_close_prepared_statement(
&self,
_query: ActionClosePreparedStatementRequest,
_request: Request<arrow_flight::Action>,
) -> Result<(), Status> {
Ok(()) // stateless — nothing to release
}
async fn get_flight_info_prepared_statement(
&self,
cmd: CommandPreparedStatementQuery,
request: Request<FlightDescriptor>,
) -> Result<Response<FlightInfo>, Status> {
let cypher = String::from_utf8(cmd.prepared_statement_handle.to_vec())
.map_err(|e| Status::invalid_argument(format!("invalid handle: {e}")))?;
self.authorize(&request, &cypher)?;
let schema = self.schema_for(&cypher).await?;
statement_info(cypher, schema, request.into_inner())
}
async fn do_get_prepared_statement(
&self,
cmd: CommandPreparedStatementQuery,
request: Request<Ticket>,
) -> Result<Response<BoxStream>, Status> {
let cypher = String::from_utf8(cmd.prepared_statement_handle.to_vec())
.map_err(|e| Status::invalid_argument(format!("invalid handle: {e}")))?;
self.authorize(&request, &cypher)?;
self.execute_cypher_stream(&cypher).await
}
async fn do_put_prepared_statement_query(
&self,
query: CommandPreparedStatementQuery,
_request: Request<PeekableFlightDataStream>,
) -> Result<DoPutPreparedStatementResult, Status> {
// No bind parameters in the Cypher dialect — echo the handle back.
Ok(DoPutPreparedStatementResult {
prepared_statement_handle: Some(query.prepared_statement_handle),
})
}
}
/// **The handshake reply: NO `HandshakeResponse` message, the bearer in the
/// `authorization` header.** The ONE copy is `nornir_flight::door::handshake_response`
/// (shared with skade-flight; the measured Dremio ODBC hang is documented
/// there). This wrapper only maps "no token configured" to the open placeholder.
pub(crate) fn handshake_response(
issued_token: Option<&str>,
) -> Result<Response<nornir_flight::door::HandshakeStream>, Status> {
nornir_flight::door::handshake_response(issued_token.unwrap_or(""))
}
const DB_SCHEMA: &str = "public";
/// Materialise the static `GetSqlInfo` payload as a `RecordBatch`. Test-only
/// hook so the lazily-built [`SQL_INFO`] can be asserted without a live server.
#[cfg(test)]
pub(crate) fn sql_info_batch() -> arrow_array::RecordBatch {
CommandGetSqlInfo::default()
.into_builder(&SQL_INFO)
.build()
.expect("sql info batch builds")
}
/// Build the `FlightInfo` for a statement (ad-hoc or prepared): one endpoint
/// whose ticket is the resolved Cypher, the query's schema, and — this is why
/// it is one function — the request's `FlightDescriptor` echoed back.
///
/// Both statement handlers built this by hand and both omitted the
/// descriptor, while every metadata handler (through [`metadata_info`])
/// carried it. Generic clients treat a `FlightInfo` with no descriptor as
/// malformed: the Flight SQL comparison's C++ client refused every query
/// result while every `GetTables` worked, and the difference was exactly this
/// field. The descriptor is part of the reply shape; it is not optional here.
pub(crate) fn statement_info(
cypher: String,
schema: arrow_schema::SchemaRef,
descriptor: FlightDescriptor,
) -> Result<Response<FlightInfo>, Status> {
let ticket_bytes = TicketStatementQuery {
statement_handle: cypher.into_bytes().into(),
}
// Pack in an `Any` (like the sibling metadata handlers' `query.as_any()`):
// arrow-flight's `do_get` decodes the ticket as `Any` and matches on the
// type_url, so a BARE `TicketStatementQuery.encode_to_vec()` mis-decodes
// (Any.type_url collides with statement_handle, tag 1) and falls through to
// `do_get_fallback` → "The defined request is invalid". `.as_any()` fixes it.
.as_any()
.encode_to_vec();
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::new(ticket_bytes)))
.with_descriptor(descriptor)
.with_total_records(-1)
.with_total_bytes(-1)
.with_ordered(false);
Ok(Response::new(info))
}
/// Build a `FlightInfo` for a metadata response: one endpoint whose ticket is
/// the request command, advertising the builder's fixed schema.
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))
}
/// Encode a single metadata `RecordBatch` (already wrapped in a `FlightError`
/// result by the builder) as a one-item Flight data stream.
fn metadata_stream(
schema: arrow_schema::SchemaRef,
batch: std::result::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())),
)
}
// ── DoPut directive parsing ───────────────────────────────────────────────────
enum PutAction {
Vertices {
label: String,
},
Edges {
src: String,
edge_type: String,
dst: String,
},
}
fn parse_put_directive(s: &str) -> Result<PutAction, FlightSqlError> {
if let Some(rest) = s.strip_prefix("VERTICES:") {
return Ok(PutAction::Vertices {
label: rest.to_string(),
});
}
if let Some(rest) = s.strip_prefix("EDGES:") {
let parts: Vec<&str> = rest.splitn(3, ':').collect();
if parts.len() == 3 {
return Ok(PutAction::Edges {
src: parts[0].to_string(),
edge_type: parts[1].to_string(),
dst: parts[2].to_string(),
});
}
}
Err(FlightSqlError::UnknownPutCommand(s.to_string()))
}