diff --git a/Cargo.lock b/Cargo.lock index a61d18ea..f45a0405 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1289,6 +1289,7 @@ dependencies = [ "base64", "beacon-config", "beacon-core", + "beacon-mcp", "bytes", "futures", "pprof", @@ -1858,6 +1859,23 @@ dependencies = [ "url", ] +[[package]] +name = "beacon-mcp" +version = "0.1.0" +dependencies = [ + "anyhow", + "arrow 58.3.0", + "arrow-json 58.3.0", + "beacon-core", + "futures", + "http", + "rmcp", + "serde", + "serde_json", + "tokio", + "tracing", +] + [[package]] name = "beacon-nd-array" version = "0.1.0" @@ -4034,6 +4052,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" @@ -7439,6 +7463,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 = "path_abs" version = "0.5.1" @@ -8231,6 +8261,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" @@ -8466,6 +8516,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" @@ -8784,6 +8863,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 = "scoped-tls" version = "1.0.1" @@ -8914,6 +9019,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" @@ -9221,6 +9337,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 c4151c02..06dadfa8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,7 +2,7 @@ # Edition-2024 members imply resolver 3; set it explicitly so features aren't # over-unified across normal/build/dev deps (fewer features compiled). resolver = "3" -members = ["beacon-api", "beacon-auth", "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-file-formats/beacon-iceberg", "beacon-file-formats/beacon-lance", "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-auth", "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-file-formats/beacon-iceberg", "beacon-file-formats/beacon-lance", "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 40a4455a..fa6c58e5 100644 --- a/beacon-api/Cargo.toml +++ b/beacon-api/Cargo.toml @@ -45,6 +45,7 @@ prost = {workspace = true} # Local dependencies beacon-config = { path = "../beacon-config" } beacon-core = { path = "../beacon-core" } +beacon-mcp = { path = "../beacon-mcp" } [dev-dependencies] # Enable beacon-core's test-only helper (ephemeral in-memory auth runtime) for the @@ -53,4 +54,4 @@ beacon-core = { path = "../beacon-core", features = ["test-util"] } # `tower::ServiceExt::oneshot` for driving the axum router in the HTTP auth tests. tower = { version = "0.5", features = ["util"] } # Temp storage roots for the admin dataset file-management HTTP tests. -tempfile = { workspace = true } \ No newline at end of file +tempfile = { workspace = true } 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 f02660c1..dc9f9d97 100644 --- a/beacon-api/src/axum/admin/mod.rs +++ b/beacon-api/src/axum/admin/mod.rs @@ -14,6 +14,7 @@ mod auth; mod check; mod crawlers; mod datasets; +mod extensions; mod external_tables; mod tables; @@ -51,6 +52,10 @@ pub(crate) fn setup_admin_router() -> (Router>, utoipa::openapi::Op .routes(routes!(tables::list_table_config)) .routes(routes!(auth::list_users)) .routes(routes!(auth::list_roles)) + .routes(routes!( + extensions::set_table_extensions, + extensions::delete_table_extensions + )) .split_for_parts(); (admin_router, admin_api) diff --git a/beacon-api/src/axum/client/mod.rs b/beacon-api/src/axum/client/mod.rs index 53eb8db9..b6e74535 100644 --- a/beacon-api/src/axum/client/mod.rs +++ b/beacon-api/src/axum/client/mod.rs @@ -45,6 +45,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::default_table_schema)) .routes(routes!(functions::list_functions)) .routes(routes!(functions::list_table_functions)) diff --git a/beacon-api/src/axum/client/tables.rs b/beacon-api/src/axum/client/tables.rs index 715b2c40..ad972a79 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}; +use beacon_core::api::{SchemaFieldView, SchemaView, TableExtensions}; use beacon_core::runtime::Runtime; use utoipa::{IntoParams, ToSchema}; @@ -110,6 +110,48 @@ pub(crate) async fn list_table_schema( } } +/// 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-api/src/axum/router.rs b/beacon-api/src/axum/router.rs index 5fc2055f..901755b0 100644 --- a/beacon-api/src/axum/router.rs +++ b/beacon-api/src/axum/router.rs @@ -72,8 +72,27 @@ pub(crate) fn setup_router( // lands on the UI instead of the API docs. let web_ui = web_ui_router(&config.server.web_ui_dir); - let mut router = client_router - .merge(admin_router) + // MCP streamable-HTTP endpoint, gated by BEACON_MCP_ENABLED (default on). It + // rides the same `resolve_identity` middleware as the client API, so MCP tool + // calls execute under the caller's identity (or the anonymous principal when + // enabled) and per-user RBAC applies at query time. + let mcp_enabled = std::env::var("BEACON_MCP_ENABLED") + .map(|v| !matches!(v.trim().to_ascii_lowercase().as_str(), "false" | "0" | "off")) + .unwrap_or(true); + let mut router = client_router.merge(admin_router); + if mcp_enabled { + let mcp_router = ::axum::Router::new() + .route_service( + "/mcp", + beacon_mcp::streamable_http_service(beacon_runtime.clone()), + ) + .layer(::axum::middleware::from_fn_with_state( + beacon_runtime.clone(), + crate::axum::auth::resolve_identity, + )); + router = router.merge(mcp_router); + } + let mut router = router .merge(Scalar::with_url("/scalar/", docs.clone())) .route( "/scalar", diff --git a/beacon-core/src/api.rs b/beacon-core/src/api.rs index 2b3fbd93..ae44d625 100644 --- a/beacon-core/src/api.rs +++ b/beacon-core/src/api.rs @@ -12,6 +12,12 @@ 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::{ + ColumnDoc, ExposedColumn, McpExtension, Preset, PresetExtension, PresetFilter, PresetOp, + 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..f164456b --- /dev/null +++ b/beacon-core/src/extensions.rs @@ -0,0 +1,590 @@ +//! 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. +/// 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")] + 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", {"name": "depth", "description": "measurement depth in meters"}] +}))] +#[serde(deny_unknown_fields)] +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 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, 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>, +} + +/// 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. +#[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 } + ] + }] +}))] +#[serde(deny_unknown_fields)] +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)] +#[serde(deny_unknown_fields)] +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)] +#[serde(deny_unknown_fields)] +pub struct PresetFilter { + /// Column the filter applies to (must exist in the table schema). + pub column: 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, +} + +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(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.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 +/// `[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(); + 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)?; + 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 { + PresetOp::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 + ); + } + PresetOp::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_at_parse() { + // An unknown operator no longer parses into the typed `op` enum at all. + let mut ext = TableExtensions::default(); + 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] + 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_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(); + 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 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(); + 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 772fb82c..b6f6b1a8 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 dataset_files; pub mod dataset_uploads; pub mod metrics; diff --git a/beacon-core/src/parser/beacon_parser.rs b/beacon-core/src/parser/beacon_parser.rs index 7da32f1a..4437a179 100644 --- a/beacon-core/src/parser/beacon_parser.rs +++ b/beacon-core/src/parser/beacon_parser.rs @@ -11,8 +11,9 @@ use beacon_auth::{Privilege, PrivilegeTarget}; use super::statement::{ AuthStatement, BeaconStatement, CreateCrawlerStatement, CreateIndexStatement, - CreateMaterializedViewStatement, DropCrawlerStatement, DropIndexStatement, RefreshStatement, - RunCrawlerStatement, ShowIndexesStatement, + CreateMaterializedViewStatement, DropCrawlerStatement, DropExtensionStatement, + DropIndexStatement, RefreshStatement, RunCrawlerStatement, SetExtensionStatement, + ShowExtensionsStatement, ShowIndexesStatement, }; /// A parser that extends `DFParser` with custom Beacon SQL syntax. @@ -57,6 +58,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(); + } + if self.is_create_index() { return self.parse_create_index(); } @@ -169,6 +182,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))) + } + fn is_create_index(&self) -> bool { let t1 = &self.df_parser.parser.peek_nth_token(0).token; let t2 = &self.df_parser.parser.peek_nth_token(1).token; @@ -863,6 +956,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 813ec710..486511c5 100644 --- a/beacon-core/src/parser/statement.rs +++ b/beacon-core/src/parser/statement.rs @@ -14,11 +14,77 @@ pub enum BeaconStatement { RunCrawler(RunCrawlerStatement), DropCrawler(DropCrawlerStatement), ShowCrawlers, + SetExtension(SetExtensionStatement), + DropExtension(DropExtensionStatement), + ShowExtensions(ShowExtensionsStatement), CreateIndex(CreateIndexStatement), DropIndex(DropIndexStatement), ShowIndexes(ShowIndexesStatement), } +/// 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('\'', "''") +} + /// Authentication and authorization management statements (users, roles, grants, denies). #[derive(Debug, Clone)] pub enum AuthStatement { @@ -236,6 +302,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}"), Self::CreateIndex(s) => write!(f, "{s}"), Self::DropIndex(s) => write!(f, "{s}"), Self::ShowIndexes(s) => write!(f, "{s}"), diff --git a/beacon-core/src/runtime.rs b/beacon-core/src/runtime.rs index e3598f9f..e5d9524a 100644 --- a/beacon-core/src/runtime.rs +++ b/beacon-core/src/runtime.rs @@ -588,6 +588,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::CreateIndex(statement) => { Ok(crate::statement_plan::create_index_plan(statement)) } @@ -820,6 +829,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) @@ -1356,6 +1389,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}")), + beacon_auth::AuthIdentity::system(), + ) + .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()), beacon_auth::AuthIdentity::system()) + .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 cc374489..ebe17592 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. @@ -588,6 +590,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 2da7baa5..e5f554f3 100644 --- a/beacon-core/src/statement_plan/mod.rs +++ b/beacon-core/src/statement_plan/mod.rs @@ -34,8 +34,8 @@ use datafusion::{ use crate::parser::statement::{ AuthStatement, CreateCrawlerStatement, CreateIndexStatement, CreateMaterializedViewStatement, - DropCrawlerStatement, DropIndexStatement, RefreshStatement, RunCrawlerStatement, - ShowIndexesStatement, + DropCrawlerStatement, DropExtensionStatement, DropIndexStatement, RefreshStatement, + RunCrawlerStatement, SetExtensionStatement, ShowExtensionsStatement, ShowIndexesStatement, }; pub(crate) use authz::authorize_logical_plan; @@ -189,6 +189,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())), + }) +} + /// Build the logical plan for `CREATE INDEX [] ON
() [USING ]`. pub(crate) fn create_index_plan(statement: CreateIndexStatement) -> LogicalPlan { LogicalPlan::Extension(Extension { diff --git a/beacon-core/src/statement_plan/physical.rs b/beacon-core/src/statement_plan/physical.rs index 63a08bac..ede31dbe 100644 --- a/beacon-core/src/statement_plan/physical.rs +++ b/beacon-core/src/statement_plan/physical.rs @@ -33,6 +33,10 @@ use super::{ }, 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 { @@ -1035,3 +1039,147 @@ impl ExecutionPlan for ShowIndexesExec { 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 54e6dca3..a9074f87 100644 --- a/beacon-core/src/statement_plan/query_planner.rs +++ b/beacon-core/src/statement_plan/query_planner.rs @@ -236,6 +236,30 @@ impl ExtensionPlanner for BeaconExtensionPlanner { )))); } + 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 ffdc30f6..e538046b 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 @@ -259,4 +345,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()); + } } diff --git a/beacon-mcp/Cargo.toml b/beacon-mcp/Cargo.toml new file mode 100644 index 00000000..e682a2b9 --- /dev/null +++ b/beacon-mcp/Cargo.toml @@ -0,0 +1,17 @@ +[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" +http = "1" diff --git a/beacon-mcp/src/catalog.rs b/beacon-mcp/src/catalog.rs new file mode 100644 index 00000000..377b5db3 --- /dev/null +++ b/beacon-mcp/src/catalog.rs @@ -0,0 +1,557 @@ +//! 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, PresetOp}; +use beacon_core::api::{SchemaFieldView, SchemaView}; +use beacon_core::runtime::Runtime; +use beacon_core::AuthIdentity; +use rmcp::model::{Tool, ToolAnnotations}; +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, + identity: AuthIdentity, +) -> 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(), identity).await + } + other => run_table_tool(runtime, other, &args, identity).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 +} + +/// 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 { + 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 { + 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 { + 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 { + 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(); + // 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(|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 { + // 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( + 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.")); + + // 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 = resolved.iter().map(|c| c.name.clone()).collect(); + let glossary: Vec = resolved + .iter() + .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 columns.".to_string() + } else { + format!( + "Columns to return, each shown as name (type): meaning. Omit for all. {}.", + 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 !column_names.is_empty() { + props.insert( + "select".into(), + json!({ + "type": "array", + "items": { "type": "string", "enum": column_names }, + "description": select_description + }), + ); + } + 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)." }), + ); + + 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( + runtime: &Arc, + tool_name: &str, + args: &Map, + identity: AuthIdentity, +) -> 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, identity).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_column_names(); + + 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.as_str()), + "column '{col}' is not exposed by this tool" + ); + } + } + cols.iter().map(|c| quote_ident(c)).collect::>().join(", ") + } + _ => default_select(exposed.as_deref()), + }; + + 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<&[&str]>) -> 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 { + PresetOp::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])? + )) + } + PresetOp::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(", "))) + } + op => Ok(format!("{col} {} {}", op.as_sql(), render_scalar(&filter.value)?)), + } +} + +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 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, + tool_name: None, + description: None, + title: None, + exposed_columns: cols.map(|c| { + c.into_iter() + .map(|s| beacon_core::extensions::ExposedColumn::Name(s.to_string())) + .collect() + }), + } + } + + #[test] + fn builds_select_with_preset_between() { + let p = preset( + "shallow", + vec![PresetFilter { + column: "depth".into(), + op: PresetOp::Between, + 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: PresetOp::In, + 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..31deab7e --- /dev/null +++ b/beacon-mcp/src/result.rs @@ -0,0 +1,65 @@ +//! 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 beacon_core::AuthIdentity; +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, + identity: AuthIdentity, +) -> anyhow::Result { + let result = runtime.run_query(Query::sql(sql), identity).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..e48831e0 --- /dev/null +++ b/beacon-mcp/src/server.rs @@ -0,0 +1,82 @@ +//! The [`ServerHandler`] implementation: capabilities, tool listing, dispatch. + +use std::sync::Arc; + +use beacon_core::runtime::Runtime; +use beacon_core::AuthIdentity; +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(); + let identity = identity_from_context(&context); + match crate::catalog::dispatch(&self.runtime, request.name.as_ref(), args, identity).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())])), + } + } +} + +/// Recover the caller's [`AuthIdentity`], resolved by the `resolve_identity` +/// 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 { + let mut identity = context + .extensions + .get::() + .and_then(|parts| parts.extensions.get::().cloned()) + .unwrap_or_else(AuthIdentity::empty); + identity.is_super_user = false; + identity +} 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 new file mode 100644 index 00000000..9a862999 --- /dev/null +++ b/docs/mcp.md @@ -0,0 +1,172 @@ +# MCP server + +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 + +The tool set is generated dynamically from the runtime on every `tools/list`: + +- **`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. Every tool carries `annotations.readOnlyHint: true`. Results +are capped (1000 rows) to keep tool output bounded. + +## Exposing a table to MCP + +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 '{ + "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", "description": "Surface layer", + "filters": [{"column": "depth", "op": "<=", "value": 10}]} + ] +}'; +``` + +`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`. + +### How the `mcp` fields map to the MCP `Tool` standard + +| 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`. | + +`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; +``` + +> 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": { + "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"}}}' +``` + +A `200` with an `initialize` result means the credential was accepted; `401` +means it was rejected.