diff --git a/beacon-api/src/axum/admin/extensions.rs b/beacon-api/src/axum/admin/extensions.rs new file mode 100644 index 00000000..015bd4da --- /dev/null +++ b/beacon-api/src/axum/admin/extensions.rs @@ -0,0 +1,63 @@ +//! Admin endpoints for managing a table's downstream extensions. + +use std::sync::Arc; + +use ::axum::{ + extract::{Path, State}, + http::StatusCode, + Json, +}; +use beacon_core::api::TableExtensions; +use beacon_core::runtime::Runtime; + +use super::bad_request; + +/// Replaces the named table's extensions document (MCP descriptor, query +/// presets). The document is validated against the table schema; an empty body +/// (`{}`) clears all extensions. +#[tracing::instrument(level = "info", skip(state, extensions))] +#[utoipa::path( + tag = "admin", + put, + path = "/api/admin/table-extensions/{table_name}", + params(("table_name" = String, Path, description = "Registered table name")), + request_body = TableExtensions, + responses( + (status = 200, description = "Extensions updated"), + (status = 400, description = "Invalid request, validation failed, or table not found") + ), + security(("basic-auth" = [])) +)] +pub(crate) async fn set_table_extensions( + State(state): State>, + Path(table_name): Path, + Json(extensions): Json, +) -> Result<(), (StatusCode, String)> { + state + .set_table_extensions(table_name, extensions) + .await + .map_err(bad_request) +} + +/// Removes all extensions from the named table. +#[tracing::instrument(level = "info", skip(state))] +#[utoipa::path( + tag = "admin", + delete, + path = "/api/admin/table-extensions/{table_name}", + params(("table_name" = String, Path, description = "Registered table name")), + responses( + (status = 200, description = "Extensions removed"), + (status = 400, description = "Table not found or removal failed") + ), + security(("basic-auth" = [])) +)] +pub(crate) async fn delete_table_extensions( + State(state): State>, + Path(table_name): Path, +) -> Result<(), (StatusCode, String)> { + state + .delete_table_extensions(table_name) + .await + .map_err(bad_request) +} diff --git a/beacon-api/src/axum/admin/mod.rs b/beacon-api/src/axum/admin/mod.rs index 3d88a551..d14a9453 100644 --- a/beacon-api/src/axum/admin/mod.rs +++ b/beacon-api/src/axum/admin/mod.rs @@ -14,6 +14,7 @@ use crate::axum::auth::basic_auth; mod check; mod crawlers; +mod extensions; mod external_tables; /// OpenAPI document marker for the admin surface. @@ -37,6 +38,10 @@ pub(crate) fn setup_admin_router() -> (Router>, utoipa::openapi::Op )) .routes(routes!(crawlers::run_crawler)) .routes(routes!(external_tables::create_external_table)) + .routes(routes!( + extensions::set_table_extensions, + extensions::delete_table_extensions + )) .layer(::axum::middleware::from_fn(basic_auth)) .split_for_parts(); diff --git a/beacon-api/src/axum/client/mod.rs b/beacon-api/src/axum/client/mod.rs index 2d1b4a74..237ed597 100644 --- a/beacon-api/src/axum/client/mod.rs +++ b/beacon-api/src/axum/client/mod.rs @@ -41,6 +41,7 @@ pub(crate) fn setup_client_router() -> (Router>, utoipa::openapi::O .routes(routes!(tables::list_tables_with_schema)) .routes(routes!(tables::default_table)) .routes(routes!(tables::list_table_schema)) + .routes(routes!(tables::list_table_extensions)) .routes(routes!(tables::list_table_config)) .routes(routes!(tables::default_table_schema)) .routes(routes!(functions::list_functions)) diff --git a/beacon-api/src/axum/client/tables.rs b/beacon-api/src/axum/client/tables.rs index 5947909e..e5df3929 100644 --- a/beacon-api/src/axum/client/tables.rs +++ b/beacon-api/src/axum/client/tables.rs @@ -7,7 +7,7 @@ use ::axum::{ http::StatusCode, Json, }; -use beacon_core::api::{SchemaFieldView, SchemaView, TableConfigView}; +use beacon_core::api::{SchemaFieldView, SchemaView, TableConfigView, TableExtensions}; use beacon_core::runtime::Runtime; use utoipa::{IntoParams, ToSchema}; @@ -152,6 +152,48 @@ pub(crate) async fn list_table_config( } } +/// Query parameters for [`list_table_extensions`]. +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, ToSchema, IntoParams)] +pub struct ListTableExtensionsQuery { + /// Name of the registered table whose extensions to return. + pub table_name: String, +} + +/// Returns the downstream extensions (MCP descriptor, query presets) attached to +/// the named table, or 404 if the table is not registered. A table with no +/// extensions returns an empty object. +#[tracing::instrument(level = "info", skip(state))] +#[utoipa::path( + tag = "tables", + get, + path = "/api/table-extensions", + params(ListTableExtensionsQuery), + responses( + (status = 200, description = "The table's extensions", body = TableExtensions), + (status = 404, description = "Table not found"), + ), + security( + (), + ("basic-auth" = []), + ("bearer" = []) + ) +)] +pub(crate) async fn list_table_extensions( + State(state): State>, + Query(query): Query, +) -> Result, (StatusCode, String)> { + match state.get_table_extensions(query.table_name.clone()).await { + Ok(extensions) => Ok(Json(extensions)), + Err(error) => { + tracing::error!(?error, "error listing table extensions"); + Err(( + StatusCode::NOT_FOUND, + format!("Table {} not found", query.table_name), + )) + } + } +} + /// Returns the Arrow schema of the runtime's default table. #[tracing::instrument(level = "info", skip(state))] #[utoipa::path( diff --git a/beacon-core/src/api.rs b/beacon-core/src/api.rs index 5d5e283c..a5161d70 100644 --- a/beacon-core/src/api.rs +++ b/beacon-core/src/api.rs @@ -12,6 +12,11 @@ use crate::metrics::ConsolidatedMetrics; use serde_json::{Map, Value}; use utoipa::ToSchema; +/// Re-exported typed table-extension contracts (see [`crate::extensions`]). +pub use crate::extensions::{ + McpExtension, Preset, PresetExtension, PresetFilter, TableExtensions, +}; + /// A single parameter of a registered function. #[derive(Debug, Clone, serde::Serialize, serde::Deserialize, ToSchema)] pub struct FunctionParameterInfo { diff --git a/beacon-core/src/extensions.rs b/beacon-core/src/extensions.rs new file mode 100644 index 00000000..6f99ea7a --- /dev/null +++ b/beacon-core/src/extensions.rs @@ -0,0 +1,433 @@ +//! Typed, consumer-facing table extensions (MCP descriptor, query presets). +//! +//! Extensions are metadata *about how to use* a table — distinct from its storage +//! definition and from format `options`. They are stored decoupled from the table +//! definition in a `tables:///extensions.json` sidecar, so they: +//! +//! - apply uniformly to every table type (listing/Iceberg/Delta/SQL/remote/view), +//! - can be edited without rebuilding the provider, +//! - survive provider re-registration (materialized-view refresh, Iceberg alter), +//! - are removed automatically on `DROP TABLE` (the table directory is deleted). +//! +//! This module owns the typed contract, schema validation, and the +//! read/modify/write logic shared by the SQL DDL path (`SET/DROP EXTENSION`, +//! `SHOW EXTENSIONS`) and the REST path (`Runtime` methods). + +use std::sync::{Arc, OnceLock}; + +use anyhow::Context; +use arrow::array::StringArray; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use beacon_data_lake::{SchemaPersistenceService, TABLES_OBJECT_STORE_URL}; +use datafusion::prelude::SessionContext; +use serde::{Deserialize, Serialize}; +use utoipa::ToSchema; + +/// The comparison operators a [`PresetFilter`] may use. +pub const PRESET_OPS: [&str; 8] = ["=", "!=", "<", "<=", ">", ">=", "between", "in"]; + +/// The full set of extensions attached to a table — the `extensions.json` +/// document. Missing kinds are omitted from the serialized form. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, ToSchema)] +pub struct TableExtensions { + /// MCP descriptor: how downstream MCP servers should surface this table. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub mcp: Option, + /// Named, predefined filter sets consumers can apply downstream. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub preset: Option, +} + +/// MCP descriptor: how a downstream MCP server should expose this table as a +/// tool/resource. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)] +#[schema(example = json!({ + "enabled": true, + "tool_name": "query_ocean_observations", + "description": "Argo float observations by location, depth, and time.", + "exposed_columns": ["lat", "lon", "depth", "temperature", "time"] +}))] +pub struct McpExtension { + /// Whether downstream MCP servers should expose this table at all. + #[serde(default)] + pub enabled: bool, + /// Tool name to expose. Downstream may default to the table name if unset. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tool_name: Option, + /// Human-readable description for the MCP tool/resource. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub description: Option, + /// Columns to expose. `None` (omitted) exposes all columns. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub exposed_columns: Option>, +} + +/// A set of named, predefined filters consumers can apply. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)] +#[schema(example = json!({ + "presets": [{ + "name": "north_atlantic_surface", + "description": "Surface measurements in the North Atlantic", + "filters": [ + { "column": "lat", "op": "between", "value": [0, 60] }, + { "column": "depth", "op": "<=", "value": 10 } + ] + }] +}))] +pub struct PresetExtension { + /// The named presets. + pub presets: Vec, +} + +/// A single named preset: a bundle of filters applied together. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)] +pub struct Preset { + /// Unique (within the table) preset name. + pub name: String, + /// Optional human-readable description. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub description: Option, + /// The filters that make up the preset. + pub filters: Vec, +} + +/// A single predefined filter within a [`Preset`]. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)] +pub struct PresetFilter { + /// Column the filter applies to (must exist in the table schema). + pub column: String, + /// Comparison operator — one of [`PRESET_OPS`]. + #[schema(example = "between")] + pub op: String, + /// Filter value: a scalar, `[lo, hi]` for `between`, or `[..]` for `in`. + #[schema(value_type = Object)] + pub value: serde_json::Value, +} + +impl TableExtensions { + /// Whether no extensions are set. + pub fn is_empty(&self) -> bool { + self.mcp.is_none() && self.preset.is_none() + } + + /// Parse the JSON payload for one extension `kind` and splice it into the + /// document, leaving the other kinds untouched. + pub fn set_kind(&mut self, kind: &str, json: &str) -> anyhow::Result<()> { + match kind.to_ascii_lowercase().as_str() { + "mcp" => { + self.mcp = Some( + serde_json::from_str(json).context("invalid 'mcp' extension payload")?, + ); + } + "preset" => { + self.preset = Some( + serde_json::from_str(json).context("invalid 'preset' extension payload")?, + ); + } + other => anyhow::bail!( + "unknown extension kind '{other}'; expected one of: mcp, preset" + ), + } + Ok(()) + } + + /// Remove one extension `kind` from the document. + pub fn drop_kind(&mut self, kind: &str) -> anyhow::Result<()> { + match kind.to_ascii_lowercase().as_str() { + "mcp" => self.mcp = None, + "preset" => self.preset = None, + other => anyhow::bail!( + "unknown extension kind '{other}'; expected one of: mcp, preset" + ), + } + Ok(()) + } + + /// Validate every set extension against the table's Arrow schema. + pub fn validate(&self, schema: &Schema) -> anyhow::Result<()> { + if let Some(mcp) = &self.mcp { + mcp.validate(schema)?; + } + if let Some(preset) = &self.preset { + preset.validate(schema)?; + } + Ok(()) + } +} + +impl McpExtension { + fn validate(&self, schema: &Schema) -> anyhow::Result<()> { + if let Some(columns) = &self.exposed_columns { + for column in columns { + ensure_column(schema, column)?; + } + } + Ok(()) + } +} + +impl PresetExtension { + fn validate(&self, schema: &Schema) -> anyhow::Result<()> { + let mut seen = std::collections::HashSet::new(); + for preset in &self.presets { + if !seen.insert(preset.name.as_str()) { + anyhow::bail!("duplicate preset name '{}'", preset.name); + } + for filter in &preset.filters { + ensure_column(schema, &filter.column)?; + if !PRESET_OPS.contains(&filter.op.as_str()) { + anyhow::bail!( + "preset '{}' uses unsupported operator '{}'; expected one of: {}", + preset.name, + filter.op, + PRESET_OPS.join(", ") + ); + } + validate_filter_value_shape(&preset.name, filter)?; + } + } + Ok(()) + } +} + +/// `between` requires a two-element array; `in` requires a non-empty array. +fn validate_filter_value_shape(preset: &str, filter: &PresetFilter) -> anyhow::Result<()> { + match filter.op.as_str() { + "between" => { + let ok = filter.value.as_array().is_some_and(|a| a.len() == 2); + anyhow::ensure!( + ok, + "preset '{preset}' filter on '{}' uses 'between' but value is not a two-element array", + filter.column + ); + } + "in" => { + let ok = filter.value.as_array().is_some_and(|a| !a.is_empty()); + anyhow::ensure!( + ok, + "preset '{preset}' filter on '{}' uses 'in' but value is not a non-empty array", + filter.column + ); + } + _ => {} + } + Ok(()) +} + +fn ensure_column(schema: &Schema, column: &str) -> anyhow::Result<()> { + anyhow::ensure!( + schema.column_with_name(column).is_some(), + "column '{column}' does not exist in the table schema" + ); + Ok(()) +} + +/// Arrow schema produced by `SHOW EXTENSIONS FOR `: a single JSON column. +pub fn show_extensions_arrow_schema() -> SchemaRef { + static SCHEMA: OnceLock = OnceLock::new(); + SCHEMA + .get_or_init(|| Arc::new(Schema::new(vec![Field::new("extensions", DataType::Utf8, false)]))) + .clone() +} + +/// The table's live Arrow schema, erroring if the table is not registered. +async fn table_schema(ctx: &Arc, name: &str) -> anyhow::Result { + let provider = ctx + .table_provider(name) + .await + .map_err(|_| anyhow::anyhow!("table '{name}' not found"))?; + Ok(provider.schema()) +} + +fn persistence(ctx: &Arc) -> SchemaPersistenceService { + SchemaPersistenceService::new(ctx.clone(), TABLES_OBJECT_STORE_URL.clone()) +} + +/// Load a table's extensions, returning an empty set if none are stored. +pub async fn get_table_extensions( + ctx: &Arc, + name: &str, +) -> anyhow::Result { + anyhow::ensure!(ctx.table_exist(name)?, "table '{name}' not found"); + match persistence(ctx).load_table_extensions_json(name).await? { + Some(json) => Ok(serde_json::from_str(&json) + .context("stored table extensions are not valid")?), + None => Ok(TableExtensions::default()), + } +} + +/// Set (or replace) a single extension kind from its JSON payload, validating it +/// against the table schema before persisting. +pub async fn set_table_extension( + ctx: &Arc, + name: &str, + kind: &str, + json: &str, +) -> anyhow::Result<()> { + let schema = table_schema(ctx, name).await?; + let mut extensions = get_table_extensions(ctx, name).await?; + extensions.set_kind(kind, json)?; + extensions.validate(&schema)?; + persistence(ctx) + .persist_table_extensions_json(name, serde_json::to_string_pretty(&extensions)?) + .await?; + Ok(()) +} + +/// Remove a single extension kind. The sidecar is deleted if nothing remains. +pub async fn drop_table_extension( + ctx: &Arc, + name: &str, + kind: &str, +) -> anyhow::Result<()> { + let mut extensions = get_table_extensions(ctx, name).await?; + extensions.drop_kind(kind)?; + write_or_remove(ctx, name, &extensions).await +} + +/// Replace the entire extensions document (the REST `PUT` surface), validating it +/// against the table schema. An empty document removes the sidecar. +pub async fn set_table_extensions( + ctx: &Arc, + name: &str, + extensions: TableExtensions, +) -> anyhow::Result<()> { + let schema = table_schema(ctx, name).await?; + extensions.validate(&schema)?; + write_or_remove(ctx, name, &extensions).await +} + +/// Remove all extensions for a table. +pub async fn delete_table_extensions( + ctx: &Arc, + name: &str, +) -> anyhow::Result<()> { + anyhow::ensure!(ctx.table_exist(name)?, "table '{name}' not found"); + persistence(ctx).remove_table_extensions_json(name).await?; + Ok(()) +} + +async fn write_or_remove( + ctx: &Arc, + name: &str, + extensions: &TableExtensions, +) -> anyhow::Result<()> { + let service = persistence(ctx); + if extensions.is_empty() { + service.remove_table_extensions_json(name).await?; + } else { + service + .persist_table_extensions_json(name, serde_json::to_string_pretty(extensions)?) + .await?; + } + Ok(()) +} + +/// Build the single-row `SHOW EXTENSIONS FOR
` result. +pub async fn show_table_extensions_batch( + ctx: &Arc, + name: &str, +) -> anyhow::Result { + let extensions = get_table_extensions(ctx, name).await?; + let json = serde_json::to_string_pretty(&extensions)?; + let column = StringArray::from(vec![Some(json)]); + RecordBatch::try_new(show_extensions_arrow_schema(), vec![Arc::new(column)]) + .context("failed to build SHOW EXTENSIONS batch") +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::{DataType, Field}; + + fn schema() -> Schema { + Schema::new(vec![ + Field::new("lat", DataType::Float64, true), + Field::new("lon", DataType::Float64, true), + Field::new("depth", DataType::Float64, true), + ]) + } + + #[test] + fn validates_preset_against_schema() { + let mut ext = TableExtensions::default(); + ext.set_kind( + "preset", + r#"{"presets":[{"name":"p","filters":[{"column":"lat","op":"between","value":[0,60]}]}]}"#, + ) + .unwrap(); + assert!(ext.validate(&schema()).is_ok()); + } + + #[test] + fn rejects_unknown_column() { + let mut ext = TableExtensions::default(); + ext.set_kind( + "preset", + r#"{"presets":[{"name":"p","filters":[{"column":"nope","op":"=","value":1}]}]}"#, + ) + .unwrap(); + let err = ext.validate(&schema()).unwrap_err().to_string(); + assert!(err.contains("does not exist"), "unexpected: {err}"); + } + + #[test] + fn rejects_bad_operator() { + let mut ext = TableExtensions::default(); + ext.set_kind( + "preset", + r#"{"presets":[{"name":"p","filters":[{"column":"lat","op":"~~","value":1}]}]}"#, + ) + .unwrap(); + assert!(ext.validate(&schema()).unwrap_err().to_string().contains("unsupported operator")); + } + + #[test] + fn rejects_between_without_pair() { + let mut ext = TableExtensions::default(); + ext.set_kind( + "preset", + r#"{"presets":[{"name":"p","filters":[{"column":"lat","op":"between","value":5}]}]}"#, + ) + .unwrap(); + assert!(ext.validate(&schema()).unwrap_err().to_string().contains("between")); + } + + #[test] + fn rejects_duplicate_preset_names() { + let mut ext = TableExtensions::default(); + ext.set_kind( + "preset", + r#"{"presets":[{"name":"p","filters":[]},{"name":"p","filters":[]}]}"#, + ) + .unwrap(); + assert!(ext.validate(&schema()).unwrap_err().to_string().contains("duplicate preset")); + } + + #[test] + fn mcp_exposed_columns_must_exist() { + let mut ext = TableExtensions::default(); + ext.set_kind("mcp", r#"{"enabled":true,"exposed_columns":["lat","ghost"]}"#) + .unwrap(); + assert!(ext.validate(&schema()).unwrap_err().to_string().contains("does not exist")); + } + + #[test] + fn unknown_kind_is_rejected() { + let mut ext = TableExtensions::default(); + assert!(ext.set_kind("bogus", "{}").is_err()); + assert!(ext.drop_kind("bogus").is_err()); + } + + #[test] + fn set_and_drop_kinds_are_independent() { + let mut ext = TableExtensions::default(); + ext.set_kind("mcp", r#"{"enabled":true}"#).unwrap(); + ext.set_kind("preset", r#"{"presets":[]}"#).unwrap(); + ext.drop_kind("mcp").unwrap(); + assert!(ext.mcp.is_none()); + assert!(ext.preset.is_some()); + assert!(!ext.is_empty()); + } +} diff --git a/beacon-core/src/lib.rs b/beacon-core/src/lib.rs index d85b8bcc..9e128743 100644 --- a/beacon-core/src/lib.rs +++ b/beacon-core/src/lib.rs @@ -1,4 +1,5 @@ pub mod api; +pub mod extensions; pub mod metrics; pub mod parser; pub mod query; diff --git a/beacon-core/src/parser/beacon_parser.rs b/beacon-core/src/parser/beacon_parser.rs index ceba62b0..d315f22a 100644 --- a/beacon-core/src/parser/beacon_parser.rs +++ b/beacon-core/src/parser/beacon_parser.rs @@ -8,7 +8,8 @@ use datafusion::sql::{ use super::statement::{ BeaconStatement, CreateCrawlerStatement, CreateMaterializedViewStatement, DropCrawlerStatement, - RefreshStatement, RunCrawlerStatement, + DropExtensionStatement, RefreshStatement, RunCrawlerStatement, SetExtensionStatement, + ShowExtensionsStatement, }; /// A parser that extends `DFParser` with custom Beacon SQL syntax. @@ -49,6 +50,18 @@ impl<'a> BeaconParser<'a> { return self.parse_show_crawlers(); } + if self.is_set_extension() { + return self.parse_set_extension(); + } + + if self.is_drop_extension() { + return self.parse_drop_extension(); + } + + if self.is_show_extensions() { + return self.parse_show_extensions(); + } + let df_statement = Box::new(self.df_parser.parse_statement()?); Ok(BeaconStatement::DFStatement(df_statement)) @@ -149,6 +162,86 @@ impl<'a> BeaconParser<'a> { Ok(BeaconStatement::ShowCrawlers) } + /// Whether the next two tokens are ` EXTENSION`, where `KW1` matches + /// `first` (used for `SET EXTENSION` and `DROP EXTENSION`). + fn is_keyword_then_extension(&self, first: impl Fn(&Token) -> bool) -> bool { + let t1 = &self.df_parser.parser.peek_nth_token(0).token; + let t2 = &self.df_parser.parser.peek_nth_token(1).token; + first(t1) && matches!(t2, Token::Word(w) if w.value.to_uppercase() == "EXTENSION") + } + + fn is_set_extension(&self) -> bool { + self.is_keyword_then_extension(|t| matches!(t, Token::Word(w) if w.keyword == Keyword::SET)) + } + + fn is_drop_extension(&self) -> bool { + self.is_keyword_then_extension(|t| matches!(t, Token::Word(w) if w.keyword == Keyword::DROP)) + } + + fn is_show_extensions(&self) -> bool { + let t1 = &self.df_parser.parser.peek_nth_token(0).token; + let t2 = &self.df_parser.parser.peek_nth_token(1).token; + matches!(t1, Token::Word(w) if w.value.to_uppercase() == "SHOW") + && matches!(t2, Token::Word(w) if w.value.to_uppercase() == "EXTENSIONS") + } + + /// Parse: SET EXTENSION '' FOR
TO '' + fn parse_set_extension(&mut self) -> Result { + self.df_parser.parser.next_token(); // SET + self.df_parser.parser.next_token(); // EXTENSION + let kind = self.parse_string_value()?; + self.expect_keyword(Keyword::FOR)?; + let table = self.parse_object_name()?; + self.expect_keyword(Keyword::TO)?; + let json = self.parse_string_value()?; + Ok(BeaconStatement::SetExtension(SetExtensionStatement { + kind, + table, + json, + })) + } + + /// Parse: DROP EXTENSION '' FOR
+ fn parse_drop_extension(&mut self) -> Result { + self.df_parser.parser.next_token(); // DROP + self.df_parser.parser.next_token(); // EXTENSION + let kind = self.parse_string_value()?; + self.expect_keyword(Keyword::FOR)?; + let table = self.parse_object_name()?; + Ok(BeaconStatement::DropExtension(DropExtensionStatement { + kind, + table, + })) + } + + /// Parse: SHOW EXTENSIONS FOR
+ fn parse_show_extensions(&mut self) -> Result { + self.df_parser.parser.next_token(); // SHOW + self.df_parser.parser.next_token(); // EXTENSIONS + self.expect_keyword(Keyword::FOR)?; + let table = self.parse_object_name()?; + Ok(BeaconStatement::ShowExtensions(ShowExtensionsStatement { + table, + })) + } + + /// Consume the expected keyword or error. + fn expect_keyword(&mut self, keyword: Keyword) -> Result<()> { + self.df_parser + .parser + .expect_keyword(keyword) + .map(|_| ()) + .map_err(|e| DataFusionError::External(Box::new(e))) + } + + /// Parse a (possibly schema-qualified) object name. + fn parse_object_name(&mut self) -> Result { + self.df_parser + .parser + .parse_object_name(false) + .map_err(|e| DataFusionError::External(Box::new(e))) + } + /// Read a single string value (single-quoted string, identifier, or number). fn parse_string_value(&mut self) -> Result { let token = self.df_parser.parser.next_token(); @@ -414,6 +507,58 @@ mod tests { } } + #[test] + fn test_parse_set_extension() { + let sql = "SET EXTENSION 'preset' FOR obs TO '{\"presets\":[]}'"; + let mut p = BeaconParser::new(sql).unwrap(); + match p.parse_statement().unwrap() { + BeaconStatement::SetExtension(s) => { + assert_eq!(s.kind, "preset"); + assert_eq!(s.table.to_string(), "obs"); + assert_eq!(s.json, "{\"presets\":[]}"); + } + other => panic!("expected SetExtension, got {other:?}"), + } + } + + #[test] + fn test_parse_drop_and_show_extensions() { + let mut p = BeaconParser::new("DROP EXTENSION 'mcp' FOR obs").unwrap(); + match p.parse_statement().unwrap() { + BeaconStatement::DropExtension(s) => { + assert_eq!(s.kind, "mcp"); + assert_eq!(s.table.to_string(), "obs"); + } + other => panic!("expected DropExtension, got {other:?}"), + } + + let mut p = BeaconParser::new("SHOW EXTENSIONS FOR schema.obs").unwrap(); + match p.parse_statement().unwrap() { + BeaconStatement::ShowExtensions(s) => assert_eq!(s.table.to_string(), "schema.obs"), + other => panic!("expected ShowExtensions, got {other:?}"), + } + } + + #[test] + fn test_extension_ddl_does_not_shadow_standard_sql() { + // SET , DROP TABLE, and SHOW TABLES must still reach DataFusion. + for sql in ["SET timezone = 'UTC'", "DROP TABLE t", "SHOW TABLES"] { + let mut p = BeaconParser::new(sql).unwrap(); + assert!( + matches!(p.parse_statement().unwrap(), BeaconStatement::DFStatement(_)), + "`{sql}` should be a DataFusion statement" + ); + } + } + + #[test] + fn test_set_extension_display_roundtrip() { + let sql = "SET EXTENSION 'preset' FOR obs TO '{\"presets\":[]}'"; + let mut p = BeaconParser::new(sql).unwrap(); + let stmt = p.parse_statement().unwrap(); + assert_eq!(stmt.to_string(), sql); + } + #[test] fn test_create_crawler_display_roundtrip() { let sql = "CREATE CRAWLER argo ON 'argo/' WITH ('format' 'parquet')"; diff --git a/beacon-core/src/parser/statement.rs b/beacon-core/src/parser/statement.rs index e08c2489..19eb1df3 100644 --- a/beacon-core/src/parser/statement.rs +++ b/beacon-core/src/parser/statement.rs @@ -12,6 +12,72 @@ pub enum BeaconStatement { RunCrawler(RunCrawlerStatement), DropCrawler(DropCrawlerStatement), ShowCrawlers, + SetExtension(SetExtensionStatement), + DropExtension(DropExtensionStatement), + ShowExtensions(ShowExtensionsStatement), +} + +/// SET EXTENSION '' FOR
TO '' +#[derive(Debug, Clone)] +pub struct SetExtensionStatement { + /// Extension kind (e.g. `mcp`, `preset`). + pub kind: String, + /// Target table. + pub table: ObjectName, + /// The extension payload as a JSON string literal. + pub json: String, +} + +impl Display for SetExtensionStatement { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "SET EXTENSION '{}' FOR {} TO '{}'", + escape_sql_literal(&self.kind), + self.table, + escape_sql_literal(&self.json) + ) + } +} + +/// DROP EXTENSION '' FOR
+#[derive(Debug, Clone)] +pub struct DropExtensionStatement { + /// Extension kind to remove. + pub kind: String, + /// Target table. + pub table: ObjectName, +} + +impl Display for DropExtensionStatement { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "DROP EXTENSION '{}' FOR {}", + escape_sql_literal(&self.kind), + self.table + ) + } +} + +/// SHOW EXTENSIONS FOR
+#[derive(Debug, Clone)] +pub struct ShowExtensionsStatement { + /// Target table. + pub table: ObjectName, +} + +impl Display for ShowExtensionsStatement { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "SHOW EXTENSIONS FOR {}", self.table) + } +} + +/// Escape a value for embedding in a single-quoted SQL string literal so the +/// `Display` form re-parses to the same value (the tokenizer turns `''` back into +/// `'`). +fn escape_sql_literal(value: &str) -> String { + value.replace('\'', "''") } /// CREATE CRAWLER [ON ''] [WITH (k 'v', ...)] @@ -108,6 +174,9 @@ impl Display for BeaconStatement { Self::RunCrawler(s) => write!(f, "{s}"), Self::DropCrawler(s) => write!(f, "{s}"), Self::ShowCrawlers => write!(f, "SHOW CRAWLERS"), + Self::SetExtension(s) => write!(f, "{s}"), + Self::DropExtension(s) => write!(f, "{s}"), + Self::ShowExtensions(s) => write!(f, "{s}"), } } } diff --git a/beacon-core/src/runtime.rs b/beacon-core/src/runtime.rs index ab503001..a085436b 100644 --- a/beacon-core/src/runtime.rs +++ b/beacon-core/src/runtime.rs @@ -361,6 +361,15 @@ impl Runtime { Ok(crate::statement_plan::drop_crawler_plan(statement)) } BeaconStatement::ShowCrawlers => Ok(crate::statement_plan::show_crawlers_plan()), + BeaconStatement::SetExtension(statement) => { + Ok(crate::statement_plan::set_extension_plan(statement)) + } + BeaconStatement::DropExtension(statement) => { + Ok(crate::statement_plan::drop_extension_plan(statement)) + } + BeaconStatement::ShowExtensions(statement) => { + Ok(crate::statement_plan::show_extensions_plan(statement)) + } BeaconStatement::DFStatement(statement) => { crate::statement_plan::lower_df_statement(&self.session_ctx, *statement).await } @@ -531,6 +540,30 @@ impl Runtime { } } + /// Load a table's downstream extensions (MCP descriptor, query presets). + /// Returns an empty set if the table has none. + pub async fn get_table_extensions( + &self, + table_name: String, + ) -> anyhow::Result { + crate::extensions::get_table_extensions(&self.session_ctx, &table_name).await + } + + /// Replace a table's extensions document, validating it against the table + /// schema. An empty document removes the stored extensions. + pub async fn set_table_extensions( + &self, + table_name: String, + extensions: crate::extensions::TableExtensions, + ) -> anyhow::Result<()> { + crate::extensions::set_table_extensions(&self.session_ctx, &table_name, extensions).await + } + + /// Remove all of a table's extensions. + pub async fn delete_table_extensions(&self, table_name: String) -> anyhow::Result<()> { + crate::extensions::delete_table_extensions(&self.session_ctx, &table_name).await + } + pub async fn list_table_schema(&self, table_name: String) -> Option { self.session_ctx .table(table_name) @@ -907,6 +940,91 @@ mod client_query_tests { assert_eq!(batches[0].num_columns(), 2); } + /// `SET EXTENSION` / `SHOW EXTENSIONS` / `DROP EXTENSION` round-trip end to + /// end, and an extension referencing a missing column is rejected. + #[tokio::test(flavor = "multi_thread")] + async fn table_extensions_sql_round_trip() { + let runtime = Runtime::new(std::sync::Arc::new(beacon_config::Config::load().unwrap())) + .await + .expect("runtime should start"); + let suffix = uuid::Uuid::new_v4().simple(); + let table = format!("ext_{suffix}"); + + run_sql(&runtime, &format!("CREATE TABLE {table} (lat BIGINT, depth BIGINT)")).await; + + // SET a preset via SQL, then read it back through the typed API. + run_sql( + &runtime, + &format!( + "SET EXTENSION 'preset' FOR {table} TO '{{\"presets\":[{{\"name\":\"shallow\",\"filters\":[{{\"column\":\"depth\",\"op\":\"<=\",\"value\":10}}]}}]}}'" + ), + ) + .await; + + let ext = runtime + .get_table_extensions(table.clone()) + .await + .expect("extensions should load"); + let preset = ext.preset.expect("preset extension should be set"); + assert_eq!(preset.presets[0].name, "shallow"); + + // SHOW EXTENSIONS returns one JSON row mentioning the preset. + let batches = runtime + .run_query( + crate::query::Query::sql(format!("SHOW EXTENSIONS FOR {table}")), + true, + ) + .await + .expect("show extensions should run") + .into_record_stream() + .expect("streamed result") + .try_collect::>() + .await + .expect("stream should drain"); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 1, "SHOW EXTENSIONS returns one row"); + let json = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("extensions column is Utf8") + .value(0); + assert!(json.contains("shallow"), "SHOW output should include the preset: {json}"); + + // An extension over a non-existent column is rejected by validation. + let rejected = try_run_sql( + &runtime, + &format!( + "SET EXTENSION 'preset' FOR {table} TO '{{\"presets\":[{{\"name\":\"x\",\"filters\":[{{\"column\":\"ghost\",\"op\":\"=\",\"value\":1}}]}}]}}'" + ), + ) + .await; + assert!(rejected.is_err(), "preset over a missing column should be rejected"); + + // DROP removes it; the document becomes empty. + run_sql(&runtime, &format!("DROP EXTENSION 'preset' FOR {table}")).await; + assert!( + runtime + .get_table_extensions(table.clone()) + .await + .expect("extensions should load") + .is_empty(), + "dropping the only extension leaves an empty document" + ); + } + + /// Like `run_sql`, but returns the result (draining the stream) so callers can + /// assert on failures from side-effecting statements. + async fn try_run_sql(runtime: &Runtime, sql: &str) -> anyhow::Result<()> { + runtime + .run_query(crate::query::Query::sql(sql.to_string()), true) + .await? + .into_record_stream()? + .try_collect::>() + .await?; + Ok(()) + } + /// A query with an `output` format is written to a file and returned as a /// file download. #[tokio::test(flavor = "multi_thread")] diff --git a/beacon-core/src/statement_plan/logical.rs b/beacon-core/src/statement_plan/logical.rs index 79df3b6b..8a163a35 100644 --- a/beacon-core/src/statement_plan/logical.rs +++ b/beacon-core/src/statement_plan/logical.rs @@ -23,6 +23,8 @@ use datafusion::{ }, }; +use crate::extensions::show_extensions_arrow_schema; + /// Shared empty schema returned by beacon's side-effecting statement nodes, /// which produce no rows. `schema()` must return a reference, so the schema is /// stored once rather than rebuilt per node. @@ -417,6 +419,127 @@ impl UserDefinedLogicalNodeCore for AlterTableNode { } } +fn show_extensions_df_schema() -> &'static DFSchemaRef { + static SCHEMA: OnceLock = OnceLock::new(); + SCHEMA.get_or_init(|| { + Arc::new( + DFSchema::try_from(show_extensions_arrow_schema().as_ref().clone()) + .expect("SHOW EXTENSIONS schema is valid"), + ) + }) +} + +/// Logical node for `SET EXTENSION '' FOR
TO ''`. +#[derive(Debug, PartialEq, Eq, PartialOrd, Hash)] +pub(crate) struct SetExtensionNode { + pub(crate) kind: String, + pub(crate) table: String, + pub(crate) json: String, +} + +impl SetExtensionNode { + pub(crate) fn new(kind: String, table: String, json: String) -> Self { + Self { kind, table, json } + } +} + +impl UserDefinedLogicalNodeCore for SetExtensionNode { + fn name(&self) -> &str { + "SetExtension" + } + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + fn schema(&self) -> &DFSchemaRef { + empty_schema() + } + fn expressions(&self) -> Vec { + vec![] + } + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "SetExtension: table={} kind={}", self.table, self.kind) + } + fn with_exprs_and_inputs(&self, _exprs: Vec, _inputs: Vec) -> Result { + Ok(Self { + kind: self.kind.clone(), + table: self.table.clone(), + json: self.json.clone(), + }) + } +} + +/// Logical node for `DROP EXTENSION '' FOR
`. +#[derive(Debug, PartialEq, Eq, PartialOrd, Hash)] +pub(crate) struct DropExtensionNode { + pub(crate) kind: String, + pub(crate) table: String, +} + +impl DropExtensionNode { + pub(crate) fn new(kind: String, table: String) -> Self { + Self { kind, table } + } +} + +impl UserDefinedLogicalNodeCore for DropExtensionNode { + fn name(&self) -> &str { + "DropExtension" + } + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + fn schema(&self) -> &DFSchemaRef { + empty_schema() + } + fn expressions(&self) -> Vec { + vec![] + } + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "DropExtension: table={} kind={}", self.table, self.kind) + } + fn with_exprs_and_inputs(&self, _exprs: Vec, _inputs: Vec) -> Result { + Ok(Self { + kind: self.kind.clone(), + table: self.table.clone(), + }) + } +} + +/// Logical node for `SHOW EXTENSIONS FOR
`. Produces one JSON row. +#[derive(Debug, PartialEq, Eq, PartialOrd, Hash)] +pub(crate) struct ShowExtensionsNode { + pub(crate) table: String, +} + +impl ShowExtensionsNode { + pub(crate) fn new(table: String) -> Self { + Self { table } + } +} + +impl UserDefinedLogicalNodeCore for ShowExtensionsNode { + fn name(&self) -> &str { + "ShowExtensions" + } + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + fn schema(&self) -> &DFSchemaRef { + show_extensions_df_schema() + } + fn expressions(&self) -> Vec { + vec![] + } + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "ShowExtensions: table={}", self.table) + } + fn with_exprs_and_inputs(&self, _exprs: Vec, _inputs: Vec) -> Result { + Ok(Self { + table: self.table.clone(), + }) + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/beacon-core/src/statement_plan/mod.rs b/beacon-core/src/statement_plan/mod.rs index f3b6718e..6c413158 100644 --- a/beacon-core/src/statement_plan/mod.rs +++ b/beacon-core/src/statement_plan/mod.rs @@ -30,8 +30,9 @@ use datafusion::{ }; use crate::parser::statement::{ - CreateCrawlerStatement, CreateMaterializedViewStatement, DropCrawlerStatement, RefreshStatement, - RunCrawlerStatement, + CreateCrawlerStatement, CreateMaterializedViewStatement, DropCrawlerStatement, + DropExtensionStatement, RefreshStatement, RunCrawlerStatement, SetExtensionStatement, + ShowExtensionsStatement, }; pub(crate) use lower::lower_df_statement; @@ -145,6 +146,34 @@ pub(crate) fn show_crawlers_plan() -> LogicalPlan { }) } +/// Build the logical plan for `SET EXTENSION '' FOR
TO ''`. +pub(crate) fn set_extension_plan(statement: SetExtensionStatement) -> LogicalPlan { + LogicalPlan::Extension(Extension { + node: Arc::new(logical::SetExtensionNode::new( + statement.kind, + statement.table.to_string(), + statement.json, + )), + }) +} + +/// Build the logical plan for `DROP EXTENSION '' FOR
`. +pub(crate) fn drop_extension_plan(statement: DropExtensionStatement) -> LogicalPlan { + LogicalPlan::Extension(Extension { + node: Arc::new(logical::DropExtensionNode::new( + statement.kind, + statement.table.to_string(), + )), + }) +} + +/// Build the logical plan for `SHOW EXTENSIONS FOR
`. +pub(crate) fn show_extensions_plan(statement: ShowExtensionsStatement) -> LogicalPlan { + LogicalPlan::Extension(Extension { + node: Arc::new(logical::ShowExtensionsNode::new(statement.table.to_string())), + }) +} + /// Plan and execute a beacon statement logical plan through the single /// `create_physical_plan` -> `execute_stream` pipeline, coalescing the result the /// same way the legacy statement executor does. diff --git a/beacon-core/src/statement_plan/physical.rs b/beacon-core/src/statement_plan/physical.rs index a6f9509c..1b9e3ad7 100644 --- a/beacon-core/src/statement_plan/physical.rs +++ b/beacon-core/src/statement_plan/physical.rs @@ -30,6 +30,10 @@ use super::{ logical::{count_arrow_schema, show_crawlers_arrow_schema, AlterTableSpec}, materialized_view, SessionCell, }; +use crate::extensions::{ + drop_table_extension, set_table_extension, show_extensions_arrow_schema, + show_table_extensions_batch, +}; /// `PlanProperties` for a single-partition node producing `schema`. fn plan_properties(schema: SchemaRef) -> PlanProperties { @@ -852,3 +856,147 @@ impl ExecutionPlan for ShowCrawlersExec { Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } } + +/// Physical node for `SET EXTENSION '' FOR
TO ''`. +#[derive(Debug)] +pub(crate) struct SetExtensionExec { + kind: String, + table: String, + json: String, + session: SessionCell, + cache: Arc, +} + +impl SetExtensionExec { + pub(crate) fn new(kind: String, table: String, json: String, session: SessionCell) -> Self { + Self { + kind, + table, + json, + session, + cache: Arc::new(side_effect_properties()), + } + } + fn fmt_label(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "SetExtensionExec: table={} kind={}", self.table, self.kind) + } +} + +side_effect_exec!( + SetExtensionExec, + "SetExtensionExec", + |exec: &SetExtensionExec| { + let session = upgrade_session(&exec.session)?; + let kind = exec.kind.clone(); + let table = exec.table.clone(); + let json = exec.json.clone(); + Ok(side_effect_stream(async move { + set_table_extension(&session, &table, &kind, &json) + .await + .map_err(to_df_err) + })) + } +); + +/// Physical node for `DROP EXTENSION '' FOR
`. +#[derive(Debug)] +pub(crate) struct DropExtensionExec { + kind: String, + table: String, + session: SessionCell, + cache: Arc, +} + +impl DropExtensionExec { + pub(crate) fn new(kind: String, table: String, session: SessionCell) -> Self { + Self { + kind, + table, + session, + cache: Arc::new(side_effect_properties()), + } + } + fn fmt_label(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "DropExtensionExec: table={} kind={}", self.table, self.kind) + } +} + +side_effect_exec!( + DropExtensionExec, + "DropExtensionExec", + |exec: &DropExtensionExec| { + let session = upgrade_session(&exec.session)?; + let kind = exec.kind.clone(); + let table = exec.table.clone(); + Ok(side_effect_stream(async move { + drop_table_extension(&session, &table, &kind) + .await + .map_err(to_df_err) + })) + } +); + +/// Physical node for `SHOW EXTENSIONS FOR
`. Produces one JSON row. +#[derive(Debug)] +pub(crate) struct ShowExtensionsExec { + table: String, + session: SessionCell, + cache: Arc, +} + +impl ShowExtensionsExec { + pub(crate) fn new(table: String, session: SessionCell) -> Self { + Self { + table, + session, + cache: Arc::new(plan_properties(show_extensions_arrow_schema())), + } + } +} + +impl DisplayAs for ShowExtensionsExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, "ShowExtensionsExec: table={}", self.table) + } + DisplayFormatType::TreeRender => write!(f, "ShowExtensionsExec"), + } + } +} + +impl ExecutionPlan for ShowExtensionsExec { + fn name(&self) -> &str { + "ShowExtensionsExec" + } + fn as_any(&self) -> &dyn Any { + self + } + fn properties(&self) -> &Arc { + &self.cache + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + let session = upgrade_session(&self.session)?; + let table = self.table.clone(); + let schema = show_extensions_arrow_schema(); + let stream = futures::stream::once(async move { + show_table_extensions_batch(&session, &table) + .await + .map_err(to_df_err) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } +} diff --git a/beacon-core/src/statement_plan/query_planner.rs b/beacon-core/src/statement_plan/query_planner.rs index dd89edf1..2951c7a9 100644 --- a/beacon-core/src/statement_plan/query_planner.rs +++ b/beacon-core/src/statement_plan/query_planner.rs @@ -203,6 +203,30 @@ impl ExtensionPlanner for BeaconExtensionPlanner { return Ok(Some(Arc::new(physical::ShowCrawlersExec::new(session)))); } + if let Some(set) = any.downcast_ref::() { + return Ok(Some(Arc::new(physical::SetExtensionExec::new( + set.kind.clone(), + set.table.clone(), + set.json.clone(), + session, + )))); + } + + if let Some(drop) = any.downcast_ref::() { + return Ok(Some(Arc::new(physical::DropExtensionExec::new( + drop.kind.clone(), + drop.table.clone(), + session, + )))); + } + + if let Some(show) = any.downcast_ref::() { + return Ok(Some(Arc::new(physical::ShowExtensionsExec::new( + show.table.clone(), + session, + )))); + } + // Unrecognized node: let the default planner handle it. Ok(None) } diff --git a/beacon-data-lake/src/lib.rs b/beacon-data-lake/src/lib.rs index 0fde5a0c..27b2e71b 100644 --- a/beacon-data-lake/src/lib.rs +++ b/beacon-data-lake/src/lib.rs @@ -15,7 +15,7 @@ pub use files::temp_output_file::TempOutputFile; pub use files::{create_listing_url, create_temp_output_file, list_dataset_schema, list_datasets}; pub use table_runtime::init_tables; pub use table_runtime::persistent_schema_provider::PersistentSchemaProvider; -pub use table_runtime::schema_persistence::definition_from_provider; +pub use table_runtime::schema_persistence::{definition_from_provider, SchemaPersistenceService}; pub mod prelude { pub use super::files::*; diff --git a/beacon-data-lake/src/table_runtime/schema_persistence.rs b/beacon-data-lake/src/table_runtime/schema_persistence.rs index 77e238cb..97ac43c0 100644 --- a/beacon-data-lake/src/table_runtime/schema_persistence.rs +++ b/beacon-data-lake/src/table_runtime/schema_persistence.rs @@ -37,6 +37,92 @@ impl SchemaPersistenceService { .await } + /// Persist a table's extensions sidecar to `tables:///extensions.json`. + /// + /// Extensions are stored separately from `table.json` so they apply to every + /// table type uniformly and can be edited without rebuilding the provider. + pub async fn persist_table_extensions_json( + &self, + table_name: &str, + extensions_json: String, + ) -> datafusion::error::Result<()> { + let path = object_store::path::Path::from(format!("{}/extensions.json", table_name)); + let table_object_store = self.table_object_store(table_name)?; + table_object_store + .put(&path, extensions_json.into_bytes().into()) + .await + .map_err(|error| { + DataFusionError::Plan(format!( + "Failed to store table extensions for table {}: {}", + table_name, error + )) + })?; + Ok(()) + } + + /// Load a table's extensions sidecar, or `None` if it has none. + pub async fn load_table_extensions_json( + &self, + table_name: &str, + ) -> datafusion::error::Result> { + let path = object_store::path::Path::from(format!("{}/extensions.json", table_name)); + let table_object_store = self.table_object_store(table_name)?; + match table_object_store.get(&path).await { + Ok(result) => { + let bytes = result.bytes().await.map_err(|error| { + DataFusionError::Plan(format!( + "Failed to read table extensions for table {}: {}", + table_name, error + )) + })?; + let text = String::from_utf8(bytes.to_vec()).map_err(|error| { + DataFusionError::Plan(format!( + "Table extensions for table {} are not valid UTF-8: {}", + table_name, error + )) + })?; + Ok(Some(text)) + } + Err(object_store::Error::NotFound { .. }) => Ok(None), + Err(error) => Err(DataFusionError::Plan(format!( + "Failed to load table extensions for table {}: {}", + table_name, error + ))), + } + } + + /// Remove a table's extensions sidecar. A missing sidecar is not an error. + pub async fn remove_table_extensions_json( + &self, + table_name: &str, + ) -> datafusion::error::Result<()> { + let path = object_store::path::Path::from(format!("{}/extensions.json", table_name)); + let table_object_store = self.table_object_store(table_name)?; + match table_object_store.delete(&path).await { + Ok(()) | Err(object_store::Error::NotFound { .. }) => Ok(()), + Err(error) => Err(DataFusionError::Plan(format!( + "Failed to remove table extensions for table {}: {}", + table_name, error + ))), + } + } + + /// Resolve the object store backing `tables://` table definitions. + fn table_object_store( + &self, + table_name: &str, + ) -> datafusion::error::Result> { + self.session_context + .runtime_env() + .object_store(&self.table_directory_store_url) + .map_err(|error| { + DataFusionError::Plan(format!( + "Failed to get table object store for table {}: {}", + table_name, error + )) + }) + } + pub async fn remove_persisted_table(&self, table_name: &str) -> datafusion::error::Result<()> { let path = object_store::path::Path::from(table_name); let table_object_store = self @@ -257,4 +343,65 @@ mod tests { assert!(err.to_string().contains("requires a SQL definition")); } + + #[tokio::test] + async fn extensions_sidecar_round_trip_and_cleanup() { + let (service, _ctx, table_store, _url) = test_service(); + + // Missing sidecar reads as None. + assert!(service + .load_table_extensions_json("obs") + .await + .expect("load should succeed") + .is_none()); + + // Persist then load returns the stored JSON. + let payload = r#"{"mcp":{"enabled":true}}"#.to_string(); + service + .persist_table_extensions_json("obs", payload.clone()) + .await + .expect("persist should succeed"); + assert_eq!( + service + .load_table_extensions_json("obs") + .await + .expect("load should succeed") + .as_deref(), + Some(payload.as_str()) + ); + assert!(table_store + .get(&Path::from("obs/extensions.json")) + .await + .is_ok()); + + // Explicit removal clears it (and is a no-op when already absent). + service + .remove_table_extensions_json("obs") + .await + .expect("remove should succeed"); + assert!(service + .load_table_extensions_json("obs") + .await + .expect("load should succeed") + .is_none()); + service + .remove_table_extensions_json("obs") + .await + .expect("removing an absent sidecar is not an error"); + + // Dropping the whole table directory removes the sidecar too. + service + .persist_table_extensions_json("obs", payload) + .await + .expect("persist should succeed"); + service + .remove_persisted_table("obs") + .await + .expect("table removal should succeed"); + assert!(service + .load_table_extensions_json("obs") + .await + .expect("load should succeed") + .is_none()); + } }