diff --git a/.cargo/config.toml b/.cargo/config.toml new file mode 100644 index 0000000..8a2e879 --- /dev/null +++ b/.cargo/config.toml @@ -0,0 +1,13 @@ +# macOS: prefer `brew install duckdb` (ships libduckdb + pkg-config). Linker search paths +# below match default Homebrew prefixes. Without brew, run `export DUCKDB_DOWNLOAD_LIB=1` +# so libduckdb-sys downloads the official osx-universal dylib. +# +# Linux: do not set global rustflags here; CI installs `libduckdb.so` system-wide and sets +# `DUCKDB_LIB_DIR` / `DUCKDB_INCLUDE_DIR`, while `Makefile` only falls back to +# `DUCKDB_DOWNLOAD_LIB=1` when `DUCKDB_LIB_DIR` is unset. + +[target.aarch64-apple-darwin] +rustflags = ["-L", "/opt/homebrew/lib"] + +[target.x86_64-apple-darwin] +rustflags = ["-L", "/usr/local/lib"] diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f073ada..d59016b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -22,10 +22,12 @@ env: RUST_TOOLCHAIN_VERSION: "1.91.1" CARGO_TERM_COLOR: always RUST_BACKTRACE: 1 - # DuckDB Rust crate: download prebuilt libduckdb in CI (no system libduckdb). - DUCKDB_DOWNLOAD_LIB: "1" - # Do not set RUSTFLAGS to -fuse-ld=mold: mold then fails to resolve -lduckdb from - # libduckdb-sys's rustc-link-search (see mold: fatal: library not found: duckdb). + # libduckdb-sys: point at system-installed lib so rust-lld resolves -lduckdb reliably. + # The "Install libduckdb" step in each job puts the .so + headers there. + DUCKDB_LIB_DIR: /usr/local/lib + DUCKDB_INCLUDE_DIR: /usr/local/include + # Must match duckdb crate version: duckdb 1.85.1 → libduckdb-sys 1.10501.0 → DuckDB v1.5.1. + DUCKDB_VERSION: "1.5.1" jobs: # Fast, no PyO3 / no compile — runs in parallel with clippy + unit tests. @@ -46,6 +48,15 @@ jobs: steps: - uses: actions/checkout@44c2b7a8a4ea60a981eaca3cf939b5f4305c123b # v4.1.5 + - name: Install libduckdb + run: | + curl -fsSL "https://github.com/duckdb/duckdb/releases/download/v${DUCKDB_VERSION}/libduckdb-linux-amd64.zip" -o /tmp/duckdb.zip + unzip -o /tmp/duckdb.zip -d /tmp/duckdb + sudo cp /tmp/duckdb/libduckdb.so /usr/local/lib/ + sudo cp /tmp/duckdb/duckdb.h /usr/local/include/ + sudo cp /tmp/duckdb/duckdb.hpp /usr/local/include/ 2>/dev/null || true + sudo ldconfig + - uses: dtolnay/rust-toolchain@d8352f6b1d2e870bc5716e7a6d9b65c4cc244a1a with: toolchain: ${{ env.RUST_TOOLCHAIN_VERSION }} @@ -76,6 +87,15 @@ jobs: steps: - uses: actions/checkout@44c2b7a8a4ea60a981eaca3cf939b5f4305c123b # v4.1.5 + - name: Install libduckdb + run: | + curl -fsSL "https://github.com/duckdb/duckdb/releases/download/v${DUCKDB_VERSION}/libduckdb-linux-amd64.zip" -o /tmp/duckdb.zip + unzip -o /tmp/duckdb.zip -d /tmp/duckdb + sudo cp /tmp/duckdb/libduckdb.so /usr/local/lib/ + sudo cp /tmp/duckdb/duckdb.h /usr/local/include/ + sudo cp /tmp/duckdb/duckdb.hpp /usr/local/include/ 2>/dev/null || true + sudo ldconfig + - uses: dtolnay/rust-toolchain@d8352f6b1d2e870bc5716e7a6d9b65c4cc244a1a with: toolchain: ${{ env.RUST_TOOLCHAIN_VERSION }} @@ -96,7 +116,8 @@ jobs: echo "PYO3_PYTHON=${GITHUB_WORKSPACE}/.venv/bin/python3" >> "$GITHUB_ENV" echo "PYTHONPATH=$(.venv/bin/python3 -c 'import site; print(site.getsitepackages()[0])')" >> "$GITHUB_ENV" - - run: cargo test --tests --workspace --exclude queryflux-e2e-tests --exclude queryflux-bench + - name: Unit & integration tests (make test) + run: make test e2e: name: E2E tests (Docker) @@ -118,6 +139,15 @@ jobs: large-packages: true swap-storage: true + - name: Install libduckdb + run: | + curl -fsSL "https://github.com/duckdb/duckdb/releases/download/v${DUCKDB_VERSION}/libduckdb-linux-amd64.zip" -o /tmp/duckdb.zip + unzip -o /tmp/duckdb.zip -d /tmp/duckdb + sudo cp /tmp/duckdb/libduckdb.so /usr/local/lib/ + sudo cp /tmp/duckdb/duckdb.h /usr/local/include/ + sudo cp /tmp/duckdb/duckdb.hpp /usr/local/include/ 2>/dev/null || true + sudo ldconfig + - uses: dtolnay/rust-toolchain@d8352f6b1d2e870bc5716e7a6d9b65c4cc244a1a with: toolchain: ${{ env.RUST_TOOLCHAIN_VERSION }} @@ -138,27 +168,17 @@ jobs: echo "PYO3_PYTHON=${GITHUB_WORKSPACE}/.venv/bin/python3" >> "$GITHUB_ENV" echo "PYTHONPATH=$(.venv/bin/python3 -c 'import site; print(site.getsitepackages()[0])')" >> "$GITHUB_ENV" - # E2E compiles deps + integration tests via `cargo test` below. The `queryflux` binary is - # already covered by the clippy job (`--all-targets`). - - - name: Start e2e Docker stack - run: docker compose -f docker/docker-compose.test.yml --project-directory . up -d --wait trino starrocks sentinel + # Same as `make test-e2e` locally (compose up → e2e crate → compose down on success). + - name: E2E tests (make test-e2e) + run: make test-e2e - name: Engine logs (debug on failure) if: failure() run: | - docker compose -f docker/docker-compose.test.yml --project-directory . logs --no-color --tail=200 trino || true - docker compose -f docker/docker-compose.test.yml --project-directory . logs --no-color --tail=200 starrocks || true - - - name: E2E tests - run: | - TRINO_URL=http://localhost:18081 \ - STARROCKS_URL=mysql://root@localhost:9030 \ - LAKEKEEPER_URL=http://localhost:18181 \ - MINIO_ENDPOINT=localhost:19000 \ - DUCKDB_DOWNLOAD_LIB=1 \ - cargo test -p queryflux-e2e-tests -- --test-threads=1 --include-ignored + docker compose -f docker/test/docker-compose.test.yml --project-directory . logs --no-color --tail=200 trino || true + docker compose -f docker/test/docker-compose.test.yml --project-directory . logs --no-color --tail=200 starrocks || true + docker compose -f docker/test/docker-compose.test.yml --project-directory . logs --no-color --tail=200 fakesnow || true - name: Stop e2e Docker stack if: always() - run: docker compose -f docker/docker-compose.test.yml --project-directory . down + run: docker compose -f docker/test/docker-compose.test.yml --project-directory . down diff --git a/.github/workflows/docker-verify.yml b/.github/workflows/docker-verify.yml index 7fa053b..8229098 100644 --- a/.github/workflows/docker-verify.yml +++ b/.github/workflows/docker-verify.yml @@ -6,6 +6,10 @@ name: Docker — QueryFlux (verify) on: pull_request: + paths-ignore: + - 'website/**' + - 'docs/**' + - '*.md' jobs: docker-verify: diff --git a/Cargo.lock b/Cargo.lock index 7059b09..23b8769 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" @@ -857,6 +880,12 @@ dependencies = [ "tokio", ] +[[package]] +name = "base16ct" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf" + [[package]] name = "base64" version = "0.21.7" @@ -946,6 +975,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 +1073,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 +1145,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 +1241,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" @@ -1298,6 +1372,18 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "crypto-bigint" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0dc92fb57ca44df6db8059111ab3af99a63d5d0f8375d9972e319a379c6bab76" +dependencies = [ + "generic-array", + "rand_core 0.6.4", + "subtle", + "zeroize", +] + [[package]] name = "crypto-common" version = "0.1.7" @@ -1329,6 +1415,33 @@ dependencies = [ "memchr", ] +[[package]] +name = "curve25519-dalek" +version = "4.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be" +dependencies = [ + "cfg-if", + "cpufeatures", + "curve25519-dalek-derive", + "digest", + "fiat-crypto", + "rustc_version", + "subtle", + "zeroize", +] + +[[package]] +name = "curve25519-dalek-derive" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "darling" version = "0.20.11" @@ -1504,6 +1617,44 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" +[[package]] +name = "ecdsa" +version = "0.16.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca" +dependencies = [ + "der", + "digest", + "elliptic-curve", + "rfc6979", + "signature", + "spki", +] + +[[package]] +name = "ed25519" +version = "2.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" +dependencies = [ + "pkcs8", + "signature", +] + +[[package]] +name = "ed25519-dalek" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" +dependencies = [ + "curve25519-dalek", + "ed25519", + "serde", + "sha2", + "subtle", + "zeroize", +] + [[package]] name = "either" version = "1.15.0" @@ -1513,6 +1664,27 @@ dependencies = [ "serde", ] +[[package]] +name = "elliptic-curve" +version = "0.13.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5e6043086bf7973472e0c7dff2142ea0b680d30e18d9cc40f267efbf222bd47" +dependencies = [ + "base16ct", + "crypto-bigint", + "digest", + "ff", + "generic-array", + "group", + "hkdf", + "pem-rfc7468", + "pkcs8", + "rand_core 0.6.4", + "sec1", + "subtle", + "zeroize", +] + [[package]] name = "equivalent" version = "1.0.2" @@ -1569,6 +1741,22 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" +[[package]] +name = "ff" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0b50bfb653653f9ca9095b427bed08ab8d75a137839d9ad64eb11810d5b6393" +dependencies = [ + "rand_core 0.6.4", + "subtle", +] + +[[package]] +name = "fiat-crypto" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d" + [[package]] name = "filetime" version = "0.2.27" @@ -1815,6 +2003,7 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", + "zeroize", ] [[package]] @@ -1875,6 +2064,17 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "group" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63" +dependencies = [ + "ff", + "rand_core 0.6.4", + "subtle", +] + [[package]] name = "h2" version = "0.3.27" @@ -2338,6 +2538,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" @@ -2466,6 +2676,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" @@ -2473,11 +2698,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0529410abe238729a60b108898784df8984c87f6054c9c4fcacc47e4803c1ce1" dependencies = [ "base64 0.22.1", + "ed25519-dalek", "getrandom 0.2.17", + "hmac", "js-sys", + "p256", + "p384", "pem", + "rand 0.8.5", + "rsa", "serde", "serde_json", + "sha2", "signature", "simple_asn1", ] @@ -3066,6 +3298,30 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" +[[package]] +name = "p256" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9863ad85fa8f4460f9c48cb909d38a0d689dba1f6f6988a5e3e0d31071bcd4b" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + +[[package]] +name = "p384" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe42f1670a52a47d448f14b6a5c61dd78fce51856e68edaa38f7ae3a46b8d6b6" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + [[package]] name = "parking" version = "2.2.1" @@ -3101,6 +3357,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" @@ -3187,6 +3453,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" @@ -3194,6 +3475,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" dependencies = [ "der", + "pkcs5", + "rand_core 0.6.4", "spki", ] @@ -3249,6 +3532,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "primeorder" +version = "0.13.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "353e1ca18966c16d9deb1c69278edbc5f194139612772bd9537af60ac231e1e6" +dependencies = [ + "elliptic-curve", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -3449,7 +3741,7 @@ version = "0.1.0" dependencies = [ "async-trait", "dashmap", - "jsonwebtoken", + "jsonwebtoken 10.3.0", "ldap3", "queryflux-core", "reqwest 0.12.28", @@ -3527,9 +3819,11 @@ dependencies = [ "reqwest 0.12.28", "serde", "serde_json", + "snowflake-connector-rs", "tokio", "tracing-subscriber", "trino-rust-client", + "url", ] [[package]] @@ -3551,9 +3845,11 @@ dependencies = [ "reqwest 0.12.28", "serde", "serde_json", + "snowflake-connector-rs", "tokio", "tokio-stream", "tracing", + "url", ] [[package]] @@ -3565,9 +3861,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 +3878,10 @@ dependencies = [ "queryflux-routing", "queryflux-translation", "reqwest 0.12.28", + "rsa", "serde", "serde_json", + "sha2", "tokio", "tokio-stream", "tonic", @@ -3587,6 +3889,7 @@ dependencies = [ "tracing", "urlencoding", "utoipa", + "uuid", ] [[package]] @@ -3921,6 +4224,16 @@ dependencies = [ "web-sys", ] +[[package]] +name = "rfc6979" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2" +dependencies = [ + "hmac", + "subtle", +] + [[package]] name = "ring" version = "0.17.14" @@ -4157,6 +4470,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" @@ -4187,6 +4509,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" @@ -4203,6 +4536,20 @@ version = "4.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1c107b6f4780854c8b126e228ea8869f4d7b71260f962fefb57b996b8959ba6b" +[[package]] +name = "sec1" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc" +dependencies = [ + "base16ct", + "der", + "generic-array", + "pkcs8", + "subtle", + "zeroize", +] + [[package]] name = "security-framework" version = "3.7.0" @@ -4413,6 +4760,28 @@ dependencies = [ "serde", ] +[[package]] +name = "snowflake-connector-rs" +version = "0.9.0" +source = "git+https://github.com/lakeops-org/snowflake-connector-rs?branch=main#e828836333508f68f8ffaf801ffd632d80eda0bb" +dependencies = [ + "base64 0.22.1", + "chrono", + "flate2", + "http 1.4.0", + "jsonwebtoken 10.3.0", + "pkcs8", + "reqwest 0.12.28", + "rsa", + "serde", + "serde_json", + "sha2", + "thiserror 2.0.18", + "tokio", + "url", + "uuid", +] + [[package]] name = "socket2" version = "0.5.10" @@ -5065,13 +5434,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/Cargo.toml b/Cargo.toml index 18c7a7f..5a709f3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -81,3 +81,6 @@ mysql_async = { version = "0.36.1", default-features = false, features = ["defau # OpenAPI documentation utoipa = { version = "5.4.0", features = ["axum_extras", "chrono"] } + +# Snowflake engine backend — lakeops fork includes `QueryExecutor::snowflake_columns` for empty rowsets. +snowflake-connector-rs = { git = "https://github.com/lakeops-org/snowflake-connector-rs", branch = "main" } diff --git a/Makefile b/Makefile index b7686ad..3418f1d 100644 --- a/Makefile +++ b/Makefile @@ -1,6 +1,18 @@ 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 . + +# site-packages for `.venv` (any Python 3.x); expanded when recipes run after `make setup`. +PYTHONPATH_VENV = $(shell .venv/bin/python3 -c 'import site; print(site.getsitepackages()[0])') + +# DuckDB: on Linux (local dev, not CI), let libduckdb-sys download the prebuilt .so. +# CI installs libduckdb system-wide instead (see .github/workflows/ci.yml). +# On macOS use `brew install duckdb` (see .cargo/config.toml) or `export DUCKDB_DOWNLOAD_LIB=1`. +ifeq ($(shell uname -s),Linux) + ifndef DUCKDB_LIB_DIR + export DUCKDB_DOWNLOAD_LIB := 1 + endif +endif # 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). @@ -27,9 +39,8 @@ env: server: PYO3_PYTHON=$(shell pwd)/.venv/bin/python3 \ - PYTHONPATH=$(shell pwd)/.venv/lib/python3.13/site-packages \ + PYTHONPATH=$(PYTHONPATH_VENV) \ RUST_LOG=queryflux=info,queryflux_frontend=info \ - DUCKDB_DOWNLOAD_LIB=1 \ $(CARGO) run --bin queryflux -- --config config.local.yaml ## Stop Docker services and any running QueryFlux process stop: @@ -50,42 +61,44 @@ clippy: $(CARGO) clippy --all-targets --all-features -- -D warnings ## Run unit/integration tests (no external services needed). +## Same command as CI `.github/workflows/ci.yml` (`make test`). ## PYO3_PYTHON + PYTHONPATH: PyO3 (routing + translation). The venv must include `sqlglot` ## (`pip install -r requirements.txt` via `make setup`) for `queryflux-translation` transform tests. test: @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 \ - $(CARGO) test --tests --workspace --exclude queryflux-e2e-tests + PYTHONPATH=$(PYTHONPATH_VENV) \ + $(CARGO) test --tests --workspace --exclude queryflux-e2e-tests --exclude queryflux-bench ## Micro-benchmark: mock Trino + StarRocks backends vs QueryFlux (release build). ## Optional: QUERYFLUX_BENCH_WARMUP, QUERYFLUX_BENCH_ITERATIONS, QUERYFLUX_BENCH_TRINO_POLL. benchmark: @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 \ + PYTHONPATH=$(PYTHONPATH_VENV) \ $(CARGO) build --release --bin queryflux PYO3_PYTHON=$(shell pwd)/.venv/bin/python3 \ - PYTHONPATH=$(shell pwd)/.venv/lib/python3.13/site-packages \ + PYTHONPATH=$(PYTHONPATH_VENV) \ $(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. +## Run E2E tests. Spins up Trino + StarRocks + fakesnow + Lakekeeper via Docker. +## Same command as CI `.github/workflows/ci.yml` (`make test-e2e`). +## 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). 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 + PYTHONPATH=$(PYTHONPATH_VENV) \ + $(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 \ + PYTHONPATH=$(PYTHONPATH_VENV) \ 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 \ $(CARGO) test -p queryflux-e2e-tests --manifest-path Cargo.toml -- --test-threads=1 --include-ignored --nocapture $(COMPOSE_TEST) down 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-cluster-manager/src/cluster_state.rs b/crates/queryflux-cluster-manager/src/cluster_state.rs index 0246f69..3e9f532 100644 --- a/crates/queryflux-cluster-manager/src/cluster_state.rs +++ b/crates/queryflux-cluster-manager/src/cluster_state.rs @@ -95,6 +95,10 @@ impl ClusterState { self.running_queries.store(clamped, Ordering::Relaxed); } + pub fn set_queued_queries(&self, count: u64) { + self.queued_queries.store(count, Ordering::Relaxed); + } + pub fn increment_running(&self) { self.running_queries.fetch_add(1, Ordering::Relaxed); } diff --git a/crates/queryflux-core/src/config.rs b/crates/queryflux-core/src/config.rs index 2c4da40..abd9617 100644 --- a/crates/queryflux-core/src/config.rs +++ b/crates/queryflux-core/src/config.rs @@ -292,6 +292,14 @@ pub struct QueryFluxConfig { pub query_history_retention_days: Option, #[serde(default)] pub metrics: MetricsConfig, + /// When `true`, refuse to start if the Snowflake HTTP frontend is enabled but + /// [`FrontendConfig::session_affinity_acknowledged`] is not set on that frontend. + /// + /// Use this in multi-replica deployments after configuring the load balancer for + /// **session affinity** (sticky routing) to one QueryFlux instance per Snowflake client token. + /// Default `false` keeps single-replica and dev setups unchanged. + #[serde(default)] + pub enforce_snowflake_http_session_affinity: bool, } impl QueryFluxConfig { @@ -322,6 +330,8 @@ pub struct FrontendsConfig { pub clickhouse_http: Option, #[serde(default)] pub flight_sql: Option, + #[serde(default)] + pub snowflake_http: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -330,6 +340,26 @@ pub struct FrontendConfig { #[serde(default = "default_true")] pub enabled: bool, pub port: u16, + /// Operator assertion: the load balancer provides **session affinity** for this frontend + /// so a given client’s connections reach the same QueryFlux instance. + /// + /// **Snowflake HTTP wire protocol** (`snowflakeHttp`) keeps login sessions in process memory; + /// without sticky routing, requests after login may hit another replica and fail with + /// “session not found”. Other frontends ignore this flag. + /// + /// Set to `true` only after configuring affinity (e.g. consistent hash on `Authorization` / + /// Snowflake token, or cookie-based stickiness). Pair with + /// [`QueryFluxConfig::enforce_snowflake_http_session_affinity`] to fail fast if unset. + #[serde(default)] + pub session_affinity_acknowledged: bool, + /// **Snowflake HTTP wire only** — max session lifetime in seconds (from login). + /// Omitted → 86400 (24h). `0` = no max-age limit. + #[serde(default)] + pub snowflake_session_max_age_secs: Option, + /// **Snowflake HTTP wire only** — idle timeout in seconds since the last validated request + /// (heartbeat, token refresh, query). Omitted → 14400 (4h). `0` = no idle limit. + #[serde(default)] + pub snowflake_session_idle_timeout_secs: Option, } impl Default for FrontendConfig { @@ -337,6 +367,9 @@ impl Default for FrontendConfig { Self { enabled: true, port: 8080, + session_affinity_acknowledged: false, + snowflake_session_max_age_secs: None, + snowflake_session_idle_timeout_secs: None, } } } @@ -505,6 +538,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)] @@ -530,8 +565,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. @@ -648,6 +696,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/engine_registry.rs b/crates/queryflux-core/src/engine_registry.rs index 71fbbd1..f94cd54 100644 --- a/crates/queryflux-core/src/engine_registry.rs +++ b/crates/queryflux-core/src/engine_registry.rs @@ -10,7 +10,7 @@ use serde::Serialize; -use crate::config::{ClusterAuth, ClusterConfig, EngineConfig}; +use crate::config::{ClusterAuth, ClusterConfig, EngineConfig, QueryAuthConfig}; use crate::query::EngineType; // --------------------------------------------------------------------------- @@ -237,6 +237,76 @@ 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. +/// +/// - Missing / empty `authType` → `Ok(None)`. +/// - Known `authType` with missing required fields → `Err` (so callers fail fast instead of +/// building adapters with empty credentials). +pub fn parse_auth_from_config_json( + json: &serde_json::Value, +) -> Result, String> { + let s = + |key: &str| -> Option { json.get(key).and_then(|v| v.as_str()).map(String::from) }; + let require = |key: &str| -> Result { + s(key) + .filter(|v| !v.is_empty()) + .ok_or_else(|| format!("missing or empty '{key}' for this authType")) + }; + match s("authType").as_deref() { + None | Some("") => Ok(None), + Some("basic") => Ok(Some(ClusterAuth::Basic { + username: require("authUsername")?, + password: require("authPassword")?, + })), + Some("bearer") => Ok(Some(ClusterAuth::Bearer { + token: require("authToken")?, + })), + Some("keyPair") => Ok(Some(ClusterAuth::KeyPair { + username: require("authUsername")?, + private_key_pem: require("authPassword")?, + private_key_passphrase: s("authToken"), + })), + Some("accessKey") => Ok(Some(ClusterAuth::AccessKey { + access_key_id: require("authUsername")?, + secret_access_key: require("authPassword")?, + session_token: s("authToken"), + })), + Some("roleArn") => Ok(Some(ClusterAuth::RoleArn { + role_arn: require("authUsername")?, + external_id: s("authToken"), + })), + Some(other) => Err(format!("unsupported authType: '{other}'")), + } +} + +/// Extract per-query auth (`queryAuth` / Type 2) from the cluster `config` JSONB blob. +/// +/// Same JSON shape as YAML `queryAuth` on [`ClusterConfig`] (written on upsert from YAML +/// and preserved in Postgres `cluster_configs.config`). +pub fn parse_query_auth_from_config_json(json: &serde_json::Value) -> Option { + json.get("queryAuth") + .filter(|v| !v.is_null()) + .and_then(|v| serde_json::from_value::(v.clone()).ok()) +} + +/// 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 +321,7 @@ pub fn engine_key(engine: &EngineConfig) -> &'static str { EngineConfig::StarRocks => "starRocks", EngineConfig::ClickHouse => "clickHouse", EngineConfig::Athena => "athena", + EngineConfig::Snowflake => "snowflake", } } @@ -263,6 +334,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 +348,26 @@ impl From<&EngineConfig> for EngineType { EngineConfig::StarRocks => EngineType::StarRocks, EngineConfig::ClickHouse => EngineType::ClickHouse, EngineConfig::Athena => EngineType::Athena, + EngineConfig::Snowflake => EngineType::Snowflake, } } } + +#[cfg(test)] +mod query_auth_parse_tests { + use super::*; + use crate::config::QueryAuthConfig; + + #[test] + fn parse_query_auth_impersonate() { + let blob = serde_json::json!({ "queryAuth": { "type": "impersonate" } }); + let parsed = parse_query_auth_from_config_json(&blob).unwrap(); + assert!(matches!(parsed, QueryAuthConfig::Impersonate)); + } + + #[test] + fn parse_query_auth_omitted_is_none() { + let blob = serde_json::json!({ "endpoint": "http://t:8080" }); + assert!(parse_query_auth_from_config_json(&blob).is_none()); + } +} diff --git a/crates/queryflux-core/src/query.rs b/crates/queryflux-core/src/query.rs index 5f04f92..b7f6a23 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::Snowflake, + FrontendProtocol::SnowflakeSqlApi => SqlDialect::Snowflake, } } } @@ -93,6 +99,8 @@ pub enum EngineType { ClickHouse, /// Amazon Athena (Presto/Trino-compatible SQL over S3). Athena, + /// Snowflake cloud data warehouse. + Snowflake, } impl EngineType { @@ -103,6 +111,7 @@ impl EngineType { EngineType::DuckDb | EngineType::DuckDbHttp => SqlDialect::DuckDb, EngineType::StarRocks => SqlDialect::StarRocks, EngineType::ClickHouse => SqlDialect::ClickHouse, + EngineType::Snowflake => SqlDialect::Snowflake, } } } @@ -115,6 +124,7 @@ pub enum SqlDialect { DuckDb, StarRocks, ClickHouse, + Snowflake, MySql, Postgres, Generic, @@ -140,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/Cargo.toml b/crates/queryflux-e2e-tests/Cargo.toml index 34bfcc1..e884dd8 100644 --- a/crates/queryflux-e2e-tests/Cargo.toml +++ b/crates/queryflux-e2e-tests/Cargo.toml @@ -43,3 +43,7 @@ tracing-subscriber = { workspace = true } # Async trait async-trait = { workspace = true } +# Snowflake wire client — same dependency as engine-adapters / fakesnow e2e +snowflake-connector-rs = { workspace = true } +url = { workspace = true } + diff --git a/crates/queryflux-e2e-tests/src/harness.rs b/crates/queryflux-e2e-tests/src/harness.rs index b440b41..101468f 100644 --- a/crates/queryflux-e2e-tests/src/harness.rs +++ b/crates/queryflux-e2e-tests/src/harness.rs @@ -23,14 +23,17 @@ 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, SnowflakeFrontend}, state::LiveConfig, trino_http::{state::AppState, TrinoHttpFrontend}, }; @@ -58,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"; @@ -177,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." )); } @@ -230,9 +290,12 @@ 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(Default::default()), }); - let router: Router = TrinoHttpFrontend::new(state, port).router(); + let trino_fe = TrinoHttpFrontend::new(state.clone(), port); + let snowflake_fe = SnowflakeFrontend::new(state, port); + let router: Router = trino_fe.router().merge(snowflake_fe.router()); let listener = TcpListener::bind(format!("127.0.0.1:{port}")).await?; let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>(); tokio::spawn(async move { @@ -288,7 +351,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()); } @@ -332,3 +395,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/lib.rs b/crates/queryflux-e2e-tests/src/lib.rs index e1179e0..c79115a 100644 --- a/crates/queryflux-e2e-tests/src/lib.rs +++ b/crates/queryflux-e2e-tests/src/lib.rs @@ -1,3 +1,4 @@ pub mod harness; pub mod iceberg_seed; +pub mod snowflake_rs_client; pub mod trino_client; diff --git a/crates/queryflux-e2e-tests/src/snowflake_rs_client.rs b/crates/queryflux-e2e-tests/src/snowflake_rs_client.rs new file mode 100644 index 0000000..9013263 --- /dev/null +++ b/crates/queryflux-e2e-tests/src/snowflake_rs_client.rs @@ -0,0 +1,43 @@ +//! Build a [`snowflake_connector_rs::SnowflakeClient`] pointed at fakesnow (same defaults as +//! [`crate::harness::TestHarness`] Snowflake cluster config). + +use std::time::Duration; + +use anyhow::Context; +use snowflake_connector_rs::{ + SnowflakeAuthMethod, SnowflakeClient, SnowflakeClientConfig, SnowflakeEndpointConfig, + SnowflakeQueryConfig, SnowflakeSession, SnowflakeSessionConfig, +}; +use url::Url; + +/// Matches [`crate::harness`] / docker `FAKESNOW_URL` default. +pub fn fakesnow_http_url() -> String { + std::env::var("FAKESNOW_URL").unwrap_or_else(|_| "http://localhost:18085".to_string()) +} + +/// Login used by the harness `SnowflakeAdapter::try_from_cluster_config` fakesnow block. +pub fn fakesnow_login() -> (&'static str, &'static str, &'static str) { + ("fake", "snow", "fakesnow") +} + +/// New authenticated session against fakesnow. +pub async fn fakesnow_session() -> anyhow::Result { + let base = fakesnow_http_url(); + let url = Url::parse(&base).with_context(|| format!("invalid FAKESNOW_URL: {base}"))?; + let (user, pass, account) = fakesnow_login(); + + let session_cfg = SnowflakeSessionConfig::default(); + let query_cfg = SnowflakeQueryConfig::default() + .with_async_query_completion_timeout(Duration::from_secs(120)); + + let cfg = SnowflakeClientConfig::new(user, account, SnowflakeAuthMethod::Password(pass.into())) + .with_session(session_cfg) + .with_query(query_cfg) + .with_endpoint(SnowflakeEndpointConfig::custom_base_url(url)); + + let client = SnowflakeClient::new(cfg).map_err(|e| anyhow::anyhow!("{e}"))?; + client + .create_session() + .await + .map_err(|e| anyhow::anyhow!("{e}")) +} diff --git a/crates/queryflux-e2e-tests/src/trino_client.rs b/crates/queryflux-e2e-tests/src/trino_client.rs index 9424934..8ca321a 100644 --- a/crates/queryflux-e2e-tests/src/trino_client.rs +++ b/crates/queryflux-e2e-tests/src/trino_client.rs @@ -30,14 +30,19 @@ 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)] +#[serde(default)] struct TrinoStats { + #[allow(dead_code)] state: String, } @@ -113,10 +118,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/routing_tests.rs b/crates/queryflux-e2e-tests/tests/routing_tests.rs index 90f72bb..02cb7a8 100644 --- a/crates/queryflux-e2e-tests/tests/routing_tests.rs +++ b/crates/queryflux-e2e-tests/tests/routing_tests.rs @@ -1,7 +1,7 @@ /// Header routing and Trino↔StarRocks targeting via `X-Qf-Group`. /// /// Requires at least one backend (same as [`TestHarness::new`]). Run with -/// `make test-e2e` or after `docker compose -f docker/docker-compose.test.yml up`. +/// `make test-e2e` or after `docker compose -f docker/test/docker-compose.test.yml up`. /// /// Run: cargo test -p queryflux-e2e-tests --test routing_tests use std::sync::OnceLock; diff --git a/crates/queryflux-e2e-tests/tests/snowflake_connector_rs_tests.rs b/crates/queryflux-e2e-tests/tests/snowflake_connector_rs_tests.rs new file mode 100644 index 0000000..d2d99fe --- /dev/null +++ b/crates/queryflux-e2e-tests/tests/snowflake_connector_rs_tests.rs @@ -0,0 +1,111 @@ +//! `snowflake-connector-rs` against **fakesnow** (same stack as `snowflake_tests.rs` via QueryFlux +//! adapter). Exercises `session::query`, `execute` + `fetch_all`, chunked `fetch_next_chunk`, and +//! `QueryExecutor::snowflake_columns` on an empty rowset. +//! +//! Run with fakesnow up (e.g. docker compose) and: +//! cargo test -p queryflux-e2e-tests --test snowflake_connector_rs_tests -- --include-ignored + +use std::sync::OnceLock; + +use queryflux_e2e_tests::{ + harness::{TestHarness, GROUP_SNOWFLAKE}, + snowflake_rs_client::fakesnow_session, +}; + +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") + }) +} + +macro_rules! require_snowflake { + () => { + if !harness().has_group(GROUP_SNOWFLAKE) { + eprintln!( + "SKIP: fakesnow not available (set FAKESNOW_URL or start docker/test/docker-compose.test.yml)" + ); + return; + } + }; +} + +#[tokio::test] +#[ignore = "requires fakesnow — run with: make test-e2e"] +async fn snowflake_rs_session_query_fetchall() { + require_snowflake!(); + let session = fakesnow_session().await.expect("fakesnow_session"); + let rows = session + .query("SELECT 1 + 1 AS result") + .await + .expect("query"); + assert_eq!(rows.len(), 1, "expected one row"); + let v: i64 = rows[0].get("RESULT").expect("RESULT column"); + assert_eq!(v, 2); +} + +#[tokio::test] +#[ignore = "requires fakesnow — run with: make test-e2e"] +async fn snowflake_rs_execute_fetch_all() { + require_snowflake!(); + let session = fakesnow_session().await.expect("fakesnow_session"); + let sql = "SELECT 'rs-fetch-all' AS label"; + let executor = session.execute(sql).await.expect("execute"); + let rows = executor.fetch_all().await.expect("fetch_all"); + assert_eq!(rows.len(), 1); + let label: String = rows[0].get("LABEL").expect("LABEL"); + assert_eq!(label, "rs-fetch-all"); +} + +#[tokio::test] +#[ignore = "requires fakesnow — run with: make test-e2e"] +async fn snowflake_rs_execute_fetch_next_chunk_matches_fetch_all_row_count() { + require_snowflake!(); + let session = fakesnow_session().await.expect("fakesnow_session"); + let sql = "SELECT 1 AS v UNION ALL SELECT 2 UNION ALL SELECT 3"; + + let all = { + let ex = session.execute(sql).await.expect("execute"); + ex.fetch_all().await.expect("fetch_all") + }; + assert_eq!(all.len(), 3); + + let ex2 = session.execute(sql).await.expect("execute"); + let mut total = 0usize; + loop { + let chunk = ex2.fetch_next_chunk().await.expect("fetch_next_chunk"); + let Some(part) = chunk else { break }; + total += part.len(); + } + assert_eq!(total, 3, "sum of chunk row counts should match fetch_all"); +} + +#[tokio::test] +#[ignore = "requires fakesnow — run with: make test-e2e"] +async fn snowflake_rs_snowflake_columns_when_no_rows() { + require_snowflake!(); + let session = fakesnow_session().await.expect("fakesnow_session"); + let executor = session + .execute("SELECT 99 AS n WHERE 1 = 0") + .await + .expect("execute"); + + let cols = executor.snowflake_columns(); + assert_eq!( + cols.len(), + 1, + "metadata must be present even when rowset is empty" + ); + assert_eq!(cols[0].name(), "N"); + + let rows = executor.fetch_all().await.expect("fetch_all"); + assert!(rows.is_empty(), "no data rows expected for WHERE 1 = 0"); +} 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..6acbbe2 --- /dev/null +++ b/crates/queryflux-e2e-tests/tests/snowflake_tests.rs @@ -0,0 +1,215 @@ +/// 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(); + let marker = format!( + "qf_metric_{}_{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0) + ); + let sql = format!("SELECT 999 AS metric_test -- {marker}"); + c.execute_on(&sql, GROUP_SNOWFLAKE).await.expect("query"); + + let record = h + .wait_for_record(|r| { + r.cluster_group.0 == GROUP_SNOWFLAKE && r.sql_preview.contains(&marker) + }) + .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..09e3472 100644 --- a/crates/queryflux-engine-adapters/Cargo.toml +++ b/crates/queryflux-engine-adapters/Cargo.toml @@ -13,7 +13,7 @@ serde = { workspace = true } serde_json = { workspace = true } reqwest = { workspace = true } tracing = { workspace = true } -duckdb = { version = "1.85.1"} +duckdb = { version = "1.85.1" } mysql_async = { workspace = true } arrow = { workspace = true } futures = { workspace = true } @@ -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 = { workspace = true } +url = "2" diff --git a/crates/queryflux-engine-adapters/src/athena/mod.rs b/crates/queryflux-engine-adapters/src/athena/mod.rs index c70417b..78a6195 100644 --- a/crates/queryflux-engine-adapters/src/athena/mod.rs +++ b/crates/queryflux-engine-adapters/src/athena/mod.rs @@ -169,6 +169,56 @@ 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).map_err(|e| { + QueryFluxError::Engine(format!("cluster '{cluster_name_str}': invalid auth ({e})")) + })?; + let auth = match auth { + Some(a @ ClusterAuth::AccessKey { .. }) | Some(a @ ClusterAuth::RoleArn { .. }) => { + Some(a) + } + Some(_) => { + return Err(QueryFluxError::Engine(format!( + "cluster '{cluster_name_str}': Athena only supports accessKey or roleArn auth" + ))); + } + None => None, + }; + 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 +779,28 @@ 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, + ) -> Result> { + let name = cluster_name.0.clone(); + Ok(Arc::new( + AthenaAdapter::try_from_config_json(cluster_name, group, json, name.as_str()).await?, + )) + } +} diff --git a/crates/queryflux-engine-adapters/src/duckdb/http.rs b/crates/queryflux-engine-adapters/src/duckdb/http.rs index 6d4d389..d745204 100644 --- a/crates/queryflux-engine-adapters/src/duckdb/http.rs +++ b/crates/queryflux-engine-adapters/src/duckdb/http.rs @@ -140,6 +140,29 @@ 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).map_err(|e| { + QueryFluxError::Engine(format!("cluster '{cluster_name_str}': invalid auth ({e})")) + })?; + 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 +569,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, + ) -> Result> { + let name = cluster_name.0.clone(); + Ok(Arc::new(DuckDbHttpAdapter::try_from_config_json( + cluster_name, + group, + json, + name.as_str(), + )?)) + } +} diff --git a/crates/queryflux-engine-adapters/src/duckdb/mod.rs b/crates/queryflux-engine-adapters/src/duckdb/mod.rs index e554c22..27a1513 100644 --- a/crates/queryflux-engine-adapters/src/duckdb/mod.rs +++ b/crates/queryflux-engine-adapters/src/duckdb/mod.rs @@ -67,6 +67,35 @@ 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).map_err(|e| { + QueryFluxError::Engine(format!("cluster '{cluster_name_str}': invalid auth ({e})")) + })?; + 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 +429,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, + ) -> Result> { + let name = cluster_name.0.clone(); + Ok(Arc::new(DuckDbAdapter::try_from_config_json( + cluster_name, + group, + json, + name.as_str(), + )?)) + } +} diff --git a/crates/queryflux-engine-adapters/src/lib.rs b/crates/queryflux-engine-adapters/src/lib.rs index e3169b8..8f03c97 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,28 @@ 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, + ) -> 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..3d3749c --- /dev/null +++ b/crates/queryflux-engine-adapters/src/snowflake/mod.rs @@ -0,0 +1,614 @@ +use std::sync::Arc; +use std::time::Duration; + +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 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, SnowflakeColumn, + SnowflakeColumnType, SnowflakeEndpointConfig, SnowflakeQueryConfig, SnowflakeRow, + SnowflakeSessionConfig, +}; +use tokio_stream::wrappers::UnboundedReceiverStream; +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) + .map_err(|e| { + QueryFluxError::Engine(format!("cluster '{cluster_name_str}': invalid auth ({e})")) + })? + .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 = build_snowflake_client(SnowflakeClientParams { + cluster_name_str, + account: account.clone(), + username, + sf_auth, + warehouse: json_str(json, "warehouse"), + database: json_str(json, "catalog"), + schema: json_str(json, "schema"), + role: json_str(json, "role"), + endpoint: json_str(json, "endpoint"), + })?; + + 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 = build_snowflake_client(SnowflakeClientParams { + cluster_name_str, + account: account.clone(), + username, + sf_auth, + warehouse: cfg.warehouse.clone(), + database: cfg.catalog.clone(), + schema: cfg.schema.clone(), + role: cfg.role.clone(), + endpoint: cfg.endpoint.clone(), + })?; + + 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()) + } +} + +struct SnowflakeClientParams<'a> { + cluster_name_str: &'a str, + account: String, + username: String, + sf_auth: SnowflakeAuthMethod, + warehouse: Option, + database: Option, + schema: Option, + role: Option, + endpoint: Option, +} + +fn build_snowflake_client(p: SnowflakeClientParams<'_>) -> Result { + let mut session = SnowflakeSessionConfig::default(); + if let Some(w) = p.warehouse { + session = session.with_warehouse(w); + } + if let Some(d) = p.database { + session = session.with_database(d); + } + if let Some(s) = p.schema { + session = session.with_schema(s); + } + if let Some(r) = p.role { + session = session.with_role(r); + } + + let query = SnowflakeQueryConfig::default() + .with_async_query_completion_timeout(Duration::from_secs(300)); + + let mut cfg = SnowflakeClientConfig::new(p.username, p.account, p.sf_auth) + .with_session(session) + .with_query(query); + + if let Some(ep) = p.endpoint { + let url = Url::parse(&ep).map_err(|e| { + QueryFluxError::Engine(format!( + "cluster '{}': invalid endpoint URL: {e}", + p.cluster_name_str + )) + })?; + cfg = cfg.with_endpoint(SnowflakeEndpointConfig::custom_base_url(url)); + } + + SnowflakeClient::new(cfg).map_err(|e| { + QueryFluxError::Engine(format!( + "cluster '{}': failed to create Snowflake client: {e}", + p.cluster_name_str + )) + }) +} + +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 executor = sf_session + .execute(sql) + .await + .map_err(|e| QueryFluxError::Engine(format!("Snowflake query failed: {e}")))?; + + let (tx, rx) = tokio::sync::mpsc::unbounded_channel::>(); + + tokio::spawn(async move { + // Metadata is on the first response even when `rowSet` / chunks are empty (e.g. LIMIT 0). + let col_types = executor.snowflake_columns(); + 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)); + if schema.fields().is_empty() { + return; + } + + let mut emitted_rows = false; + loop { + let chunk = match executor.fetch_next_chunk().await { + Ok(c) => c, + Err(e) => { + let _ = tx.send(Err(QueryFluxError::Engine(format!( + "Snowflake query failed: {e}" + )))); + return; + } + }; + let Some(rows) = chunk else { break }; + + if rows.is_empty() { + continue; + } + emitted_rows = true; + match build_snowflake_record_batch(Arc::clone(&schema), &col_types, &rows) { + Ok(batch) => { + let _ = tx.send(Ok(batch)); + } + Err(e) => { + let _ = tx.send(Err(e)); + return; + } + } + } + + if !emitted_rows { + let _ = tx.send(Ok(RecordBatch::new_empty(schema))); + } + }); + + Ok(Box::pin(UnboundedReceiverStream::new(rx))) + } + + 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, + ) -> Result> { + let name = cluster_name.0.clone(); + Ok(Arc::new(SnowflakeAdapter::try_from_config_json( + cluster_name, + group, + json, + name.as_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_snowflake_record_batch( + schema: Arc, + col_types: &[SnowflakeColumn], + rows: &[SnowflakeRow], +) -> Result { + let mut columns: Vec = Vec::with_capacity(col_types.len()); + for (col_idx, sf_col) in col_types.iter().enumerate() { + let dt = schema.field(col_idx).data_type(); + let col = build_arrow_column(dt, sf_col.column_type(), rows, col_idx)?; + columns.push(col); + } + RecordBatch::try_new(schema, columns) + .map_err(|e| QueryFluxError::Engine(format!("Snowflake RecordBatch failed: {e}"))) +} + +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..a00a9a7 100644 --- a/crates/queryflux-engine-adapters/src/starrocks/mod.rs +++ b/crates/queryflux-engine-adapters/src/starrocks/mod.rs @@ -43,14 +43,35 @@ pub struct StarRocksAdapter { } impl StarRocksAdapter { + /// When `auth` is [`Some`](Option::Some), it must be [`ClusterAuth::Basic`]; other variants error. pub fn new( cluster_name: ClusterName, group_name: ClusterGroupName, endpoint: String, auth: Option, ) -> Result { - let base_opts = Opts::from_url(&endpoint) - .map_err(|e| QueryFluxError::Engine(format!("StarRocks invalid URL: {e}")))?; + if let Some(ref a) = auth { + let unsupported = match a { + ClusterAuth::Basic { .. } => None, + ClusterAuth::Bearer { .. } => Some("bearer"), + ClusterAuth::KeyPair { .. } => Some("keyPair"), + ClusterAuth::AccessKey { .. } => Some("accessKey"), + ClusterAuth::RoleArn { .. } => Some("roleArn"), + }; + if let Some(kind) = unsupported { + return Err(QueryFluxError::Engine(format!( + "cluster '{}': StarRocks only supports ClusterAuth::basic (MySQL username/password); unsupported auth type '{kind}'", + cluster_name.0 + ))); + } + } + + let base_opts = Opts::from_url(&endpoint).map_err(|e| { + QueryFluxError::Engine(format!( + "cluster '{}': StarRocks invalid endpoint URL: {e}", + cluster_name.0 + )) + })?; // Always disable Unix socket preference — StarRocks doesn't support the // `@@socket` system variable that mysql_async queries when prefer_socket=true. @@ -71,6 +92,24 @@ 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).map_err(|e| { + QueryFluxError::Engine(format!("cluster '{cluster_name_str}': invalid auth ({e})")) + })?; + Self::new(cluster_name, group_name, endpoint, auth) + } + /// Build from persisted / YAML [`ClusterConfig`]. pub fn try_from_cluster_config( cluster_name: ClusterName, @@ -81,11 +120,7 @@ impl StarRocksAdapter { let endpoint = cfg.endpoint.clone().ok_or_else(|| { QueryFluxError::Engine(format!("cluster '{cluster_name_str}': missing endpoint")) })?; - Self::new(cluster_name, group_name, endpoint, cfg.auth.clone()).map_err(|e| { - QueryFluxError::Engine(format!( - "cluster '{cluster_name_str}': failed to create StarRocks adapter ({e})" - )) - }) + Self::new(cluster_name, group_name, endpoint, cfg.auth.clone()) } /// Execute a DDL/setup statement that returns no rows (CREATE EXTERNAL CATALOG, etc.). @@ -561,3 +596,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, + ) -> Result> { + let name = cluster_name.0.clone(); + Ok(Arc::new(StarRocksAdapter::try_from_config_json( + cluster_name, + group, + json, + name.as_str(), + )?)) + } +} diff --git a/crates/queryflux-engine-adapters/src/trino/mod.rs b/crates/queryflux-engine-adapters/src/trino/mod.rs index b0051ca..11da127 100644 --- a/crates/queryflux-engine-adapters/src/trino/mod.rs +++ b/crates/queryflux-engine-adapters/src/trino/mod.rs @@ -93,6 +93,31 @@ 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).map_err(|e| { + QueryFluxError::Engine(format!("cluster '{cluster_name_str}': invalid auth ({e})")) + })?; + 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 +892,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, + ) -> Result> { + let name = cluster_name.0.clone(); + Ok(Arc::new(TrinoAdapter::try_from_config_json( + cluster_name, + group, + json, + name.as_str(), + )?)) + } +} 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..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,6 +250,18 @@ pub fn build_frontends_status( "Arrow Flight SQL / gRPC-style access (driver-dependent).", frontends.flight_sql.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..05d949d 100644 --- a/crates/queryflux-frontend/src/dispatch.rs +++ b/crates/queryflux-frontend/src/dispatch.rs @@ -13,11 +13,13 @@ use queryflux_core::tags::{merge_tags, QueryTags}; use queryflux_core::{ error::{QueryFluxError, Result}, query::{ - ClusterGroupName, ClusterName, ExecutingQuery, FrontendProtocol, ProxyQueryId, - QueryExecution, QueryStats, QueryStatus, QueuedQuery, + ClusterGroupName, ClusterName, EngineType, ExecutingQuery, FrontendProtocol, ProxyQueryId, + 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,28 @@ 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)"); + // Trino-specific: FINISHED on first POST with no nextUri. Other async engines must not run + // this path — it parses Trino JSON and would corrupt metrics / persistence if misapplied. + if next_uri.is_none() { + if let Some(ref ib) = initial_body { + if adapter.engine_type() == EngineType::Trino { + 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 +281,140 @@ 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, + // Original SQL: when translated, `ExecutingQuery.translated_sql` holds it; otherwise `sql` is original. + 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 + .metrics + .on_query_finished(&executing.cluster_group.0, &executing.cluster_name.0); + 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, + }, + ); + } + + 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-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..47de9ab --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/format.rs @@ -0,0 +1,372 @@ +/// 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 (i, batch) in sf_batches.iter().enumerate() { + if let Err(e) = writer.write(batch) { + tracing::warn!("Failed to write Arrow batch {i} to IPC stream: {e}"); + } + } + if let Err(e) = writer.finish() { + tracing::warn!("Failed to finish Arrow IPC stream: {e}"); + } + } + 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..880c224 --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/handlers/common.rs @@ -0,0 +1,163 @@ +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. + // Drop hop-by-hop / transport headers (same idea as `proxy_response`). + if matches!( + name.as_str(), + "authorization" + | "host" + | "content-length" + | "connection" + | "keep-alive" + | "te" + | "trailer" + | "transfer-encoding" + | "upgrade" + | "proxy-authenticate" + | "proxy-authorization" + ) { + 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..64a443a --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/handlers/query.rs @@ -0,0 +1,223 @@ +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::{extract_snowflake_token, parse_snowflake_json_body, 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 body_json: Value = match parse_snowflake_json_body(&headers, &body) { + Ok(v) => v, + Err(_) => { + return sf_error( + StatusCode::BAD_REQUEST, + 390000, + "Invalid query request body", + ); + } + }; + let sql = match body_json + .get("sqlText") + .and_then(|v| v.as_str()) + .filter(|s| !s.trim().is_empty()) + { + Some(s) => s.to_string(), + None => { + return sf_error( + StatusCode::BAD_REQUEST, + 390000, + "Missing or invalid sqlText", + ); + } + }; + + // Validate TTL/idle, bump last_seen, clone fields (must not hold DashMap guard across await). + let snapshot = match state + .snowflake_sessions + .validate_snowflake_session(&qf_token) + { + Some(v) => v.snapshot, + None => { + return sf_error( + StatusCode::UNAUTHORIZED, + 390390, + "Session not found or expired", + ); + } + }; + let auth_ctx = snapshot.auth_ctx; + let group = snapshot.group; + let user = snapshot.user; + let database = snapshot.database.unwrap_or_default(); + let schema = snapshot.schema.unwrap_or_default(); + + let session_ctx = SessionContext::MySqlWire { + user, + schema: Some(if schema.is_empty() { + database.clone() + } else { + schema.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..307ebde --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/handlers/session.rs @@ -0,0 +1,187 @@ +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(); + let now = Instant::now(); + 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: now, + last_seen: 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 + .validate_snowflake_session(&token) + .is_none() + { + 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..85c4a53 --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/handlers/token.rs @@ -0,0 +1,61 @@ +use std::sync::Arc; + +use crate::state::AppState; +use axum::extract::State; +use axum::http::{HeaderMap, StatusCode}; +use axum::response::{IntoResponse, Response}; +use serde_json::json; + +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 validates the session and returns `validityInSecondsST` as the remaining +/// seconds until max session age or idle timeout (see `SnowflakeHttpSessionPolicy`), omitting +/// that field when both limits are disabled. +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", + ) + } + }; + let validated = match state.snowflake_sessions.validate_snowflake_session(&token) { + Some(v) => v, + None => { + return sf_error( + StatusCode::UNAUTHORIZED, + 390104, + "Session not found or expired", + ); + } + }; + + let data = if let Some(secs) = validated.validity_in_seconds_st { + json!({ + "sessionToken": token, + "validityInSecondsST": secs, + }) + } else { + json!({ + "sessionToken": token, + }) + }; + + ( + StatusCode::OK, + axum::Json(json!({ + "data": data, + "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..e813b8d --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/http/session_store.rs @@ -0,0 +1,258 @@ +//! In-memory Snowflake HTTP wire sessions (login token → routing/auth context). +//! +//! **Not replicated across QueryFlux processes.** With multiple replicas or during rolling +//! upgrades, clients must stick to one instance (load balancer affinity on the Snowflake token / +//! `Authorization` header) or sessions will not resolve. Configure +//! `queryflux.enforceSnowflakeHttpSessionAffinity` + `sessionAffinityAcknowledged` on the +//! Snowflake HTTP frontend once sticky routing is in place. +//! +//! Sessions honor [`SnowflakeHttpSessionPolicy`] (max age + idle timeout) on every authenticated +//! request (`validate_snowflake_session`). + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use dashmap::DashMap; +use queryflux_auth::AuthContext; +use queryflux_core::config::FrontendConfig; +use queryflux_core::query::ClusterGroupName; + +/// Wall-clock and idle limits for Snowflake HTTP wire sessions. +#[derive(Debug, Clone)] +pub struct SnowflakeHttpSessionPolicy { + /// Maximum lifetime since login. `None` = no limit. + pub max_session_age: Option, + /// Maximum time since last successful [`SnowflakeSessionStore::validate_snowflake_session`]. + /// `None` = no idle eviction. + pub idle_timeout: Option, +} + +impl Default for SnowflakeHttpSessionPolicy { + fn default() -> Self { + Self { + max_session_age: Some(Duration::from_secs(24 * 3600)), + idle_timeout: Some(Duration::from_secs(4 * 3600)), + } + } +} + +impl SnowflakeHttpSessionPolicy { + /// Build policy from `frontends.snowflakeHttp` YAML. Omitted fields use defaults (24h / 4h). + /// `0` disables that limit (matches “no max age” / “no idle timeout”). + pub fn from_frontend_config(cfg: &FrontendConfig) -> Self { + Self { + max_session_age: match cfg.snowflake_session_max_age_secs { + None => Some(Duration::from_secs(24 * 3600)), + Some(0) => None, + Some(s) => Some(Duration::from_secs(s)), + }, + idle_timeout: match cfg.snowflake_session_idle_timeout_secs { + None => Some(Duration::from_secs(4 * 3600)), + Some(0) => None, + Some(s) => Some(Duration::from_secs(s)), + }, + } + } +} + +/// Snapshot returned after a successful [`SnowflakeSessionStore::validate_snowflake_session`]. +#[derive(Debug, Clone)] +pub struct SnowflakeSessionSnapshot { + pub auth_ctx: AuthContext, + pub group: ClusterGroupName, + pub user: Option, + pub database: Option, + pub schema: Option, +} + +/// Successful validation: session fields plus optional `validityInSecondsST` for the token response. +#[derive(Debug, Clone)] +pub struct ValidatedSnowflakeSession { + pub snapshot: SnowflakeSessionSnapshot, + /// Remaining seconds until **max session age** or **idle timeout**, whichever is sooner. + /// `None` when both policy limits are disabled (unbounded session). + pub validity_in_seconds_st: Option, +} + +/// Stores active QueryFlux Snowflake wire-protocol sessions keyed by the qf_token +/// issued to the client at login. **Process-local** — no backend Snowflake account is needed. +pub struct SnowflakeSessionStore { + sessions: DashMap, + policy: SnowflakeHttpSessionPolicy, +} + +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, + /// Last successful [`SnowflakeSessionStore::validate_snowflake_session`] (or login). + pub last_seen: Instant, +} + +impl SnowflakeSessionStore { + pub fn new(policy: SnowflakeHttpSessionPolicy) -> Arc { + Arc::new(Self { + sessions: DashMap::new(), + policy, + }) + } + + 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); + } + + /// Look up the session, enforce [`SnowflakeHttpSessionPolicy`], bump `last_seen` on success, + /// and remove the entry when expired or missing. + pub fn validate_snowflake_session(&self, token: &str) -> Option { + let mut guard = self.sessions.get_mut(token)?; + let now = Instant::now(); + + if let Some(max) = self.policy.max_session_age { + if guard.created_at.elapsed() > max { + drop(guard); + self.sessions.remove(token); + return None; + } + } + if let Some(idle) = self.policy.idle_timeout { + if guard.last_seen.elapsed() > idle { + drop(guard); + self.sessions.remove(token); + return None; + } + } + + let age_remaining = self + .policy + .max_session_age + .map(|max| max.saturating_sub(guard.created_at.elapsed()).as_secs()); + let idle_remaining = self + .policy + .idle_timeout + .map(|idle| idle.saturating_sub(guard.last_seen.elapsed()).as_secs()); + let validity_in_seconds_st = match (age_remaining, idle_remaining) { + (Some(a), Some(i)) => Some(a.min(i)), + (Some(a), None) => Some(a), + (None, Some(i)) => Some(i), + (None, None) => None, + }; + + guard.last_seen = now; + Some(ValidatedSnowflakeSession { + validity_in_seconds_st, + snapshot: SnowflakeSessionSnapshot { + auth_ctx: guard.auth_ctx.clone(), + group: guard.group.clone(), + user: guard.user.clone(), + database: guard.database.clone(), + schema: guard.schema.clone(), + }, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use queryflux_core::query::ClusterGroupName; + use std::thread; + + fn dummy_session(token: &str, created: Instant, last_seen: Instant) -> SnowflakeSession { + SnowflakeSession { + qf_token: token.to_string(), + user: Some("u".into()), + auth_ctx: AuthContext { + user: "u".into(), + groups: vec![], + roles: vec![], + raw_token: None, + }, + group: ClusterGroupName("g".into()), + database: None, + schema: None, + created_at: created, + last_seen, + } + } + + #[test] + fn validate_updates_last_seen_and_allows_repeated_checks_within_idle() { + let store = SnowflakeSessionStore::new(SnowflakeHttpSessionPolicy { + max_session_age: None, + idle_timeout: Some(Duration::from_secs(3600)), + }); + let now = Instant::now(); + store.insert("t1".into(), dummy_session("t1", now, now)); + assert!(store.validate_snowflake_session("t1").is_some()); + assert!(store.validate_snowflake_session("t1").is_some()); + } + + #[test] + fn idle_timeout_evicts_session() { + let store = SnowflakeSessionStore::new(SnowflakeHttpSessionPolicy { + max_session_age: None, + idle_timeout: Some(Duration::from_millis(40)), + }); + let now = Instant::now(); + store.insert("t2".into(), dummy_session("t2", now, now)); + assert!(store.validate_snowflake_session("t2").is_some()); + thread::sleep(Duration::from_millis(80)); + assert!(store.validate_snowflake_session("t2").is_none()); + assert!(store.get("t2").is_none()); + } + + #[test] + fn max_session_age_evicts_even_if_recently_touched() { + let store = SnowflakeSessionStore::new(SnowflakeHttpSessionPolicy { + max_session_age: Some(Duration::from_millis(40)), + idle_timeout: Some(Duration::from_secs(3600)), + }); + let Some(created) = Instant::now().checked_sub(Duration::from_millis(100)) else { + return; + }; + store.insert("t3".into(), dummy_session("t3", created, Instant::now())); + assert!(store.validate_snowflake_session("t3").is_none()); + assert!(store.get("t3").is_none()); + } + + #[test] + fn validate_reports_remaining_ttl_min_of_age_and_idle() { + let store = SnowflakeSessionStore::new(SnowflakeHttpSessionPolicy { + max_session_age: Some(Duration::from_secs(10_000)), + idle_timeout: Some(Duration::from_secs(100)), + }); + let now = Instant::now(); + store.insert("t4".into(), dummy_session("t4", now, now)); + let v = store.validate_snowflake_session("t4").unwrap(); + assert!(v.validity_in_seconds_st.unwrap() <= 100); + } + + #[test] + fn validate_omits_validity_when_both_limits_disabled() { + let store = SnowflakeSessionStore::new(SnowflakeHttpSessionPolicy { + max_session_age: None, + idle_timeout: None, + }); + let now = Instant::now(); + store.insert("t5".into(), dummy_session("t5", now, now)); + let v = store.validate_snowflake_session("t5").unwrap(); + assert!(v.validity_in_seconds_st.is_none()); + } +} diff --git a/crates/queryflux-frontend/src/snowflake/mod.rs b/crates/queryflux-frontend/src/snowflake/mod.rs new file mode 100644 index 0000000..c473aef --- /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 } + } + + pub 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..931174a --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/proxy.rs @@ -0,0 +1,76 @@ +use std::collections::HashMap; +use std::time::Duration; + +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() + .timeout(Duration::from_secs(300)) + .connect_timeout(Duration::from_secs(30)) + .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..af5756c --- /dev/null +++ b/crates/queryflux-frontend/src/snowflake/sql_api/handlers.rs @@ -0,0 +1,302 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use arrow::array::{Array, StringArray}; +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; +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<()> { + let cast_columns: Vec = (0..batch.num_columns()) + .map(|col_idx| CastColumn::new(batch.column(col_idx))) + .collect(); + + for row_idx in 0..batch.num_rows() { + let row: Vec = cast_columns + .iter() + .map(|col| col.value_at(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() + } +} + +/// A column pre-cast to Utf8 so the conversion happens once per batch, not once per cell. +enum CastColumn { + Strings(Arc), + /// Values we cannot stringify for JSON without corrupting data. + Unsupported, +} + +impl CastColumn { + fn new(arr: &Arc) -> Self { + if *arr.data_type() == DataType::Utf8 { + return Self::Strings(Arc::clone(arr)); + } + match arrow_cast(arr, &DataType::Utf8) { + Ok(casted) => Self::Strings(casted), + Err(_) => Self::Unsupported, + } + } + + fn value_at(&self, row: usize) -> Value { + match self { + Self::Strings(arr) => { + if arr.is_null(row) { + return Value::Null; + } + let str_arr = arr.as_any().downcast_ref::().unwrap(); + Value::String(str_arr.value(row).to_string()) + } + Self::Unsupported => Value::Null, + } + } +} + +// --------------------------------------------------------------------------- +// SQL API v2 error helper — preserves the real HTTP status code +// --------------------------------------------------------------------------- + +fn sql_api_error(status: StatusCode, code: &str, message: &str) -> Response { + ( + status, + axum::Json(json!({ + "code": code, + "message": message, + "sqlState": "P0001", + "statementHandle": "" + })), + ) + .into_response() +} + +// --------------------------------------------------------------------------- +// 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 sql_api_error(StatusCode::BAD_REQUEST, "390000", "Invalid JSON body"), + }; + let Some(sql) = body_json["statement"] + .as_str() + .filter(|s| !s.trim().is_empty()) + else { + return sql_api_error( + StatusCode::BAD_REQUEST, + "390000", + "Missing or empty statement", + ); + }; + let sql = sql.to_string(); + + // Stateless auth: Bearer token in Authorization header. + let auth_ctx = match authenticate(&state, &headers).await { + Ok(ctx) => ctx, + Err(e) => return sql_api_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 sql_api_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..198e233 --- /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(Default::default()), + }); + + 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..048e111 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,21 @@ pub struct AppState { pub authorization: Arc, /// Resolves per-user `QueryCredentials` from `AuthContext` + cluster `queryAuth` config. pub identity_resolver: Arc, + /// Active Snowflake **HTTP wire** sessions (Snowflake connector “Form 1”), keyed by the + /// token issued at login. + /// + /// **Process-local only** — not shared across QueryFlux replicas. Multi-instance deployments + /// must use load-balancer **session affinity** (sticky routing) so login and all follow-up + /// requests hit the same instance, or sessions will fail with “not found”. Rolling restarts + /// drop in-memory sessions unless clients reconnect. See `queryflux.enforceSnowflakeHttpSessionAffinity` + /// and `frontends.snowflakeHttp.sessionAffinityAcknowledged` in config to assert affinity is configured. + /// + /// (Shared persistence for these sessions is not implemented yet.) + /// + /// Session lifetime is enforced in-process via `SnowflakeSessionStore::validate_snowflake_session` + /// (max age + idle timeout; see `frontends.snowflakeHttp.snowflakeSessionMaxAgeSecs` / + /// `snowflakeSessionIdleTimeoutSecs` in YAML). + pub snowflake_sessions: Arc, } /// Stable per-query metadata that does not change across the query's lifecycle. diff --git a/crates/queryflux-persistence/src/cluster_config.rs b/crates/queryflux-persistence/src/cluster_config.rs index d6f0597..b6476d7 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 }) => { @@ -188,6 +200,12 @@ impl UpsertClusterConfig { None => {} } + if let Some(qa) = &cfg.query_auth { + if let Ok(v) = serde_json::to_value(qa) { + config.insert("queryAuth".into(), v); + } + } + Some(Self { engine_key: engine_key.to_owned(), enabled: cfg.enabled, @@ -221,71 +239,10 @@ 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 +// built from the JSONB config blob via `try_from_config_json()` on each adapter. +// Type 1 auth uses `parse_auth_from_config_json`; Type 2 (`queryAuth`) uses +// `parse_query_auth_from_config_json` — both in `queryflux_core::engine_registry`. impl ClusterGroupConfigRecord { pub fn to_core(&self) -> ClusterGroupConfig { 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..3033f21 100644 --- a/crates/queryflux/src/main.rs +++ b/crates/queryflux/src/main.rs @@ -20,6 +20,10 @@ use queryflux_frontend::{ flight_sql::FlightSqlFrontend, mysql_wire::MysqlWireFrontend, postgres_wire::PostgresWireFrontend, + snowflake::{ + http::session_store::{SnowflakeHttpSessionPolicy, SnowflakeSessionStore}, + SnowflakeFrontend, + }, state::LiveConfig, trino_http::{state::AppState, TrinoHttpFrontend}, FrontendListenerTrait, @@ -126,6 +130,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 @@ -155,23 +163,60 @@ 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, group resolution, and + // `BackendIdentityResolver` (`queryAuth`). Adapters are still built from the + // raw JSONB 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: match queryflux_core::engine_registry::parse_auth_from_config_json( + &r.config, + ) { + Ok(a) => a, + Err(e) => { + tracing::warn!( + cluster = %r.name, + "invalid auth in cluster config JSON: {e}" + ); + None + } + }, + query_auth: + queryflux_core::engine_registry::parse_query_auth_from_config_json( + &r.config, + ), + }, + )) }) - .collect::>()?; + .collect(); + startup_cluster_records = Some(db_cluster_records); let group_records = pg .list_group_configs() @@ -278,23 +323,43 @@ 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, + ) + .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. @@ -392,6 +457,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 +468,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 { @@ -548,6 +621,29 @@ async fn main() -> Result<()> { } } + // --- Snowflake HTTP: sessions are in-memory on this process only --- + if let Some(sf) = config.queryflux.frontends.snowflake_http.as_ref() { + if sf.enabled { + if config.queryflux.enforce_snowflake_http_session_affinity + && !sf.session_affinity_acknowledged + { + anyhow::bail!( + "Snowflake HTTP is enabled and queryflux.enforceSnowflakeHttpSessionAffinity is true, \ + but frontends.snowflakeHttp.sessionAffinityAcknowledged is false. \ + Wire sessions live in process memory; configure your load balancer for session affinity \ + to the same QueryFlux replica for all requests that reuse the Snowflake login token \ + (e.g. consistent hash on the Authorization header), then set sessionAffinityAcknowledged: true. \ + For a single-replica deployment, omit enforceSnowflakeHttpSessionAffinity." + ); + } + tracing::info!( + "Snowflake HTTP frontend: login sessions are stored in this process only; \ + multi-replica setups require load balancer session affinity to the same instance per client token. \ + Set queryflux.enforceSnowflakeHttpSessionAffinity: true with sessionAffinityAcknowledged: true after configuring routing." + ); + } + } + let identity_resolver = Arc::new(BackendIdentityResolver::new()); let cluster_configs = config.clusters.clone(); @@ -580,22 +676,60 @@ async fn main() -> Result<()> { group_translation_scripts, group_default_tags, }; - let adapter_reload_cache = Arc::new(tokio::sync::Mutex::new(AdapterReloadCache { - adapters: live_config.adapters.clone(), - config_json: live_config + // Seed the reload cache. When Postgres is active, fingerprint `engine_key` + JSONB config + // (same format as `build_live_config` on reload) so an engine change rebuilds adapters even + // when the config blob shape is unchanged. For YAML-only, fold canonical `engine_key` + `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.engine_key.as_str(), &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(), + .map(|(k, v)| { + let ek = v + .engine + .as_ref() + .map(queryflux_core::engine_registry::engine_key) + .unwrap_or(""); + ( + k.clone(), + serde_json::to_string(&(ek, v)).unwrap_or_default(), + ) + }) + .collect() + }; + let adapter_reload_cache = Arc::new(tokio::sync::Mutex::new(AdapterReloadCache { + adapters: live_config.adapters.clone(), + 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 .iter() .map(|(_, s)| (s.cluster_name.0.clone(), s.clone())) .collect(), + routing_fallback: config.routing_fallback.clone(), + routers_cfg: config.routers.clone(), })); let live = Arc::new(tokio::sync::RwLock::new(live_config)); + let snowflake_session_policy = config + .queryflux + .frontends + .snowflake_http + .as_ref() + .map(SnowflakeHttpSessionPolicy::from_frontend_config) + .unwrap_or_default(); + let app_state = Arc::new(AppState { external_address: external_address.clone(), live: live.clone(), @@ -605,6 +739,7 @@ async fn main() -> Result<()> { auth_provider, authorization, identity_resolver, + snowflake_sessions: SnowflakeSessionStore::new(snowflake_session_policy), }); // --- Start admin server (Prometheus /metrics + future /admin/* endpoints) --- @@ -939,12 +1074,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(()) @@ -963,7 +1110,8 @@ type GroupStatesMap = HashMap< >; /// Holds adapter instances between DB reloads. Adapters are recreated when the -/// serialized [`ClusterConfig`] for a cluster changes so pools pick up new endpoints/credentials. +/// reload fingerprint changes (`engine_key` + config JSON), so engine switches and +/// endpoint/credential updates rebuild adapters. struct AdapterReloadCache { adapters: HashMap>, config_json: HashMap, @@ -971,6 +1119,10 @@ struct AdapterReloadCache { /// Preserved across reloads so that health status and running-query counters /// are not reset to their initial values every time the config is reloaded. cluster_states: HashMap>, + /// Last-known routing from DB (or YAML at startup). Used when `load_routing_config` returns + /// `Ok(None)` so periodic reload does not wipe routing. + routing_fallback: String, + routers_cfg: Vec, } fn health_targets_from_groups( @@ -995,15 +1147,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). +/// only when the fingerprint of `engine_key` + JSONB config matches the previous reload; +/// otherwise they are rebuilt (e.g. engine switch, 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, @@ -1016,32 +1170,49 @@ async fn build_live_config( cluster_state::ClusterState, simple::SimpleClusterGroupManager, strategy::strategy_from_config, }; + use queryflux_core::engine_registry::{ + json_str, parse_auth_from_config_json, parse_engine_key, parse_query_auth_from_config_json, + }; 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.engine_key.as_str(), &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, - cluster_name_str, + &record.engine_key, + &record.config, ) .await { @@ -1054,10 +1225,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(); @@ -1081,8 +1254,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; @@ -1092,59 +1265,48 @@ 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 - // is_healthy and running_queries are not reset across reloads. - // 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 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. - prev.clone() - } else { - Arc::new(ClusterState::new( - ClusterName(member_name.clone()), - group_key.clone(), - cluster_cid, - group_cid, - engine_type, - cluster_cfg.endpoint.clone(), - max_q, - cluster_cfg.enabled, - )) - } - } else { - let s = Arc::new(ClusterState::new( - ClusterName(member_name.clone()), - group_key.clone(), - cluster_cid, - group_cid, - engine_type, - cluster_cfg.endpoint.clone(), - max_q, - cluster_cfg.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) { - s.set_healthy(prev.is_healthy()); + // When the JSONB + engine_key fingerprint is unchanged, rebuild `ClusterState` from + // the current record anyway (group membership, IDs, endpoint, max_q may still change) + // but copy health and queue counters from the previous generation. + let cfg_json = serde_json::to_string(&(record.engine_key.as_str(), &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 = Arc::new(ClusterState::new( + ClusterName(member_name.clone()), + group_key.clone(), + cluster_cid, + group_cid, + engine_type, + endpoint, + max_q, + record.enabled, + )); + if let Some(prev) = cache.cluster_states.get(member_name.as_str()) { + let snap = prev.snapshot(); + state.set_healthy(snap.is_healthy); + if config_unchanged { + state.set_running_queries(snap.running_queries); + state.set_queued_queries(snap.queued_queries); } - s - }; + } states.push(state); } @@ -1155,13 +1317,50 @@ 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 (`queryAuth` from JSONB). + 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: match parse_auth_from_config_json(&r.config) { + Ok(a) => a, + Err(e) => { + tracing::warn!( + cluster = %r.name, + "reload: invalid auth in cluster config JSON: {e}" + ); + None + } + }, + query_auth: parse_query_auth_from_config_json(&r.config), + }, + )) + }) + .collect(); + // Build router chain. let fallback = ClusterGroupName(routing_fallback.to_string()); let mut routers: Vec> = Vec::new(); @@ -1173,6 +1372,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 +1384,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())), }, )); } @@ -1264,7 +1473,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, @@ -1274,6 +1483,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, @@ -1288,16 +1499,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() @@ -1315,7 +1516,7 @@ async fn reload_live_config( .map(|r| (r.name.clone(), r.to_core())) .collect(); - // Load routing from DB if present; otherwise fall back to empty defaults. + // Load routing from DB if present; otherwise keep last-known routing (startup YAML or previous DB load). let (routing_fallback, routers_cfg) = match pg.load_routing_config().await { Ok(Some(loaded)) => { let mut routers = Vec::new(); @@ -1327,9 +1528,11 @@ async fn reload_live_config( } } } + cache.routing_fallback = loaded.routing_fallback.clone(); + cache.routers_cfg.clone_from(&routers); (loaded.routing_fallback, routers) } - Ok(None) => (String::new(), Vec::new()), + Ok(None) => (cache.routing_fallback.clone(), cache.routers_cfg.clone()), Err(e) => { return Err(anyhow::anyhow!("reload: load_routing_config: {e}")); } @@ -1344,7 +1547,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..1265d28 100644 --- a/crates/queryflux/src/registered_engines.rs +++ b/crates/queryflux/src/registered_engines.rs @@ -8,29 +8,63 @@ 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, +) -> 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) + .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 +122,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/queryflux/Dockerfile b/docker/queryflux/Dockerfile index dabccfa..f741f8a 100644 --- a/docker/queryflux/Dockerfile +++ b/docker/queryflux/Dockerfile @@ -5,7 +5,7 @@ # # PyO3/sqlglot: builder + runtime both python:3.12.13-slim-trixie (bookworm has no python3.12 apt). # Rust: rustup 1.91.1 on top of that image. -# DuckDB: prebuilt lib (DUCKDB_DOWNLOAD_LIB). The .so is copied into the runtime stage; see LD_LIBRARY_PATH. +# DuckDB: DUCKDB_DOWNLOAD_LIB + DUCKDB_STATIC (mold/lld fail on -lduckdb dylib). .so may still be copied for loader path. # # Build (context MUST be this repo root so Cargo.toml + Cargo.lock match): # cd /path/to/queryflux && docker build -f docker/queryflux/Dockerfile -t queryflux:local . 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/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/adding-engine-support.md b/docs/adding-engine-support.md deleted file mode 100644 index 4d74e14..0000000 --- a/docs/adding-engine-support.md +++ /dev/null @@ -1,310 +0,0 @@ -# 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**. - ---- - -## 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`](../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`](../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)). - -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`](../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 [architecture.md](architecture.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 - -- [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) - -**Rust files referenced above** - -- [`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 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..fc0abab 100644 --- a/queryflux-studio/app/protocols/page.tsx +++ b/queryflux-studio/app/protocols/page.tsx @@ -1,26 +1,44 @@ 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: 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 +95,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/queryflux-studio/components/cluster-config/config-field-row.tsx b/queryflux-studio/components/cluster-config/config-field-row.tsx index f3927e2..a76f8d7 100644 --- a/queryflux-studio/components/cluster-config/config-field-row.tsx +++ b/queryflux-studio/components/cluster-config/config-field-row.tsx @@ -53,7 +53,17 @@ export function ConfigFieldRow({ {supportedAuth.map((a) => ( ))} 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" + /> +
+
+ +