feat(sql): add a statement extension seam behind the sql feature

An embedder can extend the SQL dialect with its own statements, but has
had nowhere to say what a statement *is*: its grammar, the label it
reports to an audit log, and the access it needs are three separate
facts, and describing them separately means a host ends up with parallel
downcast chains that must be kept in step by hand. Adding a statement to
only two of the three is a silent gap rather than a compile error.

`lancedb::sql` gains the seam that keeps them together:

- `CustomSqlHandler` contributes a grammar; `route_custom_sql` picks the
  one that owns a statement and leaves the rest to DataFusion.
- `SqlStatement` pairs a planned node with its audit label and the
  `AccessRequirement`s it needs. The vocabulary names what is reached
  for -- read, write, own, create, database, namespace, system -- rather
  than a privilege, so no access-control model has to live in the
  dialect.
- `StatementRegistry` holds both, with the consultation order it is
  given. Registration is front-insertion so an extension can get ahead
  of a catch-all that would otherwise swallow its keyword.
- `WriteObserver` reports a committed write without the statement
  knowing how its host represents that.
- `DmlResult` carries what a DML statement did as a one-row batch.

The feature adds no dependency that is not already required, so it is on
by default; the flag is there so an embedder that does not want the
surface can opt out.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

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"]

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"]

python/python/lancedb/remote/table.py

@@ -1008,7 +1008,7 @@ class RemoteTable(Table):
         return LOOP.run(self._table.drop_columns(columns))
 
     def set_unenforced_primary_key(self, columns: Union[str, Iterable[str]]) -> None:
-        """Set the unenforced primary key for this table to a single column."""
+        """Not supported on LanceDB Cloud."""
         return LOOP.run(self._table.set_unenforced_primary_key(columns))
 
     def set_lsm_write_spec(self, spec: "LsmWriteSpec") -> None:

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",

rust/lancedb/src/remote/table.rs

@@ -32,7 +32,6 @@ use crate::table::Tags;
 use crate::table::UpdateResult;
 use crate::table::lsm_stats::GetLsmStatsResponse;
 use crate::table::merge::MergeFilter;
-use crate::table::primary_key;
 use crate::table::query::create_multi_vector_plan;
 use crate::table::write_progress::FinishOnDrop;
 use crate::table::{
@@ -69,7 +68,6 @@ use lance::arrow::json::{JsonDataType, JsonSchema};
 use lance::dataset::refs::TagContents;
 use lance::dataset::scanner::DatasetRecordBatchStream;
 use lance::dataset::{ColumnAlteration, NewColumnTransform, Version};
-use lance_core::datatypes::Schema as LanceSchema;
 use lance_datafusion::exec::{OneShotExec, execute_plan};
 use reqwest::{RequestBuilder, Response};
 use serde::{Deserialize, Serialize};
@@ -2994,27 +2992,10 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
         }
     }
 
-    /// The unenforced primary key is Lance schema field metadata, so this
-    /// installs it through the `update_field_metadata` endpoint. The commit
-    /// layer behind that endpoint is what actually installs and validates the
-    /// key, exactly as on a native table; the checks here only fail fast with
-    /// the same messages a native table gives.
-    async fn set_unenforced_primary_key(&self, columns: &[&str]) -> Result<()> {
-        self.check_mutable().await?;
-
-        let arrow_schema = self.schema().await?;
-        let schema = LanceSchema::try_from(arrow_schema.as_ref()).map_err(|e| Error::Schema {
-            message: format!("Invalid schema: {}", e),
-        })?;
-        primary_key::validate(&schema, columns)?;
-
-        self.update_field_metadata(&[FieldMetadataUpdate {
-            path: columns[0].to_string(),
-            metadata: primary_key::install_edit(),
-            replace: false,
-        }])
-        .await?;
-        Ok(())
+    async fn set_unenforced_primary_key(&self, _columns: &[&str]) -> Result<()> {
+        Err(Error::NotSupported {
+            message: "set_unenforced_primary_key is not supported on LanceDB cloud.".into(),
+        })
     }
 
     async fn flush_lsm(&self) -> Result<()> {
@@ -11851,120 +11832,6 @@ mod tests {
         assert_eq!(result.version, 7);
     }
 
-    /// The unenforced primary key is field metadata, so the remote table
-    /// installs it through the `update_field_metadata` endpoint.
-    #[tokio::test]
-    async fn test_set_unenforced_primary_key() {
-        let table = Table::new_with_handler("my_table", |request| {
-            assert_eq!(request.method(), "POST");
-            match request.url().path() {
-                "/v1/table/my_table/describe/" => {
-                    let schema = Schema::new(vec![
-                        Field::new("id", DataType::Int64, false),
-                        Field::new("name", DataType::Utf8, true),
-                    ]);
-                    http::Response::builder()
-                        .status(200)
-                        .body(describe_response(&schema))
-                        .unwrap()
-                }
-                "/v1/table/my_table/update_field_metadata/" => {
-                    let body = request_body_json(&request);
-                    assert_eq!(body["updates"].as_array().unwrap().len(), 1);
-                    let update = &body["updates"][0];
-                    assert_eq!(update["path"], "id");
-                    assert_eq!(update["replace"], json!(false));
-                    assert_eq!(
-                        update["metadata"]["lance-schema:unenforced-primary-key:position"],
-                        "1"
-                    );
-                    assert_eq!(
-                        update["metadata"]["lance-schema:unenforced-primary-key"],
-                        json!(null)
-                    );
-                    http::Response::builder()
-                        .status(200)
-                        .body(r#"{"version": 3, "fields": {}}"#.to_string())
-                        .unwrap()
-                }
-                path => panic!("Unexpected path: {}", path),
-            }
-        });
-
-        table.set_unenforced_primary_key(["id"]).await.unwrap();
-    }
-
-    /// Requests the native table rejects are rejected here too, before any
-    /// write reaches the server.
-    #[tokio::test]
-    async fn test_set_unenforced_primary_key_rejects_invalid_requests() {
-        let table = Table::new_with_handler("my_table", |request| match request.url().path() {
-            "/v1/table/my_table/describe/" => {
-                let schema = Schema::new(vec![
-                    Field::new("id", DataType::Int64, false),
-                    Field::new("score", DataType::Float32, true),
-                ]);
-                http::Response::builder()
-                    .status(200)
-                    .body(describe_response(&schema))
-                    .unwrap()
-            }
-            path => panic!("Unexpected path: {}", path),
-        });
-
-        for columns in [
-            vec![],
-            vec!["id", "score"],
-            vec!["nonexistent"],
-            vec!["score"],
-        ] {
-            let err = table
-                .set_unenforced_primary_key(columns.clone())
-                .await
-                .unwrap_err();
-            assert!(
-                matches!(err, Error::InvalidInput { .. }),
-                "unexpected error for {:?}: {:?}",
-                columns,
-                err
-            );
-        }
-    }
-
-    /// The key is immutable once set, and the schema the server already
-    /// reports is enough to say so.
-    #[tokio::test]
-    async fn test_set_unenforced_primary_key_already_set() {
-        let table = Table::new_with_handler("my_table", |request| match request.url().path() {
-            "/v1/table/my_table/describe/" => {
-                let schema = Schema::new(vec![
-                    Field::new("id", DataType::Int64, false).with_metadata(HashMap::from([(
-                        "lance-schema:unenforced-primary-key:position".to_string(),
-                        "1".to_string(),
-                    )])),
-                    Field::new("name", DataType::Utf8, false),
-                ]);
-                http::Response::builder()
-                    .status(200)
-                    .body(describe_response(&schema))
-                    .unwrap()
-            }
-            path => panic!("Unexpected path: {}", path),
-        });
-
-        for column in ["name", "id"] {
-            let err = table
-                .set_unenforced_primary_key([column])
-                .await
-                .unwrap_err();
-            assert!(
-                err.to_string().contains("already set"),
-                "unexpected error: {:?}",
-                err
-            );
-        }
-    }
-
     // ----- Branch support -----
 
     /// Parse a request's in-memory JSON body. Only valid for JSON-body ops

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<Self> {
+        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<String>,
+        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<Self> {
+        if *batch.schema().as_ref() != dml_result_schema() || batch.num_rows() != 1 {
+            return None;
+        }
+        Some(Self {
+            table: batch
+                .column(0)
+                .as_any()
+                .downcast_ref::<StringArray>()?
+                .value(0)
+                .to_string(),
+            operation: DmlOperation::parse(
+                batch
+                    .column(1)
+                    .as_any()
+                    .downcast_ref::<StringArray>()?
+                    .value(0),
+            )?,
+            rows_affected: batch
+                .column(2)
+                .as_any()
+                .downcast_ref::<Int64Array>()?
+                .value(0),
+            version: batch
+                .column(3)
+                .as_any()
+                .downcast_ref::<Int64Array>()?
+                .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);
+    }
+}

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

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<String>,
+    /// 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<dyn WriteObserver>>,
+    database: &str,
+    schema: &str,
+    table: &str,
+    table_uri: Option<String>,
+    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;
+}

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<Option<LogicalPlan>> {
+    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<AtomicUsize>,
+    }
+
+    enum Outcome {
+        Plans,
+        Declines,
+    }
+
+    impl Handler {
+        fn new(keyword: Keyword, outcome: Outcome) -> (Arc<Self>, Arc<AtomicUsize>) {
+            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<Option<LogicalPlan>> {
+            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(&registry, "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(&registry, 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(&registry, "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(&registry, "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(&registry, "EXPLAIN t").unwrap().is_some());
+        assert_eq!(first_parsed.load(Ordering::SeqCst), 1);
+        assert_eq!(second_parsed.load(Ordering::SeqCst), 0);
+    }
+}

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<Option<LogicalPlan>>;
+}
+
+/// 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<String>) -> 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<Vec<AccessRequirement>>;
+}
+
+/// 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<Arc<dyn CustomSqlHandler>>,
+    statements: Vec<Arc<dyn SqlStatement>>,
+}
+
+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<Arc<dyn CustomSqlHandler>>,
+        statements: Vec<Arc<dyn SqlStatement>>,
+    ) -> 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<dyn CustomSqlHandler>) -> &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<dyn SqlStatement>) -> &mut Self {
+        self.statements.insert(0, statement);
+        self
+    }
+
+    /// The grammars, in the order they are consulted.
+    pub fn parsers(&self) -> &[Arc<dyn CustomSqlHandler>] {
+        &self.parsers
+    }
+
+    /// The statement owning this planned node, if any.
+    pub fn claim(&self, node: &dyn Any) -> Option<&Arc<dyn SqlStatement>> {
+        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::<u8>()
+        }
+
+        fn audit_operation(&self) -> &'static str {
+            self.label
+        }
+
+        fn access_requirements(
+            &self,
+            _node: &dyn Any,
+            _context: &RequirementContext<'_>,
+        ) -> DfResult<Vec<AccessRequirement>> {
+            Ok(vec![])
+        }
+    }
+
+    fn claimant(label: &'static str, claims_everything: bool) -> Arc<dyn SqlStatement> {
+        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());
+    }
+}

rust/lancedb/src/table.rs

@@ -78,7 +78,7 @@ pub mod delete;
 pub mod lsm_stats;
 pub mod merge;
 pub mod optimize;
-pub(crate) mod primary_key;
+mod primary_key;
 pub mod query;
 pub mod refresh;
 pub mod schema_evolution;

rust/lancedb/src/table/primary_key.rs

@@ -11,26 +11,24 @@
 //! Only a single-column primary key is supported, and the key cannot be
 //! changed once set.
 
-use std::collections::HashMap;
-
 use arrow_schema::DataType;
-use lance_core::datatypes::{
-    Field as LanceField, LANCE_UNENFORCED_PRIMARY_KEY, LANCE_UNENFORCED_PRIMARY_KEY_POSITION,
-    Schema as LanceSchema,
-};
+use lance_core::datatypes::{LANCE_UNENFORCED_PRIMARY_KEY, LANCE_UNENFORCED_PRIMARY_KEY_POSITION};
 
 use crate::error::{Error, Result};
 use crate::table::NativeTable;
 
-/// Validate a `set_unenforced_primary_key` request against `schema`, returning
-/// the field the key would be installed on.
+/// Set the unenforced primary key on `table` to the single column in `columns`.
 ///
-/// Shared by [`NativeTable`] and the remote table so both reject the same
-/// requests with the same messages. Fails if `columns` is not exactly one
-/// column (compound primary keys are not supported), if the column does not
-/// exist or has an unsupported dtype, or if the table already has an
-/// unenforced primary key (changing the primary key is not supported).
-pub fn validate<'a>(schema: &'a LanceSchema, columns: &[&str]) -> Result<&'a LanceField> {
+/// Fails if `columns` is not exactly one column (compound primary keys are not
+/// supported), if the column does not exist or has an unsupported dtype, or if
+/// the table already has an unenforced primary key (changing the primary key
+/// is not supported).
+pub(super) async fn set_unenforced_primary_key(
+    table: &NativeTable,
+    columns: &[&str],
+) -> Result<()> {
+    table.dataset.ensure_mutable()?;
+
     if columns.is_empty() {
         return Err(Error::InvalidInput {
             message: "set_unenforced_primary_key: a column is required".into(),
@@ -46,71 +44,43 @@ pub fn validate<'a>(schema: &'a LanceSchema, columns: &[&str]) -> Result<&'a Lan
     }
     let column = columns[0];
 
-    // The primary key is immutable once set. The Lance commit layer is the
-    // source of truth for this (it also covers the concurrent-writer race);
-    // this check just fails fast with a clear message.
-    if !schema.unenforced_primary_key().is_empty() {
-        return Err(Error::InvalidInput {
-            message: "set_unenforced_primary_key: an unenforced primary key is already set on this table; changing it is not supported".into(),
-        });
-    }
-
-    let field = schema.field(column).ok_or_else(|| Error::InvalidInput {
-        message: format!(
-            "set_unenforced_primary_key: column '{}' not found on table",
-            column
-        ),
-    })?;
-    if !is_supported_pk_dtype(&field.data_type()) {
-        return Err(Error::InvalidInput {
-            message: format!(
-                "set_unenforced_primary_key: column '{}' has dtype {:?} which is not supported as a primary key. Supported: Int32, Int64, Utf8, LargeUtf8, Binary, LargeBinary, FixedSizeBinary",
-                column,
-                field.data_type()
-            ),
-        });
-    }
-    Ok(field)
-}
-
-/// The field metadata edit that installs the primary key on a field: keys to
-/// set (`Some`) or delete (`None`).
-///
-/// Position metadata is 1-indexed; `Schema::unenforced_primary_key` treats
-/// position 0 as a legacy "no specific position" fallback, so the legacy
-/// boolean key is cleared and only the position governs.
-pub fn install_edit() -> HashMap<String, Option<String>> {
-    HashMap::from([
-        (LANCE_UNENFORCED_PRIMARY_KEY.to_string(), None),
-        (
-            LANCE_UNENFORCED_PRIMARY_KEY_POSITION.to_string(),
-            Some("1".to_string()),
-        ),
-    ])
-}
-
-/// Set the unenforced primary key on `table` to the single column in `columns`.
-pub(super) async fn set_unenforced_primary_key(
-    table: &NativeTable,
-    columns: &[&str],
-) -> Result<()> {
-    table.dataset.ensure_mutable()?;
-
     let updates = {
         let dataset = table.dataset.get().await?;
-        let field = validate(dataset.schema(), columns)?;
+        let schema = dataset.schema();
 
-        let mut metadata = field.metadata.clone();
-        for (key, value) in install_edit() {
-            match value {
-                Some(value) => {
-                    metadata.insert(key, value);
-                }
-                None => {
-                    metadata.remove(&key);
-                }
-            }
+        // The primary key is immutable once set. The Lance commit layer is the
+        // source of truth for this (it also covers the concurrent-writer race);
+        // this check just fails fast with a clear message.
+        if !schema.unenforced_primary_key().is_empty() {
+            return Err(Error::InvalidInput {
+                message: "set_unenforced_primary_key: an unenforced primary key is already set on this table; changing it is not supported".into(),
+            });
         }
+
+        let field = schema.field(column).ok_or_else(|| Error::InvalidInput {
+            message: format!(
+                "set_unenforced_primary_key: column '{}' not found on table",
+                column
+            ),
+        })?;
+        if !is_supported_pk_dtype(&field.data_type()) {
+            return Err(Error::InvalidInput {
+                message: format!(
+                    "set_unenforced_primary_key: column '{}' has dtype {:?} which is not supported as a primary key. Supported: Int32, Int64, Utf8, LargeUtf8, Binary, LargeBinary, FixedSizeBinary",
+                    column,
+                    field.data_type()
+                ),
+            });
+        }
+
+        // Position metadata is 1-indexed; `Schema::unenforced_primary_key`
+        // treats position 0 as a legacy "no specific position" fallback.
+        let mut metadata = field.metadata.clone();
+        metadata.remove(LANCE_UNENFORCED_PRIMARY_KEY);
+        metadata.insert(
+            LANCE_UNENFORCED_PRIMARY_KEY_POSITION.to_string(),
+            "1".to_string(),
+        );
         vec![(field_id_to_u32(field.id, &field.name)?, metadata)]
     };