diff --git a/Cargo.lock b/Cargo.lock index b445be24..390f9340 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -301,7 +301,7 @@ version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f633dbfdf39c039ada1bf9e34c694816eb71fbb7dc78f613993b7245e078a1ed" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", ] [[package]] @@ -501,9 +501,9 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" [[package]] name = "bitflags" -version = "2.9.1" +version = "2.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b8e56985ec62d17e9c1001dc89c88ecd7dc08e47eba5ec7c29c7b5eeecde967" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" [[package]] name = "bitvec" @@ -526,6 +526,56 @@ dependencies = [ "generic-array", ] +[[package]] +name = "bollard" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97ccca1260af6a459d75994ad5acc1651bcabcbdbc41467cc9786519ab854c30" +dependencies = [ + "base64", + "bollard-stubs", + "bytes", + "futures-core", + "futures-util", + "hex", + "home", + "http", + "http-body-util", + "hyper", + "hyper-named-pipe", + "hyper-rustls", + "hyper-util", + "hyperlocal", + "log", + "pin-project-lite", + "rustls", + "rustls-native-certs", + "rustls-pemfile", + "rustls-pki-types", + "serde", + "serde_derive", + "serde_json", + "serde_repr", + "serde_urlencoded", + "thiserror 2.0.12", + "tokio", + "tokio-util", + "tower-service", + "url", + "winapi", +] + +[[package]] +name = "bollard-stubs" +version = "1.47.1-rc.27.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f179cfbddb6e77a5472703d4b30436bff32929c0aa8a9008ecf23d1d3cdd0da" +dependencies = [ + "serde", + "serde_repr", + "serde_with", +] + [[package]] name = "borsh" version = "1.7.0" @@ -664,7 +714,7 @@ dependencies = [ "num-traits", "serde", "wasm-bindgen", - "windows-link", + "windows-link 0.1.3", ] [[package]] @@ -765,7 +815,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] @@ -941,7 +991,7 @@ version = "0.28.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "829d955a0bb380ef178a640b91779e3987da38c9aea133b20614cfed8cdea9c6" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "crossterm_winapi", "parking_lot", "rustix 0.38.44", @@ -954,7 +1004,7 @@ version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "crossterm_winapi", "derive_more 2.1.1", "document-features", @@ -1310,6 +1360,17 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fea41bba32d969b513997752735605054bc0dfa92b4c56bf1189f2e174be7a10" +[[package]] +name = "docker_credential" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29547a1dc60885a552306986316bc9701ba120c1a8db6769fa68691529ad373d" +dependencies = [ + "base64", + "serde", + "serde_json", +] + [[package]] name = "document-features" version = "0.2.12" @@ -1393,6 +1454,17 @@ dependencies = [ "str-buf", ] +[[package]] +name = "etcetera" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26c7b13d0780cb82722fd59f6f57f925e143427e4a75313a6c77243bf5326ae6" +dependencies = [ + "cfg-if", + "home", + "windows-sys 0.59.0", +] + [[package]] name = "eventsource-stream" version = "0.2.3" @@ -1717,7 +1789,7 @@ version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0bf760ebf69878d9fd8f110c89703d90ce35095324d1f1edcb595c63945ee757" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "ignore", "walkdir", ] @@ -1785,6 +1857,15 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "http" version = "1.3.1" @@ -1857,6 +1938,20 @@ dependencies = [ "want", ] +[[package]] +name = "hyper-named-pipe" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fab3637d6b04a8037af8a266fdf6cf92ea957e8c53981a2bf6136572531025bf" +dependencies = [ + "hex", + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-rustls" version = "0.27.7" @@ -1898,6 +1993,21 @@ dependencies = [ "tracing", ] +[[package]] +name = "hyperlocal" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "986c5ce3b994526b3cd75578e62554abd09f0899d6206de48b3e96ab34ccc8c7" +dependencies = [ + "hex", + "http-body-util", + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "iana-time-zone" version = "0.1.63" @@ -2107,7 +2217,7 @@ version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6654738b8024300cf062d04a1c13c10c8e2cea598ec1c47dc9b6641159429756" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "crossterm 0.29.0", "dyn-clone", "fuzzy-matcher", @@ -2130,7 +2240,7 @@ version = "0.7.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d93587f37623a1a17d94ef2bc9ada592f5465fe7732084ab7beefabe5c77c0c4" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "cfg-if", "libc", ] @@ -2338,7 +2448,7 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4488594b9328dee448adb906d8b126d9b7deb7cf5c22161ee591610bb1be83c0" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "libc", "redox_syscall 0.5.15", ] @@ -2517,7 +2627,7 @@ version = "0.30.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "cfg-if", "cfg_aliases", "libc", @@ -2653,7 +2763,7 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "900831247d2fe1a09a683278e5384cfb8c80c79fe6b166f9d14bfdde0ea1b03c" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "objc2", ] @@ -2684,6 +2794,12 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + [[package]] name = "option-ext" version = "0.2.0" @@ -2725,6 +2841,31 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "parse-display" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "914a1c2265c98e2446911282c6ac86d8524f495792c38c5bd884f80499c7538a" +dependencies = [ + "parse-display-derive", + "regex", + "regex-syntax 0.8.5", +] + +[[package]] +name = "parse-display-derive" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ae7800a4c974efd12df917266338e79a7a74415173caf7e70aa0a0707345281" +dependencies = [ + "proc-macro2", + "quote", + "regex", + "regex-syntax 0.8.5", + "structmeta", + "syn 2.0.104", +] + [[package]] name = "paste" version = "1.0.15" @@ -3182,7 +3323,7 @@ version = "0.5.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e8af0dde094006011e6a740d4879319439489813bd0bcdc7d821beaeeff48ec" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", ] [[package]] @@ -3506,11 +3647,11 @@ version = "0.38.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "errno", "libc", "linux-raw-sys 0.4.15", - "windows-sys 0.59.0", + "windows-sys 0.52.0", ] [[package]] @@ -3519,7 +3660,7 @@ version = "1.0.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "11181fbabf243db407ef8df94a6ce0b2f9a733bd8be4ad02b4eda9602296cac8" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "errno", "libc", "linux-raw-sys 0.9.4", @@ -3540,6 +3681,27 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + +[[package]] +name = "rustls-pemfile" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "rustls-pki-types" version = "1.12.0" @@ -3606,6 +3768,15 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "schemars" version = "0.9.0" @@ -3668,6 +3839,29 @@ dependencies = [ "syn 2.0.104", ] +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags 2.13.1", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "semver" version = "1.0.28" @@ -4029,6 +4223,29 @@ version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "structmeta" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e1575d8d40908d70f6fd05537266b90ae71b15dbbe7a8b7dffa2b759306d329" +dependencies = [ + "proc-macro2", + "quote", + "structmeta-derive", + "syn 2.0.104", +] + +[[package]] +name = "structmeta-derive" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "152a0b65a590ff6c3da95cabe2353ee04e6167c896b28e3b14478c2636c922fc" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.104", +] + [[package]] name = "strum" version = "0.24.1" @@ -4208,6 +4425,35 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "testcontainers" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23bb7577dca13ad86a78e8271ef5d322f37229ec83b8d98da6d996c588a1ddb1" +dependencies = [ + "async-trait", + "bollard", + "bollard-stubs", + "bytes", + "docker_credential", + "either", + "etcetera", + "futures", + "log", + "memchr", + "parse-display", + "pin-project-lite", + "serde", + "serde_json", + "serde_with", + "thiserror 2.0.12", + "tokio", + "tokio-stream", + "tokio-tar", + "tokio-util", + "url", +] + [[package]] name = "testutils" version = "0.3.70-rc.1" @@ -4590,6 +4836,7 @@ dependencies = [ "snafu", "spinners", "tempfile", + "testcontainers", "testutils", "tokio", "tokio-test", @@ -4612,7 +4859,7 @@ version = "0.6.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adc82fd73de2a9722ac5da747f12383d2bfdb93591ee6c58486e0097890f05f2" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", "bytes", "futures-util", "http", @@ -5130,7 +5377,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.48.0", ] [[package]] @@ -5147,7 +5394,7 @@ checksum = "c0fdd3ddb90610c7638aa2b3a3ab2904fb9e5cdbecc643ddb3647212781c4ae3" dependencies = [ "windows-implement", "windows-interface", - "windows-link", + "windows-link 0.1.3", "windows-result", "windows-strings", ] @@ -5180,13 +5427,19 @@ version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + [[package]] name = "windows-result" version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "56f42bd332cc6c8eac5af113fc0c1fd6a8fd2aa08a0119358686e5160d0586c6" dependencies = [ - "windows-link", + "windows-link 0.1.3", ] [[package]] @@ -5195,7 +5448,7 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "56e6c93f3a0c3b36176cb1327a4958a0353d5d166c2a35cb268ace15e91d3b57" dependencies = [ - "windows-link", + "windows-link 0.1.3", ] [[package]] @@ -5243,6 +5496,15 @@ dependencies = [ "windows-targets 0.53.2", ] +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link 0.2.1", +] + [[package]] name = "windows-targets" version = "0.42.2" @@ -5506,7 +5768,7 @@ version = "0.39.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1" dependencies = [ - "bitflags 2.9.1", + "bitflags 2.13.1", ] [[package]] diff --git a/crates/tower-cmd/Cargo.toml b/crates/tower-cmd/Cargo.toml index 6d5bcf5e..db90795f 100644 --- a/crates/tower-cmd/Cargo.toml +++ b/crates/tower-cmd/Cargo.toml @@ -49,3 +49,4 @@ uuid = { version = "1.0", features = ["v4"] } futures = { workspace = true } cucumber = { version = "0.21", features = ["macros"] } tokio-test = "0.4" +testcontainers = { version = "0.24", features = ["blocking"] } diff --git a/crates/tower-cmd/src/catalogs.rs b/crates/tower-cmd/src/catalogs.rs index eba2a6e0..003e5e41 100644 --- a/crates/tower-cmd/src/catalogs.rs +++ b/crates/tower-cmd/src/catalogs.rs @@ -10,6 +10,19 @@ use crate::{api, beta, output, util::cmd}; const STORAGE_CATALOG_TYPE: &str = "tower-catalog"; +/// Row cap for agent-issued queries (the MCP path). Rows past this are dropped +/// and the caller is told the set was truncated, so a model can't pull an +/// unbounded table into memory or its context. Human CLI queries are uncapped. +pub(crate) const AGENT_MAX_ROWS: usize = 1_000; + +/// Leading keywords that write data/schema or repoint the session. An +/// agent-issued query starting with one of these is refused — defence in depth +/// on top of the read-only credentials it is given. +const WRITE_LEADING_KEYWORDS: &[&str] = &[ + "insert", "update", "delete", "merge", "create", "drop", "alter", "truncate", "replace", + "copy", "attach", "detach", +]; + pub fn catalogs_cmd() -> Command { Command::new("catalogs") .about(format!( @@ -309,9 +322,10 @@ pub async fn do_show(out: &output::Out, config: Config, args: &ArgMatches) { out.text(&human, &json_data); } -/// Lists the tables in a catalog by attaching it in DuckDB. Unlike the query -/// path this never exits: `show` should still render catalog details when the -/// tables can't be fetched. +/// Lists the tables in a catalog by attaching it in DuckDB, wrapping +/// [`list_catalog_tables`] with the CLI spinner. Unlike the query path this +/// never exits: `show` should still render catalog details when the tables +/// can't be fetched. async fn fetch_catalog_tables( out: &output::Out, config: &Config, @@ -319,21 +333,34 @@ async fn fetch_catalog_tables( env: &str, ) -> Result { let mut spinner = out.spinner("Listing tables..."); + match list_catalog_tables(config, name, env).await { + Ok(query_result) => { + spinner.success(out); + Ok(query_result) + } + Err(err) => { + spinner.failure(out); + Err(err) + } + } +} - let response = match api::vend_catalog_credentials( +/// Attaches a storage catalog read-only and returns its (namespace, table) +/// rows via `SHOW ALL TABLES`. Shared by the CLI `show` and the MCP server. +/// Errors are returned with the OAuth token redacted. +pub(crate) async fn list_catalog_tables( + config: &Config, + name: &str, + env: &str, +) -> Result { + let response = api::vend_catalog_credentials( config, name, env, vend_catalog_credentials_body::Mode::Read, ) .await - { - Ok(response) => response, - Err(err) => { - spinner.failure(out); - return Err(err.to_string()); - } - }; + .map_err(|err| err.to_string())?; let token = response.credentials.oauth_token.clone(); let setup = attach_statements( @@ -348,22 +375,14 @@ async fn fetch_catalog_tables( &setup, "SELECT \"schema\", name FROM (SHOW ALL TABLES) WHERE database = ? ORDER BY \"schema\", name", duckdb::params![db_name], + None, ) }) .await; match result { - Ok(Ok(query_result)) => { - spinner.success(out); - Ok(query_result) - } - Ok(Err(err)) => { - spinner.failure(out); - Err(redact_token(&err.to_string(), &token)) - } - Err(err) => { - spinner.failure(out); - Err(redact_token(&err.to_string(), &token)) - } + Ok(Ok(query_result)) => Ok(query_result), + Ok(Err(err)) => Err(redact_token(&err.to_string(), &token)), + Err(err) => Err(redact_token(&err.to_string(), &token)), } } @@ -393,74 +412,254 @@ pub async fn do_query(out: &output::Out, config: Config, args: &ArgMatches) { out.die("No SQL statement provided. Pass one with --sql or pipe it via stdin."); } - let response = match api::describe_catalog(&config, name, &env).await { - Ok(response) => response, - Err(err) => out.tower_error_and_die(err, "Fetching catalog details failed"), - }; - if !is_storage_catalog_type(Some(&response.catalog.r#type)) { - out.die(&format!( - "Querying is only supported for {} catalogs; '{}' has type '{}'.", - STORAGE_CATALOG_TYPE, name, response.catalog.r#type - )); + let write = cmd::get_bool_flag(args, "write"); + + let mut spinner = out.spinner("Running query..."); + match query_catalog(&config, name, &env, sql, QueryOptions::cli(write)).await { + Ok(query_result) => { + spinner.success(out); + output_query_result(out, &query_result); + } + Err(err) => { + spinner.failure(out); + out.die(&err); + } } +} - let write = cmd::get_bool_flag(args, "write"); - let query_result = execute_catalog_query(out, &config, name, &env, sql, write).await; - output_query_result(out, &query_result); +/// How a query should be run: which credentials to vend, whether to lock the +/// DuckDB session down for an untrusted (agent) caller, and how many rows to +/// return. +pub(crate) struct QueryOptions { + /// Vend read-write credentials and attach the catalog read-write. + write: bool, + /// Lock the session for an untrusted caller: reject multi-statement and + /// write/DDL SQL, block local-filesystem and community-extension access, + /// and prevent the query from unwinding those settings. + sandbox: bool, + /// Drop rows past this count and flag the result truncated; `None` is + /// unbounded. + max_rows: Option, } -/// Vends credentials for the catalog, attaches it in an in-memory DuckDB, and -/// runs `sql` against it. Read-only unless `write` is set, in which case -/// read-write credentials are vended and the attach allows writes. Dies with a -/// user-facing error on failure. -async fn execute_catalog_query( - out: &output::Out, +impl QueryOptions { + /// Human-operated CLI query: unsandboxed and uncapped, honouring `--write`. + pub(crate) fn cli(write: bool) -> Self { + Self { + write, + sandbox: false, + max_rows: None, + } + } + + /// Agent-operated MCP query: always read-only, sandboxed, and row-capped. + pub(crate) fn agent() -> Self { + Self { + write: false, + sandbox: true, + max_rows: Some(AGENT_MAX_ROWS), + } + } +} + +/// Result-returning core shared by the CLI (`do_query`) and the MCP server. +/// Validates that the catalog is a Tower-managed storage catalog, vends +/// short-lived credentials, attaches it in an in-memory DuckDB, and runs `sql` +/// against it. When `options.sandbox` is set the SQL is gated to a single +/// read-only statement and the session is hardened before it runs. Errors are +/// returned (with the OAuth token redacted) rather than printed. +pub(crate) async fn query_catalog( config: &Config, name: &str, env: &str, sql: String, - write: bool, -) -> QueryResult { - let mode = if write { + options: QueryOptions, +) -> Result { + if options.sandbox { + if contains_multiple_statements(&sql) { + return Err("Only a single SQL statement can be run per query.".to_string()); + } + if is_write_statement(&sql) { + return Err( + "This query is read-only; write and DDL statements are not allowed.".to_string(), + ); + } + } + + let response = api::describe_catalog(config, name, env) + .await + .map_err(|err| format!("Fetching catalog details failed: {err}"))?; + if !is_storage_catalog_type(Some(&response.catalog.r#type)) { + return Err(format!( + "Querying is only supported for {} catalogs; '{}' has type '{}'.", + STORAGE_CATALOG_TYPE, name, response.catalog.r#type + )); + } + + let mode = if options.write { vend_catalog_credentials_body::Mode::ReadWrite } else { vend_catalog_credentials_body::Mode::Read }; - let mut spinner = out.spinner("Running query..."); - - let response = match api::vend_catalog_credentials(config, name, env, mode).await { - Ok(response) => response, - Err(err) => { - spinner.failure(out); - out.tower_error_and_die(err, "Running query failed"); - } - }; + let response = api::vend_catalog_credentials(config, name, env, mode) + .await + .map_err(|err| format!("Running query failed: {err}"))?; let token = response.credentials.oauth_token.clone(); - let setup = attach_statements(name, &response.credentials, mode); - let result = tokio::task::spawn_blocking(move || run_duckdb_query(&setup, &sql, [])).await; + let mut setup = attach_statements(name, &response.credentials, mode); + if options.sandbox { + setup.extend(agent_hardening_statements()); + } + let max_rows = options.max_rows; + let result = + tokio::task::spawn_blocking(move || run_duckdb_query(&setup, &sql, [], max_rows)).await; match result { - Ok(Ok(query_result)) => { - spinner.success(out); - query_result - } - Ok(Err(err)) => { - spinner.failure(out); - out.die(&format!( - "Query failed: {}", - redact_token(&err.to_string(), &token) - )); + Ok(Ok(query_result)) => Ok(query_result), + Ok(Err(err)) => Err(format!( + "Query failed: {}", + redact_token(&err.to_string(), &token) + )), + Err(err) => Err(format!( + "Query execution panicked: {}", + redact_token(&err.to_string(), &token) + )), + } +} + +/// Locks an attached DuckDB session down before an untrusted query runs: no +/// local-filesystem access (so `read_csv('/etc/passwd')` and friends fail), no +/// community extensions, and a configuration lock so the query can't unwind any +/// of it. These run after the catalog is attached — the attach itself installs +/// extensions and reaches the network. Network access stays on (httpfs is how +/// the Iceberg catalog reads its data), so this narrows but does not eliminate +/// the query surface; full isolation needs OS/container/WASM sandboxing. +fn agent_hardening_statements() -> Vec { + vec![ + "SET disabled_filesystems = 'LocalFileSystem'".to_string(), + "SET allow_community_extensions = false".to_string(), + "SET lock_configuration = true".to_string(), + ] +} + +/// The leading SQL keyword, lowercased, after skipping leading whitespace and +/// `--` / `/* */` comments. +fn first_sql_keyword(sql: &str) -> String { + let mut s = sql.trim_start(); + loop { + if let Some(rest) = s.strip_prefix("--") { + match rest.find('\n') { + Some(nl) => s = rest[nl + 1..].trim_start(), + None => return String::new(), + } + } else if let Some(rest) = s.strip_prefix("/*") { + match rest.find("*/") { + Some(end) => s = rest[end + 2..].trim_start(), + None => return String::new(), + } + } else { + break; } - Err(err) => { - spinner.failure(out); - out.die(&format!( - "Query execution panicked: {}", - redact_token(&err.to_string(), &token) - )); + } + s.chars() + .take_while(|c| c.is_ascii_alphabetic()) + .collect::() + .to_lowercase() +} + +/// True when `sql` starts with a write/DDL keyword. +fn is_write_statement(sql: &str) -> bool { + WRITE_LEADING_KEYWORDS.contains(&first_sql_keyword(sql).as_str()) +} + +/// True when `sql` holds more than one statement. A `;` inside a string literal +/// or comment is data, not a separator, so those spans are skipped. This gate +/// matters because duckdb-rs `prepare` runs every statement but the last as a +/// side effect, so unguarded multi-statement SQL would execute its leading +/// statements even though only the final one is returned. +fn contains_multiple_statements(sql: &str) -> bool { + #[derive(PartialEq)] + enum State { + Normal, + Single, + Double, + Line, + Block, + } + + let mut state = State::Normal; + let mut statements = 0usize; + let mut current_has_content = false; + let mut chars = sql.chars().peekable(); + + while let Some(c) = chars.next() { + match state { + State::Normal => match c { + '\'' => { + state = State::Single; + current_has_content = true; + } + '"' => { + state = State::Double; + current_has_content = true; + } + '-' if chars.peek() == Some(&'-') => { + chars.next(); + state = State::Line; + } + '/' if chars.peek() == Some(&'*') => { + chars.next(); + state = State::Block; + } + ';' => { + if current_has_content { + statements += 1; + if statements > 1 { + return true; + } + } + current_has_content = false; + } + c if c.is_whitespace() => {} + _ => current_has_content = true, + }, + State::Single => { + if c == '\'' { + if chars.peek() == Some(&'\'') { + chars.next(); + } else { + state = State::Normal; + } + } + } + State::Double => { + if c == '"' { + if chars.peek() == Some(&'"') { + chars.next(); + } else { + state = State::Normal; + } + } + } + State::Line => { + if c == '\n' { + state = State::Normal; + } + } + State::Block => { + if c == '*' && chars.peek() == Some(&'/') { + chars.next(); + state = State::Normal; + } + } } } + + if current_has_content { + statements += 1; + } + statements > 1 } fn read_sql_from_stdin(out: &output::Out) -> String { @@ -475,9 +674,12 @@ fn read_sql_from_stdin(out: &output::Out) -> String { sql } -struct QueryResult { - columns: Vec, - rows: Vec>, +#[derive(Debug)] +pub(crate) struct QueryResult { + pub(crate) columns: Vec, + pub(crate) rows: Vec>, + /// Rows were dropped to honour a caller-supplied row cap. + pub(crate) truncated: bool, } /// Statements that install the Iceberg support and attach the catalog under @@ -515,11 +717,13 @@ fn attach_statements( /// Runs `setup` statements one at a time, then `query` as a prepared statement /// with `params` bound. Values that fit a bind position should go through -/// `params` rather than into the query text. +/// `params` rather than into the query text. When `max_rows` is set, rows past +/// it are dropped and the result is flagged truncated. fn run_duckdb_query( setup: &[String], query: &str, params: P, + max_rows: Option, ) -> Result { let conn = duckdb::Connection::open_in_memory()?; for statement in setup { @@ -529,6 +733,7 @@ fn run_duckdb_query( let mut stmt = conn.prepare(query)?; let mut columns: Vec = Vec::new(); let mut rows = Vec::new(); + let mut truncated = false; { let mut result_rows = stmt.query(params)?; @@ -536,6 +741,10 @@ fn run_duckdb_query( if columns.is_empty() { columns = row.as_ref().column_names(); } + if max_rows.is_some_and(|max| rows.len() >= max) { + truncated = true; + break; + } let mut record = Vec::with_capacity(columns.len()); for idx in 0..columns.len() { let value: duckdb::types::Value = row.get(idx)?; @@ -550,7 +759,11 @@ fn run_duckdb_query( columns = stmt.column_names(); } - Ok(QueryResult { columns, rows }) + Ok(QueryResult { + columns, + rows, + truncated, + }) } fn duckdb_value_to_json(value: duckdb::types::Value) -> serde_json::Value { @@ -657,7 +870,14 @@ fn output_query_result(out: &output::Out, result: &QueryResult) { .collect(); out.table(result.columns.clone(), data, Some(&json_rows)); - out.note(&format!("\n{} row(s)\n", result.rows.len())); + if result.truncated { + out.note(&format!( + "\n{} row(s) (truncated; more rows matched)\n", + result.rows.len() + )); + } else { + out.note(&format!("\n{} row(s)\n", result.rows.len())); + } } fn json_value_to_cell(value: &serde_json::Value) -> String { @@ -668,7 +888,7 @@ fn json_value_to_cell(value: &serde_json::Value) -> String { } } -fn is_storage_catalog_type(catalog_type: Option<&str>) -> bool { +pub(crate) fn is_storage_catalog_type(catalog_type: Option<&str>) -> bool { catalog_type == Some(STORAGE_CATALOG_TYPE) } @@ -920,8 +1140,9 @@ fn snippets( #[cfg(test)] mod tests { use super::{ - attach_statements, catalogs_cmd, duckdb_value_to_json, is_storage_catalog_type, parse_mode, - run_duckdb_query, snippets, token_export_command, + agent_hardening_statements, attach_statements, catalogs_cmd, contains_multiple_statements, + duckdb_value_to_json, first_sql_keyword, is_storage_catalog_type, is_write_statement, + parse_mode, run_duckdb_query, snippets, token_export_command, }; use tower_api::models::{vend_catalog_credentials_body, CatalogCredentials}; @@ -1139,6 +1360,7 @@ mod tests { &setup, "SELECT \"schema\", name FROM (SHOW ALL TABLES) WHERE database = ? ORDER BY \"schema\", name", duckdb::params!["memory"], + None, ) .expect("query should succeed"); @@ -1287,7 +1509,7 @@ mod tests { "CREATE TABLE t (id INTEGER, name VARCHAR); INSERT INTO t VALUES (1, 'a'), (2, NULL);" .to_string(), ]; - let result = run_duckdb_query(&setup, "SELECT id, name FROM t ORDER BY id", []) + let result = run_duckdb_query(&setup, "SELECT id, name FROM t ORDER BY id", [], None) .expect("query should succeed"); assert_eq!(result.columns, vec!["id", "name"]); @@ -1308,6 +1530,7 @@ mod tests { &[], "SELECT [1, 2] AS l, {'a': 1, 'b': 'x'} AS s, MAP {'k': 2} AS m", [], + None, ) .expect("query should succeed"); @@ -1322,15 +1545,274 @@ mod tests { ); } + /// Polls MinIO's health endpoint over a raw socket until it answers 200. + fn wait_for_minio(port: u16) -> bool { + use std::io::{Read, Write}; + for _ in 0..60 { + if let Ok(mut stream) = std::net::TcpStream::connect(("127.0.0.1", port)) { + let _ = stream.set_read_timeout(Some(std::time::Duration::from_secs(2))); + let request = + "GET /minio/health/live HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n"; + let mut response = String::new(); + if stream.write_all(request.as_bytes()).is_ok() + && stream.read_to_string(&mut response).is_ok() + && response.contains(" 200 ") + { + return true; + } + } + std::thread::sleep(std::time::Duration::from_millis(500)); + } + false + } + + /// Sandbox integration test against real object storage. It starts a MinIO + /// container (via testcontainers), seeds a real Iceberg table on it, and + /// reads it back through the hardened setup. This is the coverage the + /// local-only tests can't give: it proves the DuckDB hardening does NOT + /// break a real object-store Iceberg read (the reader uses S3FileSystem, + /// which the hardening leaves enabled) while local-filesystem access and + /// configuration changes stay blocked in the same session. Needs a Docker + /// daemon; it self-skips when one isn't available, so a plain `cargo test` + /// still passes on a machine without Docker. + #[test] + fn sandbox_holds_over_object_store_iceberg() { + use testcontainers::core::{IntoContainerPort, WaitFor}; + use testcontainers::runners::SyncRunner; + use testcontainers::{GenericImage, ImageExt}; + + let bucket = "warehouse"; + // MinIO exposes a top-level data-dir subdirectory as a bucket, so + // pre-creating one in the container command needs no S3 client or `mc`. + let image = GenericImage::new("minio/minio", "latest") + .with_wait_for(WaitFor::seconds(1)) + .with_exposed_port(9000.tcp()) + .with_entrypoint("sh") + .with_env_var("MINIO_ROOT_USER", "minioadmin") + .with_env_var("MINIO_ROOT_PASSWORD", "minioadmin") + .with_cmd([ + "-c".to_string(), + format!("mkdir -p /data/{bucket} && exec minio server /data"), + ]); + + let container = match image.start() { + Ok(container) => container, + Err(err) => { + eprintln!( + "skipping sandbox_holds_over_object_store_iceberg (no Docker daemon?): {err}" + ); + return; + } + }; + let port = container + .get_host_port_ipv4(9000.tcp()) + .expect("mapped MinIO port"); + assert!(wait_for_minio(port), "MinIO did not become healthy"); + + let secret = format!( + "CREATE SECRET s3sec (TYPE s3, KEY_ID 'minioadmin', SECRET 'minioadmin', \ + ENDPOINT '127.0.0.1:{port}', URL_STYLE 'path', USE_SSL false, REGION 'us-east-1')" + ); + let table = format!("s3://{bucket}/tbl"); + let extensions = || { + vec![ + "INSTALL httpfs".to_string(), + "LOAD httpfs".to_string(), + "INSTALL iceberg".to_string(), + "LOAD iceberg".to_string(), + ] + }; + + // Seed a real Iceberg table on object storage. + let seed = duckdb::Connection::open_in_memory().expect("open duckdb"); + for stmt in extensions() { + seed.execute_batch(&stmt).expect("load extension"); + } + seed.execute_batch(&secret).expect("create s3 secret"); + seed.execute_batch(&format!( + "COPY (SELECT * FROM (VALUES (1,'a'),(2,'b'),(3,NULL)) t(id,name)) \ + TO '{table}' (FORMAT iceberg)" + )) + .expect("seed iceberg table on object storage"); + + // Hardened setup mirrors the agent query path: extensions + secret, then lockdown. + let mut setup = extensions(); + setup.push(secret); + setup.extend(super::agent_hardening_statements()); + + // 1) The real object-store read still works under the sandbox. + let read = run_duckdb_query( + &setup, + &format!("SELECT id, name FROM iceberg_scan('{table}') ORDER BY id"), + [], + None, + ) + .expect("hardening must not break object-store Iceberg reads"); + assert_eq!(read.columns, vec!["id", "name"]); + assert_eq!( + read.rows, + vec![ + vec![serde_json::json!(1), serde_json::json!("a")], + vec![serde_json::json!(2), serde_json::json!("b")], + vec![serde_json::json!(3), serde_json::Value::Null], + ] + ); + + // 2) Local filesystem access stays blocked in the same session. + let local = run_duckdb_query(&setup, "SELECT * FROM read_csv('/etc/hostname')", [], None) + .expect_err("local filesystem reads must stay blocked"); + assert!( + local.to_string().to_lowercase().contains("disabled"), + "unexpected error: {local}" + ); + + // 3) Configuration stays locked. + let cfg = run_duckdb_query(&setup, "SET memory_limit = '1GB'", [], None) + .expect_err("configuration must stay locked"); + let cfg = cfg.to_string().to_lowercase(); + assert!( + cfg.contains("lock") || cfg.contains("configuration"), + "unexpected error: {cfg}" + ); + } + + /// Regression check against real Iceberg data. The other `run_duckdb_query` + /// tests use plain in-memory tables; this one writes a small Iceberg table + /// with DuckDB's `COPY … (FORMAT iceberg)` (real metadata + manifests + + /// parquet) and reads it back through `run_duckdb_query`, so the iceberg + /// reader and our column/row extraction are exercised end to end, NULLs and + /// the row cap included. Needs the `iceberg` DuckDB extension, which is + /// fetched on first run and then cached; it is not exercised under the + /// sandbox because `iceberg_scan` of a local path uses the LocalFileSystem + /// that the hardening deliberately disables (production reads S3). + #[test] + fn iceberg_scan_reads_written_table_through_run_duckdb_query() { + let dir = std::env::temp_dir().join(format!("tower_iceberg_test_{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).expect("create temp dir"); + let table = dir.join("events").to_string_lossy().replace('\'', "''"); + + // Generate a real Iceberg table on disk. + let conn = duckdb::Connection::open_in_memory().expect("open duckdb"); + conn.execute_batch("INSTALL iceberg").expect("install iceberg"); + conn.execute_batch("LOAD iceberg").expect("load iceberg"); + conn.execute_batch(&format!( + "COPY (SELECT * FROM (VALUES (1, 'a'), (2, 'b'), (3, NULL)) AS t(id, name)) \ + TO '{table}' (FORMAT iceberg)" + )) + .expect("write iceberg table"); + + let setup = vec!["INSTALL iceberg".to_string(), "LOAD iceberg".to_string()]; + + // Columns, positional rows, and a NULL survive the round trip. + let result = run_duckdb_query( + &setup, + &format!("SELECT id, name FROM iceberg_scan('{table}') ORDER BY id"), + [], + None, + ) + .expect("iceberg_scan should read the table"); + assert_eq!(result.columns, vec!["id", "name"]); + assert_eq!( + result.rows, + vec![ + vec![serde_json::json!(1), serde_json::json!("a")], + vec![serde_json::json!(2), serde_json::json!("b")], + vec![serde_json::json!(3), serde_json::Value::Null], + ] + ); + assert!(!result.truncated); + + // The row cap applies to Iceberg reads too. + let capped = run_duckdb_query( + &setup, + &format!("SELECT id FROM iceberg_scan('{table}') ORDER BY id"), + [], + Some(2), + ) + .expect("iceberg_scan should read the table"); + assert_eq!(capped.rows.len(), 2); + assert!(capped.truncated); + + let _ = std::fs::remove_dir_all(&dir); + } + #[test] fn run_duckdb_query_reports_columns_for_empty_results() { - let result = - run_duckdb_query(&[], "SELECT 1 AS x WHERE 1 = 0", []).expect("query should succeed"); + let result = run_duckdb_query(&[], "SELECT 1 AS x WHERE 1 = 0", [], None) + .expect("query should succeed"); assert_eq!(result.columns, vec!["x"]); assert!(result.rows.is_empty()); } + #[test] + fn run_duckdb_query_caps_rows_and_flags_truncation() { + let capped = run_duckdb_query(&[], "SELECT * FROM range(5) AS t(i)", [], Some(3)) + .expect("query should succeed"); + assert_eq!(capped.rows.len(), 3); + assert!(capped.truncated); + + let exact = run_duckdb_query(&[], "SELECT * FROM range(3) AS t(i)", [], Some(3)) + .expect("query should succeed"); + assert_eq!(exact.rows.len(), 3); + assert!(!exact.truncated); + } + + #[test] + fn contains_multiple_statements_ignores_separators_in_strings_and_comments() { + assert!(!contains_multiple_statements("SELECT 1")); + assert!(!contains_multiple_statements("SELECT 1;")); + assert!(!contains_multiple_statements(" SELECT 1 ; ")); + assert!(!contains_multiple_statements("SELECT 'a;b'")); + assert!(!contains_multiple_statements("SELECT 1 -- ; not a statement")); + assert!(!contains_multiple_statements("SELECT 1; -- trailing comment")); + assert!(!contains_multiple_statements("SELECT /* ; */ 1")); + + assert!(contains_multiple_statements("SELECT 1; SELECT 2")); + assert!(contains_multiple_statements("SELECT 1; DROP TABLE t")); + assert!(contains_multiple_statements("SELECT 'a;b'; SELECT 2")); + } + + #[test] + fn write_statements_are_detected_through_case_and_comments() { + assert_eq!(first_sql_keyword(" SELECT 1"), "select"); + assert_eq!(first_sql_keyword("/* c */ INSERT INTO t VALUES (1)"), "insert"); + assert_eq!(first_sql_keyword("-- lead\nDELETE FROM t"), "delete"); + + assert!(!is_write_statement("SELECT * FROM t")); + assert!(!is_write_statement("WITH x AS (SELECT 1) SELECT * FROM x")); + assert!(is_write_statement("insert into t values (1)")); + assert!(is_write_statement(" DROP TABLE t")); + assert!(is_write_statement("COPY t TO 'out.csv'")); + assert!(is_write_statement("ATTACH 'x' AS y")); + } + + #[test] + fn agent_hardening_allows_selects_but_blocks_local_files_and_config_changes() { + let setup = agent_hardening_statements(); + + let ok = run_duckdb_query(&setup, "SELECT 1 AS x", [], None) + .expect("a plain select should still run under the hardened session"); + assert_eq!(ok.columns, vec!["x"]); + assert_eq!(ok.rows, vec![vec![serde_json::json!(1)]]); + + let fs_err = run_duckdb_query(&setup, "SELECT * FROM read_csv('Cargo.toml')", [], None) + .expect_err("local filesystem access should be blocked"); + assert!( + fs_err.to_string().to_lowercase().contains("disabled"), + "unexpected error: {fs_err}" + ); + + let cfg_err = run_duckdb_query(&setup, "SET memory_limit = '1GB'", [], None) + .expect_err("configuration should be locked"); + let cfg_msg = cfg_err.to_string().to_lowercase(); + assert!( + cfg_msg.contains("lock") || cfg_msg.contains("configuration"), + "unexpected error: {cfg_err}" + ); + } + #[test] fn token_export_command_fetches_token_without_printing_it() { let credentials = CatalogCredentials::new( diff --git a/crates/tower-cmd/src/mcp.rs b/crates/tower-cmd/src/mcp.rs index fc301457..4c55dc85 100644 --- a/crates/tower-cmd/src/mcp.rs +++ b/crates/tower-cmd/src/mcp.rs @@ -199,6 +199,18 @@ struct ShowCatalogRequest { environment: Option, } +#[derive(Debug, Deserialize, JsonSchema)] +struct QueryCatalogRequest { + /// Name of the Tower-managed storage catalog to query + name: String, + /// A single read-only SQL statement. Tables must be fully qualified as + /// .. (there is no default schema). Writes, DDL, + /// and multiple statements are rejected. + sql: String, + /// The environment the catalog belongs to (defaults to "default") + environment: Option, +} + pub fn mcp_cmd() -> Command { Command::new("mcp-server") .about("Runs an MCP server for LLM interaction") @@ -719,7 +731,9 @@ impl TowerService { } } - #[tool(description = "Show details for a catalog, including its property names")] + #[tool( + description = "Show a catalog's details: its property names and, for Tower-managed storage catalogs, the namespaces and tables you can query." + )] async fn tower_catalogs_show( &self, Parameters(request): Parameters, @@ -739,17 +753,88 @@ impl TowerService { }) }) .collect(); + + // Only Tower-managed storage catalogs expose queryable tables; + // for anything else `tables` stays null. A listing failure is + // surfaced in `tables_error` without failing the whole call. + let (tables, tables_error) = + if crate::catalogs::is_storage_catalog_type(Some(&catalog.r#type)) { + match crate::catalogs::list_catalog_tables( + &self.config, + &request.name, + environment, + ) + .await + { + Ok(result) => ( + Value::Array( + result + .rows + .iter() + .map(|row| { + json!({ + "namespace": row.first().cloned().unwrap_or(Value::Null), + "table": row.get(1).cloned().unwrap_or(Value::Null), + }) + }) + .collect(), + ), + Value::Null, + ), + Err(e) => (Value::Null, Value::String(e)), + } + } else { + (Value::Null, Value::Null) + }; + Self::json_success(json!({ "name": catalog.name, "type": catalog.r#type, "environment": catalog.environment, "properties": properties, + "tables": tables, + "tables_error": tables_error, })) } Err(e) => Self::error_result("Failed to show catalog", e), } } + #[tool( + description = "Run one read-only SQL statement against a Tower-managed storage (Iceberg) catalog and return the columns plus rows as positional arrays (one value per column, in column order). Fully qualify tables as \"\".\"\".\"
\"; call tower_catalogs_show first to list a catalog's namespaces and tables. Only a single statement runs per call, and writes/DDL are rejected. Results are capped at 1000 rows — when more matched, the response sets \"truncated\": true, so narrow with WHERE/LIMIT or aggregate." + )] + async fn tower_catalogs_query( + &self, + Parameters(request): Parameters, + ) -> Result { + let environment = request.environment.as_deref().unwrap_or("default"); + let sql = request.sql.trim().to_string(); + if sql.is_empty() { + return Self::text_error("No SQL statement provided.".to_string()); + } + + // Agents get read-only, sandboxed, row-capped access; positional rows + // preserve duplicate column names (e.g. from joins) that an object keyed + // by name would silently collapse. + match crate::catalogs::query_catalog( + &self.config, + &request.name, + environment, + sql, + crate::catalogs::QueryOptions::agent(), + ) + .await + { + Ok(result) => Self::json_success(json!({ + "columns": result.columns, + "rows": result.rows, + "row_count": result.rows.len(), + "truncated": result.truncated, + })), + Err(e) => Self::text_error(e), + } + } + #[tool(description = "List teams you belong to")] async fn tower_teams_list(&self) -> Result { if self.config.api_key.is_some() { diff --git a/tests/integration/features/environment.py b/tests/integration/features/environment.py index b6fb2d71..c3227b4b 100644 --- a/tests/integration/features/environment.py +++ b/tests/integration/features/environment.py @@ -16,6 +16,25 @@ def before_all(context): def before_scenario(context, scenario): + # Catalog scenarios hit the real attach -> Iceberg -> query path, which the + # mock server can't stand up. Skip them unless a storage catalog is + # configured (and rely on a real TOWER_URL + session being provided too). + if "catalogs" in scenario.effective_tags: + catalog = os.environ.get("TOWER_TEST_CATALOG") + if not catalog: + scenario.skip( + "set TOWER_TEST_CATALOG (and a real TOWER_URL) to run catalog tests" + ) + return + context.test_catalog = catalog + context.test_catalog_env = os.environ.get("TOWER_TEST_CATALOG_ENV", "default") + context.test_catalog_table = os.environ.get("TOWER_TEST_CATALOG_TABLE") + if "catalog-data" in scenario.effective_tags and not context.test_catalog_table: + scenario.skip( + "set TOWER_TEST_CATALOG_TABLE to run the catalog data query test" + ) + return + # Create a temporary working directory for this scenario context.temp_dir = tempfile.mkdtemp(prefix="tower_test_") context.original_cwd = os.getcwd() diff --git a/tests/integration/features/mcp_catalogs.feature b/tests/integration/features/mcp_catalogs.feature new file mode 100644 index 00000000..97671112 --- /dev/null +++ b/tests/integration/features/mcp_catalogs.feature @@ -0,0 +1,35 @@ +@catalogs +Feature: MCP catalog querying + As a developer using the Tower MCP server with an agent + I want to query Tower-managed storage catalogs safely + So that an agent can explore data without write access or unbounded results + + # These scenarios exercise the real attach -> Iceberg -> query path, so they + # are skipped unless a storage catalog is configured (see environment.py): + # TOWER_TEST_CATALOG name of a tower-catalog storage catalog + # TOWER_URL a real server (not the mock) with a valid session + # TOWER_TEST_CATALOG_ENV environment the catalog lives in (default: default) + # TOWER_TEST_CATALOG_TABLE optional "cat"."ns"."tbl" for the data query below + + Scenario: Show a storage catalog lists its tables + When I show the test catalog via MCP + Then I should receive a success response + And the catalog response should list tables + + Scenario: Run a read-only query against the catalog + When I query the test catalog with SQL "SELECT 1 AS one" via MCP + Then I should receive query results + And the query result columns should include "one" + + Scenario: Write statements are rejected + When I query the test catalog with SQL "DROP TABLE does_not_exist" via MCP + Then I should receive an error response about a read-only query + + Scenario: Multiple statements are rejected + When I query the test catalog with SQL "SELECT 1; SELECT 2" via MCP + Then I should receive an error response about a single statement + + @catalog-data + Scenario: Query a configured table returns positional rows + When I query the configured catalog table via MCP + Then I should receive query results diff --git a/tests/integration/features/steps/mcp_steps.py b/tests/integration/features/steps/mcp_steps.py index 1fa07c3d..6973aaa5 100644 --- a/tests/integration/features/steps/mcp_steps.py +++ b/tests/integration/features/steps/mcp_steps.py @@ -904,3 +904,120 @@ def step_then_receive_workflow_help_stdio(context): @given('I have a simple hello world application named "{app_name}"') def step_create_hello_world_app_named(context, app_name): create_towerfile(context, app_name=app_name) + + +# --- Catalog querying (gated on TOWER_TEST_CATALOG; see environment.py) ------- + + +def _first_text_content(response): + """Return the first text block of an MCP response, or ''.""" + for item in response.get("content", []): + if isinstance(item, dict): + if item.get("type") == "text": + return item.get("text", "") + elif getattr(item, "type", None) == "text": + return getattr(item, "text", "") + return "" + + +def _parse_json_content(response): + """Parse the response's text block as JSON, or None if it isn't JSON.""" + try: + return json.loads(_first_text_content(response)) + except (ValueError, TypeError): + return None + + +@when("I show the test catalog via MCP") +@async_run_until_complete +async def step_show_test_catalog(context): + await call_mcp_tool( + context, + "tower_catalogs_show", + {"name": context.test_catalog, "environment": context.test_catalog_env}, + ) + + +@when('I query the test catalog with SQL "{sql}" via MCP') +@async_run_until_complete +async def step_query_test_catalog(context, sql): + await call_mcp_tool( + context, + "tower_catalogs_query", + { + "name": context.test_catalog, + "environment": context.test_catalog_env, + "sql": sql, + }, + ) + + +@when("I query the configured catalog table via MCP") +@async_run_until_complete +async def step_query_configured_table(context): + await call_mcp_tool( + context, + "tower_catalogs_query", + { + "name": context.test_catalog, + "environment": context.test_catalog_env, + "sql": f"SELECT * FROM {context.test_catalog_table} LIMIT 5", + }, + ) + + +@then("the catalog response should list tables") +def step_catalog_lists_tables(context): + data = _parse_json_content(context.mcp_response) + assert ( + data is not None + ), f"Expected JSON content, got: {context.mcp_response.get('content')}" + assert "tables" in data, f"Response should include a 'tables' field, got: {data}" + + +@then("I should receive query results") +def step_receive_query_results(context): + assert context.mcp_response.get( + "success", False + ), f"Expected a successful query, got: {context.mcp_response}" + data = _parse_json_content(context.mcp_response) + assert ( + data is not None + ), f"Expected JSON content, got: {context.mcp_response.get('content')}" + assert isinstance( + data.get("columns"), list + ), f"Expected a 'columns' list, got: {data}" + assert isinstance(data.get("rows"), list), f"Expected a 'rows' list, got: {data}" + # Rows are positional arrays (one value per column), not objects keyed by name. + for row in data["rows"]: + assert isinstance(row, list), f"Expected positional row arrays, got: {row}" + + +@then('the query result columns should include "{column}"') +def step_query_columns_include(context, column): + data = _parse_json_content(context.mcp_response) + assert ( + data is not None and column in data.get("columns", []) + ), f"Expected column '{column}', got: {data}" + + +@then("I should receive an error response about a read-only query") +def step_error_read_only(context): + assert is_error_response( + context.mcp_response + ), f"Expected an error response, got: {context.mcp_response}" + text = str(context.mcp_response).lower() + assert ( + "read-only" in text or "not allowed" in text + ), f"Error should mention the read-only restriction, got: {context.mcp_response}" + + +@then("I should receive an error response about a single statement") +def step_error_single_statement(context): + assert is_error_response( + context.mcp_response + ), f"Expected an error response, got: {context.mcp_response}" + text = str(context.mcp_response).lower() + assert ( + "single" in text and "statement" in text + ), f"Error should mention the single-statement rule, got: {context.mcp_response}"