From 4b34c16ba1441cf5171ee900af012b5864602995 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Tue, 23 Jun 2026 00:36:19 +0200 Subject: [PATCH 1/8] Add typed table extensions (MCP + presets) for downstream consumers Tables can now carry consumer-facing metadata, decoupled from their storage definition: an MCP descriptor (how a downstream MCP server should surface the table) and named query presets (predefined filter sets). Extensions are typed and validated against the live table schema. Storage is a `tables:///extensions.json` sidecar, separate from `table.json`, so extensions apply uniformly to every table type, can be edited without rebuilding the provider, survive provider re-registration (MV refresh, Iceberg alter), and are removed automatically on DROP TABLE. Manage via SQL: SET EXTENSION '' FOR TO '' DROP EXTENSION '' FOR
SHOW EXTENSIONS FOR
or REST: public GET /api/table-extensions; admin PUT/DELETE /api/admin/table-extensions/{name}. All OpenAPI-documented. Validation rejects unknown columns, unsupported operators, malformed between/in values, and duplicate preset names. Tests: extension validation + parser unit tests, a persistence round-trip and cleanup test, and an end-to-end runtime test (CREATE TABLE -> SET/SHOW/DROP EXTENSION -> validation rejection). --- beacon-api/src/axum/admin/extensions.rs | 63 +++ beacon-api/src/axum/admin/mod.rs | 5 + beacon-api/src/axum/client/mod.rs | 1 + beacon-api/src/axum/client/tables.rs | 44 +- beacon-core/src/api.rs | 5 + beacon-core/src/extensions.rs | 433 ++++++++++++++++++ beacon-core/src/lib.rs | 1 + beacon-core/src/parser/beacon_parser.rs | 147 +++++- beacon-core/src/parser/statement.rs | 69 +++ beacon-core/src/runtime.rs | 118 +++++ beacon-core/src/statement_plan/logical.rs | 123 +++++ beacon-core/src/statement_plan/mod.rs | 33 +- beacon-core/src/statement_plan/physical.rs | 148 ++++++ .../src/statement_plan/query_planner.rs | 24 + beacon-data-lake/src/lib.rs | 2 +- .../src/table_runtime/schema_persistence.rs | 147 ++++++ 16 files changed, 1358 insertions(+), 5 deletions(-) create mode 100644 beacon-api/src/axum/admin/extensions.rs create mode 100644 beacon-core/src/extensions.rs 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()); + } } From b3db3282eafdca0e1b9434ca25e5afb25f7292bb Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Tue, 30 Jun 2026 11:39:38 +0200 Subject: [PATCH 2/8] Add MCP server (beacon-mcp) exposing tables and queries over streamable HTTP Introduces a `beacon-mcp` crate that turns beacon into a Model Context Protocol server, mounted at `/mcp` in beacon-api via rmcp's streamable-HTTP transport. MCP clients (e.g. Claude) can discover tables and run read-only queries. Tools are generated from the runtime: - generic: `list_tables`, `describe_table`, `run_sql` (SELECT-only) - one tool per table whose `mcp` table extension is enabled, with inputs derived from the extension metadata (exposed_columns -> `select`, presets -> a `preset` enum that expands to the stored filters), built on the table-extensions feature. Execution runs through Runtime::run_query as a non-super-user (read-only), with results serialized to JSON and capped to 1000 rows. Identifiers are quoted and preset values rendered as SQL literals. Pinned rmcp =1.8.0 (2.0.0 was one day old at time of writing). Tests: catalog/SQL-builder unit tests (preset expansion, column-exposure enforcement, value escaping). Verified end to end against a running server: initialize -> tools/list -> tools/call run_sql returns JSON rows. Docs: docs/mcp.md (tools, exposing a table, connecting Claude). --- Cargo.lock | 128 +++++++++++ Cargo.toml | 2 +- beacon-api/Cargo.toml | 3 +- beacon-api/src/axum/router.rs | 6 + beacon-mcp/Cargo.toml | 16 ++ beacon-mcp/src/catalog.rs | 406 ++++++++++++++++++++++++++++++++++ beacon-mcp/src/lib.rs | 35 +++ beacon-mcp/src/result.rs | 60 +++++ beacon-mcp/src/server.rs | 61 +++++ docs/mcp.md | 43 ++++ 10 files changed, 758 insertions(+), 2 deletions(-) create mode 100644 beacon-mcp/Cargo.toml create mode 100644 beacon-mcp/src/catalog.rs create mode 100644 beacon-mcp/src/lib.rs create mode 100644 beacon-mcp/src/result.rs create mode 100644 beacon-mcp/src/server.rs create mode 100644 docs/mcp.md diff --git a/Cargo.lock b/Cargo.lock index 8dd567cb..2d4273ff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1230,6 +1230,7 @@ dependencies = [ "base64", "beacon-config", "beacon-core", + "beacon-mcp", "bytes", "futures", "futures-util", @@ -1781,6 +1782,22 @@ dependencies = [ "typetag", ] +[[package]] +name = "beacon-mcp" +version = "0.1.0" +dependencies = [ + "anyhow", + "arrow 58.3.0", + "arrow-json 58.3.0", + "beacon-core", + "futures", + "rmcp", + "serde", + "serde_json", + "tokio", + "tracing", +] + [[package]] name = "beacon-nd-array" version = "0.1.0" @@ -3903,6 +3920,12 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "earcutr" version = "0.4.3" @@ -6599,6 +6622,12 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pastey" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ee67f1008b1ba2321834326597b8e186293b049a023cdef258527550b9935b4" + [[package]] name = "pathdiff" version = "0.2.3" @@ -7325,6 +7354,26 @@ dependencies = [ "thiserror 2.0.18", ] +[[package]] +name = "ref-cast" +version = "1.0.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f354300ae66f76f1c85c5f84693f0ce81d747e2c3f21a45fef496d89c960bf7d" +dependencies = [ + "ref-cast-impl", +] + +[[package]] +name = "ref-cast-impl" +version = "1.0.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "regex" version = "1.12.3" @@ -7556,6 +7605,35 @@ dependencies = [ "libc", ] +[[package]] +name = "rmcp" +version = "1.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d1f571c72940a19d9532fe52dbea8bc9912bf1d766c2970bb824056b86f3f59" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "futures", + "http", + "http-body", + "http-body-util", + "pastey", + "pin-project-lite", + "rand 0.10.1", + "schemars", + "serde", + "serde_json", + "sse-stream", + "thiserror 2.0.18", + "tokio", + "tokio-stream", + "tokio-util", + "tower-service", + "tracing", + "uuid", +] + [[package]] name = "rmp" version = "0.8.15" @@ -7850,6 +7928,32 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "schemars" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2b42f36aa1cd011945615b92222f6bf73c599a102a300334cd7f8dbeec726cc" +dependencies = [ + "chrono", + "dyn-clone", + "ref-cast", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d115b50f4aaeea07e79c1912f645c7513d81715d0420f8bc77a18c6260b307f" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn 2.0.117", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -7974,6 +8078,17 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "serde_derive_internals" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "serde_json" version = "1.0.150" @@ -8269,6 +8384,19 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "sse-stream" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3962b63f038885f15bce2c6e02c0e7925c072f1ac86bb60fd44c5c6b762fb72" +dependencies = [ + "bytes", + "futures-util", + "http-body", + "http-body-util", + "pin-project-lite", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" diff --git a/Cargo.toml b/Cargo.toml index f2ada597..5b6ecdea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["beacon-api", "beacon-file-formats/beacon-arrow-netcdf", "beacon-file-formats/beacon-arrow-odv", "beacon-common", "beacon-config", "beacon-core", "beacon-functions", "beacon-data-lake", "beacon-sql-databases", "beacon-file-formats/beacon-delta", "beacon-file-formats/beacon-arrow-zarr", "beacon-file-formats/beacon-binary-format", "beacon-object-storage", "beacon-file-formats/beacon-nd-arrow", "beacon-datafusion-ext", "beacon-iceberg", "beacon-file-formats/beacon-nd-array", "beacon-file-formats/beacon-arrow-tiff", "beacon-file-formats/beacon-arrow-atlas", "beacon-file-formats/beacon-arrow-geoparquet", "beacon-file-formats/beacon-arrow-bbf", "beacon-file-formats/beacon-arrow-ipc", "beacon-file-formats/beacon-arrow-csv", "beacon-file-formats/beacon-arrow-parquet"] +members = ["beacon-api", "beacon-mcp", "beacon-file-formats/beacon-arrow-netcdf", "beacon-file-formats/beacon-arrow-odv", "beacon-common", "beacon-config", "beacon-core", "beacon-functions", "beacon-data-lake", "beacon-sql-databases", "beacon-file-formats/beacon-delta", "beacon-file-formats/beacon-arrow-zarr", "beacon-file-formats/beacon-binary-format", "beacon-object-storage", "beacon-file-formats/beacon-nd-arrow", "beacon-datafusion-ext", "beacon-iceberg", "beacon-file-formats/beacon-nd-array", "beacon-file-formats/beacon-arrow-tiff", "beacon-file-formats/beacon-arrow-atlas", "beacon-file-formats/beacon-arrow-geoparquet", "beacon-file-formats/beacon-arrow-bbf", "beacon-file-formats/beacon-arrow-ipc", "beacon-file-formats/beacon-arrow-csv", "beacon-file-formats/beacon-arrow-parquet"] exclude = ["beacon-file-formats/beacon-binary-format-toolbox"] [workspace.dependencies] diff --git a/beacon-api/Cargo.toml b/beacon-api/Cargo.toml index 3d9093d5..c2e16a9c 100644 --- a/beacon-api/Cargo.toml +++ b/beacon-api/Cargo.toml @@ -39,4 +39,5 @@ prost = "0.14.1" # Local dependencies beacon-config = { path = "../beacon-config" } -beacon-core = { path = "../beacon-core" } \ No newline at end of file +beacon-core = { path = "../beacon-core" } +beacon-mcp = { path = "../beacon-mcp" } \ No newline at end of file diff --git a/beacon-api/src/axum/router.rs b/beacon-api/src/axum/router.rs index bb719768..17a1c757 100644 --- a/beacon-api/src/axum/router.rs +++ b/beacon-api/src/axum/router.rs @@ -56,6 +56,12 @@ pub(crate) fn setup_router( let router = client_router .merge(admin_router) + // MCP streamable-HTTP endpoint (read-only; tools generated from table + // `mcp` extensions). Mounted as a tower service on a single path. + .route_service( + "/mcp", + beacon_mcp::streamable_http_service(beacon_runtime.clone()), + ) .merge(Scalar::with_url("/scalar/", docs.clone())) .route( "/scalar", diff --git a/beacon-mcp/Cargo.toml b/beacon-mcp/Cargo.toml new file mode 100644 index 00000000..56a86b24 --- /dev/null +++ b/beacon-mcp/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "beacon-mcp" +version = "0.1.0" +edition = "2021" + +[dependencies] +beacon-core = { path = "../beacon-core" } +rmcp = { version = "=1.8.0", default-features = false, features = ["server", "transport-streamable-http-server"] } +tokio = { workspace = true } +futures = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +anyhow = { workspace = true } +tracing = { workspace = true } +arrow = { workspace = true } +arrow-json = "58" diff --git a/beacon-mcp/src/catalog.rs b/beacon-mcp/src/catalog.rs new file mode 100644 index 00000000..4f374995 --- /dev/null +++ b/beacon-mcp/src/catalog.rs @@ -0,0 +1,406 @@ +//! Generates the MCP tool catalog from the runtime and dispatches tool calls. +//! +//! Three generic tools are always present (`list_tables`, `describe_table`, +//! `run_sql`). In addition, every table whose `mcp` extension is enabled becomes +//! its own tool whose input schema is derived from the extension metadata +//! (exposed columns + named presets). + +use std::sync::Arc; + +use beacon_core::extensions::{McpExtension, PresetExtension, PresetFilter}; +use beacon_core::runtime::Runtime; +use rmcp::model::Tool; +use serde_json::{json, Map, Value}; + +use crate::result::{run_sql_to_json, MAX_ROWS}; + +/// Build the full tool list: generic tools + per-table tools from extensions. +pub async fn build_tools(runtime: &Arc) -> anyhow::Result> { + let mut tools = vec![list_tables_tool(), describe_table_tool(), run_sql_tool()]; + for table in runtime.list_tables() { + let ext = runtime + .get_table_extensions(table.clone()) + .await + .unwrap_or_default(); + if ext.mcp.as_ref().is_some_and(|mcp| mcp.enabled) { + tools.push(table_tool(runtime, &table, &ext.mcp.unwrap(), ext.preset.as_ref()).await); + } + } + Ok(tools) +} + +/// Route a tool call to its handler. +pub async fn dispatch( + runtime: &Arc, + name: &str, + args: Map, +) -> anyhow::Result { + match name { + "list_tables" => list_tables_json(runtime).await, + "describe_table" => describe_table_json(runtime, &args).await, + "run_sql" => { + let sql = args + .get("sql") + .and_then(Value::as_str) + .ok_or_else(|| anyhow::anyhow!("missing required 'sql' argument"))?; + run_sql_to_json(runtime, sql.to_string()).await + } + other => run_table_tool(runtime, other, &args).await, + } +} + +// ---- generic tools ------------------------------------------------------- + +fn object_schema(props: Value, required: &[&str]) -> Map { + let mut schema = Map::new(); + schema.insert("type".into(), json!("object")); + schema.insert("properties".into(), props); + if !required.is_empty() { + schema.insert("required".into(), json!(required)); + } + schema +} + +fn list_tables_tool() -> Tool { + Tool::new( + "list_tables", + "List the tables registered in beacon, with their MCP exposure status.", + object_schema(json!({}), &[]), + ) +} + +fn describe_table_tool() -> Tool { + Tool::new( + "describe_table", + "Return a table's column schema and its attached extensions (MCP descriptor, presets).", + object_schema( + json!({ "table_name": { "type": "string", "description": "Name of the table." } }), + &["table_name"], + ), + ) +} + +fn run_sql_tool() -> Tool { + Tool::new( + "run_sql", + "Run a read-only SQL query (SELECT only) against beacon and return JSON rows.", + object_schema( + json!({ "sql": { "type": "string", "description": "A read-only SELECT statement." } }), + &["sql"], + ), + ) +} + +async fn list_tables_json(runtime: &Arc) -> anyhow::Result { + let mut out = Vec::new(); + for table in runtime.list_tables() { + let ext = runtime + .get_table_extensions(table.clone()) + .await + .unwrap_or_default(); + out.push(json!({ + "name": table, + "mcp_enabled": ext.mcp.as_ref().map(|m| m.enabled).unwrap_or(false), + "description": ext.mcp.as_ref().and_then(|m| m.description.clone()), + })); + } + Ok(serde_json::to_string_pretty(&out)?) +} + +async fn describe_table_json( + runtime: &Arc, + args: &Map, +) -> anyhow::Result { + let table = args + .get("table_name") + .and_then(Value::as_str) + .ok_or_else(|| anyhow::anyhow!("missing required 'table_name' argument"))?; + let schema = runtime + .list_table_schema_view(table.to_string()) + .await + .ok_or_else(|| anyhow::anyhow!("table '{table}' not found"))?; + let ext = runtime + .get_table_extensions(table.to_string()) + .await + .unwrap_or_default(); + let columns: Vec = schema + .fields + .iter() + .map(|f| json!({ "name": f.name, "data_type": f.data_type, "nullable": f.nullable })) + .collect(); + Ok(serde_json::to_string_pretty( + &json!({ "name": table, "columns": columns, "extensions": ext }), + )?) +} + +// ---- per-table tools ----------------------------------------------------- + +fn default_tool_name(table: &str) -> String { + format!("query_{table}") +} + +async fn table_tool( + runtime: &Arc, + table: &str, + mcp: &McpExtension, + preset: Option<&PresetExtension>, +) -> Tool { + let name = mcp + .tool_name + .clone() + .unwrap_or_else(|| default_tool_name(table)); + let description = mcp + .description + .clone() + .unwrap_or_else(|| format!("Query the '{table}' table.")); + + // Columns offered to the model: the curated `exposed_columns` if set, + // otherwise the table's full schema. + let columns: Vec = match &mcp.exposed_columns { + Some(cols) => cols.clone(), + None => runtime + .list_table_schema_view(table.to_string()) + .await + .map(|s| s.fields.into_iter().map(|f| f.name).collect()) + .unwrap_or_default(), + }; + let preset_names: Vec = preset + .map(|p| p.presets.iter().map(|x| x.name.clone()).collect()) + .unwrap_or_default(); + + let mut props = Map::new(); + if !columns.is_empty() { + props.insert( + "select".into(), + json!({ + "type": "array", + "items": { "type": "string", "enum": columns }, + "description": "Columns to return. Omit for all exposed columns." + }), + ); + } + if !preset_names.is_empty() { + props.insert( + "preset".into(), + json!({ + "type": "string", + "enum": preset_names, + "description": "Apply a predefined, named filter set." + }), + ); + } + props.insert( + "limit".into(), + json!({ "type": "integer", "description": "Maximum rows to return (default 100)." }), + ); + + Tool::new(name, description, object_schema(Value::Object(props), &[])) +} + +async fn run_table_tool( + runtime: &Arc, + tool_name: &str, + args: &Map, +) -> anyhow::Result { + for table in runtime.list_tables() { + let ext = runtime + .get_table_extensions(table.clone()) + .await + .unwrap_or_default(); + let Some(mcp) = ext.mcp.as_ref().filter(|m| m.enabled) else { + continue; + }; + let name = mcp + .tool_name + .clone() + .unwrap_or_else(|| default_tool_name(&table)); + if name != tool_name { + continue; + } + let sql = build_table_sql(&table, mcp, ext.preset.as_ref(), args)?; + return run_sql_to_json(runtime, sql).await; + } + anyhow::bail!("unknown tool '{tool_name}'") +} + +/// Build a `SELECT` for a per-table tool from its args, expanding a chosen preset +/// into `WHERE` clauses. Identifiers are quoted and values rendered as literals. +fn build_table_sql( + table: &str, + mcp: &McpExtension, + preset: Option<&PresetExtension>, + args: &Map, +) -> anyhow::Result { + let exposed = mcp.exposed_columns.as_ref(); + + let select = match args.get("select").and_then(Value::as_array) { + Some(arr) if !arr.is_empty() => { + let cols: Vec = arr + .iter() + .filter_map(|v| v.as_str().map(String::from)) + .collect(); + if let Some(exp) = exposed { + for col in &cols { + anyhow::ensure!(exp.contains(col), "column '{col}' is not exposed by this tool"); + } + } + cols.iter().map(|c| quote_ident(c)).collect::>().join(", ") + } + _ => default_select(exposed), + }; + + let mut clauses = Vec::new(); + if let Some(name) = args.get("preset").and_then(Value::as_str) { + let preset = preset + .and_then(|p| p.presets.iter().find(|x| x.name == name)) + .ok_or_else(|| anyhow::anyhow!("unknown preset '{name}'"))?; + for filter in &preset.filters { + clauses.push(render_filter(filter)?); + } + } + + let limit = args + .get("limit") + .and_then(Value::as_u64) + .unwrap_or(100) + .min(MAX_ROWS as u64); + + let mut sql = format!("SELECT {select} FROM {}", quote_ident(table)); + if !clauses.is_empty() { + sql.push_str(" WHERE "); + sql.push_str(&clauses.join(" AND ")); + } + sql.push_str(&format!(" LIMIT {limit}")); + Ok(sql) +} + +fn default_select(exposed: Option<&Vec>) -> String { + match exposed { + Some(cols) if !cols.is_empty() => { + cols.iter().map(|c| quote_ident(c)).collect::>().join(", ") + } + _ => "*".to_string(), + } +} + +/// Render a stored preset filter into a SQL boolean expression. +fn render_filter(filter: &PresetFilter) -> anyhow::Result { + let col = quote_ident(&filter.column); + match filter.op.as_str() { + "=" | "!=" | "<" | "<=" | ">" | ">=" => { + Ok(format!("{col} {} {}", filter.op, render_scalar(&filter.value)?)) + } + "between" => { + let arr = filter + .value + .as_array() + .filter(|a| a.len() == 2) + .ok_or_else(|| anyhow::anyhow!("'between' requires a two-element array"))?; + Ok(format!( + "{col} BETWEEN {} AND {}", + render_scalar(&arr[0])?, + render_scalar(&arr[1])? + )) + } + "in" => { + let arr = filter + .value + .as_array() + .filter(|a| !a.is_empty()) + .ok_or_else(|| anyhow::anyhow!("'in' requires a non-empty array"))?; + let vals = arr.iter().map(render_scalar).collect::>>()?; + Ok(format!("{col} IN ({})", vals.join(", "))) + } + other => anyhow::bail!("unsupported operator '{other}'"), + } +} + +fn render_scalar(value: &Value) -> anyhow::Result { + Ok(match value { + Value::Number(n) => n.to_string(), + Value::String(s) => format!("'{}'", s.replace('\'', "''")), + Value::Bool(b) => if *b { "TRUE".into() } else { "FALSE".into() }, + Value::Null => "NULL".into(), + other => anyhow::bail!("unsupported filter value: {other}"), + }) +} + +/// Quote a SQL identifier, escaping embedded double quotes. +fn quote_ident(ident: &str) -> String { + format!("\"{}\"", ident.replace('"', "\"\"")) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn preset(name: &str, filters: Vec) -> PresetExtension { + PresetExtension { + presets: vec![beacon_core::extensions::Preset { + name: name.to_string(), + description: None, + filters, + }], + } + } + + fn mcp(cols: Option>) -> McpExtension { + McpExtension { + enabled: true, + tool_name: None, + description: None, + exposed_columns: cols.map(|c| c.into_iter().map(String::from).collect()), + } + } + + #[test] + fn builds_select_with_preset_between() { + let p = preset( + "shallow", + vec![PresetFilter { + column: "depth".into(), + op: "between".into(), + value: serde_json::json!([0, 10]), + }], + ); + let mut args = Map::new(); + args.insert("preset".into(), Value::String("shallow".into())); + let sql = build_table_sql("obs", &mcp(Some(vec!["lat", "depth"])), Some(&p), &args).unwrap(); + assert_eq!( + sql, + r#"SELECT "lat", "depth" FROM "obs" WHERE "depth" BETWEEN 0 AND 10 LIMIT 100"# + ); + } + + #[test] + fn rejects_unexposed_select_column() { + let mut args = Map::new(); + args.insert("select".into(), serde_json::json!(["ghost"])); + let err = build_table_sql("obs", &mcp(Some(vec!["lat"])), None, &args).unwrap_err(); + assert!(err.to_string().contains("not exposed"), "{err}"); + } + + #[test] + fn unknown_preset_errors() { + let mut args = Map::new(); + args.insert("preset".into(), Value::String("nope".into())); + let err = build_table_sql("obs", &mcp(None), None, &args).unwrap_err(); + assert!(err.to_string().contains("unknown preset"), "{err}"); + } + + #[test] + fn in_and_string_values_are_escaped() { + let f = PresetFilter { + column: "basin".into(), + op: "in".into(), + value: serde_json::json!(["a'b", "c"]), + }; + assert_eq!(render_filter(&f).unwrap(), r#""basin" IN ('a''b', 'c')"#); + } + + #[test] + fn default_select_is_star_without_exposed() { + let sql = build_table_sql("obs", &mcp(None), None, &Map::new()).unwrap(); + assert_eq!(sql, r#"SELECT * FROM "obs" LIMIT 100"#); + } +} diff --git a/beacon-mcp/src/lib.rs b/beacon-mcp/src/lib.rs new file mode 100644 index 00000000..51e1de13 --- /dev/null +++ b/beacon-mcp/src/lib.rs @@ -0,0 +1,35 @@ +//! Model Context Protocol (MCP) server for beacon. +//! +//! Exposes beacon as an MCP server over the streamable-HTTP transport so MCP +//! clients (e.g. Claude) can discover tables and run read-only queries. The tool +//! surface is generated from the runtime: a few generic tools plus one tool per +//! table that opts in via its `mcp` table extension (see +//! [`beacon_core::extensions`]). Presets declared on a table become a typed +//! `preset` parameter that expands to the stored filters. +//! +//! All execution flows through [`beacon_core::runtime::Runtime::run_query`] as a +//! non-super-user, so only read-only `SELECT`s are permitted. + +mod catalog; +mod result; +mod server; + +use std::sync::Arc; + +use beacon_core::runtime::Runtime; +use rmcp::transport::streamable_http_server::session::local::LocalSessionManager; +use rmcp::transport::{StreamableHttpServerConfig, StreamableHttpService}; + +pub use server::BeaconMcpServer; + +/// Build the MCP streamable-HTTP tower service, ready to mount in an axum router +/// (e.g. `Router::route_service("/mcp", beacon_mcp::streamable_http_service(rt))`). +pub fn streamable_http_service( + runtime: Arc, +) -> StreamableHttpService { + StreamableHttpService::new( + move || Ok(BeaconMcpServer::new(runtime.clone())), + Arc::new(LocalSessionManager::default()), + StreamableHttpServerConfig::default(), + ) +} diff --git a/beacon-mcp/src/result.rs b/beacon-mcp/src/result.rs new file mode 100644 index 00000000..86e88a8b --- /dev/null +++ b/beacon-mcp/src/result.rs @@ -0,0 +1,60 @@ +//! Query execution and result serialization for MCP tools. + +use std::sync::Arc; + +use arrow::record_batch::RecordBatch; +use beacon_core::query::Query; +use beacon_core::runtime::Runtime; +use futures::TryStreamExt; + +/// Hard cap on rows returned to the model, to keep tool output bounded. +pub const MAX_ROWS: usize = 1000; + +/// Run a read-only SQL query and return the rows as a JSON string. Executes as a +/// non-super-user, so DDL/DML is rejected by the planner. +pub async fn run_sql_to_json(runtime: &Arc, sql: String) -> anyhow::Result { + let result = runtime.run_query(Query::sql(sql), false).await?; + let mut stream = result.into_record_stream()?; + + let mut batches: Vec = Vec::new(); + let mut total = 0usize; + let mut truncated = false; + while let Some(batch) = stream.try_next().await? { + let remaining = MAX_ROWS - total; + if batch.num_rows() > remaining { + batches.push(batch.slice(0, remaining)); + truncated = true; + break; + } + total += batch.num_rows(); + batches.push(batch); + if total >= MAX_ROWS { + // Peek no further; assume more may exist. + truncated = stream.try_next().await?.is_some(); + break; + } + } + + let rows = batches_to_json(&batches)?; + if truncated { + Ok(format!( + "{{\"truncated\":true,\"max_rows\":{MAX_ROWS},\"rows\":{rows}}}" + )) + } else { + Ok(rows) + } +} + +/// Serialize record batches to a JSON array of row objects. +fn batches_to_json(batches: &[RecordBatch]) -> anyhow::Result { + if batches.iter().all(|b| b.num_rows() == 0) { + return Ok("[]".to_string()); + } + let mut buf = Vec::new(); + let mut writer = arrow_json::ArrayWriter::new(&mut buf); + for batch in batches { + writer.write(batch)?; + } + writer.finish()?; + Ok(String::from_utf8(buf)?) +} diff --git a/beacon-mcp/src/server.rs b/beacon-mcp/src/server.rs new file mode 100644 index 00000000..3d3aa1c9 --- /dev/null +++ b/beacon-mcp/src/server.rs @@ -0,0 +1,61 @@ +//! The [`ServerHandler`] implementation: capabilities, tool listing, dispatch. + +use std::sync::Arc; + +use beacon_core::runtime::Runtime; +use rmcp::handler::server::ServerHandler; +use rmcp::model::{ + CallToolRequestParams, CallToolResult, Content, ListToolsResult, PaginatedRequestParams, + ServerCapabilities, ServerInfo, +}; +use rmcp::service::RequestContext; +use rmcp::{ErrorData, RoleServer}; + +/// MCP server backed by a beacon [`Runtime`]. Cloned per session by the +/// transport; the runtime handle is shared. +#[derive(Clone)] +pub struct BeaconMcpServer { + runtime: Arc, +} + +impl BeaconMcpServer { + pub fn new(runtime: Arc) -> Self { + Self { runtime } + } +} + +impl ServerHandler for BeaconMcpServer { + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()).with_instructions( + "Beacon data lake. Call `list_tables` to discover tables, \ + `describe_table` for a table's schema and available presets, the \ + per-table tools to query curated datasets (optionally via a named \ + preset), and `run_sql` for read-only SQL (SELECT only).", + ) + } + + async fn list_tools( + &self, + _request: Option, + _context: RequestContext, + ) -> Result { + let tools = crate::catalog::build_tools(&self.runtime) + .await + .map_err(|error| ErrorData::internal_error(error.to_string(), None))?; + Ok(ListToolsResult::with_all_items(tools)) + } + + async fn call_tool( + &self, + request: CallToolRequestParams, + _context: RequestContext, + ) -> Result { + let args = request.arguments.unwrap_or_default(); + match crate::catalog::dispatch(&self.runtime, request.name.as_ref(), args).await { + Ok(text) => Ok(CallToolResult::success(vec![Content::text(text)])), + // Surface tool failures as an error result (not a protocol error) so + // the model can read and react to the message. + Err(error) => Ok(CallToolResult::error(vec![Content::text(error.to_string())])), + } + } +} diff --git a/docs/mcp.md b/docs/mcp.md new file mode 100644 index 00000000..71da32f4 --- /dev/null +++ b/docs/mcp.md @@ -0,0 +1,43 @@ +# MCP server + +Beacon ships an [MCP](https://modelcontextprotocol.io) server over the +streamable-HTTP transport at `POST/GET/DELETE /mcp`, so MCP clients (e.g. Claude) +can discover tables and run read-only queries. + +## Tools + +Generated dynamically from the runtime: + +- `list_tables` — registered tables and their MCP exposure status. +- `describe_table` — a table's column schema plus its extensions (MCP descriptor, presets). +- `run_sql` — run a read-only `SELECT` and get JSON rows. +- **one tool per table** whose `mcp` extension is enabled (see below). Its inputs are + derived from the extension: `select` (restricted to `exposed_columns`), `preset` + (an enum of the table's preset names, expanded to filters), and `limit`. + +All execution runs as a non-super-user, so only `SELECT` is permitted. Results are +capped (1000 rows) to keep tool output bounded. + +## Exposing a table to MCP + +Use the table-extensions surface (SQL or REST): + +```sql +SET EXTENSION 'mcp' FOR obs TO '{"enabled":true,"tool_name":"query_obs","description":"Ocean observations","exposed_columns":["lat","lon","depth","temperature"]}'; +SET EXTENSION 'preset' FOR obs TO '{"presets":[{"name":"shallow","filters":[{"column":"depth","op":"<=","value":10}]}]}'; +``` + +`query_obs` then appears as an MCP tool with a `preset: "shallow"` option. + +## Connecting Claude + +**HTTP (Claude Code / API):** point the client at `http://:/mcp`. + +**Claude Desktop** (via an HTTP-capable MCP entry): + +```json +{ "mcpServers": { "beacon": { "url": "http://localhost:5001/mcp" } } } +``` + +> The endpoint is currently unauthenticated and read-only. Put it behind your +> existing auth/proxy if exposing beyond localhost. From 53d11d767ea5573c8d658f4aa9ae79d205157e71 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Tue, 30 Jun 2026 13:14:40 +0200 Subject: [PATCH 3/8] Enforce read-only MCP: clear super-user on every tool call MCP tool execution now always runs with is_super_user cleared, so the query planner rejects DDL/DML for any caller (defense-in-depth, independent of how the identity was resolved). The caller's roles are preserved so per-user read grants still apply. Docs updated. --- beacon-mcp/src/server.rs | 11 +++++++++-- docs/mcp.md | 13 +++++++++---- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/beacon-mcp/src/server.rs b/beacon-mcp/src/server.rs index f5ece076..e48831e0 100644 --- a/beacon-mcp/src/server.rs +++ b/beacon-mcp/src/server.rs @@ -66,10 +66,17 @@ impl ServerHandler for BeaconMcpServer { /// middleware and carried in the HTTP request parts that the streamable-HTTP /// transport injects into the MCP request context. Falls back to a role-less /// identity (no access) when absent. +/// +/// The MCP surface is strictly read-only: the returned identity always has +/// `is_super_user` cleared, so the query planner rejects any DDL/DML regardless +/// of the caller's privileges. The caller's `roles` are preserved so per-user +/// read grants (RBAC) still apply. fn identity_from_context(context: &RequestContext) -> AuthIdentity { - context + let mut identity = context .extensions .get::() .and_then(|parts| parts.extensions.get::().cloned()) - .unwrap_or_else(AuthIdentity::empty) + .unwrap_or_else(AuthIdentity::empty); + identity.is_super_user = false; + identity } diff --git a/docs/mcp.md b/docs/mcp.md index 71da32f4..7f6803d5 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -15,8 +15,11 @@ Generated dynamically from the runtime: derived from the extension: `select` (restricted to `exposed_columns`), `preset` (an enum of the table's preset names, expanded to filters), and `limit`. -All execution runs as a non-super-user, so only `SELECT` is permitted. Results are -capped (1000 rows) to keep tool output bounded. +The MCP surface is **strictly read-only**: every tool call executes with +`is_super_user` cleared, so the query planner rejects any DDL/DML (`CREATE`, +`INSERT`, `UPDATE`, `DELETE`, `SET EXTENSION`, …) regardless of who connects — +only `SELECT` runs. The caller's roles are preserved, so per-user read grants +(RBAC) still apply. Results are capped (1000 rows) to keep tool output bounded. ## Exposing a table to MCP @@ -39,5 +42,7 @@ SET EXTENSION 'preset' FOR obs TO '{"presets":[{"name":"shallow","filters":[{"co { "mcpServers": { "beacon": { "url": "http://localhost:5001/mcp" } } } ``` -> The endpoint is currently unauthenticated and read-only. Put it behind your -> existing auth/proxy if exposing beyond localhost. +> The endpoint rides the same `resolve_identity` middleware as the client API: +> requests authenticate via `Authorization` (resolving to that user's roles), or +> the anonymous principal when enabled, or a role-less identity otherwise. It is +> read-only regardless. Gate it with `BEACON_MCP_ENABLED=false` to disable. From 80a52c890c19bb2c4889028e17448b3a51265ddd Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Tue, 30 Jun 2026 14:18:26 +0200 Subject: [PATCH 4/8] Validate preset payloads strictly: typed op enum + deny_unknown_fields Preset/MCP extension payloads now parse against a strict structure: - PresetFilter.op is a typed PresetOp enum (= != < <= > >= between in), serialized with the symbolic spelling; any other operator is rejected at parse time with a clear error instead of only at schema-validation. - All extension structs use serde(deny_unknown_fields), so typos/extra keys are rejected rather than silently dropped. This guarantees stored extension JSON conforms to a known shape that consumers can reliably parse. Updated tests + API re-export. --- beacon-core/src/api.rs | 2 +- beacon-core/src/extensions.rs | 92 ++++++++++++++++++++++++++--------- beacon-mcp/src/catalog.rs | 17 +++---- 3 files changed, 78 insertions(+), 33 deletions(-) diff --git a/beacon-core/src/api.rs b/beacon-core/src/api.rs index 2b5f7d23..5d1fe3dd 100644 --- a/beacon-core/src/api.rs +++ b/beacon-core/src/api.rs @@ -14,7 +14,7 @@ use utoipa::ToSchema; /// Re-exported typed table-extension contracts (see [`crate::extensions`]). pub use crate::extensions::{ - McpExtension, Preset, PresetExtension, PresetFilter, TableExtensions, + McpExtension, Preset, PresetExtension, PresetFilter, PresetOp, TableExtensions, }; /// A single parameter of a registered function. diff --git a/beacon-core/src/extensions.rs b/beacon-core/src/extensions.rs index 6f99ea7a..52eda801 100644 --- a/beacon-core/src/extensions.rs +++ b/beacon-core/src/extensions.rs @@ -25,11 +25,49 @@ use serde::{Deserialize, Serialize}; use utoipa::ToSchema; /// The comparison operators a [`PresetFilter`] may use. -pub const PRESET_OPS: [&str; 8] = ["=", "!=", "<", "<=", ">", ">=", "between", "in"]; +/// The comparison operators a [`PresetFilter`] may use. Serialized with the +/// symbolic/SQL spelling shown, so stored preset JSON stays human-readable, and +/// any other value is rejected at parse time with a clear error. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, ToSchema)] +pub enum PresetOp { + #[serde(rename = "=")] + Eq, + #[serde(rename = "!=")] + Ne, + #[serde(rename = "<")] + Lt, + #[serde(rename = "<=")] + Lte, + #[serde(rename = ">")] + Gt, + #[serde(rename = ">=")] + Gte, + #[serde(rename = "between")] + Between, + #[serde(rename = "in")] + In, +} + +impl PresetOp { + /// The SQL spelling of this operator. + pub fn as_sql(self) -> &'static str { + match self { + PresetOp::Eq => "=", + PresetOp::Ne => "!=", + PresetOp::Lt => "<", + PresetOp::Lte => "<=", + PresetOp::Gt => ">", + PresetOp::Gte => ">=", + PresetOp::Between => "BETWEEN", + PresetOp::In => "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)] +#[serde(deny_unknown_fields)] pub struct TableExtensions { /// MCP descriptor: how downstream MCP servers should surface this table. #[serde(default, skip_serializing_if = "Option::is_none")] @@ -48,6 +86,7 @@ pub struct TableExtensions { "description": "Argo float observations by location, depth, and time.", "exposed_columns": ["lat", "lon", "depth", "temperature", "time"] }))] +#[serde(deny_unknown_fields)] pub struct McpExtension { /// Whether downstream MCP servers should expose this table at all. #[serde(default)] @@ -75,6 +114,7 @@ pub struct McpExtension { ] }] }))] +#[serde(deny_unknown_fields)] pub struct PresetExtension { /// The named presets. pub presets: Vec, @@ -82,6 +122,7 @@ pub struct PresetExtension { /// A single named preset: a bundle of filters applied together. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)] +#[serde(deny_unknown_fields)] pub struct Preset { /// Unique (within the table) preset name. pub name: String, @@ -94,12 +135,12 @@ pub struct Preset { /// A single predefined filter within a [`Preset`]. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)] +#[serde(deny_unknown_fields)] 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, + /// Comparison operator. + pub op: PresetOp, /// Filter value: a scalar, `[lo, hi]` for `between`, or `[..]` for `in`. #[schema(value_type = Object)] pub value: serde_json::Value, @@ -176,14 +217,6 @@ impl PresetExtension { } 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)?; } } @@ -193,8 +226,8 @@ impl PresetExtension { /// `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" => { + match filter.op { + PresetOp::Between => { let ok = filter.value.as_array().is_some_and(|a| a.len() == 2); anyhow::ensure!( ok, @@ -202,7 +235,7 @@ fn validate_filter_value_shape(preset: &str, filter: &PresetFilter) -> anyhow::R filter.column ); } - "in" => { + PresetOp::In => { let ok = filter.value.as_array().is_some_and(|a| !a.is_empty()); anyhow::ensure!( ok, @@ -373,14 +406,29 @@ mod tests { } #[test] - fn rejects_bad_operator() { + fn rejects_bad_operator_at_parse() { + // An unknown operator no longer parses into the typed `op` enum at all. 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")); + let err = ext + .set_kind( + "preset", + r#"{"presets":[{"name":"p","filters":[{"column":"lat","op":"~~","value":1}]}]}"#, + ) + .unwrap_err(); + assert!(err.to_string().contains("preset"), "unexpected: {err}"); + } + + #[test] + fn rejects_unknown_field_at_parse() { + // `deny_unknown_fields` rejects typos/extra keys instead of dropping them. + let mut ext = TableExtensions::default(); + assert!(ext + .set_kind( + "preset", + r#"{"presets":[{"name":"p","filters":[],"bogus":1}]}"#, + ) + .is_err()); + assert!(ext.set_kind("mcp", r#"{"enabled":true,"nope":1}"#).is_err()); } #[test] diff --git a/beacon-mcp/src/catalog.rs b/beacon-mcp/src/catalog.rs index 2850a224..9cc1dd43 100644 --- a/beacon-mcp/src/catalog.rs +++ b/beacon-mcp/src/catalog.rs @@ -7,7 +7,7 @@ use std::sync::Arc; -use beacon_core::extensions::{McpExtension, PresetExtension, PresetFilter}; +use beacon_core::extensions::{McpExtension, PresetExtension, PresetFilter, PresetOp}; use beacon_core::runtime::Runtime; use beacon_core::AuthIdentity; use rmcp::model::Tool; @@ -289,11 +289,8 @@ fn default_select(exposed: Option<&Vec>) -> String { /// Render a stored preset filter into a SQL boolean expression. fn render_filter(filter: &PresetFilter) -> anyhow::Result { let col = quote_ident(&filter.column); - match filter.op.as_str() { - "=" | "!=" | "<" | "<=" | ">" | ">=" => { - Ok(format!("{col} {} {}", filter.op, render_scalar(&filter.value)?)) - } - "between" => { + match filter.op { + PresetOp::Between => { let arr = filter .value .as_array() @@ -305,7 +302,7 @@ fn render_filter(filter: &PresetFilter) -> anyhow::Result { render_scalar(&arr[1])? )) } - "in" => { + PresetOp::In => { let arr = filter .value .as_array() @@ -314,7 +311,7 @@ fn render_filter(filter: &PresetFilter) -> anyhow::Result { let vals = arr.iter().map(render_scalar).collect::>>()?; Ok(format!("{col} IN ({})", vals.join(", "))) } - other => anyhow::bail!("unsupported operator '{other}'"), + op => Ok(format!("{col} {} {}", op.as_sql(), render_scalar(&filter.value)?)), } } @@ -362,7 +359,7 @@ mod tests { "shallow", vec![PresetFilter { column: "depth".into(), - op: "between".into(), + op: PresetOp::Between, value: serde_json::json!([0, 10]), }], ); @@ -395,7 +392,7 @@ mod tests { fn in_and_string_values_are_escaped() { let f = PresetFilter { column: "basin".into(), - op: "in".into(), + op: PresetOp::In, value: serde_json::json!(["a'b", "c"]), }; assert_eq!(render_filter(&f).unwrap(), r#""basin" IN ('a''b', 'c')"#); From 439485e40cbfe554809dad97dce6dda767005225 Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Tue, 30 Jun 2026 14:35:46 +0200 Subject: [PATCH 5/8] Align MCP extension to the MCP Tool standard + validate tool_name - Validate mcp tool_name against MCP/Anthropic rules (1-64 chars [A-Za-z0-9_-]) so a bad name can't break a client's tools/list; sanitize generated defaults. - Add optional `title` (maps to Tool.title) and emit annotations.readOnlyHint on every generated tool, matching the MCP Tool descriptor. - McpExtension already parses strictly (deny_unknown_fields from prior change). - Tests + docs updated. --- beacon-core/src/extensions.rs | 35 ++++++++++++++++++++++++++++++ beacon-mcp/src/catalog.rs | 41 +++++++++++++++++++++++++++-------- docs/mcp.md | 7 ++++++ 3 files changed, 74 insertions(+), 9 deletions(-) diff --git a/beacon-core/src/extensions.rs b/beacon-core/src/extensions.rs index 52eda801..1d5223d8 100644 --- a/beacon-core/src/extensions.rs +++ b/beacon-core/src/extensions.rs @@ -97,6 +97,9 @@ pub struct McpExtension { /// Human-readable description for the MCP tool/resource. #[serde(default, skip_serializing_if = "Option::is_none")] pub description: Option, + /// Human-readable title for the generated tool (MCP `Tool.title`). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub title: Option, /// Columns to expose. `None` (omitted) exposes all columns. #[serde(default, skip_serializing_if = "Option::is_none")] pub exposed_columns: Option>, @@ -199,6 +202,13 @@ impl TableExtensions { impl McpExtension { fn validate(&self, schema: &Schema) -> anyhow::Result<()> { + if let Some(name) = &self.tool_name { + anyhow::ensure!( + is_valid_tool_name(name), + "mcp tool_name '{name}' must be 1-64 characters of letters, digits, '_' or '-' \ + (MCP/Anthropic tool-name rules)" + ); + } if let Some(columns) = &self.exposed_columns { for column in columns { ensure_column(schema, column)?; @@ -208,6 +218,17 @@ impl McpExtension { } } +/// Whether `name` satisfies MCP/Anthropic tool-name rules: 1-64 characters of +/// `[A-Za-z0-9_-]`. A non-conforming name can make a client reject the entire +/// tool list, so it is rejected when the extension is set. +pub fn is_valid_tool_name(name: &str) -> bool { + !name.is_empty() + && name.len() <= 64 + && name + .chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-') +} + impl PresetExtension { fn validate(&self, schema: &Schema) -> anyhow::Result<()> { let mut seen = std::collections::HashSet::new(); @@ -453,6 +474,20 @@ mod tests { assert!(ext.validate(&schema()).unwrap_err().to_string().contains("duplicate preset")); } + #[test] + fn mcp_rejects_invalid_tool_name() { + let mut ext = TableExtensions::default(); + ext.set_kind("mcp", r#"{"enabled":true,"tool_name":"bad name!"}"#) + .unwrap(); + let err = ext.validate(&schema()).unwrap_err().to_string(); + assert!(err.contains("tool_name"), "unexpected: {err}"); + // A clean name passes. + let mut ok = TableExtensions::default(); + ok.set_kind("mcp", r#"{"enabled":true,"tool_name":"query_obs"}"#) + .unwrap(); + assert!(ok.validate(&schema()).is_ok()); + } + #[test] fn mcp_exposed_columns_must_exist() { let mut ext = TableExtensions::default(); diff --git a/beacon-mcp/src/catalog.rs b/beacon-mcp/src/catalog.rs index 9cc1dd43..17b9368c 100644 --- a/beacon-mcp/src/catalog.rs +++ b/beacon-mcp/src/catalog.rs @@ -10,7 +10,7 @@ use std::sync::Arc; use beacon_core::extensions::{McpExtension, PresetExtension, PresetFilter, PresetOp}; use beacon_core::runtime::Runtime; use beacon_core::AuthIdentity; -use rmcp::model::Tool; +use rmcp::model::{Tool, ToolAnnotations}; use serde_json::{json, Map, Value}; use crate::result::{run_sql_to_json, MAX_ROWS}; @@ -63,34 +63,40 @@ fn object_schema(props: Value, required: &[&str]) -> Map { schema } +/// Mark a tool as read-only via the MCP `Tool.annotations.readOnlyHint`, so +/// clients know it never mutates state. Every beacon MCP tool is read-only. +fn read_only(tool: Tool) -> Tool { + tool.with_annotations(ToolAnnotations::new().read_only(true)) +} + fn list_tables_tool() -> Tool { - Tool::new( + read_only(Tool::new( "list_tables", "List the tables registered in beacon, with their MCP exposure status.", object_schema(json!({}), &[]), - ) + )) } fn describe_table_tool() -> Tool { - Tool::new( + read_only(Tool::new( "describe_table", "Return a table's column schema and its attached extensions (MCP descriptor, presets).", object_schema( json!({ "table_name": { "type": "string", "description": "Name of the table." } }), &["table_name"], ), - ) + )) } fn run_sql_tool() -> Tool { - Tool::new( + read_only(Tool::new( "run_sql", "Run a read-only SQL query (SELECT only) against beacon and return JSON rows.", object_schema( json!({ "sql": { "type": "string", "description": "A read-only SELECT statement." } }), &["sql"], ), - ) + )) } async fn list_tables_json(runtime: &Arc) -> anyhow::Result { @@ -138,7 +144,15 @@ async fn describe_table_json( // ---- per-table tools ----------------------------------------------------- fn default_tool_name(table: &str) -> String { - format!("query_{table}") + // Sanitize to MCP-safe characters so a table name with dots/spaces still + // yields a valid tool name (see `beacon_core::extensions::is_valid_tool_name`). + let sanitized: String = table + .chars() + .map(|c| if c.is_ascii_alphanumeric() || c == '_' || c == '-' { c } else { '_' }) + .collect(); + let mut name = format!("query_{sanitized}"); + name.truncate(64); + name } async fn table_tool( @@ -196,7 +210,15 @@ async fn table_tool( json!({ "type": "integer", "description": "Maximum rows to return (default 100)." }), ); - Tool::new(name, description, object_schema(Value::Object(props), &[])) + let mut tool = read_only(Tool::new( + name, + description, + object_schema(Value::Object(props), &[]), + )); + if let Some(title) = mcp.title.clone() { + tool = tool.with_title(title); + } + tool } async fn run_table_tool( @@ -349,6 +371,7 @@ mod tests { enabled: true, tool_name: None, description: None, + title: None, exposed_columns: cols.map(|c| c.into_iter().map(String::from).collect()), } } diff --git a/docs/mcp.md b/docs/mcp.md index 7f6803d5..58ff0c29 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -32,6 +32,13 @@ SET EXTENSION 'preset' FOR obs TO '{"presets":[{"name":"shallow","filters":[{"co `query_obs` then appears as an MCP tool with a `preset: "shallow"` option. +The `mcp` descriptor maps to the MCP `Tool` standard: `tool_name` → `Tool.name` +(validated to MCP/Anthropic rules — 1–64 chars of `[A-Za-z0-9_-]`; the generated +default is sanitized), `title` → `Tool.title`, `description` → `Tool.description`, +`exposed_columns` constrain the generated `inputSchema`, and every tool carries +`annotations.readOnlyHint: true`. Payloads are parsed strictly +(`deny_unknown_fields`), so unknown keys are rejected rather than ignored. + ## Connecting Claude **HTTP (Claude Code / API):** point the client at `http://:/mcp`. From f110bb00fcf02af1d44b235e6531407115ebabbe Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 1 Jul 2026 09:56:24 +0200 Subject: [PATCH 6/8] MCP: document columns and table meaning in the mcp extension exposed_columns entries may now be either a bare name or {name, description}, so curators can describe what each column means. Descriptions are folded into the generated tool's 'select' help and returned by describe_table; the table's own meaning continues via 'description' -> Tool.description. New ExposedColumn/ ColumnDoc types (strict-parsed), validation by column name, tests + docs. --- beacon-core/src/api.rs | 3 +- beacon-core/src/extensions.rs | 84 ++++++++++++++++++++++++++++++++--- beacon-mcp/src/catalog.rs | 55 ++++++++++++++++------- docs/mcp.md | 20 +++++++-- 4 files changed, 138 insertions(+), 24 deletions(-) diff --git a/beacon-core/src/api.rs b/beacon-core/src/api.rs index 5d1fe3dd..ae44d625 100644 --- a/beacon-core/src/api.rs +++ b/beacon-core/src/api.rs @@ -14,7 +14,8 @@ use utoipa::ToSchema; /// Re-exported typed table-extension contracts (see [`crate::extensions`]). pub use crate::extensions::{ - McpExtension, Preset, PresetExtension, PresetFilter, PresetOp, TableExtensions, + ColumnDoc, ExposedColumn, McpExtension, Preset, PresetExtension, PresetFilter, PresetOp, + TableExtensions, }; /// A single parameter of a registered function. diff --git a/beacon-core/src/extensions.rs b/beacon-core/src/extensions.rs index 1d5223d8..f164456b 100644 --- a/beacon-core/src/extensions.rs +++ b/beacon-core/src/extensions.rs @@ -84,7 +84,7 @@ pub struct TableExtensions { "enabled": true, "tool_name": "query_ocean_observations", "description": "Argo float observations by location, depth, and time.", - "exposed_columns": ["lat", "lon", "depth", "temperature", "time"] + "exposed_columns": ["lat", "lon", {"name": "depth", "description": "measurement depth in meters"}] }))] #[serde(deny_unknown_fields)] pub struct McpExtension { @@ -94,15 +94,59 @@ pub struct McpExtension { /// 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. + /// Human-readable description of what the table contains / means (maps to the + /// MCP `Tool.description`, so the model knows what the table is). #[serde(default, skip_serializing_if = "Option::is_none")] pub description: Option, /// Human-readable title for the generated tool (MCP `Tool.title`). #[serde(default, skip_serializing_if = "Option::is_none")] pub title: Option, - /// Columns to expose. `None` (omitted) exposes all columns. + /// Columns to expose, optionally documented. `None` (omitted) exposes all + /// columns. Each entry is either a bare column name (`"lat"`) or an object + /// `{ "name": "depth", "description": "measurement depth in meters" }` + /// describing what the column means. #[serde(default, skip_serializing_if = "Option::is_none")] - pub exposed_columns: Option>, + pub exposed_columns: Option>, +} + +/// A column surfaced through the MCP tool — a bare name, or a name plus a +/// human-readable description of what it represents. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)] +#[serde(untagged)] +pub enum ExposedColumn { + /// Just the column name. + Name(String), + /// A column name together with a description of its meaning. + Documented(ColumnDoc), +} + +/// A documented column: its name and what it represents. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)] +#[serde(deny_unknown_fields)] +pub struct ColumnDoc { + /// Column name (must exist in the table schema). + pub name: String, + /// What the column means / represents. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub description: Option, +} + +impl ExposedColumn { + /// The column name. + pub fn name(&self) -> &str { + match self { + ExposedColumn::Name(name) => name, + ExposedColumn::Documented(doc) => &doc.name, + } + } + + /// The column's description, if documented. + pub fn description(&self) -> Option<&str> { + match self { + ExposedColumn::Name(_) => None, + ExposedColumn::Documented(doc) => doc.description.as_deref(), + } + } } /// A set of named, predefined filters consumers can apply. @@ -211,11 +255,18 @@ impl McpExtension { } if let Some(columns) = &self.exposed_columns { for column in columns { - ensure_column(schema, column)?; + ensure_column(schema, column.name())?; } } Ok(()) } + + /// The names of the curated exposed columns, if any are set. + pub fn exposed_column_names(&self) -> Option> { + self.exposed_columns + .as_ref() + .map(|cols| cols.iter().map(ExposedColumn::name).collect()) + } } /// Whether `name` satisfies MCP/Anthropic tool-name rules: 1-64 characters of @@ -496,6 +547,29 @@ mod tests { assert!(ext.validate(&schema()).unwrap_err().to_string().contains("does not exist")); } + #[test] + fn mcp_accepts_documented_columns() { + // Columns may be bare names or {name, description} objects, mixed freely. + let mut ext = TableExtensions::default(); + ext.set_kind( + "mcp", + r#"{"enabled":true,"exposed_columns":["lat",{"name":"depth","description":"meters"}]}"#, + ) + .unwrap(); + assert!(ext.validate(&schema()).is_ok()); + let cols = ext.mcp.as_ref().unwrap().exposed_columns.as_ref().unwrap(); + assert_eq!((cols[0].name(), cols[0].description()), ("lat", None)); + assert_eq!((cols[1].name(), cols[1].description()), ("depth", Some("meters"))); + // A documented column with an unknown key is rejected (deny_unknown_fields). + let mut bad = TableExtensions::default(); + assert!(bad + .set_kind( + "mcp", + r#"{"enabled":true,"exposed_columns":[{"name":"lat","unit":"deg"}]}"#, + ) + .is_err()); + } + #[test] fn unknown_kind_is_rejected() { let mut ext = TableExtensions::default(); diff --git a/beacon-mcp/src/catalog.rs b/beacon-mcp/src/catalog.rs index 17b9368c..a472d2da 100644 --- a/beacon-mcp/src/catalog.rs +++ b/beacon-mcp/src/catalog.rs @@ -170,28 +170,46 @@ async fn table_tool( .clone() .unwrap_or_else(|| format!("Query the '{table}' table.")); - // Columns offered to the model: the curated `exposed_columns` if set, - // otherwise the table's full schema. - let columns: Vec = match &mcp.exposed_columns { - Some(cols) => cols.clone(), + // Columns offered to the model: the curated `exposed_columns` (with optional + // per-column descriptions) if set, otherwise the table's full schema. + let columns: Vec<(String, Option)> = match &mcp.exposed_columns { + Some(cols) => cols + .iter() + .map(|c| (c.name().to_string(), c.description().map(String::from))) + .collect(), None => runtime .list_table_schema_view(table.to_string()) .await - .map(|s| s.fields.into_iter().map(|f| f.name).collect()) + .map(|s| s.fields.into_iter().map(|f| (f.name, None)).collect()) .unwrap_or_default(), }; + let column_names: Vec = columns.iter().map(|(name, _)| name.clone()).collect(); + // Fold any per-column descriptions into the `select` help so the model knows + // what each column means. + let glossary: Vec = columns + .iter() + .filter_map(|(name, desc)| desc.as_ref().map(|d| format!("{name}: {d}"))) + .collect(); + let select_description = if glossary.is_empty() { + "Columns to return. Omit for all exposed columns.".to_string() + } else { + format!( + "Columns to return. Omit for all exposed columns. Column meanings — {}.", + glossary.join("; ") + ) + }; let preset_names: Vec = preset .map(|p| p.presets.iter().map(|x| x.name.clone()).collect()) .unwrap_or_default(); let mut props = Map::new(); - if !columns.is_empty() { + if !column_names.is_empty() { props.insert( "select".into(), json!({ "type": "array", - "items": { "type": "string", "enum": columns }, - "description": "Columns to return. Omit for all exposed columns." + "items": { "type": "string", "enum": column_names }, + "description": select_description }), ); } @@ -256,7 +274,7 @@ fn build_table_sql( preset: Option<&PresetExtension>, args: &Map, ) -> anyhow::Result { - let exposed = mcp.exposed_columns.as_ref(); + let exposed = mcp.exposed_column_names(); let select = match args.get("select").and_then(Value::as_array) { Some(arr) if !arr.is_empty() => { @@ -264,14 +282,17 @@ fn build_table_sql( .iter() .filter_map(|v| v.as_str().map(String::from)) .collect(); - if let Some(exp) = exposed { + if let Some(exp) = &exposed { for col in &cols { - anyhow::ensure!(exp.contains(col), "column '{col}' is not exposed by this tool"); + anyhow::ensure!( + exp.contains(&col.as_str()), + "column '{col}' is not exposed by this tool" + ); } } cols.iter().map(|c| quote_ident(c)).collect::>().join(", ") } - _ => default_select(exposed), + _ => default_select(exposed.as_deref()), }; let mut clauses = Vec::new(); @@ -299,10 +320,10 @@ fn build_table_sql( Ok(sql) } -fn default_select(exposed: Option<&Vec>) -> String { +fn default_select(exposed: Option<&[&str]>) -> String { match exposed { Some(cols) if !cols.is_empty() => { - cols.iter().map(|c| quote_ident(c)).collect::>().join(", ") + cols.iter().map(|&c| quote_ident(c)).collect::>().join(", ") } _ => "*".to_string(), } @@ -372,7 +393,11 @@ mod tests { tool_name: None, description: None, title: None, - exposed_columns: cols.map(|c| c.into_iter().map(String::from).collect()), + exposed_columns: cols.map(|c| { + c.into_iter() + .map(|s| beacon_core::extensions::ExposedColumn::Name(s.to_string())) + .collect() + }), } } diff --git a/docs/mcp.md b/docs/mcp.md index 58ff0c29..0d613744 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -26,7 +26,18 @@ only `SELECT` runs. The caller's roles are preserved, so per-user read grants Use the table-extensions surface (SQL or REST): ```sql -SET EXTENSION 'mcp' FOR obs TO '{"enabled":true,"tool_name":"query_obs","description":"Ocean observations","exposed_columns":["lat","lon","depth","temperature"]}'; +SET EXTENSION 'mcp' FOR obs TO '{ + "enabled": true, + "tool_name": "query_obs", + "title": "Ocean observations", + "description": "Argo float profiles: temperature and salinity by location, depth and time.", + "exposed_columns": [ + {"name": "lat", "description": "latitude in decimal degrees"}, + {"name": "lon", "description": "longitude in decimal degrees"}, + {"name": "depth", "description": "measurement depth in meters"}, + "temperature" + ] +}'; SET EXTENSION 'preset' FOR obs TO '{"presets":[{"name":"shallow","filters":[{"column":"depth","op":"<=","value":10}]}]}'; ``` @@ -34,8 +45,11 @@ SET EXTENSION 'preset' FOR obs TO '{"presets":[{"name":"shallow","filters":[{"co The `mcp` descriptor maps to the MCP `Tool` standard: `tool_name` → `Tool.name` (validated to MCP/Anthropic rules — 1–64 chars of `[A-Za-z0-9_-]`; the generated -default is sanitized), `title` → `Tool.title`, `description` → `Tool.description`, -`exposed_columns` constrain the generated `inputSchema`, and every tool carries +default is sanitized), `title` → `Tool.title`, and `description` describes **what +the table means** → `Tool.description`. `exposed_columns` constrain the generated +`inputSchema`; each entry is either a bare name or `{"name", "description"}` — the +per-column meanings are folded into the `select` parameter help and returned by +`describe_table`, so the model knows what each field represents. Every tool carries `annotations.readOnlyHint: true`. Payloads are parsed strictly (`deny_unknown_fields`), so unknown keys are rejected rather than ignored. From 186e09c3b673c08835ef2a5f2df62536e6f03b6a Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 1 Jul 2026 10:15:57 +0200 Subject: [PATCH 7/8] MCP: merge schema types with column descriptions for the model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit describe_table and each per-table tool now present a merged per-column view — name + data_type + nullable + description — scoped to exposed_columns (in order) when set, or all columns otherwise. Descriptions come from the extension, with a fallback to the Arrow field's description/comment metadata. The select parameter help lists each column as 'name (type): meaning'. Adds resolve_columns() + test. --- beacon-mcp/src/catalog.rs | 147 ++++++++++++++++++++++++++++++++------ docs/mcp.md | 11 ++- 2 files changed, 133 insertions(+), 25 deletions(-) diff --git a/beacon-mcp/src/catalog.rs b/beacon-mcp/src/catalog.rs index a472d2da..377b5db3 100644 --- a/beacon-mcp/src/catalog.rs +++ b/beacon-mcp/src/catalog.rs @@ -8,6 +8,7 @@ use std::sync::Arc; use beacon_core::extensions::{McpExtension, PresetExtension, PresetFilter, PresetOp}; +use beacon_core::api::{SchemaFieldView, SchemaView}; use beacon_core::runtime::Runtime; use beacon_core::AuthIdentity; use rmcp::model::{Tool, ToolAnnotations}; @@ -131,16 +132,78 @@ async fn describe_table_json( .get_table_extensions(table.to_string()) .await .unwrap_or_default(); - let columns: Vec = schema - .fields + // Merge schema types with per-column descriptions, scoped to the exposed + // columns (or all columns when none are curated). + let columns: Vec = resolve_columns(&schema, ext.mcp.as_ref()) .iter() - .map(|f| json!({ "name": f.name, "data_type": f.data_type, "nullable": f.nullable })) + .map(|c| { + json!({ + "name": c.name, + "data_type": c.data_type, + "nullable": c.nullable, + "description": c.description, + }) + }) .collect(); Ok(serde_json::to_string_pretty( &json!({ "name": table, "columns": columns, "extensions": ext }), )?) } +/// A column resolved for the model: its schema type merged with any description. +struct ResolvedColumn { + name: String, + data_type: String, + nullable: bool, + description: Option, +} + +/// Merge the table schema with the mcp extension's per-column descriptions. +/// Scoped to `exposed_columns` (in that order) when set, otherwise every column. +/// A column's description comes from the extension entry, falling back to the +/// Arrow field's `description`/`comment` metadata when present. +fn resolve_columns(schema: &SchemaView, mcp: Option<&McpExtension>) -> Vec { + match mcp.and_then(|m| m.exposed_columns.as_ref()) { + Some(cols) => { + let by_name: std::collections::HashMap<&str, &SchemaFieldView> = + schema.fields.iter().map(|f| (f.name.as_str(), f)).collect(); + cols.iter() + .filter_map(|c| { + let field = by_name.get(c.name())?; + Some(ResolvedColumn { + name: c.name().to_string(), + data_type: field.data_type.clone(), + nullable: field.nullable, + description: c + .description() + .map(String::from) + .or_else(|| field_description(field)), + }) + }) + .collect() + } + None => schema + .fields + .iter() + .map(|f| ResolvedColumn { + name: f.name.clone(), + data_type: f.data_type.clone(), + nullable: f.nullable, + description: field_description(f), + }) + .collect(), + } +} + +/// A column description carried in the Arrow field metadata, if any. +fn field_description(field: &SchemaFieldView) -> Option { + field + .metadata + .get("description") + .or_else(|| field.metadata.get("comment")) + .cloned() +} + // ---- per-table tools ----------------------------------------------------- fn default_tool_name(table: &str) -> String { @@ -170,31 +233,26 @@ async fn table_tool( .clone() .unwrap_or_else(|| format!("Query the '{table}' table.")); - // Columns offered to the model: the curated `exposed_columns` (with optional - // per-column descriptions) if set, otherwise the table's full schema. - let columns: Vec<(String, Option)> = match &mcp.exposed_columns { - Some(cols) => cols - .iter() - .map(|c| (c.name().to_string(), c.description().map(String::from))) - .collect(), - None => runtime - .list_table_schema_view(table.to_string()) - .await - .map(|s| s.fields.into_iter().map(|f| (f.name, None)).collect()) - .unwrap_or_default(), + // Merge the table schema (types) with the extension's per-column descriptions, + // scoped to `exposed_columns` when set, or all columns otherwise, so the model + // sees name + data type + meaning for every queryable column. + let resolved = match runtime.list_table_schema_view(table.to_string()).await { + Some(schema) => resolve_columns(&schema, Some(mcp)), + None => Vec::new(), }; - let column_names: Vec = columns.iter().map(|(name, _)| name.clone()).collect(); - // Fold any per-column descriptions into the `select` help so the model knows - // what each column means. - let glossary: Vec = columns + let column_names: Vec = resolved.iter().map(|c| c.name.clone()).collect(); + let glossary: Vec = resolved .iter() - .filter_map(|(name, desc)| desc.as_ref().map(|d| format!("{name}: {d}"))) + .map(|c| match &c.description { + Some(desc) => format!("{} ({}): {}", c.name, c.data_type, desc), + None => format!("{} ({})", c.name, c.data_type), + }) .collect(); let select_description = if glossary.is_empty() { - "Columns to return. Omit for all exposed columns.".to_string() + "Columns to return. Omit for all columns.".to_string() } else { format!( - "Columns to return. Omit for all exposed columns. Column meanings — {}.", + "Columns to return, each shown as name (type): meaning. Omit for all. {}.", glossary.join("; ") ) }; @@ -387,6 +445,51 @@ mod tests { } } + fn field(name: &str, data_type: &str) -> SchemaFieldView { + SchemaFieldView { + name: name.into(), + data_type: data_type.into(), + nullable: true, + metadata: Default::default(), + } + } + + #[test] + fn resolve_columns_merges_types_and_descriptions() { + use beacon_core::extensions::{ColumnDoc, ExposedColumn}; + let schema = SchemaView { + fields: vec![field("lat", "Float64"), field("depth", "Float64"), field("x", "Int64")], + metadata: Default::default(), + }; + + // No exposed_columns -> all columns, types included, no descriptions. + let all = resolve_columns(&schema, Some(&mcp(None))); + assert_eq!(all.len(), 3); + assert_eq!((all[0].name.as_str(), all[0].data_type.as_str()), ("lat", "Float64")); + + // Exposed subset (in order), merging schema type + entry description. + let ext = McpExtension { + enabled: true, + tool_name: None, + title: None, + description: None, + exposed_columns: Some(vec![ + ExposedColumn::Documented(ColumnDoc { + name: "depth".into(), + description: Some("meters".into()), + }), + ExposedColumn::Name("lat".into()), + ]), + }; + let cols = resolve_columns(&schema, Some(&ext)); + assert_eq!(cols.len(), 2); + assert_eq!( + (cols[0].name.as_str(), cols[0].data_type.as_str(), cols[0].description.as_deref()), + ("depth", "Float64", Some("meters")) + ); + assert_eq!((cols[1].name.as_str(), cols[1].description.as_deref()), ("lat", None)); + } + fn mcp(cols: Option>) -> McpExtension { McpExtension { enabled: true, diff --git a/docs/mcp.md b/docs/mcp.md index 0d613744..0ed3b776 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -9,11 +9,16 @@ can discover tables and run read-only queries. Generated dynamically from the runtime: - `list_tables` — registered tables and their MCP exposure status. -- `describe_table` — a table's column schema plus its extensions (MCP descriptor, presets). +- `describe_table` — one merged view per column (`name`, `data_type`, `nullable`, + `description`), scoped to `exposed_columns` when set or all columns otherwise, + plus the raw extensions. Descriptions come from the extension, falling back to + the Arrow field's `description`/`comment` metadata. - `run_sql` — run a read-only `SELECT` and get JSON rows. - **one tool per table** whose `mcp` extension is enabled (see below). Its inputs are - derived from the extension: `select` (restricted to `exposed_columns`), `preset` - (an enum of the table's preset names, expanded to filters), and `limit`. + derived from the extension: `select` (restricted to `exposed_columns`, and its + help lists each column as `name (type): meaning` so the model knows types and + meanings), `preset` (an enum of the table's preset names, expanded to filters), + and `limit`. The MCP surface is **strictly read-only**: every tool call executes with `is_super_user` cleared, so the query planner rejects any DDL/DML (`CREATE`, From 98fa84bc8d9595468fcc49f9629cd4e86a0de46b Mon Sep 17 00:00:00 2001 From: Robin Kooyman Date: Wed, 1 Jul 2026 10:21:50 +0200 Subject: [PATCH 8/8] docs: expand MCP user guide + add how-it-works architecture doc - docs/mcp.md: config/env vars, agent authentication (create user, header examples for Claude Code / Desktop / SDKs), Tool-standard mapping table, quick check. - docs/mcp-architecture.md: request lifecycle, tool generation, identity flow + read-only enforcement, tool-call->SQL, result handling, extension points. --- docs/mcp-architecture.md | 143 ++++++++++++++++++++++++++++++++ docs/mcp.md | 174 ++++++++++++++++++++++++++++++--------- 2 files changed, 279 insertions(+), 38 deletions(-) create mode 100644 docs/mcp-architecture.md diff --git a/docs/mcp-architecture.md b/docs/mcp-architecture.md new file mode 100644 index 00000000..e1193c6d --- /dev/null +++ b/docs/mcp-architecture.md @@ -0,0 +1,143 @@ +# How the MCP server works + +This explains the internals of beacon's MCP server — the request path, how tools +are generated, how identity and read-only enforcement work, and how a tool call +becomes a query. For usage (enabling, exposing tables, connecting clients) see +[mcp.md](mcp.md). + +## Overview + +The MCP server is a thin **protocol adapter** in front of the existing +`Runtime`. It adds no query engine of its own: every tool call is translated into +a normal `Runtime::run_query`, so MCP inherits beacon's planner, catalog, +metrics, and RBAC. The tool set is not hard-coded — it is generated from the live +catalog and each table's `mcp`/`preset` extension, so curators shape the MCP +surface with SQL/REST, no code changes. + +``` +Claude ──MCP JSON-RPC over streamable HTTP──▶ beacon-api /mcp + │ resolve_identity (per request) + ▼ + beacon-mcp (rmcp ServerHandler) + list_tools / call_tool + │ + ▼ + Runtime ── get_table_extensions / list_table_schema_view + ── run_query(Query, identity[super-user cleared]) +``` + +## Crate layout (`beacon-mcp`) + +| File | Responsibility | +|---|---| +| `server.rs` | The rmcp `ServerHandler`: `get_info`, `list_tools`, `call_tool`, and identity extraction. | +| `catalog.rs` | Tool generation (`resolve_columns`, per-table tools), preset→SQL, argument→query mapping. | +| `result.rs` | Runs a query and serializes the rows to JSON. | +| `lib.rs` | `streamable_http_service(runtime)` — builds the tower service mounted at `/mcp`. | + +It depends on `rmcp` (the official Rust MCP SDK, pinned `=1.8.0`) for the protocol +and transport, and on `beacon-core` for `Runtime`, the extension types, and +`AuthIdentity`. + +## Transport & mounting + +`beacon_mcp::streamable_http_service(runtime)` returns an rmcp +`StreamableHttpService` (a `tower::Service`) wrapping a fresh `BeaconMcpServer` +per session. `beacon-api`'s `router.rs` mounts it: + +``` +route_service("/mcp", beacon_mcp::streamable_http_service(runtime)) + .layer(from_fn_with_state(runtime, resolve_identity)) // same as client API +``` + +behind a `BEACON_MCP_ENABLED` check. Streamable HTTP uses one endpoint for the +whole session (POST for requests, GET for the SSE stream, DELETE to end it); rmcp +tracks sessions with an in-memory `LocalSessionManager`. + +## Request lifecycle + +1. **`initialize`** — client and server negotiate; `get_info` advertises the + `tools` capability and server instructions. +2. **`tools/list`** — `list_tools` builds the current tool set (see below). It is + rebuilt per call, so newly-exposed tables appear without a restart. +3. **`tools/call`** — `call_tool` resolves the caller's identity, dispatches by + tool name, executes, and returns the result as MCP text content (`isError` set + on failure). + +## Tool generation (`list_tools`) + +Two groups are assembled: + +**Generic tools** — always present: `list_tables`, `describe_table`, `run_sql`. + +**Per-table tools** — for each table whose `mcp` extension has `enabled: true`, +one tool is generated from the extension metadata: + +- `name` ← `tool_name` (or a sanitized `query_
`), `title`, `description`. +- `inputSchema` is built from `resolve_columns` + presets: + - `select` — an array whose `enum` is the exposed column names; its description + lists each column as `name (type): meaning`. + - `preset` — an `enum` of the table's preset names (if any). + - `limit` — integer, default 100. +- `annotations.readOnlyHint = true`. + +`resolve_columns` merges the table's Arrow schema (types) with the extension's +per-column descriptions, scoped to `exposed_columns` (in order) when set or all +columns otherwise, with a fallback to the Arrow field's `description`/`comment` +metadata. The same function feeds `describe_table`, so the tool schema and the +description tool agree. + +## Executing a call (`call_tool`) + +``` +call_tool(request, context): + identity = identity_from_context(context) # super-user cleared + dispatch(runtime, request.name, request.arguments, identity): + "list_tables" -> catalog listing (JSON) + "describe_table" -> merged columns + extensions (JSON) + "run_sql" -> run_query(SELECT, identity) -> JSON rows +
-> build_table_sql(...) -> run_query(..., identity) -> JSON rows +``` + +**Per-table tools → SQL.** `build_table_sql` turns the validated arguments into a +`SELECT`: `select` (checked against `exposed_columns`, identifiers quoted), the +chosen `preset` expanded to a `WHERE` from its stored filters, and `limit`. +Preset operators are the typed `PresetOp` enum; values are rendered safely +(scalars/arrays escaped) — the model never supplies raw SQL through these tools. + +## Identity & read-only enforcement + +The `resolve_identity` middleware authenticates each HTTP request (Basic/Bearer → +a user's roles, or the anonymous principal, or an empty identity) and inserts the +`AuthIdentity` into the request extensions. The streamable-HTTP transport injects +the request `http::request::Parts` into the MCP `RequestContext`, so +`identity_from_context` recovers that `AuthIdentity`. + +Before use it **clears `is_super_user`**: + +```rust +let mut identity = /* from request parts, or AuthIdentity::empty() */; +identity.is_super_user = false; // MCP is read-only, always +``` + +This is defense-in-depth: the query planner gates DDL/DML on super-user, so a +non-super identity can only run `SELECT`. The caller's `roles` are preserved, so +when `BEACON_AUTH_ENFORCE=true` per-user read grants still apply. (Beacon's +super-user is config-only and never a client identity, so this only matters if +that ever changes.) + +## Results + +Rows are collected from the query stream, capped at 1000, and serialized to JSON +(`result.rs`) as the tool's text content. Errors are returned as MCP tool errors +(`isError: true`) with the message, rather than failing the JSON-RPC call, so the +model can read and react to them. + +## Extending it + +- **New generic tool** — add a builder in `catalog.rs`, list it in `list_tools`, + and add a `dispatch` arm. +- **Richer per-table inputs** (e.g. server-side filters, sort) — extend the + per-table `inputSchema` and `build_table_sql`. +- **New extension-driven behavior** — add fields to the `mcp` extension in + `beacon-core`'s `extensions` module; they flow here via `get_table_extensions`. diff --git a/docs/mcp.md b/docs/mcp.md index 0ed3b776..9a862999 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -1,34 +1,53 @@ # MCP server -Beacon ships an [MCP](https://modelcontextprotocol.io) server over the -streamable-HTTP transport at `POST/GET/DELETE /mcp`, so MCP clients (e.g. Claude) -can discover tables and run read-only queries. +Beacon ships an [MCP](https://modelcontextprotocol.io) server so MCP clients +(e.g. Claude) can discover beacon's tables and run **read-only** queries against +them. It is served over the streamable-HTTP transport at `POST/GET/DELETE /mcp` +by the `beacon-mcp` crate, mounted alongside the REST API. + +For how it works internally, see [mcp-architecture.md](mcp-architecture.md). + +## Enabling & configuration + +The endpoint is mounted by default. Relevant environment variables: + +| Variable | Default | Effect | +|---|---|---| +| `BEACON_MCP_ENABLED` | `true` | Mount `/mcp`. Set `false`/`0`/`off` to disable. | +| `BEACON_AUTH_ANONYMOUS_ENABLED` | `true` | Unauthenticated requests resolve to the anonymous principal. | +| `BEACON_AUTH_ENFORCE` | `false` | Apply per-role read grants at query time. | + +With the defaults, `/mcp` is on and open (anonymous, read-only). To lock it down, +set `BEACON_AUTH_ENFORCE=true` and `BEACON_AUTH_ANONYMOUS_ENABLED=false`, then +issue each agent a credential (below). ## Tools -Generated dynamically from the runtime: +The tool set is generated dynamically from the runtime on every `tools/list`: -- `list_tables` — registered tables and their MCP exposure status. -- `describe_table` — one merged view per column (`name`, `data_type`, `nullable`, - `description`), scoped to `exposed_columns` when set or all columns otherwise, - plus the raw extensions. Descriptions come from the extension, falling back to - the Arrow field's `description`/`comment` metadata. -- `run_sql` — run a read-only `SELECT` and get JSON rows. -- **one tool per table** whose `mcp` extension is enabled (see below). Its inputs are - derived from the extension: `select` (restricted to `exposed_columns`, and its - help lists each column as `name (type): meaning` so the model knows types and - meanings), `preset` (an enum of the table's preset names, expanded to filters), - and `limit`. +- **`list_tables`** — registered tables and their MCP exposure status. +- **`describe_table`** — a merged per-column view (`name`, `data_type`, + `nullable`, `description`), scoped to `exposed_columns` when set or all columns + otherwise, plus the raw extensions. Descriptions come from the extension, + falling back to the Arrow field's `description`/`comment` metadata. +- **`run_sql`** — run a read-only `SELECT` and get JSON rows. +- **one tool per table** whose `mcp` extension is enabled. Inputs are derived from + the extension: `select` (restricted to `exposed_columns`; its help lists each + column as `name (type): meaning`), `preset` (an enum of the table's preset + names, expanded to filters at query time), and `limit`. The MCP surface is **strictly read-only**: every tool call executes with `is_super_user` cleared, so the query planner rejects any DDL/DML (`CREATE`, `INSERT`, `UPDATE`, `DELETE`, `SET EXTENSION`, …) regardless of who connects — only `SELECT` runs. The caller's roles are preserved, so per-user read grants -(RBAC) still apply. Results are capped (1000 rows) to keep tool output bounded. +(RBAC) still apply. Every tool carries `annotations.readOnlyHint: true`. Results +are capped (1000 rows) to keep tool output bounded. ## Exposing a table to MCP -Use the table-extensions surface (SQL or REST): +Tables become MCP tools via the table-extensions surface (SQL or the admin REST +API). The `mcp` extension describes the table; an optional `preset` extension +adds named, predefined filter sets. ```sql SET EXTENSION 'mcp' FOR obs TO '{ @@ -37,38 +56,117 @@ SET EXTENSION 'mcp' FOR obs TO '{ "title": "Ocean observations", "description": "Argo float profiles: temperature and salinity by location, depth and time.", "exposed_columns": [ - {"name": "lat", "description": "latitude in decimal degrees"}, - {"name": "lon", "description": "longitude in decimal degrees"}, + {"name": "lat", "description": "latitude in decimal degrees"}, + {"name": "lon", "description": "longitude in decimal degrees"}, {"name": "depth", "description": "measurement depth in meters"}, "temperature" ] }'; -SET EXTENSION 'preset' FOR obs TO '{"presets":[{"name":"shallow","filters":[{"column":"depth","op":"<=","value":10}]}]}'; + +SET EXTENSION 'preset' FOR obs TO '{ + "presets": [ + {"name": "shallow", "description": "Surface layer", + "filters": [{"column": "depth", "op": "<=", "value": 10}]} + ] +}'; ``` -`query_obs` then appears as an MCP tool with a `preset: "shallow"` option. +`query_obs` then appears as an MCP tool with a `preset: "shallow"` option. Read +back or remove with `SHOW EXTENSIONS FOR obs` / `DROP EXTENSION 'mcp' FOR obs`. -The `mcp` descriptor maps to the MCP `Tool` standard: `tool_name` → `Tool.name` -(validated to MCP/Anthropic rules — 1–64 chars of `[A-Za-z0-9_-]`; the generated -default is sanitized), `title` → `Tool.title`, and `description` describes **what -the table means** → `Tool.description`. `exposed_columns` constrain the generated -`inputSchema`; each entry is either a bare name or `{"name", "description"}` — the -per-column meanings are folded into the `select` parameter help and returned by -`describe_table`, so the model knows what each field represents. Every tool carries -`annotations.readOnlyHint: true`. Payloads are parsed strictly -(`deny_unknown_fields`), so unknown keys are rejected rather than ignored. +### How the `mcp` fields map to the MCP `Tool` standard -## Connecting Claude +| Extension field | MCP `Tool` | Notes | +|---|---|---| +| `tool_name` | `name` | Validated to 1–64 chars of `[A-Za-z0-9_-]`; the generated default (`query_
`) is sanitized. | +| `title` | `title` | Human-readable label. | +| `description` | `description` | What the **table** means. | +| `exposed_columns` | `inputSchema` | Constrains `select`; per-column meanings feed its help + `describe_table`. | +| — | `annotations.readOnlyHint` | Always `true`. | -**HTTP (Claude Code / API):** point the client at `http://:/mcp`. +`exposed_columns` entries are either a bare name (`"lat"`) or +`{"name": ..., "description": ...}`. Payloads parse strictly +(`deny_unknown_fields`, typed operators), so typos/extra keys are rejected rather +than silently dropped. See the table-extensions docs for the full schema. + +## Authenticating an agent + +`/mcp` authenticates via the HTTP `Authorization` header (the same +`resolve_identity` path as the client API): + +- **Basic** — `Authorization: Basic base64(user:pass)` → a beacon user's roles. +- **Bearer** — `Authorization: Bearer ` → an OIDC/OAuth2 JWT. +- **No header** → the anonymous principal (if enabled), else no access. + +MCP is read-only regardless of identity; the identity only decides *which reads* +are allowed (when `BEACON_AUTH_ENFORCE=true`). + +Create a read-only user to hand to an agent (as a super-user, via SQL or the +admin API): + +```sql +CREATE USER agent WITH PASSWORD 's3cret'; +-- when enforcing, grant reads and assign a role: +GRANT SELECT ON obs TO ROLE readers; +GRANT ROLE readers TO USER agent; +``` -**Claude Desktop** (via an HTTP-capable MCP entry): +> Beacon's built-in super-user is **config-only** (`BEACON_ADMIN_*`) and is not a +> client identity, so those admin credentials do **not** authenticate on `/mcp`. +> Use a `CREATE USER` account or an OIDC token. + +## Connecting a client + +**Claude Code (CLI)** — streamable HTTP with an auth header: + +```bash +claude mcp add --transport http beacon https://your-host/mcp \ + --header "Authorization: Basic $(printf 'agent:s3cret' | base64)" +# or: --header "Authorization: Bearer " +``` + +**Claude Desktop** — for a static token, bridge with `mcp-remote`: ```json -{ "mcpServers": { "beacon": { "url": "http://localhost:5001/mcp" } } } +{ + "mcpServers": { + "beacon": { + "command": "npx", + "args": ["mcp-remote", "https://your-host/mcp", + "--header", "Authorization: Bearer "] + } + } +} +``` + +(For an open/anonymous local instance you can point a client straight at the URL +with no header: `{ "mcpServers": { "beacon": { "url": "http://localhost:5001/mcp" } } }`.) + +**Programmatic (MCP SDKs)** — set the header on the streamable-HTTP transport: + +```ts +new StreamableHTTPClientTransport(new URL("https://your-host/mcp"), { + requestInit: { headers: { Authorization: "Bearer " } }, +}); +``` + +```python +streamablehttp_client("https://your-host/mcp", + headers={"Authorization": "Bearer "}) +``` + +The transport attaches the header to every request, which is what beacon needs — +it authenticates per request, even within a long-lived MCP session. + +## Quick check + +```bash +curl -s -X POST http://127.0.0.1:5001/mcp \ + -H "Content-Type: application/json" \ + -H "Accept: application/json, text/event-stream" \ + -H "Authorization: Basic $(printf 'agent:s3cret' | base64)" \ + -d '{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2024-11-05","capabilities":{},"clientInfo":{"name":"c","version":"0"}}}' ``` -> The endpoint rides the same `resolve_identity` middleware as the client API: -> requests authenticate via `Authorization` (resolving to that user's roles), or -> the anonymous principal when enabled, or a role-less identity otherwise. It is -> read-only regardless. Gate it with `BEACON_MCP_ENABLED=false` to disable. +A `200` with an `initialize` result means the credential was accepted; `401` +means it was rejected.