From a0a99495f78b3bcda19c59d04d33f468221c749f Mon Sep 17 00:00:00 2001 From: WaylandYang Date: Thu, 3 Sep 2026 10:51:33 +0800 Subject: [PATCH] The lakehouse is one protocol away: Trino, Databricks and Snowflake as query engines Co-Authored-By: Claude Fable 5.1 --- Cargo.lock | 111 +++++- README.md | 4 +- crates/utopia-server/Cargo.toml | 6 + crates/utopia-server/src/api/chat.rs | 4 +- .../src/api/datasource_routes.rs | 19 +- crates/utopia-server/src/api/tools.rs | 3 +- crates/utopia-server/src/query_engine.rs | 205 ---------- crates/utopia-server/src/query_engine/conn.rs | 253 +++++++++++++ .../src/query_engine/databricks.rs | 271 ++++++++++++++ crates/utopia-server/src/query_engine/mod.rs | 349 ++++++++++++++++++ .../src/query_engine/postgres.rs | 93 +++++ .../src/query_engine/snowflake.rs | 257 +++++++++++++ .../utopia-server/src/query_engine/trino.rs | 266 +++++++++++++ crates/utopia-store/src/datasources.rs | 21 +- ...0018-the-lakehouse-is-one-protocol-away.md | 58 +++ docs/decisions/README.md | 1 + migrations/0020_lakehouse_engines.sql | 7 + web/src/i18n/en.ts | 8 +- web/src/i18n/zh.ts | 7 +- web/src/pages/Settings.tsx | 10 +- 20 files changed, 1716 insertions(+), 237 deletions(-) delete mode 100644 crates/utopia-server/src/query_engine.rs create mode 100644 crates/utopia-server/src/query_engine/conn.rs create mode 100644 crates/utopia-server/src/query_engine/databricks.rs create mode 100644 crates/utopia-server/src/query_engine/mod.rs create mode 100644 crates/utopia-server/src/query_engine/postgres.rs create mode 100644 crates/utopia-server/src/query_engine/snowflake.rs create mode 100644 crates/utopia-server/src/query_engine/trino.rs create mode 100644 docs/decisions/0018-the-lakehouse-is-one-protocol-away.md create mode 100644 migrations/0020_lakehouse_engines.sql diff --git a/Cargo.lock b/Cargo.lock index 4fd495552..83ce8b63b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -116,6 +116,16 @@ dependencies = [ "password-hash", ] +[[package]] +name = "assert-json-diff" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "async-compression" version = "0.4.43" @@ -324,6 +334,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64ct" version = "1.8.3" @@ -934,6 +950,24 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c286de4e81ea2590afc24d754e0f83810c566f50a1388fa75ebd57928c0d9745" +[[package]] +name = "deadpool" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b" +dependencies = [ + "deadpool-runtime", + "lazy_static", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" + [[package]] name = "debug_unsafe" version = "0.1.4" @@ -1354,6 +1388,21 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" +[[package]] +name = "futures" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.34" @@ -1427,6 +1476,7 @@ version = "0.3.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" dependencies = [ + "futures-channel", "futures-core", "futures-io", "futures-macro", @@ -1569,6 +1619,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17592d60ebacc7d5e169f4663c5f84f9161cc90328abcfe8456f41e4dfcb284" + [[package]] name = "hex" version = "0.4.3" @@ -1727,7 +1783,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -2128,7 +2184,7 @@ version = "10.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eba32bfb4ffdeaca3e34431072faf01745c9b26d25504aa7a6cf5684334fc4fc" dependencies = [ - "base64", + "base64 0.22.1", "ed25519-dalek", "getrandom 0.2.17", "hmac", @@ -2555,6 +2611,16 @@ dependencies = [ "libm", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "object" version = "0.39.1" @@ -2572,7 +2638,7 @@ checksum = "d354792e39fa5f0009e47623cf8b15b099bf9a652fa55c6f817fe28ac84fea50" dependencies = [ "async-trait", "aws-lc-rs", - "base64", + "base64 0.22.1", "bytes", "chrono", "crc-fast", @@ -2806,7 +2872,7 @@ version = "3.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" dependencies = [ - "base64", + "base64 0.22.1", "serde_core", ] @@ -3305,7 +3371,7 @@ version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04e9018c9d814e5f30cc16a0f03271aeab3571e609612d9fe78c1aa8d11c2f62" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-core", @@ -3875,7 +3941,7 @@ version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ee6798b1838b6a0f69c007c133b8df5866302197e404e8b6ee8ed3e3a5e68dc6" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "chrono", "crc", @@ -3952,7 +4018,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aa003f0038df784eb8fecbbac13affe3da23b45194bd57dba231c8f48199c526" dependencies = [ "atoi", - "base64", + "base64 0.22.1", "bitflags", "byteorder", "bytes", @@ -3996,7 +4062,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db58fcd5a53cf07c184b154801ff91347e4c30d17a3562a635ff028ad5deda46" dependencies = [ "atoi", - "base64", + "base64 0.22.1", "bitflags", "byteorder", "chrono", @@ -4212,7 +4278,7 @@ checksum = "edde6a10743fff00a4e1a8c9ef020bf5f3cbad301b7d2d39f2b07f123c4eac07" dependencies = [ "aho-corasick", "arc-swap", - "base64", + "base64 0.22.1", "bitpacking", "bon", "byteorder", @@ -4922,6 +4988,7 @@ dependencies = [ "async-trait", "axum", "axum-extra", + "base64 0.23.1", "chrono", "dotenvy", "feed-rs", @@ -4929,6 +4996,7 @@ dependencies = [ "futures-util", "jsonwebtoken", "object_store", + "percent-encoding", "quick-xml 0.42.0", "reqwest", "serde", @@ -4940,6 +5008,7 @@ dependencies = [ "tower-http", "tracing", "tracing-subscriber", + "url", "utopia-core", "utopia-extract", "utopia-ingest", @@ -4947,6 +5016,7 @@ dependencies = [ "utopia-search", "utopia-store", "uuid", + "wiremock", ] [[package]] @@ -5506,6 +5576,29 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" +[[package]] +name = "wiremock" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08db1edfb05d9b3c1542e521aea074442088292f00b5f28e435c714a98f85031" +dependencies = [ + "assert-json-diff", + "base64 0.22.1", + "deadpool", + "futures", + "http", + "http-body-util", + "hyper", + "hyper-util", + "log", + "once_cell", + "regex", + "serde", + "serde_json", + "tokio", + "url", +] + [[package]] name = "wit-bindgen" version = "0.57.1" diff --git a/README.md b/README.md index a46c5af17..a71c26b23 100644 --- a/README.md +++ b/README.md @@ -61,7 +61,7 @@ One Rust binary and one Postgres. Full-text search is embedded in the binary, ve | **Entity resolution and review** | Duplicates are resolved in three stages: exact name or alias, embedding similarity, then a model's call on the doubtful pairs. Every merge can be undone. Uncertain cases go to a review queue: low-confidence extractions, suspected duplicates and cardinality conflicts. | | **Reasoning and derivation** | Ontology axioms compile into rules: transitivity, symmetry, inverses and relation hierarchy derive new facts by forward chaining. Derivation is off by default, since a wrong axiom derives wrong facts. A derived fact is marked as such on the graph, carries validity and confidence like any other, and shows what it was derived from. When it contradicts an asserted fact, the asserted one stands. | | **Conflict detection** | Three kinds of conflict, three sets of choices. A new fact that clashes with an older one: close the old, keep both, or reject the new. Data that breaks an axiom (self-loop, asymmetry, transitive cycle, cardinality): retract the fact, relax the axiom, or accept both. The ontology itself is checked first, because violations of a self-contradictory ontology are noise. | -| **Ontology-driven querying** | Mount a Postgres database on a base and chat can query it alongside the documents. The agent proposes how its tables map onto the ontology, and you confirm. The method behind it, [Ontology2SQL](https://github.com/deeplethe/ontology2sql), is state of the art on BIRD Mini-Dev for SQLite and PostgreSQL ([submission](https://github.com/bird-bench/bird-bench.github.io/pull/218)). | +| **Ontology-driven querying** | Mount a database on a base (Postgres, Trino for Iceberg / Delta Lake / Hive, Databricks, Snowflake) and chat can query it alongside the documents. The agent proposes how its tables map onto the ontology, and you confirm. The method behind it, [Ontology2SQL](https://github.com/deeplethe/ontology2sql), is state of the art on BIRD Mini-Dev for SQLite and PostgreSQL ([submission](https://github.com/bird-bench/bird-bench.github.io/pull/218)). | | **Multi-user and permissions** | Each knowledge base has its own members and roles: owner, admin, editor and viewer. Open bases are readable by everyone in the deployment, restricted ones only by invitation. The first account registered becomes the system administrator. | | **Decision ledger** | Confirming or rejecting a fact, merging or reverting an entity, rebuilding the graph: each leaves a record of who, when, and what the object looked like at the time. The ledger is append-only, and a record outlives its object, even the base it belonged to. | | **[Decision intelligence (in development)](#roadmap)** | Record a decision, replay both what was understood and the course it took, and reason over overlaid scenarios. | @@ -103,7 +103,7 @@ cd web && pnpm install && pnpm dev - [ ] **Decision reasoning**: constraint computation, and replaying a decision after the fact - [ ] **Execution gate**: checking an agent's calls against ontology rules and symbolic logic -- [ ] **Lakehouse for mapping and querying**: mapping exploration and Ontology2SQL over Iceberg / Delta Lake, Databricks, Snowflake and MaxCompute +- [ ] **MaxCompute**: mapping exploration and Ontology2SQL over Alibaba Cloud MaxCompute (Iceberg / Delta Lake via Trino, Databricks and Snowflake are in, awaiting a run against a real cluster) - [ ] **More sources**: MySQL, ClickHouse and Doris drivers; S3, WebDAV, Notion and Feishu connectors - [ ] **Time to the moment**: an `instant` precision beside year / month / day, for sources that carry a real timestamp. Today a connector rounds it to a UTC day, which can shift an event across midnight by one day - [ ] **Agent memory over MCP**: episode writes, the retrieve endpoint, and the MCP server diff --git a/crates/utopia-server/Cargo.toml b/crates/utopia-server/Cargo.toml index ffd2b5462..628e03dbf 100644 --- a/crates/utopia-server/Cargo.toml +++ b/crates/utopia-server/Cargo.toml @@ -41,3 +41,9 @@ async-trait = "0.1.92" sqlparser = "0.62.0" object_store = { version = "0.14.1", features = ["aws", "azure", "gcp"] } quick-xml.workspace = true +url = "2.5.8" +percent-encoding = "2.3.2" +base64 = "0.23.1" + +[dev-dependencies] +wiremock = "0.6.5" diff --git a/crates/utopia-server/src/api/chat.rs b/crates/utopia-server/src/api/chat.rs index 2c73d2696..f6ef8237c 100644 --- a/crates/utopia-server/src/api/chat.rs +++ b/crates/utopia-server/src/api/chat.rs @@ -71,7 +71,7 @@ fn tools_schema(can_write: bool, data_source_names: &[String]) -> serde_json::Va }, "sql": { "type": "string", - "description": "One SELECT/WITH statement (PostgreSQL dialect)." + "description": "One SELECT/WITH statement in the source's own SQL dialect (PostgreSQL, Trino, Databricks or Snowflake; the schema document names the engine)." }, "purpose": { "type": "string", @@ -435,7 +435,7 @@ pub async fn chat( }; if !ds_names.is_empty() { system_prompt.push_str(&format!( - "\nData: query_data runs read-only SQL (PostgreSQL dialect) against: {}. \ + "\nData: query_data runs read-only SQL (in each source's own dialect) against: {}. \ For questions about numbers/metrics, search for the source's schema document \ first, then query. State units and the time range you used in the answer.", ds_names.join(", ") diff --git a/crates/utopia-server/src/api/datasource_routes.rs b/crates/utopia-server/src/api/datasource_routes.rs index 9a3e9e822..8956a4ed0 100644 --- a/crates/utopia-server/src/api/datasource_routes.rs +++ b/crates/utopia-server/src/api/datasource_routes.rs @@ -53,10 +53,25 @@ pub async fn create( Json(body): Json, ) -> ApiResult> { require_admin(&user)?; + // 引擎跟着 scheme 走，界面只有一个连接串输入框；body.engine 只为兼容旧调用留着 + let engine = crate::query_engine::engine_from_conn(&body.conn_string).ok_or_else(|| { + utopia_core::AppError::invalid( + "unsupported_conn_scheme", + format!( + "Connection string must start with one of: postgres://, trino://, databricks://, snowflake:// (engines: {})", + crate::query_engine::ENGINES.join(", ") + ), + ) + })?; + let _ = &body.engine; + // 连接串的形状在登记时就校验（缺令牌、缺 warehouse……），错误信息里带写法； + // 否则要等到「测试」才知道，而那一步只回 ok:false + crate::query_engine::engine_for(engine, &body.conn_string) + .map_err(|e| utopia_core::AppError::invalid("bad_conn_string", e.to_string()))?; let id = utopia_store::datasources::create( &state.pool, &body.name, - &body.engine, + engine, &body.conn_string, user.id, ) @@ -228,7 +243,7 @@ async fn sync_schema_doc(state: &AppState, kb_id: Uuid, ds_id: Uuid) -> anyhow:: .await?; let mut md = format!( - "# Data source: {name}\n\nTables and columns available for SQL queries against this source.\n" + "# Data source: {name}\n\nEngine: {engine}. Tables and columns available for SQL queries against this source; write SQL in this engine's dialect.\n" ); let mut current = String::new(); let mut tables = 0usize; diff --git a/crates/utopia-server/src/api/tools.rs b/crates/utopia-server/src/api/tools.rs index 22c41d9f4..98a52a2c1 100644 --- a/crates/utopia-server/src/api/tools.rs +++ b/crates/utopia-server/src/api/tools.rs @@ -410,8 +410,9 @@ pub(super) fn charter_source_json(n: usize, h: &utopia_search::DocsSection) -> s /// 问数执行：安全闸（解析白名单）→ 引擎执行（只读会话 + 强制 LIMIT + 超时）→ JSON 行。 async fn run_query(state: &AppState, ds_id: Uuid, sql: &str) -> anyhow::Result { - let guarded = crate::query_engine::guard_sql(sql)?; let (engine, conn) = utopia_store::datasources::engine_and_conn(&state.pool, ds_id).await?; + // 闸门按引擎选方言：Databricks 的反引号、Snowflake 的 :: 转型都得先过得了解析 + let guarded = crate::query_engine::guard_sql_for(&engine, sql)?; let result = crate::query_engine::engine_for(&engine, &conn)? .execute(&guarded) .await?; diff --git a/crates/utopia-server/src/query_engine.rs b/crates/utopia-server/src/query_engine.rs deleted file mode 100644 index 01a015bc2..000000000 --- a/crates/utopia-server/src/query_engine.rs +++ /dev/null @@ -1,205 +0,0 @@ -//! 问数查询引擎：trait 接缝（BlobStore 同手法）+ 引擎无关的安全闸。 -//! -//! 引擎按协议族扩，不按产品名扩：postgres（本文件）→ mysql 线协议族（白捡 -//! TiDB/OceanBase/Doris/StarRocks）→ HTTP 族（ClickHouse、Trino——后者一个顶起 -//! Iceberg/Delta/Hive 整个湖仓生态）。挂载模型与注册表引擎无关，加引擎零迁移。 -//! -//! 安全闸（纵深防御，不信任模型）： -//! 1. sqlparser 解析：仅放行单条 SELECT/WITH（含 CTE），拒绝 DML/DDL/多语句/SELECT INTO -//! 2. 强制外包一层 LIMIT（cap+1 探测截断） -//! 3. 会话级只读 + 语句超时（引擎各自的机制，parser 万一漏网也写不进去） -//! 4. 结果统一为 JSON Lines（各引擎都有原生 JSON 行输出，也是模型最好消化的格式） - -use sqlparser::ast::Statement; -use sqlparser::dialect::PostgreSqlDialect; -use sqlparser::parser::Parser; -use sqlx::postgres::PgPoolOptions; -use sqlx::Row; -use std::time::Duration; - -/// 行数上限（外包 LIMIT cap+1，第 201 行只用来判断截断）。 -pub const ROW_CAP: usize = 200; -const STATEMENT_TIMEOUT_SECS: u32 = 10; - -pub struct QueryResult { - /// 每行一个 JSON 对象文本（键序 = 查询列序） - pub rows: Vec, - pub truncated: bool, -} - -#[derive(Debug)] -pub struct SchemaColumn { - pub schema: String, - pub table: String, - pub column: String, - pub data_type: String, - pub comment: Option, -} - -#[async_trait::async_trait] -pub trait QueryEngine: Send + Sync { - async fn test(&self) -> anyhow::Result<()>; - async fn fetch_schema(&self) -> anyhow::Result>; - /// 执行已过闸的 SELECT。实现自身仍需强制只读会话与超时（纵深防御）。 - async fn execute(&self, sql: &str) -> anyhow::Result; -} - -/// 引擎工厂。conn 凭据只在服务端流转。 -pub fn engine_for(engine: &str, conn: &str) -> anyhow::Result> { - match engine { - "postgres" => Ok(Box::new(PostgresEngine { - conn: conn.to_string(), - })), - other => anyhow::bail!("Unsupported engine: {other}"), - } -} - -/// 安全闸第 1 层：解析并校验，返回规整后的语句文本。 -pub fn guard_sql(sql: &str) -> anyhow::Result { - let cleaned = sql.trim().trim_end_matches(';').trim(); - if cleaned.is_empty() { - anyhow::bail!("Empty SQL"); - } - let statements = Parser::parse_sql(&PostgreSqlDialect {}, cleaned) - .map_err(|e| anyhow::anyhow!("SQL parse error: {e}"))?; - if statements.len() != 1 { - anyhow::bail!("Exactly one statement is allowed"); - } - match &statements[0] { - Statement::Query(_) => Ok(cleaned.to_string()), - other => anyhow::bail!( - "Read-only: only SELECT/WITH queries are allowed (got {})", - statement_kind(other) - ), - } -} - -fn statement_kind(s: &Statement) -> &'static str { - match s { - Statement::Insert { .. } => "INSERT", - Statement::Update { .. } => "UPDATE", - Statement::Delete { .. } => "DELETE", - Statement::Drop { .. } => "DROP", - Statement::CreateTable { .. } | Statement::CreateView { .. } => "CREATE", - Statement::AlterTable { .. } => "ALTER", - Statement::Truncate { .. } => "TRUNCATE", - Statement::Copy { .. } => "COPY", - _ => "a non-SELECT statement", - } -} - -// --------------------------------------------------------------------------- -// Postgres 族（顺带覆盖 Greenplum/Timescale 等 PG 兼容系） -// --------------------------------------------------------------------------- - -pub struct PostgresEngine { - conn: String, -} - -impl PostgresEngine { - async fn pool(&self) -> anyhow::Result { - Ok(PgPoolOptions::new() - .max_connections(1) - .acquire_timeout(Duration::from_secs(5)) - .connect(&self.conn) - .await?) - } -} - -#[async_trait::async_trait] -impl QueryEngine for PostgresEngine { - async fn test(&self) -> anyhow::Result<()> { - let pool = self.pool().await?; - sqlx::query("SELECT 1").execute(&pool).await?; - pool.close().await; - Ok(()) - } - - async fn fetch_schema(&self) -> anyhow::Result> { - let pool = self.pool().await?; - let rows: Vec<(String, String, String, String, Option)> = sqlx::query_as( - "SELECT c.table_schema, c.table_name, c.column_name, - c.data_type, pgd.description - FROM information_schema.columns c - LEFT JOIN pg_catalog.pg_statio_all_tables st - ON st.schemaname = c.table_schema AND st.relname = c.table_name - LEFT JOIN pg_catalog.pg_description pgd - ON pgd.objoid = st.relid AND pgd.objsubid = c.ordinal_position - WHERE c.table_schema NOT IN ('pg_catalog', 'information_schema') - ORDER BY c.table_schema, c.table_name, c.ordinal_position", - ) - .fetch_all(&pool) - .await?; - pool.close().await; - Ok(rows - .into_iter() - .map(|(schema, table, column, data_type, comment)| SchemaColumn { - schema, - table, - column, - data_type, - comment, - }) - .collect()) - } - - async fn execute(&self, sql: &str) -> anyhow::Result { - let pool = self.pool().await?; - // 纵深防御第 3 层：会话级只读 + 超时（parser 漏网也写不进去、跑不死库） - sqlx::query("SET default_transaction_read_only = on") - .execute(&pool) - .await?; - sqlx::query(&format!( - "SET statement_timeout = '{STATEMENT_TIMEOUT_SECS}s'" - )) - .execute(&pool) - .await?; - // 第 2 层：外包 LIMIT；row_to_json 让 PG 全权处理类型→JSON（文本键序保留列序） - let wrapped = format!( - "SELECT row_to_json(_q)::text AS _j FROM ( {sql} ) AS _q LIMIT {}", - ROW_CAP + 1 - ); - let fetched = sqlx::query(&wrapped).fetch_all(&pool).await?; - pool.close().await; - - let truncated = fetched.len() > ROW_CAP; - let rows = fetched - .into_iter() - .take(ROW_CAP) - .map(|r| r.try_get::("_j").unwrap_or_else(|_| "{}".into())) - .collect(); - Ok(QueryResult { rows, truncated }) - } -} - -#[cfg(test)] -mod tests { - use super::guard_sql; - - #[test] - fn allows_select_and_cte() { - assert!(guard_sql("SELECT region, sum(amount) FROM orders GROUP BY 1").is_ok()); - assert!(guard_sql("WITH t AS (SELECT 1 AS x) SELECT * FROM t;").is_ok()); - } - - #[test] - fn rejects_writes_and_ddl() { - for bad in [ - "UPDATE orders SET amount = 0", - "DELETE FROM orders", - "INSERT INTO orders (region) VALUES ('east')", - "DROP TABLE orders", - "TRUNCATE orders", - "CREATE TABLE t (id int)", - "ALTER TABLE orders ADD COLUMN x int", - ] { - assert!(guard_sql(bad).is_err(), "should reject: {bad}"); - } - } - - #[test] - fn rejects_multi_statement() { - assert!(guard_sql("SELECT 1; DROP TABLE orders").is_err()); - assert!(guard_sql("").is_err()); - } -} diff --git a/crates/utopia-server/src/query_engine/conn.rs b/crates/utopia-server/src/query_engine/conn.rs new file mode 100644 index 000000000..1d3cd5192 --- /dev/null +++ b/crates/utopia-server/src/query_engine/conn.rs @@ -0,0 +1,253 @@ +//! 连接串解析。一个输入框、四种 scheme；这里把 URL 拆成各引擎要的字段。 +//! +//! 写法沿用 `postgres://user:pass@host/db` 的形状：凭据在 userinfo 里，HTTP 族的 +//! 令牌放 password 位（`databricks://:TOKEN@…`），路径是「目录 / 库 / schema」， +//! 引擎特有的开关走 query。`ssl=false` 让 HTTP 族走明文——给本地代理与测试用， +//! 线上的三家都只认 https。 + +use percent_encoding::percent_decode_str; +use url::Url; + +fn decode(s: &str) -> String { + percent_decode_str(s).decode_utf8_lossy().into_owned() +} + +fn query(u: &Url, key: &str) -> Option { + u.query_pairs() + .find(|(k, _)| k == key) + .map(|(_, v)| v.into_owned()) + .filter(|v| !v.is_empty()) +} + +fn ssl_off(u: &Url) -> bool { + matches!( + query(u, "ssl").as_deref(), + Some("false") | Some("0") | Some("off") | Some("no") + ) +} + +fn segments(u: &Url) -> Vec { + u.path_segments() + .map(|s| s.filter(|x| !x.is_empty()).map(decode).collect()) + .unwrap_or_default() +} + +/// 令牌：password 位优先；没有 password 时 username 位也算（`databricks://TOKEN@host` +/// 少打一个冒号是最常见的手滑）；最后看 `?token=` +fn token_of(u: &Url) -> Option { + u.password() + .map(decode) + .filter(|s| !s.is_empty()) + .or_else(|| Some(decode(u.username())).filter(|s| !s.is_empty())) + .or_else(|| query(u, "token")) +} + +fn base_of(u: &Url, https: bool, default_port: u16) -> anyhow::Result { + let host = u + .host_str() + .ok_or_else(|| anyhow::anyhow!("{}://: a host is required", u.scheme()))?; + let port = u.port().unwrap_or(default_port); + let scheme = if https { "https" } else { "http" }; + // 默认端口不写进 URL：reqwest 照样能连，日志里也干净 + let explicit = match (https, port) { + (true, 443) | (false, 80) => String::new(), + _ => format!(":{port}"), + }; + Ok(format!("{scheme}://{host}{explicit}")) +} + +/// `trino://user[:password]@host[:port]/[catalog[/schema]][?ssl=true|false]` +/// +/// 明文 http 是 Trino 的默认（8080）；带密码、`ssl=true`、或端口 443 / 8443 时走 https—— +/// Trino 自己也拒绝在明文上收密码。`presto://` 是同一个协议的旧名。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TrinoConn { + pub base: String, + pub user: String, + pub password: Option, + pub catalog: Option, + pub schema: Option, +} + +impl TrinoConn { + pub fn parse(conn: &str) -> anyhow::Result { + let u = Url::parse(conn.trim())?; + let user = decode(u.username()); + if user.is_empty() { + anyhow::bail!("trino://: a user is required (it becomes X-Trino-User), e.g. trino://alice@host:8080/hive/default"); + } + let password = u.password().map(decode).filter(|s| !s.is_empty()); + let https = !ssl_off(&u) + && (password.is_some() + || query(&u, "ssl").as_deref() == Some("true") + || matches!(u.port(), Some(443) | Some(8443))); + let base = base_of(&u, https, if https { 443 } else { 8080 })?; + let segs = segments(&u); + Ok(Self { + base, + user, + password, + catalog: segs.first().cloned(), + schema: segs.get(1).cloned(), + }) + } +} + +/// `databricks://:TOKEN@workspace-host/sql/1.0/warehouses/WAREHOUSE_ID[?catalog=main&schema=default]` +/// +/// 路径就是 JDBC 里的 httpPath，从控制台复制过来不用改；`?warehouse=ID` 也认。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DatabricksConn { + pub base: String, + pub token: String, + pub warehouse_id: String, + pub catalog: Option, + pub schema: Option, +} + +impl DatabricksConn { + pub fn parse(conn: &str) -> anyhow::Result { + let u = Url::parse(conn.trim())?; + let base = base_of(&u, !ssl_off(&u), if ssl_off(&u) { 80 } else { 443 })?; + let token = token_of(&u).ok_or_else(|| { + anyhow::anyhow!("databricks://: a personal access token is required, e.g. databricks://:TOKEN@host/sql/1.0/warehouses/ID") + })?; + let segs = segments(&u); + let from_path = segs + .iter() + .position(|s| s == "warehouses") + .and_then(|i| segs.get(i + 1).cloned()); + let warehouse_id = query(&u, "warehouse").or(from_path).ok_or_else(|| { + anyhow::anyhow!("databricks://: a SQL warehouse is required — the /sql/1.0/warehouses/ID path or ?warehouse=ID") + })?; + Ok(Self { + base, + token, + warehouse_id, + catalog: query(&u, "catalog"), + schema: query(&u, "schema"), + }) + } +} + +/// `snowflake://:TOKEN@account.snowflakecomputing.com/[DATABASE[/SCHEMA]][?warehouse=WH&role=R&token_type=pat|oauth]` +/// +/// SQL API 不收密码，只收令牌：默认当作 programmatic access token，`token_type=oauth` +/// 换成 OAuth 令牌。密钥对 JWT 要本地签名，这一版不做。 +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SnowflakeConn { + pub base: String, + pub token: String, + /// `X-Snowflake-Authorization-Token-Type` 的值 + pub token_type: &'static str, + pub database: Option, + pub schema: Option, + pub warehouse: Option, + pub role: Option, +} + +impl SnowflakeConn { + pub fn parse(conn: &str) -> anyhow::Result { + let u = Url::parse(conn.trim())?; + let base = base_of(&u, !ssl_off(&u), if ssl_off(&u) { 80 } else { 443 })?; + let token = token_of(&u).ok_or_else(|| { + anyhow::anyhow!("snowflake://: a programmatic access token or OAuth token is required, e.g. snowflake://:TOKEN@account.snowflakecomputing.com/DB/SCHEMA?warehouse=WH") + })?; + let token_type = match query(&u, "token_type") + .as_deref() + .map(str::to_ascii_lowercase) + .as_deref() + { + None | Some("pat") | Some("programmatic_access_token") => "PROGRAMMATIC_ACCESS_TOKEN", + Some("oauth") => "OAUTH", + Some(other) => { + anyhow::bail!("snowflake://: unknown token_type '{other}' (pat or oauth)") + } + }; + let segs = segments(&u); + Ok(Self { + base, + token, + token_type, + database: segs.first().cloned(), + schema: segs.get(1).cloned(), + warehouse: query(&u, "warehouse"), + role: query(&u, "role"), + }) + } +} + +#[cfg(test)] +mod tests { + use super::{DatabricksConn, SnowflakeConn, TrinoConn}; + + #[test] + fn trino_defaults_to_plain_http_and_upgrades_when_it_must() { + let c = TrinoConn::parse("trino://alice@lake.internal:8080/iceberg/sales").unwrap(); + assert_eq!(c.base, "http://lake.internal:8080"); + assert_eq!(c.user, "alice"); + assert_eq!(c.catalog.as_deref(), Some("iceberg")); + assert_eq!(c.schema.as_deref(), Some("sales")); + + let c = TrinoConn::parse("trino://alice:s%40cret@lake.internal/hive").unwrap(); + assert_eq!(c.base, "https://lake.internal"); + assert_eq!(c.password.as_deref(), Some("s@cret")); + + let c = TrinoConn::parse("trino://alice@lake.internal:8443/hive").unwrap(); + assert_eq!(c.base, "https://lake.internal:8443"); + + let c = TrinoConn::parse("presto://bob@127.0.0.1:9000?ssl=false").unwrap(); + assert_eq!(c.base, "http://127.0.0.1:9000"); + assert_eq!(c.catalog, None); + + assert!(TrinoConn::parse("trino://lake.internal/hive").is_err()); + } + + #[test] + fn databricks_reads_the_http_path_from_the_console() { + let c = DatabricksConn::parse( + "databricks://:dapi123@dbc-abc.cloud.databricks.com/sql/1.0/warehouses/9f2a?catalog=main&schema=sales", + ) + .unwrap(); + assert_eq!(c.base, "https://dbc-abc.cloud.databricks.com"); + assert_eq!(c.token, "dapi123"); + assert_eq!(c.warehouse_id, "9f2a"); + assert_eq!(c.catalog.as_deref(), Some("main")); + assert_eq!(c.schema.as_deref(), Some("sales")); + + let c = DatabricksConn::parse("databricks://dapi123@host?warehouse=w1").unwrap(); + assert_eq!(c.token, "dapi123"); + assert_eq!(c.warehouse_id, "w1"); + + assert!(DatabricksConn::parse("databricks://host/sql/1.0/warehouses/w1").is_err()); + assert!(DatabricksConn::parse("databricks://:t@host").is_err()); + } + + #[test] + fn snowflake_takes_a_token_and_the_session_knobs() { + let c = SnowflakeConn::parse( + "snowflake://:tok@xy12345.eu-central-1.snowflakecomputing.com/ANALYTICS/PUBLIC?warehouse=WH&role=ANALYST", + ) + .unwrap(); + assert_eq!( + c.base, + "https://xy12345.eu-central-1.snowflakecomputing.com" + ); + assert_eq!(c.token, "tok"); + assert_eq!(c.token_type, "PROGRAMMATIC_ACCESS_TOKEN"); + assert_eq!(c.database.as_deref(), Some("ANALYTICS")); + assert_eq!(c.schema.as_deref(), Some("PUBLIC")); + assert_eq!(c.warehouse.as_deref(), Some("WH")); + assert_eq!(c.role.as_deref(), Some("ANALYST")); + + let c = + SnowflakeConn::parse("snowflake://:tok@acct.snowflakecomputing.com?token_type=oauth") + .unwrap(); + assert_eq!(c.token_type, "OAUTH"); + assert!(SnowflakeConn::parse( + "snowflake://:tok@acct.snowflakecomputing.com?token_type=jwt" + ) + .is_err()); + assert!(SnowflakeConn::parse("snowflake://acct.snowflakecomputing.com/DB").is_err()); + } +} diff --git a/crates/utopia-server/src/query_engine/databricks.rs b/crates/utopia-server/src/query_engine/databricks.rs new file mode 100644 index 000000000..7f9eb3793 --- /dev/null +++ b/crates/utopia-server/src/query_engine/databricks.rs @@ -0,0 +1,271 @@ +//! Databricks SQL Statement Execution API（`/api/2.0/sql/statements`）。 +//! 一个 SQL warehouse 后面是 Unity Catalog 的整个湖仓（Delta 为主）， +//! 令牌是 personal access token。结果要 INLINE + JSON_ARRAY：值全是字符串， +//! 按 manifest 里的列类型还原成数与布尔。 + +use super::conn::DatabricksConn; +use super::{ + coerce, rows_to_json_lines, sql_literal, truncate_rows, wrap_limit, QueryEngine, QueryResult, + SchemaColumn, HTTP_POLL_BUDGET, ROW_CAP, +}; +use serde::Deserialize; +use serde_json::json; +use std::time::{Duration, Instant}; + +pub struct DatabricksEngine { + conn: DatabricksConn, +} + +#[derive(Deserialize)] +struct StatementResponse { + statement_id: Option, + status: Status, + manifest: Option, + result: Option, +} + +#[derive(Deserialize)] +struct Status { + state: String, + error: Option, +} + +#[derive(Deserialize)] +struct StatusError { + message: Option, + error_code: Option, +} + +#[derive(Deserialize)] +struct Manifest { + schema: Option, +} + +#[derive(Deserialize)] +struct Schema { + columns: Vec, +} + +#[derive(Deserialize)] +struct ColumnInfo { + name: String, + type_text: Option, +} + +#[derive(Deserialize)] +struct ResultData { + data_array: Option>>, +} + +impl DatabricksEngine { + pub fn new(conn: DatabricksConn) -> Self { + Self { conn } + } + + async fn run(&self, sql: &str) -> anyhow::Result<(Vec, Vec>)> { + let client = super::http()?; + let mut body = json!({ + "warehouse_id": self.conn.warehouse_id, + "statement": sql, + "wait_timeout": "30s", + "on_wait_timeout": "CONTINUE", + "disposition": "INLINE", + "format": "JSON_ARRAY", + "row_limit": ROW_CAP + 1, + }); + if let Some(c) = &self.conn.catalog { + body["catalog"] = json!(c); + } + if let Some(s) = &self.conn.schema { + body["schema"] = json!(s); + } + let mut resp: StatementResponse = client + .post(format!("{}/api/2.0/sql/statements", self.conn.base)) + .bearer_auth(&self.conn.token) + .json(&body) + .send() + .await? + .error_for_status()? + .json() + .await?; + let started = Instant::now(); + loop { + match resp.status.state.as_str() { + "SUCCEEDED" => break, + "PENDING" | "RUNNING" => { + let id = resp + .statement_id + .clone() + .ok_or_else(|| anyhow::anyhow!("Databricks returned no statement_id"))?; + if started.elapsed() > HTTP_POLL_BUDGET { + anyhow::bail!( + "Databricks statement did not finish within {}s", + HTTP_POLL_BUDGET.as_secs() + ); + } + tokio::time::sleep(Duration::from_secs(1)).await; + resp = client + .get(format!("{}/api/2.0/sql/statements/{id}", self.conn.base)) + .bearer_auth(&self.conn.token) + .send() + .await? + .error_for_status()? + .json() + .await?; + } + other => { + let e = resp.status.error.as_ref(); + let code = e + .and_then(|e| e.error_code.clone()) + .map(|c| format!("{c}: ")) + .unwrap_or_default(); + let msg = e + .and_then(|e| e.message.clone()) + .unwrap_or_else(|| format!("statement ended in state {other}")); + anyhow::bail!("{code}{msg}"); + } + } + } + let columns: Vec = resp + .manifest + .and_then(|m| m.schema) + .map(|s| s.columns) + .unwrap_or_default(); + let raw_rows = resp.result.and_then(|r| r.data_array).unwrap_or_default(); + let rows = raw_rows + .into_iter() + .map(|row| { + row.iter() + .enumerate() + .map(|(i, v)| { + let ty = columns + .get(i) + .and_then(|c| c.type_text.as_deref()) + .unwrap_or(""); + coerce(ty, v) + }) + .collect() + }) + .collect(); + Ok((columns.into_iter().map(|c| c.name).collect(), rows)) + } +} + +#[async_trait::async_trait] +impl QueryEngine for DatabricksEngine { + async fn test(&self) -> anyhow::Result<()> { + self.run("SELECT 1").await.map(|_| ()) + } + + async fn fetch_schema(&self) -> anyhow::Result> { + // 带 catalog 就查那个 catalog 的 information_schema；不带就是会话默认的 + let prefix = self + .conn + .catalog + .as_deref() + .map(|c| format!("`{}`.", c.replace('`', "``"))) + .unwrap_or_default(); + let schema_filter = self + .conn + .schema + .as_deref() + .map(|s| format!(" AND table_schema = {}", sql_literal(s))) + .unwrap_or_default(); + let sql = format!( + "SELECT table_schema, table_name, column_name, data_type, comment \ + FROM {prefix}information_schema.columns \ + WHERE table_schema <> 'information_schema'{schema_filter} \ + ORDER BY table_schema, table_name, ordinal_position" + ); + let (_, rows) = self.run(&sql).await?; + Ok(rows.into_iter().map(super::trino::schema_row).collect()) + } + + async fn execute(&self, sql: &str) -> anyhow::Result { + let (columns, rows) = self.run(&wrap_limit(sql)).await?; + let (rows, truncated) = truncate_rows(rows); + Ok(QueryResult { + rows: rows_to_json_lines(&columns, &rows), + truncated, + }) + } +} + +#[cfg(test)] +mod tests { + use super::super::conn::DatabricksConn; + use super::super::QueryEngine; + use super::DatabricksEngine; + use serde_json::json; + use wiremock::matchers::{header, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + fn conn(server: &MockServer) -> DatabricksConn { + DatabricksConn::parse(&format!( + "databricks://:dapi-test@{}/sql/1.0/warehouses/wh1?catalog=main&ssl=false", + server.uri().trim_start_matches("http://") + )) + .unwrap() + } + + #[tokio::test] + async fn polls_until_succeeded_and_restores_types() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/2.0/sql/statements")) + .and(header("authorization", "Bearer dapi-test")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "statement_id": "s1", + "status": { "state": "PENDING" } + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/api/2.0/sql/statements/s1")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "statement_id": "s1", + "status": { "state": "SUCCEEDED" }, + "manifest": { "schema": { "columns": [ + { "name": "region", "type_text": "STRING", "position": 0 }, + { "name": "total", "type_text": "DECIMAL(12,2)", "position": 1 }, + { "name": "active", "type_text": "BOOLEAN", "position": 2 } + ] } }, + "result": { "data_array": [ ["east", "12.50", "true"], ["west", null, "false"] ] } + }))) + .expect(1) + .mount(&server) + .await; + + let out = DatabricksEngine::new(conn(&server)) + .execute("SELECT region, total, active FROM orders") + .await + .unwrap(); + assert_eq!( + out.rows, + vec![ + r#"{"region":"east","total":12.5,"active":true}"#, + r#"{"region":"west","total":null,"active":false}"# + ] + ); + } + + #[tokio::test] + async fn a_failed_statement_reports_the_message() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/2.0/sql/statements")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "statement_id": "s2", + "status": { "state": "FAILED", "error": { "error_code": "BAD_REQUEST", "message": "TABLE_OR_VIEW_NOT_FOUND: nope" } } + }))) + .mount(&server) + .await; + let err = DatabricksEngine::new(conn(&server)) + .execute("SELECT * FROM nope") + .await + .unwrap_err() + .to_string(); + assert!(err.contains("TABLE_OR_VIEW_NOT_FOUND"), "{err}"); + } +} diff --git a/crates/utopia-server/src/query_engine/mod.rs b/crates/utopia-server/src/query_engine/mod.rs new file mode 100644 index 000000000..2bd171322 --- /dev/null +++ b/crates/utopia-server/src/query_engine/mod.rs @@ -0,0 +1,349 @@ +//! 问数查询引擎：trait 接缝（BlobStore 同手法）+ 引擎无关的安全闸。 +//! +//! 引擎按协议族扩，不按产品名扩：postgres 线协议（`postgres.rs`）→ HTTP 族—— +//! `trino.rs` 一个顶起 Iceberg / Delta / Hive 整个湖仓生态，`databricks.rs`、 +//! `snowflake.rs` 各走自家的 SQL REST API。挂载模型与注册表引擎无关，加引擎只放宽 +//! 一条 CHECK。连接串是唯一的输入：引擎由 scheme 决定（[`engine_from_conn`]）， +//! 剩下的部分各引擎自己拆（`conn.rs`），凭据只在服务端流转。 +//! +//! 安全闸（纵深防御，不信任模型）： +//! 1. sqlparser 解析：仅放行单条 SELECT/WITH（含 CTE），拒绝 DML/DDL/多语句/SELECT INTO。 +//! 按引擎选方言；sqlparser 没有 Trino 方言，Generic 是它的超集 +//! 2. 强制外包一层 LIMIT（cap+1 探测截断） +//! 3. 会话级只读 + 语句超时（引擎各自的机制，parser 万一漏网也写不进去）。 +//! HTTP 族没有会话，只有语句超时——只读靠第 1 层，这是它们比线协议少的那一层 +//! 4. 结果统一为 JSON Lines：PG 让库自己转；HTTP 族拿到列名与值后在这里拼，列序保留 + +mod conn; +mod databricks; +mod postgres; +mod snowflake; +mod trino; + +use sqlparser::ast::Statement; +use sqlparser::dialect::{DatabricksDialect, GenericDialect, PostgreSqlDialect, SnowflakeDialect}; +use sqlparser::parser::Parser; +use std::time::Duration; + +/// 行数上限（外包 LIMIT cap+1，第 201 行只用来判断截断）。 +pub const ROW_CAP: usize = 200; +pub(crate) const STATEMENT_TIMEOUT_SECS: u32 = 10; +/// HTTP 族：单次请求的超时，与整条语句从提交到拿完结果的轮询预算 +pub(crate) const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(20); +pub(crate) const HTTP_POLL_BUDGET: Duration = Duration::from_secs(30); + +/// 注册表里 `engine` 列的取值。迁移里的 CHECK 与这张表要一致 +pub const ENGINES: &[&str] = &["postgres", "trino", "databricks", "snowflake"]; + +#[derive(Debug)] +pub struct QueryResult { + /// 每行一个 JSON 对象文本（键序 = 查询列序） + pub rows: Vec, + pub truncated: bool, +} + +#[derive(Debug)] +pub struct SchemaColumn { + pub schema: String, + pub table: String, + pub column: String, + pub data_type: String, + pub comment: Option, +} + +#[async_trait::async_trait] +pub trait QueryEngine: Send + Sync { + async fn test(&self) -> anyhow::Result<()>; + async fn fetch_schema(&self) -> anyhow::Result>; + /// 执行已过闸的 SELECT。实现自身仍需强制只读会话与超时（纵深防御）。 + async fn execute(&self, sql: &str) -> anyhow::Result; +} + +/// scheme → 引擎名。界面只有一个连接串输入框，这里是它唯一的分派点。 +pub fn engine_from_conn(conn: &str) -> Option<&'static str> { + let scheme = conn.trim().split("://").next()?.to_ascii_lowercase(); + match scheme.as_str() { + "postgres" | "postgresql" => Some("postgres"), + "trino" | "presto" => Some("trino"), + "databricks" => Some("databricks"), + "snowflake" => Some("snowflake"), + _ => None, + } +} + +/// 引擎工厂。conn 凭据只在服务端流转。 +pub fn engine_for(engine: &str, conn: &str) -> anyhow::Result> { + match engine { + "postgres" => Ok(Box::new(postgres::PostgresEngine::new(conn))), + "trino" => Ok(Box::new(trino::TrinoEngine::new(conn::TrinoConn::parse( + conn, + )?))), + "databricks" => Ok(Box::new(databricks::DatabricksEngine::new( + conn::DatabricksConn::parse(conn)?, + ))), + "snowflake" => Ok(Box::new(snowflake::SnowflakeEngine::new( + conn::SnowflakeConn::parse(conn)?, + ))), + other => anyhow::bail!("Unsupported engine: {other}"), + } +} + +/// 安全闸第 1 层：按引擎方言解析并校验，返回规整后的语句文本。 +pub fn guard_sql_for(engine: &str, sql: &str) -> anyhow::Result { + let cleaned = sql.trim().trim_end_matches(';').trim(); + if cleaned.is_empty() { + anyhow::bail!("Empty SQL"); + } + let parsed = match engine { + "databricks" => Parser::parse_sql(&DatabricksDialect {}, cleaned), + "snowflake" => Parser::parse_sql(&SnowflakeDialect {}, cleaned), + "trino" => Parser::parse_sql(&GenericDialect {}, cleaned), + _ => Parser::parse_sql(&PostgreSqlDialect {}, cleaned), + }; + let statements = parsed.map_err(|e| anyhow::anyhow!("SQL parse error: {e}"))?; + if statements.len() != 1 { + anyhow::bail!("Exactly one statement is allowed"); + } + match &statements[0] { + Statement::Query(_) => Ok(cleaned.to_string()), + other => anyhow::bail!( + "Read-only: only SELECT/WITH queries are allowed (got {})", + statement_kind(other) + ), + } +} + +fn statement_kind(s: &Statement) -> &'static str { + match s { + Statement::Insert { .. } => "INSERT", + Statement::Update { .. } => "UPDATE", + Statement::Delete { .. } => "DELETE", + Statement::CreateTable { .. } => "CREATE TABLE", + Statement::Drop { .. } => "DROP", + Statement::AlterTable { .. } => "ALTER TABLE", + Statement::Truncate { .. } => "TRUNCATE", + _ => "a non-SELECT statement", + } +} + +/// 第 2 层：外包一层 LIMIT。三个 HTTP 引擎都认这个写法；PG 有自己的 row_to_json 版本 +pub(crate) fn wrap_limit(sql: &str) -> String { + format!("SELECT * FROM ( {sql} ) AS _q LIMIT {}", ROW_CAP + 1) +} + +/// 第 201 行只用来判断截断，不交给模型 +pub(crate) fn truncate_rows(mut rows: Vec) -> (Vec, bool) { + let truncated = rows.len() > ROW_CAP; + rows.truncate(ROW_CAP); + (rows, truncated) +} + +/// HTTP 族共用：「列名 + 行值」拼成 JSON Lines。手拼而不是 `serde_json::Map`， +/// 后者不开 `preserve_order` 就按键排序，而列序是查询写下的顺序，模型读表靠它 +pub(crate) fn rows_to_json_lines( + columns: &[String], + rows: &[Vec], +) -> Vec { + rows.iter() + .map(|row| { + let mut line = String::from("{"); + for (i, col) in columns.iter().enumerate() { + if i > 0 { + line.push(','); + } + line.push_str(&serde_json::to_string(col).unwrap_or_else(|_| "\"?\"".into())); + line.push(':'); + let value = row.get(i).cloned().unwrap_or(serde_json::Value::Null); + line.push_str(&value.to_string()); + } + line.push('}'); + line + }) + .collect() +} + +/// Databricks 的 JSON_ARRAY 与 Snowflake 的 data 把每个值都给成字符串（或 null）。 +/// 按列类型把数与布尔还原，其余留字符串——模型对 `"42"` 和 `42` 的算术不一样 +pub(crate) fn coerce(type_name: &str, raw: &serde_json::Value) -> serde_json::Value { + let serde_json::Value::String(s) = raw else { + return raw.clone(); + }; + let ty = type_name.to_ascii_uppercase(); + const NUMERIC: &[&str] = &[ + "INT", "LONG", "SHORT", "BYTE", "FLOAT", "DOUBLE", "DECIMAL", "NUMBER", "FIXED", "REAL", + "NUMERIC", + ]; + // INTERVAL 也含 "INT"：解析不成数就原样留下，不会误伤 + if NUMERIC.iter().any(|k| ty.contains(k)) { + if let Ok(n) = s.parse::() { + return n.into(); + } + if let Ok(f) = s.parse::() { + if let Some(n) = serde_json::Number::from_f64(f) { + return serde_json::Value::Number(n); + } + } + } + if ty.starts_with("BOOL") { + match s.as_str() { + "true" | "TRUE" => return true.into(), + "false" | "FALSE" => return false.into(), + _ => {} + } + } + raw.clone() +} + +/// 单引号字面量的转义：schema 名进 information_schema 的 WHERE 子句 +pub(crate) fn sql_literal(s: &str) -> String { + format!("'{}'", s.replace('\'', "''")) +} + +/// HTTP 族共用的客户端。 +/// +/// **代理策略是显式的**：回环地址与 `NO_PROXY` 里的主机直连，其余按 `HTTPS_PROXY` / +/// `HTTP_PROXY` / `ALL_PROXY` 走。不用 reqwest 的系统代理探测——Windows 上它读注册表， +/// 而注册表里 `127.*` 这种绕过写法它认不全，本机的替身服务会被送进代理拿回 502。 +/// 服务进程该看环境变量，这条规矩与 docker-compose 里的写法一致 +pub(crate) fn http() -> anyhow::Result { + Ok(reqwest::Client::builder() + .timeout(HTTP_REQUEST_TIMEOUT) + .user_agent("utopia") + .proxy(reqwest::Proxy::custom(|url: &reqwest::Url| proxy_for(url))) + .build()?) +} + +fn proxy_for(url: &reqwest::Url) -> Option { + let host = url.host_str()?; + let loopback = host.eq_ignore_ascii_case("localhost") + || host + .trim_matches(|c| c == '[' || c == ']') + .parse::() + .map(|ip| ip.is_loopback()) + .unwrap_or(false); + if loopback || no_proxy_matches(host) { + return None; + } + let keys: &[&str] = if url.scheme() == "https" { + &["HTTPS_PROXY", "https_proxy", "ALL_PROXY", "all_proxy"] + } else { + &["HTTP_PROXY", "http_proxy", "ALL_PROXY", "all_proxy"] + }; + keys.iter() + .find_map(|k| std::env::var(k).ok()) + .filter(|v| !v.trim().is_empty()) + .and_then(|v| reqwest::Url::parse(v.trim()).ok()) +} + +/// `NO_PROXY=localhost,127.0.0.1,.internal,corp.example` 的常见写法：整名相等， +/// 或者以点开头的后缀匹配 +fn no_proxy_matches(host: &str) -> bool { + let raw = std::env::var("NO_PROXY") + .or_else(|_| std::env::var("no_proxy")) + .unwrap_or_default(); + raw.split(',') + .map(str::trim) + .filter(|p| !p.is_empty() && *p != "*") + .any(|p| { + let p = p.trim_start_matches('.'); + host.eq_ignore_ascii_case(p) + || host + .to_ascii_lowercase() + .ends_with(&format!(".{}", p.to_ascii_lowercase())) + }) + || raw.split(',').any(|p| p.trim() == "*") +} + +#[cfg(test)] +mod tests { + use super::{coerce, engine_from_conn, guard_sql_for, rows_to_json_lines}; + use serde_json::json; + + fn guard_sql(sql: &str) -> anyhow::Result { + guard_sql_for("postgres", sql) + } + + #[test] + fn allows_select_and_cte() { + assert!(guard_sql("SELECT region, sum(amount) FROM orders GROUP BY 1").is_ok()); + assert!(guard_sql("WITH t AS (SELECT 1 AS x) SELECT * FROM t;").is_ok()); + } + + #[test] + fn rejects_writes_and_ddl() { + for bad in [ + "UPDATE orders SET amount = 0", + "DELETE FROM orders", + "INSERT INTO orders (region) VALUES ('east')", + "DROP TABLE orders", + "TRUNCATE orders", + "CREATE TABLE t (id int)", + "ALTER TABLE orders ADD COLUMN x int", + ] { + assert!(guard_sql(bad).is_err(), "should reject: {bad}"); + } + } + + #[test] + fn rejects_multi_statement() { + assert!(guard_sql("SELECT 1; DROP TABLE orders").is_err()); + assert!(guard_sql("").is_err()); + } + + #[test] + fn every_dialect_keeps_the_same_gate() { + for engine in ["postgres", "trino", "databricks", "snowflake"] { + assert!( + guard_sql_for(engine, "SELECT a FROM t WHERE b > 1").is_ok(), + "{engine}" + ); + assert!(guard_sql_for(engine, "DELETE FROM t").is_err(), "{engine}"); + assert!( + guard_sql_for(engine, "SELECT 1; SELECT 2").is_err(), + "{engine}" + ); + } + // 各家的方言细节：反引号、双冒号转型，都要过得去 + assert!(guard_sql_for("databricks", "SELECT `region` FROM main.sales.orders").is_ok()); + assert!(guard_sql_for("snowflake", "SELECT amount::number FROM db.public.orders").is_ok()); + assert!(guard_sql_for("trino", "SELECT count(*) FROM hive.default.orders").is_ok()); + } + + #[test] + fn engine_follows_the_scheme() { + assert_eq!(engine_from_conn("postgres://u:p@h/db"), Some("postgres")); + assert_eq!(engine_from_conn("postgresql://u:p@h/db"), Some("postgres")); + assert_eq!(engine_from_conn("trino://u@h:8443/hive"), Some("trino")); + assert_eq!(engine_from_conn("presto://u@h/hive"), Some("trino")); + assert_eq!( + engine_from_conn("databricks://:t@h/sql/1.0/warehouses/x"), + Some("databricks") + ); + assert_eq!( + engine_from_conn("snowflake://:t@a.snowflakecomputing.com/db"), + Some("snowflake") + ); + assert_eq!(engine_from_conn("mysql://u@h/db"), None); + assert_eq!(engine_from_conn("garbage"), None); + } + + #[test] + fn json_lines_keep_column_order() { + let cols = vec!["zeta".to_string(), "alpha".to_string()]; + let rows = vec![vec![json!(1), json!("x")], vec![json!(null)]]; + assert_eq!( + rows_to_json_lines(&cols, &rows), + vec![r#"{"zeta":1,"alpha":"x"}"#, r#"{"zeta":null,"alpha":null}"#] + ); + } + + #[test] + fn strings_come_back_as_numbers_when_the_column_says_so() { + assert_eq!(coerce("DOUBLE", &json!("12.5")), json!(12.5)); + assert_eq!(coerce("fixed", &json!("42")), json!(42)); + assert_eq!(coerce("BOOLEAN", &json!("true")), json!(true)); + assert_eq!(coerce("STRING", &json!("42")), json!("42")); + assert_eq!(coerce("INTERVAL", &json!("1 day")), json!("1 day")); + assert_eq!(coerce("DOUBLE", &json!(null)), json!(null)); + } +} diff --git a/crates/utopia-server/src/query_engine/postgres.rs b/crates/utopia-server/src/query_engine/postgres.rs new file mode 100644 index 000000000..5bba3a4ef --- /dev/null +++ b/crates/utopia-server/src/query_engine/postgres.rs @@ -0,0 +1,93 @@ +//! Postgres 族（顺带覆盖 Greenplum / Timescale 等 PG 兼容系）。线协议直连， +//! 是四个引擎里唯一有会话可设只读的那个。 + +use super::{QueryEngine, QueryResult, SchemaColumn, ROW_CAP, STATEMENT_TIMEOUT_SECS}; +use sqlx::postgres::PgPoolOptions; +use sqlx::Row; +use std::time::Duration; + +pub struct PostgresEngine { + conn: String, +} + +impl PostgresEngine { + pub fn new(conn: &str) -> Self { + Self { + conn: conn.to_string(), + } + } + + async fn pool(&self) -> anyhow::Result { + Ok(PgPoolOptions::new() + .max_connections(1) + .acquire_timeout(Duration::from_secs(5)) + .connect(&self.conn) + .await?) + } +} + +#[async_trait::async_trait] +impl QueryEngine for PostgresEngine { + async fn test(&self) -> anyhow::Result<()> { + let pool = self.pool().await?; + sqlx::query("SELECT 1").execute(&pool).await?; + pool.close().await; + Ok(()) + } + + async fn fetch_schema(&self) -> anyhow::Result> { + let pool = self.pool().await?; + let rows: Vec<(String, String, String, String, Option)> = sqlx::query_as( + "SELECT c.table_schema, c.table_name, c.column_name, + c.data_type, pgd.description + FROM information_schema.columns c + LEFT JOIN pg_catalog.pg_statio_all_tables st + ON st.schemaname = c.table_schema AND st.relname = c.table_name + LEFT JOIN pg_catalog.pg_description pgd + ON pgd.objoid = st.relid AND pgd.objsubid = c.ordinal_position + WHERE c.table_schema NOT IN ('pg_catalog', 'information_schema') + ORDER BY c.table_schema, c.table_name, c.ordinal_position", + ) + .fetch_all(&pool) + .await?; + pool.close().await; + Ok(rows + .into_iter() + .map(|(schema, table, column, data_type, comment)| SchemaColumn { + schema, + table, + column, + data_type, + comment, + }) + .collect()) + } + + async fn execute(&self, sql: &str) -> anyhow::Result { + let pool = self.pool().await?; + // 纵深防御第 3 层：会话级只读 + 超时（parser 漏网也写不进去、跑不死库） + sqlx::query("SET default_transaction_read_only = on") + .execute(&pool) + .await?; + sqlx::query(&format!( + "SET statement_timeout = '{STATEMENT_TIMEOUT_SECS}s'" + )) + .execute(&pool) + .await?; + // 第 2 层：外包 LIMIT；row_to_json 让 PG 全权处理类型→JSON（文本键序保留列序） + let wrapped = format!( + "SELECT row_to_json(_q)::text AS _j FROM ( {sql} ) AS _q LIMIT {}", + ROW_CAP + 1 + ); + let fetched = sqlx::query(&wrapped).fetch_all(&pool).await?; + pool.close().await; + + let truncated = fetched.len() > ROW_CAP; + let rows = fetched + .into_iter() + .take(ROW_CAP) + .map(|r| r.try_get::("_j").unwrap_or_else(|_| "{}".into())) + .collect(); + Ok(QueryResult { rows, truncated }) + } +} diff --git a/crates/utopia-server/src/query_engine/snowflake.rs b/crates/utopia-server/src/query_engine/snowflake.rs new file mode 100644 index 000000000..ff8106117 --- /dev/null +++ b/crates/utopia-server/src/query_engine/snowflake.rs @@ -0,0 +1,257 @@ +//! Snowflake SQL API v2（`/api/v2/statements`）。同步提交（`async=false`）拿不完的 +//! 语句回 202，沿 statementHandle 轮询。值全是字符串，按 rowType 还原数与布尔。 +//! +//! 只收令牌，不收密码：programmatic access token 或 OAuth。密钥对 JWT 要本地签名， +//! 这一版不做——见 `conn.rs`。 + +use super::conn::SnowflakeConn; +use super::{ + coerce, rows_to_json_lines, sql_literal, truncate_rows, wrap_limit, QueryEngine, QueryResult, + SchemaColumn, HTTP_POLL_BUDGET, STATEMENT_TIMEOUT_SECS, +}; +use reqwest::StatusCode; +use serde::Deserialize; +use serde_json::json; +use std::time::{Duration, Instant}; + +pub struct SnowflakeEngine { + conn: SnowflakeConn, +} + +#[derive(Deserialize)] +struct StatementResponse { + #[serde(rename = "resultSetMetaData")] + meta: Option, + data: Option>>, + message: Option, + code: Option, + #[serde(rename = "statementHandle")] + handle: Option, +} + +#[derive(Deserialize)] +struct Meta { + #[serde(rename = "rowType")] + row_type: Vec, +} + +#[derive(Deserialize)] +struct RowType { + name: String, + #[serde(rename = "type")] + ty: String, +} + +impl SnowflakeEngine { + pub fn new(conn: SnowflakeConn) -> Self { + Self { conn } + } + + fn request(&self, r: reqwest::RequestBuilder) -> reqwest::RequestBuilder { + r.bearer_auth(&self.conn.token) + .header("X-Snowflake-Authorization-Token-Type", self.conn.token_type) + .header("Accept", "application/json") + } + + async fn run(&self, sql: &str) -> anyhow::Result<(Vec, Vec>)> { + let client = super::http()?; + let mut body = json!({ + "statement": sql, + "timeout": STATEMENT_TIMEOUT_SECS, + "parameters": { "MULTI_STATEMENT_COUNT": "1" }, + }); + for (key, value) in [ + ("database", &self.conn.database), + ("schema", &self.conn.schema), + ("warehouse", &self.conn.warehouse), + ("role", &self.conn.role), + ] { + if let Some(v) = value { + body[key] = json!(v); + } + } + let mut http = self + .request(client.post(format!("{}/api/v2/statements?async=false", self.conn.base))) + .json(&body) + .send() + .await?; + let started = Instant::now(); + // 202 = 还在跑；其余非 2xx 的 body 里带 message + while http.status() == StatusCode::ACCEPTED { + let partial: StatementResponse = http.json().await?; + let handle = partial.handle.ok_or_else(|| { + anyhow::anyhow!("Snowflake returned 202 without a statementHandle") + })?; + if started.elapsed() > HTTP_POLL_BUDGET { + anyhow::bail!( + "Snowflake statement did not finish within {}s", + HTTP_POLL_BUDGET.as_secs() + ); + } + tokio::time::sleep(Duration::from_secs(1)).await; + http = self + .request(client.get(format!("{}/api/v2/statements/{handle}", self.conn.base))) + .send() + .await?; + } + if !http.status().is_success() { + let status = http.status(); + let text = http.text().await.unwrap_or_default(); + let msg = serde_json::from_str::(&text) + .ok() + .and_then(|r| r.message) + .unwrap_or(text); + anyhow::bail!("Snowflake {status}: {msg}"); + } + let resp: StatementResponse = http.json().await?; + if let (Some(code), Some(message)) = (&resp.code, &resp.message) { + // 2xx 里也可能带业务错误码；090001 是 "statement executed successfully" + if code != "090001" && resp.meta.is_none() { + anyhow::bail!("Snowflake {code}: {message}"); + } + } + let types: Vec = resp.meta.map(|m| m.row_type).unwrap_or_default(); + let rows = resp + .data + .unwrap_or_default() + .into_iter() + .map(|row| { + row.iter() + .enumerate() + .map(|(i, v)| coerce(types.get(i).map(|t| t.ty.as_str()).unwrap_or(""), v)) + .collect() + }) + .collect(); + Ok((types.into_iter().map(|t| t.name).collect(), rows)) + } +} + +#[async_trait::async_trait] +impl QueryEngine for SnowflakeEngine { + async fn test(&self) -> anyhow::Result<()> { + self.run("SELECT 1").await.map(|_| ()) + } + + async fn fetch_schema(&self) -> anyhow::Result> { + let database = self.conn.database.as_deref().ok_or_else(|| { + anyhow::anyhow!("snowflake://: put the database in the connection string (snowflake://:TOKEN@account/DATABASE) so the schema can be read") + })?; + let schema_filter = self + .conn + .schema + .as_deref() + .map(|s| format!(" AND table_schema = {}", sql_literal(s))) + .unwrap_or_default(); + let sql = format!( + "SELECT table_schema, table_name, column_name, data_type, comment \ + FROM \"{}\".information_schema.columns \ + WHERE table_schema <> 'INFORMATION_SCHEMA'{schema_filter} \ + ORDER BY table_schema, table_name, ordinal_position", + database.replace('"', "\"\"") + ); + let (_, rows) = self.run(&sql).await?; + Ok(rows.into_iter().map(super::trino::schema_row).collect()) + } + + async fn execute(&self, sql: &str) -> anyhow::Result { + let (columns, rows) = self.run(&wrap_limit(sql)).await?; + let (rows, truncated) = truncate_rows(rows); + Ok(QueryResult { + rows: rows_to_json_lines(&columns, &rows), + truncated, + }) + } +} + +#[cfg(test)] +mod tests { + use super::super::conn::SnowflakeConn; + use super::super::QueryEngine; + use super::SnowflakeEngine; + use serde_json::json; + use wiremock::matchers::{header, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + fn conn(server: &MockServer) -> SnowflakeConn { + SnowflakeConn::parse(&format!( + "snowflake://:pat-test@{}/ANALYTICS/PUBLIC?warehouse=WH&ssl=false", + server.uri().trim_start_matches("http://") + )) + .unwrap() + } + + #[tokio::test] + async fn a_synchronous_answer_is_typed_by_row_type() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/v2/statements")) + .and(header("authorization", "Bearer pat-test")) + .and(header( + "X-Snowflake-Authorization-Token-Type", + "PROGRAMMATIC_ACCESS_TOKEN", + )) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "resultSetMetaData": { "numRows": 1, "rowType": [ + { "name": "REGION", "type": "text" }, + { "name": "TOTAL", "type": "fixed", "scale": 2 } + ] }, + "data": [ ["east", "42.10"] ], + "code": "090001", + "statementHandle": "h1", + "message": "Statement executed successfully." + }))) + .expect(1) + .mount(&server) + .await; + let out = SnowflakeEngine::new(conn(&server)) + .execute("SELECT region, total FROM orders") + .await + .unwrap(); + assert_eq!(out.rows, vec![r#"{"REGION":"east","TOTAL":42.1}"#]); + } + + #[tokio::test] + async fn a_202_is_polled_until_the_answer_arrives() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/v2/statements")) + .respond_with(ResponseTemplate::new(202).set_body_json(json!({ + "code": "333334", "statementHandle": "h2", "message": "Asynchronous execution in progress." + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/api/v2/statements/h2")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "resultSetMetaData": { "rowType": [ { "name": "N", "type": "fixed" } ] }, + "data": [ ["1"] ], "code": "090001", "statementHandle": "h2" + }))) + .expect(1) + .mount(&server) + .await; + let out = SnowflakeEngine::new(conn(&server)) + .execute("SELECT 1 AS n") + .await + .unwrap(); + assert_eq!(out.rows, vec![r#"{"N":1}"#]); + } + + #[tokio::test] + async fn an_error_body_is_surfaced() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/api/v2/statements")) + .respond_with(ResponseTemplate::new(422).set_body_json(json!({ + "code": "002003", "message": "SQL compilation error: Object 'NOPE' does not exist" + }))) + .mount(&server) + .await; + let err = SnowflakeEngine::new(conn(&server)) + .execute("SELECT * FROM nope") + .await + .unwrap_err() + .to_string(); + assert!(err.contains("does not exist"), "{err}"); + } +} diff --git a/crates/utopia-server/src/query_engine/trino.rs b/crates/utopia-server/src/query_engine/trino.rs new file mode 100644 index 000000000..4b2c40dc3 --- /dev/null +++ b/crates/utopia-server/src/query_engine/trino.rs @@ -0,0 +1,266 @@ +//! Trino（旧名 Presto）：REST 协议 `POST /v1/statement`，然后沿 `nextUri` 一页页取。 +//! 一个引擎顶起整个湖仓——Iceberg / Delta / Hive / Hudi 都是它的 catalog， +//! 换格式不换协议。Starburst 同协议。 +//! +//! 没有会话可设只读：超时靠 `X-Trino-Session: query_max_execution_time`， +//! 只读靠 `guard_sql_for`。 + +use super::conn::TrinoConn; +use super::{ + rows_to_json_lines, sql_literal, truncate_rows, wrap_limit, QueryEngine, QueryResult, + SchemaColumn, HTTP_POLL_BUDGET, STATEMENT_TIMEOUT_SECS, +}; +use base64::Engine as _; +use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION}; +use serde::Deserialize; +use std::time::Instant; + +pub struct TrinoEngine { + conn: TrinoConn, +} + +#[derive(Deserialize)] +struct Column { + name: String, +} + +#[derive(Deserialize)] +struct Page { + #[serde(rename = "nextUri")] + next_uri: Option, + columns: Option>, + data: Option>>, + error: Option, +} + +#[derive(Deserialize)] +struct TrinoError { + message: String, + #[serde(rename = "errorName")] + error_name: Option, +} + +impl TrinoEngine { + pub fn new(conn: TrinoConn) -> Self { + Self { conn } + } + + fn headers(&self) -> anyhow::Result { + let mut h = HeaderMap::new(); + h.insert("X-Trino-User", HeaderValue::from_str(&self.conn.user)?); + h.insert("X-Trino-Source", HeaderValue::from_static("utopia")); + h.insert( + "X-Trino-Session", + HeaderValue::from_str(&format!( + "query_max_execution_time={STATEMENT_TIMEOUT_SECS}s" + ))?, + ); + if let Some(c) = &self.conn.catalog { + h.insert("X-Trino-Catalog", HeaderValue::from_str(c)?); + } + if let Some(s) = &self.conn.schema { + h.insert("X-Trino-Schema", HeaderValue::from_str(s)?); + } + if let Some(p) = &self.conn.password { + let raw = format!("{}:{p}", self.conn.user); + let token = base64::engine::general_purpose::STANDARD.encode(raw); + h.insert( + AUTHORIZATION, + HeaderValue::from_str(&format!("Basic {token}"))?, + ); + } + Ok(h) + } + + /// 提交并沿 nextUri 收完：列在第一个带 columns 的页上，数据分页累积 + async fn run(&self, sql: &str) -> anyhow::Result<(Vec, Vec>)> { + let client = super::http()?; + let headers = self.headers()?; + let mut page: Page = client + .post(format!("{}/v1/statement", self.conn.base)) + .headers(headers.clone()) + .body(sql.to_string()) + .send() + .await? + .error_for_status()? + .json() + .await?; + let started = Instant::now(); + let mut columns: Option> = None; + let mut rows = Vec::new(); + loop { + if let Some(e) = page.error { + let name = e.error_name.map(|n| format!("{n}: ")).unwrap_or_default(); + anyhow::bail!("{name}{}", e.message); + } + if columns.is_none() { + columns = page + .columns + .take() + .map(|cs| cs.into_iter().map(|c| c.name).collect()); + } + if let Some(d) = page.data.take() { + rows.extend(d); + } + let Some(next) = page.next_uri.take() else { + break; + }; + if started.elapsed() > HTTP_POLL_BUDGET { + anyhow::bail!( + "Trino query did not finish within {}s", + HTTP_POLL_BUDGET.as_secs() + ); + } + page = client + .get(&next) + .headers(headers.clone()) + .send() + .await? + .error_for_status()? + .json() + .await?; + } + Ok((columns.unwrap_or_default(), rows)) + } +} + +#[async_trait::async_trait] +impl QueryEngine for TrinoEngine { + async fn test(&self) -> anyhow::Result<()> { + self.run("SELECT 1").await.map(|_| ()) + } + + async fn fetch_schema(&self) -> anyhow::Result> { + let catalog = self.conn.catalog.as_deref().ok_or_else(|| { + anyhow::anyhow!("trino://: put the catalog in the connection string (trino://user@host/CATALOG) so the schema can be read") + })?; + let schema_filter = self + .conn + .schema + .as_deref() + .map(|s| format!(" AND table_schema = {}", sql_literal(s))) + .unwrap_or_default(); + let sql = format!( + "SELECT table_schema, table_name, column_name, data_type, comment \ + FROM \"{}\".information_schema.columns \ + WHERE table_schema <> 'information_schema'{schema_filter} \ + ORDER BY table_schema, table_name, ordinal_position", + catalog.replace('"', "\"\"") + ); + let (_, rows) = self.run(&sql).await?; + Ok(rows.into_iter().map(schema_row).collect()) + } + + async fn execute(&self, sql: &str) -> anyhow::Result { + let (columns, rows) = self.run(&wrap_limit(sql)).await?; + let (rows, truncated) = truncate_rows(rows); + Ok(QueryResult { + rows: rows_to_json_lines(&columns, &rows), + truncated, + }) + } +} + +/// information_schema 的一行 → SchemaColumn（值可能是 null，comment 常是） +pub(crate) fn schema_row(row: Vec) -> SchemaColumn { + let text = |i: usize| -> String { + row.get(i) + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string() + }; + SchemaColumn { + schema: text(0), + table: text(1), + column: text(2), + data_type: text(3), + comment: row + .get(4) + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .map(str::to_string), + } +} + +#[cfg(test)] +mod tests { + use super::super::conn::TrinoConn; + use super::super::QueryEngine; + use super::TrinoEngine; + use serde_json::json; + use wiremock::matchers::{body_string_contains, header, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + #[tokio::test] + async fn follows_next_uri_and_keeps_column_order() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/statement")) + .and(header("X-Trino-User", "alice")) + .and(header("X-Trino-Catalog", "hive")) + .and(body_string_contains("LIMIT 201")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "q1", + "nextUri": format!("{}/v1/statement/q1/1", server.uri()), + "stats": { "state": "QUEUED" } + }))) + .expect(1) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/v1/statement/q1/1")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "q1", + "columns": [ { "name": "region", "type": "varchar" }, { "name": "total", "type": "double" } ], + "data": [ ["east", 12.5], ["west", 3] ], + "stats": { "state": "FINISHED" } + }))) + .expect(1) + .mount(&server) + .await; + + let uri = server.uri(); + let conn = TrinoConn::parse(&format!( + "trino://alice@{}/hive/default?ssl=false", + uri.trim_start_matches("http://") + )) + .unwrap(); + let out = TrinoEngine::new(conn) + .execute("SELECT region, total FROM orders") + .await + .unwrap(); + assert_eq!( + out.rows, + vec![ + r#"{"region":"east","total":12.5}"#, + r#"{"region":"west","total":3}"# + ] + ); + assert!(!out.truncated); + } + + #[tokio::test] + async fn a_trino_error_page_becomes_an_error() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/statement")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "q2", + "error": { "message": "line 1:8: Table 'hive.default.nope' does not exist", "errorName": "TABLE_NOT_FOUND" }, + "stats": { "state": "FAILED" } + }))) + .mount(&server) + .await; + let conn = TrinoConn::parse(&format!( + "trino://alice@{}/hive?ssl=false", + server.uri().trim_start_matches("http://") + )) + .unwrap(); + let err = TrinoEngine::new(conn) + .execute("SELECT * FROM nope") + .await + .unwrap_err() + .to_string(); + assert!(err.contains("TABLE_NOT_FOUND"), "{err}"); + } +} diff --git a/crates/utopia-store/src/datasources.rs b/crates/utopia-store/src/datasources.rs index 5fad63897..6e3ec3a9e 100644 --- a/crates/utopia-store/src/datasources.rs +++ b/crates/utopia-store/src/datasources.rs @@ -18,17 +18,16 @@ type DataSourceRow = ( Option, ); -/// 连接串 → 无凭据摘要（host:port/db）。解析失败给占位符，绝不回显原串。 +/// 连接串 → 无凭据摘要（host[:port]/path）。解析失败给占位符，绝不回显原串。 +/// 端口没写就不补：四种 scheme 的默认端口各不相同，补错比不补更误导 pub fn conn_summary(conn: &str) -> String { url::Url::parse(conn) .ok() .map(|u| { format!( - "{}:{}{}", + "{}{}{}", u.host_str().unwrap_or("?"), - u.port() - .map(|p| p.to_string()) - .unwrap_or_else(|| "5432".into()), + u.port().map(|p| format!(":{p}")).unwrap_or_default(), u.path() ) }) @@ -77,16 +76,12 @@ pub async fn create( "Data source name is required", )); } - if engine != "postgres" { - return Err(AppError::invalid( - "only_postgres", - "Only the postgres engine is supported for now", - )); - } - if !conn_string.starts_with("postgres://") && !conn_string.starts_with("postgresql://") { + // 引擎由调用方按连接串的 scheme 定（`query_engine::engine_from_conn`）； + // 允许的取值在迁移 0020 的 CHECK 里，这里不再复制一份 + if engine.is_empty() || conn_string.trim().is_empty() { return Err(AppError::invalid( "bad_conn_string", - "Connection string must start with postgres://", + "A connection string is required", )); } let id = Uuid::now_v7(); diff --git a/docs/decisions/0018-the-lakehouse-is-one-protocol-away.md b/docs/decisions/0018-the-lakehouse-is-one-protocol-away.md new file mode 100644 index 000000000..2bbdcf6c1 --- /dev/null +++ b/docs/decisions/0018-the-lakehouse-is-one-protocol-away.md @@ -0,0 +1,58 @@ +# 0018 · The lakehouse is one protocol away + +- **Status**: implemented · `trino` / `databricks` / `snowflake` engines in `query_engine/` (migration `0020` widens the `engine` CHECK) · every engine is covered by protocol replays only (wiremock); **none has run against a real cluster** (#240, #241, #242) · MaxCompute is not done, see the last section +- **Written**: 2026-09-03 (conventions in the [README](README.md)) +- **Related**: [0011](0011-a-mapping-is-not-a-fact.md) placed data sources at the deployment level and mounts at the base level; this record leaves that layer alone. [0016](0016-close-the-open-seams-before-cutting-new-ones.md) D4 put the MySQL wire protocol ahead of the lakehouse; the first section explains why the order flipped + +> The roadmap line reads "Iceberg / Delta Lake, Databricks, Snowflake and MaxCompute": four names, three kinds of thing. The first two are table formats, the next two are services, the last is a service on another cloud. Treating them as four engines would be writing code per product name. This record pins down what an engine is first, then decides which ones to build. + +## Engines follow protocols + +The header of `query_engine` (written for 0011) already gave the direction: the Postgres wire protocol, then the MySQL wire family, then the HTTP family. This step lands the HTTP family and skips MySQL. 0016 D4's "TiDB / OceanBase / Doris / StarRocks for free" is a fine list, but the lakehouse is what is wanted now, and on the protocol axis it is closer than it looks: + +| Wanted | What it is | Which protocol that is for us | +|---|---|---| +| Iceberg, Delta Lake, Hive, Hudi | table formats plus a catalog, with no query endpoint of their own | one **Trino** catalog each; `POST /v1/statement` with `nextUri` paging | +| Databricks | a Delta lakehouse behind a SQL warehouse | its own **SQL Statement Execution API** (`/api/2.0/sql/statements`) | +| Snowflake | a cloud warehouse that also reads Iceberg | its own **SQL API v2** (`/api/v2/statements`) | +| MaxCompute | Alibaba Cloud's warehouse | signed REST, asynchronous instances, results through Tunnel | + +The first three rows are all "JSON in, JSON out, Bearer or Basic auth". `reqwest` is already a dependency; each engine is about two hundred lines. The binary still carries no native database driver, which is the promise in the README's first sentence. It is also why the answer to Iceberg is Trino rather than an Iceberg reader: reading Iceberg directly pulls in Arrow, Parquet, object-storage SDKs and a query planner, and that is a different product. + +## The connection string is the only input + +The data-source page has a name and a connection string. Three new engines add no dropdown: **the scheme picks the engine** (`engine_from_conn`), and each engine parses the rest (`conn.rs`). The shape follows `postgres://user:pass@host/db`: credentials in the userinfo, the path is "catalog / database / schema", engine-specific switches go in the query string: + +``` +trino://alice[:password]@host[:8080]/catalog[/schema][?ssl=true] +databricks://:TOKEN@workspace-host/sql/1.0/warehouses/ID[?catalog=main&schema=default] +snowflake://:TOKEN@account.snowflakecomputing.com/DB[/SCHEMA][?warehouse=WH&role=R&token_type=pat|oauth] +``` + +The Databricks path is the httpPath shown in the console, so it can be pasted as is. All three tokens sit in the password position; `TOKEN@` with the colon missing is the most common slip, so the username position is accepted too. The shape is validated at registration, and the error carries the expected form. `ssl=false` exists for local proxies and stand-ins; the three services themselves only speak https. + +Left out on purpose: Snowflake key-pair JWT (local RSA signing, a dependency for a second login method, wait for someone to need it) and Trino Kerberos / OAuth2 (same reasoning). Passwords and tokens are the whole surface. + +## The HTTP family has three of the four gates + +0011 set up defense in depth: parse and admit only SELECT, wrap a LIMIT, a read-only session with a timeout, JSON Lines out. The HTTP family **has no session**, so the third gate is a timeout alone (Trino's `query_max_execution_time` session property, Databricks' `wait_timeout`, Snowflake's `timeout`), and read-only rests entirely on the first gate. The first gate therefore parses with each engine's dialect: `DatabricksDialect`, `SnowflakeDialect`, and `GenericDialect` for Trino (sqlparser has no Trino dialect; Generic is a superset). One test runs the same three checks under all four dialects: SELECT passes, DELETE fails, two statements fail. + +The fourth gate is assembled here for the HTTP family. Databricks' `JSON_ARRAY` and Snowflake's `data` return every value as a string; `coerce` restores numbers and booleans from the manifest / rowType column types, otherwise a model handed `"42"` stops doing arithmetic. Key order is assembled by hand instead of through `serde_json::Map`, which sorts keys unless `preserve_order` is on, and column order is the order the query wrote. + +## Loopback goes direct + +reqwest's system-proxy detection on Windows reads the registry and did not honor a `127.*` bypass, so a stand-in on the loopback address went through the proxy and came back as 502. The engine client now carries an explicit policy: loopback and `NO_PROXY` hosts go direct, everything else follows `HTTPS_PROXY` / `HTTP_PROXY` / `ALL_PROXY`. A server process reads its environment; that matches how docker-compose configures it. The other connectors keep reqwest's default. + +## Replays are the only tests so far + +No stand-in for any of the three runs on this machine: Docker Hub cannot be reached here (see memory), and Databricks and Snowflake are cloud-only anyway. The tests replay each vendor's documented protocol with wiremock: Trino's two `nextUri` pages, Databricks' PENDING → SUCCEEDED polling, Snowflake's 202 → 200, and each error body. **They prove our reading of the protocol, and only that.** Until one real cluster has answered, the README keeps these three marked as awaiting a real run, handled the way #214 / #215 handle GCS and Notion: an issue per engine labeled help wanted. + +## MaxCompute waits + +It is the one name of the four that is not "JSON in, JSON out": requests are signed with an AccessKey (HMAC-SHA1 over canonicalized headers), SQL runs as an asynchronous instance, and results come either through Tunnel (another protocol) or `GetInstanceResult` as CSV capped at ten thousand rows. Together that is a connector's worth of work, and this machine has no account that could sign a request, so the result could only be "probably like this". It stays on the roadmap until someone with an account arrives, or until its MySQL-compatible entry (MCQA) can ride the MySQL wire protocol of 0016 D4. + +## Open questions + +- **How much schema to fetch.** All three expose `information_schema.columns`, and a lakehouse catalog can hold thousands of tables; `sync_schema_doc` caps at 200. With a schema in the connection string only that schema is read, otherwise the whole catalog. Whether that is enough waits for a real cluster. +- **The type-restoration table** in `coerce` is hand-written from the three vendors' docs. Snowflake's `fixed` with a scale returns `"42.10"`, which becomes 42.1 and loses the trailing zero; harmless for a model, possibly not for an "exact definition". Revisit when the semantic layer keeps evidence (0016 D1) and decide whether to keep the raw string alongside. +- **Trino's `ssl` inference**: a password, `ssl=true`, or port 443 / 8443 means https, anything else is plaintext. That is trino-python's rule, and someone who gets it wrong sees a TLS error instead of a hint. diff --git a/docs/decisions/README.md b/docs/decisions/README.md index 81142c040..16e5c85c3 100644 --- a/docs/decisions/README.md +++ b/docs/decisions/README.md @@ -42,6 +42,7 @@ | 0014 | [身份跟着人，范围跟着令牌](0014-identity-from-the-person-scope-from-the-token.md) | 已实施（#180）· MCP 只读五工具 · 令牌页在账户层（A2）· 误导性的占位 crate 已删 | | 0015 | [记下一句话，不等于断言一个事实](0015-recording-a-sentence-is-not-asserting-a-fact.md) | 已实施 · 记忆抽出的事实进 `pending_facts`，Review 新档 + 跟在 remember 步骤后的确认卡 · `remember` 重新打开 · MCP 放开写是下一刀 | | 0016 | [先把开着的口子收上，再开新的](0016-close-the-open-seams-before-cutting-new-ones.md) | 规划中 · v0.1.0 之后的排期：A 收口 → B 推理机 ∥ C 尺子与本体 → D 语义层 → E 企业交付；模拟引擎后置 | +| 0018 | [The lakehouse is one protocol away](0018-the-lakehouse-is-one-protocol-away.md) | 已实施 · 问数引擎扩到 HTTP 族：Trino（Iceberg / Delta / Hive）、Databricks、Snowflake，scheme 决定引擎，只有回放测试 · MaxCompute 未做 | ## 不是决策记录的那些 diff --git a/migrations/0020_lakehouse_engines.sql b/migrations/0020_lakehouse_engines.sql new file mode 100644 index 000000000..3781a6a69 --- /dev/null +++ b/migrations/0020_lakehouse_engines.sql @@ -0,0 +1,7 @@ +-- 问数引擎扩到 HTTP 协议族：trino（Iceberg / Delta / Hive 都是它的 catalog）、 +-- databricks（SQL Statement API）、snowflake（SQL API v2）。 +-- 挂载模型与注册表引擎无关（0006 的判断仍成立），这里只放宽 engine 的取值； +-- 允许的名字与 `query_engine::ENGINES` 同一张表。 +ALTER TABLE data_sources DROP CONSTRAINT data_sources_engine_check; +ALTER TABLE data_sources ADD CONSTRAINT data_sources_engine_check + CHECK (engine IN ('postgres', 'trino', 'databricks', 'snowflake')); diff --git a/web/src/i18n/en.ts b/web/src/i18n/en.ts index c2cce2958..e08ce9918 100644 --- a/web/src/i18n/en.ts +++ b/web/src/i18n/en.ts @@ -876,7 +876,13 @@ export const en = { "Read-only database connections for asking questions about your data in Chat. " + "Register connections here; each knowledge base mounts the ones it may query.", name: "Name", - connString: "Connection string (postgres://user:pass@host:5432/db)", + connString: "Connection string — the scheme picks the engine", + // 四种写法各一行；令牌放 password 位，Databricks 的路径就是控制台里的 httpPath + connSchemes: + "postgres://user:pass@host:5432/db\n" + + "trino://user[:pass]@host:8080/catalog[/schema] (Iceberg, Delta Lake, Hive)\n" + + "databricks://:TOKEN@host/sql/1.0/warehouses/ID?catalog=main\n" + + "snowflake://:TOKEN@account.snowflakecomputing.com/DB/SCHEMA?warehouse=WH", add: "Add data source", test: "Test", testOk: "Connected", diff --git a/web/src/i18n/zh.ts b/web/src/i18n/zh.ts index c5915ed25..f834b4cee 100644 --- a/web/src/i18n/zh.ts +++ b/web/src/i18n/zh.ts @@ -800,7 +800,12 @@ export const zh: Strings = { "只读的数据库连接，用于在「对话」里就你的数据提问。" + "在这里登记连接；每个知识库各自挂载允许查询的那些。", name: "名称", - connString: "连接串（postgres://user:pass@host:5432/db）", + connString: "连接串，前缀决定引擎", + connSchemes: + "postgres://user:pass@host:5432/db\n" + + "trino://user[:pass]@host:8080/catalog[/schema] （Iceberg、Delta Lake、Hive）\n" + + "databricks://:TOKEN@host/sql/1.0/warehouses/ID?catalog=main\n" + + "snowflake://:TOKEN@account.snowflakecomputing.com/DB/SCHEMA?warehouse=WH", add: "添加数据源", test: "测试", testOk: "已连接", diff --git a/web/src/pages/Settings.tsx b/web/src/pages/Settings.tsx index 4a4aa9c09..f61040aa9 100644 --- a/web/src/pages/Settings.tsx +++ b/web/src/pages/Settings.tsx @@ -639,7 +639,12 @@ function DataSourcesAdmin() {
-
{d.name}
+
+ {d.name} + + {d.engine} + +
{d.summary}
@@ -695,6 +700,9 @@ function DataSourcesAdmin() { value={conn} onChange={(e) => setConn(e.target.value)} /> +

+ {S.settings.datasources.connSchemes} +