From 1dd8e6bb4ca4b6cc34cb9f6d0bc951841849c383 Mon Sep 17 00:00:00 2001 From: amitgilad Date: Thu, 2 Apr 2026 19:13:39 +0300 Subject: [PATCH 01/14] feat: add snowflake frontend protocol --- Cargo.lock | 24 +- crates/queryflux-core/src/config.rs | 10 + crates/queryflux-core/src/query.rs | 6 + crates/queryflux-e2e-tests/src/harness.rs | 2 + crates/queryflux-frontend/Cargo.toml | 7 + crates/queryflux-frontend/src/admin.rs | 12 + crates/queryflux-frontend/src/lib.rs | 1 + .../src/snowflake/http/format.rs | 368 ++++++++ .../src/snowflake/http/handlers/common.rs | 149 ++++ .../src/snowflake/http/handlers/mod.rs | 4 + .../src/snowflake/http/handlers/query.rs | 196 +++++ .../src/snowflake/http/handlers/session.rs | 181 ++++ .../src/snowflake/http/handlers/token.rs | 43 + .../src/snowflake/http/mod.rs | 33 + .../src/snowflake/http/session_store.rs | 52 ++ .../queryflux-frontend/src/snowflake/mod.rs | 50 ++ .../queryflux-frontend/src/snowflake/proxy.rs | 73 ++ .../src/snowflake/sql_api/auth.rs | 70 ++ .../src/snowflake/sql_api/handlers.rs | 257 ++++++ .../src/snowflake/sql_api/mod.rs | 23 + .../queryflux-frontend/src/snowflake/tests.rs | 816 ++++++++++++++++++ crates/queryflux-frontend/src/state.rs | 4 + .../src/implementations/protocol_based.rs | 9 +- .../src/implementations/python_script.rs | 2 + .../queryflux-routing/tests/router_tests.rs | 71 +- crates/queryflux/src/main.rs | 42 +- docs/README.md | 18 - docs/architecture.md | 354 -------- docs/auth-authz-design.md | 733 ---------------- docs/motivation-and-goals.md | 43 - docs/observability.md | 117 --- docs/query-translation.md | 149 ---- docs/routing-and-clusters.md | 158 ---- queryflux-studio/README.md | 2 +- queryflux-studio/app/protocols/page.tsx | 53 +- .../architecture/adding-engine-support.md | 313 +------ .../architecture/adding-support/backend.md | 85 +- .../architecture/adding-support/frontend.md | 43 + .../architecture/adding-support/overview.md | 32 + .../docs/architecture/frontends/flight-sql.md | 71 ++ .../docs/architecture/frontends/mysql-wire.md | 104 +++ .../docs/architecture/frontends/overview.md | 115 +++ .../architecture/frontends/postgres-wire.md | 84 ++ .../docs/architecture/frontends/snowflake.md | 191 ++++ .../docs/architecture/frontends/trino-http.md | 84 ++ website/docs/architecture/overview.md | 3 +- .../docs/architecture/routing-and-clusters.md | 2 +- website/docs/architecture/system-map.md | 7 +- website/docs/configuration.md | 5 + website/docs/roadmap.md | 1 + website/sidebars.ts | 24 +- .../architecture/adding-engine-support.md | 313 +------ .../architecture/adding-support/backend.md | 275 ++++++ .../architecture/adding-support/frontend.md | 43 + .../architecture/adding-support/overview.md | 32 + .../architecture/frontends/flight-sql.md | 71 ++ .../architecture/frontends/mysql-wire.md | 104 +++ .../architecture/frontends/overview.md | 115 +++ .../architecture/frontends/postgres-wire.md | 84 ++ .../architecture/frontends/snowflake.md | 191 ++++ .../architecture/frontends/trino-http.md | 84 ++ .../version-0.1.0/architecture/overview.md | 3 +- .../architecture/routing-and-clusters.md | 2 +- .../version-0.1.0/architecture/system-map.md | 7 +- .../version-0.1.0/configuration.md | 5 + .../versioned_docs/version-0.1.0/roadmap.md | 1 + .../version-0.1.0-sidebars.json | 24 +- 67 files changed, 4366 insertions(+), 2284 deletions(-) create mode 100644 crates/queryflux-frontend/src/snowflake/http/format.rs create mode 100644 crates/queryflux-frontend/src/snowflake/http/handlers/common.rs create mode 100644 crates/queryflux-frontend/src/snowflake/http/handlers/mod.rs create mode 100644 crates/queryflux-frontend/src/snowflake/http/handlers/query.rs create mode 100644 crates/queryflux-frontend/src/snowflake/http/handlers/session.rs create mode 100644 crates/queryflux-frontend/src/snowflake/http/handlers/token.rs create mode 100644 crates/queryflux-frontend/src/snowflake/http/mod.rs create mode 100644 crates/queryflux-frontend/src/snowflake/http/session_store.rs create mode 100644 crates/queryflux-frontend/src/snowflake/mod.rs create mode 100644 crates/queryflux-frontend/src/snowflake/proxy.rs create mode 100644 crates/queryflux-frontend/src/snowflake/sql_api/auth.rs create mode 100644 crates/queryflux-frontend/src/snowflake/sql_api/handlers.rs create mode 100644 crates/queryflux-frontend/src/snowflake/sql_api/mod.rs create mode 100644 crates/queryflux-frontend/src/snowflake/tests.rs delete mode 100644 docs/README.md delete mode 100644 docs/architecture.md delete mode 100644 docs/auth-authz-design.md delete mode 100644 docs/motivation-and-goals.md delete mode 100644 docs/observability.md delete mode 100644 docs/query-translation.md delete mode 100644 docs/routing-and-clusters.md rename docs/adding-engine-support.md => website/docs/architecture/adding-support/backend.md (81%) create mode 100644 website/docs/architecture/adding-support/frontend.md create mode 100644 website/docs/architecture/adding-support/overview.md create mode 100644 website/docs/architecture/frontends/flight-sql.md create mode 100644 website/docs/architecture/frontends/mysql-wire.md create mode 100644 website/docs/architecture/frontends/overview.md create mode 100644 website/docs/architecture/frontends/postgres-wire.md create mode 100644 website/docs/architecture/frontends/snowflake.md create mode 100644 website/docs/architecture/frontends/trino-http.md create mode 100644 website/versioned_docs/version-0.1.0/architecture/adding-support/backend.md create mode 100644 website/versioned_docs/version-0.1.0/architecture/adding-support/frontend.md create mode 100644 website/versioned_docs/version-0.1.0/architecture/adding-support/overview.md create mode 100644 website/versioned_docs/version-0.1.0/architecture/frontends/flight-sql.md create mode 100644 website/versioned_docs/version-0.1.0/architecture/frontends/mysql-wire.md create mode 100644 website/versioned_docs/version-0.1.0/architecture/frontends/overview.md create mode 100644 website/versioned_docs/version-0.1.0/architecture/frontends/postgres-wire.md create mode 100644 website/versioned_docs/version-0.1.0/architecture/frontends/snowflake.md create mode 100644 website/versioned_docs/version-0.1.0/architecture/frontends/trino-http.md diff --git a/Cargo.lock b/Cargo.lock index 7059b09..794eb77 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2466,6 +2466,21 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonwebtoken" +version = "9.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a87cc7a48537badeae96744432de36f4be2b4a34a05a5ef32e9dd8a1c169dde" +dependencies = [ + "base64 0.22.1", + "js-sys", + "pem", + "ring", + "serde", + "serde_json", + "simple_asn1", +] + [[package]] name = "jsonwebtoken" version = "10.3.0" @@ -3449,7 +3464,7 @@ version = "0.1.0" dependencies = [ "async-trait", "dashmap", - "jsonwebtoken", + "jsonwebtoken 10.3.0", "ldap3", "queryflux-core", "reqwest 0.12.28", @@ -3565,9 +3580,13 @@ dependencies = [ "arrow-ipc", "async-trait", "axum", + "base64 0.22.1", "bytes", "chrono", + "dashmap", + "flate2", "futures", + "jsonwebtoken 9.3.1", "prost", "queryflux-auth", "queryflux-cluster-manager", @@ -3578,8 +3597,10 @@ dependencies = [ "queryflux-routing", "queryflux-translation", "reqwest 0.12.28", + "rsa", "serde", "serde_json", + "sha2", "tokio", "tokio-stream", "tonic", @@ -3587,6 +3608,7 @@ dependencies = [ "tracing", "urlencoding", "utoipa", + "uuid", ] [[package]] diff --git a/crates/queryflux-core/src/config.rs b/crates/queryflux-core/src/config.rs index 2c4da40..5e4cc2c 100644 --- a/crates/queryflux-core/src/config.rs +++ b/crates/queryflux-core/src/config.rs @@ -322,6 +322,10 @@ pub struct FrontendsConfig { pub clickhouse_http: Option, #[serde(default)] pub flight_sql: Option, + #[serde(default)] + pub snowflake_http: Option, + #[serde(default)] + pub snowflake_sql_api: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -648,6 +652,12 @@ pub enum RouterConfig { mysql_wire: Option, #[serde(default)] clickhouse_http: Option, + #[serde(default)] + flight_sql: Option, + #[serde(default)] + snowflake_http: Option, + #[serde(default)] + snowflake_sql_api: Option, }, #[serde(rename_all = "camelCase")] Header { diff --git a/crates/queryflux-core/src/query.rs b/crates/queryflux-core/src/query.rs index 5f04f92..2dc669f 100644 --- a/crates/queryflux-core/src/query.rs +++ b/crates/queryflux-core/src/query.rs @@ -67,6 +67,10 @@ pub enum FrontendProtocol { MySqlWire, ClickHouseHttp, FlightSql, + /// Snowflake internal wire protocol used by JDBC/ODBC/Python/Go/Node.js connectors. + SnowflakeHttp, + /// Snowflake public SQL REST API v2 (/api/v2/statements). + SnowflakeSqlApi, } impl FrontendProtocol { @@ -78,6 +82,8 @@ impl FrontendProtocol { FrontendProtocol::MySqlWire => SqlDialect::MySql, FrontendProtocol::ClickHouseHttp => SqlDialect::ClickHouse, FrontendProtocol::FlightSql => SqlDialect::Generic, + FrontendProtocol::SnowflakeHttp => SqlDialect::Generic, + FrontendProtocol::SnowflakeSqlApi => SqlDialect::Generic, } } } diff --git a/crates/queryflux-e2e-tests/src/harness.rs b/crates/queryflux-e2e-tests/src/harness.rs index b440b41..4eacecc 100644 --- a/crates/queryflux-e2e-tests/src/harness.rs +++ b/crates/queryflux-e2e-tests/src/harness.rs @@ -31,6 +31,7 @@ use queryflux_engine_adapters::{ starrocks::StarRocksAdapter, trino::TrinoAdapter, EngineAdapterTrait, }; use queryflux_frontend::{ + snowflake::http::session_store::SnowflakeSessionStore, state::LiveConfig, trino_http::{state::AppState, TrinoHttpFrontend}, }; @@ -230,6 +231,7 @@ impl TestHarness { auth_provider: Arc::new(NoneAuthProvider::new(false)) as Arc, authorization: Arc::new(AllowAllAuthorization) as Arc, identity_resolver: Arc::new(BackendIdentityResolver::new()), + snowflake_sessions: SnowflakeSessionStore::new(), }); let router: Router = TrinoHttpFrontend::new(state, port).router(); diff --git a/crates/queryflux-frontend/Cargo.toml b/crates/queryflux-frontend/Cargo.toml index fd65091..51c06d7 100644 --- a/crates/queryflux-frontend/Cargo.toml +++ b/crates/queryflux-frontend/Cargo.toml @@ -32,3 +32,10 @@ arrow-flight = { version = "58.1.0", features = ["flight-sql"] } arrow-ipc = "58.1.0" tonic = { version = "0.14.5", features = ["transport"] } prost = "0.14.3" +dashmap = { workspace = true } +jsonwebtoken = "9" +rsa = "0.9" +sha2 = "0.10" +uuid = { version = "1", features = ["v4"] } +base64 = "0.22" +flate2 = "1" diff --git a/crates/queryflux-frontend/src/admin.rs b/crates/queryflux-frontend/src/admin.rs index b1d112d..c137c12 100644 --- a/crates/queryflux-frontend/src/admin.rs +++ b/crates/queryflux-frontend/src/admin.rs @@ -250,6 +250,18 @@ pub fn build_frontends_status( "Arrow Flight SQL / gRPC-style access (driver-dependent).", frontends.flight_sql.as_ref(), ), + opt_fe( + "snowflake_http", + "Snowflake HTTP (wire)", + "Snowflake internal HTTP wire protocol (JDBC/ODBC/Python/Go/Node.js connectors).", + frontends.snowflake_http.as_ref(), + ), + opt_fe( + "snowflake_sql_api", + "Snowflake SQL API", + "Snowflake REST SQL API v2 (POST /api/v2/statements).", + frontends.snowflake_sql_api.as_ref(), + ), ]; FrontendsStatusDto { diff --git a/crates/queryflux-frontend/src/lib.rs b/crates/queryflux-frontend/src/lib.rs index 4343421..3058113 100644 --- a/crates/queryflux-frontend/src/lib.rs +++ b/crates/queryflux-frontend/src/lib.rs @@ -3,6 +3,7 @@ pub mod dispatch; pub mod flight_sql; pub mod mysql_wire; pub mod postgres_wire; +pub mod snowflake; pub mod state; pub mod trino_http; diff --git a/crates/queryflux-frontend/src/snowflake/http/format.rs b/crates/queryflux-frontend/src/snowflake/http/format.rs new file mode 100644 index 0000000..8c203de --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/format.rs @@ -0,0 +1,368 @@ +/// Snowflake wire-protocol response formatting. +/// +/// Converts Arrow RecordBatches into the Snowflake JSON response format: +/// - `rowtype`: column metadata array (matches Snowflake's JSON schema) +/// - `rowsetBase64`: base64-encoded Arrow IPC stream, with Snowflake field metadata and +/// data transformations that the nanoarrow_arrow_iterator expects. +/// +/// Key references (fakesnow + Snowflake connector source): +/// https://github.com/snowflakedb/snowflake-connector-python/blob/main/src/snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowTableIterator.cpp +/// https://github.com/snowflakedb/snowflake-connector-python/blob/main/src/snowflake/connector/nanoarrow_cpp/ArrowIterator/SnowflakeType.cpp +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, Int32Array, Int64Array, StructArray, TimestampNanosecondArray, +}; +use arrow::datatypes::{DataType, Field, Fields, Schema, TimeUnit}; +use arrow::record_batch::RecordBatch; +use arrow_ipc::writer::StreamWriter; +use base64::Engine; +use serde_json::{json, Value}; + +// --------------------------------------------------------------------------- +// Arrow DataType → Snowflake type string + metadata +// --------------------------------------------------------------------------- + +struct SfTypeInfo { + logical_type: &'static str, + precision: u32, + scale: u32, + char_length: u32, + byte_length: u32, +} + +fn sf_type_info(dt: &DataType) -> SfTypeInfo { + match dt { + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 => SfTypeInfo { + logical_type: "FIXED", + precision: 38, + scale: 0, + char_length: 0, + byte_length: 8, + }, + DataType::Decimal128(p, s) | DataType::Decimal256(p, s) => SfTypeInfo { + logical_type: "FIXED", + precision: *p as u32, + scale: *s as u32, + char_length: 0, + byte_length: 16, + }, + DataType::Float16 | DataType::Float32 | DataType::Float64 => SfTypeInfo { + logical_type: "REAL", + precision: 0, + scale: 0, + char_length: 0, + byte_length: 8, + }, + DataType::Utf8 | DataType::LargeUtf8 => SfTypeInfo { + logical_type: "TEXT", + precision: 0, + scale: 0, + char_length: 16_777_216, + byte_length: 16_777_216, + }, + DataType::Boolean => SfTypeInfo { + logical_type: "BOOLEAN", + precision: 0, + scale: 0, + char_length: 0, + byte_length: 1, + }, + DataType::Date32 | DataType::Date64 => SfTypeInfo { + logical_type: "DATE", + precision: 0, + scale: 0, + char_length: 0, + byte_length: 4, + }, + DataType::Time32(_) | DataType::Time64(_) => SfTypeInfo { + logical_type: "TIME", + precision: 0, + scale: 9, + char_length: 0, + byte_length: 8, + }, + DataType::Timestamp(_, Some(_)) => SfTypeInfo { + logical_type: "TIMESTAMP_TZ", + precision: 0, + scale: 9, + char_length: 0, + byte_length: 16, + }, + DataType::Timestamp(_, None) => SfTypeInfo { + logical_type: "TIMESTAMP_NTZ", + precision: 0, + scale: 9, + char_length: 0, + byte_length: 16, + }, + DataType::Binary | DataType::LargeBinary | DataType::FixedSizeBinary(_) => SfTypeInfo { + logical_type: "BINARY", + precision: 0, + scale: 0, + char_length: 0, + byte_length: 8_388_608, + }, + _ => SfTypeInfo { + logical_type: "VARIANT", + precision: 0, + scale: 0, + char_length: 0, + byte_length: 0, + }, + } +} + +// --------------------------------------------------------------------------- +// rowtype JSON (for the HTTP response body) +// --------------------------------------------------------------------------- + +pub fn schema_to_rowtype(schema: &Schema) -> Value { + let cols: Vec = schema + .fields() + .iter() + .map(|f| { + let info = sf_type_info(f.data_type()); + json!({ + "name": f.name(), + "database": "", + "schema": "", + "table": "", + "nullable": f.is_nullable(), + "type": info.logical_type.to_lowercase(), + "byteLength": if info.byte_length > 0 { Some(info.byte_length) } else { None:: }, + "length": if info.char_length > 0 { Some(info.char_length) } else { None:: }, + "scale": if info.scale > 0 { Some(info.scale) } else { None:: }, + "precision": if info.precision > 0 { Some(info.precision) } else { None:: }, + "collation": null + }) + }) + .collect(); + json!(cols) +} + +// --------------------------------------------------------------------------- +// Arrow schema → Snowflake-annotated schema (adds required field metadata) +// --------------------------------------------------------------------------- + +/// The nanoarrow_arrow_iterator reads `metadata.at("logicalType")` (and others) from every +/// Arrow field. Missing metadata causes `unordered_map::at: key not found`. +fn sf_arrow_field(field: &Field) -> Field { + let info = sf_type_info(field.data_type()); + // Timestamps become structs in Snowflake Arrow format. + let sf_type = match field.data_type() { + DataType::Timestamp(_, tz) => { + let mut fields = vec![ + Field::new("epoch", DataType::Int64, false), + Field::new("fraction", DataType::Int32, false), + ]; + if tz.is_some() { + fields.push(Field::new("timezone", DataType::Int32, false)); + } + DataType::Struct(Fields::from(fields)) + } + // Time64 → int64 (nanoseconds) + DataType::Time64(_) => DataType::Int64, + DataType::Time32(_) => DataType::Int64, + // UInt64 → int64 (connector expects signed) + DataType::UInt64 => DataType::Int64, + other => other.clone(), + }; + + let metadata = std::collections::HashMap::from([ + ("logicalType".to_string(), info.logical_type.to_string()), + ("precision".to_string(), info.precision.to_string()), + ("scale".to_string(), info.scale.to_string()), + ("charLength".to_string(), info.char_length.to_string()), + ]); + + Field::new(field.name(), sf_type, field.is_nullable()).with_metadata(metadata) +} + +fn sf_arrow_schema(schema: &Schema) -> Schema { + let fields: Vec = schema.fields().iter().map(|f| sf_arrow_field(f)).collect(); + Schema::new(fields) +} + +// --------------------------------------------------------------------------- +// Data transformation: convert columns to Snowflake Arrow wire format +// --------------------------------------------------------------------------- + +/// Cast a column to Snowflake's expected Arrow wire type. +fn to_sf_array(arr: &ArrayRef) -> ArrayRef { + match arr.data_type() { + DataType::Timestamp(unit, _tz) => timestamp_to_sf_struct(arr, unit), + DataType::Time64(unit) => { + let ns = match unit { + TimeUnit::Nanosecond => arr.clone(), + TimeUnit::Microsecond => { + let cast = + arrow::compute::cast(arr, &DataType::Int64).unwrap_or_else(|_| arr.clone()); + // µs → ns: multiply by 1000 + let ns_arr: Int64Array = cast + .as_any() + .downcast_ref::() + .map(|a| Int64Array::from_iter(a.iter().map(|v| v.map(|x| x * 1000)))) + .unwrap_or_else(|| Int64Array::from(vec![0i64; arr.len()])); + Arc::new(ns_arr) + } + _ => arrow::compute::cast(arr, &DataType::Int64).unwrap_or_else(|_| arr.clone()), + }; + ns + } + DataType::Time32(_) => { + // Cast to ns int64 + arrow::compute::cast(arr, &DataType::Int64).unwrap_or_else(|_| arr.clone()) + } + DataType::UInt64 => { + arrow::compute::cast(arr, &DataType::Int64).unwrap_or_else(|_| arr.clone()) + } + _ => arr.clone(), + } +} + +/// Convert a Timestamp array to Snowflake's `{epoch: i64, fraction: i32}` struct. +fn timestamp_to_sf_struct(arr: &ArrayRef, unit: &TimeUnit) -> ArrayRef { + let len = arr.len(); + + // Normalize to nanosecond timestamps for uniform epoch/fraction extraction. + let ns_arr = match unit { + TimeUnit::Second => { + arrow::compute::cast(arr, &DataType::Timestamp(TimeUnit::Nanosecond, None)) + } + TimeUnit::Millisecond => { + arrow::compute::cast(arr, &DataType::Timestamp(TimeUnit::Nanosecond, None)) + } + TimeUnit::Microsecond => { + arrow::compute::cast(arr, &DataType::Timestamp(TimeUnit::Nanosecond, None)) + } + TimeUnit::Nanosecond => Ok(arr.clone()), + }; + + let (epochs, fractions): (Vec>, Vec>) = match ns_arr { + Ok(ns) => { + if let Some(ts) = ns.as_any().downcast_ref::() { + (0..len) + .map(|i| { + if ts.is_null(i) { + (None, None) + } else { + let nanos = ts.value(i); + let epoch = nanos / 1_000_000_000; + let fraction = (nanos % 1_000_000_000) as i32; + (Some(epoch), Some(fraction)) + } + }) + .unzip() + } else { + (vec![None; len], vec![None; len]) + } + } + Err(_) => (vec![None; len], vec![None; len]), + }; + + let epoch_arr = Arc::new(Int64Array::from(epochs)) as ArrayRef; + let fraction_arr = Arc::new(Int32Array::from(fractions)) as ArrayRef; + + let has_tz = matches!(arr.data_type(), DataType::Timestamp(_, Some(_))); + + let struct_arr = if has_tz { + let timezone_arr = Arc::new(Int32Array::from(vec![1440i32; len])) as ArrayRef; + StructArray::from(vec![ + ( + Arc::new(Field::new("epoch", DataType::Int64, false)), + epoch_arr, + ), + ( + Arc::new(Field::new("fraction", DataType::Int32, false)), + fraction_arr, + ), + ( + Arc::new(Field::new("timezone", DataType::Int32, false)), + timezone_arr, + ), + ]) + } else { + StructArray::from(vec![ + ( + Arc::new(Field::new("epoch", DataType::Int64, false)), + epoch_arr, + ), + ( + Arc::new(Field::new("fraction", DataType::Int32, false)), + fraction_arr, + ), + ]) + }; + + Arc::new(struct_arr) +} + +// --------------------------------------------------------------------------- +// Arrow IPC stream → base64 +// --------------------------------------------------------------------------- + +pub fn batches_to_arrow_base64(schema: &Arc, batches: &[RecordBatch]) -> String { + let sf_schema = Arc::new(sf_arrow_schema(schema)); + + let sf_batches: Vec = batches + .iter() + .filter_map(|batch| { + let sf_columns: Vec = batch.columns().iter().map(to_sf_array).collect(); + RecordBatch::try_new(sf_schema.clone(), sf_columns).ok() + }) + .collect(); + + let mut buf = Vec::new(); + if let Ok(mut writer) = StreamWriter::try_new(&mut buf, &sf_schema) { + for batch in &sf_batches { + let _ = writer.write(batch); + } + let _ = writer.finish(); + } + base64::engine::general_purpose::STANDARD.encode(&buf) +} + +// --------------------------------------------------------------------------- +// Full Snowflake query success response +// --------------------------------------------------------------------------- + +pub fn sf_query_response( + schema: &Arc, + batches: &[RecordBatch], + total_rows: u64, + query_id: &str, + database: &str, + schema_name: &str, +) -> Value { + let rowtype = schema_to_rowtype(schema); + let rowset_base64 = batches_to_arrow_base64(schema, batches); + + json!({ + "data": { + "parameters": [ + {"name": "TIMEZONE", "value": "Etc/UTC"}, + {"name": "CLIENT_RESULT_CHUNK_SIZE", "value": 160}, + {"name": "CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY", "value": 3600} + ], + "rowtype": rowtype, + "rowsetBase64": rowset_base64, + "total": total_rows, + "returned": total_rows, + "queryId": query_id, + "queryResultFormat": "arrow", + "finalDatabaseName": database, + "finalSchemaName": schema_name + }, + "success": true, + "code": null, + "message": null + }) +} diff --git a/crates/queryflux-frontend/src/snowflake/http/handlers/common.rs b/crates/queryflux-frontend/src/snowflake/http/handlers/common.rs new file mode 100644 index 0000000..00d8b5a --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/handlers/common.rs @@ -0,0 +1,149 @@ +use std::collections::HashMap; +use std::io::Read; + +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; +use bytes::Bytes; +use flate2::read::GzDecoder; +use serde_json::json; +use serde_json::Value; + +/// Snowflake clients (Python connector, JDBC, etc.) send `Content-Encoding: gzip` and gzip the +/// JSON body for most POSTs. Axum does not decompress automatically — decode before `serde_json`. +pub fn decode_snowflake_request_body(headers: &HeaderMap, body: &Bytes) -> Result, String> { + let gzip = headers + .get("content-encoding") + .and_then(|v| v.to_str().ok()) + .map(|s| { + let s = s.trim(); + s.eq_ignore_ascii_case("gzip") || s.eq_ignore_ascii_case("x-gzip") + }) + .unwrap_or(false); + if !gzip { + return Ok(body.to_vec()); + } + let mut decoder = GzDecoder::new(std::io::Cursor::new(body.as_ref())); + let mut out = Vec::new(); + decoder + .read_to_end(&mut out) + .map_err(|e| format!("gzip decompress: {e}"))?; + Ok(out) +} + +/// Parse JSON from a request body, after optional gzip decompression. +pub fn parse_snowflake_json_body(headers: &HeaderMap, body: &Bytes) -> Result { + let decoded = decode_snowflake_request_body(headers, body)?; + serde_json::from_slice(&decoded).map_err(|e| e.to_string()) +} + +/// Extract the Snowflake token from the Authorization header. +/// Expected format: `Authorization: Snowflake Token="{token}"` +pub fn extract_snowflake_token(headers: &HeaderMap) -> Option { + let auth = headers.get("authorization")?.to_str().ok()?; + // Handle both `Snowflake Token="..."` and `Snowflake Token=...` + let rest = auth.strip_prefix("Snowflake Token=")?; + let token = rest.trim_matches('"'); + if token.is_empty() { + None + } else { + Some(token.to_string()) + } +} + +/// Build a Snowflake-style JSON error response. +/// +/// **Always uses HTTP 200** with `success: false` in the body. The official Snowflake Python +/// connector treats many non-2xx status codes as *retryable* during `login-request` (including +/// 400, 403, and all 5xx — see `is_retryable_http_code` in `snowflake.connector.network`). +/// Returning 502/400 for configuration errors caused errno **251012** ("Login request is retryable") +/// and then **250001** after retries. Real Snowflake often responds with 200 + JSON `success: false`. +pub fn sf_error(_status: StatusCode, code: u64, message: &str) -> Response { + ( + StatusCode::OK, + axum::Json(json!({ + "data": null, + "code": code.to_string(), + "message": message, + "success": false + })), + ) + .into_response() +} + +/// Forward a reqwest response as an Axum response, preserving status and headers. +pub fn proxy_response( + status: reqwest::StatusCode, + resp_headers: &reqwest::header::HeaderMap, + body: Bytes, +) -> Response { + let axum_status = + StatusCode::from_u16(status.as_u16()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + + let mut builder = axum::response::Response::builder().status(axum_status); + + for (name, value) in resp_headers { + // Skip hop-by-hop headers that don't make sense to forward. + let name_lower = name.as_str().to_lowercase(); + if matches!( + name_lower.as_str(), + "transfer-encoding" | "connection" | "keep-alive" | "te" | "trailer" | "upgrade" + ) { + continue; + } + if let Ok(val_str) = value.to_str() { + builder = builder.header(name.as_str(), val_str); + } + } + + builder + .body(axum::body::Body::from(body)) + .unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response()) +} + +/// Extract non-auth headers to pass through to the warehouse. +pub fn passthrough_headers(headers: &HeaderMap) -> HashMap { + headers + .iter() + .filter_map(|(k, v)| { + let name = k.as_str().to_lowercase(); + // Never forward Authorization — we inject the service-account token ourselves. + if name == "authorization" || name == "host" || name == "content-length" { + return None; + } + v.to_str() + .ok() + .map(|val| (k.as_str().to_string(), val.to_string())) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::http::HeaderValue; + use flate2::write::GzEncoder; + use flate2::Compression; + use std::io::Write; + + #[test] + fn decodes_gzip_json_body() { + let json = br#"{"data":{"LOGIN_NAME":"u"}}"#; + let mut enc = GzEncoder::new(Vec::new(), Compression::default()); + enc.write_all(json).unwrap(); + let gz = enc.finish().unwrap(); + let bytes = Bytes::from(gz); + let mut headers = HeaderMap::new(); + headers.insert("content-encoding", HeaderValue::from_static("gzip")); + let out = decode_snowflake_request_body(&headers, &bytes).unwrap(); + assert_eq!(out.as_slice(), json); + } + + #[test] + fn passthrough_plain_json_without_gzip_header() { + let raw = br#"{"a":1}"#; + let bytes = Bytes::from_static(raw); + let headers = HeaderMap::new(); + let out = decode_snowflake_request_body(&headers, &bytes).unwrap(); + assert_eq!(out.as_slice(), raw); + } +} diff --git a/crates/queryflux-frontend/src/snowflake/http/handlers/mod.rs b/crates/queryflux-frontend/src/snowflake/http/handlers/mod.rs new file mode 100644 index 0000000..da8a32e --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/handlers/mod.rs @@ -0,0 +1,4 @@ +pub mod common; +pub mod query; +pub mod session; +pub mod token; diff --git a/crates/queryflux-frontend/src/snowflake/http/handlers/query.rs b/crates/queryflux-frontend/src/snowflake/http/handlers/query.rs new file mode 100644 index 0000000..73a42a5 --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/handlers/query.rs @@ -0,0 +1,196 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use arrow::datatypes::Schema; +use arrow::record_batch::RecordBatch; +use async_trait::async_trait; +use axum::extract::State; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; +use bytes::Bytes; +use queryflux_core::{ + error::Result, + query::{FrontendProtocol, QueryStats}, + session::SessionContext, + tags::QueryTags, +}; +use serde_json::Value; +use tracing::warn; +use uuid::Uuid; + +use crate::dispatch::{execute_to_sink, ResultSink}; +use crate::snowflake::http::format::sf_query_response; +use crate::state::AppState; + +use super::common::{decode_snowflake_request_body, extract_snowflake_token, sf_error}; + +// --------------------------------------------------------------------------- +// ResultSink that accumulates Arrow batches into Snowflake JSON format +// --------------------------------------------------------------------------- + +struct SnowflakeSink { + schema: Option>, + batches: Vec, + total_rows: u64, + error: Option, +} + +impl SnowflakeSink { + fn new() -> Self { + Self { + schema: None, + batches: Vec::new(), + total_rows: 0, + error: None, + } + } +} + +#[async_trait] +impl ResultSink for SnowflakeSink { + async fn on_schema(&mut self, schema: &Schema) -> Result<()> { + self.schema = Some(Arc::new(schema.clone())); + Ok(()) + } + + async fn on_batch(&mut self, batch: &RecordBatch) -> Result<()> { + self.total_rows += batch.num_rows() as u64; + self.batches.push(batch.clone()); + Ok(()) + } + + async fn on_complete(&mut self, _stats: &QueryStats) -> Result<()> { + Ok(()) + } + + async fn on_error(&mut self, message: &str) -> Result<()> { + self.error = Some(message.to_string()); + Ok(()) + } +} + +impl SnowflakeSink { + fn into_response(self, query_id: &str, database: &str, schema_name: &str) -> Response { + if let Some(err) = self.error { + return ( + StatusCode::OK, + axum::Json(serde_json::json!({ + "data": { + "errorCode": "100183", + "sqlState": "P0001" + }, + "code": "100183", + "message": err, + "success": false + })), + ) + .into_response(); + } + + let schema = self.schema.unwrap_or_else(|| Arc::new(Schema::empty())); + + let body = sf_query_response( + &schema, + &self.batches, + self.total_rows, + query_id, + database, + schema_name, + ); + (StatusCode::OK, axum::Json(body)).into_response() + } +} + +// --------------------------------------------------------------------------- +// Handlers +// --------------------------------------------------------------------------- + +/// POST /queries/v1/query-request — execute SQL +pub async fn query_request( + State(state): State>, + headers: HeaderMap, + body: Bytes, +) -> Response { + let qf_token = match extract_snowflake_token(&headers) { + Some(t) => t, + None => return sf_error(StatusCode::UNAUTHORIZED, 390001, "Missing token"), + }; + + // Parse SQL from body (body may be gzip per Snowflake Python connector). + let sql: String = decode_snowflake_request_body(&headers, &body) + .ok() + .and_then(|raw| serde_json::from_slice::(&raw).ok()) + .and_then(|v| v["sqlText"].as_str().map(|s| s.to_string())) + .unwrap_or_default(); + + // Clone fields out of the session (must not hold DashMap ref across await). + let (auth_ctx, group, database, schema) = { + match state.snowflake_sessions.get(&qf_token) { + Some(s) => ( + s.auth_ctx.clone(), + s.group.clone(), + s.database.clone().unwrap_or_default(), + s.schema.clone().unwrap_or_default(), + ), + None => return sf_error(StatusCode::UNAUTHORIZED, 390390, "Session not found"), + } + }; + + let session_ctx = SessionContext::MySqlWire { + user: state + .snowflake_sessions + .get(&qf_token) + .and_then(|s| s.user.clone()), + schema: Some(database.clone()), + session_vars: HashMap::new(), + tags: QueryTags::default(), + }; + + let query_id = Uuid::new_v4().to_string(); + let mut sink = SnowflakeSink::new(); + + if let Err(e) = execute_to_sink( + &state, + sql, + session_ctx, + FrontendProtocol::SnowflakeHttp, + group, + &mut sink, + &auth_ctx, + ) + .await + { + warn!(query_id = %query_id, "execute_to_sink error: {e}"); + sink.error = Some(e.to_string()); + } + + sink.into_response(&query_id, &database, &schema) +} + +/// GET /queries/v1/query-monitoring-request — async status poll (stub) +/// +/// The Snowflake Python connector polls this when `asyncExec: true`. +/// For now all queries are executed synchronously; this returns a not-found +/// body (empty queries array) which causes the connector to stop polling. +pub async fn query_monitoring_request( + State(_state): State>, + _headers: HeaderMap, +) -> Response { + ( + StatusCode::OK, + axum::Json(serde_json::json!({ + "data": {"queries": []}, + "success": true + })), + ) + .into_response() +} + +/// DELETE /queries/v1/:query_id — cancel query (no-op for sync execution) +pub async fn cancel_query(State(_state): State>, _headers: HeaderMap) -> Response { + ( + StatusCode::OK, + axum::Json(serde_json::json!({"success": true, "data": null})), + ) + .into_response() +} diff --git a/crates/queryflux-frontend/src/snowflake/http/handlers/session.rs b/crates/queryflux-frontend/src/snowflake/http/handlers/session.rs new file mode 100644 index 0000000..4b638f5 --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/handlers/session.rs @@ -0,0 +1,181 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Instant; + +use axum::extract::{Query, State}; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; +use bytes::Bytes; +use queryflux_auth::Credentials; +use queryflux_core::{query::FrontendProtocol, session::SessionContext, tags::QueryTags}; +use serde_json::{json, Value}; +use tracing::{info, warn}; +use uuid::Uuid; + +use crate::snowflake::http::session_store::SnowflakeSession; +use crate::state::AppState; + +use super::common::{extract_snowflake_token, parse_snowflake_json_body, sf_error}; + +/// POST /session/v1/login-request +/// +/// Authenticates the client against QueryFlux's auth provider, resolves a cluster +/// group via the router chain, and creates a local QF session. No backend Snowflake +/// account is contacted — QueryFlux terminates the Snowflake wire protocol and +/// dispatches SQL to any configured engine (Trino, StarRocks, DuckDB, etc.). +pub async fn login_request( + State(state): State>, + Query(params): Query>, + headers: HeaderMap, + body: Bytes, +) -> Response { + let body_json: Value = match parse_snowflake_json_body(&headers, &body) { + Ok(v) => v, + Err(_) => return sf_error(StatusCode::BAD_REQUEST, 390000, "Invalid JSON body"), + }; + + let data = &body_json["data"]; + let username = data["LOGIN_NAME"].as_str().unwrap_or("").to_string(); + let password = data["PASSWORD"].as_str().map(|s| s.to_string()); + + // Database/schema hints from query params or session parameters. + let database = params.get("databaseName").cloned().or_else(|| { + data["SESSION_PARAMETERS"]["DATABASE"] + .as_str() + .map(|s| s.to_string()) + }); + let schema = params.get("schemaName").cloned().or_else(|| { + data["SESSION_PARAMETERS"]["SCHEMA"] + .as_str() + .map(|s| s.to_string()) + }); + + // Authenticate via QueryFlux auth provider. + let creds = Credentials { + username: Some(username.clone()), + password: password.clone(), + bearer_token: None, + }; + let auth_ctx = match state.auth_provider.authenticate(&creds).await { + Ok(ctx) => ctx, + Err(e) => { + warn!(user = %username, "Snowflake HTTP login auth failed: {e}"); + return sf_error( + StatusCode::UNAUTHORIZED, + 390100, + "Incorrect username or password", + ); + } + }; + + // Route to find the cluster group for this session. + let session_ctx = SessionContext::MySqlWire { + user: Some(username.clone()), + schema: database.clone(), + session_vars: HashMap::new(), + tags: QueryTags::default(), + }; + let group = { + let live = state.live.read().await; + live.router_chain + .route( + "", + &session_ctx, + &FrontendProtocol::SnowflakeHttp, + Some(&auth_ctx), + ) + .await + }; + let group = match group { + Ok(g) => g, + Err(e) => { + warn!(user = %username, "Snowflake HTTP routing failed at login: {e}"); + return sf_error( + StatusCode::INTERNAL_SERVER_ERROR, + 390000, + &format!("Routing error: {e}"), + ); + } + }; + + let qf_token = Uuid::new_v4().to_string(); + state.snowflake_sessions.insert( + qf_token.clone(), + SnowflakeSession { + qf_token: qf_token.clone(), + user: Some(username.clone()), + auth_ctx, + group, + database: database.clone(), + schema: schema.clone(), + created_at: Instant::now(), + }, + ); + + info!(user = %username, token = %&qf_token[..8], "Snowflake HTTP session created"); + + ( + StatusCode::OK, + axum::Json(json!({ + "data": { + "token": qf_token, + "masterToken": qf_token, + "parameters": [ + {"name": "AUTOCOMMIT", "value": true}, + {"name": "CLIENT_SESSION_KEEP_ALIVE_HEARTBEAT_FREQUENCY", "value": 3600}, + {"name": "CLIENT_RESULT_CHUNK_SIZE", "value": 160}, + {"name": "QUERY_RESULT_FORMAT", "value": "ARROW_FORCE"}, + {"name": "TIMEZONE", "value": "Etc/UTC"} + ], + "sessionInfo": { + "databaseName": database.unwrap_or_default(), + "schemaName": schema.unwrap_or_default(), + "warehouseName": "", + "roleName": "PUBLIC" + } + }, + "success": true, + "code": null, + "message": null + })), + ) + .into_response() +} + +/// DELETE /session — log out +pub async fn logout(State(state): State>, headers: HeaderMap) -> Response { + if let Some(token) = extract_snowflake_token(&headers) { + state.snowflake_sessions.remove(&token); + } + ( + StatusCode::OK, + axum::Json(json!({"success": true, "code": null, "message": null, "data": null})), + ) + .into_response() +} + +/// GET /session/heartbeat — keep-alive check +pub async fn heartbeat(State(state): State>, headers: HeaderMap) -> Response { + let token = match extract_snowflake_token(&headers) { + Some(t) => t, + None => { + return sf_error( + StatusCode::UNAUTHORIZED, + 390101, + "Authorization header not found", + ) + } + }; + if !state.snowflake_sessions.contains(&token) { + return sf_error( + StatusCode::UNAUTHORIZED, + 390104, + "Session not found or expired", + ); + } + ( + StatusCode::OK, + axum::Json(json!({"success": true, "code": null, "message": null, "data": null})), + ) + .into_response() +} diff --git a/crates/queryflux-frontend/src/snowflake/http/handlers/token.rs b/crates/queryflux-frontend/src/snowflake/http/handlers/token.rs new file mode 100644 index 0000000..e914bc5 --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/handlers/token.rs @@ -0,0 +1,43 @@ +use std::sync::Arc; + +use crate::state::AppState; +use axum::extract::State; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; + +use super::common::{extract_snowflake_token, sf_error}; + +/// POST /session/token-request +/// +/// In Design B (protocol bridge), QueryFlux issues its own tokens and manages +/// sessions locally — there are no upstream Snowflake warehouse tokens to renew. +/// This endpoint simply validates the existing session and returns success. +pub async fn token_request(State(state): State>, headers: HeaderMap) -> Response { + let token = match extract_snowflake_token(&headers) { + Some(t) => t, + None => { + return sf_error( + StatusCode::UNAUTHORIZED, + 390101, + "Authorization header not found", + ) + } + }; + if !state.snowflake_sessions.contains(&token) { + return sf_error( + StatusCode::UNAUTHORIZED, + 390104, + "Session not found or expired", + ); + } + ( + StatusCode::OK, + axum::Json(serde_json::json!({ + "data": {"sessionToken": token, "validityInSecondsST": 3600}, + "success": true, + "code": null, + "message": null + })), + ) + .into_response() +} diff --git a/crates/queryflux-frontend/src/snowflake/http/mod.rs b/crates/queryflux-frontend/src/snowflake/http/mod.rs new file mode 100644 index 0000000..03a2faf --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/mod.rs @@ -0,0 +1,33 @@ +//! Snowflake Wire Protocol frontend (Form 1) — Design B: protocol bridge. +//! +//! Exposes `routes()` — a stateless `Router>` that can be merged with +//! other route sets and have state injected at the top level. + +use std::sync::Arc; + +use axum::{ + routing::{delete, get, post}, + Router, +}; + +use crate::state::AppState; + +use handlers::{query, session, token}; + +pub mod format; +pub mod handlers; +pub mod session_store; + +pub fn routes() -> Router> { + Router::new() + .route("/session/v1/login-request", post(session::login_request)) + .route("/session", delete(session::logout)) + .route("/session/heartbeat", get(session::heartbeat)) + .route("/session/token-request", post(token::token_request)) + .route("/queries/v1/query-request", post(query::query_request)) + .route( + "/queries/v1/query-monitoring-request", + get(query::query_monitoring_request), + ) + .route("/queries/v1/{query_id}", delete(query::cancel_query)) +} diff --git a/crates/queryflux-frontend/src/snowflake/http/session_store.rs b/crates/queryflux-frontend/src/snowflake/http/session_store.rs new file mode 100644 index 0000000..7b2990a --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/session_store.rs @@ -0,0 +1,52 @@ +use std::sync::Arc; +use std::time::Instant; + +use dashmap::DashMap; +use queryflux_auth::AuthContext; +use queryflux_core::query::ClusterGroupName; + +/// Stores active QueryFlux Snowflake wire-protocol sessions keyed by the qf_token +/// issued to the client at login. Sessions are local to QueryFlux — no backend +/// Snowflake account is needed. +pub struct SnowflakeSessionStore { + sessions: DashMap, +} + +pub struct SnowflakeSession { + pub qf_token: String, + pub user: Option, + pub auth_ctx: AuthContext, + /// Cluster group resolved at login time (via the router chain). + pub group: ClusterGroupName, + /// Database/schema hints from the login request (SESSION_PARAMETERS or query params). + pub database: Option, + pub schema: Option, + pub created_at: Instant, +} + +impl SnowflakeSessionStore { + pub fn new() -> Arc { + Arc::new(Self { + sessions: DashMap::new(), + }) + } + + pub fn insert(&self, token: String, session: SnowflakeSession) { + self.sessions.insert(token, session); + } + + pub fn get( + &self, + token: &str, + ) -> Option> { + self.sessions.get(token) + } + + pub fn remove(&self, token: &str) { + self.sessions.remove(token); + } + + pub fn contains(&self, token: &str) -> bool { + self.sessions.contains_key(token) + } +} diff --git a/crates/queryflux-frontend/src/snowflake/mod.rs b/crates/queryflux-frontend/src/snowflake/mod.rs new file mode 100644 index 0000000..2028013 --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/mod.rs @@ -0,0 +1,50 @@ +pub mod http; +pub mod sql_api; + +use std::sync::Arc; + +use axum::Router; +use queryflux_core::error::{QueryFluxError, Result}; +use tracing::info; + +use crate::state::AppState; +use crate::FrontendListenerTrait; + +/// Combined Snowflake frontend — wire protocol (Form 1) + SQL REST API v2 (Form 2) +/// on a single port with a single shared `Arc`. +pub struct SnowflakeFrontend { + state: Arc, + port: u16, +} + +impl SnowflakeFrontend { + pub fn new(state: Arc, port: u16) -> Self { + Self { state, port } + } + + fn router(&self) -> Router { + http::routes() + .merge(sql_api::routes()) + .with_state(self.state.clone()) + } +} + +#[async_trait::async_trait] +impl FrontendListenerTrait for SnowflakeFrontend { + async fn listen(&self) -> Result<()> { + let addr: std::net::SocketAddr = format!("0.0.0.0:{}", self.port) + .parse() + .map_err(|e: std::net::AddrParseError| QueryFluxError::Other(e.into()))?; + + info!("Snowflake frontend (wire + SQL API) listening on {addr}"); + + axum::serve( + tokio::net::TcpListener::bind(addr) + .await + .map_err(|e| QueryFluxError::Other(e.into()))?, + self.router(), + ) + .await + .map_err(|e| QueryFluxError::Other(e.into())) + } +} diff --git a/crates/queryflux-frontend/src/snowflake/proxy.rs b/crates/queryflux-frontend/src/snowflake/proxy.rs new file mode 100644 index 0000000..b2e594c --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/proxy.rs @@ -0,0 +1,73 @@ +use std::collections::HashMap; + +use bytes::Bytes; +use reqwest::{Client, Method, Response}; + +/// Parameters for [`SnowflakeProxy::forward`]. +pub struct SnowflakeForward<'a> { + pub method: Method, + pub warehouse_base_url: &'a str, + pub path: &'a str, + pub sf_token: Option<&'a str>, + pub query_string: Option<&'a str>, + pub body: Option, + pub passthrough_headers: HashMap, +} + +/// Shared HTTP client for forwarding Snowflake wire-protocol requests to backend warehouses +/// with service-account token injection. +pub struct SnowflakeProxy { + client: Client, +} + +impl Default for SnowflakeProxy { + fn default() -> Self { + Self::new() + } +} + +impl SnowflakeProxy { + pub fn new() -> Self { + Self { + client: Client::builder() + .build() + .expect("failed to build SnowflakeProxy reqwest client"), + } + } + + /// Forward a request to a Snowflake warehouse. + /// + /// When `sf_token` is `Some`, injects `Authorization: Snowflake Token="{sf_token}"`, + /// replacing any Authorization header in `passthrough_headers`. + /// When `sf_token` is `None` (e.g. during login), no Authorization header is injected. + pub async fn forward(&self, req: SnowflakeForward<'_>) -> reqwest::Result { + let url = match req.query_string { + Some(qs) if !qs.is_empty() => { + format!("{}{}?{qs}", req.warehouse_base_url, req.path) + } + _ => format!("{}{}", req.warehouse_base_url, req.path), + }; + + let mut http = self.client.request(req.method, &url); + + // Pass through headers, skipping Authorization (we inject below if token provided). + for (k, v) in &req.passthrough_headers { + if k.to_lowercase() != "authorization" { + http = http.header(k.as_str(), v.as_str()); + } + } + + if let Some(token) = req.sf_token { + http = http.header("Authorization", format!("Snowflake Token=\"{token}\"")); + http = http.header("X-Snowflake-Authorization-Token-Type", "TOKEN"); + } + + if let Some(body_bytes) = req.body { + http = http + .header("Content-Type", "application/json") + .body(body_bytes); + } + + http.send().await + } +} diff --git a/crates/queryflux-frontend/src/snowflake/sql_api/auth.rs b/crates/queryflux-frontend/src/snowflake/sql_api/auth.rs new file mode 100644 index 0000000..55d00a7 --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/sql_api/auth.rs @@ -0,0 +1,70 @@ +//! Service-account JWT generation for the Snowflake SQL REST API v2. +//! +//! Snowflake SQL API v2 uses key-pair authentication: the caller signs a JWT +//! with their RSA private key and presents it as `Authorization: Bearer {jwt}`. +//! +//! JWT claims (per Snowflake docs): +//! iss: "{account}.{user}.SHA256:{public_key_fingerprint}" +//! sub: "{account}.{user}" +//! iat: +//! exp: +//! +//! The `public_key_fingerprint` is the SHA256 hash of the SPKI DER-encoded public key, +//! base64-encoded without padding. + +use jsonwebtoken::{encode, Algorithm, EncodingKey, Header}; +use rsa::pkcs8::{DecodePrivateKey, EncodePublicKey}; +use rsa::RsaPrivateKey; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; + +#[derive(Debug, Serialize, Deserialize)] +struct SnowflakeJwtClaims { + iss: String, + sub: String, + iat: i64, + exp: i64, +} + +/// Generate a Snowflake SQL API v2 service-account JWT from the cluster's key-pair credentials. +/// +/// - `account`: Snowflake account identifier (e.g. `"xy12345"`) +/// - `username`: Service-account username (e.g. `"QUERYFLUX_SA"`) +/// - `private_key_pem`: PKCS#8 PEM-encoded RSA private key +pub fn generate_service_account_jwt( + account: &str, + username: &str, + private_key_pem: &str, +) -> Result { + let private_key = RsaPrivateKey::from_pkcs8_pem(private_key_pem) + .map_err(|e| format!("Failed to parse RSA private key: {e}"))?; + + let public_key = private_key.to_public_key(); + let public_key_der = public_key + .to_public_key_der() + .map_err(|e| format!("Failed to encode public key as DER: {e}"))?; + + // SHA256 fingerprint of the SPKI DER public key, base64-encoded (no padding). + let fingerprint = { + let hash = Sha256::digest(public_key_der.as_bytes()); + use base64::Engine; + base64::engine::general_purpose::STANDARD_NO_PAD.encode(hash) + }; + + let now = chrono::Utc::now().timestamp(); + let account_upper = account.to_uppercase(); + let user_upper = username.to_uppercase(); + + let claims = SnowflakeJwtClaims { + iss: format!("{account_upper}.{user_upper}.SHA256:{fingerprint}"), + sub: format!("{account_upper}.{user_upper}"), + iat: now, + exp: now + 60, + }; + + let encoding_key = EncodingKey::from_rsa_pem(private_key_pem.as_bytes()) + .map_err(|e| format!("Failed to build encoding key: {e}"))?; + + encode(&Header::new(Algorithm::RS256), &claims, &encoding_key) + .map_err(|e| format!("JWT encoding failed: {e}")) +} diff --git a/crates/queryflux-frontend/src/snowflake/sql_api/handlers.rs b/crates/queryflux-frontend/src/snowflake/sql_api/handlers.rs new file mode 100644 index 0000000..b362f63 --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/sql_api/handlers.rs @@ -0,0 +1,257 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use arrow::array::Array; +use arrow::compute::cast as arrow_cast; +use arrow::datatypes::{DataType, Schema}; +use arrow::record_batch::RecordBatch; +use async_trait::async_trait; +use axum::extract::State; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; +use bytes::Bytes; +use queryflux_auth::Credentials; +use queryflux_core::{ + error::Result, + query::{FrontendProtocol, QueryStats}, + session::SessionContext, + tags::QueryTags, +}; +use serde_json::{json, Value}; +use tracing::warn; +use uuid::Uuid; + +use crate::dispatch::{execute_to_sink, ResultSink}; +use crate::snowflake::http::format::schema_to_rowtype; +use crate::snowflake::http::handlers::common::{parse_snowflake_json_body, sf_error}; +use crate::state::AppState; + +// --------------------------------------------------------------------------- +// ResultSink that accumulates Arrow batches into SQL API v2 jsonv2 format +// --------------------------------------------------------------------------- + +struct SqlApiSink { + schema: Option>, + rows: Vec>, + error: Option, +} + +impl SqlApiSink { + fn new() -> Self { + Self { + schema: None, + rows: Vec::new(), + error: None, + } + } +} + +#[async_trait] +impl ResultSink for SqlApiSink { + async fn on_schema(&mut self, schema: &Schema) -> Result<()> { + self.schema = Some(Arc::new(schema.clone())); + Ok(()) + } + + async fn on_batch(&mut self, batch: &RecordBatch) -> Result<()> { + for row_idx in 0..batch.num_rows() { + let row: Vec = (0..batch.num_columns()) + .map(|col_idx| array_value_to_json(batch.column(col_idx).as_ref(), row_idx)) + .collect(); + self.rows.push(row); + } + Ok(()) + } + + async fn on_complete(&mut self, _stats: &QueryStats) -> Result<()> { + Ok(()) + } + + async fn on_error(&mut self, message: &str) -> Result<()> { + self.error = Some(message.to_string()); + Ok(()) + } +} + +impl SqlApiSink { + fn into_response(self, handle: &str) -> Response { + if let Some(err) = self.error { + return ( + StatusCode::OK, + axum::Json(json!({ + "code": "002043", + "message": err, + "sqlState": "P0001", + "statementHandle": handle + })), + ) + .into_response(); + } + + let schema = self.schema.unwrap_or_else(|| Arc::new(Schema::empty())); + let num_rows = self.rows.len() as u64; + let rowtype = schema_to_rowtype(&schema); + + ( + StatusCode::OK, + axum::Json(json!({ + "statementHandle": handle, + "message": "Statement executed successfully.", + "createdOn": chrono::Utc::now().timestamp_millis(), + "statementStatusUrl": format!("/api/v2/statements/{handle}"), + "resultSetMetaData": { + "numRows": num_rows, + "format": "jsonv2", + "rowType": rowtype, + "partitionInfo": [{"rowCount": num_rows, "uncompressedSize": 0}] + }, + "data": self.rows + })), + ) + .into_response() + } +} + +/// Serialize a single array value at `row` to a SQL API v2 JSON value (string or null). +fn array_value_to_json(arr: &dyn Array, row: usize) -> Value { + if arr.is_null(row) { + return Value::Null; + } + // Cast to Utf8 for uniform string serialization. + let string_arr = arrow_cast(arr, &DataType::Utf8); + match string_arr { + Ok(s) => { + if let Some(str_arr) = s.as_any().downcast_ref::() { + Value::String(str_arr.value(row).to_string()) + } else { + Value::Null + } + } + Err(_) => Value::String(format!("{:?}", arr.data_type())), + } +} + +// --------------------------------------------------------------------------- +// Handlers +// --------------------------------------------------------------------------- + +/// POST /api/v2/statements — submit SQL, execute synchronously, return jsonv2 +pub async fn submit_statement( + State(state): State>, + headers: HeaderMap, + body: Bytes, +) -> Response { + let body_json: Value = match parse_snowflake_json_body(&headers, &body) { + Ok(v) => v, + Err(_) => return sf_error(StatusCode::BAD_REQUEST, 390000, "Invalid JSON body"), + }; + let sql = body_json["statement"].as_str().unwrap_or("").to_string(); + + // Stateless auth: Bearer token in Authorization header. + let auth_ctx = match authenticate(&state, &headers).await { + Ok(ctx) => ctx, + Err(e) => return sf_error(StatusCode::UNAUTHORIZED, 390002, &e.to_string()), + }; + + let session_ctx = SessionContext::MySqlWire { + user: Some(auth_ctx.user.clone()), + schema: None, + session_vars: HashMap::new(), + tags: QueryTags::default(), + }; + let group = { + let live = state.live.read().await; + live.router_chain + .route( + &sql, + &session_ctx, + &FrontendProtocol::SnowflakeSqlApi, + Some(&auth_ctx), + ) + .await + }; + let group = match group { + Ok(g) => g, + Err(e) => return sf_error(StatusCode::BAD_GATEWAY, 390000, &e.to_string()), + }; + + let handle = Uuid::new_v4().to_string(); + let mut sink = SqlApiSink::new(); + + if let Err(e) = execute_to_sink( + &state, + sql, + session_ctx, + FrontendProtocol::SnowflakeSqlApi, + group, + &mut sink, + &auth_ctx, + ) + .await + { + warn!(handle = %handle, "SQL API execute_to_sink error: {e}"); + sink.error = Some(e.to_string()); + } + + sink.into_response(&handle) +} + +/// GET /api/v2/statements/:handle — stub (sync execution, nothing to poll) +pub async fn get_statement( + State(_state): State>, + _headers: HeaderMap, + axum::extract::Path(handle): axum::extract::Path, + _raw_query: axum::extract::RawQuery, +) -> Response { + ( + StatusCode::NOT_FOUND, + axum::Json(json!({ + "code": "390142", + "message": format!("Statement handle {handle} not found or already complete."), + "sqlState": "02000", + "statementHandle": handle + })), + ) + .into_response() +} + +/// DELETE /api/v2/statements/:handle — stub (sync execution, nothing to cancel) +pub async fn cancel_statement( + State(_state): State>, + _headers: HeaderMap, + axum::extract::Path(handle): axum::extract::Path, +) -> Response { + ( + StatusCode::OK, + axum::Json(json!({ + "statementHandle": handle, + "message": "Statement aborted.", + })), + ) + .into_response() +} + +// --------------------------------------------------------------------------- +// Auth helper +// --------------------------------------------------------------------------- + +async fn authenticate( + state: &Arc, + headers: &HeaderMap, +) -> std::result::Result { + let bearer = headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.strip_prefix("Bearer ")) + .map(|s| s.to_string()); + + state + .auth_provider + .authenticate(&Credentials { + username: None, + password: None, + bearer_token: bearer, + }) + .await + .map_err(|e| e.to_string()) +} diff --git a/crates/queryflux-frontend/src/snowflake/sql_api/mod.rs b/crates/queryflux-frontend/src/snowflake/sql_api/mod.rs new file mode 100644 index 0000000..70febe4 --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/sql_api/mod.rs @@ -0,0 +1,23 @@ +//! Snowflake SQL REST API v2 frontend (Form 2) — Design B: protocol bridge. +//! +//! Exposes `routes()` — a stateless `Router>`. + +use std::sync::Arc; + +use axum::{ + routing::{get, post}, + Router, +}; + +use crate::state::AppState; + +pub mod handlers; + +pub fn routes() -> Router> { + Router::new() + .route("/api/v2/statements", post(handlers::submit_statement)) + .route( + "/api/v2/statements/{handle}", + get(handlers::get_statement).delete(handlers::cancel_statement), + ) +} diff --git a/crates/queryflux-frontend/src/snowflake/tests.rs b/crates/queryflux-frontend/src/snowflake/tests.rs new file mode 100644 index 0000000..d088ed0 --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/tests.rs @@ -0,0 +1,816 @@ +//! Unit tests for the Snowflake wire-protocol frontend (Design B — protocol bridge). +//! +//! Architecture: every test spins up an in-process Axum server on a random port using +//! the real `SnowflakeHttpFrontend::router()`. A plain `reqwest` client sends requests; +//! tests assert on JSON bodies without mocking any HTTP path. The backend is a +//! `MockAdapter` defined below — no native engine library is needed. +//! +//! Coverage: +//! - Session lifecycle: login → heartbeat → token renewal → logout +//! - Auth: missing token, invalid token after logout +//! - gzip body decoding (Python connector always sends gzip POSTs) +//! - `sf_error` always returns HTTP 200 (prevents 251012 retry loop) +//! - Query execute: success, missing token, stale token, syntax error, gzip body +//! - Query monitoring stub (empty queries array) +//! - Cancel no-op (always returns success: true) +//! - `schema_to_rowtype` mapping for Int64, Utf8, Float64 +//! - `batches_to_arrow_base64` round-trips through Arrow IPC + +#[cfg(test)] +mod snowflake_frontend { + use std::collections::HashMap; + use std::io::Write; + use std::sync::Arc; + use std::time::Duration; + + use arrow::array::{Float64Array, Int64Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use arrow_ipc::reader::StreamReader; + use async_trait::async_trait; + use axum::Router; + use base64::Engine as _; + use bytes::Bytes; + use flate2::write::GzEncoder; + use flate2::Compression; + use futures::stream; + use queryflux_auth::{ + AllowAllAuthorization, AuthProvider, AuthorizationChecker, BackendIdentityResolver, + NoneAuthProvider, QueryCredentials, + }; + use queryflux_cluster_manager::{ + cluster_state::ClusterState, simple::SimpleClusterGroupManager, + strategy::strategy_from_config, + }; + use queryflux_core::{ + catalog::TableSchema as CoreTableSchema, + error::{QueryFluxError, Result as QfResult}, + query::{ + BackendQueryId, ClusterGroupName, ClusterName, EngineType, QueryExecution, + QueryPollResult, + }, + session::SessionContext, + tags::QueryTags, + }; + use queryflux_engine_adapters::{ArrowStream, EngineAdapterTrait}; + use queryflux_metrics::{ClusterSnapshot, MetricsStore, QueryRecord}; + use queryflux_persistence::in_memory::InMemoryPersistence; + use queryflux_routing::{ + chain::RouterChain, implementations::protocol_based::ProtocolBasedRouter, + }; + use queryflux_translation::TranslationService; + use serde_json::Value; + use tokio::net::TcpListener; + + use crate::snowflake::http::{ + format::{batches_to_arrow_base64, schema_to_rowtype}, + session_store::SnowflakeSessionStore, + SnowflakeHttpFrontend, + }; + use crate::state::{AppState, LiveConfig}; + + // ------------------------------------------------------------------------- + // Noop MetricsStore + // ------------------------------------------------------------------------- + + struct NoopMetrics; + + #[async_trait] + impl MetricsStore for NoopMetrics { + async fn record_query(&self, _r: QueryRecord) -> QfResult<()> { + Ok(()) + } + async fn record_cluster_snapshot(&self, _s: ClusterSnapshot) -> QfResult<()> { + Ok(()) + } + } + + // ------------------------------------------------------------------------- + // MockAdapter — returns one row `{n: 1}` for any query; rejects "SELEKT" + // ------------------------------------------------------------------------- + + struct MockAdapter; + + #[async_trait] + impl EngineAdapterTrait for MockAdapter { + async fn submit_query( + &self, + _sql: &str, + _session: &SessionContext, + _credentials: &QueryCredentials, + _tags: &QueryTags, + ) -> QfResult { + Err(QueryFluxError::Engine("submit_query not used".into())) + } + + async fn poll_query( + &self, + _id: &BackendQueryId, + _next_uri: Option<&str>, + ) -> QfResult { + Err(QueryFluxError::Engine("poll_query not used".into())) + } + + async fn cancel_query(&self, _id: &BackendQueryId) -> QfResult<()> { + Ok(()) + } + + async fn health_check(&self) -> bool { + true + } + + fn engine_type(&self) -> EngineType { + EngineType::DuckDb + } + + fn supports_async(&self) -> bool { + false + } + + /// Returns one row `{n: 1}` unless SQL contains "SELEKT" (simulated syntax error). + async fn execute_as_arrow( + &self, + sql: &str, + _session: &SessionContext, + _credentials: &QueryCredentials, + _tags: &QueryTags, + ) -> QfResult { + if sql.to_uppercase().contains("SELEKT") { + return Err(QueryFluxError::Engine("syntax error near 'SELEKT'".into())); + } + + let schema = Arc::new(Schema::new(vec![Field::new("n", DataType::Int64, false)])); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int64Array::from(vec![1i64]))]) + .map_err(|e| QueryFluxError::Engine(e.to_string()))?; + + let s = stream::once(async move { Ok(batch) }); + Ok(Box::pin(s)) + } + + async fn list_catalogs(&self) -> QfResult> { + Ok(vec![]) + } + + async fn list_databases(&self, _catalog: &str) -> QfResult> { + Ok(vec![]) + } + + async fn list_tables(&self, _catalog: &str, _database: &str) -> QfResult> { + Ok(vec![]) + } + + async fn describe_table( + &self, + _catalog: &str, + _database: &str, + _table: &str, + ) -> QfResult> { + Ok(None) + } + } + + // ------------------------------------------------------------------------- + // Server bootstrap + // ------------------------------------------------------------------------- + + /// Spins up a `SnowflakeHttpFrontend` on a random port backed by `MockAdapter`. + /// Returns `(port, shutdown_guard)` — drop the guard to stop the server. + async fn start_server() -> (u16, tokio::sync::oneshot::Sender<()>) { + let group = ClusterGroupName("mock".to_string()); + let cluster = ClusterName("mock-1".to_string()); + + let adapter = Arc::new(MockAdapter) as Arc; + + let state = Arc::new(ClusterState::new( + cluster.clone(), + group.clone(), + None, + None, + EngineType::DuckDb, + None, + 16, + true, + )); + + let mut group_states = HashMap::new(); + group_states.insert(group.clone(), (vec![state], strategy_from_config(None))); + + let mut adapters = HashMap::new(); + adapters.insert(cluster.0.clone(), adapter); + + let mut group_members = HashMap::new(); + group_members.insert("mock".to_string(), vec![cluster.0.clone()]); + + let protocol_router: Box = + Box::new(ProtocolBasedRouter { + trino_http: None, + postgres_wire: None, + mysql_wire: None, + clickhouse_http: None, + flight_sql: None, + snowflake_http: Some(group.clone()), + snowflake_sql_api: None, + }); + + let live_config = LiveConfig { + router_chain: RouterChain::new(vec![protocol_router], group.clone()), + cluster_manager: Arc::new(SimpleClusterGroupManager::new(group_states)), + adapters, + health_check_targets: vec![], + cluster_configs: HashMap::new(), + group_members, + group_order: vec!["mock".to_string()], + group_translation_scripts: HashMap::new(), + group_default_tags: HashMap::new(), + }; + + let app_state = Arc::new(AppState { + external_address: "http://127.0.0.1".to_string(), + live: Arc::new(tokio::sync::RwLock::new(live_config)), + persistence: Arc::new(InMemoryPersistence::new()), + translation: Arc::new(TranslationService::disabled()), + metrics: Arc::new(NoopMetrics), + auth_provider: Arc::new(NoneAuthProvider::new(false)) as Arc, + authorization: Arc::new(AllowAllAuthorization) as Arc, + identity_resolver: Arc::new(BackendIdentityResolver::new()), + snowflake_sessions: SnowflakeSessionStore::new(), + }); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + + let router: Router = SnowflakeHttpFrontend::new(app_state, port).router(); + let (tx, rx) = tokio::sync::oneshot::channel::<()>(); + tokio::spawn(async move { + axum::serve(listener, router) + .with_graceful_shutdown(async { + let _ = rx.await; + }) + .await + .ok(); + }); + + tokio::time::sleep(Duration::from_millis(20)).await; + (port, tx) + } + + // ------------------------------------------------------------------------- + // Helpers + // ------------------------------------------------------------------------- + + fn base_url(port: u16) -> String { + format!("http://127.0.0.1:{port}") + } + + fn auth_header(token: &str) -> String { + format!("Snowflake Token=\"{token}\"") + } + + fn login_body() -> serde_json::Value { + serde_json::json!({ + "data": { + "CLIENT_APP_ID": "test", + "CLIENT_APP_VERSION": "1.0", + "LOGIN_NAME": "testuser", + "PASSWORD": "testpass", + "AUTHENTICATOR": "SNOWFLAKE" + } + }) + } + + async fn do_login(client: &reqwest::Client, base: &str) -> Value { + client + .post(format!("{base}/session/v1/login-request")) + .json(&login_body()) + .send() + .await + .unwrap() + .json() + .await + .unwrap() + } + + // ------------------------------------------------------------------------- + // sf_error HTTP-status invariant + // ------------------------------------------------------------------------- + + /// `sf_error` must always produce HTTP 200 regardless of the `StatusCode` + /// argument — the Snowflake Python connector retries on 4xx/5xx (errno 251012). + #[tokio::test] + async fn sf_error_always_returns_http_200() { + use crate::snowflake::http::handlers::common::sf_error; + use axum::http::StatusCode; + use axum::response::IntoResponse; + + let resp = sf_error(StatusCode::BAD_GATEWAY, 390000, "test error").into_response(); + assert_eq!(resp.status(), StatusCode::OK); + } + + // ------------------------------------------------------------------------- + // Login + // ------------------------------------------------------------------------- + + #[tokio::test] + async fn login_success_returns_token() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + let body = do_login(&client, &base_url(port)).await; + + assert_eq!(body["success"], true); + let token = body["data"]["token"].as_str().unwrap(); + assert!(!token.is_empty()); + } + + #[tokio::test] + async fn login_includes_required_session_parameters() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + let body = do_login(&client, &base_url(port)).await; + + let names: Vec<&str> = body["data"]["parameters"] + .as_array() + .unwrap() + .iter() + .filter_map(|p| p["name"].as_str()) + .collect(); + assert!(names.contains(&"AUTOCOMMIT")); + assert!(names.contains(&"QUERY_RESULT_FORMAT")); + } + + #[tokio::test] + async fn login_accepts_gzip_body() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + + let json_bytes = serde_json::to_vec(&login_body()).unwrap(); + let mut enc = GzEncoder::new(Vec::new(), Compression::default()); + enc.write_all(&json_bytes).unwrap(); + let gz = enc.finish().unwrap(); + + let resp = client + .post(format!("{}/session/v1/login-request", base_url(port))) + .header("Content-Type", "application/json") + .header("Content-Encoding", "gzip") + .body(Bytes::from(gz)) + .send() + .await + .unwrap(); + + assert_eq!(resp.status(), 200); + let body: Value = resp.json().await.unwrap(); + assert_eq!(body["success"], true); + } + + /// A malformed body must return HTTP 200 with `success: false` — not HTTP 400. + #[tokio::test] + async fn login_malformed_body_returns_200_failure() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + + let resp = client + .post(format!("{}/session/v1/login-request", base_url(port))) + .header("Content-Type", "application/json") + .body("not valid json !!") + .send() + .await + .unwrap(); + + assert_eq!(resp.status(), 200, "error must be 200 not 4xx"); + let body: Value = resp.json().await.unwrap(); + assert_eq!(body["success"], false); + } + + // ------------------------------------------------------------------------- + // Heartbeat + // ------------------------------------------------------------------------- + + #[tokio::test] + async fn heartbeat_valid_session_returns_success() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + let login = do_login(&client, &base_url(port)).await; + let token = login["data"]["token"].as_str().unwrap(); + + let body: Value = client + .get(format!("{}/session/heartbeat", base_url(port))) + .header("Authorization", auth_header(token)) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!(body["success"], true); + } + + #[tokio::test] + async fn heartbeat_unknown_token_returns_200_failure() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + + let resp = client + .get(format!("{}/session/heartbeat", base_url(port))) + .header("Authorization", auth_header("not-a-real-token")) + .send() + .await + .unwrap(); + + assert_eq!(resp.status(), 200, "must be 200 not 401"); + let body: Value = resp.json().await.unwrap(); + assert_eq!(body["success"], false); + } + + #[tokio::test] + async fn heartbeat_missing_auth_returns_failure() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + + let body: Value = client + .get(format!("{}/session/heartbeat", base_url(port))) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!(body["success"], false); + } + + // ------------------------------------------------------------------------- + // Logout + // ------------------------------------------------------------------------- + + #[tokio::test] + async fn logout_removes_session() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + let login = do_login(&client, &base_url(port)).await; + let token = login["data"]["token"].as_str().unwrap(); + + let logout: Value = client + .delete(format!("{}/session", base_url(port))) + .header("Authorization", auth_header(token)) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(logout["success"], true); + + // Heartbeat must fail after logout. + let hb: Value = client + .get(format!("{}/session/heartbeat", base_url(port))) + .header("Authorization", auth_header(token)) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + assert_eq!(hb["success"], false, "session must be gone after logout"); + } + + // ------------------------------------------------------------------------- + // Token renewal + // ------------------------------------------------------------------------- + + #[tokio::test] + async fn token_renewal_returns_same_token() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + let login = do_login(&client, &base_url(port)).await; + let token = login["data"]["token"].as_str().unwrap(); + + let body: Value = client + .post(format!("{}/session/token-request", base_url(port))) + .header("Authorization", auth_header(token)) + .json(&serde_json::json!({})) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!(body["success"], true); + assert_eq!(body["data"]["sessionToken"].as_str().unwrap(), token); + assert!(body["data"]["validityInSecondsST"].as_u64().unwrap() > 0); + } + + #[tokio::test] + async fn token_renewal_with_invalid_token_returns_failure() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + + let body: Value = client + .post(format!("{}/session/token-request", base_url(port))) + .header("Authorization", auth_header("bogus")) + .json(&serde_json::json!({})) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!(body["success"], false); + } + + // ------------------------------------------------------------------------- + // Query execute + // ------------------------------------------------------------------------- + + #[tokio::test] + async fn query_select_returns_correct_structure() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + let login = do_login(&client, &base_url(port)).await; + let token = login["data"]["token"].as_str().unwrap(); + + let body: Value = client + .post(format!("{}/queries/v1/query-request", base_url(port))) + .header("Authorization", auth_header(token)) + .json(&serde_json::json!({"sqlText": "SELECT 1 AS n", "sequenceId": 1})) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!(body["success"], true, "query failed: {body}"); + assert_eq!(body["data"]["total"], 1, "expected one row"); + + let rowtype = body["data"]["rowtype"].as_array().unwrap(); + assert_eq!(rowtype.len(), 1); + assert_eq!(rowtype[0]["name"].as_str().unwrap(), "n"); + + let b64 = body["data"]["rowsetBase64"].as_str().unwrap(); + assert!(!b64.is_empty(), "rowsetBase64 must be present"); + } + + #[tokio::test] + async fn query_accepts_gzip_body() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + let login = do_login(&client, &base_url(port)).await; + let token = login["data"]["token"].as_str().unwrap(); + + let json = serde_json::to_vec(&serde_json::json!({"sqlText": "SELECT 1 AS val"})).unwrap(); + let mut enc = GzEncoder::new(Vec::new(), Compression::default()); + enc.write_all(&json).unwrap(); + let gz = enc.finish().unwrap(); + + let body: Value = client + .post(format!("{}/queries/v1/query-request", base_url(port))) + .header("Authorization", auth_header(token)) + .header("Content-Encoding", "gzip") + .body(Bytes::from(gz)) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!(body["success"], true, "gzip query must succeed: {body}"); + assert_eq!(body["data"]["total"], 1); + } + + #[tokio::test] + async fn query_missing_token_returns_failure() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + + let body: Value = client + .post(format!("{}/queries/v1/query-request", base_url(port))) + .json(&serde_json::json!({"sqlText": "SELECT 1"})) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!(body["success"], false); + } + + #[tokio::test] + async fn query_stale_token_returns_failure() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + + let body: Value = client + .post(format!("{}/queries/v1/query-request", base_url(port))) + .header("Authorization", auth_header("expired-or-wrong")) + .json(&serde_json::json!({"sqlText": "SELECT 1"})) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!(body["success"], false); + } + + /// SQL that triggers a backend error returns `success: false` with an + /// error message, never a 5xx status. + #[tokio::test] + async fn query_backend_error_returns_graceful_failure() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + let login = do_login(&client, &base_url(port)).await; + let token = login["data"]["token"].as_str().unwrap(); + + let body: Value = client + .post(format!("{}/queries/v1/query-request", base_url(port))) + .header("Authorization", auth_header(token)) + .json(&serde_json::json!({"sqlText": "SELEKT * FORM bad"})) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!(body["success"], false); + let msg = body["message"].as_str().unwrap_or(""); + assert!(!msg.is_empty(), "error message must be non-empty"); + } + + // ------------------------------------------------------------------------- + // Monitoring + Cancel stubs + // ------------------------------------------------------------------------- + + #[tokio::test] + async fn query_monitoring_stub_returns_empty_queries() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + let login = do_login(&client, &base_url(port)).await; + let token = login["data"]["token"].as_str().unwrap(); + + let body: Value = client + .get(format!( + "{}/queries/v1/query-monitoring-request", + base_url(port) + )) + .header("Authorization", auth_header(token)) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!(body["success"], true); + assert_eq!(body["data"]["queries"].as_array().unwrap().len(), 0); + } + + #[tokio::test] + async fn cancel_is_a_noop_and_returns_success() { + let (port, _guard) = start_server().await; + let client = reqwest::Client::new(); + let login = do_login(&client, &base_url(port)).await; + let token = login["data"]["token"].as_str().unwrap(); + + let body: Value = client + .delete(format!( + "{}/queries/v1/some-query-id-that-does-not-exist", + base_url(port) + )) + .header("Authorization", auth_header(token)) + .send() + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!(body["success"], true); + } + + // ------------------------------------------------------------------------- + // format.rs: schema_to_rowtype + // ------------------------------------------------------------------------- + + #[test] + fn schema_to_rowtype_maps_basic_types() { + let schema = Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + Field::new("score", DataType::Float64, true), + ]); + + let rowtype = schema_to_rowtype(&schema); + let cols = rowtype.as_array().unwrap(); + + assert_eq!(cols.len(), 3); + assert_eq!(cols[0]["name"], "id"); + assert_eq!(cols[0]["type"], "fixed"); + assert_eq!(cols[0]["nullable"], false); + assert_eq!(cols[1]["name"], "name"); + assert_eq!(cols[1]["type"], "text"); + assert_eq!(cols[1]["nullable"], true); + assert_eq!(cols[2]["name"], "score"); + assert_eq!(cols[2]["type"], "real"); + } + + // ------------------------------------------------------------------------- + // format.rs: batches_to_arrow_base64 IPC round-trip + // ------------------------------------------------------------------------- + + #[test] + fn arrow_ipc_round_trip_preserves_row_count() { + let schema = Arc::new(Schema::new(vec![ + Field::new("x", DataType::Int64, false), + Field::new("label", DataType::Utf8, true), + ])); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["a", "b", "c"])), + ], + ) + .unwrap(); + + let b64 = batches_to_arrow_base64(&schema, &[batch]); + assert!(!b64.is_empty()); + + let raw = base64::engine::general_purpose::STANDARD + .decode(&b64) + .expect("valid base64"); + let reader = StreamReader::try_new(std::io::Cursor::new(raw), None).unwrap(); + let total_rows: usize = reader.flatten().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 3); + } + + #[test] + fn arrow_ipc_empty_batches_produces_valid_stream() { + let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)])); + let b64 = batches_to_arrow_base64(&schema, &[]); + assert!(!b64.is_empty(), "must emit at least the IPC schema message"); + + let raw = base64::engine::general_purpose::STANDARD + .decode(&b64) + .unwrap(); + let reader = StreamReader::try_new(std::io::Cursor::new(raw), None).unwrap(); + let batches: Vec = reader.flatten().collect(); + assert_eq!(batches.len(), 0); + } + + #[test] + fn arrow_ipc_float64_survives_round_trip() { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Float64Array::from(vec![1.5, 2.5, 3.5]))], + ) + .unwrap(); + + let b64 = batches_to_arrow_base64(&schema, &[batch]); + let raw = base64::engine::general_purpose::STANDARD + .decode(&b64) + .unwrap(); + let reader = StreamReader::try_new(std::io::Cursor::new(raw), None).unwrap(); + let back: Vec = reader.flatten().collect(); + assert_eq!(back[0].num_rows(), 3); + } + + // ------------------------------------------------------------------------- + // common.rs: gzip decode (complementary to the existing tests in common.rs) + // ------------------------------------------------------------------------- + + #[test] + fn plain_body_is_returned_unchanged() { + use crate::snowflake::http::handlers::common::decode_snowflake_request_body; + use axum::http::HeaderMap; + + let raw = br#"{"data":{"LOGIN_NAME":"u"}}"#; + let out = + decode_snowflake_request_body(&HeaderMap::new(), &Bytes::from_static(raw)).unwrap(); + assert_eq!(out.as_slice(), raw); + } + + #[test] + fn gzip_body_is_decompressed_correctly() { + use crate::snowflake::http::handlers::common::decode_snowflake_request_body; + use axum::http::{HeaderMap, HeaderValue}; + + let json = br#"{"data":{"LOGIN_NAME":"tester"}}"#; + let mut enc = GzEncoder::new(Vec::new(), Compression::default()); + enc.write_all(json).unwrap(); + let gz = Bytes::from(enc.finish().unwrap()); + + let mut headers = HeaderMap::new(); + headers.insert("content-encoding", HeaderValue::from_static("gzip")); + let out = decode_snowflake_request_body(&headers, &gz).unwrap(); + assert_eq!(out.as_slice(), json); + } +} diff --git a/crates/queryflux-frontend/src/state.rs b/crates/queryflux-frontend/src/state.rs index d0a4f92..d0cb59e 100644 --- a/crates/queryflux-frontend/src/state.rs +++ b/crates/queryflux-frontend/src/state.rs @@ -19,6 +19,8 @@ use queryflux_persistence::Persistence; use queryflux_routing::chain::{RouterChain, RoutingTrace}; use queryflux_translation::TranslationService; +use crate::snowflake::http::session_store::SnowflakeSessionStore; + /// Everything that can be hot-reloaded from the DB without restarting the proxy. /// /// Wrapped in `Arc>` inside `AppState` so @@ -65,6 +67,8 @@ pub struct AppState { pub authorization: Arc, /// Resolves per-user `QueryCredentials` from `AuthContext` + cluster `queryAuth` config. pub identity_resolver: Arc, + /// Active Snowflake wire-protocol sessions (Form 1). Shared across all handlers. + pub snowflake_sessions: Arc, } /// Stable per-query metadata that does not change across the query's lifecycle. diff --git a/crates/queryflux-routing/src/implementations/protocol_based.rs b/crates/queryflux-routing/src/implementations/protocol_based.rs index c98a325..57a9a9d 100644 --- a/crates/queryflux-routing/src/implementations/protocol_based.rs +++ b/crates/queryflux-routing/src/implementations/protocol_based.rs @@ -9,12 +9,15 @@ use crate::RouterTrait; /// Routes based on which frontend protocol the client used. /// Useful for directing MySQL-wire clients (StarRocks) to a different group -/// than Trino HTTP clients. +/// than Trino HTTP clients, or Snowflake clients to a dedicated Snowflake group. pub struct ProtocolBasedRouter { pub trino_http: Option, pub postgres_wire: Option, pub mysql_wire: Option, pub clickhouse_http: Option, + pub flight_sql: Option, + pub snowflake_http: Option, + pub snowflake_sql_api: Option, } #[async_trait] @@ -35,7 +38,9 @@ impl RouterTrait for ProtocolBasedRouter { FrontendProtocol::PostgresWire => self.postgres_wire.clone(), FrontendProtocol::MySqlWire => self.mysql_wire.clone(), FrontendProtocol::ClickHouseHttp => self.clickhouse_http.clone(), - FrontendProtocol::FlightSql => None, + FrontendProtocol::FlightSql => self.flight_sql.clone(), + FrontendProtocol::SnowflakeHttp => self.snowflake_http.clone(), + FrontendProtocol::SnowflakeSqlApi => self.snowflake_sql_api.clone(), }; Ok(group) } diff --git a/crates/queryflux-routing/src/implementations/python_script.rs b/crates/queryflux-routing/src/implementations/python_script.rs index 6c2301f..b5b63e9 100644 --- a/crates/queryflux-routing/src/implementations/python_script.rs +++ b/crates/queryflux-routing/src/implementations/python_script.rs @@ -72,6 +72,8 @@ fn protocol_camel(p: FrontendProtocol) -> &'static str { FrontendProtocol::MySqlWire => "mysqlWire", FrontendProtocol::ClickHouseHttp => "clickHouseHttp", FrontendProtocol::FlightSql => "flightSql", + FrontendProtocol::SnowflakeHttp => "snowflakeHttp", + FrontendProtocol::SnowflakeSqlApi => "snowflakeSqlApi", } } diff --git a/crates/queryflux-routing/tests/router_tests.rs b/crates/queryflux-routing/tests/router_tests.rs index 511e736..b6b9b06 100644 --- a/crates/queryflux-routing/tests/router_tests.rs +++ b/crates/queryflux-routing/tests/router_tests.rs @@ -146,6 +146,9 @@ async fn protocol_router_trino_http() { postgres_wire: Some(group("pg-group")), mysql_wire: None, clickhouse_http: None, + flight_sql: None, + snowflake_http: None, + snowflake_sql_api: None, }; let session = trino_session(&[]); let result = router @@ -162,6 +165,9 @@ async fn protocol_router_postgres_wire() { postgres_wire: Some(group("pg-group")), mysql_wire: None, clickhouse_http: None, + flight_sql: None, + snowflake_http: None, + snowflake_sql_api: None, }; let result = router .route( @@ -182,6 +188,9 @@ async fn protocol_router_unconfigured() { postgres_wire: None, // not configured mysql_wire: None, clickhouse_http: None, + flight_sql: None, + snowflake_http: None, + snowflake_sql_api: None, }; let result = router .route( @@ -202,6 +211,9 @@ async fn protocol_router_mysql_wire() { postgres_wire: None, mysql_wire: Some(group("mysql-group")), clickhouse_http: None, + flight_sql: None, + snowflake_http: None, + snowflake_sql_api: None, }; let result = router .route( @@ -222,6 +234,9 @@ async fn protocol_router_clickhouse_http() { postgres_wire: None, mysql_wire: None, clickhouse_http: Some(group("ch-group")), + flight_sql: None, + snowflake_http: None, + snowflake_sql_api: None, }; let result = router .route( @@ -236,12 +251,15 @@ async fn protocol_router_clickhouse_http() { } #[tokio::test] -async fn protocol_router_flight_sql_not_routed() { +async fn protocol_router_flight_sql() { let router = ProtocolBasedRouter { trino_http: Some(group("trino-group")), postgres_wire: Some(group("pg-group")), mysql_wire: Some(group("mysql-group")), clickhouse_http: Some(group("ch-group")), + flight_sql: Some(group("sf-analytics")), + snowflake_http: None, + snowflake_sql_api: None, }; let result = router .route( @@ -252,7 +270,53 @@ async fn protocol_router_flight_sql_not_routed() { ) .await .unwrap(); - assert_eq!(result, None); + assert_eq!(result, Some(group("sf-analytics"))); +} + +#[tokio::test] +async fn protocol_router_snowflake_http() { + let router = ProtocolBasedRouter { + trino_http: None, + postgres_wire: None, + mysql_wire: None, + clickhouse_http: None, + flight_sql: None, + snowflake_http: Some(group("sf-group")), + snowflake_sql_api: None, + }; + let result = router + .route( + "SELECT 1", + &trino_session(&[]), + &FrontendProtocol::SnowflakeHttp, + None, + ) + .await + .unwrap(); + assert_eq!(result, Some(group("sf-group"))); +} + +#[tokio::test] +async fn protocol_router_snowflake_sql_api() { + let router = ProtocolBasedRouter { + trino_http: None, + postgres_wire: None, + mysql_wire: None, + clickhouse_http: None, + flight_sql: None, + snowflake_http: None, + snowflake_sql_api: Some(group("sf-api-group")), + }; + let result = router + .route( + "SELECT 1", + &trino_session(&[]), + &FrontendProtocol::SnowflakeSqlApi, + None, + ) + .await + .unwrap(); + assert_eq!(result, Some(group("sf-api-group"))); } // --------------------------------------------------------------------------- @@ -1043,6 +1107,9 @@ async fn chain_protocol_based_then_header_fallback() { postgres_wire: Some(group("pg-from-protocol")), mysql_wire: None, clickhouse_http: None, + flight_sql: None, + snowflake_http: None, + snowflake_sql_api: None, }), Box::new(HeaderRouter::new( "x-tenant".to_string(), diff --git a/crates/queryflux/src/main.rs b/crates/queryflux/src/main.rs index 2fea8d2..aa2673d 100644 --- a/crates/queryflux/src/main.rs +++ b/crates/queryflux/src/main.rs @@ -20,6 +20,7 @@ use queryflux_frontend::{ flight_sql::FlightSqlFrontend, mysql_wire::MysqlWireFrontend, postgres_wire::PostgresWireFrontend, + snowflake::{http::session_store::SnowflakeSessionStore, SnowflakeFrontend}, state::LiveConfig, trino_http::{state::AppState, TrinoHttpFrontend}, FrontendListenerTrait, @@ -392,6 +393,9 @@ async fn main() -> Result<()> { postgres_wire, mysql_wire, clickhouse_http, + flight_sql, + snowflake_http, + snowflake_sql_api, } => { routers.push(Box::new(ProtocolBasedRouter { trino_http: trino_http.as_ref().map(|s| ClusterGroupName(s.clone())), @@ -400,6 +404,11 @@ async fn main() -> Result<()> { clickhouse_http: clickhouse_http .as_ref() .map(|s| ClusterGroupName(s.clone())), + flight_sql: flight_sql.as_ref().map(|s| ClusterGroupName(s.clone())), + snowflake_http: snowflake_http.as_ref().map(|s| ClusterGroupName(s.clone())), + snowflake_sql_api: snowflake_sql_api + .as_ref() + .map(|s| ClusterGroupName(s.clone())), })); } RouterConfig::Header { @@ -605,6 +614,7 @@ async fn main() -> Result<()> { auth_provider, authorization, identity_resolver, + snowflake_sessions: SnowflakeSessionStore::new(), }); // --- Start admin server (Prometheus /metrics + future /admin/* endpoints) --- @@ -939,12 +949,24 @@ async fn main() -> Result<()> { } }; + let snowflake_future = async { + match &config.queryflux.frontends.snowflake_http { + Some(cfg) if cfg.enabled => { + SnowflakeFrontend::new(app_state.clone(), cfg.port) + .listen() + .await + } + _ => std::future::pending::>().await, + } + }; + tokio::select! { - r = frontend.listen() => r.map_err(|e| anyhow::anyhow!("{e}"))?, - r = admin.listen() => r.map_err(|e| anyhow::anyhow!("{e}"))?, - r = mysql_future => r.map_err(|e| anyhow::anyhow!("{e}"))?, - r = postgres_future => r.map_err(|e| anyhow::anyhow!("{e}"))?, - r = flight_sql_future => r.map_err(|e| anyhow::anyhow!("{e}"))?, + r = frontend.listen() => r.map_err(|e| anyhow::anyhow!("{e}"))?, + r = admin.listen() => r.map_err(|e| anyhow::anyhow!("{e}"))?, + r = mysql_future => r.map_err(|e| anyhow::anyhow!("{e}"))?, + r = postgres_future => r.map_err(|e| anyhow::anyhow!("{e}"))?, + r = flight_sql_future => r.map_err(|e| anyhow::anyhow!("{e}"))?, + r = snowflake_future => r.map_err(|e| anyhow::anyhow!("{e}"))?, } Ok(()) @@ -1173,6 +1195,9 @@ async fn build_live_config( postgres_wire, mysql_wire, clickhouse_http, + flight_sql, + snowflake_http, + snowflake_sql_api, } => { routers.push(Box::new( queryflux_routing::implementations::protocol_based::ProtocolBasedRouter { @@ -1182,6 +1207,13 @@ async fn build_live_config( clickhouse_http: clickhouse_http .as_ref() .map(|s| ClusterGroupName(s.clone())), + flight_sql: flight_sql.as_ref().map(|s| ClusterGroupName(s.clone())), + snowflake_http: snowflake_http + .as_ref() + .map(|s| ClusterGroupName(s.clone())), + snowflake_sql_api: snowflake_sql_api + .as_ref() + .map(|s| ClusterGroupName(s.clone())), }, )); } diff --git a/docs/README.md b/docs/README.md deleted file mode 100644 index f012376..0000000 --- a/docs/README.md +++ /dev/null @@ -1,18 +0,0 @@ -# QueryFlux architecture documentation - -

- QueryFlux — multi-engine SQL query proxy and router in Rust -

- -This folder describes how QueryFlux is put together: why it exists, how SQL is translated, and how traffic is routed to **cluster groups** and individual **clusters**. - -| Document | What it covers | -|----------|----------------| -| [motivation-and-goals.md](motivation-and-goals.md) | Problem statement, goals, and how QueryFlux fits multi-engine estates. | -| [architecture.md](architecture.md) | End-to-end query lifecycle, major crates, and component status (high level). | -| [query-translation.md](query-translation.md) | Dialect detection, sqlglot integration, when translation runs, and schema-aware mode. | -| [routing-and-clusters.md](routing-and-clusters.md) | Router chain, `routingFallback`, cluster groups, load-balancing strategies, and queueing. | -| [observability.md](observability.md) | Prometheus metrics, Grafana dashboard, QueryFlux Studio, and the Admin REST API. | -| [adding-engine-support.md](adding-engine-support.md) | Checklist for new **backend engines** (Rust adapters), **Studio** (`lib/studio-engines/` manifest + catalog slots), and **frontend wire protocols** (e.g. Postgres wire). | - -Start with [motivation-and-goals.md](motivation-and-goals.md) if you are new to the project; use [architecture.md](architecture.md) as the single-page map of the system. diff --git a/docs/architecture.md b/docs/architecture.md deleted file mode 100644 index 78f2c03..0000000 --- a/docs/architecture.md +++ /dev/null @@ -1,354 +0,0 @@ -# QueryFlux — Architecture Overview - -QueryFlux is a universal SQL query proxy and router. It accepts queries from clients over multiple protocols (Trino HTTP, PostgreSQL wire, MySQL wire, Arrow Flight SQL), routes them to the appropriate backend engine, optionally translates the SQL dialect, and streams results back in the client's native format. - -**More documentation:** [docs/README.md](README.md) indexes deeper topics — [motivation-and-goals.md](motivation-and-goals.md) (why the project exists), [query-translation.md](query-translation.md) (sqlglot and dialects), [routing-and-clusters.md](routing-and-clusters.md) (routers, groups, load balancing), [observability.md](observability.md) (Prometheus, Grafana, Studio, Admin API), [adding-engine-support.md](adding-engine-support.md) (new engines, Studio, and client protocols). - ---- - -## High-Level Flow - -``` -Client (Trino CLI / psql / mysql / DBI) - │ native protocol - ▼ -┌───────────────────┐ -│ Frontend Listener │ ← speaks the client's wire protocol -└────────┬──────────┘ - │ SQL + SessionContext - ▼ -┌───────────────────┐ -│ Router Chain │ ← selects target cluster group -└────────┬──────────┘ - │ ClusterGroupName - ▼ -┌───────────────────┐ -│ ClusterGroupManager│ ← load-balances across clusters; queues if at capacity -└────────┬──────────┘ - │ ClusterName - ▼ -┌───────────────────┐ -│ Translation Service│ ← sqlglot via PyO3; skipped when dialects match -└────────┬──────────┘ - │ translated SQL - ▼ -┌───────────────────┐ -│ Engine Adapter │ ← speaks the backend engine's native protocol -└────────┬──────────┘ - │ QueryExecution (Async | Sync) - ▼ -┌───────────────────┐ -│ Persistence │ ← stores in-flight state for async engines -└───────────────────┘ -``` - -The frontend never knows which engine it's talking to. The engine adapter never knows which client protocol was used. The dispatch layer in the middle is the only place that bridges them. - ---- - -## Workspace Layout - -``` -queryflux/ -├── crates/ -│ ├── queryflux/ # main binary — wires everything together -│ ├── queryflux-core/ # shared types: ProxyQueryId, SessionContext, QueryPollResult, … -│ ├── queryflux-config/ # ConfigProvider trait + YamlFileConfigProvider -│ ├── queryflux-frontend/ # FrontendListenerTrait + protocol implementations -│ ├── queryflux-engine-adapters/ # EngineAdapterTrait + per-engine implementations -│ ├── queryflux-routing/ # RouterTrait + RouterChain + all router implementations -│ ├── queryflux-cluster-manager/ # ClusterGroupManager: load balancing + queueing -│ ├── queryflux-persistence/ # Persistence + MetricsStore + ClusterConfigStore traits + impls -│ ├── queryflux-metrics/ # PrometheusMetrics, BufferedMetricsStore, MultiMetricsStore -│ ├── queryflux-translation/ # TranslatorTrait + SqlglotTranslator (PyO3) -│ └── queryflux-e2e-tests/ # Integration tests -├── ui/queryflux-studio/ # Next.js management UI (cluster monitoring, query history) -├── prometheus/ # Prometheus scrape config -├── grafana/ # Grafana provisioning + dashboards -├── docker/ # Docker Compose files -│ ├── docker-compose.yml # Local dev: Trino + Postgres + Prometheus + Grafana -│ └── docker-compose.test.yml # E2E test stack (isolated ports) -├── config.local.yaml # Example config for local development -└── Makefile # build / run / test shortcuts -``` - ---- - -## Core Abstractions - -### SessionContext (`queryflux-core`) - -Carries protocol-specific metadata that travels with a query from frontend through routing and into the engine adapter. Each variant holds what that protocol actually provides. - -```rust -pub enum SessionContext { - TrinoHttp { headers: HashMap }, - PostgresWire { user: Option, database: Option, session_params: HashMap }, - MySqlWire { user: Option, schema: Option, session_vars: HashMap }, - ClickHouseHttp { headers: HashMap, query_params: HashMap }, -} -``` - -### QueryExecution (`queryflux-core`) - -Engines fall into two models. The adapter declares which model it uses; dispatch handles both uniformly. - -``` -QueryExecution::Async { backend_query_id, next_uri, initial_body } - → dispatcher stores handle in Persistence - → client polls proxy until complete - -QueryExecution::Sync { result: QueryPollResult } - → dispatcher returns result immediately - → no Persistence needed -``` - -| Engine | Model | Notes | -|---|---|---| -| Trino | Async | Submit → poll `nextUri` until done | -| DuckDB | Sync | Runs on `spawn_blocking`, result available immediately | -| StarRocks | Sync | MySQL protocol, single round-trip | -| ClickHouse | — | Planned | - -### EngineAdapterTrait (`queryflux-engine-adapters`) - -```rust -pub trait EngineAdapterTrait: Send + Sync { - async fn submit_query(&self, sql: &str, session: &SessionContext) -> Result; - async fn poll_query(&self, backend_id: &BackendQueryId, next_uri: Option<&str>) -> Result; - async fn cancel_query(&self, backend_id: &BackendQueryId) -> Result<()>; - async fn health_check(&self) -> bool; - fn engine_type(&self) -> EngineType; - - // Catalog discovery — feeds schema context for translation - async fn list_catalogs(&self) -> Result>; - async fn list_databases(&self, catalog: &str) -> Result>; - async fn list_tables(&self, catalog: &str, db: &str) -> Result>; - async fn describe_table(&self, catalog: &str, db: &str, table: &str) -> Result>; -} -``` - -### RouterTrait (`queryflux-routing`) - -```rust -pub trait RouterTrait: Send + Sync { - fn type_name(&self) -> &'static str; - async fn route( - &self, - sql: &str, - session: &SessionContext, - frontend_protocol: &FrontendProtocol, - ) -> Result>; -} -``` - -`RouterChain` evaluates routers in config order. First `Ok(Some(group))` wins. Falls back to `routingFallback` if every router returns `Ok(None)`. `route_with_trace` builds a `RoutingTrace` for debugging and observability. - ---- - -## Implemented Components - -### Frontends - -| Protocol | Status | Port | -|---|---|---| -| Trino HTTP | **Done** | 8080 | -| PostgreSQL wire | **Done** | 5432 | -| MySQL wire | **Done** | 3306 | -| Arrow Flight SQL | **Done** (query execution) | — | -| Admin / Prometheus metrics | **Done** | 9000 | -| ClickHouse HTTP | Planned | 8123 | - -**Trino HTTP routes:** - -| Method | Path | Description | -|---|---|---| -| `POST` | `/v1/statement` | Submit a new query | -| `GET` | `/v1/statement/qf/queued/{id}/{seq}` | Poll a queued query (with backoff) | -| `GET` | `/v1/statement/qf/executing/{id}` | Poll an executing query | -| `DELETE` | `/v1/statement/qf/executing/{id}` | Cancel a running query | - -### Engine Adapters - -| Engine | Status | Execution model | -|---|---|---| -| Trino | **Done** | Async HTTP — transparent `nextUri` proxying | -| DuckDB | **Done** | Sync embedded — `spawn_blocking` + Arrow result set | -| StarRocks | **Done** | MySQL protocol — sync Arrow path via `execute_as_arrow` | -| ClickHouse | Planned | — | - -### Routers - -| Router | Matching criteria | -|---|---| -| `protocolBased` | Which frontend protocol the client used | -| `header` | HTTP header value (Trino HTTP only) | -| `queryRegex` | Regex patterns against SQL text | -| `clientTags` | Trino `X-Trino-Client-Tags` header | -| `pythonScript` | Custom Python function (`def route(query, ctx) -> str | None`) — see [routing-and-clusters.md](routing-and-clusters.md#python-script-router-pythonscript) | - -### Persistence - -| Store | Status | Use case | -|---|---|---| -| In-memory (`DashMap`) | **Done** | Single-instance dev | -| PostgreSQL (JSONB) | **Done** | Production / HA | -| Redis | Planned | Distributed | - -### Metrics - -| Store | Status | Purpose | -|---|---|---| -| `PrometheusMetrics` | **Done** | Real-time operational metrics at `/metrics` | -| `NoopMetricsStore` | **Done** | Default — zero overhead | -| `PostgresStore` (MetricsStore) | **Done** | Historical query records for the management UI | -| `BufferedMetricsStore` | **Done** | Async write buffer wrapping any MetricsStore | - -**Prometheus metrics exposed:** - -| Metric | Type | Labels | -|---|---|---| -| `queryflux_queries_total` | Counter | `engine_type`, `cluster_group`, `status`, `protocol` | -| `queryflux_query_duration_seconds` | Histogram | `engine_type`, `cluster_group` | -| `queryflux_translated_queries_total` | Counter | `src_dialect`, `tgt_dialect` | -| `queryflux_running_queries` | Gauge | `cluster_group`, `cluster_name` | -| `queryflux_queued_queries` | Gauge | `cluster_group` | - ---- - -## SQL Translation - -Translation is handled by [sqlglot](https://github.com/tobymao/sqlglot) (Python, 31+ dialects) called via PyO3. - -**When translation runs:** only when the incoming client dialect differs from the target engine's dialect. Trino client → Trino cluster = zero overhead passthrough. - -**Two translation modes** (both implemented in `queryflux-translation`; see [query-translation.md](query-translation.md)): - -1. **Dialect-only** (empty `SchemaContext`): `sqlglot.transpile(sql, read=src, write=tgt)` — this is what the main dispatch path uses today (`SchemaContext::default()`). -2. **Schema-aware** (non-empty `SchemaContext`): parse → `sqlglot.optimizer.optimize` with `MappingSchema` → emit in target dialect, with fallback to dialect-only if optimization fails. - -Source dialect is inferred from the frontend protocol (`TrinoHttp` → Trino, `PostgresWire` → Postgres, etc.). Target dialect comes from the selected cluster’s **engine type** (via the adapter). - -Translation gracefully degrades: if sqlglot is unavailable at startup, the service disables itself and SQL passes through untranslated. - ---- - -## Configuration - -```yaml -queryflux: - externalAddress: http://localhost:8080 - frontends: - trinoHttp: { enabled: true, port: 8080 } - postgresWire: { enabled: false, port: 5432 } - mysqlWire: { enabled: false, port: 3306 } - persistence: - inMemory: {} # or: postgres: { databaseUrl: "postgres://..." } - adminApi: - port: 9000 - -clusters: - trino-1: - engine: trino - endpoint: http://trino:8080 - enabled: true - duckdb-1: - engine: duckDb - enabled: true - databasePath: /data/analytics.duckdb # omit for in-memory - -clusterGroups: - trino-default: - enabled: true - maxRunningQueries: 100 - members: [trino-1] - - duckdb-local: - enabled: true - maxRunningQueries: 4 - members: [duckdb-1] - -translation: - errorOnUnsupported: false - -routers: - - type: protocolBased - trinoHttp: trino-default - - - type: header - headerName: X-Target-Engine - headerValueToGroup: - duckdb: duckdb-local - - - type: pythonScript - script: | - def route(query, ctx): - if "big_table" in query: - return "trino-default" - return None - -routingFallback: duckdb-local -``` - ---- - -## Local Development - -### Prerequisites - -- Rust (stable) -- Docker + Docker Compose -- Python 3.10+ - -### Setup - -```bash -# Install Python dependencies (sqlglot) -make setup - -# Export Python path for PyO3 -export PYO3_PYTHON=$(pwd)/.venv/bin/python3 - -# Start Trino + Postgres + Prometheus + Grafana, then run the proxy -make dev -``` - -### Services - -| Service | URL | Credentials | -|---|---|---| -| QueryFlux (Trino HTTP) | http://localhost:8080 | — | -| Prometheus metrics | http://localhost:9000/metrics | — | -| Trino (direct) | http://localhost:8081 | — | -| Prometheus | http://localhost:9090 | — | -| Grafana | http://localhost:3000 | admin / admin | -| PostgreSQL | localhost:5433 | queryflux / queryflux | - -### Send a query - -```bash -# Via Trino CLI -trino --server http://localhost:8080 --execute "SELECT 42" - -# Via curl -curl -s -X POST http://localhost:8080/v1/statement \ - -H "X-Trino-User: dev" \ - -d "SELECT current_date" -``` - ---- - -## Roadmap - -| Phase | Feature | Status | -|---|---|---| -| P1 | Trino HTTP frontend + DuckDB/Trino backends | **Done** | -| P1 | sqlglot translation (dialect-only) | **Done** | -| P1 | Prometheus metrics | **Done** | -| P1 | Postgres persistence + query history | **Done** | -| P1 | PostgreSQL wire frontend | **Done** | -| P1 | MySQL wire frontend + StarRocks backend | **Done** | -| P1 | Arrow Flight SQL frontend | **Done** | -| P1 | QueryFlux Studio — management UI | **Done** | -| P2 | Wire `SchemaContext` from catalog into dispatch | Planned | -| P3 | ClickHouse HTTP backend + frontend | Planned | diff --git a/docs/auth-authz-design.md b/docs/auth-authz-design.md deleted file mode 100644 index 03005d4..0000000 --- a/docs/auth-authz-design.md +++ /dev/null @@ -1,733 +0,0 @@ -# QueryFlux Auth/AuthZ Design - -## The Two-Credential Model (Core Principle) - -Every backend cluster has **two distinct credential relationships**, both configured per-cluster in `ClusterConfig`: - -**Credential Type 1 — Service Credentials** (`auth`, existing `ClusterAuth`) -- QueryFlux's own service account for the backend -- Used for: health checks, schema/catalog discovery, cluster management -- Static, ops-owned; ideally stored in Secrets Manager (not inline in config) -- Auth types: `basic`, `bearer`, `keyPair` (RSA — new, for Snowflake/Databricks) -- Never changes at request time; independent of which user is running a query - -**Credential Type 2 — Query Execution Credentials** (`queryAuth`, new field in `ClusterConfig`) -- The credentials used to execute a specific user's query on the backend -- Configured per-cluster — operators choose which mode their engine supports -- Resolved per-request from `AuthContext` (verified identity) + the configured `queryAuth` type -- Validated at startup: each engine accepts only its supported modes -- Default when omitted: `serviceAccount` (falls back to Type 1 for everything) - -`queryAuth` has exactly **three explicit types**: `serviceAccount | impersonate | tokenExchange` - -There is no `passthrough` type. For same-engine routing (Trino HTTP → Trino backend), the Trino adapter already forwards all `SessionContext::TrinoHttp { headers }` verbatim — including the client's `Authorization` header — without any special config. This implicit client header passthrough is the default today. - -Health checks always use Type 1 (`auth`) directly, never `queryAuth`. This ensures they work even when a user's token is expired or missing. - -```yaml -# config.yaml — per-cluster dual credentials (camelCase to match existing serde config) -clusters: - trino-prod: - engine: trino - endpoint: https://trino.internal:8443 - auth: # Type 1 — service credentials - type: basic - username: qf_svc - password: "..." - queryAuth: # Type 2 — query execution mode - type: impersonate # service account + X-Trino-User header - - clickhouse-prod: - engine: clickHouse - endpoint: http://clickhouse:8123 - auth: - type: basic - username: qf_svc - password: "..." - queryAuth: - type: serviceAccount # only viable option for ClickHouse - - snowflake-prod: # future adapter - engine: snowflake - endpoint: https://myaccount.snowflakecomputing.com - auth: - type: keyPair # RSA key-pair, Snowflake standard - username: QF_SVC - privateKeyPem: "..." - queryAuth: - type: tokenExchange - tokenEndpoint: https://keycloak.internal/realms/my-realm/protocol/openid-connect/token - clientId: queryflux-gateway - clientSecret: "..." - - # Trino→Trino same-IdP: no queryAuth needed. - # SessionContext headers (including Authorization) are forwarded implicitly. - trino-analytics: - engine: trino - endpoint: https://trino-analytics.internal:8443 - auth: - type: basic - username: qf_svc - password: "..." - # queryAuth omitted → serviceAccount default - # but Authorization header still forwarded by Trino adapter via SessionContext -``` - -**`queryAuth` engine compatibility** (startup validation rejects unsupported combinations): - -| Engine | `serviceAccount` | `impersonate` | `tokenExchange` | -|--------|:---:|:---:|:---:| -| Trino | ✅ | ✅ `X-Trino-User` (needs Trino file-based ACL) | — | -| ClickHouse | ✅ | ❌ no trusted proxy mechanism | — | -| StarRocks (MySQL wire) | ✅ | ❌ no wire mechanism | — | -| StarRocks (HTTP, future) | ✅ | — | — | -| Snowflake (future) | ✅ key-pair | ❌ | ✅ external OAuth | -| Databricks (future) | ✅ | ❌ | ✅ OAuth U2M | -| DuckDB | ✅ (no-op) | — | — | - -**`BackendIdentityResolver` pseudocode:** - -At this point the caller has already: selected **`ClusterGroupMember`**, merged **`connection`** hints for the adapter, and resolved **which `ClusterAuth` / profile** supplies Type 1 material. - -``` -fn resolve(auth_ctx: &AuthContext, cluster: &ClusterConfig, type1: &ClusterAuth) -> QueryCredentials { - // If no user identity available, always fall back to service account - if auth_ctx is NoneIdentity { - return ServiceAccountCreds(type1.clone()) - } - - match cluster.queryAuth.type { - serviceAccount => - ServiceAccountCreds(type1.clone()) - - impersonate => - // Use Type 1 credentials on the wire; inject user identity separately. - // IMPORTANT: suppress the client's Authorization header — do NOT forward it. - // Only Type 1 (service account) auth reaches the backend. - // User identity is injected via engine-specific header AFTER authentication. - ImpersonateCreds { - service_auth: type1.clone(), // resolved profile or cluster.auth - user: auth_ctx.user.clone(), // injected as X-Trino-User (Trino only) - } - - tokenExchange => - // Exchange auth_ctx.raw_token at the configured OAuth endpoint. - // Falls back to serviceAccount if raw_token is None. - // Per-provider contract: see Layer 3 → tokenExchange section. - exchange_token(auth_ctx.raw_token?, cluster.queryAuth.token_exchange_config) - } -} -``` - -**Mixed `queryAuth` within one cluster group:** Allowed — each cluster carries its own config. If a group uses `engineAffinity` or `weighted` strategy with members of different `queryAuth` types, the resolver uses whichever cluster was selected. Operators should ensure all members of a group use the same `queryAuth` type unless they explicitly want per-cluster behaviour; a startup warning is emitted if a group has members with mixed types. - ---- - -## Cluster group membership: array of connection options - -**Problem:** `members: [ "cluster-a", "cluster-b" ]` is not enough when the **same logical cluster** (same endpoint) participates in **multiple groups** with different **team defaults** (e.g. Snowflake **role/warehouse**, or which **auth profile** to prefer). - -**Model:** `clusterGroups[].members` becomes an **array of objects** — each entry is **one connection option** in the group’s pool (ordering preserved for failover / round-robin / weighted). - -Each **`ClusterGroupMember`** (name TBD) contains at minimum: - -- **`cluster`** — required; name of an entry in `clusters`. -- **`connection`** — optional, **engine-specific** non-secret hints merged at dispatch after a member is selected (or used to disambiguate defaults). Validated at startup: fields must match the **referenced cluster’s `engine`**; unknown or wrong-engine fields are **rejected** (fail fast). -- **`weight`** — optional; for weighted strategies within the group. -- **`defaultAuthProfile`** — optional; names a profile defined on that cluster (see below). Supplies the **group-scoped default** when this member is the path into the cluster (“team A uses role ANALYST”). - -**Backward compatibility:** Config loader may accept **either** a bare string (`"trino-prod"`) or a full object, so existing YAML keeps working during migration. - -### Per-engine `connection` shapes (all supported types) - -Each engine exposes a **closed set** of connection option types. Implement as a **serde tagged enum** `EngineConnectionOptions` (or nested `Option` structs with validation) so only valid combinations deserialize. - -| Engine | Purpose of group-level `connection` | Typical fields (non-secret) | `ClusterAuth` (Type 1) variants | -|--------|-------------------------------------|------------------------------|----------------------------------| -| **Trino** | Session context defaults for this group’s path | `catalog`, `schema`, optional `sessionProperties` map | `basic`, `bearer` | -| **ClickHouse** | Default database / settings | `database`, optional `role` (if using CH RBAC features) | `basic` (maps to `X-ClickHouse-User` / `Key`) | -| **StarRocks (MySQL wire)** | Default catalog/db context | `database` | `basic` (user/password) | -| **StarRocks (HTTP, future)** | JWT-forwarding session hints | TBD aligned with StarRocks HTTP API | `basic`, `bearer` | -| **Snowflake** | Same account URL, different **team slice** | `role`, `warehouse`, `database`, `schema` | `keyPair` (recommended), future password-based if needed | -| **Databricks (future)** | Warehouse / HTTP context | `warehouseId`, `catalog`, `httpPath` (product-specific) | `bearer`, OAuth client creds via `queryAuth` | -| **DuckDB** | Usually none (file path on cluster) | rarely `attach` hints | N/A (embedded) | - -**Rule:** Secrets (**passwords, PEMs, client secrets**) stay on **`clusters[].auth`** or **`authProfiles`** via **inline (dev) or `secretRef` (prod)** — not duplicated per group. Group `connection` carries **which role/warehouse/catalog** to use, not private keys. - -**Teams / groups pattern:** Create **one cluster group per team** (or workload). Each group lists the same Snowflake `cluster` name once, with different `connection.role` / `warehouse` and/or `defaultAuthProfile`. Combined with **allowGroups** / OpenFGA, “Alice may only hit `group:team-analytics`” implies she only gets **that group’s** Snowflake role default. - ---- - -## Auth profiles (cluster-scoped, named Type 1 variants) - -When one cluster needs **multiple service identities or Snowflake logins** (different key-pair users, or same user + different static contexts), define **`authProfiles`** on **`ClusterConfig`**: - -```yaml -clusters: - snowflake-prod: - engine: snowflake - endpoint: https://xyz.snowflakecomputing.com - defaultAuthProfile: svc_readonly - authProfiles: - svc_readonly: - type: keyPair - username: QF_READONLY - privateKeySecretRef: { provider: vault, path: secret/data/qf/sf-readonly, field: key } - svc_etl: - type: keyPair - username: QF_ETL - privateKeySecretRef: { provider: vault, path: secret/data/qf/sf-etl, field: key } - queryAuth: - type: serviceAccount # or tokenExchange — profile picks which Type 1 for serviceAccount path -``` - -**Resolution order after cluster + group member are known:** - -1. **`defaultAuthProfile`** on the **group member** entry (if set) — team/unit-specific service account. -2. Else **`defaultAuthProfile`** on the **cluster** (if set) — cluster-level default. -3. Else legacy single **`auth`** block on the cluster. - -The client **never** influences which auth profile is used. Clients influence *routing* (which group to target) via the existing router chain — headers, client tags, regex, protocol — but the auth/authz layer must approve access to that group. Once a group is selected, the auth profile is entirely determined by operator config. This separation means the group IS the privilege boundary: being authorized for `team-a-snowflake` group means you get the `svc_readonly` service account; being authorized for `team-etl` means you get `svc_etl`. No escalation possible from the client side. - -Config is invalid (startup error) if a `defaultAuthProfile` name references a profile not defined on that cluster's `authProfiles`. - ---- - -## Default routing (when no router matches) - -When the router chain evaluates all configured routers (protocol-based, header, user-group, query-regex, client-tags, python-script) and **none produces a group selection**, QueryFlux applies a two-step fallback rather than blindly using the static `routingFallback` config key: - -1. **Authorization-aware first-fit**: enumerate all cluster groups in config order. For each group, check whether `AuthContext` is authorized (OpenFGA check or `allowGroups`/`allowUsers` match). Pick the **first group the user is authorized for**. -2. **Static fallback** (`routingFallback`): if the user is not authorized for any group (or has no identity), fall back to the static `routingFallback` group — same behavior as today, but only reached when step 1 finds nothing. - -**Why this order?** Clients route implicitly via router rules when they send headers/tags/regex-matching SQL. When they send nothing, they still belong to some team — the authorization layer already knows which groups they may access. Picking the first authorized group gives users a deterministic default without requiring them to always specify routing hints. The static `routingFallback` remains for unauthenticated or unauthorized requests (e.g. health probers, legacy clients with no identity). - -**Startup constraint:** If `auth.required: true`, an unauthenticated request is rejected before routing — `routingFallback` is not reached. If `auth.required: false`, `NoneAuthProvider` still derives a user from `sessionCtx.user()`; the first-fit check runs against that identity. - -**Config remains unchanged**: `routingFallback` is still a required top-level string. No new config key needed — the behavior is implicit when the router chain produces no match and an `AuthContext` is available. - -``` -RouterChain result = None - → for each group in config order: - if authz.check(auth_ctx, group) == allowed → use this group - → if none found → use routingFallback -``` - ---- - -## End-to-end dispatch (single query) - -1. **Authenticate** → `AuthContext`. -2. **Route** → cluster **group** name (router chain → authorization-aware first-fit → `routingFallback`). -3. **Authorize** → user may use that group (OpenFGA or `allowUsers` / `allowGroups`). -4. **ClusterManager** → pick **one member** of the group (strategy: RR, weighted, failover, engine affinity). -5. **Merge connection context** → load `ClusterConfig` for `member.cluster`, apply `member.connection` (engine-validated) + `member.defaultAuthProfile`. -6. **Resolve profile** → pick `auth` material (single `auth` or named `authProfiles`). -7. **`BackendIdentityResolver`** → `QueryCredentials` from `queryAuth` + `AuthContext` (token exchange, impersonate, service account, implicit header forward). -8. **Adapter** → submit query with merged **wire auth + engine session hints** (role, catalog, etc., per adapter). - -Audit logs should record: **`auth_ctx.user`**, **group**, **cluster**, **resolved profile**, **member index or id** (if useful). - ---- - -## Secret storage (operations) - -| Approach | Use when | -|----------|----------| -| **Vault / cloud Secrets Manager** (`secretRef` on `auth` / `authProfiles`) | Production default; rotation and audit at the secrets layer. | -| **Envelope encryption in Postgres** (ciphertext in DB, DEK wrapped by **KMS**) | Policy requires all config in DB; avoid a single static app-wide passphrase without rotation. | -| **Plain YAML / plain DB columns** | Dev and test only. | - -QueryFlux should resolve `secretRef` at **startup or config reload**, not on every query, unless operators explicitly need dynamic secrets. - ---- - -## Context - -QueryFlux is a universal SQL proxy routing queries across heterogeneous backends (Trino, DuckDB, StarRocks, ClickHouse, and future cloud platforms). Today there is **no verified frontend authentication** — on Trino HTTP, client headers may be forwarded to a Trino backend; there is no gateway-level JWT validation or OpenFGA. As it grows to multi-tenant use, it needs: - -1. **Frontend auth**: verify who the user is (AuthProvider — pluggable) -2. **Authorization**: decide what they can access (OpenFGA or simple policy) -3. **Backend identity**: propagate the right credentials to each engine per its capabilities - -### Multi-engine routing (frontend A → backend D) - -The design fits **heterogeneous routing**: any supported frontend (Trino HTTP, Postgres wire, …) can target any supported backend cluster type, as long as routers and SQL translation allow it. **Gateway auth and authz** depend only on `AuthContext` and cluster group — not on whether the backend is Trino or ClickHouse. **Backend identity** is always resolved **per selected cluster** via `queryAuth` + engine capabilities: the same user may hit Trino with forwarded JWT and ClickHouse with a **service account** in the same deployment. - -### Operator choices: static backend creds vs forwarding client creds - -| Approach | Meaning | When to use | -|----------|---------|-------------| -| **Static Type 1 only** (`queryAuth` omitted or `serviceAccount`) | No per-request resolution for the wire: every query uses `clusters[].auth`. User identity may still exist in `AuthContext` for audit, authz, and metrics. | Default for ClickHouse, StarRocks MySQL wire, DuckDB; safe baseline everywhere. | -| **Implicit header forwarding** (Trino adapter today) | Client `Authorization` / `X-Trino-*` from `SessionContext` are applied after cluster auth — client's `Authorization` **wins** if present. No separate `queryAuth` type; not “free security.” | Same-IdP Trino→Trino, dev, or locked-down networks where routing is narrow. | -| **`impersonate`** | Type 1 only on the wire + `X-Trino-User`; client `Authorization` **must** be suppressed. | Trino with file-based ACL when JWT passthrough is not used. | - -**Authorization (`provider: none`) and passthrough are independent.** Turning off OpenFGA/simple lists does **not** make forwarded client creds a substitute for gateway policy: anyone who can reach the gateway may get queries routed per router rules, and the **backend** decides what those creds allow. Do **not** auto-enable “forward everything” based solely on `authorization: none`. Prefer **explicit** per-cluster behavior (`queryAuth` + adapter rules). Emit a **startup warning** when `authorization.provider: none` and implicit `Authorization` forwarding is active on a frontend that routes to **multiple** cluster groups (broad blast radius). - -**`auth.required: true` with `NoneAuthProvider`** still means **no cryptographic proof** of identity — only network trust. Document clearly for operators. - ---- - -## Architecture Overview - -``` -Client (any protocol) - ↓ -Frontend Listener - ├─ Extract Credentials (protocol-specific) ← raw material for AuthProvider - └─ Build SessionContext (unverified, as today) - ↓ -AuthProvider.authenticate(credentials) → AuthContext - { user, groups, roles, raw_token } - Pluggable: None | Static | OIDC | LDAP - ↓ -RouterChain → ClusterGroup selection - (routers can inspect AuthContext.user/groups) - ↓ -OpenFGA / Policy check - "can user X execute queries on cluster group Y?" → allowed | 403 - ↓ -ClusterManager → pick GroupMember (cluster name + connection options + optional defaultAuthProfile) - ↓ -Merge ClusterConfig + member.connection (engine-specific hints) + resolved auth profile - ↓ -BackendIdentityResolver(AuthContext, cluster.queryAuth) → QueryCredentials - serviceAccount → cluster.auth (Type 1) - impersonate → cluster.auth + user identity header (suppress client Authorization) - tokenExchange → exchange raw_token at OAuth endpoint - (implicit: Trino adapter forwards SessionContext headers unchanged when no suppression needed) - ↓ -Adapter.submit_query(sql, QueryCredentials) ← Type 2 used here -Adapter.health_check() uses cluster.auth (Type 1) ← always independent - ↓ -Backend Engine -``` - ---- - -## Layer 1: Frontend Authentication - -### `AuthProvider` trait (new `queryflux-auth` crate) - -```rust -trait AuthProvider: Send + Sync { - async fn authenticate(&self, creds: &Credentials) -> Result; -} - -struct Credentials { - username: Option, - password: Option, // from Basic auth or wire handshake - bearer_token: Option, // from Authorization: Bearer - // Future: extensible fields or a sealed enum for mTLS principal, Kerberos, IAM delegation, etc. -} - -struct AuthContext { - user: String, - groups: Vec, - roles: Vec, - raw_token: Option, // original JWT, needed for tokenExchange -} -``` - -**Why gateway auth if clients already send credentials?** Client material (Basic, Bearer, wire username) is **input**. **AuthProvider** answers: is it **valid** (signature, LDAP bind, static password), and what is the **canonical subject** for policy? **Authorization** answers: what may that subject do **at QueryFlux** (which cluster groups)? **Query resolution** answers: what credentials go **on the wire to this engine** (often Type 1 only). Unverified headers (e.g. `X-Trino-User` alone) are trivial to forge from any client that can reach the gateway — so multi-tenant or untrusted networks need **verified** auth, not only forwarding. - -**Implementations:** -- `NoneAuthProvider` — derives identity from `sessionCtx.user()` only; no cryptographic verification. `auth.required: true` with this provider does **not** add JWT/signature checks — it only enforces that a username is present unless paired with **network trust** (VPC, mTLS at the load balancer). Make this explicit in operator docs. -- `StaticAuthProvider` — user/password map in config (dev/simple deployments) -- `OidcAuthProvider` — validates JWT signature against JWKS endpoint; extracts groups/roles from claims -- `LdapAuthProvider` — binds with user credentials to verify; extracts group membership from DN - -**Credential extraction per protocol (no password verification for wire protocols):** -- `TrinoHttp`: parse `Authorization` header → Basic or Bearer → `Credentials` -- `PostgresWire`: capture `user` from startup message → `Credentials { username, .. }` -- `MySqlWire`: capture `user` from handshake → `Credentials { username, .. }` -- `ArrowFlightSQL`: gRPC metadata bearer token → `Credentials { bearer_token, .. }` - -### Auth config block (in `queryflux-core/src/config.rs`): - -```yaml -auth: - provider: none | static | oidc | ldap - required: true # with NoneProvider: network-trust only, not cryptographic assurance - oidc: - issuer: https://... - jwksUri: https://... - audience: queryflux - groupsClaim: groups - rolesClaim: roles - ldap: - url: ldap://... - bindDn: cn=svc,... - userSearchBase: ou=users,... - static: - users: - alice: { password: "...", groups: [analysts] } -``` - -### Keycloak as OIDC provider - -Keycloak maps directly onto `OidcAuthProvider` — no special code: - -```yaml -auth: - provider: oidc - oidc: - issuer: https://keycloak.internal/realms/my-realm - jwksUri: https://keycloak.internal/realms/my-realm/protocol/openid-connect/certs - audience: queryflux-client - groupsClaim: groups # requires "Group Membership" token mapper on Keycloak client - rolesClaim: realm_access.roles -``` - -Keycloak also enables `tokenExchange` for backends: QueryFlux exchanges the user's access token for a backend-scoped token (requires Keycloak `token-exchange` preview feature and "Token Exchange" permission on the target client). This is configured in `clusters[].queryAuth`, not here. - ---- - -## Layer 2: Authorization via OpenFGA - -OpenFGA implements Google Zanzibar-style fine-grained authorization, stored and managed outside QueryFlux code. - -**Scope:** OpenFGA (and simple allowlists) answer **gateway** questions — e.g. “may this subject run queries against **cluster group** G?” They do **not** replace **engine-native** RBAC (Trino system access control, ClickHouse users, Ranger on StarRocks, etc.). Table/column policies remain on the engines unless the model is extended and kept in sync deliberately. - -**Authorization Model:** - -``` -type user -type group - relations - define member: [user] -type cluster_group - relations - define reader: [user, group#member] - define writer: [user, group#member] - define admin: [user, group#member] -``` - -**Check at dispatch time** (after routing, before query execution): - -```rust -openfga_client.check( - user: format!("user:{}", auth_ctx.user), - relation: "reader", - object: format!("cluster_group:{}", selected_group), -).await? // → allowed | 403 -``` - -**Tuple lifecycle (who writes authorization data):** -- **Bootstrap**: an init script or migration tool writes tuples from a seed file when QueryFlux first starts against a new OpenFGA store -- **Admin API**: `POST /admin/authz/tuples` (new endpoint) allows operators to grant/revoke access at runtime without redeploy -- **IdP sync (optional)**: a background task reads group memberships from the IdP (LDAP, Keycloak) and syncs group-member tuples into OpenFGA on a configured interval -- **Manual**: operators use the OpenFGA CLI or Playground directly against the OpenFGA store - -**Config:** - -```yaml -authorization: - provider: openfga | none - openfga: - url: http://openfga:8080 - storeId: "..." - credentials: - method: api_key - apiKey: "..." -``` - -**Fallback when `provider: none`**: simple `allowGroups`/`allowUsers` lists on each `clusterGroup` (same as trino-gateway's role approach). No external dependency. - -```yaml -clusterGroups: - analytics: - members: - - cluster: trino-prod - connection: - type: trino - catalog: hive - schema: default - - cluster: clickhouse-prod - connection: - type: clickHouse - database: analytics - authorization: # used only when provider: none - allowGroups: [analysts, admins] - allowUsers: [svc-etl] - - team-a-snowflake: - members: - - cluster: snowflake-prod - defaultAuthProfile: svc_readonly - connection: - type: snowflake - role: ANALYST_TEAM_A - warehouse: WH_TEAM_A - authorization: - allowGroups: [team-a] - - team-b-snowflake: - members: - - cluster: snowflake-prod - defaultAuthProfile: svc_etl - connection: - type: snowflake - role: ETL_TEAM_B - warehouse: WH_TEAM_B - authorization: - allowGroups: [team-b] -``` - -**Note:** `connection.type` should align with the cluster’s engine for that member; startup validation rejects mismatches. Bare strings in `members` remain supported for backward compatibility during migration. - ---- - -## Layer 3: Backend Identity (`queryAuth` modes) - -All modes configured under `clusters[].queryAuth` (per-cluster, not per-group). - -### Implicit header forwarding (Trino HTTP → Trino, no config needed) - -The Trino HTTP adapter forwards `SessionContext::TrinoHttp { headers }` verbatim to the backend — including `Authorization` and `X-Trino-User`. No separate `queryAuth` entry is required for this path; the default `serviceAccount` fallback does not suppress these headers in the Trino adapter because the Trino adapter applies session headers after cluster auth. - -However: when `queryAuth: impersonate` is set on a Trino cluster, the adapter **must suppress the client's `Authorization` header** and use only Type 1 credentials for authentication. The `X-Trino-User` injection happens after the service account auth is applied. Failing to suppress the client `Authorization` would cause the backend to see conflicting auth credentials. - -### Mode: `serviceAccount` - -Use Type 1 credentials (`cluster.auth`) for query execution. User identity is known to QueryFlux (logged in audit/metrics) but the backend sees only the service account. - -Works for all engines. Default when `queryAuth` is omitted. - -### Mode: `impersonate` (Trino only) - -Service account authenticates to the backend; user identity injected via `X-Trino-User` header. - -**Authorization header handling:** -1. Remove client's `Authorization` header from the outgoing request -2. Apply `cluster.auth` (Type 1, Basic or Bearer) as the backend authentication -3. Set `X-Trino-User: {auth_ctx.user}` header - -**Trino-side requirement** — Trino's built-in access control **prohibits impersonation by default**. File-based access control must be configured: - -```json -{ "impersonation": [{ "original_user": "qf_svc", "new_user": ".*", "allow": true }] } -``` -```properties -http-server.access-control.config-files=/etc/trino/rules.json -``` - -This is high operator burden. For OIDC deployments where Trino is configured with JWT auth pointing to the same IdP, prefer omitting `queryAuth` (implicit header forwarding) over `impersonate`. - -**Only Trino supports `impersonate`.** ClickHouse's `X-ClickHouse-User` is an auth username requiring a matching password — it is not an impersonation header and has no trusted-proxy mechanism. StarRocks has no equivalent over MySQL wire. Startup validation rejects `impersonate` for any other engine type. - -### Mode: `tokenExchange` (Snowflake, Databricks — future adapters) - -QueryFlux exchanges the user's OIDC JWT (`auth_ctx.raw_token`) for a backend-specific OAuth access token. Requires `OidcAuthProvider` on the frontend (so `raw_token` is populated). Falls back to `serviceAccount` if `raw_token` is absent. - -**Per-provider contract:** - -| Provider | Grant type | Subject token type | Audience / scope | -|----------|-----------|-------------------|-----------------| -| Keycloak token exchange | `urn:ietf:params:oauth:grant-type:token-exchange` | `urn:ietf:params:oauth:token-type:access_token` | `audience: ` | -| Snowflake external OAuth | `urn:ietf:params:oauth:grant-type:token-exchange` | `urn:ietf:params:oauth:token-type:access_token` | `scope: session:role:` | -| Databricks OAuth U2M | `urn:ietf:params:oauth:grant-type:token-exchange` | `urn:ietf:params:oauth:token-type:access_token` | `scope: all-apis` | - -Each provider must be registered as an OAuth client in the same IdP as QueryFlux. The exchanged token is used as `Authorization: Bearer ` in the adapter request. Token caching (with TTL from `expires_in`) should be implemented to avoid an exchange call on every query. - -```yaml -clusters: - snowflake-prod: - engine: snowflake # future adapter - auth: - type: keyPair - username: QF_SVC - privateKeyPem: "..." - queryAuth: - type: tokenExchange - tokenEndpoint: https://keycloak.internal/realms/my-realm/protocol/openid-connect/token - clientId: queryflux-gateway - clientSecret: "..." - # provider-specific extras: - targetAudience: snowflake-client # for Keycloak exchange - # scope: session:role:ANALYST # for direct Snowflake OAuth -``` - ---- - -## Engine-Specific Notes - -### ClickHouse -- No JWT/OIDC support; no impersonation mechanism -- `X-ClickHouse-User` + `X-ClickHouse-Key` are full auth credentials (username + password), not impersonation headers -- Only viable `queryAuth`: `serviceAccount` -- ClickHouse Cloud: further restricted to password-only (no LDAP/Kerberos/cert) - -**Trino HTTP frontend → ClickHouse backend:** Gateway **auth** still produces `AuthContext` (who the analyst is for authz and audit). Gateway **queryAuth** for the ClickHouse cluster resolves to **Type 1 service credentials** only. ClickHouse sees the service user, not the Trino username — unless operators add a custom integration (password mirroring, external authenticator). This is expected for heterogeneous routing. - -### StarRocks -- MySQL wire (port 9030, current adapter): password-based only; `serviceAccount` is the only option -- HTTP API (ports 8030/8040, future adapter): StarRocks natively supports JWT and OAuth 2.0; a future HTTP adapter could use implicit header forwarding when StarRocks and QueryFlux share an IdP. This is a motivation for the HTTP adapter — it enables per-user identity for StarRocks Ranger policies -- No impersonation mechanism exists on either interface - -### Snowflake (future adapter) -- No header-based impersonation -- Service account should use key-pair auth (RSA JWT), not password — Snowflake's recommended pattern for automated connections (as used by Yuki) -- `tokenExchange` or `serviceAccount` are the two options -- Private keys must not be stored in config files in production; use `secretRef` to Secrets Manager - -### Trino -- Implicit header forwarding works for same-IdP deployments -- `impersonate` requires Trino file-based ACL — high operator burden, document clearly -- For `impersonate`: suppress client `Authorization`; apply service account auth; inject `X-Trino-User` - ---- - -## Snowflake Key-Pair Auth (`ClusterAuth` extension) - -The existing `ClusterAuth` only supports `Basic` and `Bearer`. A `KeyPair` variant is needed for Snowflake (and Databricks): - -```rust -pub enum ClusterAuth { - Basic { username: String, password: String }, - Bearer { token: String }, - KeyPair { // NEW - username: String, - private_key_pem: String, // PEM string or secretRef - private_key_passphrase: Option, - }, -} -``` - -Future: support `secretRef` on any auth type so private keys are fetched from AWS Secrets Manager / Vault at startup, not stored in YAML: - -```yaml -auth: - type: keyPair - username: QF_SVC - privateKeySecretRef: - provider: awsSecretsManager - secretId: "arn:aws:secretsmanager:us-east-1:123:secret:qf-snowflake-key" - field: private_key -``` - ---- - -## What Changes Where - -### New: `crates/queryflux-auth/` -- `AuthProvider` trait, `Credentials` struct, `AuthContext` struct -- `NoneAuthProvider`, `StaticAuthProvider`, `OidcAuthProvider`, `LdapAuthProvider` -- `BackendIdentityResolver` — takes `(AuthContext, QueryAuthConfig, ResolvedProfile)` → `QueryCredentials` -- **`ConnectionContextMerge`** (or inline in dispatch) — merges `ClusterGroupMember.connection` into adapter-facing session hints -- `OpenFgaAuthorizationClient` — wraps OpenFGA HTTP API -- `SimpleAuthorizationPolicy` — fallback allowGroups/allowUsers -- Optional: **`SecretResolver` trait** — resolves `secretRef` to material for `ClusterAuth` / profiles at load time - -### `queryflux-core/src/config.rs` -- Add `AuthConfig` (provider + per-provider sub-configs) -- Add `AuthorizationConfig` (openFga | none) to `ProxyConfig` -- Add `QueryAuthConfig` enum (`serviceAccount | impersonate | tokenExchange`) to `ClusterConfig` -- Add **`authProfiles`** map + **`defaultAuthProfile`** optional field on `ClusterConfig`; support **`secretRef`** on credential fields (resolve at load/reload) -- Replace `ClusterGroupConfig.members: Vec` with **`Vec`**: `{ cluster, connection?: EngineConnectionOptions, weight?, defaultAuthProfile? }`; serde **untagged** or custom deserializer to accept **legacy string OR object** -- Add **`EngineConnectionOptions`** as a **tagged enum** (or per-engine struct union) listing **all supported per-engine connection types**; startup validation: each member’s `connection` matches `clusters[cluster].engine` -- Add `authorization` block (`allowGroups`/`allowUsers` fallback) to `ClusterGroupConfig` -- Extend `ClusterAuth` with `KeyPair` variant -- `QueryAuthConfig` validated at startup against engine type; error on unsupported combination - -### `queryflux-core/src/session.rs` -- No change. `SessionContext` stays as-is (unverified protocol metadata). -- `AuthContext` lives in `queryflux-auth`. - -### `queryflux-frontend/src/state.rs` -- Add `auth_provider: Arc` -- Add `authorization: Arc` - -### `queryflux-frontend/src/dispatch.rs` -- Accept `AuthContext` in `dispatch_query()` and `execute_to_sink()` -- **First step in `dispatch_query()`**: call `state.authorization.check(auth_ctx, group)` → 403 if denied (before `acquire_cluster`) -- After cluster pick: thread **`ClusterGroupMember`** (or equivalent) so adapters receive **merged** engine session hints (`connection`) + **resolved profile** (`auth` / `authProfiles`) -- Pass `QueryCredentials` (resolved by `BackendIdentityResolver`) to adapter alongside `SessionContext` (both needed until Phase 3b replaces session hints with `EngineConnectionOptions`) - -### `queryflux-routing/src/lib.rs` (RouterTrait) -- Update `RouterTrait.route()` signature to accept `Option<&AuthContext>` alongside `&SessionContext` and `&FrontendProtocol` -- `UserGroup` router and any future identity-aware router must use verified `AuthContext.user`, not `session.user()` (which is unverified) -- `NoneAuthProvider` still produces an `AuthContext` derived from `session.user()`, so the interface is consistent regardless of provider - -### `queryflux-frontend/src/trino_http/handlers.rs` -- Extract `Authorization` header → `Credentials` → `auth_provider.authenticate()` → `AuthContext` (before calling `route_with_trace()`) -- Pass `&auth_ctx` to routers -- **Default routing** (here, not in dispatch): if `route_with_trace()` returns `used_fallback == true`, iterate `state.group_configs` in config order; call `state.authorization.check(auth_ctx, group)` for each; pick first authorized group; only use static `routingFallback` if none found. `state` needs ordered group config list for this (add to `AppState`). -- Thread `AuthContext` through to dispatch - -### `queryflux-frontend/src/postgres_wire/mod.rs` -- Capture `user` from startup message → `Credentials { username, .. }` → `auth_provider.authenticate()` -- Same default routing logic as Trino HTTP handler -- Thread `AuthContext` through - -### `queryflux-engine-adapters/src/lib.rs` -- Add `QueryCredentials` enum alongside `SessionContext` — **not replacing it**. Until Phase 3b (EngineConnectionOptions), `SessionContext` still carries session hints (catalog, schema, X-Trino-* headers). Adapters need both: `QueryCredentials` for wire auth, `SessionContext` for session setup. -- Update `submit_query` / `execute_as_arrow` to accept `&QueryCredentials` as an additional parameter - -### `queryflux-engine-adapters/src/trino/mod.rs` -- `serviceAccount`: apply `cluster.auth` (Basic/Bearer); session headers forwarded as today -- `impersonate`: apply `cluster.auth`; **remove** client `Authorization` from headers; add `X-Trino-User: {user}` - -### `queryflux-engine-adapters/src/clickhouse/mod.rs` *(future — no module exists yet)* -- `serviceAccount` only: `X-ClickHouse-User` + `X-ClickHouse-Key` from `cluster.auth` - ---- - -## Phased Implementation - -### Phase 1 — Foundation: AuthContext plumbing + NoneProvider -- Define `AuthContext` / `Credentials` / `AuthProvider` / `QueryCredentials` types -- `NoneAuthProvider`: identity from `sessionCtx.user()`, no verification (current behaviour) -- Thread `AuthContext` and `QueryCredentials` through dispatch and adapter calls -- No behaviour change; all existing deployments unaffected - -### Phase 2 — Frontend auth (Trino HTTP first) -- `OidcAuthProvider`: JWT validation via JWKS, groups/roles extraction -- `StaticAuthProvider`: config-driven user/password map -- Extract `Authorization` header in Trino HTTP handlers → `Credentials` - -### Phase 3 — Authorization -- Simple `allowGroups`/`allowUsers` policy per cluster group (no external dep) -- OpenFGA client integration as optional provider -- Admin API endpoint for tuple management - -### Phase 3b — Structured group members & per-engine connection options -- Migrate `members` to `Vec` with backward-compatible deserializer for string entries -- Implement `EngineConnectionOptions` tagged enum covering **all engines** in the compatibility table; reject cross-engine field sets at startup -- Plumb **merged connection context** from selected member into dispatch and adapters (Trino catalog/schema, Snowflake role/warehouse, etc.) - -### Phase 3c — Auth profiles + secretRef *(can overlap with Phase 3b)* -- `authProfiles` / `defaultAuthProfile` on `ClusterConfig` and `ClusterGroupMember` -- Profile resolution order: group member default → cluster default → single `auth` (no client influence) -- `secretRef` resolution from Vault / AWS Secrets Manager at config load - -### Phase 4 — `impersonate` mode for Trino -- `BackendIdentityResolver` producing `ImpersonateCreds` -- Trino adapter: suppress client `Authorization`, apply service account auth, inject `X-Trino-User` -- Startup validation: reject `impersonate` for non-Trino engines - -### Phase 5 — LDAP + wire protocol auth -- `LdapAuthProvider` -- PG/MySQL wire: OIDC bearer as session parameter - -### Phase 6 — `tokenExchange` + cloud adapters -- `tokenExchange` resolver with per-provider contract and token caching -- Snowflake adapter (REST API, key-pair auth) -- Databricks adapter (SQL Warehouses REST or Arrow Flight SQL) - ---- - -## Key Files - -| File | Change | -|------|--------| -| [queryflux-core/src/config.rs](crates/queryflux-core/src/config.rs) | Add AuthConfig, QueryAuthConfig (3 variants), AuthorizationConfig; `ClusterGroupMember`, `EngineConnectionOptions` (per-engine variants); `authProfiles` + `defaultAuthProfile`; `secretRef`; extend ClusterAuth with KeyPair | -| [queryflux-core/src/session.rs](crates/queryflux-core/src/session.rs) | No change — AuthContext is in queryflux-auth | -| [queryflux-frontend/src/state.rs](crates/queryflux-frontend/src/state.rs) | Add auth_provider, authorization checker | -| [queryflux-frontend/src/dispatch.rs](crates/queryflux-frontend/src/dispatch.rs) | Thread AuthContext + QueryCredentials; authz check as first step before acquire_cluster | -| [queryflux-frontend/src/trino_http/handlers.rs](crates/queryflux-frontend/src/trino_http/handlers.rs) | Authenticate before routing; pass AuthContext to routers; authorization-aware first-fit when used_fallback==true | -| [queryflux-frontend/src/postgres_wire/mod.rs](crates/queryflux-frontend/src/postgres_wire/mod.rs) | Same as Trino HTTP: authenticate, pass AuthContext to routers, default routing | -| [queryflux-routing/src/lib.rs](crates/queryflux-routing/src/lib.rs) | Add Option<&AuthContext> to RouterTrait.route() signature; UserGroup router uses verified identity | -| [queryflux-engine-adapters/src/lib.rs](crates/queryflux-engine-adapters/src/lib.rs) | Add QueryCredentials alongside SessionContext in submit_query / execute_as_arrow signatures | -| [queryflux-engine-adapters/src/trino/mod.rs](crates/queryflux-engine-adapters/src/trino/mod.rs) | serviceAccount (current behaviour) + impersonate (suppress + inject) | -| New: `crates/queryflux-auth/` | AuthProvider trait + all implementations + BackendIdentityResolver + OpenFGA client | diff --git a/docs/motivation-and-goals.md b/docs/motivation-and-goals.md deleted file mode 100644 index c5c3dc3..0000000 --- a/docs/motivation-and-goals.md +++ /dev/null @@ -1,43 +0,0 @@ -# Motivation and goals - -## Why QueryFlux exists - -### This is why QueryFlux - -Modern data stacks are **fragmented by design**. Different engines exist because different problems demand different trade-offs — **Trino** for federation across sources, **DuckDB** for lightweight embedded analytics, **StarRocks** or **ClickHouse** for high-throughput serving layers. Using the right engine for the right job is good engineering. - -That fragmentation has a cost that compounds quietly. Every engine speaks its own wire protocol, has its own SQL dialect quirks, and needs its own connection management. When you multiply engines by clients — BI tools, notebooks, application code, CLI tools — you do not get a linear integration problem; you get a **combinatorial** one (**N×M**). Each pairing needs its own driver, dialect handling, and retry logic. Operational concerns like routing, rate limiting, and observability get reinvented in isolation, over and over. - -**Cost makes the picture worse.** Cloud analytical systems often charge either for **compute time** or for **data scanned**; query shapes skew toward **compute-bound** or **IO-bound** work, and each class tends to run cheaper under a different pricing model. In **our own benchmarking and cost modeling** while building QueryFlux — running the same analytical SQL across engines with different pricing models — we repeatedly saw a **poor match** between query shape and billing model inflate cost by large factors (on the order of **2–5×** versus a better-matched engine in those runs). When we prototyped **workload-aware routing** (steering CPU-skewed work toward compute-priced backends and scan-heavy work toward byte-priced backends where it helped), **total workload cost** in one representative suite fell by up to about **56%**, and **individual queries** sometimes dropped by up to about **90%** compared with always using a single default. That gap is not a rounding error — it is a structural inefficiency in any stack that treats every query the same. - -Without a **single proxy layer**, capturing those savings is hard: there is no one place to observe query characteristics, apply routing logic, and dispatch to the right engine. Teams end up locked into one engine or maintaining bespoke glue that cannot reason about cost or capability in one place. - -**QueryFlux** was created to cut through that. Instead of solving the N×M integration problem at the edges, it introduces **one proxy** that clients talk to over a familiar protocol. Behind it, QueryFlux handles **routing** to the right engine, **normalizes dialect** differences, and **centralizes** the operational glue that would otherwise scatter across the stack. That same layer is where **workload- and economics-aware routing** can live — matching queries to engines by capability and, over time, by how those engines bill for execution. - -The result is a **uniform client experience** — one URL, shared tooling, consistent behavior — without forcing consolidation onto a single engine or leaving money on the table by ignoring how differently engines charge for the same work. - -**QueryFlux** is a **universal SQL query proxy and router** written in Rust. It aims to: - -1. **Speak the client’s protocol** — Trino HTTP, PostgreSQL wire, MySQL wire, and Arrow Flight SQL are all implemented. -2. **Route each query to the right backend pool** — using configurable rules (protocol, headers, regex on SQL, Trino client tags, Python scripts) instead of hard-coding one engine per deployment. -3. **Translate SQL when dialects differ** — so a Trino-oriented client can still hit DuckDB or StarRocks when routing sends the query there, using [sqlglot](https://github.com/tobymao/sqlglot) behind the scenes. -4. **Operate like a serious proxy** — per-group concurrency limits, queueing when clusters are full, health-aware selection, Prometheus metrics, and optional PostgreSQL-backed state for HA-style deployments. - -In short: **one front door**, **many engines**, **explicit routing**, **automatic dialect bridging**, and **shared capacity and observability** — with a clear place to grow toward routing that respects **cost and workload shape**, grounded in what we measured above. - -## Compared to Trino-only gateways - -Some gateways optimize for **load-balanced Trino** behind one client protocol. QueryFlux targets **heterogeneous** deployments: - -- **Multiple engine types** in one deployment (Trino, DuckDB, StarRocks, …). -- **Protocol choices at the edge** — Trino HTTP, PostgreSQL wire, MySQL wire, Arrow Flight SQL — not only Trino HTTP. -- **SQL dialect translation** when the routed engine’s SQL differs from what the client naturally speaks. - -## What success looks like for operators - -- **Predictable routing**: Rules are data (YAML / DB-backed config), ordered and traceable. -- **Controlled blast radius**: Groups cap concurrent queries; full groups queue instead of overwhelming backends. -- **Observable behavior**: Metrics and admin APIs reflect group/cluster load and routing outcomes. -- **Incremental adoption**: You can start with a single group (e.g. one Trino pool) and add engines and routers as needs grow. - -For the mechanics of translation and routing, see [query-translation.md](query-translation.md) and [routing-and-clusters.md](routing-and-clusters.md). diff --git a/docs/observability.md b/docs/observability.md deleted file mode 100644 index e3ee856..0000000 --- a/docs/observability.md +++ /dev/null @@ -1,117 +0,0 @@ -# Observability - -QueryFlux exposes three observability surfaces: **Prometheus metrics** (real-time operational), a **Grafana dashboard** (visual ops view), and **QueryFlux Studio** (admin UI with query history, cluster management, and config). - ---- - -## Prometheus metrics - -QueryFlux exposes a `/metrics` endpoint (default port 9000) in standard Prometheus text format. - -**Scrape target:** `http://:9000/metrics` - -### Exposed metrics - -| Metric | Type | Labels | Description | -|--------|------|--------|-------------| -| `queryflux_queries_total` | Counter | `engine_type`, `cluster_group`, `status`, `protocol` | Total queries by outcome and engine | -| `queryflux_query_duration_seconds` | Histogram | `engine_type`, `cluster_group` | End-to-end query duration from proxy receipt to result delivery | -| `queryflux_translated_queries_total` | Counter | `src_dialect`, `tgt_dialect` | Queries where SQL dialect translation ran | -| `queryflux_running_queries` | Gauge | `cluster_group`, `cluster_name` | Currently executing queries per cluster | -| `queryflux_queued_queries` | Gauge | `cluster_group` | Queries waiting for a free cluster slot | - -The metrics pipeline uses `MultiMetricsStore` to fan out to Prometheus (real-time) and optionally Postgres (historical). `BufferedMetricsStore` wraps the Postgres store to avoid blocking query execution on I/O. - -### Prometheus config - -The `prometheus/prometheus.yml` file configures a single scrape job: - -```yaml -scrape_configs: - - job_name: queryflux - static_configs: - - targets: ["host.docker.internal:9000"] - metrics_path: /metrics -``` - -When running via `make dev`, Prometheus starts in Docker and reaches QueryFlux on the host via `host.docker.internal`. On Linux, the `docker-compose.yml` adds `extra_hosts: ["host.docker.internal:host-gateway"]` to bridge this. - ---- - -## Grafana dashboard - -A pre-built dashboard (`grafana/dashboards/queryflux.json`) is auto-provisioned when Grafana starts via Docker Compose. It auto-refreshes every 10 seconds. - -**Access:** http://localhost:3000 (credentials: `admin` / `admin`) - -### Panels - -| Panel | Type | What it shows | -|-------|------|---------------| -| Query Rate | Stat | Queries per minute, current window | -| Error Rate | Stat | Fraction of failed queries (%) | -| p95 Latency | Stat | 95th-percentile end-to-end duration | -| Translation Rate | Stat | Fraction of queries that ran through sqlglot | -| Query Throughput by Status | Time series | Success / Failed / Cancelled over time | -| Query Throughput by Engine | Time series | Per-engine query rate over time | -| Latency Percentiles (p50 / p95 / p99) | Time series | Latency distribution trends | -| SQL Translations by Dialect Pair | Time series | Translation volume per src→tgt pair | -| Query Throughput per Cluster | Time series | Per-cluster query rate | -| Queued Queries per Cluster Group | Time series | Queue depth per group — spikes indicate saturation | - -### Provisioning - -Grafana is auto-configured via two provisioning files: - -- `grafana/provisioning/datasources/prometheus.yml` — registers the Prometheus datasource pointing at `http://prometheus:9090` -- `grafana/dashboards/` — dashboards loaded automatically at startup; no manual import needed - ---- - -## QueryFlux Studio (admin UI) - -QueryFlux Studio is a Next.js web UI served separately from the proxy. It talks to the Admin REST API on port 9000. - -**Start:** `cd ui/queryflux-studio && npm run dev` (or build and serve for production) - -**Default URL:** http://localhost:3001 - -### Pages - -| Page | What it shows | -|------|---------------| -| Dashboard | Live query rate, error rate, avg latency, translation rate; cluster health grid; recent queries | -| Clusters | All cluster groups and member clusters — health, running/queued counts, enable/disable, max concurrency | -| Queries | Searchable, filterable query history — SQL, status, duration, engine, routing trace | -| Engines | Engine registry — supported engines, connection types, config fields | - -Studio requires **Postgres persistence** to be configured — query history, cluster config, and dashboard stats are read from the DB. Without Postgres, the clusters page works (from in-memory state) but history pages return empty. - ---- - -## Admin REST API - -The admin API is served on port 9000 alongside `/metrics`. An OpenAPI spec is available at `http://localhost:9000/openapi.json` and the Swagger UI at `http://localhost:9000/docs`. - -### Endpoints - -| Method | Path | Description | -|--------|------|-------------| -| `GET` | `/health` | Health check | -| `GET` | `/metrics` | Prometheus metrics | -| `GET` | `/admin/clusters` | Live cluster state (all groups) | -| `PATCH` | `/admin/clusters/{group}/{cluster}` | Update cluster (enable/disable, max concurrency) | -| `GET` | `/admin/queries` | Query history (paginated, filterable) | -| `GET` | `/admin/stats` | Aggregated stats for the last hour | -| `GET` | `/admin/engine-stats?hours=N` | Per-engine aggregated stats | -| `GET` | `/admin/group-stats?hours=N` | Per-cluster-group aggregated stats | -| `GET` | `/admin/engines` | Distinct engine types in query log | -| `GET` | `/admin/engine-registry` | Full engine descriptor catalog | -| `GET/PUT/DELETE` | `/admin/config/clusters/{name}` | Cluster config CRUD (Postgres required) | -| `GET` | `/admin/config/clusters` | List all persisted cluster configs | -| `GET/PUT/DELETE` | `/admin/config/groups/{name}` | Cluster group config CRUD (Postgres required) | -| `GET` | `/admin/config/groups` | List all persisted group configs | -| `GET` | `/openapi.json` | OpenAPI spec | -| `GET` | `/docs` | Swagger UI | - -Config CRUD endpoints (`/admin/config/*`) require Postgres persistence. The in-memory store supports reading live cluster state but not persisted config management. diff --git a/docs/query-translation.md b/docs/query-translation.md deleted file mode 100644 index 6913c90..0000000 --- a/docs/query-translation.md +++ /dev/null @@ -1,149 +0,0 @@ -# Query translation - -This document explains **how** QueryFlux converts SQL between dialects, **when** that happens, and how it fits into the query path. - -## Role in the pipeline - -Translation runs **after** routing has chosen a **cluster group** and **after** the cluster manager has selected a **concrete cluster** (adapter), but **before** the SQL is submitted or executed on the backend. - -Conceptually: - -``` -Client SQL - → routers pick cluster group - → cluster manager picks cluster (adapter) - → translate(client dialect → engine dialect) ← this document - → adapter.submit_query / execute_as_arrow -``` - -The implementation lives mainly in the `queryflux-translation` crate (`TranslationService`, `SqlglotTranslator`) and is invoked from shared dispatch code in `queryflux-frontend` (`dispatch_query`, `execute_to_sink`). - -## Source and target dialects - -- **Source dialect** comes from the **frontend protocol**: each `FrontendProtocol` has a `default_dialect()` (e.g. Trino HTTP → Trino, MySQL wire → MySQL). See `queryflux_core::query::FrontendProtocol`. -- **Target dialect** comes from the **engine type** of the chosen adapter: `EngineType::dialect()` (e.g. DuckDB → DuckDB, StarRocks → StarRocks). See `queryflux_core::query::EngineType`. - -If source and target are considered **compatible**, translation is skipped entirely (no sqlglot call). Notably, **MySQL and StarRocks** are treated as mutually compatible in `SqlDialect::is_compatible_with`, reflecting similar client SQL expectations. - -## TranslationService and sqlglot - -`TranslationService` is the façade used by the frontend: - -- **`new_sqlglot(python_scripts)`** — Verifies that Python can import `sqlglot` (via PyO3) and stores the global fixup scripts. If that fails at startup, the binary logs a warning and falls back to **`disabled()`**, which passes SQL through unchanged. -- **`maybe_translate(sql, src, tgt, schema, group_fixups)`** — If translation is disabled, or dialects are compatible, returns the original string. Otherwise it constructs a `SqlglotTranslator` with **global** YAML `translation.pythonScripts` plus **per-group** fixup scripts from Postgres (`user_scripts` rows attached to the cluster group, ordered by their position in `translation_script_ids`), then runs translation. - -`SqlglotTranslator` runs work on a **blocking thread pool** (`spawn_blocking`) because it holds the Python GIL. Inside Python it either: - -1. **Dialect-only** — When `SchemaContext` is empty: `sqlglot.transpile(sql, read=, write=)`. -2. **Schema-aware** — When tables/columns are populated: parse with `parse_one`, build a `MappingSchema`, run `sqlglot.optimizer.optimize`, then emit SQL with the target dialect. If optimization fails, it **falls back** to dialect-only behavior (with a warning). - -The Rust type `SchemaContext` (`queryflux_translation::SchemaContext`) carries optional catalog/database and a map of **table → column → SQL type string** for sqlglot's schema-aware path. - -### Current default on the hot path - -Today, dispatch passes **`SchemaContext::default()`** (empty tables). So in production query paths you get **dialect-only** transpilation. The schema-aware branch is **implemented** in `sqlglot.rs` and is ready for future wiring (e.g. catalog providers or static schema config) to populate `SchemaContext` before `maybe_translate`. - -## Passthrough and performance - -When the client dialect matches the engine (e.g. Trino client → Trino cluster), `maybe_translate` returns immediately with **no Python work**. That keeps the common "Trino in, Trino out" case cheap. - -## Configuration - -`translation` in the root config (`queryflux_core::config::TranslationConfig`) includes: - -- **`errorOnUnsupported`** — Controls strictness when sqlglot cannot translate a construct. `false` (default) passes the original construct through best-effort; `true` fails the query. -- **`pythonScripts`** — List of global Python transform scripts run after sqlglot translation. See the next section. - -See `config.local.yaml` / your deployment YAML for concrete values. - -## Python transform scripts - -After sqlglot finishes translation, QueryFlux runs each script in order — first the global `translation.pythonScripts` from YAML, then any per-group scripts attached to the cluster group via the Admin UI. This is an escape hatch for structural transformations that sqlglot does not handle on its own — things like stripping catalog prefixes, renaming functions, or applying environment-specific rewrites. - -### Script contract - -Each script must define a `transform` function: - -```python -def transform(ast, src: str, dst: str) -> None: - ... -``` - -| Parameter | Type | Description | -|-----------|------------------|-------------------------------------------------------------| -| `ast` | `sqlglot.Expression` | Root AST node of the **already-translated** SQL — mutate in-place | -| `src` | `str` | Source dialect name (sqlglot name, e.g. `"trino"`) | -| `dst` | `str` | Target dialect name (sqlglot name, e.g. `"athena"`) | - -Top-level imports and helper functions are fully supported — the script is executed as a module before `transform` is called. QueryFlux re-serializes the AST using the target dialect once, **after all scripts have run**. - -### When scripts run - -Scripts run whenever they are configured, **including when `src == dst`** (dialects are compatible). The only time scripts are skipped entirely is when no scripts are configured and the dialects are compatible — in that case the SQL is returned unchanged with zero overhead. - -When dialects are compatible but scripts are present, sqlglot dialect translation is skipped, the SQL is parsed into an AST in the target dialect, and your scripts run against that AST. Use `src`/`dst` guards inside `transform` to apply logic only to specific pairs. - -### Example — strip catalog prefix for Athena - -Trino clients use three-part names (`catalog.database.table`). Athena has no catalog layer and expects `database.table`. sqlglot preserves the catalog structurally, so a transform script is needed: - -```yaml -translation: - pythonScripts: - - | - import sqlglot.expressions as exp - - def transform(ast, src: str, dst: str) -> None: - if dst == "athena": - for table in ast.find_all(exp.Table): - table.set("catalog", None) -``` - -### Example — multiple scripts - -Scripts are composable. Each sees the same `ast` (as mutated by previous scripts), so they chain: - -```yaml -translation: - pythonScripts: - - | - import sqlglot.expressions as exp - - def transform(ast, src: str, dst: str) -> None: - # Strip catalog when targeting Athena (any source dialect) - if dst == "athena": - for table in ast.find_all(exp.Table): - table.set("catalog", None) - - | - import sqlglot.expressions as exp - - def transform(ast, src: str, dst: str) -> None: - # Force uppercase schema names in DuckDB (environment-specific convention) - if dst == "duckdb": - for table in ast.find_all(exp.Table): - db = table.args.get("db") - if db: - db.set("this", db.name.upper()) -``` - -### Per-group scripts - -In addition to global YAML scripts, you can attach **reusable scripts** to individual cluster groups via the Admin UI (**Scripts** page → **Groups** page). Per-group scripts run after the global ones and follow the same `transform(ast, src, dst)` contract. This is useful when different groups target different engines and need distinct fixups. - -### Error handling - -If a script raises a Python exception, the query fails with a `Translation` error and the SQL is **not** sent to the backend. The error message includes the script index and the Python traceback. Scripts do not affect queries that skip translation (compatible dialects). - -### Implementation notes - -- Scripts run inside a `spawn_blocking` task on Tokio's blocking thread pool because they hold the Python GIL. -- Each script is executed in its own globals dict (same approach as `PythonScriptRouter`), so imports and helper functions defined at module level work correctly. -- The `ast` is parsed from the sqlglot-translated SQL in the **target dialect** before scripts run. Mutations do not need to account for the source dialect's syntax. -- Re-serialization happens once at the end via `ast.sql(dialect=dst)`, keeping overhead independent of the number of scripts. - -## Failure modes - -- **sqlglot missing** — Startup degrades to a disabled translation service; SQL is sent as-is, which may fail on the backend if dialects differ. -- **Translation errors** — Dispatch releases the acquired cluster slot and returns an error to the client (async Trino path logs and propagates; sync `execute_to_sink` path reports via the result sink). - -For how routing picks the group and cluster **before** translation, see [routing-and-clusters.md](routing-and-clusters.md). diff --git a/docs/routing-and-clusters.md b/docs/routing-and-clusters.md deleted file mode 100644 index 77841f0..0000000 --- a/docs/routing-and-clusters.md +++ /dev/null @@ -1,158 +0,0 @@ -# Routing, cluster groups, and clusters - -QueryFlux separates **where a query should go logically** (cluster **group**) from **which physical backend instance** serves it (**cluster** / adapter). This document explains that split, how routers work, and how the cluster manager balances load and enforces limits. - -## Vocabulary - -| Term | Meaning | -|------|--------| -| **Cluster** | A named backend instance in config (`clusters.`): one engine (Trino, DuckDB, StarRocks, …), endpoint or DB path, auth, etc. At runtime it has an **engine adapter** and a **`ClusterState`** (running count, health, limits). | -| **Cluster group** | A named **pool** (`clusterGroups.`) listing **member** cluster names, a **per-group** `maxRunningQueries`, optional `maxQueuedQueries`, optional **selection strategy**, and `enabled`. Routing returns a **group**; the cluster manager then picks a **member** cluster. | -| **Router** | A rule that inspects the SQL string, session context, and frontend protocol and optionally returns a target group name. | -| **Router chain** | All routers in config order, plus **`routingFallback`** when every router returns “no match.” | - -## Two-stage placement - -1. **Routing (group selection)** - `RouterChain` evaluates each `RouterTrait` implementation in order. The **first** router that returns `Some(ClusterGroupName)` wins. If all return `None`, **`routingFallback`** is used. - Implementation: `queryflux_routing::chain::RouterChain` (`route`, `route_with_trace`). - -2. **Cluster selection (member selection)** - `ClusterGroupManager::acquire_cluster(group)` considers only clusters in that group that are **enabled**, **healthy**, and **under** `max_running_queries`. It then uses the group’s **strategy** to pick one member and increments that cluster’s running count. - If **no** member is eligible (e.g. all at capacity or unhealthy), `acquire_cluster` returns **`None`** → the query is **queued** (Trino HTTP async path) or **retried with backoff** (sync `execute_to_sink` path). - Implementation: `queryflux_cluster_manager::simple::SimpleClusterGroupManager`. - -When the query finishes (success, failure, or cancel), **`release_cluster`** decrements the running count on that cluster. - -## Router types (config → code) - -Configured under `routers:` in YAML (`queryflux_core::config::RouterConfig`). Wired in `queryflux/src/main.rs`. - -| `type` | Behavior | -|--------|----------| -| `protocolBased` | Maps the active frontend (`trinoHttp`, `postgresWire`, `mysqlWire`, `flightSql`, `clickhouseHttp`) to a group name. | -| `header` | Matches a header value to a group (useful for Trino HTTP and similar). | -| `queryRegex` | Ordered rules: first regex match on the SQL text wins. | -| `clientTags` | Trino-style client tags header mapped to groups. | -| `pythonScript` | Embedded or file-backed Python `route(query, ctx)` returning a group name or `None`. See [Python script router](#python-script-router-pythonscript) below. | - -All five router types are implemented. Unknown `type` values in config are skipped at startup with a warning. - -## Cached routing config and DB reload (Postgres) - -When **`persistence.type`** is **`postgres`**, routing rules and cluster/group definitions loaded from the database are held in memory inside **`LiveConfig`** (including the compiled **`RouterChain`**). Each request reads the current chain from that shared snapshot (`Arc>` in `queryflux-frontend`). - -- **Periodic refresh:** `queryflux.configReloadIntervalSecs` in YAML (default **30** when omitted) controls how often a background task re-reads Postgres and replaces **`LiveConfig`** in one atomic swap. Implementation: `crates/queryflux/src/main.rs` (reload task) and `reload_live_config` → `load_routing_config`. -- **`0` disables polling only:** With **`configReloadIntervalSecs: 0`**, there is no timer-driven refresh; the in-memory config stays as loaded at startup until an **immediate refresh** runs (below). -- **Immediate refresh:** After Studio/admin API writes to routing, clusters, or groups, the proxy **notifies** the same task so a reload runs without waiting for the interval (`config_reload_notify` in `admin.rs`). -- **YAML-only mode:** With **`inMemory`** persistence there is no DB reload loop; routing comes from the process config at startup until restart. - -## Python script router (`pythonScript`) - -The script must define: - -```python -def route(query: str, ctx: dict) -> str | None: - ... -``` - -- **`query`**: SQL text (the same string the router chain sees). -- **`ctx`**: plain `dict` built by QueryFlux (string keys). **`protocol`** is always set; other keys depend on the frontend and session shape. - -| Key | When | Meaning | -|-----|------|---------| -| `protocol` | Always | One of `trinoHttp`, `postgresWire`, `mysqlWire`, `clickHouseHttp`, `flightSql` (camelCase, matching config / API). | -| `headers` | Always | `dict[str, str]`. Client headers for HTTP-style frontends (Trino HTTP uses lowercase keys as stored by the proxy, e.g. `x-trino-user`). Empty `{}` for Postgres and MySQL wire. | -| `database`, `user` | Postgres wire | From startup / auth; each may be Python `None`. | -| `sessionParams` | Postgres wire | `dict[str, str]` (parameters from `SET`). | -| `schema`, `user` | MySQL wire | Current schema and user; may be `None`. | -| `sessionVars` | MySQL wire | `dict[str, str]` (`SET SESSION`). | -| `queryParams` | ClickHouse HTTP | URL query string parameters (e.g. `?database=…`). | -| `auth` | When the request was authenticated | `{"user": str, "groups": [str, …], "roles": [str, …]}`. Raw JWT / bearer tokens are **not** passed into Python. | - -**Flight SQL** reports `protocol: "flightSql"` but **`SessionContext` is still Trino-style**: `headers` are built from gRPC metadata (see `queryflux-frontend` Flight SQL service). - -**Example (Trino HTTP):** - -```python -def route(query: str, ctx: dict) -> str | None: - if ctx.get("protocol") != "trinoHttp": - return None - user = (ctx.get("headers") or {}).get("x-trino-user") - if user == "batch": - return "heavy-trino" - return None -``` - -## Routing trace - -`route_with_trace` records each router’s decision (`matched`, optional `result`) and whether the **fallback** group was used. This supports debugging and future UI/metrics (see `RoutingTrace` in `queryflux_routing::chain`). - -## Cluster group configuration (actual shape) - -Clusters are defined **once** at the top level; groups **reference** them by name: - -```yaml -clusters: - trino-1: - engine: trino - endpoint: http://trino:8080 - enabled: true - duckdb-1: - engine: duckDb - enabled: true - -clusterGroups: - trino-default: - enabled: true - maxRunningQueries: 100 - members: [trino-1] - strategy: - type: leastLoaded - duckdb-local: - enabled: true - maxRunningQueries: 4 - members: [duckdb-1] -``` - -Notes: - -- **`maxRunningQueries`** on the group applies to **each** member cluster’s `ClusterState` when those states are built (see `main.rs` pass 2). It is the cap used for **acquire** / capacity checks. -- **`members`** can list multiple clusters, including **mixed engines** (e.g. Trino and DuckDB in one group). For that, **`engineAffinity`** (or another strategy) helps express preference order across engine types (`queryflux_cluster_manager::strategy`). - -## Selection strategies - -Configured as `strategy: { type: ... }` on a group. Implemented in `strategy.rs`: - -| Strategy | Behavior | -|----------|----------| -| `roundRobin` | Rotates among eligible members (default when strategy omitted). | -| `leastLoaded` | Picks the member with the smallest `running_queries`. | -| `failover` | First eligible member in **member list order** (priority ordering in YAML). | -| `engineAffinity` | Ordered engine preference; within each engine, least loaded. | -| `weighted` | Distributes by configured weights (deterministic pseudo-random from load). | - -Eligible candidates are always **healthy**, **enabled**, and **not at capacity** before the strategy runs. - -## Health and runtime updates - -Each `ClusterState` tracks health (`is_healthy`), updated by background health checks in the main binary. Unhealthy clusters are excluded from acquisition. - -The `ClusterGroupManager` trait also supports **`update_cluster`** (enable/disable, change `max_running_queries`) for admin-driven changes. - -## Frontend dispatch: async vs sync - -After a group is chosen, the Trino HTTP handler (`post_statement`) branches: - -- If the group is considered **async-capable** (e.g. Trino-style polling), it uses **`dispatch_query`**: acquire cluster → translate → `submit_query` → persist executing state → rewrite `nextUri` to point back at QueryFlux. -- Otherwise it uses **`execute_to_sink`**: wait for capacity (backoff loop), translate, stream Arrow batches, and synthesize a Trino-compatible JSON response. - -So **routing and cluster selection** are shared concepts; the **result delivery** shape depends on engine and frontend capabilities. - -## Mental model - -- **Routers** answer: *which pool (group) should handle this query?* -- **Cluster manager + strategy** answer: *which replica/instance in that pool?* -- **Translation** (separate doc) then aligns SQL with that instance’s engine. - -See [architecture.md](architecture.md) for the full component diagram and [query-translation.md](query-translation.md) for dialect conversion details. diff --git a/queryflux-studio/README.md b/queryflux-studio/README.md index c0cad21..efdd850 100644 --- a/queryflux-studio/README.md +++ b/queryflux-studio/README.md @@ -14,7 +14,7 @@ Open [http://localhost:3000](http://localhost:3000). Postgres-backed features (q ## Adding or changing a backend in the UI -Use **Part B — QueryFlux Studio** in **[docs/adding-engine-support.md](../../docs/adding-engine-support.md)**. Short version: +Use the **QueryFlux Studio** section in **[website/docs/architecture/adding-support/backend.md](../website/docs/architecture/adding-support/backend.md)**. Short version: 1. Add **`lib/studio-engines/engines/.ts`** exporting a **`StudioEngineModule`** (descriptor + catalog metadata + optional validation and custom form id). 2. Register it in **`lib/studio-engines/manifest.ts`**. diff --git a/queryflux-studio/app/protocols/page.tsx b/queryflux-studio/app/protocols/page.tsx index 46a0cb5..b833273 100644 --- a/queryflux-studio/app/protocols/page.tsx +++ b/queryflux-studio/app/protocols/page.tsx @@ -1,26 +1,45 @@ import { getFrontendsStatus } from "@/lib/api"; import type { ProtocolFrontendDto } from "@/lib/api-types"; +import { AlertCircle } from "lucide-react"; +import type { SimpleIcon } from "simple-icons"; import { - AlertCircle, - Cylinder, - Globe2, - LayoutGrid, - Plug, - Plane, - Radio, - type LucideIcon, -} from "lucide-react"; + siApacheparquet, + siClickhouse, + siMysql, + siOpenapiinitiative, + siPostgresql, + siSnowflake, + siTrino, +} from "simple-icons"; export const revalidate = 10; -const PROTOCOL_ICONS: Record = { - trino_http: Globe2, - mysql_wire: Plug, - postgres_wire: Cylinder, - clickhouse_http: LayoutGrid, - flight_sql: Plane, +/** [Simple Icons](https://simpleicons.org/) per `protocol.id` from `GET /admin/frontends`. */ +const PROTOCOL_SIMPLE_ICONS: Record = { + trino_http: siTrino, + mysql_wire: siMysql, + postgres_wire: siPostgresql, + clickhouse_http: siClickhouse, + flight_sql: siApacheparquet, + snowflake_http: siSnowflake, + snowflake_sql_api: siSnowflake, }; +function SimpleIconSvg({ icon, className }: { icon: SimpleIcon; className?: string }) { + return ( + + {icon.title} + + + ); +} + export default async function ProtocolsPage() { let status: Awaited> | null = null; let error: string | null = null; @@ -77,8 +96,8 @@ export default async function ProtocolsPage() { } function ProtocolCard({ protocol }: { protocol: ProtocolFrontendDto }) { - const Icon = PROTOCOL_ICONS[protocol.id] ?? Radio; const on = protocol.enabled; + const icon = PROTOCOL_SIMPLE_ICONS[protocol.id] ?? siOpenapiinitiative; return (
- +
diff --git a/website/docs/architecture/adding-engine-support.md b/website/docs/architecture/adding-engine-support.md index 89a4e3d..32f999e 100644 --- a/website/docs/architecture/adding-engine-support.md +++ b/website/docs/architecture/adding-engine-support.md @@ -1,314 +1,13 @@ --- -description: How to add new backend engines and frontend client protocols to QueryFlux — adapters, frontends, and integration points. +description: This page moved — see Extending QueryFlux (overview, backend, frontend). --- # Adding engine and protocol support -This guide separates two ideas that are easy to conflate: +This guide lives under **Extending QueryFlux** in the sidebar: -| Concept | Meaning | Example | -|--------|---------|---------| -| **Backend engine** | A **cluster** type QueryFlux routes queries **to**. It has an adapter that talks to the real database (HTTP, MySQL wire, embedded library, AWS SDK, …). | Trino, DuckDB, StarRocks, Athena | -| **Frontend protocol** | How **clients connect to QueryFlux** (ingress). SQL enters with a `FrontendProtocol` and a default source dialect for translation. | Trino HTTP, **PostgreSQL wire**, MySQL wire, Flight SQL | +- **[Overview](adding-support/overview)** — concepts and index +- **[Backend](adding-support/backend)** — Rust adapter, persistence, Studio +- **[Frontend](adding-support/frontend)** — new client protocols -Adding **PostgreSQL wire** as a client entrypoint is **not** the same as adding “PostgreSQL” as a backend: today, `PostgresWire` is already a frontend in `queryflux-frontend`; traffic still lands on the shared dispatch path and is sent to whatever **backend adapter** routing chose (often Trino). - -Use the sections below depending on whether you are extending **Studio**, a **backend adapter**, or a **frontend listener**. - ---- - -## Part A — Backend engine (Rust) - -Goal: a new `engine` value in cluster config, a live adapter, validation, translation target dialect, and wiring in the binary. - -### Registration overview - -Backends are **not** loaded dynamically. Each engine is compiled in and registered explicitly. Data flow: - -1. **Postgres / YAML** → `engine_key` column + `config` JSONB → `ClusterConfigRecord::to_core()` uses **`parse_engine_key`** and JSON helpers → typed **`ClusterConfig`**. -2. **Binary** → `registered_engines::build_adapter(...)` matches **`EngineConfig`** and calls the adapter’s **`try_from_cluster_config`** (see [`crates/queryflux/src/registered_engines.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux/src/registered_engines.rs)). -3. **Adapter** → reads only the **`ClusterConfig`** fields it needs (endpoint, auth, region, …) and constructs itself; **startup** and **hot reload** both use the same factory. - -**JSONB** stores per-cluster, per-engine payload without schema migrations; **`ClusterConfig`** in core is the typed view after `to_core()`. Engine-specific wiring belongs in **`try_from_cluster_config`**, not in `main.rs`. - -### 1. Core model (`queryflux-core`) - -- **`EngineConfig`** — Add a variant in `crates/queryflux-core/src/config.rs` (serde **camelCase** in JSON/YAML, e.g. `myEngine`). -- **`EngineType`** — Add a variant in `crates/queryflux-core/src/query.rs` if the backend is distinct for metrics, translation, or dispatch. -- **`engine_registry`** (`crates/queryflux-core/src/engine_registry.rs`) — Keep these in sync when you add a variant: - - **`engine_key(&EngineConfig)`** — `EngineConfig` → stable string key (must match the adapter descriptor and Studio). - - **`parse_engine_key(&str)`** — inverse mapping for the `engine_key` column in Postgres / API. - - **`impl From<&EngineConfig> for EngineType`** — single place for config → runtime `EngineType` (cluster manager and `main.rs` use this instead of ad-hoc matches). -- **`EngineType::dialect()`** — Return the `SqlDialect` used as the **translation target** (and extend `SqlDialect` / `is_compatible_with` in translation if needed). See [query-translation.md](query-translation.md). -- **`ClusterConfig` fields** — Add any new top-level fields (region, paths, engine-specific blobs). Prefer keeping engine-specific secrets and options in `config` JSON for Postgres-backed clusters; extend the typed struct when YAML and validation need them everywhere. - -### 2. Adapter crate (`queryflux-engine-adapters`) - -- Add a module (e.g. `src/myengine/mod.rs`) implementing **`EngineAdapterTrait`** (`submit_query`, `poll_query`, `cancel_query`, `health_check`, `engine_type`, `supports_async`, and optionally `fetch_running_query_count`, `base_url`, Arrow/catalog hooks as needed). -- Implement **`descriptor() -> EngineDescriptor`** with: - - `engine_key`, `display_name`, `description`, `hex` - - `connection_type` (`Http`, `MySqlWire`, `Embedded`, `ManagedApi`) - - `supported_auth` and **`config_fields`** (these drive `/admin/engine-registry` and should stay aligned with Studio) - - `implemented: true` when the adapter is actually wired in `main` -- Export the module from `crates/queryflux-engine-adapters/src/lib.rs` and add the crate dependency if you introduce new third-party crates. - -**Factory — `try_from_cluster_config`** - -Implement on your adapter struct so all **field extraction and validation** for that engine live next to the adapter (not in `registered_engines.rs`): - -- **Sync** (most engines): - - ```text - fn try_from_cluster_config( - cluster_name: ClusterName, - group_name: ClusterGroupName, - cfg: &ClusterConfig, - cluster_name_str: &str, - ) -> queryflux_core::error::Result - ``` - -- **Async** (e.g. Athena — AWS client setup): same parameters, `async fn`, returns `Result`. - -Use **`QueryFluxError::Engine(format!(…))`** for failures; include **`cluster_name_str`** in messages so startup and reload logs identify the cluster. Reference implementations: **`TrinoAdapter`** and **`StarRocksAdapter`** ([`trino/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/trino/mod.rs), [`starrocks/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/starrocks/mod.rs)), **`DuckDbAdapter`** / **`DuckDbHttpAdapter`** ([`duckdb/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/duckdb/mod.rs), [`duckdb/http.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/duckdb/http.rs)), **`AthenaAdapter`** ([`athena/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/athena/mod.rs)). - -Keep **`pub fn new(...)`** (or **`async fn new`**) as the low-level constructor if you want tests to build adapters without a full **`ClusterConfig`**; **`try_from_cluster_config`** can delegate to **`new`** after parsing **`cfg`**. - -### 3. Binary wiring (`crates/queryflux`) - -Registration is centralized in **`crates/queryflux/src/registered_engines.rs`**: - -- **`all_descriptors()`** — Append **`MyEngineAdapter::descriptor()`** to the returned `vec!`. [`main.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux/src/main.rs) builds **`EngineRegistry::new(registered_engines::all_descriptors())`** for validation and **`GET /admin/engine-registry`**. -- **`build_adapter(cluster_name, placeholder_group, cluster_cfg, cluster_name_str).await`** — Returns **`anyhow::Result>`**. Add a **`match`** arm on **`EngineConfig::MyEngine`** that calls **`MyEngineAdapter::try_from_cluster_config(...)`**, maps **`QueryFluxError`** to **`anyhow::Error`** (same helper as other arms), and wraps **`Arc::new(...)`**. **Startup** uses **`.context(...)?`** on the result; **hot reload** in **`build_live_config`** logs a warning and **`continue`** on error — behavior stays in **`main.rs`**, not in the factory. - -Do **not** add a second adapter-construction **`match`** in **`main.rs`**. - -**Not implemented yet:** e.g. **`EngineConfig::ClickHouse`** is handled inside **`build_adapter`** with **`anyhow::bail!`** until a **`ClickHouseAdapter`** and **`try_from_cluster_config`** exist. - -- **`EngineType` for cluster state** — In **`main.rs`** and anywhere else (e.g. group member **`ClusterState`**), use **`EngineType::from(engine_config)`** from **`engine_registry.rs`**. **`queryflux-cluster-manager`** engine affinity uses the same **`From`** impl (see **`strategy.rs`**). - -- **Special rules** — Search for engine-specific checks (e.g. `queryAuth` / impersonation) and extend validation if your engine has constraints. - -### 4. Dispatch and frontends (`queryflux-frontend`) - -- Shared query execution goes through **`dispatch_query`** / **`execute_to_sink`**. Usually no change if the new engine only differs in the adapter; if you need a **special execution path** (like Trino raw HTTP), follow the existing engine-specific branches. -- Per-protocol handlers (Trino HTTP, Postgres wire, …) should keep using the shared dispatch layer unless the protocol requires a dedicated contract. - -### 5. Persistence (`queryflux-persistence`) — why touch it if config is JSON? - -The table stores **`engine_key` as its own column** plus a **`config` JSONB** blob. The DB does not load straight into the proxy as opaque JSON: code paths call **`ClusterConfigRecord::to_core()`**, which must produce a typed **`ClusterConfig`** (including **`EngineConfig`**). - -So persistence changes are **not** “because Postgres needs a JSON schema.” They are because of this **explicit conversion layer**: - -1. **`ClusterConfigRecord::to_core`** — Calls **`parse_engine_key`** from `queryflux-core` (next to `engine_key`). Extend **`parse_engine_key`** when you add an engine; you do **not** maintain a second duplicate string match in persistence. -2. **`UpsertClusterConfig::from_core`** — Uses **`engine_key(&EngineConfig)`** from core to set the `engine_key` column when seeding from YAML. - -**Extra JSON keys** that only live inside `config` and are **already** read in `to_core` (e.g. `endpoint`, `region`, `authType`, …) usually need **no** persistence change beyond the engine-key match. You only extend the `s("…")` / `b("…")` helpers in `to_core` (and the matching `from_core` inserts) if you add **new top-level persisted fields** on `ClusterConfig` that should round-trip through that JSON. - -**Hot reload** often uses `list_cluster_configs` → records → `to_core()` → `build_live_config`; the same conversion applies. - -### 6. Optional: routing config - -- If operators choose the new group via **router JSON** (`RouterConfig` variants), no change unless you add a new router **type**. -- **Protocol-based routing** maps frontend labels to **group names**; it does not list backend engines. - -### 7. Tests and docs - -- Add or extend **e2e** tests under `crates/queryflux-e2e-tests` if you have a dockerized target. -- Update [system-map.md](system-map.md) component status if you document supported engines there. - -### 8. Suggested order of work (backend only) - -1. **`EngineConfig` / `EngineType`** + **`engine_key` / `parse_engine_key` / `From<&EngineConfig> for EngineType`** + **`dialect()`** if needed. -2. **`ClusterConfig`** fields if the engine needs new top-level keys (and persistence **`to_core`** JSON extraction if those keys live in JSONB). -3. Adapter module: **`EngineAdapterTrait`**, **`descriptor()`**, **`try_from_cluster_config`**. -4. **`registered_engines.rs`**: descriptor in **`all_descriptors()`**, arm in **`build_adapter`**. -5. Run **`cargo build -p queryflux`**; exercise **YAML** and **Postgres** load + **admin upsert** if applicable. - ---- - -## Part B — QueryFlux Studio (UI, TypeScript / React) - -Studio is the Next.js app under `ui/queryflux-studio/`. It does **not** run wire protocols; it calls the **Admin API** (`ADMIN_API_URL`, default `http://localhost:9000`) for clusters, groups, routing, and scripts. - -Backend engines are registered in Studio through **`StudioEngineModule`** objects: one file per engine under **`lib/studio-engines/engines/`**, aggregated in **`lib/studio-engines/manifest.ts`**. That manifest drives **`ENGINE_REGISTRY`**, catalog slots for implemented backends, optional flat-form validation, engine-affinity dropdown entries, and extra **`findEngineByType`** aliases. - -The proxy still exposes descriptors at **`GET /admin/engine-registry`**. Studio does **not** load that at runtime yet, so **Rust `descriptor()` and each studio module’s `descriptor` field must stay aligned by hand** (same `engineKey`, `configFields` keys, auth shapes, etc.). Shared TypeScript types live in **`lib/engine-registry-types.ts`**; **`lib/engine-registry.ts`** only re-exports helpers and builds **`ENGINE_REGISTRY`** from the manifest. - -### Where users see engines - -| User action | UI entrypoint | What must know your engine | -|-------------|---------------|----------------------------| -| Create cluster | **Clusters → Add cluster** (`components/add-cluster-dialog.tsx`) | Expanded **`ENGINE_CATALOG`** (includes studio slots) + **`findEngineDescriptor`** + **`validateClusterConfig`** / **`validateEngineSpecific`** + **`toUpsertBody`** | -| Edit cluster | **Clusters** grid → cluster card → Edit (`app/clusters/clusters-grid.tsx`) | Same + **`mergeClusterConfigFromFlat`** / **`buildClusterUpsertFromForm`** + **`EngineClusterConfig`** | -| View config | Cluster detail / engine config view in `clusters-grid.tsx` | **`findEngineDescriptor`** for labels; unknown key shows “add to engine registry” warning | -| Group strategy **engine affinity** | **Engines →** group dialog → strategy (`components/group-form-dialog.tsx`) | **`ENGINE_AFFINITY_OPTIONS`** is built by **`buildEngineAffinityOptionsFromManifest()`** from each module’s **`engineAffinity`** field (omit label override, or set **`engineAffinity: false`** to exclude, e.g. Athena). | -| Live utilization cards | **Engines (Groups)** page (`app/engines/page.tsx`) | **`findEngineByType`**; studio modules contribute aliases via **`extraTypeAliases`** (merged with static dialect aliases in **`components/engine-catalog.ts`**) | - -### 1. Studio engine module (primary registration) - -**Types:** `ui/queryflux-studio/lib/studio-engines/types.ts` — **`StudioEngineModule`**. - -**Per engine:** `ui/queryflux-studio/lib/studio-engines/engines/.ts` - -Export a constant (e.g. **`trinoStudioEngine`**) with: - -- **`descriptor`** — Full **`EngineDescriptor`** (must match Rust: **`engineKey`**, **`connectionType`**, **`supportedAuth`**, **`configFields`**, **`implemented`**, branding **`hex`**, etc.). Extend **`ConnectionType`** / **`AuthType`** in **`lib/engine-registry-types.ts`** if Rust added a variant. -- **`catalog`** — **`category`**, **`simpleIconSlug`**, **`catalogDescription`** for the engines grid / picker (display name and **`supported`** come from the descriptor when the catalog is expanded). -- **`validateFlat`** (optional) — Cross-field checks before save (e.g. Trino basic vs bearer). Dispatched by **`validateEngineSpecific`** in **`lib/studio-engines/validate-flat.ts`** (re-exported from **`lib/cluster-persist-form.ts`**). -- **`customFormId`** (optional) — String key; must match an entry in **`components/cluster-config/studio-engine-forms.tsx`** if the generic **`GenericEngineClusterConfig`** is not enough. -- **`engineAffinity`** (optional) — **`false`** to omit from affinity, or **`{ label?: string }`** for a custom dropdown label (default label is **`displayName`**). -- **`extraTypeAliases`** (optional) — Map of normalized API/type strings → canonical **`EngineDef.name`** for **`findEngineByType`** (e.g. alternate spellings). - -**Manifest:** `ui/queryflux-studio/lib/studio-engines/manifest.ts` - -- Import the new module and append it to **`STUDIO_ENGINE_MODULES`** (order affects **`ENGINE_AFFINITY_OPTIONS`** and registry iteration; catalog **card order** is separate — see below). - -**Derived registry:** `ui/queryflux-studio/lib/engine-registry.ts` - -- **`ENGINE_REGISTRY`** is **`STUDIO_ENGINE_MODULES.map((m) => m.descriptor)`**. Do not duplicate descriptor arrays here. -- **`findEngineDescriptor`**, **`implementedEngines`**, **`isClusterOnboardingSelectable`**, **`validateClusterConfig`** — unchanged behavior; **`validateClusterConfig`** still uses generic required-field checks from **`configFields`** unless you extend the Rust/TS contract. - -### 2. Catalog layout (picker order and dialect-only rows) - -**File:** `ui/queryflux-studio/components/engine-catalog.ts` - -- Implemented backends appear as **studio slots**: **`{ k: "studio", engineKey: "" }`** inside **`ENGINE_CATALOG_SLOTS`**, interleaved with static **`EngineDef`** rows (dialects with **`engineKey: null`**). -- At runtime, **`expandCatalog`** replaces each studio slot with **`studioModuleToEngineDef`** from **`lib/studio-engines/catalog.ts`**. -- Static **`STATIC_ENGINE_TYPE_ALIASES`** remains for dialects without a studio module; **`buildStudioTypeAliases()`** merges in per-module aliases and the lowercase **`engineKey`** → **`displayName`** mapping. - -**`isClusterOnboardingSelectable`** still requires a catalog row with **`supported`** and **`engineKey`**; for studio-backed engines, **`supported`** is **`descriptor.implemented`** after expansion. - -### 3. Cluster config forms - -**Router:** `ui/queryflux-studio/components/cluster-config/engine-cluster-config.tsx` - -- Resolves **`getStudioEngineModule(engineKey)`**; if **`customFormId`** is set and **`STUDIO_CUSTOM_CLUSTER_FORMS[id]`** exists, renders that component; otherwise **`GenericEngineClusterConfig`** (descriptor **`configFields`**). - -**Custom form registration:** `ui/queryflux-studio/components/cluster-config/studio-engine-forms.tsx` — map **`customFormId`** → component (see Trino / StarRocks / Athena). - -**Reference components:** `trino-cluster-config.tsx`, `starrocks-cluster-config.tsx`, `athena-cluster-config.tsx`, `generic-engine-cluster-config.tsx`, `config-field-row.tsx`. - -### 4. Persisted JSON ↔ flat form (create + edit save path) - -**File:** `ui/queryflux-studio/lib/cluster-persist-form.ts` - -- Still shared across engines. If **`cluster_configs.config`** gains **new top-level JSON keys**, update: - - **`MANAGED_CONFIG_JSON_KEYS`** - - **`persistedClusterConfigToFlat`**, **`flatToPersistedConfig`**, **`mergeClusterConfigFromFlat`** - - **`buildValidateShape`** (shape expected by **`validateClusterConfig`**) -- **`validateEngineSpecific`** is implemented in **`lib/studio-engines/validate-flat.ts`** (per-module **`validateFlat`**); this file re-exports it for call sites. - -### 5. Clusters page (grid, dialog, validation) - -**File:** `ui/queryflux-studio/app/clusters/clusters-grid.tsx` - -- Uses **`findEngineDescriptor`**, **`validateClusterConfig`**, **`validateEngineSpecific`**, **`buildValidateShape`**, **`skipImplementedCheck`** where needed. No per-engine branches beyond **`EngineClusterConfig`**. - -**File:** `ui/queryflux-studio/components/add-cluster-dialog.tsx` - -- Wires catalog → descriptor → **`EngineClusterConfig`** → **`toUpsertBody`** → **`upsertClusterConfig`**. - -### 6. Group strategy (engine affinity) - -**File:** `ui/queryflux-studio/lib/cluster-group-strategy.ts` - -- **`ENGINE_AFFINITY_OPTIONS`** = **`buildEngineAffinityOptionsFromManifest()`**. To exclude an engine, set **`engineAffinity: false`** on its **`StudioEngineModule`**. To customize the label, use **`engineAffinity: { label: "…" }`**. - -### 7. Display helpers - -**File:** `ui/queryflux-studio/lib/merge-clusters-display.ts` - -- Uses **`findEngineDescriptor(p.engineKey)`**; the descriptor must exist in the manifest. - -**File:** `ui/queryflux-studio/components/ui-helpers.tsx` (**`EngineBadge`**) - -- Uses **`ENGINE_CATALOG`**; studio-expanded rows must match **`displayName`** where badges key off names. - -**File:** `ui/queryflux-studio/components/engine-icon.tsx` - -- Consumes **`EngineDef`** (re-exported from **`engine-catalog.ts`**; types in **`lib/engine-catalog-types.ts`**). - -### 8. API types (usually unchanged) - -**File:** `ui/queryflux-studio/lib/api-types.ts` - -- **`ClusterConfigRecord`** / **`UpsertClusterConfig`** stay generic unless you add typed helpers. - -### 9. Optional: fetch registry from the proxy - -A follow-up could load **`GET /admin/engine-registry`** at runtime and hydrate forms from the API. Until then, keep **Rust `descriptor()`** and **`StudioEngineModule.descriptor`** in sync manually. - -### Studio checklist (copy-paste) - -- [ ] **`lib/studio-engines/engines/.ts`** — **`StudioEngineModule`** (`descriptor` aligned with Rust, **`catalog`**, optional **`validateFlat`**, **`customFormId`**, **`engineAffinity`**, **`extraTypeAliases`**) -- [ ] **`lib/studio-engines/manifest.ts`** — import + append to **`STUDIO_ENGINE_MODULES`** -- [ ] **`lib/engine-registry-types.ts`** — extend **`ConnectionType`** / **`AuthType`** if needed -- [ ] **`components/engine-catalog.ts`** — add **`{ k: "studio", engineKey: "…" }`** to **`ENGINE_CATALOG_SLOTS`** at the desired position -- [ ] **`components/cluster-config/studio-engine-forms.tsx`** — register component if **`customFormId`** is set -- [ ] **`lib/cluster-persist-form.ts`** — only if new persisted **`config`** JSON keys (managed keys + flat ↔ JSON + **`buildValidateShape`**) -- [ ] Smoke-test: Add cluster → save → edit → save; Engines page icons; group **engine affinity** if applicable - ---- - -## Part C — Frontend protocol (e.g. “more Postgres wire”) - -Goal: clients speak a **wire protocol to QueryFlux**, not a new backend. - -### Where the code lives - -- **PostgreSQL wire:** `crates/queryflux-frontend/src/postgres_wire/` -- **MySQL wire:** `crates/queryflux-frontend/src/mysql_wire/` -- **Trino HTTP:** `crates/queryflux-frontend/src/trino_http/` -- **Flight SQL:** `crates/queryflux-frontend/src/flight_sql/` - -### Typical steps - -1. **`FrontendProtocol`** — Already defined in `queryflux_core::query::FrontendProtocol`; add a variant only for a **new** ingress protocol. -2. **`default_dialect()`** — Set the sqlglot **source** dialect for translation (see [query-translation.md](query-translation.md)). -3. **Listener** — Bind a port, parse the protocol, build **`SessionContext`** and **`InboundQuery`**, then call shared **`dispatch_query`** (or the same helpers Trino HTTP uses). -4. **Routing** — Optionally extend **protocol-based routing** in config / persisted routing so this frontend maps to the right default group. -5. **Tests** — Protocol-level tests or e2e clients as appropriate. - -Studio does **not** implement wire protocols; it only talks to the **Admin API** for config and metrics. - ---- - -## Checklist summary - -**Backend engine** - -- [ ] `EngineConfig` + `EngineType` + `engine_key()` + **`parse_engine_key()`** + **`From<&EngineConfig> for EngineType`** + dialect mapping (`engine_registry.rs` + `query.rs`) -- [ ] `EngineAdapterTrait` + `descriptor()` -- [ ] `registered_engines.rs`: **`all_descriptors()`** + **`build_adapter()`** arm calling **`try_from_cluster_config`** on the adapter -- [ ] Adapter module: **`try_from_cluster_config`** (or async equivalent) reading **`ClusterConfig`** -- [ ] `UpsertClusterConfig::from_core` / `to_core` stay aligned via **`engine_key` / `parse_engine_key`** (no extra string match in persistence) -- [ ] Translation / compatibility if dialect is new - -**Studio (UI)** - -- [ ] `lib/studio-engines/engines/.ts` — `StudioEngineModule` (descriptor + catalog + options) -- [ ] `lib/studio-engines/manifest.ts` — register module in `STUDIO_ENGINE_MODULES` -- [ ] `lib/engine-registry-types.ts` — `ConnectionType` / `AuthType` if Rust added variants -- [ ] `components/engine-catalog.ts` — `{ k: "studio", engineKey }` slot in `ENGINE_CATALOG_SLOTS` -- [ ] `components/cluster-config/studio-engine-forms.tsx` — only if using `customFormId` -- [ ] `lib/cluster-persist-form.ts` — only if new `config` JSON keys need round-tripping -- [ ] Verify add-cluster + edit-cluster, Engines page icons / `findEngineByType`, and engine affinity if used - -**New client protocol** - -- [ ] `FrontendProtocol` + dialect + listener module + dispatch integration + routing docs - ---- - -## Related reading - -- [system-map.md](system-map.md) — End-to-end flow -- [query-translation.md](query-translation.md) — Dialects and sqlglot -- [routing-and-clusters.md](routing-and-clusters.md) — Routers and groups -- [observability.md](observability.md) — Admin API (including engine registry JSON) - -**Rust files referenced above** - -- [`crates/queryflux/src/registered_engines.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux/src/registered_engines.rs) — `all_descriptors`, `build_adapter` -- [`crates/queryflux-core/src/engine_registry.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-core/src/engine_registry.rs) — `engine_key`, `parse_engine_key`, `EngineRegistry`, `From<&EngineConfig> for EngineType` -- [`crates/queryflux-persistence/src/cluster_config.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-persistence/src/cluster_config.rs) — `to_core` / `from_core` vs `engine_key` + JSONB +Existing frontends are documented in **[Frontends](frontends/overview)**. diff --git a/docs/adding-engine-support.md b/website/docs/architecture/adding-support/backend.md similarity index 81% rename from docs/adding-engine-support.md rename to website/docs/architecture/adding-support/backend.md index 4d74e14..e11c42c 100644 --- a/docs/adding-engine-support.md +++ b/website/docs/architecture/adding-support/backend.md @@ -1,28 +1,19 @@ -# Adding engine and protocol support - -This guide separates two ideas that are easy to conflate: - -| Concept | Meaning | Example | -|--------|---------|---------| -| **Backend engine** | A **cluster** type QueryFlux routes queries **to**. It has an adapter that talks to the real database (HTTP, MySQL wire, embedded library, AWS SDK, …). | Trino, DuckDB, StarRocks, Athena | -| **Frontend protocol** | How **clients connect to QueryFlux** (ingress). SQL enters with a `FrontendProtocol` and a default source dialect for translation. | Trino HTTP, **PostgreSQL wire**, MySQL wire, Flight SQL | - -Adding **PostgreSQL wire** as a client entrypoint is **not** the same as adding “PostgreSQL” as a backend: today, `PostgresWire` is already a frontend in `queryflux-frontend`; traffic still lands on the shared dispatch path and is sent to whatever **backend adapter** routing chose (often Trino). - -Use the sections below depending on whether you are extending **Studio**, a **backend adapter**, or a **frontend listener**. - --- +description: Adding a backend engine to QueryFlux — Rust adapter, registration, persistence, and QueryFlux Studio UI. +--- + +# Backend -## Part A — Backend engine (Rust) +Goal: a new `engine` value in cluster config, a live adapter, validation, translation target dialect, wiring in the binary, and (optionally) Studio UI so operators can manage clusters for that engine. -Goal: a new `engine` value in cluster config, a live adapter, validation, translation target dialect, and wiring in the binary. +See **[Extending QueryFlux](overview.md)** for how this differs from a **frontend** protocol. For existing client protocols, see **[Frontends](../frontends/overview.md)**. ### Registration overview Backends are **not** loaded dynamically. Each engine is compiled in and registered explicitly. Data flow: 1. **Postgres / YAML** → `engine_key` column + `config` JSONB → `ClusterConfigRecord::to_core()` uses **`parse_engine_key`** and JSON helpers → typed **`ClusterConfig`**. -2. **Binary** → `registered_engines::build_adapter(...)` matches **`EngineConfig`** and calls the adapter’s **`try_from_cluster_config`** (see [`crates/queryflux/src/registered_engines.rs`](../crates/queryflux/src/registered_engines.rs)). +2. **Binary** → `registered_engines::build_adapter(...)` matches **`EngineConfig`** and calls the adapter’s **`try_from_cluster_config`** (see [`crates/queryflux/src/registered_engines.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux/src/registered_engines.rs)). 3. **Adapter** → reads only the **`ClusterConfig`** fields it needs (endpoint, auth, region, …) and constructs itself; **startup** and **hot reload** both use the same factory. **JSONB** stores per-cluster, per-engine payload without schema migrations; **`ClusterConfig`** in core is the typed view after `to_core()`. Engine-specific wiring belongs in **`try_from_cluster_config`**, not in `main.rs`. @@ -35,7 +26,7 @@ Backends are **not** loaded dynamically. Each engine is compiled in and register - **`engine_key(&EngineConfig)`** — `EngineConfig` → stable string key (must match the adapter descriptor and Studio). - **`parse_engine_key(&str)`** — inverse mapping for the `engine_key` column in Postgres / API. - **`impl From<&EngineConfig> for EngineType`** — single place for config → runtime `EngineType` (cluster manager and `main.rs` use this instead of ad-hoc matches). -- **`EngineType::dialect()`** — Return the `SqlDialect` used as the **translation target** (and extend `SqlDialect` / `is_compatible_with` in translation if needed). See [query-translation.md](query-translation.md). +- **`EngineType::dialect()`** — Return the `SqlDialect` used as the **translation target** (and extend `SqlDialect` / `is_compatible_with` in translation if needed). See [query-translation.md](../query-translation.md). - **`ClusterConfig` fields** — Add any new top-level fields (region, paths, engine-specific blobs). Prefer keeping engine-specific secrets and options in `config` JSON for Postgres-backed clusters; extend the typed struct when YAML and validation need them everywhere. ### 2. Adapter crate (`queryflux-engine-adapters`) @@ -65,7 +56,7 @@ Implement on your adapter struct so all **field extraction and validation** for - **Async** (e.g. Athena — AWS client setup): same parameters, `async fn`, returns `Result`. -Use **`QueryFluxError::Engine(format!(…))`** for failures; include **`cluster_name_str`** in messages so startup and reload logs identify the cluster. Reference implementations: **`TrinoAdapter`** and **`StarRocksAdapter`** ([`trino/mod.rs`](../crates/queryflux-engine-adapters/src/trino/mod.rs), [`starrocks/mod.rs`](../crates/queryflux-engine-adapters/src/starrocks/mod.rs)), **`DuckDbAdapter`** / **`DuckDbHttpAdapter`** ([`duckdb/mod.rs`](../crates/queryflux-engine-adapters/src/duckdb/mod.rs), [`duckdb/http.rs`](../crates/queryflux-engine-adapters/src/duckdb/http.rs)), **`AthenaAdapter`** ([`athena/mod.rs`](../crates/queryflux-engine-adapters/src/athena/mod.rs)). +Use **`QueryFluxError::Engine(format!(…))`** for failures; include **`cluster_name_str`** in messages so startup and reload logs identify the cluster. Reference implementations: **`TrinoAdapter`** and **`StarRocksAdapter`** ([`trino/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/trino/mod.rs), [`starrocks/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/starrocks/mod.rs)), **`DuckDbAdapter`** / **`DuckDbHttpAdapter`** ([`duckdb/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/duckdb/mod.rs), [`duckdb/http.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/duckdb/http.rs)), **`AthenaAdapter`** ([`athena/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/athena/mod.rs)). Keep **`pub fn new(...)`** (or **`async fn new`**) as the low-level constructor if you want tests to build adapters without a full **`ClusterConfig`**; **`try_from_cluster_config`** can delegate to **`new`** after parsing **`cfg`**. @@ -73,7 +64,7 @@ Keep **`pub fn new(...)`** (or **`async fn new`**) as the low-level constructor Registration is centralized in **`crates/queryflux/src/registered_engines.rs`**: -- **`all_descriptors()`** — Append **`MyEngineAdapter::descriptor()`** to the returned `vec!`. [`main.rs`](../crates/queryflux/src/main.rs) builds **`EngineRegistry::new(registered_engines::all_descriptors())`** for validation and **`GET /admin/engine-registry`**. +- **`all_descriptors()`** — Append **`MyEngineAdapter::descriptor()`** to the returned `vec!`. [`main.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux/src/main.rs) builds **`EngineRegistry::new(registered_engines::all_descriptors())`** for validation and **`GET /admin/engine-registry`**. - **`build_adapter(cluster_name, placeholder_group, cluster_cfg, cluster_name_str).await`** — Returns **`anyhow::Result>`**. Add a **`match`** arm on **`EngineConfig::MyEngine`** that calls **`MyEngineAdapter::try_from_cluster_config(...)`**, maps **`QueryFluxError`** to **`anyhow::Error`** (same helper as other arms), and wraps **`Arc::new(...)`**. **Startup** uses **`.context(...)?`** on the result; **hot reload** in **`build_live_config`** logs a warning and **`continue`** on error — behavior stays in **`main.rs`**, not in the factory. Do **not** add a second adapter-construction **`match`** in **`main.rs`**. @@ -110,9 +101,9 @@ So persistence changes are **not** “because Postgres needs a JSON schema.” T ### 7. Tests and docs - Add or extend **e2e** tests under `crates/queryflux-e2e-tests` if you have a dockerized target. -- Update [architecture.md](architecture.md) component status if you document supported engines there. +- Update [system-map.md](../system-map.md) component status if you document supported engines there. -### 8. Suggested order of work (backend only) +### 8. Suggested order of work (Rust) 1. **`EngineConfig` / `EngineType`** + **`engine_key` / `parse_engine_key` / `From<&EngineConfig> for EngineType`** + **`dialect()`** if needed. 2. **`ClusterConfig`** fields if the engine needs new top-level keys (and persistence **`to_core`** JSON extraction if those keys live in JSONB). @@ -122,7 +113,7 @@ So persistence changes are **not** “because Postgres needs a JSON schema.” T --- -## Part B — QueryFlux Studio (UI, TypeScript / React) +## QueryFlux Studio (UI, TypeScript / React) Studio is the Next.js app under `ui/queryflux-studio/`. It does **not** run wire protocols; it calls the **Admin API** (`ADMIN_API_URL`, default `http://localhost:9000`) for clusters, groups, routing, and scripts. @@ -246,32 +237,9 @@ A follow-up could load **`GET /admin/engine-registry`** at runtime and hydrate f --- -## Part C — Frontend protocol (e.g. “more Postgres wire”) - -Goal: clients speak a **wire protocol to QueryFlux**, not a new backend. - -### Where the code lives - -- **PostgreSQL wire:** `crates/queryflux-frontend/src/postgres_wire/` -- **MySQL wire:** `crates/queryflux-frontend/src/mysql_wire/` -- **Trino HTTP:** `crates/queryflux-frontend/src/trino_http/` -- **Flight SQL:** `crates/queryflux-frontend/src/flight_sql/` - -### Typical steps - -1. **`FrontendProtocol`** — Already defined in `queryflux_core::query::FrontendProtocol`; add a variant only for a **new** ingress protocol. -2. **`default_dialect()`** — Set the sqlglot **source** dialect for translation (see [query-translation.md](query-translation.md)). -3. **Listener** — Bind a port, parse the protocol, build **`SessionContext`** and **`InboundQuery`**, then call shared **`dispatch_query`** (or the same helpers Trino HTTP uses). -4. **Routing** — Optionally extend **protocol-based routing** in config / persisted routing so this frontend maps to the right default group. -5. **Tests** — Protocol-level tests or e2e clients as appropriate. - -Studio does **not** implement wire protocols; it only talks to the **Admin API** for config and metrics. - ---- - -## Checklist summary +## Checklist (backend) -**Backend engine** +### Rust adapter - [ ] `EngineConfig` + `EngineType` + `engine_key()` + **`parse_engine_key()`** + **`From<&EngineConfig> for EngineType`** + dialect mapping (`engine_registry.rs` + `query.rs`) - [ ] `EngineAdapterTrait` + `descriptor()` @@ -280,7 +248,7 @@ Studio does **not** implement wire protocols; it only talks to the **Admin API** - [ ] `UpsertClusterConfig::from_core` / `to_core` stay aligned via **`engine_key` / `parse_engine_key`** (no extra string match in persistence) - [ ] Translation / compatibility if dialect is new -**Studio (UI)** +### Studio (UI) - [ ] `lib/studio-engines/engines/.ts` — `StudioEngineModule` (descriptor + catalog + options) - [ ] `lib/studio-engines/manifest.ts` — register module in `STUDIO_ENGINE_MODULES` @@ -290,21 +258,18 @@ Studio does **not** implement wire protocols; it only talks to the **Admin API** - [ ] `lib/cluster-persist-form.ts` — only if new `config` JSON keys need round-tripping - [ ] Verify add-cluster + edit-cluster, Engines page icons / `findEngineByType`, and engine affinity if used -**New client protocol** - -- [ ] `FrontendProtocol` + dialect + listener module + dispatch integration + routing docs - --- ## Related reading -- [architecture.md](architecture.md) — End-to-end flow -- [query-translation.md](query-translation.md) — Dialects and sqlglot -- [routing-and-clusters.md](routing-and-clusters.md) — Routers and groups -- [observability.md](observability.md) — Admin API (including engine registry JSON) +- [Extending QueryFlux — overview](overview.md) +- [Frontend](frontend.md) — new ingress protocols +- [query-translation.md](../query-translation.md) — Dialects and sqlglot +- [routing-and-clusters.md](../routing-and-clusters.md) — Routers and groups +- [observability.md](../observability.md) — Admin API -**Rust files referenced above** +**Rust files** -- [`crates/queryflux/src/registered_engines.rs`](../crates/queryflux/src/registered_engines.rs) — `all_descriptors`, `build_adapter` -- [`crates/queryflux-core/src/engine_registry.rs`](../crates/queryflux-core/src/engine_registry.rs) — `engine_key`, `parse_engine_key`, `EngineRegistry`, `From<&EngineConfig> for EngineType` -- [`crates/queryflux-persistence/src/cluster_config.rs`](../crates/queryflux-persistence/src/cluster_config.rs) — `to_core` / `from_core` vs `engine_key` + JSONB +- [`crates/queryflux/src/registered_engines.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux/src/registered_engines.rs) — `all_descriptors`, `build_adapter` +- [`crates/queryflux-core/src/engine_registry.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-core/src/engine_registry.rs) — `engine_key`, `parse_engine_key`, `EngineRegistry`, `From<&EngineConfig> for EngineType` +- [`crates/queryflux-persistence/src/cluster_config.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-persistence/src/cluster_config.rs) — `to_core` / `from_core` vs `engine_key` + JSONB diff --git a/website/docs/architecture/adding-support/frontend.md b/website/docs/architecture/adding-support/frontend.md new file mode 100644 index 0000000..3b71dc7 --- /dev/null +++ b/website/docs/architecture/adding-support/frontend.md @@ -0,0 +1,43 @@ +--- +description: Adding a new frontend client protocol to QueryFlux — FrontendProtocol, listener, dispatch, and routing. +--- + +# Frontend + +Goal: clients speak a **wire or HTTP protocol to QueryFlux** (ingress), not a new backend. + +See **[Extending QueryFlux](overview.md)** for how this differs from a **backend** engine. Documentation for **existing** frontends (Trino HTTP, Postgres wire, MySQL wire, Flight SQL, Snowflake) lives under **[Frontends](../frontends/overview.md)**. + +### Where the code lives + +- **PostgreSQL wire:** `crates/queryflux-frontend/src/postgres_wire/` +- **MySQL wire:** `crates/queryflux-frontend/src/mysql_wire/` +- **Trino HTTP:** `crates/queryflux-frontend/src/trino_http/` +- **Flight SQL:** `crates/queryflux-frontend/src/flight_sql/` +- **Snowflake:** `crates/queryflux-frontend/src/snowflake/` + +### Typical steps for a new protocol + +1. **`FrontendProtocol`** — Already defined in `queryflux_core::query::FrontendProtocol`; add a variant only for a **new** ingress protocol. +2. **`default_dialect()`** — Set the sqlglot **source** dialect for translation (see [query-translation.md](../query-translation.md)). +3. **Listener** — Bind a port, parse the protocol, build **`SessionContext`** and **`InboundQuery`**, then call shared **`dispatch_query`** (or the same helpers Trino HTTP uses). +4. **Routing** — Optionally extend **protocol-based routing** in config / persisted routing so this frontend maps to the right default group. +5. **Tests** — Protocol-level tests or e2e clients as appropriate. + +Studio does **not** implement wire protocols; it only talks to the **Admin API** for config and metrics. + +--- + +## Checklist (frontend) + +- [ ] `FrontendProtocol` + dialect + listener module + dispatch integration + routing docs + +--- + +## Related reading + +- [Extending QueryFlux — overview](overview.md) +- [Backend](backend.md) — Rust adapter and Studio +- [Frontends](../frontends/overview.md) — Existing protocol listeners +- [query-translation.md](../query-translation.md) — Dialects and sqlglot +- [routing-and-clusters.md](../routing-and-clusters.md) — Routers and groups diff --git a/website/docs/architecture/adding-support/overview.md b/website/docs/architecture/adding-support/overview.md new file mode 100644 index 0000000..2a3f6d5 --- /dev/null +++ b/website/docs/architecture/adding-support/overview.md @@ -0,0 +1,32 @@ +--- +sidebar_label: Overview +description: How to extend QueryFlux — new backend engines (Rust + Studio) vs new frontend client protocols, with links to detailed guides. +--- + +# Extending QueryFlux + +This guide separates two ideas that are easy to conflate: + +| Concept | Meaning | Example | +|--------|---------|---------| +| **Backend engine** | A **cluster** type QueryFlux routes queries **to**. It has an adapter that talks to the real database (HTTP, MySQL wire, embedded library, AWS SDK, …). | Trino, DuckDB, StarRocks, Athena | +| **Frontend protocol** | How **clients connect to QueryFlux** (ingress). SQL enters with a `FrontendProtocol` and a default source dialect for translation. | Trino HTTP, PostgreSQL wire, MySQL wire, Flight SQL | + +Adding **PostgreSQL wire** as a client entrypoint is **not** the same as adding “PostgreSQL” as a backend: today, `PostgresWire` is already a frontend in `queryflux-frontend`; traffic still lands on the shared dispatch path and is sent to whatever **backend adapter** routing chose (often Trino). + +## Guides + +| Page | What it covers | +|------|----------------| +| [Backend](backend.md) | Rust adapter (`EngineAdapterTrait`), `registered_engines`, persistence, dispatch notes — plus **QueryFlux Studio** (`StudioEngineModule`, catalog, forms). | +| [Frontend](frontend.md) | Adding a **new** ingress protocol (`FrontendProtocol`, listener, routing). For existing protocols, see **[Frontends](../frontends/overview.md)**. | + +--- + +## Related reading + +- [Frontends](../frontends/overview.md) — Trino HTTP, Postgres wire, MySQL wire, Flight SQL, Snowflake +- [system-map.md](../system-map.md) — End-to-end flow +- [query-translation.md](../query-translation.md) — Dialects and sqlglot +- [routing-and-clusters.md](../routing-and-clusters.md) — Routers and groups +- [observability.md](../observability.md) — Admin API (including engine registry JSON) diff --git a/website/docs/architecture/frontends/flight-sql.md b/website/docs/architecture/frontends/flight-sql.md new file mode 100644 index 0000000..818b43a --- /dev/null +++ b/website/docs/architecture/frontends/flight-sql.md @@ -0,0 +1,71 @@ +--- +description: Arrow Flight SQL frontend — gRPC service, implemented RPCs, metadata passthrough, and connecting with Flight SQL clients. +--- + +# Arrow Flight SQL + +The Flight SQL frontend exposes a gRPC-based [Arrow Flight SQL](https://arrow.apache.org/docs/format/FlightSql.html) service. Clients that speak Flight SQL (e.g. the JDBC Flight SQL driver, ADBC, DuckDB's `ATTACH` over Flight) can connect to QueryFlux and run SQL queries with results streamed as native Arrow record batches — zero serialization overhead for columnar consumers. + +## Configuration + +```yaml +queryflux: + frontends: + flightSql: + enabled: true + port: 50051 +``` + +Config key: `flightSql`. Protocol identifier: `FrontendProtocol::FlightSql`. Default dialect: `SqlDialect::Generic`. + +## Implemented RPCs + +| RPC | Status | Description | +|-----|--------|-------------| +| `GetFlightInfo` (statement) | Implemented | Accepts SQL via `CommandStatementQuery`, returns a `FlightInfo` with a ticket | +| `DoGet` (statement) | Implemented | Decodes SQL from the ticket, executes the query, and streams Arrow `RecordBatch` results | +| All other Flight SQL RPCs | Unimplemented | Return gRPC `Unimplemented` status | + +The two-step flow: + +1. **`GetFlightInfo`** — client sends SQL in a `CommandStatementQuery`. QueryFlux returns a `FlightInfo` containing an endpoint with a ticket (the SQL is encoded in the ticket's `statement_handle`). The IPC schema in the response is empty at this stage. +2. **`DoGet`** — client presents the ticket. QueryFlux decodes the SQL, authenticates, routes, executes via `execute_to_sink`, and streams Arrow record batches through a `FlightDataEncoder`. + +## Authentication + +Credentials are read from gRPC metadata: + +- **`authorization`** — `Bearer ` or `Basic `, depending on your configured auth provider. + +## Execution model + +Queries execute **synchronously** via `execute_to_sink`. Results are streamed as Arrow Flight data frames directly from the query's Arrow record batches — no intermediate serialization to JSON or text. + +## Metadata and routing + +ASCII metadata entries on the gRPC request (beyond `authorization`) are collected into a key/value map and passed through to the same routing and session path used for HTTP-style frontends, so routers and auth can inspect client-supplied fields consistently across protocols. + +Query tags are **not** populated from Flight SQL metadata in the current implementation. + +## Client examples + +```python +# Python (ADBC Flight SQL driver) +import adbc_driver_flightsql.dbapi as flight_sql + +conn = flight_sql.connect(uri="grpc://localhost:50051") +cur = conn.cursor() +cur.execute("SELECT 42 AS answer") +print(cur.fetchone()) +``` + +```bash +# DuckDB (if Flight SQL extension is available) +ATTACH 'grpc://localhost:50051' AS qf (TYPE flight_sql); +SELECT * FROM qf.my_catalog.my_schema.my_table LIMIT 10; +``` + +## Related + +- [Frontends overview](overview.md) — shared dispatch and session model +- [Routing and clusters](/docs/architecture/routing-and-clusters) — `protocolBased` router with `flightSql` diff --git a/website/docs/architecture/frontends/mysql-wire.md b/website/docs/architecture/frontends/mysql-wire.md new file mode 100644 index 0000000..566f3f1 --- /dev/null +++ b/website/docs/architecture/frontends/mysql-wire.md @@ -0,0 +1,104 @@ +--- +description: MySQL wire protocol frontend — handshake, COM_QUERY, session variables, schema switching, and connecting with mysql clients. +--- + +# MySQL wire + +The MySQL wire frontend lets standard MySQL clients (`mysql` CLI, JDBC, Python `mysql-connector`, Go `go-sql-driver/mysql`, etc.) connect to QueryFlux over the MySQL wire protocol. This is the natural entry point when routing traffic to MySQL-protocol backends like StarRocks. + +## Configuration + +```yaml +queryflux: + frontends: + mysqlWire: + enabled: true + port: 3306 +``` + +Config key: `mysqlWire`. Protocol identifier: `FrontendProtocol::MySqlWire`. Default dialect: `SqlDialect::MySql`. + +## Protocol support + +| Feature | Status | +|---------|--------| +| Server handshake (protocol 10) | Supported — advertises `8.0.0-queryflux` | +| `mysql_native_password` auth | Supported | +| `COM_QUERY` | Supported | +| `COM_PING` | Supported | +| `COM_INIT_DB` (schema switch) | Supported | +| `COM_QUIT` | Supported | +| `COM_FIELD_LIST` | Stub (returns empty EOF) | +| SSL/TLS handshake | Detected and rejected — use `--ssl-mode=DISABLED` | +| Prepared statements (`COM_STMT_*`) | Not supported | + +Clients must connect without TLS (`--ssl-mode=DISABLED` for `mysql` CLI, `useSSL=false` for JDBC). If the client sends an SSL request, QueryFlux returns an error and closes the connection. + +## Handshake and authentication + +1. QueryFlux sends a server handshake packet (protocol version 10, server version `8.0.0-queryflux`, `mysql_native_password`). +2. Client responds with username and optional database. +3. QueryFlux authenticates via the configured `auth_provider`. +4. On success, an OK packet is returned and the connection is ready. + +## Execution model + +Queries execute **synchronously** via `execute_to_sink`. Results stream as MySQL text protocol result sets: + +1. Column count packet. +2. Column definition packets. +3. Row data packets (text values). +4. EOF / OK packet. + +## Built-in query handling + +The MySQL frontend intercepts several common queries that clients and drivers send during connection setup: + +| Query pattern | Behavior | +|---------------|----------| +| `SET query_tags = '...'` / `SET SESSION query_tags = '...'` | Stores tags on the session for routing | +| `SET ...` (other) | Acknowledged without forwarding to backend | +| `USE ` | Updates session schema | +| `SELECT @@version` | Returns `8.0.0-queryflux` | +| `SELECT DATABASE()` | Returns current session schema | +| `SHOW VARIABLES` / `SHOW STATUS` | Returns empty result set | + +Leading `/* ... */` and `/*!...*/` conditional comments in SQL are stripped before processing. + +## Session context + +`SessionContext::MySqlWire` carries: + +| Field | Source | +|-------|--------| +| `user` | Handshake response | +| `schema` | `COM_INIT_DB` or initial database from handshake | +| `session_vars` | Currently empty (generic `SET` is not stored) | +| `tags` | From `SET query_tags` / `SET SESSION query_tags` | + +The `schema` and `user` fields are available to routers — the `pythonScript` router receives them in `ctx["schema"]` and `ctx["user"]`. + +## Client examples + +```bash +# mysql CLI +mysql -h 127.0.0.1 -P 3306 -u dev --ssl-mode=DISABLED -e "SELECT 1" + +# With database +mysql -h 127.0.0.1 -P 3306 -u dev -D my_catalog --ssl-mode=DISABLED +``` + +```python +# Python (mysql-connector) +import mysql.connector +conn = mysql.connector.connect(host="127.0.0.1", port=3306, user="dev", ssl_disabled=True) +cur = conn.cursor() +cur.execute("SELECT 42 AS answer") +print(cur.fetchone()) +``` + +## Related + +- [Frontends overview](overview.md) — shared dispatch and session model +- [Routing and clusters](/docs/architecture/routing-and-clusters) — `protocolBased` router with `mysqlWire` +- [Query tags](/docs/architecture/query-tags) — setting tags via `SET query_tags` diff --git a/website/docs/architecture/frontends/overview.md b/website/docs/architecture/frontends/overview.md new file mode 100644 index 0000000..a18c2e5 --- /dev/null +++ b/website/docs/architecture/frontends/overview.md @@ -0,0 +1,115 @@ +--- +sidebar_label: Overview +description: How QueryFlux frontends work — protocol listeners, shared dispatch, session context, and the available client protocols. +--- + +# Frontends + +A **frontend** is the entry point for client traffic into QueryFlux. Each frontend speaks a specific wire or HTTP protocol, parses incoming SQL, builds a `SessionContext`, and hands the query to the shared dispatch layer. The client never knows which backend engine actually runs the query — the frontend translates results back into its native format before responding. + +## Available frontends + +| Frontend | Config key | Default port | Protocol | Dialect | Status | +|----------|------------|-------------|----------|---------|--------| +| [Trino HTTP](trino-http.md) | `trinoHttp` | 8080 | HTTP REST (JSON) | Trino | **Done** | +| [PostgreSQL wire](postgres-wire.md) | `postgresWire` | 5432 | PostgreSQL v3 wire | Postgres | **Done** | +| [MySQL wire](mysql-wire.md) | `mysqlWire` | 3306 | MySQL wire | MySQL | **Done** | +| [Arrow Flight SQL](flight-sql.md) | `flightSql` | 50051 | gRPC (Arrow Flight) | Generic | **Done** | +| [Snowflake](snowflake.md) | `snowflakeHttp` | 8445 | HTTP REST (JSON) | Generic | **Done** | +| ClickHouse HTTP | `clickhouseHttp` | 8123 | HTTP | ClickHouse | Planned | + +## Shared architecture + +All frontends converge on the same internal pipeline. The differences are only in how SQL enters and how results leave. + +``` +Client ──(native protocol)──► Frontend Listener + │ + SessionContext + SQL + │ + ▼ + Router Chain ──► ClusterGroupName + │ + ▼ + ClusterGroupManager ──► ClusterName + │ + ▼ + Translation Service ──► translated SQL + │ + ▼ + Engine Adapter ──► results + │ + ▼ + ResultSink ──► native protocol response +``` + +### Dispatch paths + +The dispatch layer offers two execution models: + +| Path | Used by | Behavior | +|------|---------|----------| +| **`dispatch_query`** | Trino HTTP (async-capable groups) | Submit to engine, persist handle, return polling URL. Client follows `nextUri` to stream pages. | +| **`execute_to_sink`** | All other frontends + Trino HTTP sync fallback | Wait for cluster capacity (backoff), execute query to completion, stream Arrow batches through a `ResultSink` that encodes the native protocol response. | + +### SessionContext + +Each frontend builds a protocol-specific `SessionContext` that travels with the query through routing and into dispatch: + +| Variant | Fields | +|---------|--------| +| `TrinoHttp` | `headers` (lowercased HTTP headers) | +| `PostgresWire` | `user`, `database`, `session_params` | +| `MySqlWire` | `user`, `schema`, `session_vars` | +| `ClickHouseHttp` | `headers`, `query_params` | +| `FlightSql` | `headers` (from gRPC metadata) | + +Snowflake frontends build a `MySqlWire` session context internally (user from session, schema from database hint). + +### Authentication + +All frontends extract credentials from their protocol's native mechanism (HTTP headers, wire auth packets, gRPC metadata) and pass them to the shared `auth_provider`. The auth result determines whether the query proceeds and provides context for authorization checks downstream. + +### Routing + +The `FrontendProtocol` enum identifies which frontend originated a query. The `protocolBased` router uses this to map traffic from different protocols to different cluster groups: + +```yaml +routers: + - type: protocolBased + trinoHttp: trino-default + postgresWire: trino-default + mysqlWire: starrocks-group + snowflakeHttp: trino-default + snowflakeSqlApi: analytics-group +``` + +### SQL dialect and translation + +Each frontend has a **default source dialect** (`protocol.default_dialect()`). When the target engine's dialect differs, sqlglot translates the SQL automatically. When dialects match, translation is skipped entirely. + +## Configuration + +Each frontend is enabled under `queryflux.frontends` in `config.yaml`: + +```yaml +queryflux: + frontends: + trinoHttp: + enabled: true + port: 8080 + postgresWire: + enabled: true + port: 5432 + mysqlWire: + enabled: true + port: 3306 + flightSql: + enabled: true + port: 50051 + snowflakeHttp: + enabled: true + port: 8445 +``` + +Omitting a frontend block or setting `enabled: false` disables that listener entirely. diff --git a/website/docs/architecture/frontends/postgres-wire.md b/website/docs/architecture/frontends/postgres-wire.md new file mode 100644 index 0000000..4ee4611 --- /dev/null +++ b/website/docs/architecture/frontends/postgres-wire.md @@ -0,0 +1,84 @@ +--- +description: PostgreSQL wire protocol frontend — startup, simple query, session context, and connecting with psql or any Postgres driver. +--- + +# PostgreSQL wire + +The PostgreSQL wire frontend lets standard Postgres clients (`psql`, JDBC, Python `psycopg2`, Go `pgx`, etc.) connect to QueryFlux over the PostgreSQL v3 wire protocol. Queries are executed synchronously — the TCP connection stays open while the result streams back as standard Postgres wire messages. + +## Configuration + +```yaml +queryflux: + frontends: + postgresWire: + enabled: true + port: 5432 +``` + +Config key: `postgresWire`. Protocol identifier: `FrontendProtocol::PostgresWire`. Default dialect: `SqlDialect::Postgres`. + +## Protocol support + +| Feature | Status | +|---------|--------| +| Startup (protocol version 3.0) | Supported | +| Simple query (`Q` message) | Supported | +| Extended query (Parse/Bind/Execute) | Not supported — returns error | +| SSL negotiation | Declined (`N` response) — plaintext only | +| `COPY` protocol | Not supported | + +QueryFlux responds to the SSL request with `N` (no SSL), then processes the normal startup sequence. Clients that require TLS will need to connect without it or use an external TLS terminator. + +## Startup and authentication + +1. Client sends **StartupMessage** with `user`, `database`, and optional parameters. +2. QueryFlux extracts credentials from the startup `user` field and authenticates via the configured `auth_provider`. +3. On success, QueryFlux sends `AuthenticationOk`, followed by `ParameterStatus` messages (server version `16.0-queryflux`, encoding `UTF8`, etc.), `BackendKeyData`, and `ReadyForQuery`. + +## Execution model + +All queries execute **synchronously** via `execute_to_sink`. Results stream as standard Postgres wire messages: + +1. `RowDescription` — column metadata (names, types, format codes). +2. `DataRow` messages — one per row, text-format values. +3. `CommandComplete` — summary (e.g. `SELECT 3`). +4. `ReadyForQuery` — ready for the next query. + +Errors are returned as Postgres `ErrorResponse` messages with SQLSTATE codes. + +## Session context + +`SessionContext::PostgresWire` carries: + +| Field | Source | +|-------|--------| +| `user` | Startup message `user` parameter | +| `database` | Startup message `database` parameter | +| `session_params` | Initially empty | +| `tags` | `query_tags` / `query_tag` from startup parameters | + +The `database` and `user` fields are available to routers — the `pythonScript` router receives them in `ctx["database"]` and `ctx["user"]`. + +## SET handling + +`SET` statements are acknowledged with `CommandComplete` + `ReadyForQuery` without forwarding to the backend. This keeps the connection in a valid state for clients that issue `SET` during startup. + +## Client examples + +```bash +# psql +psql -h localhost -p 5432 -U dev -d my_catalog -c "SELECT 1" + +# Python (psycopg2) +import psycopg2 +conn = psycopg2.connect(host="localhost", port=5432, user="dev", dbname="my_catalog") +cur = conn.cursor() +cur.execute("SELECT 42 AS answer") +print(cur.fetchone()) +``` + +## Related + +- [Frontends overview](overview.md) — shared dispatch and session model +- [Routing and clusters](/docs/architecture/routing-and-clusters) — `protocolBased` router with `postgresWire` diff --git a/website/docs/architecture/frontends/snowflake.md b/website/docs/architecture/frontends/snowflake.md new file mode 100644 index 0000000..45c6775 --- /dev/null +++ b/website/docs/architecture/frontends/snowflake.md @@ -0,0 +1,191 @@ +--- +description: Snowflake-compatible frontend — HTTP wire protocol and SQL API v2, session management, query execution, and client connectivity. +--- + +# Snowflake + +QueryFlux can accept connections from **Snowflake clients** — JDBC, ODBC, Python connector, Go driver, Node.js driver, and the SQL API — without any driver changes on the client side. Traffic is **not** proxied to a real Snowflake account; QueryFlux terminates the Snowflake protocol locally, authenticates via its own auth providers, routes through the standard router chain, and executes queries on whichever backend engine routing selects (Trino, StarRocks, DuckDB, etc.). + +This is a **protocol bridge**, not a Snowflake proxy: the client believes it is talking to Snowflake, but the query runs on a different engine entirely. + +## Configuration + +```yaml +queryflux: + frontends: + snowflakeHttp: + enabled: true + port: 8445 +``` + +Config key: `snowflakeHttp`. Protocol identifiers: `FrontendProtocol::SnowflakeHttp` (wire) and `FrontendProtocol::SnowflakeSqlApi` (SQL API). Default dialect: `SqlDialect::Generic` (both). + +The `snowflakeHttp` frontend starts a single HTTP listener that serves **both** wire and SQL API routes. There is no separate port for `snowflakeSqlApi` — both protocol surfaces are merged onto the same listener. + +## Two protocol flavors + +The Snowflake frontend exposes two API surfaces on a **single port**: + +| Protocol | Identifier | Typical clients | Auth model | +|----------|------------|-----------------|------------| +| **Snowflake HTTP wire** | `snowflakeHttp` | JDBC, ODBC, Python connector, Go driver, Node.js driver | Session-based (login → token) | +| **Snowflake SQL API v2** | `snowflakeSqlApi` | REST clients, service accounts, `curl` | Stateless Bearer token per request | + +Both share the same listener; routing can target them independently via `protocolBased` rules: + +```yaml +routers: + - type: protocolBased + snowflakeHttp: trino-default + snowflakeSqlApi: analytics-group +``` + +--- + +## HTTP wire protocol + +The wire protocol handles session lifecycle (login, logout, heartbeat, token refresh) and synchronous query execution. This is the API surface that JDBC, ODBC, and the Python/Go/Node connectors use internally. + +### Endpoints + +| Method | Path | Description | +|--------|------|-------------| +| `POST` | `/session/v1/login-request` | Authenticate and create a session | +| `DELETE` | `/session` | Destroy a session (logout) | +| `GET` | `/session/heartbeat` | Validate that a session token is still alive | +| `POST` | `/session/token-request` | Refresh / validate a session token | +| `POST` | `/queries/v1/query-request` | Execute a SQL query | +| `GET` | `/queries/v1/query-monitoring-request` | Query monitoring (returns empty — sync execution) | +| `DELETE` | `/queries/v1/{query_id}` | Cancel a query (no-op — sync execution) | + +### Login + +`POST /session/v1/login-request` accepts JSON with `LOGIN_NAME` and `PASSWORD` fields. Optional `databaseName` and `schemaName` can be passed as query parameters or in `SESSION_PARAMETERS`. + +On success, QueryFlux: + +1. Authenticates via the configured `auth_provider`. +2. Resolves a **cluster group** by routing with `FrontendProtocol::SnowflakeHttp`. +3. Issues a QueryFlux-generated UUID as both `token` and `masterToken`. +4. Stores the session (auth context, group, database/schema hints, user) in an in-memory session store. + +The returned token is used in subsequent requests via the `Authorization: Snowflake Token=""` header. + +### Query execution + +`POST /queries/v1/query-request` reads the SQL from the `sqlText` JSON field. Gzip-compressed request bodies (`Content-Encoding: gzip`) are supported — the Python connector uses this by default. + +Execution is **synchronous**: the query runs to completion via `execute_to_sink`, and the full result set is returned in a single response. The response includes: + +- **Snowflake-compatible JSON** with `rowType` metadata (column names, types, nullability). +- **Base64-encoded Arrow IPC** in `rowsetBase64` for the Python connector's Arrow result path. +- Standard Snowflake response fields (`success`, `code`, `message`, `data`). + +Errors are returned as **HTTP 200** with `"success": false` in the JSON body — matching Snowflake's convention so connectors handle errors correctly rather than treating non-200 responses as retryable infrastructure failures. + +--- + +## SQL API v2 + +The SQL API follows Snowflake's [REST SQL API](https://docs.snowflake.com/en/developer-guide/sql-api) convention. This is the stateless, request-per-query interface typically used by service accounts and automation. + +### Endpoints + +| Method | Path | Description | +|--------|------|-------------| +| `POST` | `/api/v2/statements` | Submit a SQL statement | +| `GET` | `/api/v2/statements/{handle}` | Get statement status/results | +| `DELETE` | `/api/v2/statements/{handle}` | Cancel a statement | + +### Authentication + +SQL API requests use `Authorization: Bearer `. The token is validated against QueryFlux's `auth_provider` on every request — there is no session to maintain. + +### Submit statement + +`POST /api/v2/statements` accepts JSON with a `statement` field containing the SQL. QueryFlux authenticates the request, routes with `FrontendProtocol::SnowflakeSqlApi`, and executes synchronously. + +The response is a Snowflake SQL API v2-compatible JSON object with: + +- `statementHandle` — a unique identifier for the execution. +- `status` — `"00000"` on success. +- `rowType` — column metadata. +- `data` — rows as arrays of JSON string values. + +### Get / cancel statement + +Since execution is synchronous, `GET /api/v2/statements/{handle}` returns a **404** (no stored handles) and `DELETE /api/v2/statements/{handle}` returns a success stub. + +--- + +## Session management + +The HTTP wire protocol maintains sessions in an in-memory `SnowflakeSessionStore` (`DashMap`). Each session holds: + +| Field | Description | +|-------|-------------| +| `token` | QueryFlux-issued UUID (the session key) | +| `auth_ctx` | Authentication context from the auth provider | +| `group` | Cluster group resolved at login time | +| `user` | Authenticated username | +| `database` / `schema` | Optional hints from the login request | +| `created_at` | Session creation timestamp | + +Wire protocol requests resolve the session from the `Snowflake Token=` header. The cluster group is fixed at login — all queries in a session route to the same group. + +The SQL API does **not** use the session store; each request is independently authenticated and routed. + +--- + +## Connecting clients + +### Python connector + +```python +import snowflake.connector + +conn = snowflake.connector.connect( + user="dev", + password="password", + account="queryflux", # any value — not used for routing + host="localhost", + port=8445, + protocol="http", + database="my_catalog", + schema="my_schema", +) + +cur = conn.cursor() +cur.execute("SELECT 1 AS num, 'hello' AS greeting") +for row in cur: + print(row) +``` + +The Python connector sends requests to the wire protocol endpoints. Set `host` and `port` to point at the QueryFlux Snowflake listener. The `account` value is required by the driver but not used by QueryFlux for routing. + +### JDBC + +``` +jdbc:snowflake://localhost:8445/?account=queryflux&ssl=off&db=my_catalog&schema=my_schema +``` + +### SQL API via curl + +```bash +curl -X POST http://localhost:8445/api/v2/statements \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{"statement": "SELECT 42 AS answer"}' +``` + +--- + +## Dialect and translation + +Both Snowflake protocol flavors map to `SqlDialect::Generic` as their default source dialect. When routing lands on a Trino cluster, for example, sqlglot translates the SQL from the generic dialect to Trino. If the target engine's dialect matches, translation is skipped. + +## Related + +- [Frontends overview](overview.md) — shared dispatch and session model +- [Configuration](/docs/configuration) — `frontends.snowflakeHttp` in `config.yaml` +- [Routing and clusters](/docs/architecture/routing-and-clusters) — `protocolBased` router with `snowflakeHttp` / `snowflakeSqlApi` diff --git a/website/docs/architecture/frontends/trino-http.md b/website/docs/architecture/frontends/trino-http.md new file mode 100644 index 0000000..ad8acb6 --- /dev/null +++ b/website/docs/architecture/frontends/trino-http.md @@ -0,0 +1,84 @@ +--- +description: Trino HTTP frontend — async query polling, nextUri rewriting, session headers, tags, and mixed-engine group support. +--- + +# Trino HTTP + +The Trino HTTP frontend lets any Trino-compatible client (Trino CLI, JDBC, Python `trino`, DBeaver, etc.) connect to QueryFlux as if it were a Trino coordinator. QueryFlux accepts the standard `/v1/statement` API, routes the query, and returns Trino-shaped JSON responses — including `nextUri` polling for async engines. + +## Configuration + +```yaml +queryflux: + frontends: + trinoHttp: + enabled: true + port: 8080 +``` + +Config key: `trinoHttp`. Protocol identifier: `FrontendProtocol::TrinoHttp`. Default dialect: `SqlDialect::Trino`. + +## Endpoints + +| Method | Path | Description | +|--------|------|-------------| +| `POST` | `/v1/statement` | Submit a new SQL query | +| `GET` | `/v1/statement/qf/queued/{id}/{seq}` | Poll a queued query (proxy-side backoff) | +| `GET` | `/v1/statement/qf/executing/{id}` | Poll an executing query | +| `GET` | `/v1/statement/{*path}` | Forward Trino-style poll paths | +| `DELETE` | `/v1/statement/qf/executing/{id}` | Cancel a running query | +| `DELETE` | `/v1/statement/{*path}` | Cancel via forwarded Trino path | + +## Authentication + +Credentials are extracted in this order: + +1. `Authorization: Basic` — decoded username and password. +2. `Authorization: Bearer` — bearer token. +3. `X-Trino-User` header — username only (no password). + +The extracted credentials are passed to the configured `auth_provider`. Failures return HTTP 401 (unauthorized) or 403 (forbidden by authorization policy). + +## Execution model + +Trino HTTP is the only frontend that supports **async polling**: + +- **Async-capable group** (e.g. Trino backend): `dispatch_query` submits to the engine, persists the executing state, and returns a Trino JSON response. The `nextUri` in the response is **rewritten** to point back at QueryFlux (`externalAddress`), so subsequent polls flow through the proxy transparently. +- **At capacity**: when the group is full, QueryFlux returns a synthetic "queued" response with a `nextUri` pointing at `/v1/statement/qf/queued/{id}/{seq}`. The client polls this URL; QueryFlux retries cluster acquisition on each poll. +- **Sync engine in a Trino group** (e.g. DuckDB, StarRocks): falls back to `execute_to_sink` with a `TrinoHttpResultSink` — the query runs to completion and the result is returned as a single Trino JSON page (no polling). + +## Session context + +`SessionContext::TrinoHttp` carries all request headers (lowercased keys) as a `HashMap`. Routers and the Python script router can inspect any header (e.g. `x-trino-user`, `x-trino-catalog`, `x-trino-schema`). + +## Query tags + +Tags can be set via: + +- `X-Trino-Client-Tags` header — comma-separated key/value pairs. +- `X-Trino-Session` header — `query_tags` or `query_tag` keys (percent-decoded). +- `SET SESSION query_tags = '...'` SQL statement — intercepted by the frontend (not forwarded to the backend). The response includes `X-Trino-Set-Session` so the client tracks the change. + +Tags are used by the `tags` router for routing decisions. See [Query tags](/docs/architecture/query-tags). + +## nextUri rewriting + +When proxying Trino async queries, QueryFlux rewrites `nextUri` URLs in response JSON so the client always polls through the proxy rather than going directly to the Trino coordinator. The rewrite replaces the host/port with `externalAddress` while preserving the Trino path from `/v1/` onward. This is done via fast byte-level patching when possible, avoiding full JSON deserialization. + +## Client examples + +```bash +# Trino CLI +trino --server http://localhost:8080 --execute "SELECT 42" + +# curl +curl -X POST http://localhost:8080/v1/statement \ + -H "X-Trino-User: dev" \ + -d "SELECT current_date" +``` + +## Related + +- [Frontends overview](overview.md) — shared dispatch and session model +- [Routing and clusters](/docs/architecture/routing-and-clusters) — how `protocolBased` maps `trinoHttp` to a group +- [Query tags](/docs/architecture/query-tags) — tag-based routing from Trino sessions diff --git a/website/docs/architecture/overview.md b/website/docs/architecture/overview.md index 098fcc8..86e806b 100644 --- a/website/docs/architecture/overview.md +++ b/website/docs/architecture/overview.md @@ -14,7 +14,8 @@ This section describes how QueryFlux is put together: why it exists, how SQL is | [Query translation](query-translation.md) | Dialect detection, sqlglot integration, when translation runs, and schema-aware mode. | | [Routing and clusters](routing-and-clusters.md) | Router chain, `routingFallback`, cluster groups, load-balancing strategies, and queueing. | | [Observability](observability.md) | Prometheus metrics, Grafana dashboard, QueryFlux Studio, and the Admin REST API. | -| [Adding engine support](adding-engine-support.md) | Checklist for new **backend engines** (Rust adapters), **Studio** (`lib/studio-engines/` manifest + catalog slots), and **frontend wire protocols** (e.g. Postgres wire). | +| [Frontends](frontends/overview.md) | Protocol listeners — Trino HTTP, PostgreSQL wire, MySQL wire, Flight SQL, and Snowflake. Shared dispatch, session model, and per-protocol details. | +| [Extending QueryFlux](adding-support/overview.md) | **[Backend](adding-support/backend.md)** (Rust + Studio) and **[Frontend](adding-support/frontend.md)** (new protocols). Same sidebar category as **[Frontends](frontends/overview.md)**. | | [Auth / authz design](auth-authz-design.md) | Authentication and authorization design notes. | Start with [Motivation and goals](motivation-and-goals.md) if you are new to the project; use [System map](system-map.md) as the single-page map of the system. diff --git a/website/docs/architecture/routing-and-clusters.md b/website/docs/architecture/routing-and-clusters.md index 386e74d..ad62b2f 100644 --- a/website/docs/architecture/routing-and-clusters.md +++ b/website/docs/architecture/routing-and-clusters.md @@ -34,7 +34,7 @@ Configured under `routers:` in YAML (`queryflux_core::config::RouterConfig`). Wi | `type` | Behavior | |--------|----------| -| `protocolBased` | Maps the active frontend (`trinoHttp`, `postgresWire`, `mysqlWire`, `flightSql`, `clickhouseHttp`) to a group name. | +| `protocolBased` | Maps the active frontend (`trinoHttp`, `postgresWire`, `mysqlWire`, `flightSql`, `clickhouseHttp`, `snowflakeHttp`, `snowflakeSqlApi`) to a group name. | | `header` | Matches a header value to a group (useful for Trino HTTP and similar). | | `queryRegex` | Ordered rules: first regex match on the SQL text wins. | | `tags` | Routes based on query tags attached to the session. Each rule specifies one or more tag key/value conditions (AND logic); first matching rule wins. See [Tags router](#tags-router-tags) below. | diff --git a/website/docs/architecture/system-map.md b/website/docs/architecture/system-map.md index 38bd783..b18b99b 100644 --- a/website/docs/architecture/system-map.md +++ b/website/docs/architecture/system-map.md @@ -4,9 +4,9 @@ description: End-to-end system map of QueryFlux — protocol frontends, dispatch # QueryFlux — Architecture Overview -QueryFlux is a universal SQL query proxy and router. It accepts queries from clients over multiple protocols (Trino HTTP, PostgreSQL wire, MySQL wire, Arrow Flight SQL), routes them to the appropriate backend engine, optionally translates the SQL dialect, and streams results back in the client's native format. +QueryFlux is a universal SQL query proxy and router. It accepts queries from clients over multiple protocols (Trino HTTP, PostgreSQL wire, MySQL wire, Arrow Flight SQL, Snowflake HTTP), routes them to the appropriate backend engine, optionally translates the SQL dialect, and streams results back in the client's native format. -**More documentation:** the [architecture documentation overview](./overview.md) indexes deeper topics — [motivation-and-goals.md](motivation-and-goals.md) (why the project exists), [query-translation.md](query-translation.md) (sqlglot and dialects), [routing-and-clusters.md](routing-and-clusters.md) (routers, groups, load balancing), [observability.md](observability.md) (Prometheus, Grafana, Studio, Admin API), [adding-engine-support.md](adding-engine-support.md) (new engines, Studio, and client protocols). +**More documentation:** the [architecture documentation overview](./overview.md) indexes deeper topics — [motivation-and-goals.md](motivation-and-goals.md) (why the project exists), [query-translation.md](query-translation.md) (sqlglot and dialects), [routing-and-clusters.md](routing-and-clusters.md) (routers, groups, load balancing), [observability.md](observability.md) (Prometheus, Grafana, Studio, Admin API), [adding-support/overview.md](adding-support/overview.md) (Extending QueryFlux — backend, frontend). --- @@ -163,6 +163,7 @@ pub trait RouterTrait: Send + Sync { | PostgreSQL wire | **Done** | 5432 | | MySQL wire | **Done** | 3306 | | Arrow Flight SQL | **Done** (query execution) | — | +| Snowflake HTTP wire + SQL API | **Done** | configurable | | Admin / Prometheus metrics | **Done** | 9000 | | ClickHouse HTTP | Planned | 8123 | @@ -251,6 +252,7 @@ queryflux: trinoHttp: { enabled: true, port: 8080 } postgresWire: { enabled: false, port: 5432 } mysqlWire: { enabled: false, port: 3306 } + snowflakeHttp: { enabled: false, port: 8445 } persistence: inMemory: {} # or: postgres: { databaseUrl: "postgres://..." } adminApi: @@ -358,6 +360,7 @@ curl -s -X POST http://localhost:8080/v1/statement \ | P1 | PostgreSQL wire frontend | **Done** | | P1 | MySQL wire frontend + StarRocks backend | **Done** | | P1 | Arrow Flight SQL frontend | **Done** | +| P1 | Snowflake HTTP wire + SQL API frontend | **Done** | | P1 | QueryFlux Studio — management UI | **Done** | | P1 | Athena backend | **Done** | | P1 | Authentication / authorization (`queryflux-auth`) | **Done** | diff --git a/website/docs/configuration.md b/website/docs/configuration.md index 37a87f8..bab76d9 100644 --- a/website/docs/configuration.md +++ b/website/docs/configuration.md @@ -14,6 +14,9 @@ queryflux: trinoHttp: enabled: true port: 8080 + snowflakeHttp: + enabled: true + port: 8445 persistence: type: inMemory # or: postgres @@ -39,6 +42,8 @@ clusterGroups: routers: - type: protocolBased trinoHttp: trino-default + snowflakeHttp: trino-default + snowflakeSqlApi: trino-default - type: header headerName: x-target-engine diff --git a/website/docs/roadmap.md b/website/docs/roadmap.md index 4625e9b..37b6fe5 100644 --- a/website/docs/roadmap.md +++ b/website/docs/roadmap.md @@ -19,6 +19,7 @@ Everything below is implemented and available on the `main` branch. | | PostgreSQL wire protocol (port 5432) | | | MySQL wire protocol (port 3306) | | | Arrow Flight SQL (gRPC) | +| | Snowflake HTTP wire + SQL API v2 ([docs](/docs/architecture/frontends/snowflake)) | | | Admin REST API + OpenAPI / Swagger UI (port 9000) | | **Backends** | Trino — async HTTP polling, transparent `nextUri` proxying | | | DuckDB — embedded, in-process, Arrow result sets | diff --git a/website/sidebars.ts b/website/sidebars.ts index c2c2fbe..cff082a 100644 --- a/website/sidebars.ts +++ b/website/sidebars.ts @@ -29,7 +29,29 @@ const sidebars: SidebarsConfig = { 'architecture/routing-and-clusters', 'architecture/query-tags', 'architecture/observability', - 'architecture/adding-engine-support', + { + type: 'category', + label: 'Frontends', + collapsed: false, + items: [ + 'architecture/frontends/overview', + 'architecture/frontends/trino-http', + 'architecture/frontends/postgres-wire', + 'architecture/frontends/mysql-wire', + 'architecture/frontends/flight-sql', + 'architecture/frontends/snowflake', + ], + }, + { + type: 'category', + label: 'Extending QueryFlux', + collapsed: false, + items: [ + 'architecture/adding-support/overview', + 'architecture/adding-support/backend', + 'architecture/adding-support/frontend', + ], + }, 'architecture/auth-authz-design', ], }, diff --git a/website/versioned_docs/version-0.1.0/architecture/adding-engine-support.md b/website/versioned_docs/version-0.1.0/architecture/adding-engine-support.md index 1ea3bb4..32f999e 100644 --- a/website/versioned_docs/version-0.1.0/architecture/adding-engine-support.md +++ b/website/versioned_docs/version-0.1.0/architecture/adding-engine-support.md @@ -1,314 +1,13 @@ --- -description: Checklist for new backend engines (Rust adapters), Studio (lib/studio-engines/ manifest + catalog slots), and frontend wire protocols (e.g. Postgres wire). +description: This page moved — see Extending QueryFlux (overview, backend, frontend). --- # Adding engine and protocol support -This guide separates two ideas that are easy to conflate: +This guide lives under **Extending QueryFlux** in the sidebar: -| Concept | Meaning | Example | -|--------|---------|---------| -| **Backend engine** | A **cluster** type QueryFlux routes queries **to**. It has an adapter that talks to the real database (HTTP, MySQL wire, embedded library, AWS SDK, …). | Trino, DuckDB, StarRocks, Athena | -| **Frontend protocol** | How **clients connect to QueryFlux** (ingress). SQL enters with a `FrontendProtocol` and a default source dialect for translation. | Trino HTTP, **PostgreSQL wire**, MySQL wire, Flight SQL | +- **[Overview](adding-support/overview)** — concepts and index +- **[Backend](adding-support/backend)** — Rust adapter, persistence, Studio +- **[Frontend](adding-support/frontend)** — new client protocols -Adding **PostgreSQL wire** as a client entrypoint is **not** the same as adding “PostgreSQL” as a backend: today, `PostgresWire` is already a frontend in `queryflux-frontend`; traffic still lands on the shared dispatch path and is sent to whatever **backend adapter** routing chose (often Trino). - -Use the sections below depending on whether you are extending **Studio**, a **backend adapter**, or a **frontend listener**. - ---- - -## Part A — Backend engine (Rust) - -Goal: a new `engine` value in cluster config, a live adapter, validation, translation target dialect, and wiring in the binary. - -### Registration overview - -Backends are **not** loaded dynamically. Each engine is compiled in and registered explicitly. Data flow: - -1. **Postgres / YAML** → `engine_key` column + `config` JSONB → `ClusterConfigRecord::to_core()` uses **`parse_engine_key`** and JSON helpers → typed **`ClusterConfig`**. -2. **Binary** → `registered_engines::build_adapter(...)` matches **`EngineConfig`** and calls the adapter’s **`try_from_cluster_config`** (see [`crates/queryflux/src/registered_engines.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux/src/registered_engines.rs)). -3. **Adapter** → reads only the **`ClusterConfig`** fields it needs (endpoint, auth, region, …) and constructs itself; **startup** and **hot reload** both use the same factory. - -**JSONB** stores per-cluster, per-engine payload without schema migrations; **`ClusterConfig`** in core is the typed view after `to_core()`. Engine-specific wiring belongs in **`try_from_cluster_config`**, not in `main.rs`. - -### 1. Core model (`queryflux-core`) - -- **`EngineConfig`** — Add a variant in `crates/queryflux-core/src/config.rs` (serde **camelCase** in JSON/YAML, e.g. `myEngine`). -- **`EngineType`** — Add a variant in `crates/queryflux-core/src/query.rs` if the backend is distinct for metrics, translation, or dispatch. -- **`engine_registry`** (`crates/queryflux-core/src/engine_registry.rs`) — Keep these in sync when you add a variant: - - **`engine_key(&EngineConfig)`** — `EngineConfig` → stable string key (must match the adapter descriptor and Studio). - - **`parse_engine_key(&str)`** — inverse mapping for the `engine_key` column in Postgres / API. - - **`impl From<&EngineConfig> for EngineType`** — single place for config → runtime `EngineType` (cluster manager and `main.rs` use this instead of ad-hoc matches). -- **`EngineType::dialect()`** — Return the `SqlDialect` used as the **translation target** (and extend `SqlDialect` / `is_compatible_with` in translation if needed). See [query-translation.md](query-translation.md). -- **`ClusterConfig` fields** — Add any new top-level fields (region, paths, engine-specific blobs). Prefer keeping engine-specific secrets and options in `config` JSON for Postgres-backed clusters; extend the typed struct when YAML and validation need them everywhere. - -### 2. Adapter crate (`queryflux-engine-adapters`) - -- Add a module (e.g. `src/myengine/mod.rs`) implementing **`EngineAdapterTrait`** (`submit_query`, `poll_query`, `cancel_query`, `health_check`, `engine_type`, `supports_async`, and optionally `fetch_running_query_count`, `base_url`, Arrow/catalog hooks as needed). -- Implement **`descriptor() -> EngineDescriptor`** with: - - `engine_key`, `display_name`, `description`, `hex` - - `connection_type` (`Http`, `MySqlWire`, `Embedded`, `ManagedApi`) - - `supported_auth` and **`config_fields`** (these drive `/admin/engine-registry` and should stay aligned with Studio) - - `implemented: true` when the adapter is actually wired in `main` -- Export the module from `crates/queryflux-engine-adapters/src/lib.rs` and add the crate dependency if you introduce new third-party crates. - -**Factory — `try_from_cluster_config`** - -Implement on your adapter struct so all **field extraction and validation** for that engine live next to the adapter (not in `registered_engines.rs`): - -- **Sync** (most engines): - - ```text - fn try_from_cluster_config( - cluster_name: ClusterName, - group_name: ClusterGroupName, - cfg: &ClusterConfig, - cluster_name_str: &str, - ) -> queryflux_core::error::Result - ``` - -- **Async** (e.g. Athena — AWS client setup): same parameters, `async fn`, returns `Result`. - -Use **`QueryFluxError::Engine(format!(…))`** for failures; include **`cluster_name_str`** in messages so startup and reload logs identify the cluster. Reference implementations: **`TrinoAdapter`** and **`StarRocksAdapter`** ([`trino/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/trino/mod.rs), [`starrocks/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/starrocks/mod.rs)), **`DuckDbAdapter`** / **`DuckDbHttpAdapter`** ([`duckdb/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/duckdb/mod.rs), [`duckdb/http.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/duckdb/http.rs)), **`AthenaAdapter`** ([`athena/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/athena/mod.rs)). - -Keep **`pub fn new(...)`** (or **`async fn new`**) as the low-level constructor if you want tests to build adapters without a full **`ClusterConfig`**; **`try_from_cluster_config`** can delegate to **`new`** after parsing **`cfg`**. - -### 3. Binary wiring (`crates/queryflux`) - -Registration is centralized in **`crates/queryflux/src/registered_engines.rs`**: - -- **`all_descriptors()`** — Append **`MyEngineAdapter::descriptor()`** to the returned `vec!`. [`main.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux/src/main.rs) builds **`EngineRegistry::new(registered_engines::all_descriptors())`** for validation and **`GET /admin/engine-registry`**. -- **`build_adapter(cluster_name, placeholder_group, cluster_cfg, cluster_name_str).await`** — Returns **`anyhow::Result>`**. Add a **`match`** arm on **`EngineConfig::MyEngine`** that calls **`MyEngineAdapter::try_from_cluster_config(...)`**, maps **`QueryFluxError`** to **`anyhow::Error`** (same helper as other arms), and wraps **`Arc::new(...)`**. **Startup** uses **`.context(...)?`** on the result; **hot reload** in **`build_live_config`** logs a warning and **`continue`** on error — behavior stays in **`main.rs`**, not in the factory. - -Do **not** add a second adapter-construction **`match`** in **`main.rs`**. - -**Not implemented yet:** e.g. **`EngineConfig::ClickHouse`** is handled inside **`build_adapter`** with **`anyhow::bail!`** until a **`ClickHouseAdapter`** and **`try_from_cluster_config`** exist. - -- **`EngineType` for cluster state** — In **`main.rs`** and anywhere else (e.g. group member **`ClusterState`**), use **`EngineType::from(engine_config)`** from **`engine_registry.rs`**. **`queryflux-cluster-manager`** engine affinity uses the same **`From`** impl (see **`strategy.rs`**). - -- **Special rules** — Search for engine-specific checks (e.g. `queryAuth` / impersonation) and extend validation if your engine has constraints. - -### 4. Dispatch and frontends (`queryflux-frontend`) - -- Shared query execution goes through **`dispatch_query`** / **`execute_to_sink`**. Usually no change if the new engine only differs in the adapter; if you need a **special execution path** (like Trino raw HTTP), follow the existing engine-specific branches. -- Per-protocol handlers (Trino HTTP, Postgres wire, …) should keep using the shared dispatch layer unless the protocol requires a dedicated contract. - -### 5. Persistence (`queryflux-persistence`) — why touch it if config is JSON? - -The table stores **`engine_key` as its own column** plus a **`config` JSONB** blob. The DB does not load straight into the proxy as opaque JSON: code paths call **`ClusterConfigRecord::to_core()`**, which must produce a typed **`ClusterConfig`** (including **`EngineConfig`**). - -So persistence changes are **not** “because Postgres needs a JSON schema.” They are because of this **explicit conversion layer**: - -1. **`ClusterConfigRecord::to_core`** — Calls **`parse_engine_key`** from `queryflux-core` (next to `engine_key`). Extend **`parse_engine_key`** when you add an engine; you do **not** maintain a second duplicate string match in persistence. -2. **`UpsertClusterConfig::from_core`** — Uses **`engine_key(&EngineConfig)`** from core to set the `engine_key` column when seeding from YAML. - -**Extra JSON keys** that only live inside `config` and are **already** read in `to_core` (e.g. `endpoint`, `region`, `authType`, …) usually need **no** persistence change beyond the engine-key match. You only extend the `s("…")` / `b("…")` helpers in `to_core` (and the matching `from_core` inserts) if you add **new top-level persisted fields** on `ClusterConfig` that should round-trip through that JSON. - -**Hot reload** often uses `list_cluster_configs` → records → `to_core()` → `build_live_config`; the same conversion applies. - -### 6. Optional: routing config - -- If operators choose the new group via **router JSON** (`RouterConfig` variants), no change unless you add a new router **type**. -- **Protocol-based routing** maps frontend labels to **group names**; it does not list backend engines. - -### 7. Tests and docs - -- Add or extend **e2e** tests under `crates/queryflux-e2e-tests` if you have a dockerized target. -- Update [system-map.md](system-map.md) component status if you document supported engines there. - -### 8. Suggested order of work (backend only) - -1. **`EngineConfig` / `EngineType`** + **`engine_key` / `parse_engine_key` / `From<&EngineConfig> for EngineType`** + **`dialect()`** if needed. -2. **`ClusterConfig`** fields if the engine needs new top-level keys (and persistence **`to_core`** JSON extraction if those keys live in JSONB). -3. Adapter module: **`EngineAdapterTrait`**, **`descriptor()`**, **`try_from_cluster_config`**. -4. **`registered_engines.rs`**: descriptor in **`all_descriptors()`**, arm in **`build_adapter`**. -5. Run **`cargo build -p queryflux`**; exercise **YAML** and **Postgres** load + **admin upsert** if applicable. - ---- - -## Part B — QueryFlux Studio (UI, TypeScript / React) - -Studio is the Next.js app under `ui/queryflux-studio/`. It does **not** run wire protocols; it calls the **Admin API** (`ADMIN_API_URL`, default `http://localhost:9000`) for clusters, groups, routing, and scripts. - -Backend engines are registered in Studio through **`StudioEngineModule`** objects: one file per engine under **`lib/studio-engines/engines/`**, aggregated in **`lib/studio-engines/manifest.ts`**. That manifest drives **`ENGINE_REGISTRY`**, catalog slots for implemented backends, optional flat-form validation, engine-affinity dropdown entries, and extra **`findEngineByType`** aliases. - -The proxy still exposes descriptors at **`GET /admin/engine-registry`**. Studio does **not** load that at runtime yet, so **Rust `descriptor()` and each studio module’s `descriptor` field must stay aligned by hand** (same `engineKey`, `configFields` keys, auth shapes, etc.). Shared TypeScript types live in **`lib/engine-registry-types.ts`**; **`lib/engine-registry.ts`** only re-exports helpers and builds **`ENGINE_REGISTRY`** from the manifest. - -### Where users see engines - -| User action | UI entrypoint | What must know your engine | -|-------------|---------------|----------------------------| -| Create cluster | **Clusters → Add cluster** (`components/add-cluster-dialog.tsx`) | Expanded **`ENGINE_CATALOG`** (includes studio slots) + **`findEngineDescriptor`** + **`validateClusterConfig`** / **`validateEngineSpecific`** + **`toUpsertBody`** | -| Edit cluster | **Clusters** grid → cluster card → Edit (`app/clusters/clusters-grid.tsx`) | Same + **`mergeClusterConfigFromFlat`** / **`buildClusterUpsertFromForm`** + **`EngineClusterConfig`** | -| View config | Cluster detail / engine config view in `clusters-grid.tsx` | **`findEngineDescriptor`** for labels; unknown key shows “add to engine registry” warning | -| Group strategy **engine affinity** | **Engines →** group dialog → strategy (`components/group-form-dialog.tsx`) | **`ENGINE_AFFINITY_OPTIONS`** is built by **`buildEngineAffinityOptionsFromManifest()`** from each module’s **`engineAffinity`** field (omit label override, or set **`engineAffinity: false`** to exclude, e.g. Athena). | -| Live utilization cards | **Engines (Groups)** page (`app/engines/page.tsx`) | **`findEngineByType`**; studio modules contribute aliases via **`extraTypeAliases`** (merged with static dialect aliases in **`components/engine-catalog.ts`**) | - -### 1. Studio engine module (primary registration) - -**Types:** `ui/queryflux-studio/lib/studio-engines/types.ts` — **`StudioEngineModule`**. - -**Per engine:** `ui/queryflux-studio/lib/studio-engines/engines/.ts` - -Export a constant (e.g. **`trinoStudioEngine`**) with: - -- **`descriptor`** — Full **`EngineDescriptor`** (must match Rust: **`engineKey`**, **`connectionType`**, **`supportedAuth`**, **`configFields`**, **`implemented`**, branding **`hex`**, etc.). Extend **`ConnectionType`** / **`AuthType`** in **`lib/engine-registry-types.ts`** if Rust added a variant. -- **`catalog`** — **`category`**, **`simpleIconSlug`**, **`catalogDescription`** for the engines grid / picker (display name and **`supported`** come from the descriptor when the catalog is expanded). -- **`validateFlat`** (optional) — Cross-field checks before save (e.g. Trino basic vs bearer). Dispatched by **`validateEngineSpecific`** in **`lib/studio-engines/validate-flat.ts`** (re-exported from **`lib/cluster-persist-form.ts`**). -- **`customFormId`** (optional) — String key; must match an entry in **`components/cluster-config/studio-engine-forms.tsx`** if the generic **`GenericEngineClusterConfig`** is not enough. -- **`engineAffinity`** (optional) — **`false`** to omit from affinity, or **`{ label?: string }`** for a custom dropdown label (default label is **`displayName`**). -- **`extraTypeAliases`** (optional) — Map of normalized API/type strings → canonical **`EngineDef.name`** for **`findEngineByType`** (e.g. alternate spellings). - -**Manifest:** `ui/queryflux-studio/lib/studio-engines/manifest.ts` - -- Import the new module and append it to **`STUDIO_ENGINE_MODULES`** (order affects **`ENGINE_AFFINITY_OPTIONS`** and registry iteration; catalog **card order** is separate — see below). - -**Derived registry:** `ui/queryflux-studio/lib/engine-registry.ts` - -- **`ENGINE_REGISTRY`** is **`STUDIO_ENGINE_MODULES.map((m) => m.descriptor)`**. Do not duplicate descriptor arrays here. -- **`findEngineDescriptor`**, **`implementedEngines`**, **`isClusterOnboardingSelectable`**, **`validateClusterConfig`** — unchanged behavior; **`validateClusterConfig`** still uses generic required-field checks from **`configFields`** unless you extend the Rust/TS contract. - -### 2. Catalog layout (picker order and dialect-only rows) - -**File:** `ui/queryflux-studio/components/engine-catalog.ts` - -- Implemented backends appear as **studio slots**: **`{ k: "studio", engineKey: "" }`** inside **`ENGINE_CATALOG_SLOTS`**, interleaved with static **`EngineDef`** rows (dialects with **`engineKey: null`**). -- At runtime, **`expandCatalog`** replaces each studio slot with **`studioModuleToEngineDef`** from **`lib/studio-engines/catalog.ts`**. -- Static **`STATIC_ENGINE_TYPE_ALIASES`** remains for dialects without a studio module; **`buildStudioTypeAliases()`** merges in per-module aliases and the lowercase **`engineKey`** → **`displayName`** mapping. - -**`isClusterOnboardingSelectable`** still requires a catalog row with **`supported`** and **`engineKey`**; for studio-backed engines, **`supported`** is **`descriptor.implemented`** after expansion. - -### 3. Cluster config forms - -**Router:** `ui/queryflux-studio/components/cluster-config/engine-cluster-config.tsx` - -- Resolves **`getStudioEngineModule(engineKey)`**; if **`customFormId`** is set and **`STUDIO_CUSTOM_CLUSTER_FORMS[id]`** exists, renders that component; otherwise **`GenericEngineClusterConfig`** (descriptor **`configFields`**). - -**Custom form registration:** `ui/queryflux-studio/components/cluster-config/studio-engine-forms.tsx` — map **`customFormId`** → component (see Trino / StarRocks / Athena). - -**Reference components:** `trino-cluster-config.tsx`, `starrocks-cluster-config.tsx`, `athena-cluster-config.tsx`, `generic-engine-cluster-config.tsx`, `config-field-row.tsx`. - -### 4. Persisted JSON ↔ flat form (create + edit save path) - -**File:** `ui/queryflux-studio/lib/cluster-persist-form.ts` - -- Still shared across engines. If **`cluster_configs.config`** gains **new top-level JSON keys**, update: - - **`MANAGED_CONFIG_JSON_KEYS`** - - **`persistedClusterConfigToFlat`**, **`flatToPersistedConfig`**, **`mergeClusterConfigFromFlat`** - - **`buildValidateShape`** (shape expected by **`validateClusterConfig`**) -- **`validateEngineSpecific`** is implemented in **`lib/studio-engines/validate-flat.ts`** (per-module **`validateFlat`**); this file re-exports it for call sites. - -### 5. Clusters page (grid, dialog, validation) - -**File:** `ui/queryflux-studio/app/clusters/clusters-grid.tsx` - -- Uses **`findEngineDescriptor`**, **`validateClusterConfig`**, **`validateEngineSpecific`**, **`buildValidateShape`**, **`skipImplementedCheck`** where needed. No per-engine branches beyond **`EngineClusterConfig`**. - -**File:** `ui/queryflux-studio/components/add-cluster-dialog.tsx` - -- Wires catalog → descriptor → **`EngineClusterConfig`** → **`toUpsertBody`** → **`upsertClusterConfig`**. - -### 6. Group strategy (engine affinity) - -**File:** `ui/queryflux-studio/lib/cluster-group-strategy.ts` - -- **`ENGINE_AFFINITY_OPTIONS`** = **`buildEngineAffinityOptionsFromManifest()`**. To exclude an engine, set **`engineAffinity: false`** on its **`StudioEngineModule`**. To customize the label, use **`engineAffinity: { label: "…" }`**. - -### 7. Display helpers - -**File:** `ui/queryflux-studio/lib/merge-clusters-display.ts` - -- Uses **`findEngineDescriptor(p.engineKey)`**; the descriptor must exist in the manifest. - -**File:** `ui/queryflux-studio/components/ui-helpers.tsx` (**`EngineBadge`**) - -- Uses **`ENGINE_CATALOG`**; studio-expanded rows must match **`displayName`** where badges key off names. - -**File:** `ui/queryflux-studio/components/engine-icon.tsx` - -- Consumes **`EngineDef`** (re-exported from **`engine-catalog.ts`**; types in **`lib/engine-catalog-types.ts`**). - -### 8. API types (usually unchanged) - -**File:** `ui/queryflux-studio/lib/api-types.ts` - -- **`ClusterConfigRecord`** / **`UpsertClusterConfig`** stay generic unless you add typed helpers. - -### 9. Optional: fetch registry from the proxy - -A follow-up could load **`GET /admin/engine-registry`** at runtime and hydrate forms from the API. Until then, keep **Rust `descriptor()`** and **`StudioEngineModule.descriptor`** in sync manually. - -### Studio checklist (copy-paste) - -- [ ] **`lib/studio-engines/engines/.ts`** — **`StudioEngineModule`** (`descriptor` aligned with Rust, **`catalog`**, optional **`validateFlat`**, **`customFormId`**, **`engineAffinity`**, **`extraTypeAliases`**) -- [ ] **`lib/studio-engines/manifest.ts`** — import + append to **`STUDIO_ENGINE_MODULES`** -- [ ] **`lib/engine-registry-types.ts`** — extend **`ConnectionType`** / **`AuthType`** if needed -- [ ] **`components/engine-catalog.ts`** — add **`{ k: "studio", engineKey: "…" }`** to **`ENGINE_CATALOG_SLOTS`** at the desired position -- [ ] **`components/cluster-config/studio-engine-forms.tsx`** — register component if **`customFormId`** is set -- [ ] **`lib/cluster-persist-form.ts`** — only if new persisted **`config`** JSON keys (managed keys + flat ↔ JSON + **`buildValidateShape`**) -- [ ] Smoke-test: Add cluster → save → edit → save; Engines page icons; group **engine affinity** if applicable - ---- - -## Part C — Frontend protocol (e.g. “more Postgres wire”) - -Goal: clients speak a **wire protocol to QueryFlux**, not a new backend. - -### Where the code lives - -- **PostgreSQL wire:** `crates/queryflux-frontend/src/postgres_wire/` -- **MySQL wire:** `crates/queryflux-frontend/src/mysql_wire/` -- **Trino HTTP:** `crates/queryflux-frontend/src/trino_http/` -- **Flight SQL:** `crates/queryflux-frontend/src/flight_sql/` - -### Typical steps - -1. **`FrontendProtocol`** — Already defined in `queryflux_core::query::FrontendProtocol`; add a variant only for a **new** ingress protocol. -2. **`default_dialect()`** — Set the sqlglot **source** dialect for translation (see [query-translation.md](query-translation.md)). -3. **Listener** — Bind a port, parse the protocol, build **`SessionContext`** and **`InboundQuery`**, then call shared **`dispatch_query`** (or the same helpers Trino HTTP uses). -4. **Routing** — Optionally extend **protocol-based routing** in config / persisted routing so this frontend maps to the right default group. -5. **Tests** — Protocol-level tests or e2e clients as appropriate. - -Studio does **not** implement wire protocols; it only talks to the **Admin API** for config and metrics. - ---- - -## Checklist summary - -**Backend engine** - -- [ ] `EngineConfig` + `EngineType` + `engine_key()` + **`parse_engine_key()`** + **`From<&EngineConfig> for EngineType`** + dialect mapping (`engine_registry.rs` + `query.rs`) -- [ ] `EngineAdapterTrait` + `descriptor()` -- [ ] `registered_engines.rs`: **`all_descriptors()`** + **`build_adapter()`** arm calling **`try_from_cluster_config`** on the adapter -- [ ] Adapter module: **`try_from_cluster_config`** (or async equivalent) reading **`ClusterConfig`** -- [ ] `UpsertClusterConfig::from_core` / `to_core` stay aligned via **`engine_key` / `parse_engine_key`** (no extra string match in persistence) -- [ ] Translation / compatibility if dialect is new - -**Studio (UI)** - -- [ ] `lib/studio-engines/engines/.ts` — `StudioEngineModule` (descriptor + catalog + options) -- [ ] `lib/studio-engines/manifest.ts` — register module in `STUDIO_ENGINE_MODULES` -- [ ] `lib/engine-registry-types.ts` — `ConnectionType` / `AuthType` if Rust added variants -- [ ] `components/engine-catalog.ts` — `{ k: "studio", engineKey }` slot in `ENGINE_CATALOG_SLOTS` -- [ ] `components/cluster-config/studio-engine-forms.tsx` — only if using `customFormId` -- [ ] `lib/cluster-persist-form.ts` — only if new `config` JSON keys need round-tripping -- [ ] Verify add-cluster + edit-cluster, Engines page icons / `findEngineByType`, and engine affinity if used - -**New client protocol** - -- [ ] `FrontendProtocol` + dialect + listener module + dispatch integration + routing docs - ---- - -## Related reading - -- [system-map.md](system-map.md) — End-to-end flow -- [query-translation.md](query-translation.md) — Dialects and sqlglot -- [routing-and-clusters.md](routing-and-clusters.md) — Routers and groups -- [observability.md](observability.md) — Admin API (including engine registry JSON) - -**Rust files referenced above** - -- [`crates/queryflux/src/registered_engines.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux/src/registered_engines.rs) — `all_descriptors`, `build_adapter` -- [`crates/queryflux-core/src/engine_registry.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-core/src/engine_registry.rs) — `engine_key`, `parse_engine_key`, `EngineRegistry`, `From<&EngineConfig> for EngineType` -- [`crates/queryflux-persistence/src/cluster_config.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-persistence/src/cluster_config.rs) — `to_core` / `from_core` vs `engine_key` + JSONB +Existing frontends are documented in **[Frontends](frontends/overview)**. diff --git a/website/versioned_docs/version-0.1.0/architecture/adding-support/backend.md b/website/versioned_docs/version-0.1.0/architecture/adding-support/backend.md new file mode 100644 index 0000000..e11c42c --- /dev/null +++ b/website/versioned_docs/version-0.1.0/architecture/adding-support/backend.md @@ -0,0 +1,275 @@ +--- +description: Adding a backend engine to QueryFlux — Rust adapter, registration, persistence, and QueryFlux Studio UI. +--- + +# Backend + +Goal: a new `engine` value in cluster config, a live adapter, validation, translation target dialect, wiring in the binary, and (optionally) Studio UI so operators can manage clusters for that engine. + +See **[Extending QueryFlux](overview.md)** for how this differs from a **frontend** protocol. For existing client protocols, see **[Frontends](../frontends/overview.md)**. + +### Registration overview + +Backends are **not** loaded dynamically. Each engine is compiled in and registered explicitly. Data flow: + +1. **Postgres / YAML** → `engine_key` column + `config` JSONB → `ClusterConfigRecord::to_core()` uses **`parse_engine_key`** and JSON helpers → typed **`ClusterConfig`**. +2. **Binary** → `registered_engines::build_adapter(...)` matches **`EngineConfig`** and calls the adapter’s **`try_from_cluster_config`** (see [`crates/queryflux/src/registered_engines.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux/src/registered_engines.rs)). +3. **Adapter** → reads only the **`ClusterConfig`** fields it needs (endpoint, auth, region, …) and constructs itself; **startup** and **hot reload** both use the same factory. + +**JSONB** stores per-cluster, per-engine payload without schema migrations; **`ClusterConfig`** in core is the typed view after `to_core()`. Engine-specific wiring belongs in **`try_from_cluster_config`**, not in `main.rs`. + +### 1. Core model (`queryflux-core`) + +- **`EngineConfig`** — Add a variant in `crates/queryflux-core/src/config.rs` (serde **camelCase** in JSON/YAML, e.g. `myEngine`). +- **`EngineType`** — Add a variant in `crates/queryflux-core/src/query.rs` if the backend is distinct for metrics, translation, or dispatch. +- **`engine_registry`** (`crates/queryflux-core/src/engine_registry.rs`) — Keep these in sync when you add a variant: + - **`engine_key(&EngineConfig)`** — `EngineConfig` → stable string key (must match the adapter descriptor and Studio). + - **`parse_engine_key(&str)`** — inverse mapping for the `engine_key` column in Postgres / API. + - **`impl From<&EngineConfig> for EngineType`** — single place for config → runtime `EngineType` (cluster manager and `main.rs` use this instead of ad-hoc matches). +- **`EngineType::dialect()`** — Return the `SqlDialect` used as the **translation target** (and extend `SqlDialect` / `is_compatible_with` in translation if needed). See [query-translation.md](../query-translation.md). +- **`ClusterConfig` fields** — Add any new top-level fields (region, paths, engine-specific blobs). Prefer keeping engine-specific secrets and options in `config` JSON for Postgres-backed clusters; extend the typed struct when YAML and validation need them everywhere. + +### 2. Adapter crate (`queryflux-engine-adapters`) + +- Add a module (e.g. `src/myengine/mod.rs`) implementing **`EngineAdapterTrait`** (`submit_query`, `poll_query`, `cancel_query`, `health_check`, `engine_type`, `supports_async`, and optionally `fetch_running_query_count`, `base_url`, Arrow/catalog hooks as needed). +- Implement **`descriptor() -> EngineDescriptor`** with: + - `engine_key`, `display_name`, `description`, `hex` + - `connection_type` (`Http`, `MySqlWire`, `Embedded`, `ManagedApi`) + - `supported_auth` and **`config_fields`** (these drive `/admin/engine-registry` and should stay aligned with Studio) + - `implemented: true` when the adapter is actually wired in `main` +- Export the module from `crates/queryflux-engine-adapters/src/lib.rs` and add the crate dependency if you introduce new third-party crates. + +**Factory — `try_from_cluster_config`** + +Implement on your adapter struct so all **field extraction and validation** for that engine live next to the adapter (not in `registered_engines.rs`): + +- **Sync** (most engines): + + ```text + fn try_from_cluster_config( + cluster_name: ClusterName, + group_name: ClusterGroupName, + cfg: &ClusterConfig, + cluster_name_str: &str, + ) -> queryflux_core::error::Result + ``` + +- **Async** (e.g. Athena — AWS client setup): same parameters, `async fn`, returns `Result`. + +Use **`QueryFluxError::Engine(format!(…))`** for failures; include **`cluster_name_str`** in messages so startup and reload logs identify the cluster. Reference implementations: **`TrinoAdapter`** and **`StarRocksAdapter`** ([`trino/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/trino/mod.rs), [`starrocks/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/starrocks/mod.rs)), **`DuckDbAdapter`** / **`DuckDbHttpAdapter`** ([`duckdb/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/duckdb/mod.rs), [`duckdb/http.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/duckdb/http.rs)), **`AthenaAdapter`** ([`athena/mod.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-engine-adapters/src/athena/mod.rs)). + +Keep **`pub fn new(...)`** (or **`async fn new`**) as the low-level constructor if you want tests to build adapters without a full **`ClusterConfig`**; **`try_from_cluster_config`** can delegate to **`new`** after parsing **`cfg`**. + +### 3. Binary wiring (`crates/queryflux`) + +Registration is centralized in **`crates/queryflux/src/registered_engines.rs`**: + +- **`all_descriptors()`** — Append **`MyEngineAdapter::descriptor()`** to the returned `vec!`. [`main.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux/src/main.rs) builds **`EngineRegistry::new(registered_engines::all_descriptors())`** for validation and **`GET /admin/engine-registry`**. +- **`build_adapter(cluster_name, placeholder_group, cluster_cfg, cluster_name_str).await`** — Returns **`anyhow::Result>`**. Add a **`match`** arm on **`EngineConfig::MyEngine`** that calls **`MyEngineAdapter::try_from_cluster_config(...)`**, maps **`QueryFluxError`** to **`anyhow::Error`** (same helper as other arms), and wraps **`Arc::new(...)`**. **Startup** uses **`.context(...)?`** on the result; **hot reload** in **`build_live_config`** logs a warning and **`continue`** on error — behavior stays in **`main.rs`**, not in the factory. + +Do **not** add a second adapter-construction **`match`** in **`main.rs`**. + +**Not implemented yet:** e.g. **`EngineConfig::ClickHouse`** is handled inside **`build_adapter`** with **`anyhow::bail!`** until a **`ClickHouseAdapter`** and **`try_from_cluster_config`** exist. + +- **`EngineType` for cluster state** — In **`main.rs`** and anywhere else (e.g. group member **`ClusterState`**), use **`EngineType::from(engine_config)`** from **`engine_registry.rs`**. **`queryflux-cluster-manager`** engine affinity uses the same **`From`** impl (see **`strategy.rs`**). + +- **Special rules** — Search for engine-specific checks (e.g. `queryAuth` / impersonation) and extend validation if your engine has constraints. + +### 4. Dispatch and frontends (`queryflux-frontend`) + +- Shared query execution goes through **`dispatch_query`** / **`execute_to_sink`**. Usually no change if the new engine only differs in the adapter; if you need a **special execution path** (like Trino raw HTTP), follow the existing engine-specific branches. +- Per-protocol handlers (Trino HTTP, Postgres wire, …) should keep using the shared dispatch layer unless the protocol requires a dedicated contract. + +### 5. Persistence (`queryflux-persistence`) — why touch it if config is JSON? + +The table stores **`engine_key` as its own column** plus a **`config` JSONB** blob. The DB does not load straight into the proxy as opaque JSON: code paths call **`ClusterConfigRecord::to_core()`**, which must produce a typed **`ClusterConfig`** (including **`EngineConfig`**). + +So persistence changes are **not** “because Postgres needs a JSON schema.” They are because of this **explicit conversion layer**: + +1. **`ClusterConfigRecord::to_core`** — Calls **`parse_engine_key`** from `queryflux-core` (next to `engine_key`). Extend **`parse_engine_key`** when you add an engine; you do **not** maintain a second duplicate string match in persistence. +2. **`UpsertClusterConfig::from_core`** — Uses **`engine_key(&EngineConfig)`** from core to set the `engine_key` column when seeding from YAML. + +**Extra JSON keys** that only live inside `config` and are **already** read in `to_core` (e.g. `endpoint`, `region`, `authType`, …) usually need **no** persistence change beyond the engine-key match. You only extend the `s("…")` / `b("…")` helpers in `to_core` (and the matching `from_core` inserts) if you add **new top-level persisted fields** on `ClusterConfig` that should round-trip through that JSON. + +**Hot reload** often uses `list_cluster_configs` → records → `to_core()` → `build_live_config`; the same conversion applies. + +### 6. Optional: routing config + +- If operators choose the new group via **router JSON** (`RouterConfig` variants), no change unless you add a new router **type**. +- **Protocol-based routing** maps frontend labels to **group names**; it does not list backend engines. + +### 7. Tests and docs + +- Add or extend **e2e** tests under `crates/queryflux-e2e-tests` if you have a dockerized target. +- Update [system-map.md](../system-map.md) component status if you document supported engines there. + +### 8. Suggested order of work (Rust) + +1. **`EngineConfig` / `EngineType`** + **`engine_key` / `parse_engine_key` / `From<&EngineConfig> for EngineType`** + **`dialect()`** if needed. +2. **`ClusterConfig`** fields if the engine needs new top-level keys (and persistence **`to_core`** JSON extraction if those keys live in JSONB). +3. Adapter module: **`EngineAdapterTrait`**, **`descriptor()`**, **`try_from_cluster_config`**. +4. **`registered_engines.rs`**: descriptor in **`all_descriptors()`**, arm in **`build_adapter`**. +5. Run **`cargo build -p queryflux`**; exercise **YAML** and **Postgres** load + **admin upsert** if applicable. + +--- + +## QueryFlux Studio (UI, TypeScript / React) + +Studio is the Next.js app under `ui/queryflux-studio/`. It does **not** run wire protocols; it calls the **Admin API** (`ADMIN_API_URL`, default `http://localhost:9000`) for clusters, groups, routing, and scripts. + +Backend engines are registered in Studio through **`StudioEngineModule`** objects: one file per engine under **`lib/studio-engines/engines/`**, aggregated in **`lib/studio-engines/manifest.ts`**. That manifest drives **`ENGINE_REGISTRY`**, catalog slots for implemented backends, optional flat-form validation, engine-affinity dropdown entries, and extra **`findEngineByType`** aliases. + +The proxy still exposes descriptors at **`GET /admin/engine-registry`**. Studio does **not** load that at runtime yet, so **Rust `descriptor()` and each studio module’s `descriptor` field must stay aligned by hand** (same `engineKey`, `configFields` keys, auth shapes, etc.). Shared TypeScript types live in **`lib/engine-registry-types.ts`**; **`lib/engine-registry.ts`** only re-exports helpers and builds **`ENGINE_REGISTRY`** from the manifest. + +### Where users see engines + +| User action | UI entrypoint | What must know your engine | +|-------------|---------------|----------------------------| +| Create cluster | **Clusters → Add cluster** (`components/add-cluster-dialog.tsx`) | Expanded **`ENGINE_CATALOG`** (includes studio slots) + **`findEngineDescriptor`** + **`validateClusterConfig`** / **`validateEngineSpecific`** + **`toUpsertBody`** | +| Edit cluster | **Clusters** grid → cluster card → Edit (`app/clusters/clusters-grid.tsx`) | Same + **`mergeClusterConfigFromFlat`** / **`buildClusterUpsertFromForm`** + **`EngineClusterConfig`** | +| View config | Cluster detail / engine config view in `clusters-grid.tsx` | **`findEngineDescriptor`** for labels; unknown key shows “add to engine registry” warning | +| Group strategy **engine affinity** | **Engines →** group dialog → strategy (`components/group-form-dialog.tsx`) | **`ENGINE_AFFINITY_OPTIONS`** is built by **`buildEngineAffinityOptionsFromManifest()`** from each module’s **`engineAffinity`** field (omit label override, or set **`engineAffinity: false`** to exclude, e.g. Athena). | +| Live utilization cards | **Engines (Groups)** page (`app/engines/page.tsx`) | **`findEngineByType`**; studio modules contribute aliases via **`extraTypeAliases`** (merged with static dialect aliases in **`components/engine-catalog.ts`**) | + +### 1. Studio engine module (primary registration) + +**Types:** `ui/queryflux-studio/lib/studio-engines/types.ts` — **`StudioEngineModule`**. + +**Per engine:** `ui/queryflux-studio/lib/studio-engines/engines/.ts` + +Export a constant (e.g. **`trinoStudioEngine`**) with: + +- **`descriptor`** — Full **`EngineDescriptor`** (must match Rust: **`engineKey`**, **`connectionType`**, **`supportedAuth`**, **`configFields`**, **`implemented`**, branding **`hex`**, etc.). Extend **`ConnectionType`** / **`AuthType`** in **`lib/engine-registry-types.ts`** if Rust added a variant. +- **`catalog`** — **`category`**, **`simpleIconSlug`**, **`catalogDescription`** for the engines grid / picker (display name and **`supported`** come from the descriptor when the catalog is expanded). +- **`validateFlat`** (optional) — Cross-field checks before save (e.g. Trino basic vs bearer). Dispatched by **`validateEngineSpecific`** in **`lib/studio-engines/validate-flat.ts`** (re-exported from **`lib/cluster-persist-form.ts`**). +- **`customFormId`** (optional) — String key; must match an entry in **`components/cluster-config/studio-engine-forms.tsx`** if the generic **`GenericEngineClusterConfig`** is not enough. +- **`engineAffinity`** (optional) — **`false`** to omit from affinity, or **`{ label?: string }`** for a custom dropdown label (default label is **`displayName`**). +- **`extraTypeAliases`** (optional) — Map of normalized API/type strings → canonical **`EngineDef.name`** for **`findEngineByType`** (e.g. alternate spellings). + +**Manifest:** `ui/queryflux-studio/lib/studio-engines/manifest.ts` + +- Import the new module and append it to **`STUDIO_ENGINE_MODULES`** (order affects **`ENGINE_AFFINITY_OPTIONS`** and registry iteration; catalog **card order** is separate — see below). + +**Derived registry:** `ui/queryflux-studio/lib/engine-registry.ts` + +- **`ENGINE_REGISTRY`** is **`STUDIO_ENGINE_MODULES.map((m) => m.descriptor)`**. Do not duplicate descriptor arrays here. +- **`findEngineDescriptor`**, **`implementedEngines`**, **`isClusterOnboardingSelectable`**, **`validateClusterConfig`** — unchanged behavior; **`validateClusterConfig`** still uses generic required-field checks from **`configFields`** unless you extend the Rust/TS contract. + +### 2. Catalog layout (picker order and dialect-only rows) + +**File:** `ui/queryflux-studio/components/engine-catalog.ts` + +- Implemented backends appear as **studio slots**: **`{ k: "studio", engineKey: "" }`** inside **`ENGINE_CATALOG_SLOTS`**, interleaved with static **`EngineDef`** rows (dialects with **`engineKey: null`**). +- At runtime, **`expandCatalog`** replaces each studio slot with **`studioModuleToEngineDef`** from **`lib/studio-engines/catalog.ts`**. +- Static **`STATIC_ENGINE_TYPE_ALIASES`** remains for dialects without a studio module; **`buildStudioTypeAliases()`** merges in per-module aliases and the lowercase **`engineKey`** → **`displayName`** mapping. + +**`isClusterOnboardingSelectable`** still requires a catalog row with **`supported`** and **`engineKey`**; for studio-backed engines, **`supported`** is **`descriptor.implemented`** after expansion. + +### 3. Cluster config forms + +**Router:** `ui/queryflux-studio/components/cluster-config/engine-cluster-config.tsx` + +- Resolves **`getStudioEngineModule(engineKey)`**; if **`customFormId`** is set and **`STUDIO_CUSTOM_CLUSTER_FORMS[id]`** exists, renders that component; otherwise **`GenericEngineClusterConfig`** (descriptor **`configFields`**). + +**Custom form registration:** `ui/queryflux-studio/components/cluster-config/studio-engine-forms.tsx` — map **`customFormId`** → component (see Trino / StarRocks / Athena). + +**Reference components:** `trino-cluster-config.tsx`, `starrocks-cluster-config.tsx`, `athena-cluster-config.tsx`, `generic-engine-cluster-config.tsx`, `config-field-row.tsx`. + +### 4. Persisted JSON ↔ flat form (create + edit save path) + +**File:** `ui/queryflux-studio/lib/cluster-persist-form.ts` + +- Still shared across engines. If **`cluster_configs.config`** gains **new top-level JSON keys**, update: + - **`MANAGED_CONFIG_JSON_KEYS`** + - **`persistedClusterConfigToFlat`**, **`flatToPersistedConfig`**, **`mergeClusterConfigFromFlat`** + - **`buildValidateShape`** (shape expected by **`validateClusterConfig`**) +- **`validateEngineSpecific`** is implemented in **`lib/studio-engines/validate-flat.ts`** (per-module **`validateFlat`**); this file re-exports it for call sites. + +### 5. Clusters page (grid, dialog, validation) + +**File:** `ui/queryflux-studio/app/clusters/clusters-grid.tsx` + +- Uses **`findEngineDescriptor`**, **`validateClusterConfig`**, **`validateEngineSpecific`**, **`buildValidateShape`**, **`skipImplementedCheck`** where needed. No per-engine branches beyond **`EngineClusterConfig`**. + +**File:** `ui/queryflux-studio/components/add-cluster-dialog.tsx` + +- Wires catalog → descriptor → **`EngineClusterConfig`** → **`toUpsertBody`** → **`upsertClusterConfig`**. + +### 6. Group strategy (engine affinity) + +**File:** `ui/queryflux-studio/lib/cluster-group-strategy.ts` + +- **`ENGINE_AFFINITY_OPTIONS`** = **`buildEngineAffinityOptionsFromManifest()`**. To exclude an engine, set **`engineAffinity: false`** on its **`StudioEngineModule`**. To customize the label, use **`engineAffinity: { label: "…" }`**. + +### 7. Display helpers + +**File:** `ui/queryflux-studio/lib/merge-clusters-display.ts` + +- Uses **`findEngineDescriptor(p.engineKey)`**; the descriptor must exist in the manifest. + +**File:** `ui/queryflux-studio/components/ui-helpers.tsx` (**`EngineBadge`**) + +- Uses **`ENGINE_CATALOG`**; studio-expanded rows must match **`displayName`** where badges key off names. + +**File:** `ui/queryflux-studio/components/engine-icon.tsx` + +- Consumes **`EngineDef`** (re-exported from **`engine-catalog.ts`**; types in **`lib/engine-catalog-types.ts`**). + +### 8. API types (usually unchanged) + +**File:** `ui/queryflux-studio/lib/api-types.ts` + +- **`ClusterConfigRecord`** / **`UpsertClusterConfig`** stay generic unless you add typed helpers. + +### 9. Optional: fetch registry from the proxy + +A follow-up could load **`GET /admin/engine-registry`** at runtime and hydrate forms from the API. Until then, keep **Rust `descriptor()`** and **`StudioEngineModule.descriptor`** in sync manually. + +### Studio checklist (copy-paste) + +- [ ] **`lib/studio-engines/engines/.ts`** — **`StudioEngineModule`** (`descriptor` aligned with Rust, **`catalog`**, optional **`validateFlat`**, **`customFormId`**, **`engineAffinity`**, **`extraTypeAliases`**) +- [ ] **`lib/studio-engines/manifest.ts`** — import + append to **`STUDIO_ENGINE_MODULES`** +- [ ] **`lib/engine-registry-types.ts`** — extend **`ConnectionType`** / **`AuthType`** if needed +- [ ] **`components/engine-catalog.ts`** — add **`{ k: "studio", engineKey: "…" }`** to **`ENGINE_CATALOG_SLOTS`** at the desired position +- [ ] **`components/cluster-config/studio-engine-forms.tsx`** — register component if **`customFormId`** is set +- [ ] **`lib/cluster-persist-form.ts`** — only if new persisted **`config`** JSON keys (managed keys + flat ↔ JSON + **`buildValidateShape`**) +- [ ] Smoke-test: Add cluster → save → edit → save; Engines page icons; group **engine affinity** if applicable + +--- + +## Checklist (backend) + +### Rust adapter + +- [ ] `EngineConfig` + `EngineType` + `engine_key()` + **`parse_engine_key()`** + **`From<&EngineConfig> for EngineType`** + dialect mapping (`engine_registry.rs` + `query.rs`) +- [ ] `EngineAdapterTrait` + `descriptor()` +- [ ] `registered_engines.rs`: **`all_descriptors()`** + **`build_adapter()`** arm calling **`try_from_cluster_config`** on the adapter +- [ ] Adapter module: **`try_from_cluster_config`** (or async equivalent) reading **`ClusterConfig`** +- [ ] `UpsertClusterConfig::from_core` / `to_core` stay aligned via **`engine_key` / `parse_engine_key`** (no extra string match in persistence) +- [ ] Translation / compatibility if dialect is new + +### Studio (UI) + +- [ ] `lib/studio-engines/engines/.ts` — `StudioEngineModule` (descriptor + catalog + options) +- [ ] `lib/studio-engines/manifest.ts` — register module in `STUDIO_ENGINE_MODULES` +- [ ] `lib/engine-registry-types.ts` — `ConnectionType` / `AuthType` if Rust added variants +- [ ] `components/engine-catalog.ts` — `{ k: "studio", engineKey }` slot in `ENGINE_CATALOG_SLOTS` +- [ ] `components/cluster-config/studio-engine-forms.tsx` — only if using `customFormId` +- [ ] `lib/cluster-persist-form.ts` — only if new `config` JSON keys need round-tripping +- [ ] Verify add-cluster + edit-cluster, Engines page icons / `findEngineByType`, and engine affinity if used + +--- + +## Related reading + +- [Extending QueryFlux — overview](overview.md) +- [Frontend](frontend.md) — new ingress protocols +- [query-translation.md](../query-translation.md) — Dialects and sqlglot +- [routing-and-clusters.md](../routing-and-clusters.md) — Routers and groups +- [observability.md](../observability.md) — Admin API + +**Rust files** + +- [`crates/queryflux/src/registered_engines.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux/src/registered_engines.rs) — `all_descriptors`, `build_adapter` +- [`crates/queryflux-core/src/engine_registry.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-core/src/engine_registry.rs) — `engine_key`, `parse_engine_key`, `EngineRegistry`, `From<&EngineConfig> for EngineType` +- [`crates/queryflux-persistence/src/cluster_config.rs`](https://github.com/lakeops-org/queryflux/blob/main/crates/queryflux-persistence/src/cluster_config.rs) — `to_core` / `from_core` vs `engine_key` + JSONB diff --git a/website/versioned_docs/version-0.1.0/architecture/adding-support/frontend.md b/website/versioned_docs/version-0.1.0/architecture/adding-support/frontend.md new file mode 100644 index 0000000..3b71dc7 --- /dev/null +++ b/website/versioned_docs/version-0.1.0/architecture/adding-support/frontend.md @@ -0,0 +1,43 @@ +--- +description: Adding a new frontend client protocol to QueryFlux — FrontendProtocol, listener, dispatch, and routing. +--- + +# Frontend + +Goal: clients speak a **wire or HTTP protocol to QueryFlux** (ingress), not a new backend. + +See **[Extending QueryFlux](overview.md)** for how this differs from a **backend** engine. Documentation for **existing** frontends (Trino HTTP, Postgres wire, MySQL wire, Flight SQL, Snowflake) lives under **[Frontends](../frontends/overview.md)**. + +### Where the code lives + +- **PostgreSQL wire:** `crates/queryflux-frontend/src/postgres_wire/` +- **MySQL wire:** `crates/queryflux-frontend/src/mysql_wire/` +- **Trino HTTP:** `crates/queryflux-frontend/src/trino_http/` +- **Flight SQL:** `crates/queryflux-frontend/src/flight_sql/` +- **Snowflake:** `crates/queryflux-frontend/src/snowflake/` + +### Typical steps for a new protocol + +1. **`FrontendProtocol`** — Already defined in `queryflux_core::query::FrontendProtocol`; add a variant only for a **new** ingress protocol. +2. **`default_dialect()`** — Set the sqlglot **source** dialect for translation (see [query-translation.md](../query-translation.md)). +3. **Listener** — Bind a port, parse the protocol, build **`SessionContext`** and **`InboundQuery`**, then call shared **`dispatch_query`** (or the same helpers Trino HTTP uses). +4. **Routing** — Optionally extend **protocol-based routing** in config / persisted routing so this frontend maps to the right default group. +5. **Tests** — Protocol-level tests or e2e clients as appropriate. + +Studio does **not** implement wire protocols; it only talks to the **Admin API** for config and metrics. + +--- + +## Checklist (frontend) + +- [ ] `FrontendProtocol` + dialect + listener module + dispatch integration + routing docs + +--- + +## Related reading + +- [Extending QueryFlux — overview](overview.md) +- [Backend](backend.md) — Rust adapter and Studio +- [Frontends](../frontends/overview.md) — Existing protocol listeners +- [query-translation.md](../query-translation.md) — Dialects and sqlglot +- [routing-and-clusters.md](../routing-and-clusters.md) — Routers and groups diff --git a/website/versioned_docs/version-0.1.0/architecture/adding-support/overview.md b/website/versioned_docs/version-0.1.0/architecture/adding-support/overview.md new file mode 100644 index 0000000..2a3f6d5 --- /dev/null +++ b/website/versioned_docs/version-0.1.0/architecture/adding-support/overview.md @@ -0,0 +1,32 @@ +--- +sidebar_label: Overview +description: How to extend QueryFlux — new backend engines (Rust + Studio) vs new frontend client protocols, with links to detailed guides. +--- + +# Extending QueryFlux + +This guide separates two ideas that are easy to conflate: + +| Concept | Meaning | Example | +|--------|---------|---------| +| **Backend engine** | A **cluster** type QueryFlux routes queries **to**. It has an adapter that talks to the real database (HTTP, MySQL wire, embedded library, AWS SDK, …). | Trino, DuckDB, StarRocks, Athena | +| **Frontend protocol** | How **clients connect to QueryFlux** (ingress). SQL enters with a `FrontendProtocol` and a default source dialect for translation. | Trino HTTP, PostgreSQL wire, MySQL wire, Flight SQL | + +Adding **PostgreSQL wire** as a client entrypoint is **not** the same as adding “PostgreSQL” as a backend: today, `PostgresWire` is already a frontend in `queryflux-frontend`; traffic still lands on the shared dispatch path and is sent to whatever **backend adapter** routing chose (often Trino). + +## Guides + +| Page | What it covers | +|------|----------------| +| [Backend](backend.md) | Rust adapter (`EngineAdapterTrait`), `registered_engines`, persistence, dispatch notes — plus **QueryFlux Studio** (`StudioEngineModule`, catalog, forms). | +| [Frontend](frontend.md) | Adding a **new** ingress protocol (`FrontendProtocol`, listener, routing). For existing protocols, see **[Frontends](../frontends/overview.md)**. | + +--- + +## Related reading + +- [Frontends](../frontends/overview.md) — Trino HTTP, Postgres wire, MySQL wire, Flight SQL, Snowflake +- [system-map.md](../system-map.md) — End-to-end flow +- [query-translation.md](../query-translation.md) — Dialects and sqlglot +- [routing-and-clusters.md](../routing-and-clusters.md) — Routers and groups +- [observability.md](../observability.md) — Admin API (including engine registry JSON) diff --git a/website/versioned_docs/version-0.1.0/architecture/frontends/flight-sql.md b/website/versioned_docs/version-0.1.0/architecture/frontends/flight-sql.md new file mode 100644 index 0000000..818b43a --- /dev/null +++ b/website/versioned_docs/version-0.1.0/architecture/frontends/flight-sql.md @@ -0,0 +1,71 @@ +--- +description: Arrow Flight SQL frontend — gRPC service, implemented RPCs, metadata passthrough, and connecting with Flight SQL clients. +--- + +# Arrow Flight SQL + +The Flight SQL frontend exposes a gRPC-based [Arrow Flight SQL](https://arrow.apache.org/docs/format/FlightSql.html) service. Clients that speak Flight SQL (e.g. the JDBC Flight SQL driver, ADBC, DuckDB's `ATTACH` over Flight) can connect to QueryFlux and run SQL queries with results streamed as native Arrow record batches — zero serialization overhead for columnar consumers. + +## Configuration + +```yaml +queryflux: + frontends: + flightSql: + enabled: true + port: 50051 +``` + +Config key: `flightSql`. Protocol identifier: `FrontendProtocol::FlightSql`. Default dialect: `SqlDialect::Generic`. + +## Implemented RPCs + +| RPC | Status | Description | +|-----|--------|-------------| +| `GetFlightInfo` (statement) | Implemented | Accepts SQL via `CommandStatementQuery`, returns a `FlightInfo` with a ticket | +| `DoGet` (statement) | Implemented | Decodes SQL from the ticket, executes the query, and streams Arrow `RecordBatch` results | +| All other Flight SQL RPCs | Unimplemented | Return gRPC `Unimplemented` status | + +The two-step flow: + +1. **`GetFlightInfo`** — client sends SQL in a `CommandStatementQuery`. QueryFlux returns a `FlightInfo` containing an endpoint with a ticket (the SQL is encoded in the ticket's `statement_handle`). The IPC schema in the response is empty at this stage. +2. **`DoGet`** — client presents the ticket. QueryFlux decodes the SQL, authenticates, routes, executes via `execute_to_sink`, and streams Arrow record batches through a `FlightDataEncoder`. + +## Authentication + +Credentials are read from gRPC metadata: + +- **`authorization`** — `Bearer ` or `Basic `, depending on your configured auth provider. + +## Execution model + +Queries execute **synchronously** via `execute_to_sink`. Results are streamed as Arrow Flight data frames directly from the query's Arrow record batches — no intermediate serialization to JSON or text. + +## Metadata and routing + +ASCII metadata entries on the gRPC request (beyond `authorization`) are collected into a key/value map and passed through to the same routing and session path used for HTTP-style frontends, so routers and auth can inspect client-supplied fields consistently across protocols. + +Query tags are **not** populated from Flight SQL metadata in the current implementation. + +## Client examples + +```python +# Python (ADBC Flight SQL driver) +import adbc_driver_flightsql.dbapi as flight_sql + +conn = flight_sql.connect(uri="grpc://localhost:50051") +cur = conn.cursor() +cur.execute("SELECT 42 AS answer") +print(cur.fetchone()) +``` + +```bash +# DuckDB (if Flight SQL extension is available) +ATTACH 'grpc://localhost:50051' AS qf (TYPE flight_sql); +SELECT * FROM qf.my_catalog.my_schema.my_table LIMIT 10; +``` + +## Related + +- [Frontends overview](overview.md) — shared dispatch and session model +- [Routing and clusters](/docs/architecture/routing-and-clusters) — `protocolBased` router with `flightSql` diff --git a/website/versioned_docs/version-0.1.0/architecture/frontends/mysql-wire.md b/website/versioned_docs/version-0.1.0/architecture/frontends/mysql-wire.md new file mode 100644 index 0000000..566f3f1 --- /dev/null +++ b/website/versioned_docs/version-0.1.0/architecture/frontends/mysql-wire.md @@ -0,0 +1,104 @@ +--- +description: MySQL wire protocol frontend — handshake, COM_QUERY, session variables, schema switching, and connecting with mysql clients. +--- + +# MySQL wire + +The MySQL wire frontend lets standard MySQL clients (`mysql` CLI, JDBC, Python `mysql-connector`, Go `go-sql-driver/mysql`, etc.) connect to QueryFlux over the MySQL wire protocol. This is the natural entry point when routing traffic to MySQL-protocol backends like StarRocks. + +## Configuration + +```yaml +queryflux: + frontends: + mysqlWire: + enabled: true + port: 3306 +``` + +Config key: `mysqlWire`. Protocol identifier: `FrontendProtocol::MySqlWire`. Default dialect: `SqlDialect::MySql`. + +## Protocol support + +| Feature | Status | +|---------|--------| +| Server handshake (protocol 10) | Supported — advertises `8.0.0-queryflux` | +| `mysql_native_password` auth | Supported | +| `COM_QUERY` | Supported | +| `COM_PING` | Supported | +| `COM_INIT_DB` (schema switch) | Supported | +| `COM_QUIT` | Supported | +| `COM_FIELD_LIST` | Stub (returns empty EOF) | +| SSL/TLS handshake | Detected and rejected — use `--ssl-mode=DISABLED` | +| Prepared statements (`COM_STMT_*`) | Not supported | + +Clients must connect without TLS (`--ssl-mode=DISABLED` for `mysql` CLI, `useSSL=false` for JDBC). If the client sends an SSL request, QueryFlux returns an error and closes the connection. + +## Handshake and authentication + +1. QueryFlux sends a server handshake packet (protocol version 10, server version `8.0.0-queryflux`, `mysql_native_password`). +2. Client responds with username and optional database. +3. QueryFlux authenticates via the configured `auth_provider`. +4. On success, an OK packet is returned and the connection is ready. + +## Execution model + +Queries execute **synchronously** via `execute_to_sink`. Results stream as MySQL text protocol result sets: + +1. Column count packet. +2. Column definition packets. +3. Row data packets (text values). +4. EOF / OK packet. + +## Built-in query handling + +The MySQL frontend intercepts several common queries that clients and drivers send during connection setup: + +| Query pattern | Behavior | +|---------------|----------| +| `SET query_tags = '...'` / `SET SESSION query_tags = '...'` | Stores tags on the session for routing | +| `SET ...` (other) | Acknowledged without forwarding to backend | +| `USE ` | Updates session schema | +| `SELECT @@version` | Returns `8.0.0-queryflux` | +| `SELECT DATABASE()` | Returns current session schema | +| `SHOW VARIABLES` / `SHOW STATUS` | Returns empty result set | + +Leading `/* ... */` and `/*!...*/` conditional comments in SQL are stripped before processing. + +## Session context + +`SessionContext::MySqlWire` carries: + +| Field | Source | +|-------|--------| +| `user` | Handshake response | +| `schema` | `COM_INIT_DB` or initial database from handshake | +| `session_vars` | Currently empty (generic `SET` is not stored) | +| `tags` | From `SET query_tags` / `SET SESSION query_tags` | + +The `schema` and `user` fields are available to routers — the `pythonScript` router receives them in `ctx["schema"]` and `ctx["user"]`. + +## Client examples + +```bash +# mysql CLI +mysql -h 127.0.0.1 -P 3306 -u dev --ssl-mode=DISABLED -e "SELECT 1" + +# With database +mysql -h 127.0.0.1 -P 3306 -u dev -D my_catalog --ssl-mode=DISABLED +``` + +```python +# Python (mysql-connector) +import mysql.connector +conn = mysql.connector.connect(host="127.0.0.1", port=3306, user="dev", ssl_disabled=True) +cur = conn.cursor() +cur.execute("SELECT 42 AS answer") +print(cur.fetchone()) +``` + +## Related + +- [Frontends overview](overview.md) — shared dispatch and session model +- [Routing and clusters](/docs/architecture/routing-and-clusters) — `protocolBased` router with `mysqlWire` +- [Query tags](/docs/architecture/query-tags) — setting tags via `SET query_tags` diff --git a/website/versioned_docs/version-0.1.0/architecture/frontends/overview.md b/website/versioned_docs/version-0.1.0/architecture/frontends/overview.md new file mode 100644 index 0000000..a18c2e5 --- /dev/null +++ b/website/versioned_docs/version-0.1.0/architecture/frontends/overview.md @@ -0,0 +1,115 @@ +--- +sidebar_label: Overview +description: How QueryFlux frontends work — protocol listeners, shared dispatch, session context, and the available client protocols. +--- + +# Frontends + +A **frontend** is the entry point for client traffic into QueryFlux. Each frontend speaks a specific wire or HTTP protocol, parses incoming SQL, builds a `SessionContext`, and hands the query to the shared dispatch layer. The client never knows which backend engine actually runs the query — the frontend translates results back into its native format before responding. + +## Available frontends + +| Frontend | Config key | Default port | Protocol | Dialect | Status | +|----------|------------|-------------|----------|---------|--------| +| [Trino HTTP](trino-http.md) | `trinoHttp` | 8080 | HTTP REST (JSON) | Trino | **Done** | +| [PostgreSQL wire](postgres-wire.md) | `postgresWire` | 5432 | PostgreSQL v3 wire | Postgres | **Done** | +| [MySQL wire](mysql-wire.md) | `mysqlWire` | 3306 | MySQL wire | MySQL | **Done** | +| [Arrow Flight SQL](flight-sql.md) | `flightSql` | 50051 | gRPC (Arrow Flight) | Generic | **Done** | +| [Snowflake](snowflake.md) | `snowflakeHttp` | 8445 | HTTP REST (JSON) | Generic | **Done** | +| ClickHouse HTTP | `clickhouseHttp` | 8123 | HTTP | ClickHouse | Planned | + +## Shared architecture + +All frontends converge on the same internal pipeline. The differences are only in how SQL enters and how results leave. + +``` +Client ──(native protocol)──► Frontend Listener + │ + SessionContext + SQL + │ + ▼ + Router Chain ──► ClusterGroupName + │ + ▼ + ClusterGroupManager ──► ClusterName + │ + ▼ + Translation Service ──► translated SQL + │ + ▼ + Engine Adapter ──► results + │ + ▼ + ResultSink ──► native protocol response +``` + +### Dispatch paths + +The dispatch layer offers two execution models: + +| Path | Used by | Behavior | +|------|---------|----------| +| **`dispatch_query`** | Trino HTTP (async-capable groups) | Submit to engine, persist handle, return polling URL. Client follows `nextUri` to stream pages. | +| **`execute_to_sink`** | All other frontends + Trino HTTP sync fallback | Wait for cluster capacity (backoff), execute query to completion, stream Arrow batches through a `ResultSink` that encodes the native protocol response. | + +### SessionContext + +Each frontend builds a protocol-specific `SessionContext` that travels with the query through routing and into dispatch: + +| Variant | Fields | +|---------|--------| +| `TrinoHttp` | `headers` (lowercased HTTP headers) | +| `PostgresWire` | `user`, `database`, `session_params` | +| `MySqlWire` | `user`, `schema`, `session_vars` | +| `ClickHouseHttp` | `headers`, `query_params` | +| `FlightSql` | `headers` (from gRPC metadata) | + +Snowflake frontends build a `MySqlWire` session context internally (user from session, schema from database hint). + +### Authentication + +All frontends extract credentials from their protocol's native mechanism (HTTP headers, wire auth packets, gRPC metadata) and pass them to the shared `auth_provider`. The auth result determines whether the query proceeds and provides context for authorization checks downstream. + +### Routing + +The `FrontendProtocol` enum identifies which frontend originated a query. The `protocolBased` router uses this to map traffic from different protocols to different cluster groups: + +```yaml +routers: + - type: protocolBased + trinoHttp: trino-default + postgresWire: trino-default + mysqlWire: starrocks-group + snowflakeHttp: trino-default + snowflakeSqlApi: analytics-group +``` + +### SQL dialect and translation + +Each frontend has a **default source dialect** (`protocol.default_dialect()`). When the target engine's dialect differs, sqlglot translates the SQL automatically. When dialects match, translation is skipped entirely. + +## Configuration + +Each frontend is enabled under `queryflux.frontends` in `config.yaml`: + +```yaml +queryflux: + frontends: + trinoHttp: + enabled: true + port: 8080 + postgresWire: + enabled: true + port: 5432 + mysqlWire: + enabled: true + port: 3306 + flightSql: + enabled: true + port: 50051 + snowflakeHttp: + enabled: true + port: 8445 +``` + +Omitting a frontend block or setting `enabled: false` disables that listener entirely. diff --git a/website/versioned_docs/version-0.1.0/architecture/frontends/postgres-wire.md b/website/versioned_docs/version-0.1.0/architecture/frontends/postgres-wire.md new file mode 100644 index 0000000..4ee4611 --- /dev/null +++ b/website/versioned_docs/version-0.1.0/architecture/frontends/postgres-wire.md @@ -0,0 +1,84 @@ +--- +description: PostgreSQL wire protocol frontend — startup, simple query, session context, and connecting with psql or any Postgres driver. +--- + +# PostgreSQL wire + +The PostgreSQL wire frontend lets standard Postgres clients (`psql`, JDBC, Python `psycopg2`, Go `pgx`, etc.) connect to QueryFlux over the PostgreSQL v3 wire protocol. Queries are executed synchronously — the TCP connection stays open while the result streams back as standard Postgres wire messages. + +## Configuration + +```yaml +queryflux: + frontends: + postgresWire: + enabled: true + port: 5432 +``` + +Config key: `postgresWire`. Protocol identifier: `FrontendProtocol::PostgresWire`. Default dialect: `SqlDialect::Postgres`. + +## Protocol support + +| Feature | Status | +|---------|--------| +| Startup (protocol version 3.0) | Supported | +| Simple query (`Q` message) | Supported | +| Extended query (Parse/Bind/Execute) | Not supported — returns error | +| SSL negotiation | Declined (`N` response) — plaintext only | +| `COPY` protocol | Not supported | + +QueryFlux responds to the SSL request with `N` (no SSL), then processes the normal startup sequence. Clients that require TLS will need to connect without it or use an external TLS terminator. + +## Startup and authentication + +1. Client sends **StartupMessage** with `user`, `database`, and optional parameters. +2. QueryFlux extracts credentials from the startup `user` field and authenticates via the configured `auth_provider`. +3. On success, QueryFlux sends `AuthenticationOk`, followed by `ParameterStatus` messages (server version `16.0-queryflux`, encoding `UTF8`, etc.), `BackendKeyData`, and `ReadyForQuery`. + +## Execution model + +All queries execute **synchronously** via `execute_to_sink`. Results stream as standard Postgres wire messages: + +1. `RowDescription` — column metadata (names, types, format codes). +2. `DataRow` messages — one per row, text-format values. +3. `CommandComplete` — summary (e.g. `SELECT 3`). +4. `ReadyForQuery` — ready for the next query. + +Errors are returned as Postgres `ErrorResponse` messages with SQLSTATE codes. + +## Session context + +`SessionContext::PostgresWire` carries: + +| Field | Source | +|-------|--------| +| `user` | Startup message `user` parameter | +| `database` | Startup message `database` parameter | +| `session_params` | Initially empty | +| `tags` | `query_tags` / `query_tag` from startup parameters | + +The `database` and `user` fields are available to routers — the `pythonScript` router receives them in `ctx["database"]` and `ctx["user"]`. + +## SET handling + +`SET` statements are acknowledged with `CommandComplete` + `ReadyForQuery` without forwarding to the backend. This keeps the connection in a valid state for clients that issue `SET` during startup. + +## Client examples + +```bash +# psql +psql -h localhost -p 5432 -U dev -d my_catalog -c "SELECT 1" + +# Python (psycopg2) +import psycopg2 +conn = psycopg2.connect(host="localhost", port=5432, user="dev", dbname="my_catalog") +cur = conn.cursor() +cur.execute("SELECT 42 AS answer") +print(cur.fetchone()) +``` + +## Related + +- [Frontends overview](overview.md) — shared dispatch and session model +- [Routing and clusters](/docs/architecture/routing-and-clusters) — `protocolBased` router with `postgresWire` diff --git a/website/versioned_docs/version-0.1.0/architecture/frontends/snowflake.md b/website/versioned_docs/version-0.1.0/architecture/frontends/snowflake.md new file mode 100644 index 0000000..45c6775 --- /dev/null +++ b/website/versioned_docs/version-0.1.0/architecture/frontends/snowflake.md @@ -0,0 +1,191 @@ +--- +description: Snowflake-compatible frontend — HTTP wire protocol and SQL API v2, session management, query execution, and client connectivity. +--- + +# Snowflake + +QueryFlux can accept connections from **Snowflake clients** — JDBC, ODBC, Python connector, Go driver, Node.js driver, and the SQL API — without any driver changes on the client side. Traffic is **not** proxied to a real Snowflake account; QueryFlux terminates the Snowflake protocol locally, authenticates via its own auth providers, routes through the standard router chain, and executes queries on whichever backend engine routing selects (Trino, StarRocks, DuckDB, etc.). + +This is a **protocol bridge**, not a Snowflake proxy: the client believes it is talking to Snowflake, but the query runs on a different engine entirely. + +## Configuration + +```yaml +queryflux: + frontends: + snowflakeHttp: + enabled: true + port: 8445 +``` + +Config key: `snowflakeHttp`. Protocol identifiers: `FrontendProtocol::SnowflakeHttp` (wire) and `FrontendProtocol::SnowflakeSqlApi` (SQL API). Default dialect: `SqlDialect::Generic` (both). + +The `snowflakeHttp` frontend starts a single HTTP listener that serves **both** wire and SQL API routes. There is no separate port for `snowflakeSqlApi` — both protocol surfaces are merged onto the same listener. + +## Two protocol flavors + +The Snowflake frontend exposes two API surfaces on a **single port**: + +| Protocol | Identifier | Typical clients | Auth model | +|----------|------------|-----------------|------------| +| **Snowflake HTTP wire** | `snowflakeHttp` | JDBC, ODBC, Python connector, Go driver, Node.js driver | Session-based (login → token) | +| **Snowflake SQL API v2** | `snowflakeSqlApi` | REST clients, service accounts, `curl` | Stateless Bearer token per request | + +Both share the same listener; routing can target them independently via `protocolBased` rules: + +```yaml +routers: + - type: protocolBased + snowflakeHttp: trino-default + snowflakeSqlApi: analytics-group +``` + +--- + +## HTTP wire protocol + +The wire protocol handles session lifecycle (login, logout, heartbeat, token refresh) and synchronous query execution. This is the API surface that JDBC, ODBC, and the Python/Go/Node connectors use internally. + +### Endpoints + +| Method | Path | Description | +|--------|------|-------------| +| `POST` | `/session/v1/login-request` | Authenticate and create a session | +| `DELETE` | `/session` | Destroy a session (logout) | +| `GET` | `/session/heartbeat` | Validate that a session token is still alive | +| `POST` | `/session/token-request` | Refresh / validate a session token | +| `POST` | `/queries/v1/query-request` | Execute a SQL query | +| `GET` | `/queries/v1/query-monitoring-request` | Query monitoring (returns empty — sync execution) | +| `DELETE` | `/queries/v1/{query_id}` | Cancel a query (no-op — sync execution) | + +### Login + +`POST /session/v1/login-request` accepts JSON with `LOGIN_NAME` and `PASSWORD` fields. Optional `databaseName` and `schemaName` can be passed as query parameters or in `SESSION_PARAMETERS`. + +On success, QueryFlux: + +1. Authenticates via the configured `auth_provider`. +2. Resolves a **cluster group** by routing with `FrontendProtocol::SnowflakeHttp`. +3. Issues a QueryFlux-generated UUID as both `token` and `masterToken`. +4. Stores the session (auth context, group, database/schema hints, user) in an in-memory session store. + +The returned token is used in subsequent requests via the `Authorization: Snowflake Token=""` header. + +### Query execution + +`POST /queries/v1/query-request` reads the SQL from the `sqlText` JSON field. Gzip-compressed request bodies (`Content-Encoding: gzip`) are supported — the Python connector uses this by default. + +Execution is **synchronous**: the query runs to completion via `execute_to_sink`, and the full result set is returned in a single response. The response includes: + +- **Snowflake-compatible JSON** with `rowType` metadata (column names, types, nullability). +- **Base64-encoded Arrow IPC** in `rowsetBase64` for the Python connector's Arrow result path. +- Standard Snowflake response fields (`success`, `code`, `message`, `data`). + +Errors are returned as **HTTP 200** with `"success": false` in the JSON body — matching Snowflake's convention so connectors handle errors correctly rather than treating non-200 responses as retryable infrastructure failures. + +--- + +## SQL API v2 + +The SQL API follows Snowflake's [REST SQL API](https://docs.snowflake.com/en/developer-guide/sql-api) convention. This is the stateless, request-per-query interface typically used by service accounts and automation. + +### Endpoints + +| Method | Path | Description | +|--------|------|-------------| +| `POST` | `/api/v2/statements` | Submit a SQL statement | +| `GET` | `/api/v2/statements/{handle}` | Get statement status/results | +| `DELETE` | `/api/v2/statements/{handle}` | Cancel a statement | + +### Authentication + +SQL API requests use `Authorization: Bearer `. The token is validated against QueryFlux's `auth_provider` on every request — there is no session to maintain. + +### Submit statement + +`POST /api/v2/statements` accepts JSON with a `statement` field containing the SQL. QueryFlux authenticates the request, routes with `FrontendProtocol::SnowflakeSqlApi`, and executes synchronously. + +The response is a Snowflake SQL API v2-compatible JSON object with: + +- `statementHandle` — a unique identifier for the execution. +- `status` — `"00000"` on success. +- `rowType` — column metadata. +- `data` — rows as arrays of JSON string values. + +### Get / cancel statement + +Since execution is synchronous, `GET /api/v2/statements/{handle}` returns a **404** (no stored handles) and `DELETE /api/v2/statements/{handle}` returns a success stub. + +--- + +## Session management + +The HTTP wire protocol maintains sessions in an in-memory `SnowflakeSessionStore` (`DashMap`). Each session holds: + +| Field | Description | +|-------|-------------| +| `token` | QueryFlux-issued UUID (the session key) | +| `auth_ctx` | Authentication context from the auth provider | +| `group` | Cluster group resolved at login time | +| `user` | Authenticated username | +| `database` / `schema` | Optional hints from the login request | +| `created_at` | Session creation timestamp | + +Wire protocol requests resolve the session from the `Snowflake Token=` header. The cluster group is fixed at login — all queries in a session route to the same group. + +The SQL API does **not** use the session store; each request is independently authenticated and routed. + +--- + +## Connecting clients + +### Python connector + +```python +import snowflake.connector + +conn = snowflake.connector.connect( + user="dev", + password="password", + account="queryflux", # any value — not used for routing + host="localhost", + port=8445, + protocol="http", + database="my_catalog", + schema="my_schema", +) + +cur = conn.cursor() +cur.execute("SELECT 1 AS num, 'hello' AS greeting") +for row in cur: + print(row) +``` + +The Python connector sends requests to the wire protocol endpoints. Set `host` and `port` to point at the QueryFlux Snowflake listener. The `account` value is required by the driver but not used by QueryFlux for routing. + +### JDBC + +``` +jdbc:snowflake://localhost:8445/?account=queryflux&ssl=off&db=my_catalog&schema=my_schema +``` + +### SQL API via curl + +```bash +curl -X POST http://localhost:8445/api/v2/statements \ + -H "Authorization: Bearer " \ + -H "Content-Type: application/json" \ + -d '{"statement": "SELECT 42 AS answer"}' +``` + +--- + +## Dialect and translation + +Both Snowflake protocol flavors map to `SqlDialect::Generic` as their default source dialect. When routing lands on a Trino cluster, for example, sqlglot translates the SQL from the generic dialect to Trino. If the target engine's dialect matches, translation is skipped. + +## Related + +- [Frontends overview](overview.md) — shared dispatch and session model +- [Configuration](/docs/configuration) — `frontends.snowflakeHttp` in `config.yaml` +- [Routing and clusters](/docs/architecture/routing-and-clusters) — `protocolBased` router with `snowflakeHttp` / `snowflakeSqlApi` diff --git a/website/versioned_docs/version-0.1.0/architecture/frontends/trino-http.md b/website/versioned_docs/version-0.1.0/architecture/frontends/trino-http.md new file mode 100644 index 0000000..ad8acb6 --- /dev/null +++ b/website/versioned_docs/version-0.1.0/architecture/frontends/trino-http.md @@ -0,0 +1,84 @@ +--- +description: Trino HTTP frontend — async query polling, nextUri rewriting, session headers, tags, and mixed-engine group support. +--- + +# Trino HTTP + +The Trino HTTP frontend lets any Trino-compatible client (Trino CLI, JDBC, Python `trino`, DBeaver, etc.) connect to QueryFlux as if it were a Trino coordinator. QueryFlux accepts the standard `/v1/statement` API, routes the query, and returns Trino-shaped JSON responses — including `nextUri` polling for async engines. + +## Configuration + +```yaml +queryflux: + frontends: + trinoHttp: + enabled: true + port: 8080 +``` + +Config key: `trinoHttp`. Protocol identifier: `FrontendProtocol::TrinoHttp`. Default dialect: `SqlDialect::Trino`. + +## Endpoints + +| Method | Path | Description | +|--------|------|-------------| +| `POST` | `/v1/statement` | Submit a new SQL query | +| `GET` | `/v1/statement/qf/queued/{id}/{seq}` | Poll a queued query (proxy-side backoff) | +| `GET` | `/v1/statement/qf/executing/{id}` | Poll an executing query | +| `GET` | `/v1/statement/{*path}` | Forward Trino-style poll paths | +| `DELETE` | `/v1/statement/qf/executing/{id}` | Cancel a running query | +| `DELETE` | `/v1/statement/{*path}` | Cancel via forwarded Trino path | + +## Authentication + +Credentials are extracted in this order: + +1. `Authorization: Basic` — decoded username and password. +2. `Authorization: Bearer` — bearer token. +3. `X-Trino-User` header — username only (no password). + +The extracted credentials are passed to the configured `auth_provider`. Failures return HTTP 401 (unauthorized) or 403 (forbidden by authorization policy). + +## Execution model + +Trino HTTP is the only frontend that supports **async polling**: + +- **Async-capable group** (e.g. Trino backend): `dispatch_query` submits to the engine, persists the executing state, and returns a Trino JSON response. The `nextUri` in the response is **rewritten** to point back at QueryFlux (`externalAddress`), so subsequent polls flow through the proxy transparently. +- **At capacity**: when the group is full, QueryFlux returns a synthetic "queued" response with a `nextUri` pointing at `/v1/statement/qf/queued/{id}/{seq}`. The client polls this URL; QueryFlux retries cluster acquisition on each poll. +- **Sync engine in a Trino group** (e.g. DuckDB, StarRocks): falls back to `execute_to_sink` with a `TrinoHttpResultSink` — the query runs to completion and the result is returned as a single Trino JSON page (no polling). + +## Session context + +`SessionContext::TrinoHttp` carries all request headers (lowercased keys) as a `HashMap`. Routers and the Python script router can inspect any header (e.g. `x-trino-user`, `x-trino-catalog`, `x-trino-schema`). + +## Query tags + +Tags can be set via: + +- `X-Trino-Client-Tags` header — comma-separated key/value pairs. +- `X-Trino-Session` header — `query_tags` or `query_tag` keys (percent-decoded). +- `SET SESSION query_tags = '...'` SQL statement — intercepted by the frontend (not forwarded to the backend). The response includes `X-Trino-Set-Session` so the client tracks the change. + +Tags are used by the `tags` router for routing decisions. See [Query tags](/docs/architecture/query-tags). + +## nextUri rewriting + +When proxying Trino async queries, QueryFlux rewrites `nextUri` URLs in response JSON so the client always polls through the proxy rather than going directly to the Trino coordinator. The rewrite replaces the host/port with `externalAddress` while preserving the Trino path from `/v1/` onward. This is done via fast byte-level patching when possible, avoiding full JSON deserialization. + +## Client examples + +```bash +# Trino CLI +trino --server http://localhost:8080 --execute "SELECT 42" + +# curl +curl -X POST http://localhost:8080/v1/statement \ + -H "X-Trino-User: dev" \ + -d "SELECT current_date" +``` + +## Related + +- [Frontends overview](overview.md) — shared dispatch and session model +- [Routing and clusters](/docs/architecture/routing-and-clusters) — how `protocolBased` maps `trinoHttp` to a group +- [Query tags](/docs/architecture/query-tags) — tag-based routing from Trino sessions diff --git a/website/versioned_docs/version-0.1.0/architecture/overview.md b/website/versioned_docs/version-0.1.0/architecture/overview.md index 098fcc8..86e806b 100644 --- a/website/versioned_docs/version-0.1.0/architecture/overview.md +++ b/website/versioned_docs/version-0.1.0/architecture/overview.md @@ -14,7 +14,8 @@ This section describes how QueryFlux is put together: why it exists, how SQL is | [Query translation](query-translation.md) | Dialect detection, sqlglot integration, when translation runs, and schema-aware mode. | | [Routing and clusters](routing-and-clusters.md) | Router chain, `routingFallback`, cluster groups, load-balancing strategies, and queueing. | | [Observability](observability.md) | Prometheus metrics, Grafana dashboard, QueryFlux Studio, and the Admin REST API. | -| [Adding engine support](adding-engine-support.md) | Checklist for new **backend engines** (Rust adapters), **Studio** (`lib/studio-engines/` manifest + catalog slots), and **frontend wire protocols** (e.g. Postgres wire). | +| [Frontends](frontends/overview.md) | Protocol listeners — Trino HTTP, PostgreSQL wire, MySQL wire, Flight SQL, and Snowflake. Shared dispatch, session model, and per-protocol details. | +| [Extending QueryFlux](adding-support/overview.md) | **[Backend](adding-support/backend.md)** (Rust + Studio) and **[Frontend](adding-support/frontend.md)** (new protocols). Same sidebar category as **[Frontends](frontends/overview.md)**. | | [Auth / authz design](auth-authz-design.md) | Authentication and authorization design notes. | Start with [Motivation and goals](motivation-and-goals.md) if you are new to the project; use [System map](system-map.md) as the single-page map of the system. diff --git a/website/versioned_docs/version-0.1.0/architecture/routing-and-clusters.md b/website/versioned_docs/version-0.1.0/architecture/routing-and-clusters.md index 8c1a02b..880765b 100644 --- a/website/versioned_docs/version-0.1.0/architecture/routing-and-clusters.md +++ b/website/versioned_docs/version-0.1.0/architecture/routing-and-clusters.md @@ -34,7 +34,7 @@ Configured under `routers:` in YAML (`queryflux_core::config::RouterConfig`). Wi | `type` | Behavior | |--------|----------| -| `protocolBased` | Maps the active frontend (`trinoHttp`, `postgresWire`, `mysqlWire`, `flightSql`, `clickhouseHttp`) to a group name. | +| `protocolBased` | Maps the active frontend (`trinoHttp`, `postgresWire`, `mysqlWire`, `flightSql`, `clickhouseHttp`, `snowflakeHttp`, `snowflakeSqlApi`) to a group name. | | `header` | Matches a header value to a group (useful for Trino HTTP and similar). | | `queryRegex` | Ordered rules: first regex match on the SQL text wins. | | `tags` | Routes based on query tags attached to the session. Each rule specifies one or more tag key/value conditions (AND logic); first matching rule wins. See [Tags router](#tags-router-tags) below. | diff --git a/website/versioned_docs/version-0.1.0/architecture/system-map.md b/website/versioned_docs/version-0.1.0/architecture/system-map.md index 6719b19..6eea3d8 100644 --- a/website/versioned_docs/version-0.1.0/architecture/system-map.md +++ b/website/versioned_docs/version-0.1.0/architecture/system-map.md @@ -4,9 +4,9 @@ description: End-to-end query lifecycle, major crates, and component status (hig # QueryFlux — Architecture Overview -QueryFlux is a universal SQL query proxy and router. It accepts queries from clients over multiple protocols (Trino HTTP, PostgreSQL wire, MySQL wire, Arrow Flight SQL), routes them to the appropriate backend engine, optionally translates the SQL dialect, and streams results back in the client's native format. +QueryFlux is a universal SQL query proxy and router. It accepts queries from clients over multiple protocols (Trino HTTP, PostgreSQL wire, MySQL wire, Arrow Flight SQL, Snowflake HTTP), routes them to the appropriate backend engine, optionally translates the SQL dialect, and streams results back in the client's native format. -**More documentation:** the [architecture documentation overview](./overview.md) indexes deeper topics — [motivation-and-goals.md](motivation-and-goals.md) (why the project exists), [query-translation.md](query-translation.md) (sqlglot and dialects), [routing-and-clusters.md](routing-and-clusters.md) (routers, groups, load balancing), [observability.md](observability.md) (Prometheus, Grafana, Studio, Admin API), [adding-engine-support.md](adding-engine-support.md) (new engines, Studio, and client protocols). +**More documentation:** the [architecture documentation overview](./overview.md) indexes deeper topics — [motivation-and-goals.md](motivation-and-goals.md) (why the project exists), [query-translation.md](query-translation.md) (sqlglot and dialects), [routing-and-clusters.md](routing-and-clusters.md) (routers, groups, load balancing), [observability.md](observability.md) (Prometheus, Grafana, Studio, Admin API), [adding-support/overview.md](adding-support/overview.md) (Extending QueryFlux — backend, frontend). --- @@ -163,6 +163,7 @@ pub trait RouterTrait: Send + Sync { | PostgreSQL wire | **Done** | 5432 | | MySQL wire | **Done** | 3306 | | Arrow Flight SQL | **Done** (query execution) | — | +| Snowflake HTTP wire + SQL API | **Done** | configurable | | Admin / Prometheus metrics | **Done** | 9000 | | ClickHouse HTTP | Planned | 8123 | @@ -251,6 +252,7 @@ queryflux: trinoHttp: { enabled: true, port: 8080 } postgresWire: { enabled: false, port: 5432 } mysqlWire: { enabled: false, port: 3306 } + snowflakeHttp: { enabled: false, port: 8445 } persistence: inMemory: {} # or: postgres: { databaseUrl: "postgres://..." } adminApi: @@ -358,6 +360,7 @@ curl -s -X POST http://localhost:8080/v1/statement \ | P1 | PostgreSQL wire frontend | **Done** | | P1 | MySQL wire frontend + StarRocks backend | **Done** | | P1 | Arrow Flight SQL frontend | **Done** | +| P1 | Snowflake HTTP wire + SQL API frontend | **Done** | | P1 | QueryFlux Studio — management UI | **Done** | | P1 | Athena backend | **Done** | | P1 | Authentication / authorization (`queryflux-auth`) | **Done** | diff --git a/website/versioned_docs/version-0.1.0/configuration.md b/website/versioned_docs/version-0.1.0/configuration.md index 37a87f8..bab76d9 100644 --- a/website/versioned_docs/version-0.1.0/configuration.md +++ b/website/versioned_docs/version-0.1.0/configuration.md @@ -14,6 +14,9 @@ queryflux: trinoHttp: enabled: true port: 8080 + snowflakeHttp: + enabled: true + port: 8445 persistence: type: inMemory # or: postgres @@ -39,6 +42,8 @@ clusterGroups: routers: - type: protocolBased trinoHttp: trino-default + snowflakeHttp: trino-default + snowflakeSqlApi: trino-default - type: header headerName: x-target-engine diff --git a/website/versioned_docs/version-0.1.0/roadmap.md b/website/versioned_docs/version-0.1.0/roadmap.md index 4625e9b..37b6fe5 100644 --- a/website/versioned_docs/version-0.1.0/roadmap.md +++ b/website/versioned_docs/version-0.1.0/roadmap.md @@ -19,6 +19,7 @@ Everything below is implemented and available on the `main` branch. | | PostgreSQL wire protocol (port 5432) | | | MySQL wire protocol (port 3306) | | | Arrow Flight SQL (gRPC) | +| | Snowflake HTTP wire + SQL API v2 ([docs](/docs/architecture/frontends/snowflake)) | | | Admin REST API + OpenAPI / Swagger UI (port 9000) | | **Backends** | Trino — async HTTP polling, transparent `nextUri` proxying | | | DuckDB — embedded, in-process, Arrow result sets | diff --git a/website/versioned_sidebars/version-0.1.0-sidebars.json b/website/versioned_sidebars/version-0.1.0-sidebars.json index cbb8451..95c48da 100644 --- a/website/versioned_sidebars/version-0.1.0-sidebars.json +++ b/website/versioned_sidebars/version-0.1.0-sidebars.json @@ -27,7 +27,29 @@ "architecture/routing-and-clusters", "architecture/query-tags", "architecture/observability", - "architecture/adding-engine-support", + { + "type": "category", + "label": "Frontends", + "collapsed": false, + "items": [ + "architecture/frontends/overview", + "architecture/frontends/trino-http", + "architecture/frontends/postgres-wire", + "architecture/frontends/mysql-wire", + "architecture/frontends/flight-sql", + "architecture/frontends/snowflake" + ] + }, + { + "type": "category", + "label": "Extending QueryFlux", + "collapsed": false, + "items": [ + "architecture/adding-support/overview", + "architecture/adding-support/backend", + "architecture/adding-support/frontend" + ] + }, "architecture/auth-authz-design" ] } From 0b43d451622d6c79a31753a9bb734270243a2c18 Mon Sep 17 00:00:00 2001 From: amitgilad Date: Thu, 2 Apr 2026 23:38:43 +0300 Subject: [PATCH 02/14] fix: create factory for all backend engines --- Cargo.lock | 155 +++++ Makefile | 7 +- README.md | 4 +- crates/queryflux-core/src/config.rs | 15 + crates/queryflux-core/src/engine_registry.rs | 51 ++ crates/queryflux-core/src/query.rs | 9 +- crates/queryflux-e2e-tests/src/harness.rs | 76 ++- .../queryflux-e2e-tests/src/trino_client.rs | 12 +- .../tests/snowflake_tests.rs | 206 +++++++ crates/queryflux-engine-adapters/Cargo.toml | 2 + .../src/athena/mod.rs | 63 ++ .../src/duckdb/http.rs | 49 ++ .../src/duckdb/mod.rs | 55 ++ crates/queryflux-engine-adapters/src/lib.rs | 30 +- .../src/snowflake/mod.rs | 557 ++++++++++++++++++ .../src/starrocks/mod.rs | 48 ++ .../src/trino/mod.rs | 51 ++ crates/queryflux-frontend/src/admin.rs | 26 +- crates/queryflux-frontend/src/dispatch.rs | 154 ++++- .../src/cluster_config.rs | 84 +-- crates/queryflux/src/main.rs | 266 ++++++--- crates/queryflux/src/registered_engines.rs | 70 ++- development.md | 2 +- docker/{ => test}/docker-compose.test.yml | 28 +- docker/test/fakesnow-apply-patches.py | 125 ++++ docker/test/fakesnow-entrypoint.sh | 11 + queryflux-studio/app/protocols/page.tsx | 3 +- .../components/cluster-config/index.ts | 1 + .../snowflake-cluster-config.tsx | 282 +++++++++ .../cluster-config/studio-engine-forms.tsx | 2 + queryflux-studio/components/engine-catalog.ts | 10 +- queryflux-studio/lib/cluster-persist-form.ts | 20 + .../lib/studio-engines/engines/snowflake.ts | 107 ++++ .../lib/studio-engines/manifest.ts | 2 + .../architecture/adding-support/backend.md | 307 +++------- .../architecture/adding-support/overview.md | 4 +- .../architecture/adding-support/backend.md | 307 +++------- .../architecture/adding-support/overview.md | 4 +- 38 files changed, 2557 insertions(+), 648 deletions(-) create mode 100644 crates/queryflux-e2e-tests/tests/snowflake_tests.rs create mode 100644 crates/queryflux-engine-adapters/src/snowflake/mod.rs rename docker/{ => test}/docker-compose.test.yml (85%) create mode 100644 docker/test/fakesnow-apply-patches.py create mode 100755 docker/test/fakesnow-entrypoint.sh create mode 100644 queryflux-studio/components/cluster-config/snowflake-cluster-config.tsx create mode 100644 queryflux-studio/lib/studio-engines/engines/snowflake.ts diff --git a/Cargo.lock b/Cargo.lock index 794eb77..735aa76 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8,6 +8,17 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + [[package]] name = "ahash" version = "0.7.8" @@ -414,6 +425,18 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "async-compression" +version = "0.4.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0f9ee0f6e02ffd7ad5816e9464499fba7b3effd01123b515c41d1697c43dad1" +dependencies = [ + "compression-codecs", + "compression-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -946,6 +969,15 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-padding" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8894febbff9f758034a5b8e12d87918f56dfc64a8e1fe757d65e29041538d93" +dependencies = [ + "generic-array", +] + [[package]] name = "borsh" version = "1.6.1" @@ -1035,6 +1067,15 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" +[[package]] +name = "cbc" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26b52a9543ae338f279b96b0b9fed9c8093744685043739079ce85cd58f289a6" +dependencies = [ + "cipher", +] + [[package]] name = "cc" version = "1.2.57" @@ -1098,6 +1139,16 @@ dependencies = [ "phf", ] +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common", + "inout", +] + [[package]] name = "clang-sys" version = "1.8.1" @@ -1184,6 +1235,23 @@ dependencies = [ "unicode-width", ] +[[package]] +name = "compression-codecs" +version = "0.4.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb7b51a7d9c967fc26773061ba86150f19c50c0d65c887cb1fbe295fd16619b7" +dependencies = [ + "compression-core", + "flate2", + "memchr", +] + +[[package]] +name = "compression-core" +version = "0.4.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75984efb6ed102a0d42db99afb6c1948f0380d1d91808d5529916e6c08b49d8d" + [[package]] name = "concurrent-queue" version = "2.5.0" @@ -2338,6 +2406,16 @@ dependencies = [ "serde_core", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "block-padding", + "generic-array", +] + [[package]] name = "ipnet" version = "2.12.0" @@ -3116,6 +3194,16 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pbkdf2" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ed6a7761f76e3b9f92dfb0a60a6a6477c61024b775147ff0973a02653abaf2" +dependencies = [ + "digest", + "hmac", +] + [[package]] name = "pem" version = "3.0.6" @@ -3202,6 +3290,21 @@ dependencies = [ "spki", ] +[[package]] +name = "pkcs5" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6" +dependencies = [ + "aes", + "cbc", + "der", + "pbkdf2", + "scrypt", + "sha2", + "spki", +] + [[package]] name = "pkcs8" version = "0.10.2" @@ -3209,6 +3312,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" dependencies = [ "der", + "pkcs5", + "rand_core 0.6.4", "spki", ] @@ -3566,9 +3671,11 @@ dependencies = [ "reqwest 0.12.28", "serde", "serde_json", + "snowflake-connector-rs", "tokio", "tokio-stream", "tracing", + "url", ] [[package]] @@ -4179,6 +4286,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "salsa20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97a22f5af31f73a954c10289c93e8a50cc23d971e80ee446f1f6f7137a088213" +dependencies = [ + "cipher", +] + [[package]] name = "same-file" version = "1.0.6" @@ -4209,6 +4325,17 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "scrypt" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0516a385866c09368f0b5bcd1caff3366aace790fcd46e2bb032697bb172fd1f" +dependencies = [ + "pbkdf2", + "salsa20", + "sha2", +] + [[package]] name = "sct" version = "0.7.1" @@ -4435,6 +4562,29 @@ dependencies = [ "serde", ] +[[package]] +name = "snowflake-connector-rs" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab6cf61174f3727f62abad67b3e407861024f959dd7cd08d23e867a8df9cec9a" +dependencies = [ + "base64 0.22.1", + "chrono", + "flate2", + "http 1.4.0", + "jsonwebtoken 9.3.1", + "pkcs8", + "reqwest 0.12.28", + "rsa", + "serde", + "serde_json", + "sha2", + "thiserror 2.0.18", + "tokio", + "url", + "uuid", +] + [[package]] name = "socket2" version = "0.5.10" @@ -5087,13 +5237,18 @@ version = "0.6.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" dependencies = [ + "async-compression", "bitflags", "bytes", + "futures-core", "futures-util", "http 1.4.0", "http-body 1.0.1", + "http-body-util", "iri-string", "pin-project-lite", + "tokio", + "tokio-util", "tower", "tower-layer", "tower-service", diff --git a/Makefile b/Makefile index b7686ad..3b92384 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,6 @@ CARGO := $(HOME)/.cargo/bin/cargo COMPOSE := docker compose -f docker/docker-compose.yml --project-directory . -COMPOSE_TEST := docker compose -f docker/docker-compose.test.yml --project-directory . +COMPOSE_TEST := docker compose -f docker/test/docker-compose.test.yml --project-directory . # Trino `tpch` schema used when loading Iceberg tables (see docker/fixtures/init.sql + data-loader). # tiny = default fast tests; sf1 ≈ 1.5M orders (long load, heavy E2E). @@ -70,7 +70,7 @@ benchmark: $(CARGO) run --release -p queryflux-bench ## Run E2E tests. Spins up Trino + StarRocks + Lakekeeper via Docker. -## Requires reachable engines; see docker/docker-compose.test.yml. +## Requires reachable engines; see docker/test/docker-compose.test.yml. ## `--test-threads=1`: StarRocks Iceberg is slow; default parallel libtest + `#[serial]` makes ## every test report libtest's 60s "slow test" spam while threads wait on the serial lock. ## Iceberg/Lakekeeper tables are created by the e2e crate (no TPC-H loader). @@ -78,11 +78,12 @@ test-e2e: @test -f .venv/bin/python3 || (echo "Run 'make setup' first" && exit 1) PYO3_PYTHON=$(shell pwd)/.venv/bin/python3 \ PYTHONPATH=$(shell pwd)/.venv/lib/python3.13/site-packages \ - $(COMPOSE_TEST) up -d --wait trino starrocks sentinel + $(COMPOSE_TEST) up -d --wait trino starrocks fakesnow sentinel PYO3_PYTHON=$(shell pwd)/.venv/bin/python3 \ PYTHONPATH=$(shell pwd)/.venv/lib/python3.13/site-packages \ TRINO_URL=http://localhost:18081 \ STARROCKS_URL=mysql://root@localhost:9030 \ + FAKESNOW_URL=http://localhost:18085 \ LAKEKEEPER_URL=http://localhost:18181 \ MINIO_ENDPOINT=localhost:19000 \ DUCKDB_DOWNLOAD_LIB=1 \ diff --git a/README.md b/README.md index 1ff3ab7..27526e5 100644 --- a/README.md +++ b/README.md @@ -189,8 +189,8 @@ queryflux/ ├── config.example.yaml ├── docker/ │ ├── docker-compose.yml # Local dev stack (`make dev`) -│ ├── docker-compose.test.yml # E2E test stack -│ ├── fixtures/ # SQL init, test data +│ ├── fixtures/ # SQL init, test data (shared with examples) +│ ├── test/ # E2E stack: docker-compose.test.yml, fakesnow helpers │ ├── queryflux/ # QueryFlux Dockerfile │ └── queryflux-studio/ # Studio Dockerfile ├── docs/ # Architecture markdown diff --git a/crates/queryflux-core/src/config.rs b/crates/queryflux-core/src/config.rs index 5e4cc2c..88e313e 100644 --- a/crates/queryflux-core/src/config.rs +++ b/crates/queryflux-core/src/config.rs @@ -509,6 +509,8 @@ pub enum EngineConfig { ClickHouse, /// Amazon Athena — serverless SQL over S3 via the AWS SDK. Athena, + /// Snowflake — cloud-native data warehouse. Connects via the Snowflake REST API. + Snowflake, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -534,8 +536,21 @@ pub struct ClusterConfig { #[serde(default)] pub workgroup: Option, /// Default Athena catalog. Defaults to `"AwsDataCatalog"` when omitted. + /// For Snowflake, used as the default database name. #[serde(default)] pub catalog: Option, + /// Snowflake account identifier (e.g. `xy12345.us-east-1`). Required when engine is `snowflake`. + #[serde(default)] + pub account: Option, + /// Default Snowflake virtual warehouse (e.g. `COMPUTE_WH`). + #[serde(default)] + pub warehouse: Option, + /// Default Snowflake role (e.g. `ANALYST`). + #[serde(default)] + pub role: Option, + /// Default Snowflake schema (e.g. `PUBLIC`). + #[serde(default)] + pub schema: Option, #[serde(default)] pub tls: Option, /// Type 1 credentials — service account used for health checks and (by default) query execution. diff --git a/crates/queryflux-core/src/engine_registry.rs b/crates/queryflux-core/src/engine_registry.rs index 71fbbd1..ac25af0 100644 --- a/crates/queryflux-core/src/engine_registry.rs +++ b/crates/queryflux-core/src/engine_registry.rs @@ -237,6 +237,54 @@ pub fn validate_cluster_config( errors } +// --------------------------------------------------------------------------- +// Config JSON helpers +// --------------------------------------------------------------------------- + +/// Extract a `ClusterAuth` from the flat DB JSON format used by persistence. +/// +/// The JSON blob stores auth as flat keys: `authType`, `authUsername`, +/// `authPassword`, `authToken`. This is the canonical format produced by +/// `UpsertClusterConfig::from_core()` and stored in the `config` JSONB column. +pub fn parse_auth_from_config_json(json: &serde_json::Value) -> Option { + let s = + |key: &str| -> Option { json.get(key).and_then(|v| v.as_str()).map(String::from) }; + match s("authType").as_deref() { + Some("basic") => Some(ClusterAuth::Basic { + username: s("authUsername").unwrap_or_default(), + password: s("authPassword").unwrap_or_default(), + }), + Some("bearer") => Some(ClusterAuth::Bearer { + token: s("authToken").unwrap_or_default(), + }), + Some("keyPair") => Some(ClusterAuth::KeyPair { + username: s("authUsername").unwrap_or_default(), + private_key_pem: s("authPassword").unwrap_or_default(), + private_key_passphrase: s("authToken"), + }), + Some("accessKey") => Some(ClusterAuth::AccessKey { + access_key_id: s("authUsername").unwrap_or_default(), + secret_access_key: s("authPassword").unwrap_or_default(), + session_token: s("authToken"), + }), + Some("roleArn") => Some(ClusterAuth::RoleArn { + role_arn: s("authUsername").unwrap_or_default(), + external_id: s("authToken"), + }), + _ => None, + } +} + +/// Extract an optional string field from a config JSON blob. +pub fn json_str(json: &serde_json::Value, key: &str) -> Option { + json.get(key).and_then(|v| v.as_str()).map(String::from) +} + +/// Extract a boolean field from a config JSON blob (defaults to `false`). +pub fn json_bool(json: &serde_json::Value, key: &str) -> bool { + json.get(key).and_then(|v| v.as_bool()).unwrap_or(false) +} + // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- @@ -251,6 +299,7 @@ pub fn engine_key(engine: &EngineConfig) -> &'static str { EngineConfig::StarRocks => "starRocks", EngineConfig::ClickHouse => "clickHouse", EngineConfig::Athena => "athena", + EngineConfig::Snowflake => "snowflake", } } @@ -263,6 +312,7 @@ pub fn parse_engine_key(s: &str) -> Result { "starRocks" => Ok(EngineConfig::StarRocks), "clickHouse" => Ok(EngineConfig::ClickHouse), "athena" => Ok(EngineConfig::Athena), + "snowflake" => Ok(EngineConfig::Snowflake), other => Err(format!("Unknown engine key: '{other}'")), } } @@ -276,6 +326,7 @@ impl From<&EngineConfig> for EngineType { EngineConfig::StarRocks => EngineType::StarRocks, EngineConfig::ClickHouse => EngineType::ClickHouse, EngineConfig::Athena => EngineType::Athena, + EngineConfig::Snowflake => EngineType::Snowflake, } } } diff --git a/crates/queryflux-core/src/query.rs b/crates/queryflux-core/src/query.rs index 2dc669f..b7f6a23 100644 --- a/crates/queryflux-core/src/query.rs +++ b/crates/queryflux-core/src/query.rs @@ -82,8 +82,8 @@ impl FrontendProtocol { FrontendProtocol::MySqlWire => SqlDialect::MySql, FrontendProtocol::ClickHouseHttp => SqlDialect::ClickHouse, FrontendProtocol::FlightSql => SqlDialect::Generic, - FrontendProtocol::SnowflakeHttp => SqlDialect::Generic, - FrontendProtocol::SnowflakeSqlApi => SqlDialect::Generic, + FrontendProtocol::SnowflakeHttp => SqlDialect::Snowflake, + FrontendProtocol::SnowflakeSqlApi => SqlDialect::Snowflake, } } } @@ -99,6 +99,8 @@ pub enum EngineType { ClickHouse, /// Amazon Athena (Presto/Trino-compatible SQL over S3). Athena, + /// Snowflake cloud data warehouse. + Snowflake, } impl EngineType { @@ -109,6 +111,7 @@ impl EngineType { EngineType::DuckDb | EngineType::DuckDbHttp => SqlDialect::DuckDb, EngineType::StarRocks => SqlDialect::StarRocks, EngineType::ClickHouse => SqlDialect::ClickHouse, + EngineType::Snowflake => SqlDialect::Snowflake, } } } @@ -121,6 +124,7 @@ pub enum SqlDialect { DuckDb, StarRocks, ClickHouse, + Snowflake, MySql, Postgres, Generic, @@ -146,6 +150,7 @@ impl SqlDialect { SqlDialect::DuckDb => "duckdb", SqlDialect::StarRocks => "starrocks", SqlDialect::ClickHouse => "clickhouse", + SqlDialect::Snowflake => "snowflake", SqlDialect::MySql => "mysql", SqlDialect::Postgres => "postgres", SqlDialect::Generic => "", diff --git a/crates/queryflux-e2e-tests/src/harness.rs b/crates/queryflux-e2e-tests/src/harness.rs index 4eacecc..74e1ce3 100644 --- a/crates/queryflux-e2e-tests/src/harness.rs +++ b/crates/queryflux-e2e-tests/src/harness.rs @@ -23,12 +23,14 @@ use queryflux_auth::{ use queryflux_cluster_manager::{ cluster_state::ClusterState, simple::SimpleClusterGroupManager, strategy::strategy_from_config, }; +use queryflux_core::config::{ClusterAuth, ClusterConfig, EngineConfig}; use queryflux_core::{ error::Result as QfResult, query::{ClusterGroupName, ClusterName, EngineType}, }; use queryflux_engine_adapters::{ - starrocks::StarRocksAdapter, trino::TrinoAdapter, EngineAdapterTrait, + snowflake::SnowflakeAdapter, starrocks::StarRocksAdapter, trino::TrinoAdapter, + EngineAdapterTrait, }; use queryflux_frontend::{ snowflake::http::session_store::SnowflakeSessionStore, @@ -59,6 +61,7 @@ impl MetricsStore for CapturingMetrics { pub const GROUP_TRINO: &str = "trino"; pub const GROUP_STARROCKS: &str = "starrocks"; +pub const GROUP_SNOWFLAKE: &str = "snowflake"; /// Set when Lakekeeper port is reachable (Iceberg tables seeded by e2e tests via Trino). pub const GROUP_LAKEKEEPER: &str = "lakekeeper"; @@ -178,10 +181,66 @@ impl TestHarness { adapters.insert(cluster.0.clone(), sr as Arc); } + // --- Snowflake (fakesnow) --- + let fakesnow_url = + std::env::var("FAKESNOW_URL").unwrap_or_else(|_| "http://localhost:18085".to_string()); + let fakesnow_available = is_fakesnow_ready(&fakesnow_url).await; + if fakesnow_available { + let group = ClusterGroupName(GROUP_SNOWFLAKE.to_string()); + let cluster = ClusterName("snowflake-1".to_string()); + let state = Arc::new(ClusterState::new( + cluster.clone(), + group.clone(), + None, + None, + EngineType::Snowflake, + Some(fakesnow_url.clone()), + 8, + true, + )); + let cfg = ClusterConfig { + engine: Some(EngineConfig::Snowflake), + enabled: true, + max_running_queries: None, + endpoint: Some(fakesnow_url), + database_path: None, + region: None, + s3_output_location: None, + workgroup: None, + catalog: None, + account: Some("fakesnow".to_string()), + warehouse: None, + role: None, + schema: None, + tls: None, + auth: Some(ClusterAuth::Basic { + username: "fake".to_string(), + password: "snow".to_string(), + }), + query_auth: None, + }; + let adapter = Arc::new( + SnowflakeAdapter::try_from_cluster_config( + cluster.clone(), + group.clone(), + &cfg, + "snowflake-1", + ) + .map_err(|e| anyhow!("Snowflake adapter: {e}"))?, + ) as Arc; + + group_states.insert(group.clone(), (vec![state], strategy_from_config(None))); + group_members.insert(GROUP_SNOWFLAKE.to_string(), vec![cluster.0.clone()]); + group_order.push(GROUP_SNOWFLAKE.to_string()); + adapters.insert(cluster.0.clone(), adapter); + available_groups.push(GROUP_SNOWFLAKE.to_string()); + header_map.insert(GROUP_SNOWFLAKE.to_string(), group); + } + if group_states.is_empty() { return Err(anyhow!( - "No backends reachable. Start docker compose (see docker/docker-compose.test.yml): \ - Trino :18081 and/or StarRocks :19030." + "No backends reachable. Start docker compose (see docker/test/docker-compose.test.yml): \ + Trino :18081 and/or StarRocks :19030 and/or fakesnow :18085." )); } @@ -290,7 +349,7 @@ impl TestHarness { } fn pick_fallback_group(group_order: &[String]) -> ClusterGroupName { - for preferred in [GROUP_TRINO, GROUP_STARROCKS] { + for preferred in [GROUP_TRINO, GROUP_STARROCKS, GROUP_SNOWFLAKE] { if group_order.iter().any(|g| g == preferred) { return ClusterGroupName(preferred.to_string()); } @@ -334,3 +393,12 @@ async fn is_lakekeeper_ready(url: &str) -> bool { let port = parsed.port().unwrap_or(8181); port_is_open(host, port).await } + +async fn is_fakesnow_ready(url: &str) -> bool { + let Ok(parsed) = reqwest::Url::parse(url) else { + return false; + }; + let host = parsed.host_str().unwrap_or("localhost"); + let port = parsed.port().unwrap_or(8085); + port_is_open(host, port).await +} diff --git a/crates/queryflux-e2e-tests/src/trino_client.rs b/crates/queryflux-e2e-tests/src/trino_client.rs index 9424934..b723fd9 100644 --- a/crates/queryflux-e2e-tests/src/trino_client.rs +++ b/crates/queryflux-e2e-tests/src/trino_client.rs @@ -30,14 +30,18 @@ pub struct ColumnInfo { struct TrinoResponse { #[serde(rename = "nextUri")] next_uri: Option, + /// Present on Trino pages; kept for JSON shape (poll loop uses `nextUri` only). + #[serde(default)] + #[allow(dead_code)] stats: Option, columns: Option>, data: Option>>, error: Option, } -#[derive(Debug, Deserialize)] +#[derive(Debug, Deserialize, Default)] struct TrinoStats { + #[allow(dead_code)] state: String, } @@ -113,10 +117,12 @@ impl TrinoClient { }); } - let state = page.stats.as_ref().map(|s| s.state.as_str()).unwrap_or(""); let next = page.next_uri.take(); - if next.is_none() || state == "FINISHED" || state == "FAILED" { + // Trino may include `stats.state: FINISHED` while `nextUri` is still set (one more poll + // may be required). Stop only when there is no next page — matches Trino clients and + // ensures QueryFlux sees a terminal GET so it can record_query. + if next.is_none() { break; } diff --git a/crates/queryflux-e2e-tests/tests/snowflake_tests.rs b/crates/queryflux-e2e-tests/tests/snowflake_tests.rs new file mode 100644 index 0000000..9b862f9 --- /dev/null +++ b/crates/queryflux-e2e-tests/tests/snowflake_tests.rs @@ -0,0 +1,206 @@ +/// Snowflake tests — require a running fakesnow instance (https://github.com/tekumara/fakesnow). +/// +/// All tests are marked `#[ignore]` and run with: make test-e2e +/// or: cargo test -p queryflux-e2e-tests --test snowflake_tests -- --include-ignored +use std::sync::OnceLock; + +use queryflux_e2e_tests::{ + harness::{TestHarness, GROUP_SNOWFLAKE}, + trino_client::TrinoClient, +}; + +static HARNESS: OnceLock = OnceLock::new(); + +fn harness() -> &'static TestHarness { + HARNESS.get_or_init(|| { + let (tx, rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + let rt = tokio::runtime::Runtime::new().expect("runtime"); + let h = rt.block_on(TestHarness::new()).expect("TestHarness::new"); + tx.send(h).expect("send harness"); + rt.block_on(std::future::pending::<()>()); + }); + rx.recv().expect("recv harness") + }) +} + +fn client() -> TrinoClient { + TrinoClient::new(&harness().base_url()) +} + +macro_rules! require_snowflake { + () => { + if !harness().has_group(GROUP_SNOWFLAKE) { + eprintln!( + "SKIP: fakesnow not available (start with docker/test/docker-compose.test.yml)" + ); + return; + } + }; +} + +// --------------------------------------------------------------------------- +// Basic queries +// --------------------------------------------------------------------------- + +#[tokio::test] +#[ignore = "requires fakesnow — run with: make test-e2e"] +async fn snowflake_select_literal() { + require_snowflake!(); + let r = client() + .execute_on("SELECT 1 + 1 AS result", GROUP_SNOWFLAKE) + .await + .expect("query"); + assert!(r.error.is_none(), "unexpected error: {:?}", r.error); + assert_eq!(r.rows.len(), 1); +} + +#[tokio::test] +#[ignore = "requires fakesnow — run with: make test-e2e"] +async fn snowflake_select_string_literal() { + require_snowflake!(); + let r = client() + .execute_on("SELECT 'hello fakesnow' AS greeting", GROUP_SNOWFLAKE) + .await + .expect("query"); + assert!(r.error.is_none(), "unexpected error: {:?}", r.error); + assert_eq!(r.rows.len(), 1); + assert_eq!(r.rows[0][0].as_str(), Some("hello fakesnow")); +} + +#[tokio::test] +#[ignore = "requires fakesnow — run with: make test-e2e"] +async fn snowflake_select_multi_row() { + require_snowflake!(); + let r = client() + .execute_on( + "SELECT 1 AS v UNION ALL SELECT 2 UNION ALL SELECT 3", + GROUP_SNOWFLAKE, + ) + .await + .expect("query"); + assert!(r.error.is_none(), "unexpected error: {:?}", r.error); + assert_eq!(r.rows.len(), 3); +} + +#[tokio::test] +#[ignore = "requires fakesnow — run with: make test-e2e"] +async fn snowflake_empty_result() { + require_snowflake!(); + let r = client() + .execute_on("SELECT 1 AS n WHERE 1 = 0", GROUP_SNOWFLAKE) + .await + .expect("query"); + assert!(r.error.is_none(), "unexpected error: {:?}", r.error); + assert_eq!(r.rows.len(), 0); +} + +// --------------------------------------------------------------------------- +// DDL + DML (fakesnow supports CREATE TABLE, INSERT, etc.) +// --------------------------------------------------------------------------- + +#[tokio::test] +#[ignore = "requires fakesnow — run with: make test-e2e"] +async fn snowflake_create_and_query_table() { + require_snowflake!(); + let c = client(); + + c.execute_on("CREATE DATABASE IF NOT EXISTS e2e_db", GROUP_SNOWFLAKE) + .await + .expect("create database"); + + c.execute_on( + "CREATE SCHEMA IF NOT EXISTS e2e_db.e2e_schema", + GROUP_SNOWFLAKE, + ) + .await + .expect("create schema"); + + c.execute_on( + "CREATE OR REPLACE TABLE e2e_db.e2e_schema.test_tbl (id INTEGER, name VARCHAR)", + GROUP_SNOWFLAKE, + ) + .await + .expect("create table"); + + c.execute_on( + "INSERT INTO e2e_db.e2e_schema.test_tbl VALUES (1, 'alice'), (2, 'bob'), (3, 'charlie')", + GROUP_SNOWFLAKE, + ) + .await + .expect("insert"); + + let r = c + .execute_on( + "SELECT id, name FROM e2e_db.e2e_schema.test_tbl ORDER BY id", + GROUP_SNOWFLAKE, + ) + .await + .expect("select"); + assert!(r.error.is_none(), "unexpected error: {:?}", r.error); + assert_eq!(r.rows.len(), 3); + assert_eq!(r.rows[0][1].as_str(), Some("alice")); + assert_eq!(r.rows[2][1].as_str(), Some("charlie")); +} + +// --------------------------------------------------------------------------- +// Type mapping +// --------------------------------------------------------------------------- + +#[tokio::test] +#[ignore = "requires fakesnow — run with: make test-e2e"] +async fn snowflake_numeric_types() { + require_snowflake!(); + let r = client() + .execute_on( + "SELECT 42 AS int_val, 3.14 AS float_val, TRUE AS bool_val", + GROUP_SNOWFLAKE, + ) + .await + .expect("query"); + assert!(r.error.is_none(), "unexpected error: {:?}", r.error); + assert_eq!(r.rows.len(), 1); +} + +// --------------------------------------------------------------------------- +// Aggregations +// --------------------------------------------------------------------------- + +#[tokio::test] +#[ignore = "requires fakesnow — run with: make test-e2e"] +async fn snowflake_aggregation() { + require_snowflake!(); + let r = client() + .execute_on( + "SELECT COUNT(*) AS cnt, SUM(v) AS total FROM (SELECT 10 AS v UNION ALL SELECT 20 UNION ALL SELECT 30) t", + GROUP_SNOWFLAKE, + ) + .await + .expect("query"); + assert!(r.error.is_none(), "unexpected error: {:?}", r.error); + assert_eq!(r.rows.len(), 1); +} + +// --------------------------------------------------------------------------- +// Metrics capture +// --------------------------------------------------------------------------- + +#[tokio::test] +#[ignore = "requires fakesnow — run with: make test-e2e"] +async fn snowflake_query_recorded_in_metrics() { + require_snowflake!(); + let h = harness(); + h.clear_records(); + let c = client(); + c.execute_on("SELECT 999 AS metric_test", GROUP_SNOWFLAKE) + .await + .expect("query"); + + let record = h + .wait_for_record(|r| r.cluster_group.0 == GROUP_SNOWFLAKE) + .await; + assert!( + record.is_some(), + "expected a query record for the snowflake group" + ); +} diff --git a/crates/queryflux-engine-adapters/Cargo.toml b/crates/queryflux-engine-adapters/Cargo.toml index eb4db27..fb9bbc3 100644 --- a/crates/queryflux-engine-adapters/Cargo.toml +++ b/crates/queryflux-engine-adapters/Cargo.toml @@ -23,3 +23,5 @@ aws-sdk-sts = "1" aws-config = { version = "1", default-features = false, features = ["rustls"] } aws-credential-types = "1" aws-types = "1" +snowflake-connector-rs = "0.9" +url = "2" diff --git a/crates/queryflux-engine-adapters/src/athena/mod.rs b/crates/queryflux-engine-adapters/src/athena/mod.rs index c70417b..31f68af 100644 --- a/crates/queryflux-engine-adapters/src/athena/mod.rs +++ b/crates/queryflux-engine-adapters/src/athena/mod.rs @@ -169,6 +169,43 @@ impl AthenaAdapter { }) } + /// Build from a DB config JSON blob (bypasses the `ClusterConfig` god struct). + pub async fn try_from_config_json( + cluster_name: ClusterName, + group_name: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result { + use queryflux_core::engine_registry::{json_str, parse_auth_from_config_json}; + + let region = json_str(json, "region").ok_or_else(|| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': missing 'region' for Athena" + )) + })?; + let s3_output = json_str(json, "s3OutputLocation").ok_or_else(|| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': missing 's3OutputLocation' for Athena" + )) + })?; + let auth = parse_auth_from_config_json(json); + Self::new( + cluster_name, + group_name, + region, + s3_output, + json_str(json, "workgroup"), + json_str(json, "catalog"), + auth, + ) + .await + .map_err(|e| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': failed to create Athena adapter ({e})" + )) + }) + } + /// Build from persisted / YAML [`ClusterConfig`]. pub async fn try_from_cluster_config( cluster_name: ClusterName, @@ -729,3 +766,29 @@ impl AthenaAdapter { } } } + +pub struct AthenaFactory; + +#[async_trait] +impl crate::EngineAdapterFactory for AthenaFactory { + fn engine_key(&self) -> &'static str { + "athena" + } + + fn descriptor(&self) -> EngineDescriptor { + AthenaAdapter::descriptor() + } + + async fn build_from_config_json( + &self, + cluster_name: ClusterName, + group: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result> { + Ok(Arc::new( + AthenaAdapter::try_from_config_json(cluster_name, group, json, cluster_name_str) + .await?, + )) + } +} diff --git a/crates/queryflux-engine-adapters/src/duckdb/http.rs b/crates/queryflux-engine-adapters/src/duckdb/http.rs index 6d4d389..9e37337 100644 --- a/crates/queryflux-engine-adapters/src/duckdb/http.rs +++ b/crates/queryflux-engine-adapters/src/duckdb/http.rs @@ -140,6 +140,27 @@ impl DuckDbHttpAdapter { }) } + /// Build from a DB config JSON blob (bypasses the `ClusterConfig` god struct). + pub fn try_from_config_json( + cluster_name: ClusterName, + group_name: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result { + use queryflux_core::engine_registry::{json_bool, json_str, parse_auth_from_config_json}; + + let endpoint = json_str(json, "endpoint").ok_or_else(|| { + QueryFluxError::Engine(format!("cluster '{cluster_name_str}': missing endpoint")) + })?; + let tls_skip = json_bool(json, "tlsInsecureSkipVerify"); + let auth = parse_auth_from_config_json(json); + Self::new(cluster_name, group_name, endpoint, tls_skip, auth).map_err(|e| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': failed to create DuckDB HTTP adapter ({e})" + )) + }) + } + /// Build from persisted / YAML [`ClusterConfig`]. pub fn try_from_cluster_config( cluster_name: ClusterName, @@ -546,3 +567,31 @@ impl DuckDbHttpAdapter { } } } + +pub struct DuckDbHttpFactory; + +#[async_trait] +impl crate::EngineAdapterFactory for DuckDbHttpFactory { + fn engine_key(&self) -> &'static str { + "duckDbHttp" + } + + fn descriptor(&self) -> EngineDescriptor { + DuckDbHttpAdapter::descriptor() + } + + async fn build_from_config_json( + &self, + cluster_name: ClusterName, + group: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result> { + Ok(Arc::new(DuckDbHttpAdapter::try_from_config_json( + cluster_name, + group, + json, + cluster_name_str, + )?)) + } +} diff --git a/crates/queryflux-engine-adapters/src/duckdb/mod.rs b/crates/queryflux-engine-adapters/src/duckdb/mod.rs index e554c22..31238fe 100644 --- a/crates/queryflux-engine-adapters/src/duckdb/mod.rs +++ b/crates/queryflux-engine-adapters/src/duckdb/mod.rs @@ -67,6 +67,33 @@ impl DuckDbAdapter { }) } + /// Build from a DB config JSON blob (bypasses the `ClusterConfig` god struct). + pub fn try_from_config_json( + cluster_name: ClusterName, + group_name: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result { + use queryflux_core::engine_registry::{json_str, parse_auth_from_config_json}; + + let database_path = json_str(json, "databasePath"); + let auth = parse_auth_from_config_json(json); + let motherduck_token = auth.and_then(|a| { + if let ClusterAuth::Bearer { token } = a { + Some(token) + } else { + None + } + }); + Self::new_with_token(cluster_name, group_name, database_path, motherduck_token).map_err( + |e| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': failed to open DuckDB ({e})" + )) + }, + ) + } + /// Build from persisted / YAML [`ClusterConfig`] (embedded path + optional MotherDuck bearer). pub fn try_from_cluster_config( cluster_name: ClusterName, @@ -400,3 +427,31 @@ impl DuckDbAdapter { } } } + +pub struct DuckDbFactory; + +#[async_trait] +impl crate::EngineAdapterFactory for DuckDbFactory { + fn engine_key(&self) -> &'static str { + "duckDb" + } + + fn descriptor(&self) -> EngineDescriptor { + DuckDbAdapter::descriptor() + } + + async fn build_from_config_json( + &self, + cluster_name: ClusterName, + group: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result> { + Ok(Arc::new(DuckDbAdapter::try_from_config_json( + cluster_name, + group, + json, + cluster_name_str, + )?)) + } +} diff --git a/crates/queryflux-engine-adapters/src/lib.rs b/crates/queryflux-engine-adapters/src/lib.rs index e3169b8..881eede 100644 --- a/crates/queryflux-engine-adapters/src/lib.rs +++ b/crates/queryflux-engine-adapters/src/lib.rs @@ -1,9 +1,11 @@ pub mod athena; pub mod duckdb; +pub mod snowflake; pub mod starrocks; pub mod trino; use std::pin::Pin; +use std::sync::Arc; use arrow::record_batch::RecordBatch; use async_trait::async_trait; @@ -11,8 +13,11 @@ use futures::Stream; use queryflux_auth::QueryCredentials; use queryflux_core::{ catalog::TableSchema, + engine_registry::EngineDescriptor, error::{QueryFluxError, Result}, - query::{BackendQueryId, EngineType, QueryExecution, QueryPollResult}, + query::{ + BackendQueryId, ClusterGroupName, ClusterName, EngineType, QueryExecution, QueryPollResult, + }, session::SessionContext, tags::QueryTags, }; @@ -20,6 +25,29 @@ use queryflux_core::{ /// A stream of Arrow RecordBatches — the universal output type for all adapters. pub type ArrowStream = Pin> + Send>>; +/// Factory for constructing engine adapters from raw configuration. +/// +/// Each engine provides a zero-sized factory struct implementing this trait. +/// This formalizes the contract that every adapter must support construction +/// from a DB config JSON blob and expose its descriptor metadata. +#[async_trait] +pub trait EngineAdapterFactory: Send + Sync { + /// The engine key string matching the DB `engine_key` column (e.g. `"trino"`, `"duckDb"`). + fn engine_key(&self) -> &'static str; + + /// Field-level schema used by the admin API and Studio UI. + fn descriptor(&self) -> EngineDescriptor; + + /// Build an adapter instance from a raw DB config JSON blob. + async fn build_from_config_json( + &self, + cluster_name: ClusterName, + group: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result>; +} + /// Implemented by each query engine backend (Trino, DuckDB, StarRocks, ClickHouse, ...). /// /// Engines that run queries synchronously return `QueryExecution::Sync` from `submit_query`. diff --git a/crates/queryflux-engine-adapters/src/snowflake/mod.rs b/crates/queryflux-engine-adapters/src/snowflake/mod.rs new file mode 100644 index 0000000..3b81fd0 --- /dev/null +++ b/crates/queryflux-engine-adapters/src/snowflake/mod.rs @@ -0,0 +1,557 @@ +use std::sync::Arc; + +use arrow::array::{ArrayRef, BooleanBuilder, Float64Builder, Int64Builder, StringBuilder}; +use arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; +use arrow::record_batch::RecordBatch; +use async_trait::async_trait; +use futures::stream; +use queryflux_auth::QueryCredentials; +use queryflux_core::{ + catalog::TableSchema, + config::{ClusterAuth, ClusterConfig}, + engine_registry::{AuthType, ConfigField, ConnectionType, EngineDescriptor, FieldType}, + error::{QueryFluxError, Result}, + query::{ + BackendQueryId, ClusterGroupName, ClusterName, EngineType, QueryExecution, QueryPollResult, + }, + session::SessionContext, + tags::QueryTags, +}; +use snowflake_connector_rs::{ + SnowflakeAuthMethod, SnowflakeClient, SnowflakeClientConfig, SnowflakeColumnType, SnowflakeRow, +}; +use url::Url; + +use crate::EngineAdapterTrait; + +pub struct SnowflakeAdapter { + pub cluster_name: ClusterName, + pub group_name: ClusterGroupName, + client: SnowflakeClient, + account: String, +} + +impl SnowflakeAdapter { + /// Build from a DB config JSON blob (bypasses the `ClusterConfig` god struct). + pub fn try_from_config_json( + cluster_name: ClusterName, + group_name: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result { + use queryflux_core::engine_registry::{json_str, parse_auth_from_config_json}; + + let account = json_str(json, "account").ok_or_else(|| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': Snowflake requires 'account' field" + )) + })?; + + let auth = parse_auth_from_config_json(json).ok_or_else(|| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': Snowflake requires 'auth' configuration" + )) + })?; + + let (username, sf_auth) = map_auth(&auth).map_err(|msg| { + QueryFluxError::Engine(format!("cluster '{cluster_name_str}': {msg}")) + })?; + + let client_cfg = SnowflakeClientConfig { + account: account.clone(), + warehouse: json_str(json, "warehouse"), + database: json_str(json, "catalog"), + schema: json_str(json, "schema"), + role: json_str(json, "role"), + timeout: Some(std::time::Duration::from_secs(300)), + }; + + let client = SnowflakeClient::new(&username, sf_auth, client_cfg).map_err(|e| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': failed to create Snowflake client: {e}" + )) + })?; + + let client = if let Some(endpoint) = json_str(json, "endpoint") { + let url = Url::parse(&endpoint).map_err(|e| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': invalid endpoint URL: {e}" + )) + })?; + let host = url.host_str().unwrap_or_default(); + let port = url.port(); + let protocol = Some(url.scheme().to_string()); + client.with_address(host, port, protocol).map_err(|e| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': failed to configure Snowflake endpoint: {e}" + )) + })? + } else { + client + }; + + Ok(Self { + cluster_name, + group_name, + client, + account, + }) + } + + pub fn try_from_cluster_config( + cluster_name: ClusterName, + group_name: ClusterGroupName, + cfg: &ClusterConfig, + cluster_name_str: &str, + ) -> Result { + let account = cfg.account.clone().ok_or_else(|| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': Snowflake requires 'account' field" + )) + })?; + + let auth = cfg.auth.clone().ok_or_else(|| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': Snowflake requires 'auth' configuration" + )) + })?; + + let (username, sf_auth) = map_auth(&auth).map_err(|msg| { + QueryFluxError::Engine(format!("cluster '{cluster_name_str}': {msg}")) + })?; + + let client_cfg = SnowflakeClientConfig { + account: account.clone(), + warehouse: cfg.warehouse.clone(), + database: cfg.catalog.clone(), + schema: cfg.schema.clone(), + role: cfg.role.clone(), + timeout: Some(std::time::Duration::from_secs(300)), + }; + + let client = SnowflakeClient::new(&username, sf_auth, client_cfg).map_err(|e| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': failed to create Snowflake client: {e}" + )) + })?; + + let client = if let Some(endpoint) = &cfg.endpoint { + let url = Url::parse(endpoint).map_err(|e| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': invalid endpoint URL: {e}" + )) + })?; + let host = url.host_str().unwrap_or_default(); + let port = url.port(); + let protocol = Some(url.scheme().to_string()); + client.with_address(host, port, protocol).map_err(|e| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': failed to configure Snowflake endpoint: {e}" + )) + })? + } else { + client + }; + + Ok(Self { + cluster_name, + group_name, + client, + account, + }) + } + + async fn run_query(&self, sql: &str) -> Result> { + let session = self.client.create_session().await.map_err(|e| { + QueryFluxError::Engine(format!("Snowflake session creation failed: {e}")) + })?; + session + .query(sql) + .await + .map_err(|e| QueryFluxError::Engine(format!("Snowflake query failed: {e}"))) + } + + async fn run_first_col(&self, sql: &str) -> Result> { + let rows = self.run_query(sql).await?; + Ok(rows + .iter() + .filter_map(|row| row.at::(0).ok()) + .collect()) + } +} + +fn map_auth(auth: &ClusterAuth) -> std::result::Result<(String, SnowflakeAuthMethod), String> { + match auth { + ClusterAuth::Basic { username, password } => Ok(( + username.clone(), + SnowflakeAuthMethod::Password(password.clone()), + )), + ClusterAuth::KeyPair { + username, + private_key_pem, + private_key_passphrase, + } => { + let method = if let Some(passphrase) = private_key_passphrase { + SnowflakeAuthMethod::KeyPair { + encrypted_pem: private_key_pem.clone(), + password: passphrase.as_bytes().to_vec(), + } + } else { + SnowflakeAuthMethod::KeyPairUnencrypted { + pem: private_key_pem.clone(), + } + }; + Ok((username.clone(), method)) + } + ClusterAuth::Bearer { token } => Ok(( + String::new(), + SnowflakeAuthMethod::Oauth { + token: token.clone(), + }, + )), + other => Err(format!( + "unsupported auth type for Snowflake: {other:?}. Use basic, keyPair, or bearer." + )), + } +} + +#[async_trait] +impl EngineAdapterTrait for SnowflakeAdapter { + async fn submit_query( + &self, + _sql: &str, + _session: &SessionContext, + _credentials: &QueryCredentials, + _tags: &QueryTags, + ) -> Result { + Err(QueryFluxError::Engine( + "Snowflake requires execute_as_arrow; use the Arrow execution path".to_string(), + )) + } + + async fn poll_query( + &self, + _backend_id: &BackendQueryId, + _next_uri: Option<&str>, + ) -> Result { + Err(QueryFluxError::Engine( + "Snowflake does not support async polling through QueryFlux".to_string(), + )) + } + + async fn cancel_query(&self, _backend_id: &BackendQueryId) -> Result<()> { + Ok(()) + } + + async fn health_check(&self) -> bool { + match self.run_query("SELECT 1").await { + Ok(_) => true, + Err(e) => { + tracing::warn!( + cluster = %self.cluster_name, + error = %e, + "Snowflake health check failed" + ); + false + } + } + } + + fn engine_type(&self) -> EngineType { + EngineType::Snowflake + } + + fn supports_async(&self) -> bool { + false + } + + fn base_url(&self) -> &str { + &self.account + } + + async fn execute_as_arrow( + &self, + sql: &str, + session: &SessionContext, + _credentials: &QueryCredentials, + _tags: &QueryTags, + ) -> Result { + let sf_session = self.client.create_session().await.map_err(|e| { + QueryFluxError::Engine(format!("Snowflake session creation failed: {e}")) + })?; + + // Apply per-query database/schema overrides from the frontend session context. + if let Some(db) = session.database() { + let use_sql = format!("USE DATABASE \"{}\"", db.replace('"', "\"\"")); + sf_session.query(use_sql.as_str()).await.map_err(|e| { + QueryFluxError::Engine(format!("Snowflake USE DATABASE failed: {e}")) + })?; + } + + let rows = sf_session + .query(sql) + .await + .map_err(|e| QueryFluxError::Engine(format!("Snowflake query failed: {e}")))?; + + if rows.is_empty() { + return Ok(Box::pin(stream::empty())); + } + + let col_types = rows[0].column_types(); + let fields: Vec = col_types + .iter() + .map(|c| { + Field::new( + c.name(), + snowflake_type_to_arrow(c.column_type()), + c.column_type().nullable(), + ) + }) + .collect(); + let schema = Arc::new(ArrowSchema::new(fields)); + + let num_cols = schema.fields().len(); + let mut columns: Vec = Vec::with_capacity(num_cols); + for (col_idx, sf_type) in col_types.iter().enumerate() { + let dt = schema.field(col_idx).data_type(); + let col = build_arrow_column(dt, sf_type.column_type(), &rows, col_idx)?; + columns.push(col); + } + + let batch = RecordBatch::try_new(schema, columns) + .map_err(|e| QueryFluxError::Engine(format!("Snowflake RecordBatch failed: {e}")))?; + + Ok(Box::pin(stream::iter(std::iter::once(Ok(batch))))) + } + + async fn list_catalogs(&self) -> Result> { + self.run_first_col("SHOW DATABASES").await + } + + async fn list_databases(&self, catalog: &str) -> Result> { + let sql = format!( + "SHOW SCHEMAS IN DATABASE \"{}\"", + catalog.replace('"', "\"\"") + ); + self.run_first_col(&sql).await + } + + async fn list_tables(&self, catalog: &str, database: &str) -> Result> { + let sql = format!( + "SHOW TABLES IN \"{}\".\"{}\"", + catalog.replace('"', "\"\""), + database.replace('"', "\"\"") + ); + self.run_first_col(&sql).await + } + + async fn describe_table( + &self, + catalog: &str, + database: &str, + table: &str, + ) -> Result> { + let qualified = format!( + "\"{}\".\"{}\".\"{table}\"", + catalog.replace('"', "\"\""), + database.replace('"', "\"\""), + ); + let rows = match self.run_query(&format!("DESCRIBE TABLE {qualified}")).await { + Ok(r) => r, + Err(_) => return Ok(None), + }; + + let columns = rows + .iter() + .filter_map(|row| { + let name: String = row.get::("name").ok()?; + let data_type = row + .get::("type") + .unwrap_or_else(|_| "VARCHAR".to_string()) + .to_uppercase(); + let nullable = row + .get::("null?") + .map(|s| s.to_uppercase() == "Y") + .unwrap_or(true); + Some(queryflux_core::catalog::ColumnDef { + name, + data_type, + nullable, + }) + }) + .collect(); + + Ok(Some(TableSchema { + catalog: catalog.to_string(), + database: database.to_string(), + table: table.to_string(), + columns, + })) + } +} + +impl SnowflakeAdapter { + pub fn descriptor() -> EngineDescriptor { + EngineDescriptor { + engine_key: "snowflake", + display_name: "Snowflake", + description: "Cloud-native data warehouse. Connects via the Snowflake REST API.", + hex: "29B5E8", + connection_type: ConnectionType::Http, + default_port: Some(443), + endpoint_example: Some("https://xy12345.us-east-1.snowflakecomputing.com"), + supported_auth: vec![AuthType::Basic, AuthType::KeyPair, AuthType::Bearer], + implemented: true, + config_fields: vec![ + ConfigField { + key: "account", + label: "Account", + description: "Snowflake account identifier (e.g. xy12345.us-east-1).", + field_type: FieldType::Text, + required: true, + example: Some("xy12345.us-east-1"), + }, + ConfigField { + key: "endpoint", + label: "Endpoint", + description: + "Custom base URL override (e.g. PrivateLink). Omit to derive from account.", + field_type: FieldType::Url, + required: false, + example: Some("https://xy12345.us-east-1.privatelink.snowflakecomputing.com"), + }, + ConfigField { + key: "warehouse", + label: "Warehouse", + description: "Default virtual warehouse for query execution.", + field_type: FieldType::Text, + required: false, + example: Some("COMPUTE_WH"), + }, + ConfigField { + key: "role", + label: "Role", + description: "Default Snowflake role.", + field_type: FieldType::Text, + required: false, + example: Some("ANALYST"), + }, + ConfigField { + key: "catalog", + label: "Database", + description: "Default Snowflake database.", + field_type: FieldType::Text, + required: false, + example: Some("MY_DATABASE"), + }, + ConfigField { + key: "schema", + label: "Schema", + description: "Default Snowflake schema.", + field_type: FieldType::Text, + required: false, + example: Some("PUBLIC"), + }, + ], + } + } +} + +pub struct SnowflakeFactory; + +#[async_trait] +impl crate::EngineAdapterFactory for SnowflakeFactory { + fn engine_key(&self) -> &'static str { + "snowflake" + } + + fn descriptor(&self) -> EngineDescriptor { + SnowflakeAdapter::descriptor() + } + + async fn build_from_config_json( + &self, + cluster_name: ClusterName, + group: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result> { + Ok(Arc::new(SnowflakeAdapter::try_from_config_json( + cluster_name, + group, + json, + cluster_name_str, + )?)) + } +} + +// --------------------------------------------------------------------------- +// Type mapping: Snowflake → Arrow +// --------------------------------------------------------------------------- + +fn snowflake_type_to_arrow(ct: &SnowflakeColumnType) -> DataType { + match ct.snowflake_type().to_ascii_lowercase().as_str() { + "fixed" => { + let scale = ct.scale().unwrap_or(0); + if scale == 0 { + DataType::Int64 + } else { + DataType::Utf8 + } + } + "real" | "float" | "double" => DataType::Float64, + "boolean" => DataType::Boolean, + _ => DataType::Utf8, + } +} + +fn build_arrow_column( + dt: &DataType, + sf_type: &SnowflakeColumnType, + rows: &[SnowflakeRow], + col_idx: usize, +) -> Result { + match dt { + DataType::Boolean => { + let mut b = BooleanBuilder::with_capacity(rows.len()); + for row in rows { + match row.at::(col_idx) { + Ok(v) => b.append_value(v), + Err(_) => b.append_null(), + } + } + Ok(Arc::new(b.finish())) + } + DataType::Int64 => { + let mut b = Int64Builder::with_capacity(rows.len()); + for row in rows { + match row.at::(col_idx) { + Ok(v) => b.append_value(v), + Err(_) => b.append_null(), + } + } + Ok(Arc::new(b.finish())) + } + DataType::Float64 => { + let mut b = Float64Builder::with_capacity(rows.len()); + for row in rows { + match row.at::(col_idx) { + Ok(v) => b.append_value(v), + Err(_) => b.append_null(), + } + } + Ok(Arc::new(b.finish())) + } + _ => { + let _ = sf_type; + let mut b = StringBuilder::with_capacity(rows.len(), rows.len() * 32); + for row in rows { + match row.at::(col_idx) { + Ok(v) => b.append_value(v), + Err(_) => b.append_null(), + } + } + Ok(Arc::new(b.finish())) + } + } +} diff --git a/crates/queryflux-engine-adapters/src/starrocks/mod.rs b/crates/queryflux-engine-adapters/src/starrocks/mod.rs index 5c69dc1..47c0450 100644 --- a/crates/queryflux-engine-adapters/src/starrocks/mod.rs +++ b/crates/queryflux-engine-adapters/src/starrocks/mod.rs @@ -71,6 +71,26 @@ impl StarRocksAdapter { }) } + /// Build from a DB config JSON blob (bypasses the `ClusterConfig` god struct). + pub fn try_from_config_json( + cluster_name: ClusterName, + group_name: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result { + use queryflux_core::engine_registry::{json_str, parse_auth_from_config_json}; + + let endpoint = json_str(json, "endpoint").ok_or_else(|| { + QueryFluxError::Engine(format!("cluster '{cluster_name_str}': missing endpoint")) + })?; + let auth = parse_auth_from_config_json(json); + Self::new(cluster_name, group_name, endpoint, auth).map_err(|e| { + QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': failed to create StarRocks adapter ({e})" + )) + }) + } + /// Build from persisted / YAML [`ClusterConfig`]. pub fn try_from_cluster_config( cluster_name: ClusterName, @@ -561,3 +581,31 @@ impl StarRocksAdapter { } } } + +pub struct StarRocksFactory; + +#[async_trait] +impl crate::EngineAdapterFactory for StarRocksFactory { + fn engine_key(&self) -> &'static str { + "starRocks" + } + + fn descriptor(&self) -> EngineDescriptor { + StarRocksAdapter::descriptor() + } + + async fn build_from_config_json( + &self, + cluster_name: ClusterName, + group: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result> { + Ok(Arc::new(StarRocksAdapter::try_from_config_json( + cluster_name, + group, + json, + cluster_name_str, + )?)) + } +} diff --git a/crates/queryflux-engine-adapters/src/trino/mod.rs b/crates/queryflux-engine-adapters/src/trino/mod.rs index b0051ca..a622db2 100644 --- a/crates/queryflux-engine-adapters/src/trino/mod.rs +++ b/crates/queryflux-engine-adapters/src/trino/mod.rs @@ -93,6 +93,29 @@ impl TrinoAdapter { )) } + /// Build from a DB config JSON blob (bypasses the `ClusterConfig` god struct). + pub fn try_from_config_json( + cluster_name: ClusterName, + group_name: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result { + use queryflux_core::engine_registry::{json_bool, json_str, parse_auth_from_config_json}; + + let endpoint = json_str(json, "endpoint").ok_or_else(|| { + QueryFluxError::Engine(format!("cluster '{cluster_name_str}': missing endpoint")) + })?; + let tls_skip = json_bool(json, "tlsInsecureSkipVerify"); + let auth = parse_auth_from_config_json(json); + Ok(Self::new( + cluster_name, + group_name, + endpoint, + tls_skip, + auth, + )) + } + /// Apply cluster-level auth credentials to a request builder. fn apply_cluster_auth(&self, builder: reqwest::RequestBuilder) -> reqwest::RequestBuilder { match &self.auth { @@ -867,3 +890,31 @@ impl TrinoAdapter { } } } + +pub struct TrinoFactory; + +#[async_trait] +impl crate::EngineAdapterFactory for TrinoFactory { + fn engine_key(&self) -> &'static str { + "trino" + } + + fn descriptor(&self) -> EngineDescriptor { + TrinoAdapter::descriptor() + } + + async fn build_from_config_json( + &self, + cluster_name: ClusterName, + group: ClusterGroupName, + json: &serde_json::Value, + cluster_name_str: &str, + ) -> Result> { + Ok(Arc::new(TrinoAdapter::try_from_config_json( + cluster_name, + group, + json, + cluster_name_str, + )?)) + } +} diff --git a/crates/queryflux-frontend/src/admin.rs b/crates/queryflux-frontend/src/admin.rs index c137c12..eb750bf 100644 --- a/crates/queryflux-frontend/src/admin.rs +++ b/crates/queryflux-frontend/src/admin.rs @@ -169,7 +169,7 @@ impl Default for SecurityConfigDto { /// One entry protocol / client surface (Trino HTTP, MySQL wire, …). #[derive(Debug, Clone, Serialize, ToSchema)] pub struct ProtocolFrontendDto { - /// Stable id: `trino_http`, `mysql_wire`, … + /// Stable id: `trino_http`, `mysql_wire`, `snowflake`, … pub id: String, pub label: String, pub short_description: String, @@ -250,18 +250,18 @@ pub fn build_frontends_status( "Arrow Flight SQL / gRPC-style access (driver-dependent).", frontends.flight_sql.as_ref(), ), - opt_fe( - "snowflake_http", - "Snowflake HTTP (wire)", - "Snowflake internal HTTP wire protocol (JDBC/ODBC/Python/Go/Node.js connectors).", - frontends.snowflake_http.as_ref(), - ), - opt_fe( - "snowflake_sql_api", - "Snowflake SQL API", - "Snowflake REST SQL API v2 (POST /api/v2/statements).", - frontends.snowflake_sql_api.as_ref(), - ), + { + // Single Axum listener: `SnowflakeFrontend` merges connector + SQL API routes; port from YAML `snowflake_http`. + let sf = frontends.snowflake_http.as_ref(); + ProtocolFrontendDto { + id: "snowflake".to_string(), + label: "Snowflake".to_string(), + short_description: "Snowflake-compatible clients on one port: connector session/query flow and REST SQL API v2 (POST /api/v2/statements)." + .to_string(), + enabled: sf.map(|c| c.enabled).unwrap_or(false), + port: sf.map(|c| c.port), + } + }, ]; FrontendsStatusDto { diff --git a/crates/queryflux-frontend/src/dispatch.rs b/crates/queryflux-frontend/src/dispatch.rs index 9f437aa..8bcd938 100644 --- a/crates/queryflux-frontend/src/dispatch.rs +++ b/crates/queryflux-frontend/src/dispatch.rs @@ -14,10 +14,12 @@ use queryflux_core::{ error::{QueryFluxError, Result}, query::{ ClusterGroupName, ClusterName, ExecutingQuery, FrontendProtocol, ProxyQueryId, - QueryExecution, QueryStats, QueryStatus, QueuedQuery, + QueryEngineStats, QueryExecution, QueryStats, QueryStatus, QueuedQuery, }, session::SessionContext, }; +use queryflux_engine_adapters::trino::api::TrinoResponse; +use queryflux_engine_adapters::EngineAdapterTrait; use queryflux_translation::SchemaContext; use tracing::{debug, info, warn}; @@ -247,9 +249,24 @@ pub async fn dispatch_query( }; // Single write per query — no updates needed between polls. // Any QueryFlux instance can serve subsequent polls using this record. - let _ = state.persistence.upsert(executing).await; + let _ = state.persistence.upsert(executing.clone()).await; info!(id = %query_id, backend = %backend_query_id, cluster = %cluster_name, "Query submitted (async)"); + if next_uri.is_none() { + if let Some(ref ib) = initial_body { + finalize_trino_async_terminal_on_submit( + state, + &cluster_manager, + &executing, + &adapter, + &session, + protocol, + ib, + ) + .await; + } + } + // Rewrite nextUri: swap Trino host → QueryFlux external address, keep full path. let proxy_next_uri = next_uri .as_deref() @@ -260,6 +277,139 @@ pub async fn dispatch_query( }) } +/// Trino may return `FINISHED` with no `nextUri` on the initial POST `/v1/statement` response. +/// Clients then never call GET `/v1/statement/...`, so [`crate::trino_http::handlers::get_executing_statement`] +/// never runs — mirror its metrics, `record_query`, and persistence cleanup here. +async fn finalize_trino_async_terminal_on_submit( + state: &Arc, + cluster_manager: &Arc, + executing: &ExecutingQuery, + adapter: &Arc, + session: &SessionContext, + protocol: FrontendProtocol, + body: &Bytes, +) { + let trino_resp: TrinoResponse = match serde_json::from_slice(body.as_ref()) { + Ok(r) => r, + Err(e) => { + warn!( + proxy_id = %executing.id, + "trino submit terminal body JSON parse failed: {e}; releasing cluster + clearing persistence" + ); + state + .metrics + .on_query_finished(&executing.cluster_group.0, &executing.cluster_name.0); + let _ = cluster_manager + .release_cluster(&executing.cluster_group, &executing.cluster_name) + .await; + let _ = state.persistence.delete(&executing.backend_query_id).await; + return; + } + }; + + if trino_resp.next_uri.is_some() { + return; + } + + let elapsed_ms = (Utc::now() - executing.creation_time) + .num_milliseconds() + .max(0) as u64; + + let was_translated = executing.translated_sql.is_some(); + let src_dialect = protocol.default_dialect(); + let ctx = QueryContext { + query_id: &executing.id, + sql: executing + .translated_sql + .as_deref() + .unwrap_or(&executing.sql), + session, + protocol, + group: &executing.cluster_group, + cluster: &executing.cluster_name, + cluster_group_config_id: executing.cluster_group_config_id, + cluster_config_id: executing.cluster_config_id, + engine_type: adapter.engine_type(), + src_dialect, + tgt_dialect: adapter.engine_type().dialect(), + was_translated, + translated_sql: if was_translated { + Some(executing.sql.clone()) + } else { + None + }, + query_tags: executing.query_tags.clone(), + }; + + let engine_stats = Some(QueryEngineStats { + engine_elapsed_time_ms: Some(trino_resp.stats.elapsed_time_millis), + cpu_time_ms: Some(trino_resp.stats.cpu_time_millis), + processed_rows: Some(trino_resp.stats.processed_rows), + processed_bytes: Some(trino_resp.stats.processed_bytes), + physical_input_bytes: Some(trino_resp.stats.physical_input_bytes), + peak_memory_bytes: Some(trino_resp.stats.peak_memory_bytes), + spilled_bytes: Some(trino_resp.stats.spilled_bytes), + total_splits: Some(trino_resp.stats.total_splits), + }); + + let backend_id = Some(executing.backend_query_id.0.clone()); + + if let Some(err) = &trino_resp.error { + state + .metrics + .on_query_finished(&executing.cluster_group.0, &executing.cluster_name.0); + state.record_query( + &ctx, + QueryOutcome { + backend_query_id: backend_id, + status: QueryStatus::Failed, + execution_ms: elapsed_ms, + rows: None, + error: Some(err.message.clone()), + routing_trace: None, + engine_stats, + }, + ); + } else if trino_resp.stats.state == "FAILED" { + state + .metrics + .on_query_finished(&executing.cluster_group.0, &executing.cluster_name.0); + state.record_query( + &ctx, + QueryOutcome { + backend_query_id: backend_id, + status: QueryStatus::Failed, + execution_ms: elapsed_ms, + rows: None, + error: Some("Trino query FAILED".to_string()), + routing_trace: None, + engine_stats, + }, + ); + } else { + state.record_query( + &ctx, + QueryOutcome { + backend_query_id: backend_id, + status: QueryStatus::Success, + execution_ms: elapsed_ms, + rows: None, + error: None, + routing_trace: None, + engine_stats, + }, + ); + state + .metrics + .on_query_finished(&executing.cluster_group.0, &executing.cluster_name.0); + } + + let _ = cluster_manager + .release_cluster(&executing.cluster_group, &executing.cluster_name) + .await; + let _ = state.persistence.delete(&executing.backend_query_id).await; +} + #[allow(clippy::too_many_arguments)] async fn persist_queued_query( state: &Arc, diff --git a/crates/queryflux-persistence/src/cluster_config.rs b/crates/queryflux-persistence/src/cluster_config.rs index d6f0597..e99bcd2 100644 --- a/crates/queryflux-persistence/src/cluster_config.rs +++ b/crates/queryflux-persistence/src/cluster_config.rs @@ -112,7 +112,7 @@ pub struct UpsertClusterGroupConfig { // --------------------------------------------------------------------------- use queryflux_core::config::{ClusterAuth, ClusterConfig, ClusterGroupConfig}; -use queryflux_core::engine_registry::{engine_key, parse_engine_key}; +use queryflux_core::engine_registry::engine_key; impl UpsertClusterConfig { pub fn from_core(cfg: &ClusterConfig) -> Option { @@ -147,6 +147,18 @@ impl UpsertClusterConfig { if let Some(v) = &cfg.catalog { config.insert("catalog".into(), v.clone().into()); } + if let Some(v) = &cfg.account { + config.insert("account".into(), v.clone().into()); + } + if let Some(v) = &cfg.warehouse { + config.insert("warehouse".into(), v.clone().into()); + } + if let Some(v) = &cfg.role { + config.insert("role".into(), v.clone().into()); + } + if let Some(v) = &cfg.schema { + config.insert("schema".into(), v.clone().into()); + } match &cfg.auth { Some(ClusterAuth::Basic { username, password }) => { @@ -221,71 +233,11 @@ impl UpsertClusterGroupConfig { // Conversion helpers: DB records → core config types (for startup loading) // --------------------------------------------------------------------------- -impl ClusterConfigRecord { - pub fn to_core(&self) -> queryflux_core::error::Result { - use queryflux_core::config::TlsConfig; - use queryflux_core::error::QueryFluxError; - - let engine = parse_engine_key(&self.engine_key).map_err(|_| { - QueryFluxError::Engine(format!("Unknown engine key in DB: '{}'", self.engine_key)) - })?; - - // Helpers to extract typed values from the config JSON. - let s = |key: &str| -> Option { - self.config - .get(key) - .and_then(|v| v.as_str()) - .map(String::from) - }; - let b = |key: &str| -> bool { - self.config - .get(key) - .and_then(|v| v.as_bool()) - .unwrap_or(false) - }; - - let auth = match s("authType").as_deref() { - Some("basic") => Some(ClusterAuth::Basic { - username: s("authUsername").unwrap_or_default(), - password: s("authPassword").unwrap_or_default(), - }), - Some("bearer") => Some(ClusterAuth::Bearer { - token: s("authToken").unwrap_or_default(), - }), - Some("accessKey") => Some(ClusterAuth::AccessKey { - access_key_id: s("authUsername").unwrap_or_default(), - secret_access_key: s("authPassword").unwrap_or_default(), - session_token: s("authToken"), - }), - Some("roleArn") => Some(ClusterAuth::RoleArn { - role_arn: s("authUsername").unwrap_or_default(), - external_id: s("authToken"), - }), - _ => None, - }; - - Ok(ClusterConfig { - engine: Some(engine), - enabled: self.enabled, - max_running_queries: self.max_running_queries.map(|v| v as u64), - endpoint: s("endpoint"), - database_path: s("databasePath"), - region: s("region"), - s3_output_location: s("s3OutputLocation"), - workgroup: s("workgroup"), - catalog: s("catalog"), - tls: if b("tlsInsecureSkipVerify") { - Some(TlsConfig { - insecure_skip_verify: true, - }) - } else { - None - }, - auth, - query_auth: None, // not persisted to DB; loaded from YAML config only - }) - } -} +// NOTE: `ClusterConfigRecord::to_core()` has been removed. Engine adapters are +// now built directly from the JSONB config blob via `try_from_config_json()` on +// each adapter. Auth extraction uses `parse_auth_from_config_json()` from +// `queryflux_core::engine_registry`. The persistence layer no longer needs to +// know about engine-specific config fields. impl ClusterGroupConfigRecord { pub fn to_core(&self) -> ClusterGroupConfig { diff --git a/crates/queryflux/src/main.rs b/crates/queryflux/src/main.rs index aa2673d..8dc2a35 100644 --- a/crates/queryflux/src/main.rs +++ b/crates/queryflux/src/main.rs @@ -127,6 +127,10 @@ async fn main() -> Result<()> { // Filled when Postgres loads cluster/group rows — used for query_history FKs on ClusterState. let mut cluster_ids_by_name: HashMap = HashMap::new(); let mut group_ids_by_name: HashMap = HashMap::new(); + // DB cluster records kept for adapter building via build_adapter_from_record. + let mut startup_cluster_records: Option< + Vec, + > = None; // --- When Postgres is active, load cluster/group config from DB --- // Merge YAML-defined clusters and groups into Postgres on **every** startup when the @@ -156,23 +160,48 @@ async fn main() -> Result<()> { // Effective config comes from Postgres (YAML above only upserts keys that appear in the file). info!("Loading cluster and group configs from Postgres"); - let cluster_records = pg + let db_cluster_records = pg .list_cluster_configs() .await .context("Load cluster configs from DB")?; - cluster_ids_by_name = cluster_records + cluster_ids_by_name = db_cluster_records .iter() .map(|r| (r.name.clone(), r.id)) .collect(); - config.clusters = cluster_records - .into_iter() - .map(|r| { - let name = r.name.clone(); - r.to_core() - .map(|c| (name, c)) - .map_err(|e| anyhow::anyhow!("{e}")) + // Build minimal ClusterConfig values for validation and group resolution. + // Engine-specific fields are NOT populated — adapters are built from the + // raw JSONB later via build_adapter_from_record. + config.clusters = db_cluster_records + .iter() + .filter_map(|r| { + let engine = + queryflux_core::engine_registry::parse_engine_key(&r.engine_key).ok()?; + Some(( + r.name.clone(), + queryflux_core::config::ClusterConfig { + engine: Some(engine), + enabled: r.enabled, + max_running_queries: r.max_running_queries.map(|v| v as u64), + endpoint: queryflux_core::engine_registry::json_str(&r.config, "endpoint"), + database_path: None, + region: None, + s3_output_location: None, + workgroup: None, + catalog: None, + account: None, + warehouse: None, + role: None, + schema: None, + tls: None, + auth: queryflux_core::engine_registry::parse_auth_from_config_json( + &r.config, + ), + query_auth: None, + }, + )) }) - .collect::>()?; + .collect(); + startup_cluster_records = Some(db_cluster_records); let group_records = pg .list_group_configs() @@ -279,23 +308,44 @@ async fn main() -> Result<()> { let mut adapters: AdapterMap = HashMap::new(); // Pass 1 — one adapter per cluster. - for (cluster_name_str, cluster_cfg) in &config.clusters { - if !cluster_cfg.enabled { - tracing::info!(cluster = %cluster_name_str, "Cluster disabled — skipping"); - continue; + // DB path: build from JSONB records directly; YAML path: build from ClusterConfig. + if let Some(records) = &startup_cluster_records { + for record in records { + if !record.enabled { + tracing::info!(cluster = %record.name, "Cluster disabled — skipping"); + continue; + } + let cluster_name = ClusterName(record.name.clone()); + let placeholder_group = ClusterGroupName("_".to_string()); + let adapter = registered_engines::build_adapter_from_record( + cluster_name, + placeholder_group, + &record.engine_key, + &record.config, + &record.name, + ) + .await + .with_context(|| format!("Failed to build adapter for cluster '{}'", record.name))?; + adapters.insert(record.name.clone(), adapter); + } + } else { + for (cluster_name_str, cluster_cfg) in &config.clusters { + if !cluster_cfg.enabled { + tracing::info!(cluster = %cluster_name_str, "Cluster disabled — skipping"); + continue; + } + let cluster_name = ClusterName(cluster_name_str.clone()); + let placeholder_group = ClusterGroupName("_".to_string()); + let adapter = registered_engines::build_adapter( + cluster_name, + placeholder_group, + cluster_cfg, + cluster_name_str, + ) + .await + .with_context(|| format!("Failed to build adapter for cluster '{cluster_name_str}'"))?; + adapters.insert(cluster_name_str.clone(), adapter); } - let cluster_name = ClusterName(cluster_name_str.clone()); - let placeholder_group = ClusterGroupName("_".to_string()); - let adapter = registered_engines::build_adapter( - cluster_name, - placeholder_group, - cluster_cfg, - cluster_name_str, - ) - .await - .with_context(|| format!("Failed to build adapter for cluster '{cluster_name_str}'"))?; - - adapters.insert(cluster_name_str.clone(), adapter); } // Pass 2 — one group entry per cluster_group, resolving member cluster names. @@ -589,13 +639,30 @@ async fn main() -> Result<()> { group_translation_scripts, group_default_tags, }; + // Seed the reload cache. When Postgres is active, use the raw JSONB config for + // fingerprinting (same format that build_live_config uses on reload). For YAML-only, + // serialize the ClusterConfig. + let initial_config_json: HashMap = + if let Some(records) = &startup_cluster_records { + records + .iter() + .map(|r| { + ( + r.name.clone(), + serde_json::to_string(&r.config).unwrap_or_default(), + ) + }) + .collect() + } else { + live_config + .cluster_configs + .iter() + .map(|(k, v)| (k.clone(), serde_json::to_string(v).unwrap_or_default())) + .collect() + }; let adapter_reload_cache = Arc::new(tokio::sync::Mutex::new(AdapterReloadCache { adapters: live_config.adapters.clone(), - config_json: live_config - .cluster_configs - .iter() - .map(|(k, v)| (k.clone(), serde_json::to_string(v).unwrap_or_default())) - .collect(), + config_json: initial_config_json, // Seed with the initial cluster states so the first reload can inherit health status. cluster_states: live_config .health_check_targets @@ -1017,15 +1084,17 @@ fn health_targets_from_groups( out } -/// Build a `LiveConfig` from the cluster/group maps and router chain components -/// that were loaded from either YAML or the database. +/// Build a `LiveConfig` from DB cluster records, group maps, and router chain components. +/// +/// This is the DB load path: adapters are built directly from the JSONB config blob +/// in each `ClusterConfigRecord`, bypassing the `ClusterConfig` god struct. /// /// `cache` holds adapter instances from the previous generation. Adapters are reused /// only when the cluster's JSON-serialized config matches the previous reload; otherwise /// they are rebuilt (e.g. endpoint or password changed). #[allow(clippy::too_many_arguments)] async fn build_live_config( - clusters: &std::collections::HashMap, + cluster_records: &[queryflux_persistence::cluster_config::ClusterConfigRecord], cluster_groups: &std::collections::HashMap, cluster_ids_by_name: &HashMap, group_ids_by_name: &HashMap, @@ -1038,31 +1107,46 @@ async fn build_live_config( cluster_state::ClusterState, simple::SimpleClusterGroupManager, strategy::strategy_from_config, }; + use queryflux_core::engine_registry::{json_str, parse_engine_key}; use queryflux_core::tags::QueryTags; + // Build a lookup map from records for group member resolution. + let records_by_name: HashMap< + &str, + &queryflux_persistence::cluster_config::ClusterConfigRecord, + > = cluster_records + .iter() + .map(|r| (r.name.as_str(), r)) + .collect(); + // Build adapters — reuse when serialized cluster config is unchanged. - for (cluster_name_str, cluster_cfg) in clusters { - if !cluster_cfg.enabled { - cache.adapters.remove(cluster_name_str); - cache.config_json.remove(cluster_name_str); + for record in cluster_records { + let cluster_name_str = &record.name; + if !record.enabled { + cache.adapters.remove(cluster_name_str.as_str()); + cache.config_json.remove(cluster_name_str.as_str()); continue; } - let cfg_json = serde_json::to_string(cluster_cfg).unwrap_or_default(); + let cfg_json = serde_json::to_string(&record.config).unwrap_or_default(); let reuse = cache.adapters.contains_key(cluster_name_str.as_str()) - && cache.config_json.get(cluster_name_str).map(String::as_str) + && cache + .config_json + .get(cluster_name_str.as_str()) + .map(String::as_str) == Some(cfg_json.as_str()); if reuse { continue; } - cache.adapters.remove(cluster_name_str); - cache.config_json.remove(cluster_name_str); + cache.adapters.remove(cluster_name_str.as_str()); + cache.config_json.remove(cluster_name_str.as_str()); let cluster_name = ClusterName(cluster_name_str.clone()); let placeholder_group = ClusterGroupName("_".to_string()); - let adapter = match registered_engines::build_adapter( + let adapter = match registered_engines::build_adapter_from_record( cluster_name, placeholder_group, - cluster_cfg, + &record.engine_key, + &record.config, cluster_name_str, ) .await @@ -1076,10 +1160,12 @@ async fn build_live_config( cache.adapters.insert(cluster_name_str.clone(), adapter); cache.config_json.insert(cluster_name_str.clone(), cfg_json); } - cache.adapters.retain(|name, _| clusters.contains_key(name)); + cache + .adapters + .retain(|name, _| records_by_name.contains_key(name.as_str())); cache .config_json - .retain(|name, _| clusters.contains_key(name)); + .retain(|name, _| records_by_name.contains_key(name.as_str())); // Build group states. let mut group_states: GroupStatesMap = HashMap::new(); @@ -1103,8 +1189,8 @@ async fn build_live_config( ); continue; } - let cluster_cfg = match clusters.get(member_name) { - Some(c) => c, + let record = match records_by_name.get(member_name.as_str()) { + Some(r) => r, None => { tracing::warn!(group = %group_name, cluster = %member_name, "Reload: group references unknown cluster"); continue; @@ -1114,15 +1200,17 @@ async fn build_live_config( tracing::info!(group = %group_name, cluster = %member_name, "Reload: skipping disabled/missing cluster in group"); continue; } - let engine = match cluster_cfg.engine.as_ref() { - Some(e) => e, - None => continue, + let engine = match parse_engine_key(&record.engine_key) { + Ok(e) => e, + Err(_) => continue, }; - let engine_type = EngineType::from(engine); - let max_q = cluster_cfg + let engine_type = EngineType::from(&engine); + let max_q = record .max_running_queries + .map(|v| v as u64) .unwrap_or(group_config.max_running_queries); - let cluster_cid = cluster_ids_by_name.get(member_name).copied(); + let endpoint = json_str(&record.config, "endpoint"); + let cluster_cid = cluster_ids_by_name.get(member_name.as_str()).copied(); let group_cid = group_ids_by_name.get(group_name.as_str()).copied(); // Reuse the previous state when the cluster config is unchanged so that @@ -1130,13 +1218,15 @@ async fn build_live_config( // When config changed or the cluster is new, create a fresh state but // still inherit is_healthy from the previous generation so the UI does not // flash healthy for 30 s until the next health-check tick. - let cfg_json = serde_json::to_string(cluster_cfg).unwrap_or_default(); - let config_unchanged = - cache.config_json.get(member_name).map(String::as_str) == Some(cfg_json.as_str()); + let cfg_json = serde_json::to_string(&record.config).unwrap_or_default(); + let config_unchanged = cache + .config_json + .get(member_name.as_str()) + .map(String::as_str) + == Some(cfg_json.as_str()); let state = if config_unchanged { - if let Some(prev) = cache.cluster_states.get(member_name) { - // Config identical — reuse the same Arc to preserve all live state. + if let Some(prev) = cache.cluster_states.get(member_name.as_str()) { prev.clone() } else { Arc::new(ClusterState::new( @@ -1145,9 +1235,9 @@ async fn build_live_config( cluster_cid, group_cid, engine_type, - cluster_cfg.endpoint.clone(), + endpoint, max_q, - cluster_cfg.enabled, + record.enabled, )) } } else { @@ -1157,12 +1247,11 @@ async fn build_live_config( cluster_cid, group_cid, engine_type, - cluster_cfg.endpoint.clone(), + endpoint, max_q, - cluster_cfg.enabled, + record.enabled, )); - // Inherit last known health so the UI doesn't flip to healthy on every reload. - if let Some(prev) = cache.cluster_states.get(member_name) { + if let Some(prev) = cache.cluster_states.get(member_name.as_str()) { s.set_healthy(prev.is_healthy()); } s @@ -1177,13 +1266,42 @@ async fn build_live_config( } let health_check_targets = health_targets_from_groups(&group_states, &cache.adapters); - // Refresh the cached states so the next reload generation can reuse/inherit from these. cache.cluster_states = health_check_targets .iter() .map(|(_, s)| (s.cluster_name.0.clone(), s.clone())) .collect(); let cluster_manager = Arc::new(SimpleClusterGroupManager::new(group_states)); + // Build minimal ClusterConfig values for BackendIdentityResolver. + // query_auth is not persisted to DB, so it's always None for DB-loaded clusters. + let cluster_configs: HashMap = cluster_records + .iter() + .filter_map(|r| { + let engine = parse_engine_key(&r.engine_key).ok()?; + Some(( + r.name.clone(), + queryflux_core::config::ClusterConfig { + engine: Some(engine), + enabled: r.enabled, + max_running_queries: r.max_running_queries.map(|v| v as u64), + endpoint: json_str(&r.config, "endpoint"), + database_path: None, + region: None, + s3_output_location: None, + workgroup: None, + catalog: None, + account: None, + warehouse: None, + role: None, + schema: None, + tls: None, + auth: None, + query_auth: None, + }, + )) + }) + .collect(); + // Build router chain. let fallback = ClusterGroupName(routing_fallback.to_string()); let mut routers: Vec> = Vec::new(); @@ -1296,7 +1414,7 @@ async fn build_live_config( cluster_manager, adapters: cache.adapters.clone(), health_check_targets, - cluster_configs: clusters.clone(), + cluster_configs, group_members, group_order, group_translation_scripts, @@ -1306,6 +1424,8 @@ async fn build_live_config( /// Load cluster/group configs + routing config from Postgres and build a fresh `LiveConfig`. /// Existing adapter instances are reused for clusters that haven't changed. +/// +/// Cluster records are passed directly to `build_live_config` — no `to_core()` conversion. async fn reload_live_config( pg: &Arc, cache: &mut AdapterReloadCache, @@ -1320,16 +1440,6 @@ async fn reload_live_config( .iter() .map(|r| (r.name.clone(), r.id)) .collect(); - let clusters: std::collections::HashMap = - cluster_records - .into_iter() - .map(|r| { - let name = r.name.clone(); - r.to_core() - .map(|c| (name, c)) - .map_err(|e| anyhow::anyhow!("{e}")) - }) - .collect::>()?; let group_records = pg .list_group_configs() @@ -1376,7 +1486,7 @@ async fn reload_live_config( }); build_live_config( - &clusters, + &cluster_records, &cluster_groups, &cluster_ids_by_name, &group_ids_by_name, diff --git a/crates/queryflux/src/registered_engines.rs b/crates/queryflux/src/registered_engines.rs index 77f5e5f..6221a7d 100644 --- a/crates/queryflux/src/registered_engines.rs +++ b/crates/queryflux/src/registered_engines.rs @@ -8,29 +8,64 @@ use queryflux_core::config::{ClusterConfig, EngineConfig}; use queryflux_core::engine_registry::EngineDescriptor; use queryflux_core::error::QueryFluxError; use queryflux_core::query::{ClusterGroupName, ClusterName}; -use queryflux_engine_adapters::athena::AthenaAdapter; -use queryflux_engine_adapters::duckdb::http::DuckDbHttpAdapter; -use queryflux_engine_adapters::duckdb::DuckDbAdapter; -use queryflux_engine_adapters::starrocks::StarRocksAdapter; -use queryflux_engine_adapters::trino::TrinoAdapter; -use queryflux_engine_adapters::EngineAdapterTrait; +use queryflux_engine_adapters::athena::{AthenaAdapter, AthenaFactory}; +use queryflux_engine_adapters::duckdb::http::{DuckDbHttpAdapter, DuckDbHttpFactory}; +use queryflux_engine_adapters::duckdb::{DuckDbAdapter, DuckDbFactory}; +use queryflux_engine_adapters::snowflake::{SnowflakeAdapter, SnowflakeFactory}; +use queryflux_engine_adapters::starrocks::{StarRocksAdapter, StarRocksFactory}; +use queryflux_engine_adapters::trino::{TrinoAdapter, TrinoFactory}; +use queryflux_engine_adapters::{EngineAdapterFactory, EngineAdapterTrait}; -/// All engine descriptors for [`queryflux_core::engine_registry::EngineRegistry`]. -pub fn all_descriptors() -> Vec { +/// All registered engine adapter factories. +/// +/// Adding a new engine means adding its factory here — the rest is driven by +/// the [`EngineAdapterFactory`] trait. +pub fn all_factories() -> Vec> { vec![ - TrinoAdapter::descriptor(), - DuckDbAdapter::descriptor(), - DuckDbHttpAdapter::descriptor(), - StarRocksAdapter::descriptor(), - AthenaAdapter::descriptor(), + Box::new(TrinoFactory), + Box::new(DuckDbFactory), + Box::new(DuckDbHttpFactory), + Box::new(StarRocksFactory), + Box::new(AthenaFactory), + Box::new(SnowflakeFactory), ] } +/// All engine descriptors for [`queryflux_core::engine_registry::EngineRegistry`]. +pub fn all_descriptors() -> Vec { + all_factories().iter().map(|f| f.descriptor()).collect() +} + fn map_qf_err(e: QueryFluxError) -> anyhow::Error { anyhow::Error::new(e) } +/// Build an adapter directly from a DB record's engine key + config JSON blob. +/// +/// This is the DB load path: `JSONB -> adapter`, bypassing the `ClusterConfig` god struct. +/// Looks up the matching [`EngineAdapterFactory`] by `engine_key`. +pub async fn build_adapter_from_record( + cluster_name: ClusterName, + group: ClusterGroupName, + engine_key: &str, + config_json: &serde_json::Value, + cluster_name_str: &str, +) -> Result> { + let factories = all_factories(); + let factory = factories + .iter() + .find(|f| f.engine_key() == engine_key) + .ok_or_else(|| anyhow::anyhow!("Unknown engine key: '{engine_key}'"))?; + + factory + .build_from_config_json(cluster_name, group, config_json, cluster_name_str) + .await + .map_err(map_qf_err) +} + /// Build an adapter for `cluster_cfg`. `cluster_name_str` is used only in error context messages. +/// +/// This is the YAML load path: `ClusterConfig -> adapter`. Kept for backward compatibility. pub async fn build_adapter( cluster_name: ClusterName, placeholder_group: ClusterGroupName, @@ -88,6 +123,15 @@ pub async fn build_adapter( .await .map_err(map_qf_err)?, ), + EngineConfig::Snowflake => Arc::new( + SnowflakeAdapter::try_from_cluster_config( + cluster_name, + placeholder_group, + cluster_cfg, + cluster_name_str, + ) + .map_err(map_qf_err)?, + ), EngineConfig::ClickHouse => { anyhow::bail!("Engine ClickHouse not yet implemented") } diff --git a/development.md b/development.md index 2a0b2a1..43f689f 100644 --- a/development.md +++ b/development.md @@ -80,6 +80,6 @@ Authoritative shapes are the serde types in `queryflux-core` (`config.rs`) and w - **PyO3 / Python not found:** Set `PYO3_PYTHON` to the venv’s `python3` and ensure `make setup` completed. - **Port conflicts:** Adjust ports in `docker/docker-compose.yml` or disable conflicting local services. -- **E2E failures:** Bring up `docker/docker-compose.test.yml` (Trino, StarRocks, Iceberg stack); see `make test-e2e`. +- **E2E failures:** Bring up `docker/test/docker-compose.test.yml` (Trino, StarRocks, Iceberg stack); see `make test-e2e`. For contribution expectations (PRs, tests, docs), see [contribute.md](contribute.md). diff --git a/docker/docker-compose.test.yml b/docker/test/docker-compose.test.yml similarity index 85% rename from docker/docker-compose.test.yml rename to docker/test/docker-compose.test.yml index 29ea39e..d0460c6 100644 --- a/docker/docker-compose.test.yml +++ b/docker/test/docker-compose.test.yml @@ -1,12 +1,12 @@ # Lean compose file for E2E tests. # -# Engines: Trino, StarRocks. +# Engines: Trino, StarRocks, Snowflake (fakesnow mock). # Iceberg stack: Postgres (Lakekeeper backend), MinIO, Lakekeeper. # -# Usage: -# docker compose -f docker-compose.test.yml up -d --wait +# Usage (repo root as project directory): +# docker compose -f docker/test/docker-compose.test.yml --project-directory . up -d --wait # cargo test -p queryflux-e2e-tests -- --include-ignored -# docker compose -f docker-compose.test.yml down +# docker compose -f docker/test/docker-compose.test.yml --project-directory . down networks: test_net: @@ -56,6 +56,25 @@ services: networks: - test_net + # fakesnow: local Snowflake-compatible HTTP server (https://github.com/tekumara/fakesnow). + # Accepts any user/password/account. Used to e2e-test the SnowflakeAdapter. + fakesnow: + image: python:3.12-slim + volumes: + - ./docker/test/fakesnow-entrypoint.sh:/entrypoint.sh:ro + - ./docker/test/fakesnow-apply-patches.py:/fakesnow-apply-patches.py:ro + entrypoint: ["/bin/sh", "/entrypoint.sh"] + ports: + - "18085:8085" + healthcheck: + test: ["CMD-SHELL", "python3 -c \"import socket; s=socket.create_connection(('localhost',8085),2); s.close()\""] + interval: 5s + timeout: 10s + retries: 30 + start_period: 60s + networks: + - test_net + # --------------------------------------------------------------------------- # Iceberg catalog stack (Lakekeeper + MinIO + Postgres) # --------------------------------------------------------------------------- @@ -207,4 +226,3 @@ services: start_period: 0s networks: - test_net - diff --git a/docker/test/fakesnow-apply-patches.py b/docker/test/fakesnow-apply-patches.py new file mode 100644 index 0000000..22507c0 --- /dev/null +++ b/docker/test/fakesnow-apply-patches.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +"""Patch installed fakesnow for snowflake-connector-rs (QueryFlux SnowflakeAdapter). + +1. Login: Rust omits data.SESSION_PARAMETERS. +2. Query: Rust only accepts queryResultFormat=json and rowset[][]; fakesnow defaults to Arrow. +""" +from __future__ import annotations + +import pathlib +import re +import sys + +HELPER = '''def _sf_json_rowset_from_cursor(cur: Any) -> list[list[str | None]]: + """snowflake-connector-rs expects JSON rowset; fakesnow serves Arrow to Python by default.""" + at = cur._arrow_table # noqa: SLF001 + if at is None or at.num_rows == 0: + return [] + names = at.column_names + + def cell(v: object) -> str | None: + if v is None: + return None + if isinstance(v, bool): + return "true" if v else "false" + return str(v) + + return [[cell(rec.get(nm)) for nm in names] for rec in at.to_pylist()] + + +''' + +# Wheel / Black-style indentation (8 spaces for `if`, 12 for body inside query_request). +ARROW_BLOCKS = ( + ( + " if cur._arrow_table: # noqa: SLF001\n" + " batch_bytes = to_ipc(to_sf(cur._arrow_table, rowtype)) # noqa: SLF001\n" + " rowset_b64 = b64encode(batch_bytes).decode(\"utf-8\")\n" + " else:\n" + " rowset_b64 = \"\"\n" + ), + ( + " if cur._arrow_table:\n" + " batch_bytes = to_ipc(to_sf(cur._arrow_table, rowtype))\n" + " rowset_b64 = b64encode(batch_bytes).decode(\"utf-8\")\n" + " else:\n" + " rowset_b64 = \"\"\n" + ), +) + +ARROW_REPLACEMENT = " json_rowset = _sf_json_rowset_from_cursor(cur)\n" + + +def main() -> None: + try: + import fakesnow.server as fsrv # noqa: PLC0415 — only after pip install in container + except ImportError as e: + sys.exit(f"fakesnow.server not importable (pip install fakesnow[server] first): {e}") + path = pathlib.Path(fsrv.__file__) + text = path.read_text() + text = text.replace("\r\n", "\n") + orig = text + + # --- 1) SESSION_PARAMETERS --- + if 'body_json["data"]["SESSION_PARAMETERS"]' in text: + text = text.replace( + 'body_json["data"]["SESSION_PARAMETERS"]', + 'body_json["data"].get("SESSION_PARAMETERS", {})', + 1, + ) + + # --- 2) Helper --- + if "_sf_json_rowset_from_cursor" not in text: + marker = "async def query_request" + if marker not in text: + sys.exit(f"patch: {marker!r} not found in {path}") + text = text.replace(marker, HELPER + marker, 1) + + # --- 3) Arrow -> json_rowset (string replace; more reliable than regex across versions) --- + if "json_rowset = _sf_json_rowset_from_cursor" not in text: + replaced = False + for block in ARROW_BLOCKS: + if block in text: + text = text.replace(block, ARROW_REPLACEMENT, 1) + replaced = True + break + if not replaced: + # Last resort: flexible regex (optional # noqa, flexible spaces) + arrow_re = re.compile( + r"^[ \t]+if cur\._arrow_table:\s*(?:# noqa: SLF001)?\s*\n" + r"[ \t]+batch_bytes = to_ipc\(to_sf\(cur\._arrow_table, rowtype\)\)\s*(?:# noqa: SLF001)?\s*\n" + r"[ \t]+rowset_b64 = b64encode\(batch_bytes\)\.decode\(\"utf-8\"\)\s*\n" + r"[ \t]+else:\s*\n" + r"[ \t]+rowset_b64 = \"\"\s*\n", + re.MULTILINE, + ) + m = arrow_re.search(text) + if not m: + sys.exit( + f"patch: could not find arrow rowset block in {path}. " + "Open an issue or extend ARROW_BLOCKS in docker/test/fakesnow-apply-patches.py" + ) + ind = re.match(r"^([ \t]+)", m.group(0), re.MULTILINE) + prefix = ind.group(1) if ind else " " + text = arrow_re.sub(f"{prefix}json_rowset = _sf_json_rowset_from_cursor(cur)\n", text, count=1) + + # --- 4) Response payload --- + text = text.replace('"rowsetBase64": rowset_b64,', '"rowset": json_rowset,', 1) + text = text.replace('"queryResultFormat": "arrow"', '"queryResultFormat": "json"') + + # --- 5) Validate (connector-rs hard-fails on arrow) --- + if '"queryResultFormat": "arrow"' in text: + sys.exit(f"patch: still contains queryResultFormat arrow after patch: {path}") + if '"rowset": json_rowset' not in text: + sys.exit(f"patch: rowset/json_rowset wiring missing after patch: {path}") + + if text == orig: + print(f"Already patched: {path}", file=sys.stderr) + return + + path.write_text(text) + print(f"Patched {path} for snowflake-connector-rs compatibility", file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/docker/test/fakesnow-entrypoint.sh b/docker/test/fakesnow-entrypoint.sh new file mode 100755 index 0000000..73b3747 --- /dev/null +++ b/docker/test/fakesnow-entrypoint.sh @@ -0,0 +1,11 @@ +#!/bin/sh +# Start fakesnow for e2e. Patches installed package for snowflake-connector-rs (see fakesnow-apply-patches.py). +set -e +if [ ! -r /fakesnow-apply-patches.py ]; then + echo "fakesnow: missing /fakesnow-apply-patches.py (compose must mount docker/test/fakesnow-apply-patches.py)" >&2 + exit 1 +fi +# Pin version so patch anchors stay stable; bump when updating docker/test/fakesnow-apply-patches.py. +pip install --quiet 'fakesnow[server]==0.11.4' +python3 /fakesnow-apply-patches.py +exec fakesnow -s -p 8085 --host 0.0.0.0 diff --git a/queryflux-studio/app/protocols/page.tsx b/queryflux-studio/app/protocols/page.tsx index b833273..fc0abab 100644 --- a/queryflux-studio/app/protocols/page.tsx +++ b/queryflux-studio/app/protocols/page.tsx @@ -21,8 +21,7 @@ const PROTOCOL_SIMPLE_ICONS: Record = { postgres_wire: siPostgresql, clickhouse_http: siClickhouse, flight_sql: siApacheparquet, - snowflake_http: siSnowflake, - snowflake_sql_api: siSnowflake, + snowflake: siSnowflake, }; function SimpleIconSvg({ icon, className }: { icon: SimpleIcon; className?: string }) { diff --git a/queryflux-studio/components/cluster-config/index.ts b/queryflux-studio/components/cluster-config/index.ts index fb4fb80..34d4046 100644 --- a/queryflux-studio/components/cluster-config/index.ts +++ b/queryflux-studio/components/cluster-config/index.ts @@ -1,5 +1,6 @@ export { EngineClusterConfig } from "./engine-cluster-config"; export { AthenaClusterConfig } from "./athena-cluster-config"; +export { SnowflakeClusterConfig } from "./snowflake-cluster-config"; export { TrinoClusterConfig } from "./trino-cluster-config"; export { StarRocksClusterConfig } from "./starrocks-cluster-config"; export { GenericEngineClusterConfig } from "./generic-engine-cluster-config"; diff --git a/queryflux-studio/components/cluster-config/snowflake-cluster-config.tsx b/queryflux-studio/components/cluster-config/snowflake-cluster-config.tsx new file mode 100644 index 0000000..e106301 --- /dev/null +++ b/queryflux-studio/components/cluster-config/snowflake-cluster-config.tsx @@ -0,0 +1,282 @@ +"use client"; + +import type { PatchClusterConfig, FlatClusterConfig } from "./types"; + +const AUTH_BASIC = "basic"; +const AUTH_KEY_PAIR = "keyPair"; +const AUTH_BEARER = "bearer"; + +export function SnowflakeClusterConfig({ + flat, + onPatch, +}: { + flat: FlatClusterConfig; + onPatch: PatchClusterConfig; +}) { + const authType = flat["auth.type"] ?? AUTH_BASIC; + + function setAuthType(next: string) { + onPatch({ + "auth.type": next, + "auth.username": "", + "auth.password": "", + "auth.token": "", + }); + } + + const labelClass = + "block text-[11px] font-semibold text-slate-500 uppercase tracking-wide mb-1.5"; + const inputClass = + "w-full text-sm border border-slate-200 rounded-lg px-3 py-2 font-mono focus:outline-none focus:ring-2 focus:ring-indigo-300 focus:border-indigo-400"; + const hintClass = "text-[10px] text-slate-400 mt-1"; + + return ( +
+

+ Connects to Snowflake via the REST API. Provide your account identifier + and credentials below. +

+ + {/* Account */} +
+ + onPatch({ account: e.target.value })} + placeholder="xy12345.us-east-1" + className={inputClass} + autoComplete="off" + /> +

+ Snowflake account identifier (e.g.{" "} + xy12345.us-east-1). +

+
+ + {/* Endpoint (optional) */} +
+ + onPatch({ endpoint: e.target.value })} + placeholder="https://xy12345.us-east-1.privatelink.snowflakecomputing.com" + className={inputClass} + autoComplete="off" + /> +

+ Custom base URL override (e.g. PrivateLink). Leave empty to derive from account. +

+
+ + {/* Warehouse */} +
+ + onPatch({ warehouse: e.target.value })} + placeholder="COMPUTE_WH" + className={inputClass} + autoComplete="off" + /> +

Default virtual warehouse for query execution.

+
+ + {/* Role */} +
+ + onPatch({ role: e.target.value })} + placeholder="ANALYST" + className={inputClass} + autoComplete="off" + /> +

Default Snowflake role.

+
+ + {/* Database (catalog) */} +
+ + onPatch({ catalog: e.target.value })} + placeholder="MY_DATABASE" + className={inputClass} + autoComplete="off" + /> +

Default Snowflake database.

+
+ + {/* Schema */} +
+ + onPatch({ schema: e.target.value })} + placeholder="PUBLIC" + className={inputClass} + autoComplete="off" + /> +

Default Snowflake schema.

+
+ + {/* Auth type selector */} +
+ + +
+ + {/* Password auth */} + {authType === AUTH_BASIC && ( +
+

+ Password credentials +

+
+ + onPatch({ "auth.username": e.target.value })} + placeholder="SVC_QUERYFLUX" + className={inputClass} + autoComplete="off" + /> +
+
+ + onPatch({ "auth.password": e.target.value })} + className={inputClass} + autoComplete="new-password" + /> +
+
+ )} + + {/* Key Pair auth */} + {authType === AUTH_KEY_PAIR && ( +
+

+ RSA Key Pair +

+
+ + onPatch({ "auth.username": e.target.value })} + placeholder="SVC_QUERYFLUX" + className={inputClass} + autoComplete="off" + /> +
+
+ +