diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index c4ef09cda..805934703 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -44,6 +44,6 @@ aws-lc-rs = "=1.16.3" napi-build = "2.3.1" [features] -default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface", "lancedb/goosefs", "lancedb/metrics-otel"] +default = ["remote", "lancedb/sql", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface", "lancedb/goosefs", "lancedb/metrics-otel"] fp16kernels = ["lancedb/fp16kernels"] remote = ["lancedb/remote"] diff --git a/python/Cargo.toml b/python/Cargo.toml index e0c72df48..9087beb23 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -47,6 +47,6 @@ libc = "0.2" pyo3-build-config = { version = "0.28", features = ["abi3-py310"] } [features] -default = ["remote", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface", "lancedb/cos", "lancedb/goosefs", "lancedb/metrics-otel"] +default = ["remote", "lancedb/sql", "lancedb/aws", "lancedb/gcs", "lancedb/azure", "lancedb/dynamodb", "lancedb/oss", "lancedb/huggingface", "lancedb/cos", "lancedb/goosefs", "lancedb/metrics-otel"] fp16kernels = ["lancedb/fp16kernels"] remote = ["lancedb/remote"] diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index c8f17d50c..7a7d111fa 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -120,7 +120,13 @@ pprof = { version = "0.14", features = ["flamegraph"] } [features] -default = [] +default = ["sql"] +# The SQL statement extension seam (`lancedb::sql`): the registry a host adds +# dialect statements through, the access vocabulary those statements declare, +# and the write-commit observer. It pulls in no dependency that is not already +# required, so it is on by default; the flag exists so an embedder that does +# not want the surface can opt out of it. +sql = [] aws = [ "lance/aws", "lance-io/aws", diff --git a/rust/lancedb/src/sql/dml.rs b/rust/lancedb/src/sql/dml.rs new file mode 100644 index 000000000..6918b67ba --- /dev/null +++ b/rust/lancedb/src/sql/dml.rs @@ -0,0 +1,150 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! The result of a DML statement, as a one-row record batch. +//! +//! A DML statement has to answer over the same channel a query does, so its +//! result is carried as an ordinary [`RecordBatch`] with a fixed schema. The +//! round trip is lossless, which is what lets a caller recover the typed form +//! after the batch has crossed a transport such as Arrow Flight. + +use std::fmt; +use std::sync::Arc; + +use arrow_array::{Int64Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; + +/// Which DML statement produced a result. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DmlOperation { + Insert, + Update, + Delete, +} + +impl DmlOperation { + pub fn as_str(self) -> &'static str { + match self { + Self::Insert => "INSERT", + Self::Update => "UPDATE", + Self::Delete => "DELETE", + } + } + + fn parse(s: &str) -> Option { + match s { + "INSERT" => Some(Self::Insert), + "UPDATE" => Some(Self::Update), + "DELETE" => Some(Self::Delete), + _ => None, + } + } +} + +impl fmt::Display for DmlOperation { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +/// What a DML statement did. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DmlResult { + pub table: String, + pub operation: DmlOperation, + pub rows_affected: i64, + pub version: i64, +} + +/// The schema every [`DmlResult`] batch carries. +pub fn dml_result_schema() -> Schema { + Schema::new(vec![ + Field::new("table", DataType::Utf8, false), + Field::new("operation", DataType::Utf8, false), + Field::new("rows_affected", DataType::Int64, false), + Field::new("version", DataType::Int64, false), + ]) +} + +impl DmlResult { + pub fn new( + table: impl Into, + operation: DmlOperation, + rows_affected: i64, + version: i64, + ) -> Self { + Self { + table: table.into(), + operation, + rows_affected, + version, + } + } + + pub fn to_record_batch(&self) -> RecordBatch { + RecordBatch::try_new( + Arc::new(dml_result_schema()), + vec![ + Arc::new(StringArray::from(vec![self.table.as_str()])), + Arc::new(StringArray::from(vec![self.operation.as_str()])), + Arc::new(Int64Array::from(vec![self.rows_affected])), + Arc::new(Int64Array::from(vec![self.version])), + ], + ) + .expect("static schema") + } + + /// Recover a result from a batch, or `None` if the batch is not one. + /// + /// A query result can arrive on the same channel, so this has to be able + /// to say "not a DML result" rather than fail. + pub fn try_from_batch(batch: &RecordBatch) -> Option { + if *batch.schema().as_ref() != dml_result_schema() || batch.num_rows() != 1 { + return None; + } + Some(Self { + table: batch + .column(0) + .as_any() + .downcast_ref::()? + .value(0) + .to_string(), + operation: DmlOperation::parse( + batch + .column(1) + .as_any() + .downcast_ref::()? + .value(0), + )?, + rows_affected: batch + .column(2) + .as_any() + .downcast_ref::()? + .value(0), + version: batch + .column(3) + .as_any() + .downcast_ref::()? + .value(0), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn roundtrip() { + let original = DmlResult::new("foo", DmlOperation::Insert, 1, 3); + let batch = original.to_record_batch(); + assert_eq!(DmlResult::try_from_batch(&batch), Some(original)); + } + + #[test] + fn non_dml_returns_none() { + let s = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let batch = RecordBatch::try_new(s, vec![Arc::new(Int64Array::from(vec![1]))]).unwrap(); + assert_eq!(DmlResult::try_from_batch(&batch), None); + } +} diff --git a/rust/lancedb/src/sql.rs b/rust/lancedb/src/sql/mod.rs similarity index 78% rename from rust/lancedb/src/sql.rs rename to rust/lancedb/src/sql/mod.rs index 7c040c51b..4a6b63659 100644 --- a/rust/lancedb/src/sql.rs +++ b/rust/lancedb/src/sql/mod.rs @@ -1,7 +1,35 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! Handles to SQL queries running on a remote database. +//! SQL: handles to queries running on a remote database, and the seam an +//! embedder extends the dialect through. +//! +//! The extension seam is behind the default-on `sql` feature. It lets a host +//! add statements to the dialect from outside this crate: a statement brings +//! its own grammar ([`CustomSqlHandler`]), and declares its audit label and +//! the access it needs ([`SqlStatement`]), so the host's authorization and +//! auditing do not have to know each statement by name. + +#[cfg(feature = "sql")] +mod dml; +#[cfg(feature = "sql")] +mod observer; +#[cfg(feature = "sql")] +mod parser; +#[cfg(feature = "sql")] +mod statement; + +#[cfg(feature = "sql")] +pub use dml::{DmlOperation, DmlResult, dml_result_schema}; +#[cfg(feature = "sql")] +pub use observer::{CommittedWrite, DmlEventKind, WriteObserver, observe_write}; +#[cfg(feature = "sql")] +pub use parser::route_custom_sql; +#[cfg(feature = "sql")] +pub use statement::{ + AccessRequirement, CreateKind, CustomSqlHandler, DatabaseScope, RelationKind, + RequirementContext, SqlStatement, StatementRegistry, SystemScope, WriteMode, +}; use std::{fmt, sync::Arc}; diff --git a/rust/lancedb/src/sql/observer.rs b/rust/lancedb/src/sql/observer.rs new file mode 100644 index 000000000..0ffc89eb9 --- /dev/null +++ b/rust/lancedb/src/sql/observer.rs @@ -0,0 +1,70 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Notification of committed writes. +//! +//! A statement that writes rows often has to tell its host that it did, so +//! that follow-up work can be scheduled. What it should *not* have to know is +//! how the host represents that notification. [`WriteObserver`] is the seam: +//! the statement reports what it wrote, and the host decides what that means +//! -- an event on a bus, a metric, or nothing at all. + +use std::sync::Arc; + +use async_trait::async_trait; + +/// Which DML operation committed. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DmlEventKind { + Insert, + Update, + Delete, +} + +/// A write that has already been made durable. +#[derive(Debug, Clone)] +pub struct CommittedWrite { + /// The database holding the table. + pub database: String, + /// The schema the statement named the table through. + pub schema: String, + /// The table written. + pub table: String, + /// The table's storage location, when the statement resolved one. + pub table_uri: Option, + /// Which operation committed. + pub kind: DmlEventKind, +} + +/// Notified after a statement's write commits. +/// +/// Implementations are best-effort by contract: the write is already durable +/// when this is called, so an observer that fails must not fail the statement. +/// That is why the method cannot report an error. +#[async_trait] +pub trait WriteObserver: Send + Sync { + async fn write_committed(&self, write: CommittedWrite); +} + +/// Report a committed write, if anything is observing. +pub async fn observe_write( + observer: Option<&Arc>, + database: &str, + schema: &str, + table: &str, + table_uri: Option, + kind: DmlEventKind, +) { + let Some(observer) = observer else { + return; + }; + observer + .write_committed(CommittedWrite { + database: database.to_string(), + schema: schema.to_string(), + table: table.to_string(), + table_uri, + kind, + }) + .await; +} diff --git a/rust/lancedb/src/sql/parser.rs b/rust/lancedb/src/sql/parser.rs new file mode 100644 index 000000000..de69d42c9 --- /dev/null +++ b/rust/lancedb/src/sql/parser.rs @@ -0,0 +1,174 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Routing a statement to the grammar that owns it. + +use datafusion::error::{DataFusionError, Result}; +use datafusion::logical_expr::LogicalPlan; +use datafusion::sql::sqlparser::{ + dialect::GenericDialect, + parser::{Parser, ParserError}, + tokenizer::{Token, Tokenizer, TokenizerError}, +}; + +use super::statement::StatementRegistry; + +/// Route a statement through the registry's grammars. +/// +/// Returns `Ok(None)` when no grammar claims the statement, which is the +/// caller's cue to hand it to DataFusion's own planner. +/// +/// The first grammar whose `matches` accepts the tokens is the only one given +/// the statement: a grammar that matches and then returns `Ok(None)` declines +/// the form rather than falling through to the next grammar. Registration +/// order therefore decides reachability, which is why [`StatementRegistry`] +/// fixes it explicitly. +pub fn route_custom_sql(registry: &StatementRegistry, sql: &str) -> Result> { + let dialect = GenericDialect {}; + let mut tokenizer = Tokenizer::new(&dialect, sql); + let tokens = tokenizer.tokenize().map_err(|e: TokenizerError| { + DataFusionError::SQL(Box::new(ParserError::TokenizerError(e.to_string())), None) + })?; + + // Handlers match on keywords, so layout must not change the decision. + let word_tokens: Vec<&Token> = tokens + .iter() + .filter(|t| !matches!(t, Token::Whitespace(_))) + .collect(); + + for handler in registry.parsers() { + if handler.matches(&word_tokens) { + // `Parser` takes ownership of the tokens, so it is built only once + // a handler has claimed the statement. + let mut parser = Parser::new(&dialect).with_tokens(tokens.clone()); + if let Some(plan) = handler.parse(&mut parser)? { + return Ok(Some(plan)); + } + break; + } + } + + Ok(None) +} + +#[cfg(test)] +mod tests { + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + + use datafusion::common::DFSchema; + use datafusion::error::Result as DfResult; + use datafusion::logical_expr::{EmptyRelation, LogicalPlan}; + use datafusion::sql::sqlparser::keywords::Keyword; + + use super::*; + use crate::sql::statement::CustomSqlHandler; + + fn empty_plan() -> LogicalPlan { + LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(DFSchema::empty()), + }) + } + + /// Matches on a leading keyword, and reports whether it was asked to parse. + struct Handler { + keyword: Keyword, + outcome: Outcome, + parsed: Arc, + } + + enum Outcome { + Plans, + Declines, + } + + impl Handler { + fn new(keyword: Keyword, outcome: Outcome) -> (Arc, Arc) { + let parsed = Arc::new(AtomicUsize::new(0)); + let handler = Arc::new(Self { + keyword, + outcome, + parsed: parsed.clone(), + }); + (handler, parsed) + } + } + + impl CustomSqlHandler for Handler { + fn matches(&self, tokens: &[&Token]) -> bool { + matches!(tokens.first(), Some(Token::Word(w)) if w.keyword == self.keyword) + } + + fn parse(&self, _parser: &mut Parser) -> DfResult> { + self.parsed.fetch_add(1, Ordering::SeqCst); + Ok(match self.outcome { + Outcome::Plans => Some(empty_plan()), + Outcome::Declines => None, + }) + } + } + + #[test] + fn an_unclaimed_statement_is_left_for_datafusion() { + let registry = StatementRegistry::new(); + assert!(route_custom_sql(®istry, "SELECT 1").unwrap().is_none()); + } + + #[test] + fn whitespace_does_not_change_which_handler_matches() { + let (handler, parsed) = Handler::new(Keyword::EXPLAIN, Outcome::Plans); + let mut registry = StatementRegistry::new(); + registry.register_parser(handler); + + for sql in ["EXPLAIN t", " EXPLAIN\n\t t "] { + assert!(route_custom_sql(®istry, sql).unwrap().is_some()); + } + assert_eq!(parsed.load(Ordering::SeqCst), 2); + } + + /// Front-insertion is what lets an extension get ahead of a catch-all that + /// would otherwise swallow the same keyword. + #[test] + fn the_last_registered_handler_is_consulted_first() { + let (first, first_parsed) = Handler::new(Keyword::EXPLAIN, Outcome::Plans); + let (second, second_parsed) = Handler::new(Keyword::EXPLAIN, Outcome::Plans); + + let mut registry = StatementRegistry::new(); + registry.register_parser(first).register_parser(second); + + assert!(route_custom_sql(®istry, "EXPLAIN t").unwrap().is_some()); + assert_eq!(second_parsed.load(Ordering::SeqCst), 1); + assert_eq!(first_parsed.load(Ordering::SeqCst), 0); + } + + /// A handler that matches and declines vetoes the statement rather than + /// letting a later handler see it. Shadowing is silent, which is why + /// registration order is part of the contract. + #[test] + fn a_handler_that_declines_shadows_the_handlers_behind_it() { + let (shadowed, shadowed_parsed) = Handler::new(Keyword::EXPLAIN, Outcome::Plans); + let (decliner, decliner_parsed) = Handler::new(Keyword::EXPLAIN, Outcome::Declines); + + let mut registry = StatementRegistry::new(); + registry.register_parser(shadowed).register_parser(decliner); + + assert!(route_custom_sql(®istry, "EXPLAIN t").unwrap().is_none()); + assert_eq!(decliner_parsed.load(Ordering::SeqCst), 1); + assert_eq!(shadowed_parsed.load(Ordering::SeqCst), 0); + } + + #[test] + fn from_parts_keeps_the_order_it_was_given() { + let (first, first_parsed) = Handler::new(Keyword::EXPLAIN, Outcome::Plans); + let (second, second_parsed) = Handler::new(Keyword::EXPLAIN, Outcome::Plans); + + let registry = StatementRegistry::from_parts(vec![first, second], vec![]); + + assert!(route_custom_sql(®istry, "EXPLAIN t").unwrap().is_some()); + assert_eq!(first_parsed.load(Ordering::SeqCst), 1); + assert_eq!(second_parsed.load(Ordering::SeqCst), 0); + } +} diff --git a/rust/lancedb/src/sql/statement.rs b/rust/lancedb/src/sql/statement.rs new file mode 100644 index 000000000..5bc521dae --- /dev/null +++ b/rust/lancedb/src/sql/statement.rs @@ -0,0 +1,365 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! The statement registry: the seam between the SQL dialect and the behaviour +//! an embedder adds to it. +//! +//! A statement owns three things that are otherwise easy to spread across +//! parallel `downcast_ref` chains: the grammar that produces its node, the +//! audit label it reports, and the access it requires. Keeping them in one +//! place is what makes a statement addable from outside this crate. +//! +//! Two registries, because the axes differ. Grammar is matched against tokens +//! before a node exists, and several statements can share one handler -- an +//! `ALTER TABLE` handler may yield a different node per subcommand. A planned +//! node, by contrast, is claimed by exactly one statement. + +use std::any::Any; +use std::sync::Arc; + +use datafusion::common::{ResolvedTableReference, TableReference}; +use datafusion::error::Result as DfResult; +use datafusion::logical_expr::LogicalPlan; +use datafusion::sql::sqlparser::{parser::Parser, tokenizer::Token}; + +/// A pluggable handler for custom SQL statements. +pub trait CustomSqlHandler: Send + Sync { + /// Whether this handler wants to handle these tokens. + /// + /// The tokens have had whitespace removed, so a handler can match on + /// leading keywords without accounting for layout. + fn matches(&self, tokens: &[&Token]) -> bool; + + /// Parse the statement into a logical plan. + /// + /// Returning `Ok(None)` declines a form this handler matched on; the + /// statement then goes to DataFusion's own planner. See + /// [`StatementRegistry`] for why that stops routing rather than falling + /// through to the next handler. + fn parse(&self, parser: &mut Parser) -> DfResult>; +} + +/// What kind of relation a requirement is about. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RelationKind { + Table, + View, +} + +/// What kind of object a DDL statement brings into existence. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CreateKind { + Table, + View, + MaterializedView, +} + +/// What a statement needs authorized before it runs. +/// +/// The vocabulary is deliberately generic: it names *what is being reached +/// for*, not the privilege that grants it. An embedder maps these onto its own +/// privilege model and audit labels, so no access-control concept has to live +/// in the dialect. +/// +/// The variants are finer-grained than a bare read/write split because the +/// distinctions are load-bearing for that mapping -- appending to a table and +/// redefining it are different grants, and collapsing them would silently +/// widen what a statement is allowed to do. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum AccessRequirement { + /// Read the contents of a relation. + Read { + relation: ResolvedTableReference, + kind: RelationKind, + }, + /// Change the rows of a relation. + Write { + relation: ResolvedTableReference, + kind: RelationKind, + mode: WriteMode, + }, + /// Change a relation's definition, or anything about it other than its + /// rows. Index and column changes land here. + Own { + relation: ResolvedTableReference, + kind: RelationKind, + }, + /// Bring a new relation into existence. + CreateIn { + relation: ResolvedTableReference, + kind: CreateKind, + }, + /// Reach the connected database itself rather than a relation in it. + Database { name: String, scope: DatabaseScope }, + /// Reach a namespace's metadata. + Namespace { database: String, namespace: String }, + /// Reach the deployment rather than any one database. + System { scope: SystemScope }, +} + +/// How an [`AccessRequirement::Write`] changes a relation's rows. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum WriteMode { + /// Add rows. + Append, + /// Change existing rows. + Modify, + /// Take rows away. + Remove, +} + +/// How far into a database an [`AccessRequirement::Database`] reaches. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DatabaseScope { + /// See that the database exists and list what is in it. + Usage, + /// Change what the database contains. + Ownership, +} + +/// How far into the deployment an [`AccessRequirement::System`] reaches. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SystemScope { + /// Observe deployment-wide state. + Usage, + /// Act on deployment-wide state. + Operate, +} + +impl AccessRequirement { + /// The relation this requirement is about, for the relation-shaped + /// variants. + /// + /// An embedder's privilege mapping is written against the variants + /// directly; this is the shortcut for the common case of needing the + /// relation without caring which shape asked for it. + pub fn relation(&self) -> Option<(&ResolvedTableReference, RelationKind)> { + match self { + Self::Read { relation, kind } + | Self::Write { relation, kind, .. } + | Self::Own { relation, kind } => Some((relation, *kind)), + Self::CreateIn { .. } + | Self::Database { .. } + | Self::Namespace { .. } + | Self::System { .. } => None, + } + } +} + +/// Everything a statement needs in order to state its requirements, without +/// reaching for the engine running it. +pub struct RequirementContext<'a> { + /// The database a bare relation name resolves against. + pub default_database: &'a str, + /// The schema a bare relation name resolves against. + pub default_schema: &'a str, +} + +impl RequirementContext<'_> { + /// Resolve a possibly-bare reference against the request's defaults. + pub fn resolve(&self, relation: TableReference) -> ResolvedTableReference { + relation.resolve(self.default_database, self.default_schema) + } + + /// Resolve a bare relation name against the request's defaults. + pub fn resolve_bare(&self, name: impl Into) -> ResolvedTableReference { + self.resolve(TableReference::bare(name.into())) + } +} + +/// One statement in the dialect: the node it plans to, what it is called in an +/// audit log, and what it needs authorized. +pub trait SqlStatement: Send + Sync { + /// Whether this statement owns the given planned node. + fn claims(&self, node: &dyn Any) -> bool; + + /// The audit label for this statement. + /// + /// This is an open string rather than an enum so that an embedder can add + /// a statement -- and a label for it -- without changing this crate. + fn audit_operation(&self) -> &'static str; + + /// What must be authorized before the node runs. + /// + /// Returning an empty set means the statement needs nothing beyond + /// whatever the engine already collects from the plan's scans. + fn access_requirements( + &self, + node: &dyn Any, + context: &RequirementContext<'_>, + ) -> DfResult>; +} + +/// The set of statements and grammars an engine knows about. +/// +/// Ordering is load-bearing on the parse side and stays explicit. A handler +/// may be a catch-all over its leading keyword -- erroring on any form of that +/// keyword it does not recognize, or matching on the first token alone -- so a +/// handler registered *after* such a one can never be reached for that +/// keyword. Extensions are therefore consulted before whatever is already +/// registered. +/// +/// A handler that matches and then returns `Ok(None)` stops routing entirely +/// rather than falling through to the next handler; the statement then goes to +/// DataFusion's own planner. That veto is intentional -- it is how a handler +/// declines a form it matched on -- but it means an overlapping handler +/// registered later is shadowed rather than reported, which is the other +/// reason ordering is explicit here. +#[derive(Default)] +pub struct StatementRegistry { + parsers: Vec>, + statements: Vec>, +} + +impl StatementRegistry { + /// An empty registry. + pub fn new() -> Self { + Self::default() + } + + /// Build a registry from an explicit, already-ordered set. + /// + /// The ordering is used as given -- unlike [`Self::register_parser`], this + /// does not reverse anything. It is how an embedder that owns the whole + /// dialect states the order once. + pub fn from_parts( + parsers: Vec>, + statements: Vec>, + ) -> Self { + Self { + parsers, + statements, + } + } + + /// Add a grammar, consulted before every grammar already registered. + /// + /// Registration is front-insertion because an existing catch-all handler + /// would otherwise shadow anything added later; see the type docs. + pub fn register_parser(&mut self, parser: Arc) -> &mut Self { + self.parsers.insert(0, parser); + self + } + + /// Add a statement, consulted before every statement already registered. + pub fn register_statement(&mut self, statement: Arc) -> &mut Self { + self.statements.insert(0, statement); + self + } + + /// The grammars, in the order they are consulted. + pub fn parsers(&self) -> &[Arc] { + &self.parsers + } + + /// The statement owning this planned node, if any. + pub fn claim(&self, node: &dyn Any) -> Option<&Arc> { + self.statements.iter().find(|s| s.claims(node)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct Claimant { + label: &'static str, + claims_everything: bool, + } + + impl SqlStatement for Claimant { + fn claims(&self, node: &dyn Any) -> bool { + self.claims_everything && node.is::() + } + + fn audit_operation(&self) -> &'static str { + self.label + } + + fn access_requirements( + &self, + _node: &dyn Any, + _context: &RequirementContext<'_>, + ) -> DfResult> { + Ok(vec![]) + } + } + + fn claimant(label: &'static str, claims_everything: bool) -> Arc { + Arc::new(Claimant { + label, + claims_everything, + }) + } + + #[test] + fn an_unclaimed_node_has_no_statement() { + let mut registry = StatementRegistry::new(); + registry.register_statement(claimant("never", false)); + assert!(registry.claim(&0u8).is_none()); + } + + /// Front-insertion on the claim side too: an extension must be able to + /// take over a node shape that something already registered also claims. + #[test] + fn the_last_registered_statement_claims_first() { + let mut registry = StatementRegistry::new(); + registry + .register_statement(claimant("first", true)) + .register_statement(claimant("second", true)); + + assert_eq!(registry.claim(&0u8).unwrap().audit_operation(), "second"); + } + + #[test] + fn from_parts_keeps_the_claim_order_it_was_given() { + let registry = + StatementRegistry::from_parts(vec![], vec![claimant("a", true), claimant("b", true)]); + assert_eq!(registry.claim(&0u8).unwrap().audit_operation(), "a"); + } + + #[test] + fn a_bare_name_resolves_against_the_request_defaults() { + let context = RequirementContext { + default_database: "db", + default_schema: "public", + }; + let resolved = context.resolve_bare("t"); + assert_eq!(&*resolved.catalog, "db"); + assert_eq!(&*resolved.schema, "public"); + assert_eq!(&*resolved.table, "t"); + } + + #[test] + fn a_qualified_name_keeps_its_own_parts() { + let context = RequirementContext { + default_database: "db", + default_schema: "public", + }; + let resolved = context.resolve(TableReference::partial("other", "t")); + assert_eq!(&*resolved.catalog, "db"); + assert_eq!(&*resolved.schema, "other"); + } + + #[test] + fn only_the_relation_shaped_requirements_name_a_relation() { + let relation = TableReference::bare("t").resolve("db", "public"); + + let read = AccessRequirement::Read { + relation: relation.clone(), + kind: RelationKind::Table, + }; + assert_eq!(read.relation().unwrap().1, RelationKind::Table); + + let create = AccessRequirement::CreateIn { + relation, + kind: CreateKind::MaterializedView, + }; + assert!(create.relation().is_none()); + + let system = AccessRequirement::System { + scope: SystemScope::Operate, + }; + assert!(system.relation().is_none()); + } +}