crates/graphar-flight/src/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()
}
}
/// What an incoming query string IS, as far as this engine can honour it.
///
/// graphar-flight executes Cypher. The only SQL it executes faithfully is the
/// one shape generic Flight SQL drivers (Power BI's ADBC connector,
/// `adbc_driver_flightsql`, the Dremio ODBC driver) emit to read a table they
/// discovered through `GetTables`: `SELECT * FROM <table>`, which is rewritten
/// to that table's registered Cypher.
///
/// # Why this is an enum and not an `Option<String>`
///
/// The previous parser grabbed the first token after `FROM` and returned it —
/// for `SELECT a, b FROM people WHERE a > 1` it returned `people`, and the
/// server sent the WHOLE table: projection dropped, predicate dropped, the
/// row count wrong, and no error. Power BI folds by default, so a card that
/// says `COUNT(*) … WHERE region = 'north'` showed the count of everything.
/// The user cannot see that from the dashboard. A query the engine cannot
/// honour has to be REJECTED, loudly, naming what would have been dropped —
/// that is [`SqlShape::Unhonoured`].
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SqlShape {
/// Not SQL (Cypher, or a string with no `FROM`): served verbatim as Cypher.
NotSql,
/// `SELECT * FROM <ident> [[AS] alias] [;]` — executed faithfully by the
/// rewrite. `table` is the identifier, quotes stripped, qualification kept
/// (see [`table_candidates`]).
Honoured { table: String },
/// SQL whose result would NOT be what the client asked for: the clauses
/// the rewrite would have silently thrown away, in source order. `table`
/// is empty when the statement has no single FROM table (a CTE).
Unhonoured { table: String, dropped: Vec<String> },
}
/// Classify `sql` — see [`SqlShape`]. Pure; no engine, no registry.
pub fn classify_sql(sql: &str) -> SqlShape {
let trimmed = sql.trim().trim_end_matches(';').trim();
let lower = trimmed.to_ascii_lowercase();
let mut dropped: Vec<String> = Vec::new();
// `WITH … SELECT` is SQL the rewrite has no arm for at all.
if lower.starts_with("with") && lower[4..].starts_with(|c: char| c.is_whitespace() || c == '(') {
// The whole statement is the CTE's; there is no one table to name.
return SqlShape::Unhonoured {
table: String::new(),
dropped: vec!["WITH (common table expression)".into()],
};
}
let Some(from) = find_keyword(&lower, "from") else {
return SqlShape::NotSql;
};
let Some(select) = find_keyword(&lower, "select") else {
return SqlShape::NotSql;
};
if select > from {
return SqlShape::NotSql;
}
// ── between SELECT and FROM: the projection ──
let projection = trimmed[select + 6..from].trim();
let proj_lower = projection.to_ascii_lowercase();
let projection = if let Some(rest) = proj_lower.strip_prefix("distinct") {
if rest.starts_with(|c: char| c.is_whitespace()) {
dropped.push("DISTINCT".into());
projection[8..].trim()
} else {
projection
}
} else {
projection
};
if projection != "*" {
dropped.push(format!("projection `{projection}` (only `*` is served)"));
}
// ── after FROM: the identifier, then whatever else ──
let after = trimmed[from + 4..].trim();
// The identifier ends at whitespace OR a comma: `people, cities` names two
// tables, and a comma glued to the first must not become part of its name.
let ident_len = after
.find(|c: char| c.is_whitespace() || c == ',')
.unwrap_or(after.len());
let ident = &after[..ident_len];
if ident.is_empty() {
return SqlShape::NotSql;
}
let table = unquote_qualified(ident);
let rest = after[ident_len..].trim();
if !rest.is_empty() && !is_bare_alias(rest) {
let rest_lower = rest.to_ascii_lowercase();
let mut found = false;
for (kw, label) in [
("join", "JOIN"),
("where", "WHERE"),
("group", "GROUP BY"),
("having", "HAVING"),
("order", "ORDER BY"),
("limit", "LIMIT"),
("offset", "OFFSET"),
("fetch", "FETCH"),
("union", "UNION"),
("intersect", "INTERSECT"),
("except", "EXCEPT"),
] {
if find_keyword(&rest_lower, kw).is_some() {
dropped.push(label.into());
found = true;
}
}
if rest.starts_with(',') {
dropped.push("a second FROM table (comma join)".into());
found = true;
}
if !found {
dropped.push(format!("trailing `{rest}`"));
}
}
if dropped.is_empty() {
SqlShape::Honoured { table }
} else {
SqlShape::Unhonoured { table, dropped }
}
}
/// `[AS] <identifier>` and nothing else — an alias changes no result.
fn is_bare_alias(rest: &str) -> bool {
let mut words = rest.split_whitespace();
let Some(first) = words.next() else { return false };
let alias = if first.eq_ignore_ascii_case("as") {
match words.next() {
Some(a) => a,
None => return false,
}
} else {
first
};
if words.next().is_some() {
return false;
}
let a = alias.trim_matches('"').trim_matches('`');
!a.is_empty()
&& a.chars().next().is_some_and(|c| c.is_alphabetic() || c == '_')
&& a.chars().all(|c| c.is_alphanumeric() || c == '_')
}
/// Byte offset of `kw` in `lower` as a whole word (bounded by whitespace,
/// parentheses, commas, or the string ends), or `None`.
fn find_keyword(lower: &str, kw: &str) -> Option<usize> {
let mut start = 0;
while let Some(pos) = lower[start..].find(kw) {
let at = start + pos;
let before_ok = at == 0
|| lower[..at]
.chars()
.next_back()
.is_some_and(|c| c.is_whitespace() || c == '(' || c == ',' || c == ')');
let end = at + kw.len();
let after_ok = end == lower.len()
|| lower[end..]
.chars()
.next()
.is_some_and(|c| c.is_whitespace() || c == '(' || c == ',' || c == ';');
if before_ok && after_ok {
return Some(at);
}
start = end;
}
None
}
/// Strip `"`, `` ` `` and `'` from every dot-separated segment of an
/// identifier, keeping the dots: `"public"."people"` → `public.people`.
fn unquote_qualified(ident: &str) -> String {
ident
.split('.')
.map(|seg| seg.trim_matches('"').trim_matches('`').trim_matches('\''))
.collect::<Vec<_>>()
.join(".")
}
/// The registry names a (possibly qualified) identifier may be registered
/// under, most specific first: the full dotted name, then the last segment.
///
/// Generic drivers qualify what `GetTables` gave them (`public.people`,
/// `"public"."people"`, `` `public`.`people` ``); the registry holds the bare
/// table name. Both must reach the same Cypher.
pub fn table_candidates(table: &str) -> Vec<String> {
let mut v = vec![table.to_string()];
if let Some((_, last)) = table.rsplit_once('.') {
if !last.is_empty() {
v.push(last.to_string());
}
}
v
}