graphar-flight 0.1.2

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, parse_select_from},
};

/// 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,
}

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(),
        })
    }

    /// 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
    }

    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.
    ///
    /// Order: a registered key or table-name alias wins; otherwise a trivial
    /// `SELECT * FROM <table>` (what generic ADBC/ODBC drivers emit) is
    /// rewritten to the aliased Cypher; otherwise the string is taken as
    /// Cypher verbatim and served through the auto-string path.
    fn target_cypher(&self, incoming: &str) -> String {
        if let Some(query) = self.registry.resolve_table(incoming) {
            return query;
        }
        if let Some(ident) = parse_select_from(incoming)
            && let Some(query) = self.registry.resolve_table(&ident)
        {
            return query;
        }
        incoming.to_string()
    }

    /// 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 both the response payload and the
    // `authorization` trailer) 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)?;

        let token = self.auth.issued_token().unwrap_or("").to_string();
        let response_msg = HandshakeResponse {
            protocol_version: 0,
            payload: token.clone().into(),
        };
        let output = futures::stream::once(async move { Ok(response_msg) });
        let mut response: Response<Pin<Box<dyn Stream<Item = _> + Send>>> =
            Response::new(Box::pin(output));
        if !token.is_empty() {
            let bearer = format!("Bearer {token}")
                .parse()
                .map_err(|_| Status::internal("token not header-safe"))?;
            response.metadata_mut().insert("authorization", bearer);
        }
        Ok(response)
    }

    // ── 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);
            }
        };

        let ticket_bytes = TicketStatementQuery {
            statement_handle: cypher.into_bytes().into(),
        }
        .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_total_records(-1)
            .with_total_bytes(-1)
            .with_ordered(false);

        Ok(Response::new(info))
    }

    // ── 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.
        let raw = request.into_inner().into_inner(); // Streaming<FlightData>
        let batches = FlightRecordBatchStream::new_from_flight_data(
            raw.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 mut loader = self
            .new_loader()
            .await
            .map_err(|e| Status::internal(e.to_string()))?;

        let n = match action {
            PutAction::Vertices { label } => loader
                .insert_vertices(&label, &batches)
                .await
                .map_err(|e| Status::internal(e.to_string()))?,
            PutAction::Edges {
                src,
                edge_type,
                dst,
            } => loader
                .insert_edges(&src, &edge_type, &dst, &batches)
                .await
                .map_err(|e| Status::internal(e.to_string()))?,
        };

        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?;

        let ticket_bytes = TicketStatementQuery {
            statement_handle: cypher.into_bytes().into(),
        }
        .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_total_records(-1)
            .with_total_bytes(-1)
            .with_ordered(false);
        Ok(Response::new(info))
    }

    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),
        })
    }
}

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 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()))
}