Skip to main content

graphar_flight/
service.rs

1use std::{
2    collections::HashMap,
3    pin::Pin,
4    sync::{Arc, LazyLock, Mutex},
5};
6
7use arrow::ipc::writer::IpcWriteOptions;
8use arrow_array::RecordBatch;
9use arrow_flight::{
10    FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest, HandshakeResponse,
11    IpcMessage, SchemaAsIpc, Ticket,
12    decode::FlightRecordBatchStream,
13    encode::FlightDataEncoderBuilder,
14    error::FlightError,
15    sql::{
16        ActionClosePreparedStatementRequest, ActionCreatePreparedStatementRequest,
17        ActionCreatePreparedStatementResult, CommandGetCatalogs, CommandGetDbSchemas,
18        CommandGetSqlInfo, CommandGetTableTypes, CommandGetTables, CommandPreparedStatementQuery,
19        CommandStatementQuery, CommandStatementUpdate, DoPutPreparedStatementResult,
20        ProstMessageExt, SqlInfo, TicketStatementQuery,
21        metadata::{SqlInfoData, SqlInfoDataBuilder},
22        server::{FlightSqlService, PeekableFlightDataStream},
23    },
24};
25use futures::{Stream, TryStreamExt};
26use graphar_falkordb::FalkorDbLoader;
27use prost::Message;
28use tonic::{Request, Response, Status, Streaming};
29
30use crate::{
31    auth::{AuthConfig, Identity},
32    error::FlightSqlError,
33    falkor::FalkorExecutor,
34    registry::{SchemaRegistry, parse_select_from},
35};
36
37/// Static `GetSqlInfo` payload: what knut advertises about itself to a
38/// generic Flight SQL driver (Power BI's ADBC connector, `adbc_driver_flightsql`).
39/// The dialect is Cypher, not SQL, so we report read-only with no DDL/DML SQL.
40static SQL_INFO: LazyLock<SqlInfoData> = LazyLock::new(|| {
41    let mut b = SqlInfoDataBuilder::new();
42    b.append(
43        SqlInfo::FlightSqlServerName,
44        "knut graphar-flight (Cypher over FalkorDB)",
45    );
46    b.append(SqlInfo::FlightSqlServerVersion, env!("CARGO_PKG_VERSION"));
47    // Arrow format version, per Schema.fbs.
48    b.append(SqlInfo::FlightSqlServerArrowVersion, "1.3");
49    b.append(SqlInfo::FlightSqlServerReadOnly, true);
50    // We speak Cypher, not SQL DDL/DML — advertise no SQL grammar.
51    b.append(SqlInfo::SqlDdlCatalog, false);
52    b.append(SqlInfo::SqlIdentifierQuoteChar, "\"");
53    b.build().expect("static SqlInfoData builds")
54});
55
56type BoxStream = Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
57
58/// Upper bound on the number of auto-schema `RecordBatch`es held between a
59/// `GetFlightInfo` and its matching `DoGet`. Each in-flight ad-hoc query needs
60/// at most one slot; the cap exists only so a client that calls `GetFlightInfo`
61/// without ever following up with `DoGet` (or disappears mid-flight) can't grow
62/// the map without bound. When full, the oldest-inserted entry is dropped.
63const SCHEMA_CACHE_CAP: usize = 16;
64
65/// Cache of inferred auto-schema results, keyed by resolved Cypher.
66///
67/// In auto-schema mode (an unregistered query) `GetFlightInfo` must execute the
68/// query to learn its schema. Rather than throw that result away and re-run the
69/// identical query in `DoGet` (H9), the whole `RecordBatch` is parked here and
70/// the matching `DoGet` removes-and-streams it. Bounded (`SCHEMA_CACHE_CAP`) and
71/// shared across cloned services, so memory stays flat and any request thread
72/// can serve the follow-up.
73#[derive(Clone, Default)]
74struct SchemaCache {
75    inner: Arc<Mutex<HashMap<String, RecordBatch>>>,
76}
77
78impl SchemaCache {
79    /// Park an inferred batch under its resolved Cypher. If the cache is at
80    /// capacity (and this key is new), one existing entry is evicted first so
81    /// the map can never exceed [`SCHEMA_CACHE_CAP`].
82    fn insert(&self, cypher: String, batch: RecordBatch) {
83        let mut map = self.inner.lock().expect("schema cache mutex");
84        if map.len() >= SCHEMA_CACHE_CAP && !map.contains_key(&cypher) {
85            // Drop an arbitrary entry — these are short-lived GetFlightInfo→DoGet
86            // hand-offs, so any abandoned one is fair game.
87            if let Some(victim) = map.keys().next().cloned() {
88                map.remove(&victim);
89            }
90        }
91        map.insert(cypher, batch);
92    }
93
94    /// Take the cached batch for a Cypher, removing it (a DoGet consumes it).
95    fn take(&self, cypher: &str) -> Option<RecordBatch> {
96        self.inner
97            .lock()
98            .expect("schema cache mutex")
99            .remove(cypher)
100    }
101}
102
103/// A Flight SQL service that accepts Cypher queries instead of SQL.
104///
105/// ## Query (DoGet)
106/// Send any key registered in [`SchemaRegistry`] as the query string.
107/// The server executes it against FalkorDB and streams Arrow `RecordBatch`es.
108///
109/// ## Insert (DoPut)
110/// Set the `CommandStatementUpdate` query to:
111/// - `VERTICES:<label>` — bulk-insert vertices from the Arrow stream
112/// - `EDGES:<src>:<type>:<dst>` — bulk-insert edges from the Arrow stream
113#[derive(Clone)]
114pub struct CypherFlightService {
115    executor: FalkorExecutor,
116    registry: Arc<SchemaRegistry>,
117    redis_url: Arc<String>,
118    graph: Arc<String>,
119    auth: Arc<AuthConfig>,
120    /// Auto-schema results parked between `GetFlightInfo` and `DoGet` (H9).
121    schema_cache: SchemaCache,
122}
123
124impl CypherFlightService {
125    pub async fn connect(
126        redis_url: impl Into<String>,
127        graph: impl Into<String>,
128        registry: Arc<SchemaRegistry>,
129    ) -> crate::error::Result<Self> {
130        let redis_url = Arc::new(redis_url.into());
131        let graph = Arc::new(graph.into());
132        let executor = FalkorExecutor::connect(redis_url.as_str(), graph.as_str()).await?;
133        Ok(Self {
134            executor,
135            registry,
136            redis_url,
137            graph,
138            auth: Arc::new(AuthConfig::default()),
139            schema_cache: SchemaCache::default(),
140        })
141    }
142
143    /// Attach an authentication config. Drives the handshake; the matching
144    /// per-call enforcement is installed by [`crate::server::run_server_with`].
145    pub fn with_auth(mut self, auth: Arc<AuthConfig>) -> Self {
146        self.auth = auth;
147        self
148    }
149
150    async fn new_loader(&self) -> crate::error::Result<FalkorDbLoader> {
151        FalkorDbLoader::connect(&self.redis_url, self.graph.as_str())
152            .await
153            .map_err(FlightSqlError::from)
154    }
155
156    /// Resolve an incoming query string to the Cypher to actually run.
157    ///
158    /// Order: a registered key or table-name alias wins; otherwise a trivial
159    /// `SELECT * FROM <table>` (what generic ADBC/ODBC drivers emit) is
160    /// rewritten to the aliased Cypher; otherwise the string is taken as
161    /// Cypher verbatim and served through the auto-string path.
162    fn target_cypher(&self, incoming: &str) -> String {
163        if let Some(query) = self.registry.resolve_table(incoming) {
164            return query;
165        }
166        if let Some(ident) = parse_select_from(incoming)
167            && let Some(query) = self.registry.resolve_table(&ident)
168        {
169            return query;
170        }
171        incoming.to_string()
172    }
173
174    /// The Arrow schema a query returns: the registered schema when known,
175    /// else inferred by executing it once (auto-string mode).
176    ///
177    /// In the auto path the executed `RecordBatch` is **not** thrown away — it
178    /// is parked in [`Self::schema_cache`] keyed by `cypher`, so the matching
179    /// `DoGet` can stream it without re-running the query (H9). Registered
180    /// queries need no execution here and cache nothing.
181    async fn schema_for(&self, cypher: &str) -> Result<arrow_schema::SchemaRef, Status> {
182        if let Some(s) = self.registry.get(cypher) {
183            Ok(s)
184        } else {
185            let batch = self
186                .executor
187                .clone()
188                .query_auto(cypher)
189                .await
190                .map_err(|e| Status::internal(e.to_string()))?;
191            let schema = batch.schema();
192            self.schema_cache.insert(cypher.to_string(), batch);
193            Ok(schema)
194        }
195    }
196
197    /// The authorization name for a resolved Cypher query: its friendly
198    /// registered table name when it has one, else the Cypher string itself.
199    /// This is the key the [`crate::auth::Authorizer`] policy matches against.
200    fn authz_name(&self, cypher: &str) -> String {
201        self.registry
202            .name_of(cypher)
203            .unwrap_or_else(|| cypher.to_string())
204    }
205
206    /// Consult the per-query authorizer for a resolved Cypher query, using the
207    /// caller [`Identity`] the interceptor stashed in the request extensions
208    /// (absent ⇒ [`Identity::Anonymous`], e.g. an open server with no
209    /// interceptor). Denial → `permission_denied`. A no-op under the default
210    /// [`crate::auth::Authorizer::AllowAll`].
211    fn authorize<T>(&self, request: &Request<T>, cypher: &str) -> Result<(), Status> {
212        let identity = request
213            .extensions()
214            .get::<Identity>()
215            .cloned()
216            .unwrap_or(Identity::Anonymous);
217        self.auth.authorize(&identity, &self.authz_name(cypher))
218    }
219
220    /// Execute a (already resolved) Cypher query against FalkorDB and stream
221    /// the result as Arrow Flight data. Shared by plain and prepared DoGet.
222    ///
223    /// For an auto-schema query whose batch was parked by a preceding
224    /// `GetFlightInfo`, the cached batch is taken and streamed as-is — no second
225    /// execution (H9). A cache miss (DoGet without a prior GetFlightInfo, or an
226    /// evicted entry) falls back to executing the query fresh. Registered
227    /// queries are never cached and always execute here against their typed
228    /// schema.
229    async fn execute_cypher_stream(&self, cypher: &str) -> Result<Response<BoxStream>, Status> {
230        let mut exec = self.executor.clone();
231        let batch = if let Some(schema) = self.registry.get(cypher) {
232            exec.query(cypher, &schema)
233                .await
234                .map_err(|e| Status::internal(e.to_string()))?
235        } else if let Some(cached) = self.schema_cache.take(cypher) {
236            cached
237        } else {
238            exec.query_auto(cypher)
239                .await
240                .map_err(|e| Status::internal(e.to_string()))?
241        };
242        let schema = batch.schema();
243        let stream: BoxStream = Box::pin(
244            FlightDataEncoderBuilder::new()
245                .with_schema(schema)
246                .build(futures::stream::iter([Ok(batch)]))
247                .map_err(|e: FlightError| Status::internal(e.to_string())),
248        );
249        Ok(Response::new(stream))
250    }
251}
252
253#[tonic::async_trait]
254impl FlightSqlService for CypherFlightService {
255    type FlightService = CypherFlightService;
256
257    // ── Required: sql info registry (we don't advertise SQL capabilities) ────
258
259    async fn register_sql_info(&self, _id: i32, _result: &SqlInfo) {}
260
261    // ── Handshake — Basic login → bearer token ───────────────────────────────
262    //
263    // The canonical Flight SQL auth flow for generic Windows clients (Power
264    // BI's ADBC connector, the Flight SQL ODBC driver): the client opens with
265    // `authorization: Basic <base64(user:pass)>`, the server validates it and
266    // returns a bearer token (in both the response payload and the
267    // `authorization` trailer) which the client replays on every later call.
268
269    async fn do_handshake(
270        &self,
271        request: Request<Streaming<HandshakeRequest>>,
272    ) -> Result<
273        Response<Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send>>>,
274        Status,
275    > {
276        let header = request
277            .metadata()
278            .get("authorization")
279            .and_then(|v| v.to_str().ok());
280        // Reuse the exact same acceptance rules as the per-call interceptor.
281        self.auth.check_header(header)?;
282
283        let token = self.auth.issued_token().unwrap_or("").to_string();
284        let response_msg = HandshakeResponse {
285            protocol_version: 0,
286            payload: token.clone().into(),
287        };
288        let output = futures::stream::once(async move { Ok(response_msg) });
289        let mut response: Response<Pin<Box<dyn Stream<Item = _> + Send>>> =
290            Response::new(Box::pin(output));
291        if !token.is_empty() {
292            let bearer = format!("Bearer {token}")
293                .parse()
294                .map_err(|_| Status::internal("token not header-safe"))?;
295            response.metadata_mut().insert("authorization", bearer);
296        }
297        Ok(response)
298    }
299
300    // ── GetFlightInfo — client calls this first to discover schema + ticket ──
301    //
302    // If the query is registered, use the known schema (fast, typed).
303    // Otherwise fall back to auto-string mode: execute the query once to read
304    // the column names from FalkorDB's response header, build an all-Utf8
305    // schema, and encode the query in the ticket for DoGet to re-execute.
306
307    async fn get_flight_info_statement(
308        &self,
309        query: CommandStatementQuery,
310        request: Request<arrow_flight::FlightDescriptor>,
311    ) -> Result<Response<FlightInfo>, Status> {
312        // Rewrite `SELECT * FROM <table>` / aliases to the underlying Cypher,
313        // then authorize the caller for it before describing its schema.
314        let cypher = self.target_cypher(&query.query);
315        self.authorize(&request, &cypher)?;
316        let schema = self.schema_for(&cypher).await?;
317
318        let ticket_bytes = TicketStatementQuery {
319            statement_handle: cypher.into_bytes().into(),
320        }
321        .encode_to_vec();
322
323        let info = FlightInfo::new()
324            .try_with_schema(schema.as_ref())
325            .map_err(|e| Status::internal(e.to_string()))?
326            .with_endpoint(FlightEndpoint::new().with_ticket(Ticket::new(ticket_bytes)))
327            .with_total_records(-1)
328            .with_total_bytes(-1)
329            .with_ordered(false);
330
331        Ok(Response::new(info))
332    }
333
334    // ── DoGet — execute Cypher, stream back Arrow data ───────────────────────
335
336    async fn do_get_statement(
337        &self,
338        ticket: TicketStatementQuery,
339        request: Request<Ticket>,
340    ) -> Result<Response<BoxStream>, Status> {
341        let query = String::from_utf8(ticket.statement_handle.to_vec())
342            .map_err(|e| Status::invalid_argument(format!("invalid ticket: {e}")))?;
343        self.authorize(&request, &query)?;
344        self.execute_cypher_stream(&query).await
345    }
346
347    // ── DoPut — receive Arrow stream, bulk-insert into FalkorDB ──────────────
348
349    async fn do_put_statement_update(
350        &self,
351        cmd: CommandStatementUpdate,
352        request: Request<PeekableFlightDataStream>,
353    ) -> Result<i64, Status> {
354        let directive = cmd.query.trim().to_string();
355        let action =
356            parse_put_directive(&directive).map_err(|e| Status::invalid_argument(e.to_string()))?;
357
358        // Decode incoming FlightData stream → RecordBatches.
359        let raw = request.into_inner().into_inner(); // Streaming<FlightData>
360        let batches = FlightRecordBatchStream::new_from_flight_data(
361            raw.map_err(|s| FlightError::Tonic(Box::new(s))),
362        )
363        .try_collect::<Vec<_>>()
364        .await
365        .map_err(|e| Status::internal(e.to_string()))?;
366
367        if batches.is_empty() {
368            return Ok(0);
369        }
370
371        let mut loader = self
372            .new_loader()
373            .await
374            .map_err(|e| Status::internal(e.to_string()))?;
375
376        let n = match action {
377            PutAction::Vertices { label } => loader
378                .insert_vertices(&label, &batches)
379                .await
380                .map_err(|e| Status::internal(e.to_string()))?,
381            PutAction::Edges {
382                src,
383                edge_type,
384                dst,
385            } => loader
386                .insert_edges(&src, &edge_type, &dst, &batches)
387                .await
388                .map_err(|e| Status::internal(e.to_string()))?,
389        };
390
391        Ok(n as i64)
392    }
393
394    // ── Metadata surface — lets generic ADBC/ODBC Flight SQL drivers ─────────
395    //    (Power BI's connector, adbc_driver_flightsql) introspect knut.
396    //    One catalog (the graph), one schema (`public`), and one "table" per
397    //    registered query.
398
399    async fn get_flight_info_catalogs(
400        &self,
401        query: CommandGetCatalogs,
402        request: Request<FlightDescriptor>,
403    ) -> Result<Response<FlightInfo>, Status> {
404        let descriptor = request.into_inner();
405        let ticket = Ticket::new(query.as_any().encode_to_vec());
406        metadata_info(descriptor, ticket, query.into_builder().schema())
407    }
408
409    async fn get_flight_info_schemas(
410        &self,
411        query: CommandGetDbSchemas,
412        request: Request<FlightDescriptor>,
413    ) -> Result<Response<FlightInfo>, Status> {
414        let descriptor = request.into_inner();
415        let ticket = Ticket::new(query.as_any().encode_to_vec());
416        metadata_info(descriptor, ticket, query.into_builder().schema())
417    }
418
419    async fn get_flight_info_tables(
420        &self,
421        query: CommandGetTables,
422        request: Request<FlightDescriptor>,
423    ) -> Result<Response<FlightInfo>, Status> {
424        let descriptor = request.into_inner();
425        let ticket = Ticket::new(query.as_any().encode_to_vec());
426        metadata_info(descriptor, ticket, query.into_builder().schema())
427    }
428
429    async fn get_flight_info_table_types(
430        &self,
431        query: CommandGetTableTypes,
432        request: Request<FlightDescriptor>,
433    ) -> Result<Response<FlightInfo>, Status> {
434        let descriptor = request.into_inner();
435        let ticket = Ticket::new(query.as_any().encode_to_vec());
436        metadata_info(descriptor, ticket, query.into_builder().schema())
437    }
438
439    async fn get_flight_info_sql_info(
440        &self,
441        query: CommandGetSqlInfo,
442        request: Request<FlightDescriptor>,
443    ) -> Result<Response<FlightInfo>, Status> {
444        let descriptor = request.into_inner();
445        let ticket = Ticket::new(query.as_any().encode_to_vec());
446        metadata_info(descriptor, ticket, query.into_builder(&SQL_INFO).schema())
447    }
448
449    async fn do_get_catalogs(
450        &self,
451        query: CommandGetCatalogs,
452        _request: Request<Ticket>,
453    ) -> Result<Response<BoxStream>, Status> {
454        let mut builder = query.into_builder();
455        builder.append(self.graph.as_str());
456        let schema = builder.schema();
457        Ok(Response::new(metadata_stream(schema, builder.build())))
458    }
459
460    async fn do_get_schemas(
461        &self,
462        query: CommandGetDbSchemas,
463        _request: Request<Ticket>,
464    ) -> Result<Response<BoxStream>, Status> {
465        let mut builder = query.into_builder();
466        builder.append(self.graph.as_str(), DB_SCHEMA);
467        let schema = builder.schema();
468        Ok(Response::new(metadata_stream(schema, builder.build())))
469    }
470
471    async fn do_get_tables(
472        &self,
473        query: CommandGetTables,
474        _request: Request<Ticket>,
475    ) -> Result<Response<BoxStream>, Status> {
476        let mut builder = query.into_builder();
477        for (name, table_schema) in self.registry.tables() {
478            builder
479                .append(
480                    self.graph.as_str(),
481                    DB_SCHEMA,
482                    &name,
483                    "TABLE",
484                    table_schema.as_ref(),
485                )
486                .map_err(|e| Status::internal(e.to_string()))?;
487        }
488        let schema = builder.schema();
489        Ok(Response::new(metadata_stream(schema, builder.build())))
490    }
491
492    async fn do_get_table_types(
493        &self,
494        query: CommandGetTableTypes,
495        _request: Request<Ticket>,
496    ) -> Result<Response<BoxStream>, Status> {
497        let mut builder = query.into_builder();
498        builder.append("TABLE");
499        let schema = builder.schema();
500        Ok(Response::new(metadata_stream(schema, builder.build())))
501    }
502
503    async fn do_get_sql_info(
504        &self,
505        query: CommandGetSqlInfo,
506        _request: Request<Ticket>,
507    ) -> Result<Response<BoxStream>, Status> {
508        let builder = query.into_builder(&SQL_INFO);
509        let schema = builder.schema();
510        Ok(Response::new(metadata_stream(schema, builder.build())))
511    }
512
513    // ── Prepared statements ──────────────────────────────────────────────────
514    //    Stateless: the handle *is* the resolved Cypher, so no server-side
515    //    statement table is needed and any node can serve any handle.
516
517    async fn do_action_create_prepared_statement(
518        &self,
519        query: ActionCreatePreparedStatementRequest,
520        request: Request<arrow_flight::Action>,
521    ) -> Result<ActionCreatePreparedStatementResult, Status> {
522        let cypher = self.target_cypher(&query.query);
523        self.authorize(&request, &cypher)?;
524        let schema = self.schema_for(&cypher).await?;
525        let IpcMessage(schema_bytes) =
526            SchemaAsIpc::new(schema.as_ref(), &IpcWriteOptions::default())
527                .try_into()
528                .map_err(|e: arrow::error::ArrowError| Status::internal(e.to_string()))?;
529        Ok(ActionCreatePreparedStatementResult {
530            prepared_statement_handle: cypher.into_bytes().into(),
531            dataset_schema: schema_bytes,
532            parameter_schema: Default::default(), // no bind parameters
533        })
534    }
535
536    async fn do_action_close_prepared_statement(
537        &self,
538        _query: ActionClosePreparedStatementRequest,
539        _request: Request<arrow_flight::Action>,
540    ) -> Result<(), Status> {
541        Ok(()) // stateless — nothing to release
542    }
543
544    async fn get_flight_info_prepared_statement(
545        &self,
546        cmd: CommandPreparedStatementQuery,
547        request: Request<FlightDescriptor>,
548    ) -> Result<Response<FlightInfo>, Status> {
549        let cypher = String::from_utf8(cmd.prepared_statement_handle.to_vec())
550            .map_err(|e| Status::invalid_argument(format!("invalid handle: {e}")))?;
551        self.authorize(&request, &cypher)?;
552        let schema = self.schema_for(&cypher).await?;
553
554        let ticket_bytes = TicketStatementQuery {
555            statement_handle: cypher.into_bytes().into(),
556        }
557        .encode_to_vec();
558
559        let info = FlightInfo::new()
560            .try_with_schema(schema.as_ref())
561            .map_err(|e| Status::internal(e.to_string()))?
562            .with_endpoint(FlightEndpoint::new().with_ticket(Ticket::new(ticket_bytes)))
563            .with_total_records(-1)
564            .with_total_bytes(-1)
565            .with_ordered(false);
566        Ok(Response::new(info))
567    }
568
569    async fn do_get_prepared_statement(
570        &self,
571        cmd: CommandPreparedStatementQuery,
572        request: Request<Ticket>,
573    ) -> Result<Response<BoxStream>, Status> {
574        let cypher = String::from_utf8(cmd.prepared_statement_handle.to_vec())
575            .map_err(|e| Status::invalid_argument(format!("invalid handle: {e}")))?;
576        self.authorize(&request, &cypher)?;
577        self.execute_cypher_stream(&cypher).await
578    }
579
580    async fn do_put_prepared_statement_query(
581        &self,
582        query: CommandPreparedStatementQuery,
583        _request: Request<PeekableFlightDataStream>,
584    ) -> Result<DoPutPreparedStatementResult, Status> {
585        // No bind parameters in the Cypher dialect — echo the handle back.
586        Ok(DoPutPreparedStatementResult {
587            prepared_statement_handle: Some(query.prepared_statement_handle),
588        })
589    }
590}
591
592const DB_SCHEMA: &str = "public";
593
594/// Materialise the static `GetSqlInfo` payload as a `RecordBatch`. Test-only
595/// hook so the lazily-built [`SQL_INFO`] can be asserted without a live server.
596#[cfg(test)]
597pub(crate) fn sql_info_batch() -> arrow_array::RecordBatch {
598    CommandGetSqlInfo::default()
599        .into_builder(&SQL_INFO)
600        .build()
601        .expect("sql info batch builds")
602}
603
604/// Build a `FlightInfo` for a metadata response: one endpoint whose ticket is
605/// the request command, advertising the builder's fixed schema.
606fn metadata_info(
607    descriptor: FlightDescriptor,
608    ticket: Ticket,
609    schema: arrow_schema::SchemaRef,
610) -> Result<Response<FlightInfo>, Status> {
611    let info = FlightInfo::new()
612        .try_with_schema(schema.as_ref())
613        .map_err(|e| Status::internal(e.to_string()))?
614        .with_endpoint(FlightEndpoint::new().with_ticket(ticket))
615        .with_descriptor(descriptor);
616    Ok(Response::new(info))
617}
618
619/// Encode a single metadata `RecordBatch` (already wrapped in a `FlightError`
620/// result by the builder) as a one-item Flight data stream.
621fn metadata_stream(
622    schema: arrow_schema::SchemaRef,
623    batch: std::result::Result<arrow_array::RecordBatch, FlightError>,
624) -> BoxStream {
625    Box::pin(
626        FlightDataEncoderBuilder::new()
627            .with_schema(schema)
628            .build(futures::stream::once(async move { batch }))
629            .map_err(|e: FlightError| Status::internal(e.to_string())),
630    )
631}
632
633// ── DoPut directive parsing ───────────────────────────────────────────────────
634
635enum PutAction {
636    Vertices {
637        label: String,
638    },
639    Edges {
640        src: String,
641        edge_type: String,
642        dst: String,
643    },
644}
645
646fn parse_put_directive(s: &str) -> Result<PutAction, FlightSqlError> {
647    if let Some(rest) = s.strip_prefix("VERTICES:") {
648        return Ok(PutAction::Vertices {
649            label: rest.to_string(),
650        });
651    }
652    if let Some(rest) = s.strip_prefix("EDGES:") {
653        let parts: Vec<&str> = rest.splitn(3, ':').collect();
654        if parts.len() == 3 {
655            return Ok(PutAction::Edges {
656                src: parts[0].to_string(),
657                edge_type: parts[1].to_string(),
658                dst: parts[2].to_string(),
659            });
660        }
661    }
662    Err(FlightSqlError::UnknownPutCommand(s.to_string()))
663}