feat(gateway): add Spark Thrift Server as a selectable lakehouse engine - #19411
Draft
ad1happy2go wants to merge 1 commit into
Draft
feat(gateway): add Spark Thrift Server as a selectable lakehouse engine#19411ad1happy2go wants to merge 1 commit into
ad1happy2go wants to merge 1 commit into
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #19411 +/- ##
============================================
- Coverage 75.27% 72.49% -2.78%
+ Complexity 32510 31431 -1079
============================================
Files 2574 2574
Lines 142991 142994 +3
Branches 17529 17542 +13
============================================
- Hits 107637 103670 -3967
- Misses 27289 31152 +3863
- Partials 8065 8172 +107
Flags with carried forward coverage won't be shown. Click here to find out more. 🚀 New features to boost your workflow:
|
ad1happy2go
force-pushed
the
spark-thrift-gateway
branch
3 times, most recently
from
July 30, 2026 08:20
e5e3c54 to
3eb759e
Compare
GATEWAY_ENGINE=spark serves the same guarded tool surface (query_lakehouse /
list_tables / describe_table) through a Spark Thrift Server (HiveServer2
protocol) -- e.g. Apache Kyuubi, verified below -- instead of Trino:
- Pluggable connector registry (tools/connector.py): a LakehouseConnector
bundles the four engine-specific concerns (client construction, tool
registration, agent prompt framing, and the /ready + /v1/info surfaces)
behind one abstraction with a @register_connector decorator. The app,
readiness/info endpoints, and agent prompt now resolve the engine through
the registry instead of `if engine == ...` branches, so adding an engine is
one self-contained, self-registering module -- no edits to shared wiring.
A LakehouseClient Protocol (execute/ping/close) is the client contract.
- SparkThriftClient (PyHive, pure-sasl): async facade with per-query hard
timeout + server-side cancel, cached ping, and last_error surfaced in
/ready details. PyHive ships as the optional 'spark' extra; a missing
dependency degrades to a clear /ready message.
- Connection pooling: HiveServer2 sessions are expensive to open, so the
client keeps a bounded LIFO pool (GATEWAY_SPARK_MAX_CONNECTIONS, also the
query concurrency bound) of reused sessions, recycles them past the idle
and lifetime limits, discards any session whose query timed out, and
retries once on a fresh connection when a pooled session died underneath a
query (server restart, idle kill) -- safe because traffic is SELECT-only.
close() drains the pool on shutdown.
- spark_tools: Spark SQL dialect guardrails (sqlglot read/write=spark),
SHOW TABLES / DESCRIBE TABLE for catalog introspection, backtick
identifier quoting, HiveServer2 error message cleanup, and the same
error-as-payload contract as the Trino backend.
- shared common module (QueryResult, result shaping, identifier validation)
now backs both engines; Trino modules re-export for compatibility.
- config: GATEWAY_ENGINE + GATEWAY_SPARK_{HOST,PORT,DATABASE,USER,AUTH,
PASSWORD,MAX_CONNECTIONS,POOL_MAX_IDLE_SECONDS,POOL_MAX_LIFETIME_SECONDS}
with fail-fast validation (ldap requires a password).
- /ready reports the active engine's check; /v1/info gains engine and
engine-neutral sql_url/web_ui_url; the chat UI connect panel adapts.
- engine-aware agent system prompt (dialect, database vs catalog.schema).
- helm chart: engine value + conditional GATEWAY_SPARK_* env (helm lint and
template verified for both engines); Docker image installs the spark extra.
- tests: full offline mirror of the Trino suite (fake Spark client, tools,
client timeout/cancel/ping-cache, connection-pool reuse/recycle/stale-retry,
auth-kwargs mapping, config validation, readiness/info, guardrail dialects)
plus gated live integration tests (GATEWAY_IT_SPARK_HOST) mirroring the
Trino ones.
Verified live against both a raw Spark Thrift Server and Apache Kyuubi
(HiveServer2, same wire protocol, connected with only a host/port change)
serving Hudi 1.3 tables on Spark 4.1: list/describe/query and agent
round-trips work, pooled sessions are reused across queries, and Spark-only
reads -- Hudi 1.x unstructured columns (BLOB struct, VECTOR embeddings via
size(embedding)) -- are queryable through the same gateway tools, which the
Trino connector cannot serve today.
ad1happy2go
force-pushed
the
spark-thrift-gateway
branch
2 times, most recently
from
July 30, 2026 13:27
9bec541 to
3d818b4
Compare
3 tasks
Collaborator
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Describe the issue this Pull Request addresses
Part of the agentic-lakehouse umbrella #19256. The v1 gateway (#19265) serves its guarded lakehouse tools exclusively through Trino — but many Hudi deployments already run Spark and have no Trino in the picture, and Spark reads everything Spark can read: MOR snapshot queries and Hudi 1.x tables carrying the unstructured types (
BLOB,VECTOR,VARIANT) that the in-repo Trino connector does not serve yet. This PR makes the lakehouse engine pluggable:GATEWAY_ENGINE=sparkserves the identical tool surface through a Spark Thrift Server (HiveServer2 protocol), so any Spark cluster with the Hudi bundle becomes an agent-queryable lakehouse with one env var.Summary and Changelog
One selectable backend behind the existing tool contract; the chat API, UI, MCP server, guardrails, and agent behavior are identical across engines — only the SQL dialect and the endpoint reported by
/v1/infochange.tools/spark_client.py— SparkThriftClient (PyHive)TrinoClient: each query runs the sync client in a dedicated bounded worker pool under a hard timeout, with best-effort server-sidecursor.cancel()on expiry;ping()is cached so readiness probes do not open a fresh Thrift session every period.pip install 'hudi-agent-gateway[spark]', pure-sasl — no C toolchain); a missing dependency degrades to a clear/readydetail instead of a crash, and the client recordslast_errorso readiness failures carry the underlying cause (connection refused vs missing dependency vs timeout).none(SASL PLAIN, the Spark Thrift default),nosasl, andldap(config validation fails fast if the password is missing); qualified result column names (tbl.col) are normalized to bare names so results look identical across engines.tools/spark_tools.py— the same three tools, Spark shapequery_lakehousevalidates and re-renders model SQL through the shared AST guardrails in thesparksqlglot dialect (single statement, SELECT-only, row cap injected as a realLIMIT); backtick identifiers parse and render correctly.information_schema:list_tablesusesSHOW TABLES IN,describe_tableusesDESCRIBE TABLE, with the same strict identifier validation (names arrive over HTTP/MCP and never reach SQL unvalidated) and the same error-as-payload contract ({"error", "hint"}) so the agent self-corrects; HiveServer2's Java-stack-trace error renderings are reduced to the first meaningful line.Shared plumbing
tools/common.pyhosts the engine-agnostic pieces (QueryResult, result shaping with row/byte truncation notices, identifier validation) withLakehouseQueryError/LakehouseTimeoutErrorbase classes both engines' errors subclass; the Trino modules re-export their existing names, so nothing downstream changes.enforce_guardrailsgains adialectparameter (defaulttrino), hints name the dialect the model should write./readyreports the active engine's check (with cause detail for spark),/v1/infogainsengineand engine-neutralsql_url/web_ui_url, the chat UI connect panel shows "Spark SQL (JDBC)" with thejdbc:hive2://URL and hides the web-UI row when the engine has none, and the agent system prompt speaks the active dialect (database.tablevscatalog.schema.table).sparkextra so one image serves both engines; the Helm chart gains anenginevalue with conditionalGATEWAY_SPARK_*env (helm lint and template verified for both engines).Tests — full mirror of the Trino suite plus client-level coverage
SHOW TABLES/DESCRIBESQL shapes, identifier-injection rejection, not-found mapping from both empty results and server errors, tool-surface parity across engines, client timeout→cancel mapping, ping caching,last_errorlifecycle, concurrency, config validation, engine-aware readiness/info.GATEWAY_IT_SPARK_HOST, mirroring the Trino live tests): list/describe/query against a real Spark Thrift Server serving the sametripsdataset the local-dev example writes, plus a live guardrail-rejection check.No code copied from other projects.
Impact
New optional engine + configs:
GATEWAY_ENGINEandGATEWAY_SPARK_{HOST,PORT,DATABASE,USER,AUTH,PASSWORD}(documented in the module README's config table, with a quickstart section). New pip extraspark. One intentional API rename while the module is days old:/v1/info'strino_url/trino_ui_urlbecome engine-neutralsql_url/web_ui_url(plus a newenginefield); the shipped UI is updated in the same PR. No existing Hudi APIs, configs, storage format, or bundles change.Risk Level
low — additive code paths; the Trino backend is untouched apart from importing shared helpers it re-exports. Verified by the full offline suite (121 passing; ruff and mypy clean), the new gated live integration tests passing 4/4 against a real Spark Thrift Server serving Hudi 1.x tables, and the existing Trino live integration tests re-run green against the local-dev stack (no regression). Additionally validated end to end: agent round-trips through
GATEWAY_ENGINE=sparkwith both a local Ollama model and a frontier model answered aggregation and filtered-lookup questions correctly, and the same gateway tools queried a Hudi 1.x unstructured table (BLOB struct columns,size(embedding)over aVECTOR(768)column) — the concrete capability the Trino connector cannot serve today.Documentation Update
Module README updated in this PR (engine row in the config table, "Using a Spark Thrift Server instead of Trino" quickstart, live integration-test instructions). Website docs for the gateway are tracked with the umbrella #19256.
Contributor's checklist