graphar_flight/

registry.rs

use std::{
    collections::HashMap,
    sync::{Arc, RwLock},
};

use arrow_schema::{DataType, Field, Schema, SchemaRef};

/// Maps Cypher query strings to their expected Arrow return schemas.
///
/// ## Two modes
///
/// **Auto-string (default)** — leave the registry empty or only register the
/// queries you care about. Unregistered queries fall back to an all-`Utf8`
/// schema inferred from FalkorDB's response header at query time.
///
/// **Registered (typed)** — pre-declare the schema so columns like `_gar_id`,
/// `src`, and `dst` come back as `Int64` instead of strings. Use the
/// convenience helpers [`register_vertex_query`] and [`register_edge_query`]
/// for the common cases, or [`register`] for full control.
///
/// [`register_vertex_query`]: SchemaRegistry::register_vertex_query
/// [`register_edge_query`]: SchemaRegistry::register_edge_query
/// [`register`]: SchemaRegistry::register
///
/// ## Example — minimal setup (auto-string only)
///
/// ```rust
/// use std::sync::Arc;
/// use graphar_flight::SchemaRegistry;
///
/// // An empty registry is valid; all queries use auto-string mode.
/// let registry = Arc::new(SchemaRegistry::new());
/// ```
///
/// ## Example — typed vertex + edge schemas
///
/// ```rust
/// use std::sync::Arc;
/// use graphar_flight::SchemaRegistry;
///
/// let registry = Arc::new(SchemaRegistry::new());
///
/// // _gar_id → Int64, extra columns → Utf8.
/// registry.register_vertex_query(
///     "MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name, n.age",
///     &["name", "age"],
/// );
///
/// // src + dst → Int64, extra columns → Utf8.
/// registry.register_edge_query(
///     "MATCH (a:Person)-[r:KNOWS]->(b:Person) \
///      RETURN a._gar_id AS src, b._gar_id AS dst, r.since",
///     &["since"],
/// );
/// ```
#[derive(Debug, Default, Clone)]
pub struct SchemaRegistry {
    inner: Arc<RwLock<HashMap<String, SchemaRef>>>,
    /// Friendly table name → the Cypher query key it stands for. Lets the
    /// Flight SQL metadata surface advertise short table names (`GetTables`)
    /// while the query itself stays the registry key. Empty for queries
    /// registered without a name.
    names: Arc<RwLock<HashMap<String, String>>>,
}

impl SchemaRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    /// Register an arbitrary schema under `key`. Overwrites any existing entry.
    ///
    /// Use this when you need full control over column types. For the common
    /// vertex / edge patterns prefer [`register_vertex_query`] and
    /// [`register_edge_query`].
    ///
    /// [`register_vertex_query`]: SchemaRegistry::register_vertex_query
    /// [`register_edge_query`]: SchemaRegistry::register_edge_query
    ///
    /// ```rust
    /// use std::sync::Arc;
    /// use graphar_flight::SchemaRegistry;
    /// use arrow_schema::{DataType, Field, Schema};
    ///
    /// let registry = SchemaRegistry::new();
    ///
    /// registry.register(
    ///     "MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name, n.age",
    ///     Arc::new(Schema::new(vec![
    ///         Field::new("_gar_id", DataType::Int64, false),
    ///         Field::new("name",    DataType::Utf8,  true),
    ///         Field::new("age",     DataType::Int64, true),   // typed!
    ///     ])),
    /// );
    /// ```
    pub fn register(&self, key: impl Into<String>, schema: SchemaRef) -> &Self {
        self.inner.write().unwrap().insert(key.into(), schema);
        self
    }

    /// Register a **vertex** query schema.
    ///
    /// Adds `_gar_id: Int64` (non-nullable) as the first column, then appends
    /// each name in `extra_cols` as a nullable `Utf8` column.
    ///
    /// The `query` string must exactly match what the Flight SQL client sends.
    /// By convention, the `RETURN` clause should alias the internal ID as
    /// `_gar_id` and list the extra columns in the same order as `extra_cols`.
    ///
    /// ## Example
    ///
    /// ```rust
    /// use std::sync::Arc;
    /// use graphar_flight::SchemaRegistry;
    ///
    /// let registry = SchemaRegistry::new();
    ///
    /// // Person vertices — id as Int64, everything else as Utf8.
    /// registry.register_vertex_query(
    ///     "MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name, n.city",
    ///     &["name", "city"],
    /// );
    ///
    /// // No extra properties — only _gar_id.
    /// registry.register_vertex_query(
    ///     "MATCH (n:Company) RETURN n._gar_id AS _gar_id",
    ///     &[],
    /// );
    /// ```
    pub fn register_vertex_query(&self, query: &str, extra_cols: &[&str]) -> &Self {
        let mut fields = vec![Field::new("_gar_id", DataType::Int64, false)];
        fields.extend(
            extra_cols
                .iter()
                .map(|&n| Field::new(n, DataType::Utf8, true)),
        );
        self.register(query, Arc::new(Schema::new(fields)))
    }

    /// Register an **edge** query schema.
    ///
    /// Adds `src: Int64` and `dst: Int64` (both non-nullable) as the first two
    /// columns, then appends each name in `extra_cols` as a nullable `Utf8` column.
    ///
    /// The `RETURN` clause should alias the source `_gar_id` as `src` and the
    /// destination `_gar_id` as `dst`, matching GraphAr's internal IDs.
    ///
    /// ## Example
    ///
    /// ```rust
    /// use std::sync::Arc;
    /// use graphar_flight::SchemaRegistry;
    ///
    /// let registry = SchemaRegistry::new();
    ///
    /// // Edges with no properties — just src + dst.
    /// registry.register_edge_query(
    ///     "MATCH (a:Person)-[:FOLLOWS]->(b:Person) \
    ///      RETURN a._gar_id AS src, b._gar_id AS dst",
    ///     &[],
    /// );
    ///
    /// // Edges with a string property.
    /// registry.register_edge_query(
    ///     "MATCH (a:Person)-[r:KNOWS]->(b:Person) \
    ///      RETURN a._gar_id AS src, b._gar_id AS dst, r.since",
    ///     &["since"],
    /// );
    ///
    /// // Multiple edge properties.
    /// registry.register_edge_query(
    ///     "MATCH (a:Person)-[r:RATED]->(b:Movie) \
    ///      RETURN a._gar_id AS src, b._gar_id AS dst, r.score, r.review",
    ///     &["score", "review"],
    /// );
    /// ```
    pub fn register_edge_query(&self, query: &str, extra_cols: &[&str]) -> &Self {
        let mut fields = vec![
            Field::new("src", DataType::Int64, false),
            Field::new("dst", DataType::Int64, false),
        ];
        fields.extend(
            extra_cols
                .iter()
                .map(|&n| Field::new(n, DataType::Utf8, true)),
        );
        self.register(query, Arc::new(Schema::new(fields)))
    }

    /// Register a query under a friendly **table name**.
    ///
    /// The schema is stored under `query` (so direct execution of the Cypher
    /// still works) and `name` becomes an alias the Flight SQL metadata
    /// surface advertises in `GetTables`. A generic ADBC/ODBC client can then
    /// `SELECT * FROM <name>` and the server rewrites it to `query`.
    ///
    /// ```rust
    /// use graphar_flight::SchemaRegistry;
    /// use std::sync::Arc;
    /// use arrow_schema::{DataType, Field, Schema};
    ///
    /// let registry = SchemaRegistry::new();
    /// registry.register_named(
    ///     "people",
    ///     "MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name",
    ///     Arc::new(Schema::new(vec![
    ///         Field::new("_gar_id", DataType::Int64, false),
    ///         Field::new("name",    DataType::Utf8,  true),
    ///     ])),
    /// );
    /// assert_eq!(
    ///     registry.resolve_table("people").as_deref(),
    ///     Some("MATCH (n:Person) RETURN n._gar_id AS _gar_id, n.name"),
    /// );
    /// ```
    pub fn register_named(
        &self,
        name: impl Into<String>,
        query: impl Into<String>,
        schema: SchemaRef,
    ) -> &Self {
        let query = query.into();
        self.names
            .write()
            .unwrap()
            .insert(name.into(), query.clone());
        self.register(query, schema)
    }

    /// Look up a schema by exact key. Returns `None` for unregistered queries
    /// (the server then falls back to auto-string mode).
    pub fn get(&self, key: &str) -> Option<SchemaRef> {
        self.inner.read().unwrap().get(key).cloned()
    }

    /// List all registered query keys.
    pub fn keys(&self) -> Vec<String> {
        self.inner.read().unwrap().keys().cloned().collect()
    }

    /// The friendly table name a Cypher query was registered under, if any.
    /// The inverse of [`resolve_table`]: used by per-query authorization so a
    /// policy can name a query by its short table name. Returns `None` for a
    /// query registered without a name (the caller then keys on the Cypher).
    ///
    /// [`resolve_table`]: SchemaRegistry::resolve_table
    pub fn name_of(&self, query: &str) -> Option<String> {
        self.names
            .read()
            .unwrap()
            .iter()
            .find(|(_, q)| q.as_str() == query)
            .map(|(name, _)| name.clone())
    }

    /// Resolve a `GetTables`-style identifier (a friendly name or the Cypher
    /// key itself) to the Cypher query to execute. Returns `None` when neither
    /// a name alias nor a registered query matches.
    pub fn resolve_table(&self, ident: &str) -> Option<String> {
        if let Some(query) = self.names.read().unwrap().get(ident) {
            return Some(query.clone());
        }
        if self.inner.read().unwrap().contains_key(ident) {
            return Some(ident.to_string());
        }
        None
    }

    /// Every registered query as `(table_name, schema)`, for the Flight SQL
    /// `GetTables` surface. Named queries report their alias; queries
    /// registered without a name report the Cypher string as the table name.
    pub fn tables(&self) -> Vec<(String, SchemaRef)> {
        let names = self.names.read().unwrap();
        let query_to_name: HashMap<&str, &str> = names
            .iter()
            .map(|(n, q)| (q.as_str(), n.as_str()))
            .collect();
        self.inner
            .read()
            .unwrap()
            .iter()
            .map(|(query, schema)| {
                let table = query_to_name
                    .get(query.as_str())
                    .map(|n| n.to_string())
                    .unwrap_or_else(|| query.clone());
                (table, schema.clone())
            })
            .collect()
    }
}

/// Extract the table identifier from a trivial `SELECT … FROM <ident>` query,
/// stripping surrounding quotes/backticks. Returns `None` for anything that
/// isn't a single-table SELECT (the caller then treats the input as Cypher).
///
/// This is deliberately minimal — generic Flight SQL drivers (Power BI's ADBC
/// connector, `adbc_driver_flightsql`) emit `SELECT * FROM "table"` to read a
/// table they discovered through `GetTables`; that is the shape we rewrite.
pub fn parse_select_from(sql: &str) -> Option<String> {
    let lower = sql.to_ascii_lowercase();
    let from = lower.find(" from ")?;
    let after = sql[from + 6..].trim();
    // Take the first whitespace/`;`-delimited token after FROM.
    let token = after
        .split(|c: char| c.is_whitespace() || c == ';')
        .next()?
        .trim();
    if token.is_empty() {
        return None;
    }
    let unquoted = token.trim_matches('"').trim_matches('`').trim_matches('\'');
    Some(unquoted.to_string())
}