This is an automated email from the ASF dual-hosted git repository.
nsivabalan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hudi.git
The following commit(s) were added to refs/heads/master by this push:
new 71069b2f5259 feat(gateway): add Spark Connect SQL backend (#20329)
71069b2f5259 is described below
commit 71069b2f52591cf275948b23e0cb0793779f5542
Author: Shuo Cheng <[email protected]>
AuthorDate: Mon Nov 2 10:16:41 2026 +0800
feat(gateway): add Spark Connect SQL backend (#20329)
hudi-agent-gateway can now run its SQL tools against the Hudi Spark
Connect endpoint instead of Trino. Set GATEWAY_SQL_BACKEND=spark-connect
or Helm sqlBackend (values-spark-connect.yaml); Trino stays the default.
list_tables, describe_table and query_lakehouse keep their interfaces
for both Agent and MCP clients.
- Uses the PySpark Connect client, installed through the optional
[spark-connect] extra. Each request gets its own session, with up to 8
concurrent queries, bounded results, timeouts, and interruptAll on
timeout or cancel.
- Spark-dialect guardrails keep the gateway read-only. They reject
script transforms, reflection, file-path and data-source tables,
IDENTIFIER(), qualified or unknown table functions, and Hudi table
functions given a path. Two-part names resolve through the catalog.
- Stopping or disconnecting a streaming chat turn cancels that turn's
tool tasks, which interrupts the running Spark query.
- The readiness check uses its own probe and does not compete with
queries. /v1/info reports the active backend.
- Chat requests have a deadline. On timeout, clients are told to start a
new session, because the old one may hold an incomplete tool call.
Part of #20232.
---
.github/workflows/agent_gateway_ci.yml | 18 +-
hudi-agent-gateway/Dockerfile | 2 +-
hudi-agent-gateway/README.md | 98 +++-
hudi-agent-gateway/pyproject.toml | 8 +-
hudi-agent-gateway/src/hudi_agent_gateway/agent.py | 10 +-
.../src/hudi_agent_gateway/api/chat.py | 85 +++-
.../src/hudi_agent_gateway/api/meta.py | 31 +-
.../src/hudi_agent_gateway/api/models.py | 12 +-
hudi-agent-gateway/src/hudi_agent_gateway/app.py | 10 +-
.../{tools/__init__.py => cancellation.py} | 36 +-
.../src/hudi_agent_gateway/config.py | 67 ++-
.../src/hudi_agent_gateway/mcp_server.py | 4 +-
.../src/hudi_agent_gateway/tools/__init__.py | 18 +-
.../src/hudi_agent_gateway/tools/guardrails.py | 138 +++++-
.../{mcp_server.py => tools/lakehouse_client.py} | 42 +-
.../src/hudi_agent_gateway/tools/registry.py | 11 +
.../tools/spark_connect_client.py | 213 +++++++++
.../src/hudi_agent_gateway/tools/trino_client.py | 34 +-
.../src/hudi_agent_gateway/tools/trino_tools.py | 91 ++--
.../src/hudi_agent_gateway/ui/app.js | 79 +++-
.../src/hudi_agent_gateway/ui/index.html | 9 +-
.../src/hudi_agent_gateway/ui/style.css | 4 +
hudi-agent-gateway/tests/conftest.py | 10 +-
.../tests/integration/test_spark_connect_live.py | 84 ++++
hudi-agent-gateway/tests/test_agent_loop.py | 14 +
hudi-agent-gateway/tests/test_api_meta.py | 4 +-
hudi-agent-gateway/tests/test_chat_sse.py | 143 +++++-
hudi-agent-gateway/tests/test_guardrails.py | 214 +++++++++
hudi-agent-gateway/tests/test_mcp.py | 2 +-
hudi-agent-gateway/tests/test_spark_connect.py | 498 +++++++++++++++++++++
hudi-lakehouse/README.md | 9 +-
.../charts/hudi-agent-gateway/templates/NOTES.txt | 7 +-
.../hudi-agent-gateway/templates/deployment.yaml | 24 +-
.../hudi-agent-gateway/values-local-dev.yaml | 1 -
.../hudi-agent-gateway/values-spark-connect.yaml | 27 +-
.../charts/hudi-agent-gateway/values.yaml | 13 +
hudi-lakehouse/charts/hudi-spark-connect/README.md | 28 +-
.../hudi-spark-connect/templates/configmap.yaml | 1 +
hudi-lakehouse/tests/test_gateway_chart.py | 63 +++
hudi-lakehouse/tests/test_spark_connect_chart.py | 1 +
40 files changed, 1952 insertions(+), 211 deletions(-)
diff --git a/.github/workflows/agent_gateway_ci.yml
b/.github/workflows/agent_gateway_ci.yml
index 768c485b4922..6e16575e3514 100644
--- a/.github/workflows/agent_gateway_ci.yml
+++ b/.github/workflows/agent_gateway_ci.yml
@@ -61,7 +61,7 @@ jobs:
cache: pip
cache-dependency-path: hudi-agent-gateway/pyproject.toml
- name: Install
- run: pip install -e ".[dev]"
+ run: pip install -e ".[dev,spark-connect]"
- name: Fetch pinned Hudi Architect baseline
run: >-
git fetch --no-tags --depth=1 origin
@@ -98,6 +98,18 @@ jobs:
../hudi-ai-operator/skills/hudi-architect/validate_flink_capabilities.py
../hudi-ai-operator/skills/hudi-architect/validate_flink_design.py
../hudi-ai-operator/skills/hudi-architect/validate_flink_existing_table.py
+ - name: Test guardrails with minimum sqlglot
+ # Read the lower bound from project metadata so this test follows
future bumps.
+ run: |
+ sqlglot_min_requirement=$(python - <<'PY'
+ import tomllib
+ with open("pyproject.toml", "rb") as f:
+ dependencies = tomllib.load(f)["project"]["dependencies"]
+ print(next(dep for dep in dependencies if
dep.startswith("sqlglot>=")).replace(">=", "==", 1))
+ PY
+ )
+ python -m pip install "$sqlglot_min_requirement"
+ python -m pytest tests/test_guardrails.py -q
flink-sql-fixtures:
name: Hudi 1.2.0 / Flink 1.20.1 SQL fixtures
@@ -145,3 +157,7 @@ jobs:
helm template "$chart" > /dev/null
fi
done
+ helm lint hudi-lakehouse/charts/hudi-agent-gateway \
+ -f
hudi-lakehouse/charts/hudi-agent-gateway/values-spark-connect.yaml
+ helm template hudi-lakehouse/charts/hudi-agent-gateway \
+ -f
hudi-lakehouse/charts/hudi-agent-gateway/values-spark-connect.yaml > /dev/null
diff --git a/hudi-agent-gateway/Dockerfile b/hudi-agent-gateway/Dockerfile
index 27fae43b0ce6..91a4b92243f2 100644
--- a/hudi-agent-gateway/Dockerfile
+++ b/hudi-agent-gateway/Dockerfile
@@ -27,7 +27,7 @@ WORKDIR /opt/hudi-agent-gateway
COPY pyproject.toml README.md ./
RUN mkdir -p src/hudi_agent_gateway \
&& printf '__version__ = "0.1.0"\n' > src/hudi_agent_gateway/__init__.py \
- && pip install --no-cache-dir . \
+ && pip install --no-cache-dir '.[spark-connect]' \
&& pip uninstall -y hudi-agent-gateway
COPY src ./src
RUN pip install --no-cache-dir --no-deps .
diff --git a/hudi-agent-gateway/README.md b/hudi-agent-gateway/README.md
index d1ca1c2435c6..eecdb2ed95c1 100644
--- a/hudi-agent-gateway/README.md
+++ b/hudi-agent-gateway/README.md
@@ -26,17 +26,26 @@ process hosts three surfaces over the same set of guarded
lakehouse tools:
| MCP server | `/mcp` (streamable HTTP) | external agents (Claude, anything
MCP) call the lakehouse tools directly |
| Chat UI | `/ui/` | first-party ChatGPT-style web UI (zero third-party code) |
-The v1 tools query the lakehouse through Trino: `query_lakehouse` (guarded,
+The tools query the lakehouse through Trino (default) or the official PySpark
+Connect client: `query_lakehouse` (guarded,
read-only SQL), `list_tables`, `describe_table`. Every model-written query
passes AST-level guardrails (single statement, SELECT-only, row cap injected
as a real `LIMIT`) and every invocation is logged as structured JSON — the
seed of the gateway's trace collection.
+For Spark, dynamic `IDENTIFIER(...)` calls and custom or qualified table
functions
+are rejected. Allowed table functions are `range`, `explode`, `explode_outer`,
+`posexplode`, `posexplode_outer`, `inline`, `inline_outer`, `stack`,
`json_tuple`,
+and `hudi_query`, `hudi_table_changes`, `hudi_query_timeline`,
`hudi_filesystem_view`,
+`hudi_metadata`. Hudi functions require a literal catalog table name such as
+`'default.orders'`; paths and computed table names are rejected.
+`hudi_filesystem_view` accepts only the table argument, without a subpath.
+
## Quickstart (local process)
```bash
cd hudi-agent-gateway
-python3.12 -m venv .venv && .venv/bin/pip install -e ".[dev]"
+python3.12 -m venv .venv && .venv/bin/pip install -e ".[dev,spark-connect]"
# point at a Trino with the Hudi connector (e.g. the local-dev stack,
port-forwarded):
GATEWAY_TRINO_HOST=localhost GATEWAY_TRINO_PORT=18080 \
@@ -61,6 +70,14 @@ ollama pull qwen3:8b
GATEWAY_LLM_PROVIDER=ollama GATEWAY_LLM_MODEL=qwen3:8b hudi-agent-gateway serve
```
+The chat UI shows progress and elapsed time while a response is running.
**Stop**
+closes the response stream and cancels the agent task. If the gateway does not
+accept the connection within 15 seconds, the UI reports a connection timeout.
+The server's chat timeout returns HTTP 504
+for JSON requests or a `chat_timeout` SSE error. Provider/SQL cancellation is
+best-effort; remote work may take time to stop. After Stop or an error the UI
+starts a new conversation on the next message to avoid incomplete tool-call
history.
+
Connect an MCP client:
```bash
@@ -85,7 +102,11 @@ Environment variables (prefix `GATEWAY_` except the
standard key names):
| `ANTHROPIC_API_KEY` / `OPENAI_API_KEY` | — | required by the matching
provider |
| `GATEWAY_OLLAMA_BASE_URL` | `http://localhost:11434` | Ollama endpoint |
| `GATEWAY_OPENAI_BASE_URL` | — | endpoint for `openai-compatible` (vLLM,
Together, …) |
-| `GATEWAY_LLM_TIMEOUT_SECONDS` | `120` | per-model-call timeout |
+| `GATEWAY_LLM_TIMEOUT_SECONDS` | `120` | HTTP operation/read timeout;
streaming data can keep a call alive |
+| `GATEWAY_CHAT_TIMEOUT_SECONDS` | `300` | total wall-clock limit for one chat
turn, including all model/tool calls |
+| `GATEWAY_SQL_BACKEND` | `trino` | `trino` or `spark-connect` |
+| `GATEWAY_SPARK_CONNECT_URI` |
`sc://hudi-spark-connect.hudi-lakehouse.svc:15002` | PySpark Connect URI;
optional `/;key=value` parameters |
+| `GATEWAY_SPARK_CONNECT_CATALOG` / `_SCHEMA` | `spark_catalog` / `default` |
Spark query and metadata defaults |
| `GATEWAY_TRINO_HOST` / `_PORT` | `hudi-trino.hudi-lakehouse.svc` / `8080` |
Trino coordinator |
| `GATEWAY_TRINO_CATALOG` / `_SCHEMA` / `_USER` | `hudi` / `default` /
`hudi-agent-gateway` | query defaults |
| `GATEWAY_SQL_ROW_CAP` | `200` | LIMIT enforced on every query |
@@ -98,10 +119,79 @@ Environment variables (prefix `GATEWAY_` except the
standard key names):
| `GATEWAY_MCP_ENABLED` | `true` | serve /mcp |
| `GATEWAY_HOST` / `GATEWAY_PORT` / `GATEWAY_LOG_LEVEL` | `0.0.0.0` / `8000` /
`INFO` | server basics |
-Startup never depends on the LLM or Trino being reachable: `/health` is
+Startup never depends on the LLM or the selected SQL backend being reachable:
`/health` is
liveness, `/ready` reports per-dependency status (and gates the Kubernetes
readiness probe).
+## Spark Connect backend
+
+Install the gateway in a Python 3.11/3.12 environment with the `spark-connect`
+extra. It includes `pyspark-client==4.1.3`, matching the chart's Spark version,
+and PyArrow 18.1.0. The gateway image includes this extra.
+The official Python Connect client sends RPCs to the remote Spark server; the
+gateway image requires no Java runtime, Spark JARs, or native ADBC driver.
+PySpark also brings pandas, gRPC, and protobuf dependencies; use a dedicated
+environment and upgrade the client alongside the server.
+
+After deploying the [Spark Connect
chart](../hudi-lakehouse/charts/hudi-spark-connect/README.md)
+and forwarding port 15002:
+
+```bash
+GATEWAY_SQL_BACKEND=spark-connect \
+GATEWAY_SPARK_CONNECT_URI='sc://localhost:15002' \
+ hudi-agent-gateway serve
+```
+
+Use Helm `sqlBackend: spark-connect` and `sparkConnect.uri/catalog/schema`,
+or apply the chart's `values-spark-connect.yaml` overlay. An optional
+`sparkConnect.existingSecret` supplies the URI via `sparkConnect.uriSecretKey`.
+`/ready` reports a `spark-connect` check (cached SELECT 1) and the LLM check;
+`/v1/info`, Agent prompts, MCP tool descriptions, and the UI reflect the
selected
+backend. Connection credentials are omitted from `/v1/info`.
+
+`query_lakehouse` accepts Spark SQL with backtick identifiers and enforces the
+same single-SELECT/AST/LIMIT policy. `list_tables` uses SHOW TABLES and reports
+`TABLE` for persistent objects (Spark's SHOW TABLES does not distinguish views
+from tables), `TEMPORARY` for temporary objects. `describe_table` uses a
zero-row
+SELECT and PySpark's schema, preserving column/type/nullability result columns.
+Types use Spark SQL names such as `int`, `bigint`, and `array<string>`. Each
call
+creates an independent remote session, sets its catalog/schema, and stops that
+session in a finally block. Sessions are not shared across tool calls.
+
+**Bounds:** SQL LIMIT and DataFrame.limit before collect cap query result rows;
+the existing JSON size cap truncates tool responses. Wide rows can still
consume
+substantial client memory, and SHOW TABLES can enumerate the whole namespace
+before returning limited rows. When capped, SHOW TABLES can return an arbitrary
+subset, since Spark does not guarantee its order. The JSON cap is a response
limit, not a
+server/memory budget. Queries use at most eight workers with no accumulating
+queue. A timed-out call keeps its capacity slot until execution and cleanup
+actually finish.
+
+**Timeout/cancellation:** gateway timeouts return promptly and request
+`SparkSession.interruptAll()` on a separate thread; task cancellation follows
+the same path. Stopping a streaming chat also cancels the tool tasks belonging
+to that turn, including those running inside the agent. Each request owns its
+session, so this interrupts only that request's operations. Spark provides a
+server interruption RPC, but network
+failures and races with operation startup mean cancellation remains
best-effort.
+A timeout does not prove Spark stopped. Hung RPCs can delay session cleanup;
+worker pools shut down without waiting, and Python cannot forcibly terminate
+worker threads. Restart the process if all eight slots remain stuck. Connect
+URIs can use `/;use_ssl=true;token=...` for an appropriately configured
endpoint;
+the chart itself has no authentication or TLS. Gateway tools expose no DML;
+direct Superset/Zeppelin writes have their own
+[tutorials and
limitations](../hudi-lakehouse/charts/hudi-spark-connect/README.md#spark-connect-client-tutorials).
+
+Opt-in live tests use an existing table and do not write data:
+
+```bash
+GATEWAY_IT_SPARK_CONNECT_URI='sc://localhost:15002' \
+GATEWAY_IT_SPARK_TABLE=bi_adbc_demo.orders \
+ .venv/bin/python -m pytest tests/integration/test_spark_connect_live.py -q
+```
+
+Add `GATEWAY_IT_LLM=1` to exercise the real Agent loop with the configured LLM.
+
## Development
```bash
diff --git a/hudi-agent-gateway/pyproject.toml
b/hudi-agent-gateway/pyproject.toml
index 8b548f2e5271..12d71c3f47bb 100644
--- a/hudi-agent-gateway/pyproject.toml
+++ b/hudi-agent-gateway/pyproject.toml
@@ -35,7 +35,8 @@ dependencies = [
"langchain-ollama>=1.0",
"fastmcp>=2.10,<3",
"trino>=0.330",
- "sqlglot>=25.0",
+ # Guardrails rely on AST shapes verified at this minimum version in CI.
+ "sqlglot>=30.21.0",
"pydantic>=2.7",
"pydantic-settings>=2.3",
"sse-starlette>=2.1",
@@ -43,6 +44,11 @@ dependencies = [
]
[project.optional-dependencies]
+spark-connect = [
+ # The session cleanup workaround relies on PySpark's private client API.
+ "pyspark-client==4.1.3",
+ "pyarrow==18.1.0",
+]
dev = [
"pytest>=8.0",
"pytest-asyncio>=0.24",
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/agent.py
b/hudi-agent-gateway/src/hudi_agent_gateway/agent.py
index eb2add9873ff..2c4217fc3cf0 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/agent.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/agent.py
@@ -31,13 +31,13 @@ from hudi_agent_gateway.tools.registry import ToolRegistry
_SYSTEM_PROMPT = """\
You are the Apache Hudi lakehouse analyst. You answer questions about data in \
-an Apache Hudi lakehouse by querying it through Trino using your tools. The \
+an Apache Hudi lakehouse by querying it through {backend} using your tools.
The \
default catalog is `{catalog}` and the default schema is `{schema}`; tables \
are Hudi tables registered in the catalog's metastore.
Tool strategy:
- For unfamiliar tables, call `list_tables` and `describe_table` BEFORE
writing SQL.
-- `query_lakehouse` accepts exactly ONE read-only SELECT statement (Trino
SQL). \
+- `query_lakehouse` accepts exactly ONE read-only SELECT statement ({dialect}
SQL). \
Qualify names as catalog.schema.table when querying outside the defaults.
- A server-side row cap of {row_cap} rows is enforced. Prefer aggregations, \
GROUP BY, and explicit LIMIT over raw row dumps.
@@ -55,8 +55,10 @@ Grounding rules:
def build_system_prompt(settings: GatewaySettings) -> str:
extra = f"\n{settings.system_prompt_extra}" if
settings.system_prompt_extra else ""
return _SYSTEM_PROMPT.format(
- catalog=settings.trino_catalog,
- schema=settings.trino_schema,
+ backend="Spark Connect" if settings.sql_backend == "spark-connect"
else "Trino",
+ dialect="Spark" if settings.sql_backend == "spark-connect" else
"Trino",
+ catalog=settings.lakehouse_catalog,
+ schema=settings.lakehouse_schema,
row_cap=settings.sql_row_cap,
extra=extra,
)
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/api/chat.py
b/hudi-agent-gateway/src/hudi_agent_gateway/api/chat.py
index 08263e256155..c6a7efcc8e5e 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/api/chat.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/api/chat.py
@@ -18,9 +18,11 @@
from __future__ import annotations
+import asyncio
import logging
import uuid
from collections.abc import AsyncIterator
+from contextlib import aclosing
from typing import Any
from fastapi import APIRouter, HTTPException, Request
@@ -33,11 +35,13 @@ from hudi_agent_gateway.api.models import (
ChatResponse,
DoneEvent,
ErrorEvent,
+ StatusEvent,
TokenEvent,
ToolCallEvent,
ToolResultEvent,
ToolTraceEntry,
)
+from hudi_agent_gateway.cancellation import ToolCancellationScope,
tool_cancellation_scope
from hudi_agent_gateway.log import log_event, request_id_var, session_id_var
from hudi_agent_gateway.tools.registry import is_tool_error
@@ -46,10 +50,10 @@ logger = logging.getLogger("hudi_agent_gateway.chat")
def _message_text(msg: Any) -> str:
text = getattr(msg, "text", None)
- if callable(text): # langchain-core < 1.x compatibility
- return text()
if isinstance(text, str):
return text
+ if callable(text): # langchain-core < 1.x compatibility
+ return text()
return str(getattr(msg, "content", ""))
@@ -97,10 +101,13 @@ async def chat(request: Request, body: ChatRequest) -> Any:
if wants_stream:
return EventSourceResponse(_stream_chat(request, body))
- agent = state.agents.get(body.model)
inputs = {"messages": [HumanMessage(body.message)]}
try:
- result = await agent.ainvoke(inputs, config=_agent_config(request,
body.session_id))
+ async with asyncio.timeout(state.settings.chat_timeout_seconds):
+ agent = state.agents.get(body.model)
+ result = await agent.ainvoke(inputs, config=_agent_config(request,
body.session_id))
+ except TimeoutError as e:
+ raise HTTPException(status_code=504, detail=_timeout_message(request))
from e
except GraphRecursionError as e:
raise HTTPException(
status_code=422,
@@ -126,27 +133,44 @@ async def chat(request: Request, body: ChatRequest) ->
Any:
async def _stream_chat(request: Request, body: ChatRequest) ->
AsyncIterator[ServerSentEvent]:
state = request.app.state
- agent = state.agents.get(body.model)
inputs = {"messages": [HumanMessage(body.message)]}
final_text_parts: list[str] = []
+ yield _status_event("Preparing response…")
+ scope = ToolCancellationScope()
+ scope_token = tool_cancellation_scope.set(scope)
try:
- async for mode, chunk in agent.astream(
- inputs,
- config=_agent_config(request, body.session_id),
- stream_mode=["messages", "updates"],
- ):
- if mode == "messages":
- msg_chunk, _meta = chunk
- if isinstance(msg_chunk, AIMessageChunk):
- text = _message_text(msg_chunk)
- if text:
- final_text_parts.append(text)
- yield ServerSentEvent(
- event="token",
data=TokenEvent(text=text).model_dump_json()
- )
- elif mode == "updates" and isinstance(chunk, dict):
- for events in _updates_to_events(chunk):
- yield events
+ async with asyncio.timeout(state.settings.chat_timeout_seconds):
+ agent = state.agents.get(body.model)
+ async with aclosing(agent.astream(
+ inputs,
+ config=_agent_config(request, body.session_id),
+ stream_mode=["messages", "updates"],
+ )) as stream:
+ async for mode, chunk in stream:
+ if mode == "messages":
+ msg_chunk, _meta = chunk
+ if isinstance(msg_chunk, AIMessageChunk):
+ text = _message_text(msg_chunk)
+ if text:
+ final_text_parts.append(text)
+ yield ServerSentEvent(
+ event="token",
data=TokenEvent(text=text).model_dump_json()
+ )
+ elif mode == "updates" and isinstance(chunk, dict):
+ for event in _updates_to_events(chunk):
+ yield event
+ except TimeoutError:
+ log_event(logger, "chat_timeout",
timeout_seconds=state.settings.chat_timeout_seconds)
+ yield ServerSentEvent(
+ event="error",
+ data=ErrorEvent(
+ code="chat_timeout", message=_timeout_message(request)
+ ).model_dump_json(),
+ )
+ return
+ except asyncio.CancelledError:
+ log_event(logger, "chat_cancelled")
+ raise
except GraphRecursionError:
yield ServerSentEvent(
event="error",
@@ -163,6 +187,11 @@ async def _stream_chat(request: Request, body:
ChatRequest) -> AsyncIterator[Ser
data=ErrorEvent(code="agent_error",
message=str(e)).model_dump_json(),
)
return
+ finally:
+ # LangGraph may run tools in tasks that outlive its streaming iterator.
+ # Cancel only this turn's tasks when the stream ends or disconnects.
+ scope.cancel()
+ tool_cancellation_scope.reset(scope_token)
yield ServerSentEvent(
event="done",
@@ -172,8 +201,20 @@ async def _stream_chat(request: Request, body:
ChatRequest) -> AsyncIterator[Ser
)
+def _timeout_message(request: Request) -> str:
+ seconds = request.app.state.settings.chat_timeout_seconds
+ return (f"The response exceeded the {seconds:g}s time limit. Start a new
session "
+ "before retrying, because this session may contain an incomplete
tool call.")
+
+
+def _status_event(message: str) -> ServerSentEvent:
+ return ServerSentEvent(event="status",
data=StatusEvent(message=message).model_dump_json())
+
+
def _updates_to_events(update: dict[str, Any]) -> list[ServerSentEvent]:
events: list[ServerSentEvent] = []
+ if "pre_model_hook" in update:
+ events.append(_status_event("Waiting for model…"))
agent_update = update.get("agent") or {}
for msg in agent_update.get("messages", []):
if isinstance(msg, AIMessage):
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/api/meta.py
b/hudi-agent-gateway/src/hudi_agent_gateway/api/meta.py
index 0273b7798859..8d881aab87d6 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/api/meta.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/api/meta.py
@@ -19,6 +19,7 @@
from __future__ import annotations
import time
+from urllib.parse import urlparse
from fastapi import APIRouter, Request, Response
@@ -46,11 +47,12 @@ async def ready(request: Request, response: Response) ->
ReadyResponse:
state = request.app.state
checks: dict[str, DependencyStatus] = {}
- trino_ok = await state.trino_client.ping()
- checks["trino"] = DependencyStatus(
- ok=trino_ok,
- detail="reachable" if trino_ok else
- f"cannot reach trino at
{state.settings.trino_host}:{state.settings.trino_port}",
+ lakehouse_ok = await state.lakehouse_client.ping()
+ backend = state.settings.sql_backend
+ checks[backend] = DependencyStatus(
+ ok=lakehouse_ok,
+ detail="reachable" if lakehouse_ok else
+ f"cannot query {backend}; check connection settings, client driver,
and server logs",
)
llm_ok, llm_detail = await state.llm_readiness.check()
@@ -86,14 +88,25 @@ async def models(request: Request) -> ModelsResponse:
@router.get("/v1/info", response_model=InfoResponse,
response_model_by_alias=True)
async def info(request: Request) -> InfoResponse:
s = request.app.state.settings
+ lakehouse_url = f"http://{s.trino_host}:{s.trino_port}";
+ if s.sql_backend == "spark-connect":
+ uri = urlparse("http" + s.spark_connect_uri[2:])
+ host = uri.hostname or ""
+ if ":" in host:
+ host = f"[{host}]"
+ lakehouse_url = f"sc://{host}:{uri.port or 15002}"
return InfoResponse(
name="hudi-agent-gateway",
version=__version__,
provider=s.llm_provider,
model=s.llm_model,
- catalog=s.trino_catalog,
- schema=s.trino_schema,
- trino_url=f"http://{s.trino_host}:{s.trino_port}";,
- trino_ui_url=f"http://{s.trino_host}:{s.trino_port}/ui/";,
+ catalog=s.lakehouse_catalog,
+ schema=s.lakehouse_schema,
+ sql_backend=s.sql_backend,
+ lakehouse_url=lakehouse_url,
+ trino_url=f"http://{s.trino_host}:{s.trino_port}"; if s.sql_backend ==
"trino" else "",
+ trino_ui_url=(
+ f"http://{s.trino_host}:{s.trino_port}/ui/"; if s.sql_backend ==
"trino" else ""
+ ),
mcp_enabled=s.mcp_enabled,
)
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/api/models.py
b/hudi-agent-gateway/src/hudi_agent_gateway/api/models.py
index 0c574e38829c..b78a65cbfcd4 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/api/models.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/api/models.py
@@ -54,6 +54,10 @@ class ChatResponse(BaseModel):
# --- SSE event payloads (each is the JSON body of one `data:` line) ---
+class StatusEvent(BaseModel):
+ message: str
+
+
class TokenEvent(BaseModel):
text: str
@@ -77,7 +81,7 @@ class DoneEvent(BaseModel):
class ErrorEvent(BaseModel):
- code: Literal["llm_error", "agent_error", "recursion_limit"]
+ code: Literal["llm_error", "agent_error", "recursion_limit",
"chat_timeout"]
message: str
@@ -105,8 +109,10 @@ class InfoResponse(BaseModel):
catalog: str
schema_: str = Field(alias="schema")
# In-cluster/network endpoints, for the UI's connect panel.
- trino_url: str
- trino_ui_url: str
+ sql_backend: str = "trino"
+ lakehouse_url: str = ""
+ trino_url: str = ""
+ trino_ui_url: str = ""
mcp_enabled: bool
model_config = {"populate_by_name": True}
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/app.py
b/hudi-agent-gateway/src/hudi_agent_gateway/app.py
index de1faf8e8a6e..32e919e8cd11 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/app.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/app.py
@@ -36,8 +36,7 @@ from hudi_agent_gateway.llm import LLMReadiness
from hudi_agent_gateway.log import log_event, setup_logging
from hudi_agent_gateway.mcp_server import build_mcp
from hudi_agent_gateway.sessions import SessionStore
-from hudi_agent_gateway.tools import build_registry
-from hudi_agent_gateway.tools.trino_client import TrinoClient
+from hudi_agent_gateway.tools import build_client, build_registry
logger = logging.getLogger("hudi_agent_gateway.app")
@@ -46,8 +45,8 @@ def create_app(settings: GatewaySettings | None = None) ->
FastAPI:
s = settings or GatewaySettings()
setup_logging(s.log_level)
- trino_client = TrinoClient(s)
- registry = build_registry(s, trino_client)
+ lakehouse_client = build_client(s)
+ registry = build_registry(s, lakehouse_client)
# NOTE: the MCP ASGI sub-app has its own lifespan that MUST run inside the
# outer app's lifespan, or FastMCP's streamable-HTTP session manager never
# starts and /mcp requests hang.
@@ -56,7 +55,7 @@ def create_app(settings: GatewaySettings | None = None) ->
FastAPI:
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
app.state.settings = s
- app.state.trino_client = trino_client
+ app.state.lakehouse_client = lakehouse_client
app.state.registry = registry
app.state.llm_readiness = LLMReadiness(s)
app.state.sessions = SessionStore(s)
@@ -79,6 +78,7 @@ def create_app(settings: GatewaySettings | None = None) ->
FastAPI:
yield
finally:
sweeper.cancel()
+ lakehouse_client.close()
app = FastAPI(title="hudi-agent-gateway", version=__version__,
lifespan=lifespan)
app.include_router(meta.router)
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/tools/__init__.py
b/hudi-agent-gateway/src/hudi_agent_gateway/cancellation.py
similarity index 52%
copy from hudi-agent-gateway/src/hudi_agent_gateway/tools/__init__.py
copy to hudi-agent-gateway/src/hudi_agent_gateway/cancellation.py
index 9e4db82cdc7b..30d8b70f5cac 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/tools/__init__.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/cancellation.py
@@ -14,24 +14,28 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""Tool assembly: the single place where tool modules get registered.
+"""Track tool tasks belonging to one streaming chat turn."""
-Adding a new tool module = write the module with a ``register(registry, ...)``
-function, then call it here.
-"""
+import asyncio
+from contextvars import ContextVar
-from __future__ import annotations
-from hudi_agent_gateway.config import GatewaySettings
-from hudi_agent_gateway.tools import trino_tools
-from hudi_agent_gateway.tools.registry import ToolRegistry
-from hudi_agent_gateway.tools.trino_client import TrinoClient
+class ToolCancellationScope:
+ def __init__(self) -> None:
+ self.tasks: set[asyncio.Task] = set()
+ self.cancelled = False
+ def add(self, task: asyncio.Task) -> None:
+ if self.cancelled:
+ raise asyncio.CancelledError
+ self.tasks.add(task)
-def build_registry(
- settings: GatewaySettings, trino_client: TrinoClient | None = None
-) -> ToolRegistry:
- registry = ToolRegistry()
- client = trino_client or TrinoClient(settings)
- trino_tools.register(registry, client, settings)
- return registry
+ def cancel(self) -> None:
+ self.cancelled = True
+ for task in tuple(self.tasks):
+ task.cancel()
+
+
+tool_cancellation_scope: ContextVar[ToolCancellationScope | None] = ContextVar(
+ "tool_cancellation_scope", default=None
+)
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/config.py
b/hudi-agent-gateway/src/hudi_agent_gateway/config.py
index 46f2d9836a08..34b300ee7d8a 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/config.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/config.py
@@ -25,6 +25,7 @@ through ``/ready`` instead.
from __future__ import annotations
from typing import Literal
+from urllib.parse import urlparse
from pydantic import AliasChoices, Field, model_validator
from pydantic_settings import BaseSettings, SettingsConfigDict
@@ -33,9 +34,7 @@ LLMProvider = Literal["anthropic", "openai", "ollama",
"openai-compatible"]
class GatewaySettings(BaseSettings):
- model_config = SettingsConfigDict(
- env_prefix="GATEWAY_", extra="ignore", populate_by_name=True
- )
+ model_config = SettingsConfigDict(env_prefix="GATEWAY_", extra="ignore",
populate_by_name=True)
# --- server ---
host: str = "0.0.0.0"
@@ -58,7 +57,13 @@ class GatewaySettings(BaseSettings):
ollama_base_url: str = "http://localhost:11434";
openai_base_url: str = "" # required for provider=openai-compatible (vLLM
etc.)
- # --- Trino / lakehouse ---
+ # --- lakehouse ---
+ sql_backend: Literal["trino", "spark-connect"] = "trino"
+ spark_connect_uri: str = "sc://hudi-spark-connect.hudi-lakehouse.svc:15002"
+ spark_connect_catalog: str = "spark_catalog"
+ spark_connect_schema: str = "default"
+
+ # --- Trino ---
trino_host: str = "hudi-trino.hudi-lakehouse.svc"
trino_port: int = 8080
trino_user: str = "hudi-agent-gateway"
@@ -66,12 +71,13 @@ class GatewaySettings(BaseSettings):
trino_schema: str = "default"
# --- guardrails / shaping ---
- sql_row_cap: int = 200
- sql_timeout_seconds: float = 120.0
- tool_result_max_bytes: int = 50_000
+ sql_row_cap: int = Field(default=200, gt=0)
+ sql_timeout_seconds: float = Field(default=120.0, gt=0)
+ tool_result_max_bytes: int = Field(default=50_000, ge=80)
# --- agent / sessions ---
agent_max_iterations: int = 25
+ chat_timeout_seconds: float = Field(default=300.0, gt=0)
session_ttl_seconds: float = 3600.0
max_sessions: int = 1000
max_messages_per_session: int = 40
@@ -80,6 +86,53 @@ class GatewaySettings(BaseSettings):
# --- surfaces ---
mcp_enabled: bool = True
+ @property
+ def sql_dialect(self) -> str:
+ return "spark" if self.sql_backend == "spark-connect" else "trino"
+
+ @property
+ def lakehouse_catalog(self) -> str:
+ return (
+ self.spark_connect_catalog
+ if self.sql_backend == "spark-connect"
+ else self.trino_catalog
+ )
+
+ @property
+ def lakehouse_schema(self) -> str:
+ return (
+ self.spark_connect_schema if self.sql_backend == "spark-connect"
else self.trino_schema
+ )
+
+ @model_validator(mode="after")
+ def _validate_spark_connect(self) -> GatewaySettings:
+ if self.sql_backend == "spark-connect":
+ try:
+ # PySpark parses semicolon parameters with an HTTP URI parser.
+ uri = urlparse("http" + self.spark_connect_uri[2:])
+ port = uri.port
+ except ValueError as e:
+ raise ValueError(
+ "GATEWAY_SPARK_CONNECT_URI requires a valid URI with a
port from 1 to 65535"
+ ) from e
+ if port is not None and not 1 <= port <= 65535:
+ raise ValueError("GATEWAY_SPARK_CONNECT_URI port must be from
1 to 65535")
+ if (
+ not self.spark_connect_uri.startswith("sc://")
+ or not uri.hostname
+ or uri.username is not None
+ or uri.password is not None
+ or uri.path not in ("", "/")
+ or uri.query
+ or uri.fragment
+ or (uri.params and any(len(p.split("=")) != 2 for p in
uri.params.split(";")))
+ ):
+ raise ValueError(
+ "GATEWAY_SPARK_CONNECT_URI requires sc://host:port "
+ "with optional /;key=value parameters"
+ )
+ return self
+
@model_validator(mode="after")
def _validate_provider_requirements(self) -> GatewaySettings:
if self.llm_provider == "anthropic" and not self.anthropic_api_key:
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/mcp_server.py
b/hudi-agent-gateway/src/hudi_agent_gateway/mcp_server.py
index f1d17fd55365..308e42f02705 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/mcp_server.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/mcp_server.py
@@ -24,9 +24,9 @@ from hudi_agent_gateway import __version__
from hudi_agent_gateway.tools.registry import ToolRegistry
_INSTRUCTIONS = (
- "Tools for querying an Apache Hudi lakehouse through Trino. Use
list_tables "
+ "Tools for querying an Apache Hudi lakehouse. Use list_tables "
"and describe_table to discover the schema, then query_lakehouse with a "
- "single read-only SELECT statement (Trino SQL). Results are row-capped and
"
+ "single read-only SELECT statement in the tool's SQL dialect. Results are
row-capped and "
"may be truncated (indicated in the payload)."
)
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/tools/__init__.py
b/hudi-agent-gateway/src/hudi_agent_gateway/tools/__init__.py
index 9e4db82cdc7b..91ed10c1b4c9 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/tools/__init__.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/tools/__init__.py
@@ -24,14 +24,28 @@ from __future__ import annotations
from hudi_agent_gateway.config import GatewaySettings
from hudi_agent_gateway.tools import trino_tools
+from hudi_agent_gateway.tools.lakehouse_client import LakehouseClient
from hudi_agent_gateway.tools.registry import ToolRegistry
from hudi_agent_gateway.tools.trino_client import TrinoClient
+def build_client(settings: GatewaySettings) -> LakehouseClient:
+ if settings.sql_backend == "spark-connect":
+ try:
+ from hudi_agent_gateway.tools.spark_connect_client import
SparkConnectClient
+ except ModuleNotFoundError as e:
+ if e.name and (e.name == "pyspark" or
e.name.startswith("pyspark.")):
+ raise RuntimeError("Spark Connect requires the spark-connect
extra") from e
+ raise
+
+ return SparkConnectClient(settings)
+ return TrinoClient(settings)
+
+
def build_registry(
- settings: GatewaySettings, trino_client: TrinoClient | None = None
+ settings: GatewaySettings, lakehouse_client: LakehouseClient | None = None
) -> ToolRegistry:
registry = ToolRegistry()
- client = trino_client or TrinoClient(settings)
+ client = lakehouse_client or build_client(settings)
trino_tools.register(registry, client, settings)
return registry
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/tools/guardrails.py
b/hudi-agent-gateway/src/hudi_agent_gateway/tools/guardrails.py
index ec261eada2ca..8bf048fe035d 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/tools/guardrails.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/tools/guardrails.py
@@ -16,7 +16,7 @@
"""SQL guardrails for model-written queries.
-AST-level (sqlglot, Trino dialect) rather than regex: read-only enforcement
+AST-level (sqlglot, selected backend dialect) rather than regex: read-only
enforcement
survives comments, CTEs, and string literals; the row cap is injected as a
real ``LIMIT``. Fail-closed: anything that does not parse is rejected.
"""
@@ -30,6 +30,7 @@ from hudi_agent_gateway.tools.registry import ToolInputError
_FORBIDDEN_NODES: tuple[type[exp.Expression], ...] = (
exp.Insert,
+ exp.Into,
exp.Update,
exp.Delete,
exp.Merge,
@@ -42,31 +43,41 @@ _FORBIDDEN_NODES: tuple[type[exp.Expression], ...] = (
exp.Command, # also covers CALL and EXPLAIN, which parse as generic
commands
exp.Use,
)
-
-_HINT = "Provide exactly one read-only SELECT statement (Trino SQL)."
-
-
-def enforce_guardrails(sql: str, row_cap: int) -> str:
+_SPARK_FORBIDDEN_FUNCTIONS = frozenset({"java_method", "reflect",
"try_reflect", "identifier"})
+_SPARK_HUDI_TABLE_FUNCTIONS = frozenset({
+ "hudi_query", "hudi_table_changes", "hudi_query_timeline",
"hudi_filesystem_view",
+ "hudi_metadata",
+})
+_SPARK_TABLE_FUNCTIONS = _SPARK_HUDI_TABLE_FUNCTIONS | frozenset({
+ "range", "explode", "explode_outer", "posexplode", "posexplode_outer",
"inline",
+ "inline_outer", "stack", "json_tuple",
+})
+
+
+def enforce_guardrails(
+ sql: str, row_cap: int, *, dialect: str = "trino", spark_catalog: str =
"spark_catalog"
+) -> str:
"""Validate ``sql`` as a single read-only SELECT and cap its LIMIT.
Returns the (possibly rewritten) SQL to execute. Raises
:class:`ToolInputError` on any violation.
"""
+ hint = f"Provide exactly one read-only SELECT statement ({dialect} SQL)."
try:
- statements = sqlglot.parse(sql, read="trino")
+ statements = sqlglot.parse(sql, read=dialect)
except sqlglot.errors.SqlglotError as e:
- raise ToolInputError(f"could not parse SQL: {e}", hint=_HINT) from e
+ raise ToolInputError(f"could not parse SQL: {e}", hint=hint) from e
statements = [s for s in statements if s is not None]
if len(statements) != 1:
raise ToolInputError(
- f"exactly one statement is allowed, got {len(statements)}",
hint=_HINT
+ f"exactly one statement is allowed, got {len(statements)}",
hint=hint
)
stmt = statements[0]
if not isinstance(stmt, (exp.Select, exp.SetOperation)):
raise ToolInputError(
- f"read-only: only SELECT is allowed, got {type(stmt).__name__}",
hint=_HINT
+ f"read-only: only SELECT is allowed, got {type(stmt).__name__}",
hint=hint
)
# Defense in depth: no write/DDL/command node anywhere in the tree
@@ -74,19 +85,108 @@ def enforce_guardrails(sql: str, row_cap: int) -> str:
for node in stmt.walk():
if isinstance(node, _FORBIDDEN_NODES):
raise ToolInputError(
- f"read-only: {type(node).__name__} is not allowed", hint=_HINT
+ f"read-only: {type(node).__name__} is not allowed", hint=hint
)
+ if dialect == "spark":
+ _validate_spark_node(node, spark_catalog=spark_catalog, hint=hint)
+
+ return _cap_limit(stmt, row_cap).sql(dialect=dialect)
+
+
+def _validate_spark_node(node: object, *, spark_catalog: str, hint: str) ->
None:
+ if isinstance(node, exp.QueryTransform):
+ raise ToolInputError("read-only: Spark script transforms are not
allowed", hint=hint)
+ if isinstance(node, exp.Table):
+ _validate_spark_table(node, spark_catalog=spark_catalog, hint=hint)
+ elif isinstance(node, exp.Lateral):
+ # LATERAL functions are not necessarily wrapped in an exp.Table.
+ if not isinstance(node.this, exp.Subquery):
+ _validate_spark_table_function(node.this)
+ elif isinstance(node, exp.Func):
+ name = _spark_function_name(node)
+ if name in _SPARK_FORBIDDEN_FUNCTIONS:
+ # Reflection invokes JVM methods; IDENTIFIER hides names from AST
checks.
+ raise ToolInputError(f"read-only: Spark function {name} is not
allowed", hint=hint)
+ if name in _SPARK_HUDI_TABLE_FUNCTIONS:
+ # Check arguments at the function node, regardless of its AST
context.
+ _validate_hudi_table_function(node, name=name)
+
+
+def _validate_spark_table(node: exp.Table, *, spark_catalog: str, hint: str)
-> None:
+ if not isinstance(node.this, exp.Identifier):
+ _validate_spark_table_function(node.this, qualified=bool(node.db or
node.catalog))
+ return
+ if any(any(char in part.name for char in "/\\:") for part in node.parts):
+ raise ToolInputError("read-only: Spark file paths are not allowed",
hint=hint)
+ if not node.catalog and node.db:
+ # Any two-part name can be interpreted as a data source and path,
+ # including custom providers. Force catalog table resolution.
+ node.set("catalog", exp.to_identifier(spark_catalog, quoted=True))
+
+
+def _spark_function_name(node: exp.Func) -> str:
+ return (node.name if isinstance(node, exp.Anonymous) else
node.sql_name()).lower()
+
+
+def _validate_spark_table_function(node: exp.Expression, *, qualified: bool =
False) -> None:
+ # Unknown/custom TVFs may read files even with
spark.sql.runSQLOnFiles=false.
+ name = _spark_function_name(node) if isinstance(node, exp.Func) else
node.sql(dialect="spark")
+ if qualified or not isinstance(node, exp.Func) or name not in
_SPARK_TABLE_FUNCTIONS:
+ raise ToolInputError(
+ f"read-only: Spark table function {name} is not allowed",
+ hint="Use an unqualified table function: "
+ + ", ".join(sorted(_SPARK_TABLE_FUNCTIONS)) + ".",
+ )
+
+
+def _validate_hudi_table_function(node: exp.Func, *, name: str) -> None:
+ args = node.expressions
+ hint = f"Pass a literal table name as the first argument to {name}, e.g.
'default.orders'."
+ _validate_hudi_table_argument(args[0] if args else None, hint=hint)
+ # The optional filesystem-view subpath is appended to the catalog location
and
+ # can traverse outside it. Only the whole-table form is exposed by the
gateway.
+ if name == "hudi_filesystem_view" and len(args) != 1:
+ raise ToolInputError(
+ "read-only: Hudi filesystem-view subpaths are not allowed",
+ hint="Use hudi_filesystem_view('default.orders') without a
subpath.",
+ )
- return _cap_limit(stmt, row_cap).sql(dialect="trino")
+
+def _validate_hudi_table_argument(table: exp.Expression | None, *, hint: str)
-> None:
+ error = "read-only: Hudi table functions require a literal catalog table
identifier"
+ # Hudi evaluates expressions before deciding whether the first argument is
a path.
+ # Reject computed values, including concat/chr/casts, rather than trying
to fold them.
+ if not isinstance(table, exp.Literal) or not table.is_string:
+ raise ToolInputError(error, hint=hint)
+ if not table.this or any(char in table.this for char in "/\\:"):
+ raise ToolInputError(error, hint=hint)
+ try:
+ identifier = sqlglot.parse_one(table.this, read="spark",
into=exp.Table)
+ except sqlglot.errors.SqlglotError as e:
+ raise ToolInputError(error, hint=hint) from e
+ if (not all(isinstance(part, exp.Identifier) for part in identifier.parts)
+ or identifier.catalog
+ or any(key not in {"this", "db", "catalog"} for key in
identifier.args)):
+ raise ToolInputError(error, hint=hint)
def _cap_limit(stmt: exp.Query, row_cap: int) -> exp.Query:
limit = stmt.args.get("limit")
- if limit is not None:
- try:
- existing = int(limit.expression.name)
- except (TypeError, ValueError):
- existing = None
- if existing is not None and existing <= row_cap:
- return stmt
+ if isinstance(limit, (exp.Limit, exp.Fetch)):
+ count = limit.args.get("count") if isinstance(limit, exp.Fetch) else
limit.expression
+ if isinstance(limit, exp.Fetch) and count is None:
+ count = exp.Literal.number(1)
+ limit.set("count", count)
+ if isinstance(count, exp.Neg):
+ raise ToolInputError(
+ "row limit must be non-negative",
+ hint="Use LIMIT or FETCH with a non-negative integer.",
+ )
+ options = limit.args.get("limit_options")
+ if options and (options.args.get("with_ties") or
options.args.get("percent")):
+ # A count with ties or a percentage is not a hard row bound. Keep
the
+ # original selection semantics, then cap the result in an outer
query.
+ return
exp.select("*").from_(stmt.subquery("_guarded_query")).limit(row_cap)
+ if isinstance(count, exp.Literal) and count.is_int:
+ return stmt.limit(min(int(count.this), row_cap))
return stmt.limit(row_cap)
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/mcp_server.py
b/hudi-agent-gateway/src/hudi_agent_gateway/tools/lakehouse_client.py
similarity index 50%
copy from hudi-agent-gateway/src/hudi_agent_gateway/mcp_server.py
copy to hudi-agent-gateway/src/hudi_agent_gateway/tools/lakehouse_client.py
index f1d17fd55365..63be3a05a23a 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/mcp_server.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/tools/lakehouse_client.py
@@ -14,28 +14,32 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""The MCP surface: the same tool registry, exposed to external agents."""
+"""Result and error contract shared by the lakehouse backends."""
-from __future__ import annotations
+from dataclasses import dataclass, field
+from typing import Any, Protocol
-from fastmcp import FastMCP
-from hudi_agent_gateway import __version__
-from hudi_agent_gateway.tools.registry import ToolRegistry
+class QueryError(Exception):
+ """A lakehouse query failed."""
-_INSTRUCTIONS = (
- "Tools for querying an Apache Hudi lakehouse through Trino. Use
list_tables "
- "and describe_table to discover the schema, then query_lakehouse with a "
- "single read-only SELECT statement (Trino SQL). Results are row-capped and
"
- "may be truncated (indicated in the payload)."
-)
+class QueryTimeoutError(QueryError):
+ pass
-def build_mcp(registry: ToolRegistry) -> FastMCP:
- mcp: FastMCP = FastMCP(
- name="hudi-lakehouse",
- version=__version__,
- instructions=_INSTRUCTIONS,
- )
- registry.attach_to_mcp(mcp)
- return mcp
+
+@dataclass
+class QueryResult:
+ columns: list[str]
+ rows: list[list[Any]]
+ row_count: int = field(init=False)
+
+ def __post_init__(self) -> None:
+ self.row_count = len(self.rows)
+
+
+class LakehouseClient(Protocol):
+ async def execute(self, sql: str, *, timeout: float, max_rows: int) ->
QueryResult: ...
+ async def describe(self, sql: str, *, timeout: float, max_rows: int) ->
QueryResult: ...
+ async def ping(self, timeout: float = 5.0, cache_seconds: float = 10.0) ->
bool: ...
+ def close(self) -> None: ...
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/tools/registry.py
b/hudi-agent-gateway/src/hudi_agent_gateway/tools/registry.py
index 92aa6106af92..6ddda97051c3 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/tools/registry.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/tools/registry.py
@@ -29,6 +29,7 @@ outcome) -- the same trace regardless of which surface called
the tool.
from __future__ import annotations
+import asyncio
import functools
import inspect
import json
@@ -41,6 +42,7 @@ from typing import Any
from langchain_core.tools import BaseTool, StructuredTool
from pydantic import BaseModel
+from hudi_agent_gateway.cancellation import tool_cancellation_scope
from hudi_agent_gateway.log import log_event
logger = logging.getLogger("hudi_agent_gateway.tools")
@@ -81,15 +83,24 @@ def _with_invocation_logging(name: str, fn: ToolHandler) ->
ToolHandler:
async def wrapped(*args: Any, **kwargs: Any) -> str:
start = time.monotonic()
outcome = "ok"
+ scope = tool_cancellation_scope.get()
+ task = asyncio.current_task()
+ if scope is not None and task is not None:
+ scope.add(task)
try:
result = await fn(*args, **kwargs)
if is_tool_error(result):
outcome = "tool_error"
return result
+ except asyncio.CancelledError:
+ outcome = "cancelled"
+ raise
except Exception:
outcome = "exception"
raise
finally:
+ if scope is not None and task is not None:
+ scope.tasks.discard(task)
log_event(
logger,
"tool_invocation",
diff --git
a/hudi-agent-gateway/src/hudi_agent_gateway/tools/spark_connect_client.py
b/hudi-agent-gateway/src/hudi_agent_gateway/tools/spark_connect_client.py
new file mode 100644
index 000000000000..247fc972e8e9
--- /dev/null
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/tools/spark_connect_client.py
@@ -0,0 +1,213 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Bounded async facade over the official PySpark Connect client.
+
+Each call owns a remote session. Timeout/cancellation requests interruptAll on
+that session off the event loop, but cannot guarantee when Spark stops.
Capacity
+is retained until the worker finishes, including session cleanup.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import threading
+import time
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass, field
+from typing import Any
+from urllib.parse import unquote, urlparse
+
+from pyspark.errors import PySparkException
+from pyspark.sql.connect.session import SparkSession
+
+from hudi_agent_gateway.config import GatewaySettings
+from hudi_agent_gateway.tools.lakehouse_client import QueryError, QueryResult,
QueryTimeoutError
+
+
+class _RequestSparkSession(SparkSession):
+ def __del__(self) -> None:
+ # PySpark 4.1 closes even an already-stopped session again during GC.
+ # close() shuts down a process-wide release executor. GC can run while
+ # another thread starts a worker under that executor's lock,
deadlocking
+ # thread startup and every subsequent executor submission (including
the
+ # web event loop). Our worker already calls stop() in its finally
block.
+ client = getattr(self, "_client", None)
+ if client is not None and client.is_closed:
+ return
+ super().__del__()
+
+
+@dataclass
+class _Operation:
+ cancelled: threading.Event = field(default_factory=threading.Event)
+ lock: Any = field(default_factory=threading.Lock)
+ session: SparkSession | None = None
+
+ def cancel(self) -> None:
+ # Serialize cancellation with session removal/close, never with
execution.
+ with self.lock:
+ if self.session is not None:
+ with contextlib.suppress(Exception):
+ self.session.interruptAll()
+
+ def check_cancelled(self) -> None:
+ if self.cancelled.is_set():
+ raise QueryError("query cancelled before execution")
+
+
+class SparkConnectClient:
+ def __init__(self, settings: GatewaySettings) -> None:
+ self._settings = settings
+ self._executor = ThreadPoolExecutor(max_workers=8,
thread_name_prefix="spark-query")
+ self._ping_executor = ThreadPoolExecutor(max_workers=1,
thread_name_prefix="spark-ping")
+ self._cancel_executor = ThreadPoolExecutor(max_workers=8,
thread_name_prefix="spark-cancel")
+ # ThreadPoolExecutor's queue is unbounded. Hold a slot until the worker
+ # finishes, even after its async caller timed out.
+ self._slots = threading.BoundedSemaphore(8)
+ self._ping_slot = threading.BoundedSemaphore(1)
+ self._ping_cache: tuple[float, bool] | None = None
+ self._ping_task: asyncio.Task[bool] | None = None
+ self._closed = False
+ self._operations_lock = threading.Lock()
+ self._operations: dict[int, _Operation] = {}
+
+ def close(self) -> None:
+ with self._operations_lock:
+ self._closed = True
+ operations = list(self._operations.values())
+ for op in operations:
+ op.cancelled.set()
+ with contextlib.suppress(RuntimeError):
+ self._cancel_executor.submit(op.cancel)
+ self._executor.shutdown(wait=False, cancel_futures=True)
+ self._ping_executor.shutdown(wait=False, cancel_futures=True)
+ self._cancel_executor.shutdown(wait=False)
+
+ def _execute_sync(
+ self, sql: str, max_rows: int, op: _Operation, describe: bool = False
+ ) -> QueryResult:
+ op.check_cancelled()
+ # Construct directly: getOrCreate() can reuse another request's
session.
+ spark =
_RequestSparkSession(connection=self._settings.spark_connect_uri)
+ with op.lock:
+ op.session = spark
+ try:
+ op.check_cancelled()
+
spark.catalog.setCurrentCatalog(self._settings.spark_connect_catalog)
+ op.check_cancelled()
+
spark.catalog.setCurrentDatabase(self._settings.spark_connect_schema)
+ op.check_cancelled()
+ df = spark.sql(sql)
+ op.check_cancelled()
+ if describe:
+ return QueryResult(
+ columns=["column_name", "data_type", "is_nullable"],
+ rows=[
+ [f.name, f.dataType.simpleString(), "YES" if
f.nullable else "NO"]
+ for f in df.schema.fields[:max_rows]
+ ],
+ )
+ # Limit before collecting, including metadata queries outside
guardrails.
+ df = df.limit(max_rows)
+ columns = df.columns
+ op.check_cancelled()
+ return QueryResult(columns=columns, rows=[list(row) for row in
df.collect()])
+ finally:
+ with op.lock:
+ op.session = None
+ spark.stop()
+
+ async def execute(self, sql: str, *, timeout: float, max_rows: int) ->
QueryResult:
+ return await self._run(sql, timeout=timeout, max_rows=max_rows)
+
+ async def describe(self, sql: str, *, timeout: float, max_rows: int) ->
QueryResult:
+ return await self._run(sql, timeout=timeout, max_rows=max_rows,
describe=True)
+
+ async def _run(
+ self, sql: str, *, timeout: float, max_rows: int, describe: bool =
False,
+ probe: bool = False,
+ ) -> QueryResult:
+ slots = self._ping_slot if probe else self._slots
+ executor = self._ping_executor if probe else self._executor
+ op = _Operation()
+ with self._operations_lock:
+ if self._closed:
+ raise QueryError("Spark Connect client is closed")
+ if not slots.acquire(blocking=False):
+ raise QueryError(
+ "Spark Connect query capacity exhausted; retry after
active queries finish"
+ )
+ try:
+ work = executor.submit(self._execute_sync, sql, max_rows, op,
describe)
+ except Exception:
+ slots.release()
+ raise
+ self._operations[id(op)] = op
+
+ def release_operation(_: Any) -> None:
+ with self._operations_lock:
+ self._operations.pop(id(op), None)
+ slots.release()
+
+ work.add_done_callback(release_operation)
+ future = asyncio.wrap_future(work)
+ try:
+ return await asyncio.wait_for(future, timeout)
+ except (TimeoutError, asyncio.CancelledError) as e:
+ op.cancelled.set()
+ # No blocking Connect RPC (including interrupt/stop) runs on the
event loop.
+ with contextlib.suppress(RuntimeError):
+ self._cancel_executor.submit(op.cancel)
+ if isinstance(e, asyncio.CancelledError):
+ raise
+ raise QueryTimeoutError(
+ f"query exceeded the {timeout:g}s timeout; Spark may still be
running it"
+ ) from e
+ except QueryError:
+ raise
+ except PySparkException as e:
+ message = str(e).replace(self._settings.spark_connect_uri, "[Spark
Connect endpoint]")
+ uri = urlparse("http" + self._settings.spark_connect_uri[2:])
+ # URI parameters can contain tokens or custom authentication
headers.
+ for parameter in uri.params.split(";"):
+ _, _, value = parameter.partition("=")
+ if value:
+ message = message.replace(value, "[redacted]").replace(
+ unquote(value), "[redacted]"
+ )
+ raise QueryError(message[:2000]) from e
+ except Exception as e:
+ raise QueryError(f"Spark Connect query failed
({type(e).__name__})") from e
+
+ async def ping(self, timeout: float = 5.0, cache_seconds: float = 10.0) ->
bool:
+ if self._ping_cache is not None and time.monotonic() -
self._ping_cache[0] < cache_seconds:
+ return self._ping_cache[1]
+
+ if self._ping_task is None or self._ping_task.done():
+ self._ping_task = asyncio.create_task(self._probe(timeout))
+ # Cancelling one readiness request must not cancel a shared probe.
+ return await asyncio.shield(self._ping_task)
+
+ async def _probe(self, timeout: float) -> bool:
+ try:
+ await self._run("SELECT 1", timeout=timeout, max_rows=1,
probe=True)
+ ok = True
+ except Exception:
+ ok = False
+ self._ping_cache = (time.monotonic(), ok)
+ return ok
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_client.py
b/hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_client.py
index 2eda8730a454..bf35924333e1 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_client.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_client.py
@@ -22,30 +22,20 @@ import asyncio
import contextlib
import time
from concurrent.futures import ThreadPoolExecutor
-from dataclasses import dataclass, field
from typing import Any
import trino
from hudi_agent_gateway.config import GatewaySettings
-
-
-class TrinoQueryError(Exception):
- """A query failed on the Trino side; message is safe to surface to the
model."""
-
-
-class TrinoTimeoutError(TrinoQueryError):
- pass
-
-
-@dataclass
-class QueryResult:
- columns: list[str]
- rows: list[list[Any]]
- row_count: int = field(init=False)
-
- def __post_init__(self) -> None:
- self.row_count = len(self.rows)
+from hudi_agent_gateway.tools.lakehouse_client import (
+ QueryError as TrinoQueryError,
+)
+from hudi_agent_gateway.tools.lakehouse_client import (
+ QueryResult as QueryResult,
+)
+from hudi_agent_gateway.tools.lakehouse_client import (
+ QueryTimeoutError as TrinoTimeoutError,
+)
class TrinoClient:
@@ -61,6 +51,9 @@ class TrinoClient:
self._executor = ThreadPoolExecutor(max_workers=8,
thread_name_prefix="trino-query")
self._ping_cache: tuple[float, bool] | None = None
+ def close(self) -> None:
+ self._executor.shutdown(wait=False, cancel_futures=True)
+
def _connect(self) -> trino.dbapi.Connection:
s = self._settings
return trino.dbapi.connect(
@@ -105,6 +98,9 @@ class TrinoClient:
except trino.exceptions.Error as e:
raise TrinoQueryError(str(e)) from e
+ async def describe(self, sql: str, *, timeout: float, max_rows: int) ->
QueryResult:
+ return await self.execute(sql, timeout=timeout, max_rows=max_rows)
+
async def ping(self, timeout: float = 5.0, cache_seconds: float = 10.0) ->
bool:
"""Cached reachability check feeding ``/ready``.
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_tools.py
b/hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_tools.py
index 95b6b33082d6..986677116488 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_tools.py
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_tools.py
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""Lakehouse tools backed by Trino.
+"""Lakehouse tools backed by Trino or Spark Connect.
Handlers return JSON strings. Expected failures (guardrail rejections, query
errors, timeouts) are returned as ``{"error": ..., "hint": ...}`` payloads
@@ -32,13 +32,13 @@ from pydantic import Field
from hudi_agent_gateway.config import GatewaySettings
from hudi_agent_gateway.tools.guardrails import enforce_guardrails
-from hudi_agent_gateway.tools.registry import ToolInputError, ToolRegistry
-from hudi_agent_gateway.tools.trino_client import (
+from hudi_agent_gateway.tools.lakehouse_client import (
+ LakehouseClient,
+ QueryError,
QueryResult,
- TrinoClient,
- TrinoQueryError,
- TrinoTimeoutError,
+ QueryTimeoutError,
)
+from hudi_agent_gateway.tools.registry import ToolInputError, ToolRegistry
_QUERY_DESC = (
"Run a single read-only SELECT statement (Trino SQL) against the lakehouse
"
@@ -87,15 +87,17 @@ def shape_result(result: QueryResult, *, max_bytes: int,
sql: str, row_cap: int)
if len(text.encode()) > max_bytes:
# Fail-safe: even zero rows can bust the limit (huge sql text, very
# many columns, or a tiny configured cap). Return a bounded payload.
- text = json.dumps({
- "sql": sql[:100],
- "column_count": len(result.columns),
- "rows": [],
- "row_count": result.row_count,
- "truncated": True,
- "notice": "Result exceeded the size limit; rows and column names
omitted. "
- "Narrow the query or raise the result size limit.",
- })
+ text = json.dumps(
+ {
+ "sql": sql[:100],
+ "column_count": len(result.columns),
+ "rows": [],
+ "row_count": result.row_count,
+ "truncated": True,
+ "notice": "Result exceeded the size limit; rows and column
names omitted. "
+ "Narrow the query or raise the result size limit.",
+ }
+ )
if len(text.encode()) > max_bytes:
text = json.dumps({"error": "result exceeded the size limit",
"truncated": True})
return text
@@ -129,9 +131,9 @@ def _validate_identifier(value: str, kind: str) -> str:
def _split_table_name(table: str, settings: GatewaySettings) -> tuple[str,
str, str]:
parts = table.split(".")
if len(parts) == 1:
- return settings.trino_catalog, settings.trino_schema, parts[0]
+ return settings.lakehouse_catalog, settings.lakehouse_schema, parts[0]
if len(parts) == 2:
- return settings.trino_catalog, parts[0], parts[1]
+ return settings.lakehouse_catalog, parts[0], parts[1]
if len(parts) == 3:
return parts[0], parts[1], parts[2]
raise ToolInputError(
@@ -139,28 +141,37 @@ def _split_table_name(table: str, settings:
GatewaySettings) -> tuple[str, str,
)
-def register(registry: ToolRegistry, client: TrinoClient, settings:
GatewaySettings) -> None:
- @registry.register("query_lakehouse", _QUERY_DESC)
+def register(registry: ToolRegistry, client: LakehouseClient, settings:
GatewaySettings) -> None:
+ spark = settings.sql_backend == "spark-connect"
+ dialect = settings.sql_dialect
+
+ @registry.register(
+ "query_lakehouse", _QUERY_DESC.replace("Trino", "Spark" if spark else
"Trino")
+ )
async def query_lakehouse(
sql: Annotated[
- str, Field(description="A single read-only SELECT statement in
Trino SQL.")
+ str,
+ Field(description="A single read-only SELECT in the configured
backend's SQL dialect."),
],
) -> str:
try:
- safe_sql = enforce_guardrails(sql, row_cap=settings.sql_row_cap)
+ safe_sql = enforce_guardrails(
+ sql, row_cap=settings.sql_row_cap, dialect=dialect,
+ spark_catalog=settings.spark_connect_catalog,
+ )
result = await client.execute(
safe_sql, timeout=settings.sql_timeout_seconds,
max_rows=settings.sql_row_cap
)
except ToolInputError as e:
return _error(str(e), e.hint)
- except TrinoTimeoutError as e:
+ except QueryTimeoutError as e:
return _error(
f"query failed: {e}",
f"The query hit the gateway's
{settings.sql_timeout_seconds:.0f}s limit -- "
"narrow it (partition filter, pre-aggregate, fewer columns) or
raise "
"GATEWAY_SQL_TIMEOUT_SECONDS.",
)
- except TrinoQueryError as e:
+ except QueryError as e:
return _error(f"query failed: {e}", "Fix the SQL and try again.")
return shape_result(
result,
@@ -179,20 +190,30 @@ def register(registry: ToolRegistry, client: TrinoClient,
settings: GatewaySetti
] = "",
) -> str:
try:
- cat = _validate_identifier(catalog or settings.trino_catalog,
"catalog")
- sch = _validate_identifier(schema_name or settings.trino_schema,
"schema")
+ cat = _validate_identifier(catalog or settings.lakehouse_catalog,
"catalog")
+ sch = _validate_identifier(schema_name or
settings.lakehouse_schema, "schema")
except ToolInputError as e:
return _error(str(e), e.hint)
sql = (
- f'SELECT table_schema, table_name, table_type FROM
"{cat}".information_schema.tables '
- f"WHERE table_schema = '{sch}' ORDER BY table_name"
+ f"SHOW TABLES IN `{cat}`.`{sch}`"
+ if spark
+ else (
+ "SELECT table_schema, table_name, table_type "
+ f'FROM "{cat}".information_schema.tables '
+ f"WHERE table_schema = '{sch}' ORDER BY table_name"
+ )
)
try:
result = await client.execute(
sql, timeout=settings.sql_timeout_seconds,
max_rows=settings.sql_row_cap
)
- except TrinoQueryError as e:
+ except QueryError as e:
return _error(f"listing tables failed: {e}")
+ if spark:
+ result = QueryResult(
+ columns=["table_schema", "table_name", "table_type"],
+ rows=[[row[0], row[1], "TEMPORARY" if row[2] else "TABLE"] for
row in result.rows],
+ )
return shape_result(
result, max_bytes=settings.tool_result_max_bytes, sql=sql,
row_cap=settings.sql_row_cap
)
@@ -211,15 +232,19 @@ def register(registry: ToolRegistry, client: TrinoClient,
settings: GatewaySetti
except ToolInputError as e:
return _error(str(e), e.hint)
sql = (
- "SELECT column_name, data_type, is_nullable "
- f'FROM "{cat}".information_schema.columns '
- f"WHERE table_schema = '{sch}' AND table_name = '{tbl}' ORDER BY
ordinal_position"
+ f"SELECT * FROM `{cat}`.`{sch}`.`{tbl}` LIMIT 0"
+ if spark
+ else (
+ "SELECT column_name, data_type, is_nullable "
+ f'FROM "{cat}".information_schema.columns '
+ f"WHERE table_schema = '{sch}' AND table_name = '{tbl}' ORDER
BY ordinal_position"
+ )
)
try:
- result = await client.execute(
+ result = await (client.describe if spark else client.execute)(
sql, timeout=settings.sql_timeout_seconds,
max_rows=settings.sql_row_cap
)
- except TrinoQueryError as e:
+ except QueryError as e:
return _error(f"describe failed: {e}")
if result.row_count == 0:
return _error(
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/ui/app.js
b/hudi-agent-gateway/src/hudi_agent_gateway/ui/app.js
index 8fa4186eccf4..0be5f6635bb1 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/ui/app.js
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/ui/app.js
@@ -27,6 +27,11 @@
var messagesEl = document.getElementById("messages");
var welcomeEl = document.getElementById("welcome");
var inputEl = document.getElementById("input");
+ var activeController = null;
+ var stopEl = document.getElementById("stop");
+ stopEl.addEventListener("click", function () {
+ if (activeController) activeController.abort();
+ });
var sendEl = document.getElementById("send");
var composerEl = document.getElementById("composer");
var modelInfoEl = document.getElementById("model-info");
@@ -49,7 +54,10 @@
fetch("../v1/info").then(function (r) { return r.json(); }).then(function
(info) {
modelInfoEl.textContent = info.provider + " · catalog " + info.catalog;
var base = window.location.origin;
- document.getElementById("ep-trino").textContent = info.trino_url;
+ document.getElementById("ep-trino").textContent = info.lakehouse_url ||
info.trino_url;
+ document.getElementById("ep-sql-label").textContent =
+ info.sql_backend === "spark-connect" ? "Spark Connect" : "Trino SQL";
+ document.getElementById("ep-trino-ui-row").hidden = !info.trino_ui_url;
var trinoUi = document.getElementById("ep-trino-ui");
trinoUi.textContent = info.trino_ui_url;
trinoUi.href = info.trino_ui_url;
@@ -158,6 +166,8 @@
function setBusy(value) {
busy = value;
sendEl.disabled = value;
+ stopEl.hidden = !value;
+ document.getElementById("clear-chat").disabled = value;
}
function send(text) {
@@ -165,6 +175,25 @@
setBusy(true);
var msg = el("div", "msg assistant", messagesEl);
+ var status = el("div", "response-status", msg);
+ status.setAttribute("role", "status");
+ var phase = "Connecting…";
+ var started = Date.now();
+ function updateStatus() {
+ status.textContent = phase + " · " + Math.floor((Date.now() - started) /
1000) + "s";
+ }
+ updateStatus();
+ var timer = setInterval(updateStatus, 1000);
+ var controller = new AbortController();
+ activeController = controller;
+ var connectionTimedOut = false;
+ var connectionTimer = setTimeout(function () {
+ connectionTimedOut = true;
+ controller.abort();
+ }, 15000);
+ var terminal = false;
+ var failed = false;
+ var stopped = false;
var bubble = el("div", "bubble md", msg);
var answer = "";
var cursor = el("span", "cursor", bubble);
@@ -181,8 +210,10 @@
fetch("../v1/chat", {
method: "POST",
headers: { "Content-Type": "application/json", "Accept":
"text/event-stream" },
+ signal: controller.signal,
body: JSON.stringify(payload)
}).then(function (resp) {
+ clearTimeout(connectionTimer);
if (!resp.ok || !resp.body) {
return resp.text().then(function (body) {
throw new Error("gateway returned " + resp.status + ": " +
body.slice(0, 300));
@@ -194,7 +225,10 @@
function pump() {
return reader.read().then(function (step) {
- if (step.done) { finish(); return null; }
+ if (step.done) {
+ if (!terminal) throw new Error("Connection closed before the
response completed.");
+ return null;
+ }
// Normalize SSE line endings (servers may emit \r\n) before framing.
buffer += decoder.decode(step.value, { stream: true
}).replace(/\r/g, "");
var frames = buffer.split("\n\n");
@@ -205,9 +239,16 @@
}
return pump();
}).catch(function (err) {
- addErrorBanner(err.message || String(err));
- finish();
- });
+ if (connectionTimedOut) {
+ failed = true;
+ addErrorBanner("Gateway did not respond within 15s. Check the service
connection and try again.");
+ } else if (err.name === "AbortError") {
+ stopped = true;
+ } else {
+ failed = true;
+ addErrorBanner(err.message || String(err));
+ }
+ }).finally(finish);
function handleFrame(frame) {
var event = "message";
@@ -220,10 +261,16 @@
var data;
try { data = JSON.parse(dataLines.join("\n")); } catch (e) { return; }
- if (event === "token") {
+ if (event === "status") {
+ phase = data.message;
+ updateStatus();
+ } else if (event === "token") {
+ phase = "Generating answer…";
answer += data.text;
renderAnswer();
} else if (event === "tool_call") {
+ phase = "Running " + data.name + "…";
+ updateStatus();
var chip = el("details", "tool-chip running", bubble.parentNode);
bubble.parentNode.insertBefore(chip, bubble);
var summary = el("summary", "", chip);
@@ -233,6 +280,8 @@
chips[data.id || data.name] = chip;
scrollToBottom();
} else if (event === "tool_result") {
+ phase = "Waiting for model…";
+ updateStatus();
var target = chips[data.id || data.name];
if (target) {
target.classList.remove("running");
@@ -244,8 +293,11 @@
}
scrollToBottom();
} else if (event === "error") {
+ terminal = true;
+ failed = true;
addErrorBanner(data.code + ": " + data.message);
} else if (event === "done") {
+ terminal = true;
if (!answer && data.message) {
answer = data.message;
renderAnswer();
@@ -254,10 +306,23 @@
}
function finish() {
+ clearTimeout(connectionTimer);
+ clearInterval(timer);
+ activeController = null;
+ phase = stopped ? "Stopped" : failed ? "Failed" : "Complete";
+ updateStatus();
+ if (stopped || failed) {
+ // A cancelled turn can leave unmatched tool calls in checkpoint
history.
+ sessionId = newSession();
+ status.textContent += " · Next message starts a new conversation.";
+ }
+ msg.querySelectorAll(".tool-chip.running").forEach(function (chip) {
+ chip.classList.remove("running");
+ });
setBusy(false);
if (cursor.parentNode) cursor.remove();
if (!answer && !msg.querySelector(".tool-chip")) {
- bubble.innerHTML = "<p><em>(no response)</em></p>";
+ bubble.textContent = stopped ? "Response stopped." : failed ?
"Response failed." : "(no response)";
} else {
bubble.innerHTML = window.HudiMarkdown.render(answer);
}
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/ui/index.html
b/hudi-agent-gateway/src/hudi_agent_gateway/ui/index.html
index d924ce4d047f..7cb68bed5d62 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/ui/index.html
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/ui/index.html
@@ -42,13 +42,13 @@ limitations under the License.
<section id="connect-panel" hidden>
<div class="endpoint">
<div>
- <strong>Trino SQL</strong>
+ <strong id="ep-sql-label">Lakehouse SQL</strong>
<small>point any SQL client / BI tool at the lakehouse directly</small>
</div>
<code id="ep-trino">…</code>
<button class="copy" data-copy="ep-trino" title="Copy">⎘</button>
</div>
- <div class="endpoint">
+ <div class="endpoint" id="ep-trino-ui-row">
<div>
<strong>Trino Web UI</strong>
<small>query monitoring & cluster overview</small>
@@ -72,7 +72,7 @@ limitations under the License.
<code id="ep-mcp">…</code>
<button class="copy" data-copy="ep-mcp" title="Copy">⎘</button>
</div>
- <p class="endpoint-note">Trino is addressed by its in-cluster service
name; from outside the
+ <p class="endpoint-note">The SQL endpoint is addressed by its in-cluster
service name; from outside the
cluster, port-forward or expose the service. The inference and MCP URLs
are this gateway.</p>
</section>
@@ -80,7 +80,7 @@ limitations under the License.
<div class="welcome" id="welcome">
<h1>Talk to your lakehouse</h1>
<p>Ask questions about your Hudi tables in plain English. The agent
- inspects the catalog, writes guarded Trino SQL, and answers from real
+ inspects the catalog, writes guarded read-only SQL, and answers from real
query results.</p>
<div class="suggestions">
<button class="suggestion">What tables are available?</button>
@@ -94,6 +94,7 @@ limitations under the License.
<form id="composer">
<textarea id="input" rows="1" placeholder="Ask about your data…"
autocomplete="off" autofocus></textarea>
+ <button id="stop" type="button" hidden>Stop</button>
<button id="send" type="submit" title="Send">➤</button>
</form>
<p class="fineprint">Answers are generated from live queries against your
lakehouse. Verify before acting.</p>
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/ui/style.css
b/hudi-agent-gateway/src/hudi_agent_gateway/ui/style.css
index b4c0e5e129a3..5340084627a5 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/ui/style.css
+++ b/hudi-agent-gateway/src/hudi_agent_gateway/ui/style.css
@@ -258,3 +258,7 @@ footer { background: var(--panel); border-top: 1px solid
var(--border); padding:
}
#send:disabled { opacity: 0.45; cursor: default; }
.fineprint { text-align: center; color: var(--muted); font-size: 11.5px;
margin: 8px 0 0; }
+
+.response-status { color: var(--muted); font-size: 12px; margin-bottom: 8px; }
+#stop { border: 1px solid var(--border); border-radius: 10px; padding: 10px
14px;
+ background: var(--bg); color: var(--text); cursor: pointer; }
diff --git a/hudi-agent-gateway/tests/conftest.py
b/hudi-agent-gateway/tests/conftest.py
index 66b8981bb5ba..7f8d64127e3e 100644
--- a/hudi-agent-gateway/tests/conftest.py
+++ b/hudi-agent-gateway/tests/conftest.py
@@ -63,7 +63,13 @@ class FakeTrinoClient:
raise self.error
return self.result
- async def ping(self, timeout: float = 5.0) -> bool:
+ async def describe(self, sql: str, *, timeout: float, max_rows: int) ->
QueryResult:
+ return await self.execute(sql, timeout=timeout, max_rows=max_rows)
+
+ def close(self) -> None:
+ pass
+
+ async def ping(self, timeout: float = 5.0, cache_seconds: float = 10.0) ->
bool:
return self.error is None
def fail_with(self, message: str) -> None:
@@ -77,7 +83,7 @@ def fake_trino() -> FakeTrinoClient:
@pytest.fixture()
def registry(settings: GatewaySettings, fake_trino: FakeTrinoClient):
- return build_registry(settings, trino_client=fake_trino) # type:
ignore[arg-type]
+ return build_registry(settings, lakehouse_client=fake_trino) # type:
ignore[arg-type]
class ScriptedChatModel(GenericFakeChatModel):
diff --git a/hudi-agent-gateway/tests/integration/test_spark_connect_live.py
b/hudi-agent-gateway/tests/integration/test_spark_connect_live.py
new file mode 100644
index 000000000000..7d961667ff41
--- /dev/null
+++ b/hudi-agent-gateway/tests/integration/test_spark_connect_live.py
@@ -0,0 +1,84 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Read-only tests on an existing table in the Helm Connect endpoint.
+
+Set GATEWAY_IT_SPARK_CONNECT_URI and GATEWAY_IT_SPARK_TABLE (schema.table).
+GATEWAY_IT_LLM=1 also enables the real Agent loop. No fixture writes here.
+"""
+
+import json
+import os
+
+import pytest
+from fastmcp import Client
+
+from hudi_agent_gateway.config import GatewaySettings
+from hudi_agent_gateway.mcp_server import build_mcp
+from hudi_agent_gateway.tools import build_client, build_registry
+
+URI = os.environ.get("GATEWAY_IT_SPARK_CONNECT_URI", "")
+TABLE = os.environ.get("GATEWAY_IT_SPARK_TABLE", "")
+pytestmark = pytest.mark.skipif(
+ not URI or not TABLE, reason="set Connect URI and existing test table"
+)
+
+
[email protected]()
+def connect_registry():
+ s = GatewaySettings(
+ sql_backend="spark-connect",
+ spark_connect_uri=URI,
+ )
+ client = build_client(s)
+ try:
+ yield s, build_registry(s, client)
+ finally:
+ client.close()
+
+
+async def test_tools_and_mcp_on_connect(connect_registry):
+ _, registry = connect_registry
+ schema, table = TABLE.split(".")
+ described = json.loads(await
registry.get("describe_table").handler(table=TABLE))
+ assert "error" not in described and described["row_count"] > 0
+ tables = json.loads(await
registry.get("list_tables").handler(schema_name=schema))
+ assert "error" not in tables and table in [row[1] for row in
tables["rows"]]
+ # describe_table validates identifiers before the SELECT is constructed.
+ sql = f"SELECT count(*) AS rows FROM `{schema}`.`{table}`"
+ out = json.loads(await registry.get("query_lakehouse").handler(sql=sql))
+ assert "error" not in out and out["row_count"] == 1
+ async with Client(build_mcp(registry)) as mcp:
+ result = await mcp.call_tool("query_lakehouse", {"sql": sql})
+ assert json.loads(result.content[0].text)["rows"] == out["rows"]
+ result = await mcp.call_tool("query_lakehouse", {"sql": f"DELETE FROM
{TABLE}"})
+ assert "read-only" in json.loads(result.content[0].text)["error"]
+
+
[email protected](os.environ.get("GATEWAY_IT_LLM") != "1", reason="set
GATEWAY_IT_LLM=1")
+async def test_agent_on_connect(connect_registry):
+ from langgraph.checkpoint.memory import InMemorySaver
+
+ from hudi_agent_gateway.agent import build_agent
+ from hudi_agent_gateway.llm import build_chat_model
+
+ s, registry = connect_registry
+ agent = build_agent(build_chat_model(s), registry, InMemorySaver(), s)
+ out = await agent.ainvoke(
+ {"messages": [("user", f"Describe the columns of {TABLE}, using your
tools.")]},
+ config={"configurable": {"thread_id": "connect-it"}},
+ )
+ assert any(getattr(m, "name", None) == "describe_table" for m in
out["messages"])
diff --git a/hudi-agent-gateway/tests/test_agent_loop.py
b/hudi-agent-gateway/tests/test_agent_loop.py
index 935497212d57..154beec3ed11 100644
--- a/hudi-agent-gateway/tests/test_agent_loop.py
+++ b/hudi-agent-gateway/tests/test_agent_loop.py
@@ -42,6 +42,20 @@ async def test_loop_executes_tool_and_answers(settings,
registry, fake_trino) ->
assert "100 trips" in messages[-1].text()
+async def test_spark_agent_loop_calls_guarded_tool(settings, fake_trino) ->
None:
+ from hudi_agent_gateway.tools import build_registry
+
+ spark = settings.model_copy(update={"sql_backend": "spark-connect"})
+ registry = build_registry(spark, fake_trino)
+ agent = build_agent(scripted_model(TOOL_CALL_THEN_ANSWER), registry,
InMemorySaver(), spark)
+ result = await agent.ainvoke(
+ {"messages": [HumanMessage("How many trips per city?")]},
+ config={"configurable": {"thread_id": "spark"}},
+ )
+ assert fake_trino.executed and "LIMIT 100" in fake_trino.executed[0]
+ assert any(isinstance(message, ToolMessage) for message in
result["messages"])
+
+
async def test_checkpointer_preserves_history_across_turns(settings, registry)
-> None:
saver = InMemorySaver()
first = scripted_model(
diff --git a/hudi-agent-gateway/tests/test_api_meta.py
b/hudi-agent-gateway/tests/test_api_meta.py
index 0aeba0c9dc26..487178a8af0c 100644
--- a/hudi-agent-gateway/tests/test_api_meta.py
+++ b/hudi-agent-gateway/tests/test_api_meta.py
@@ -29,7 +29,7 @@ async def client(settings, fake_trino, monkeypatch):
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=app), base_url="http://test";
) as c, app.router.lifespan_context(app):
- app.state.trino_client = fake_trino
+ app.state.lakehouse_client = fake_trino
yield c, app
@@ -42,7 +42,7 @@ async def test_health(client) -> None:
async def test_ready_reports_dependency_failures(client) -> None:
c, app = client
- app.state.trino_client.fail_with("connection refused")
+ app.state.lakehouse_client.fail_with("connection refused")
class NotReady:
async def check(self):
diff --git a/hudi-agent-gateway/tests/test_chat_sse.py
b/hudi-agent-gateway/tests/test_chat_sse.py
index 92f68bbdbcc4..e335b5f8b5b5 100644
--- a/hudi-agent-gateway/tests/test_chat_sse.py
+++ b/hudi-agent-gateway/tests/test_chat_sse.py
@@ -18,11 +18,12 @@
from __future__ import annotations
+import asyncio
import json
import httpx
import pytest
-from langchain_core.messages import AIMessage
+from langchain_core.messages import AIMessage, AIMessageChunk
from hudi_agent_gateway.agent import build_agent
from hudi_agent_gateway.app import create_app
@@ -132,3 +133,143 @@ async def test_validation_rejects_empty_message(client)
-> None:
c, _ = client
resp = await c.post("/v1/chat", json={"message": "", "session_id": "s"})
assert resp.status_code == 422
+
+
[email protected]("streaming", [False, True])
+async def test_whole_turn_timeout_cancels_active_model(client, streaming):
+ c, app = client
+ app.state.settings.chat_timeout_seconds = 0.05
+ cancelled = asyncio.Event()
+
+ class EndlessAgent:
+ async def ainvoke(self, *args, **kwargs):
+ try:
+ await asyncio.Event().wait()
+ finally:
+ cancelled.set()
+
+ async def astream(self, *args, **kwargs):
+ try:
+ while True:
+ # Active token delivery must not reset the total deadline.
+ yield "messages", (AIMessageChunk(content="x"), {})
+ await asyncio.sleep(0.005)
+ finally:
+ cancelled.set()
+
+ app.state.agents = FixedAgents(EndlessAgent())
+ async with asyncio.timeout(2):
+ response = await c.post("/v1/chat", json={
+ "message": "test", "session_id": "deadline", "stream": streaming,
+ })
+ assert cancelled.is_set()
+ if streaming:
+ events = _parse_sse(response.text)
+ assert events[0][0] == "status"
+ assert any(kind == "token" for kind, _ in events)
+ assert events[-1][1]["code"] == "chat_timeout"
+ assert "Start a new session" in events[-1][1]["message"]
+ assert all(kind != "done" for kind, _ in events)
+ else:
+ assert response.status_code == 504
+ assert "Start a new session" in response.json()["detail"]
+
+
+async def test_disconnected_sse_cancels_agent(client):
+ from starlette.requests import Request
+
+ c, app = client
+ started, cancelled = asyncio.Event(), asyncio.Event()
+
+ class WaitingAgent:
+ async def astream(self, *args, **kwargs):
+ started.set()
+ try:
+ await asyncio.Event().wait()
+ yield "updates", {}
+ finally:
+ cancelled.set()
+
+ from hudi_agent_gateway.api.chat import chat
+ from hudi_agent_gateway.api.models import ChatRequest
+
+ app.state.agents = FixedAgents(WaitingAgent())
+ request = Request({"type": "http", "method": "POST", "path": "/v1/chat",
+ "headers": [], "app": app})
+ response = await chat(request, ChatRequest(message="test",
session_id="cancel", stream=True))
+
+ async def receive():
+ await started.wait()
+ return {"type": "http.disconnect"}
+
+ async def send(message):
+ pass
+
+ async with asyncio.timeout(2):
+ await response(request.scope, receive, send)
+ assert cancelled.is_set()
+
+
+async def test_disconnected_sse_cancels_real_agent_tool(client, settings):
+ from starlette.requests import Request
+
+ from hudi_agent_gateway.api.chat import chat
+ from hudi_agent_gateway.api.models import ChatRequest
+ from hudi_agent_gateway.tools import build_registry
+
+ _, app = client
+ started, cancelled = asyncio.Event(), asyncio.Event()
+
+ class WaitingClient:
+ async def execute(self, *args, **kwargs):
+ started.set()
+ try:
+ await asyncio.Event().wait()
+ finally:
+ cancelled.set()
+
+ registry = build_registry(settings, lakehouse_client=WaitingClient())
+ app.state.agents = FixedAgents(build_agent(
+ scripted_model(TOOL_CALL_THEN_ANSWER), registry,
app.state.sessions.checkpointer, settings
+ ))
+ request = Request({"type": "http", "method": "POST", "path": "/v1/chat",
+ "headers": [], "app": app})
+ response = await chat(
+ request, ChatRequest(message="test", session_id="real-cancel",
stream=True)
+ )
+
+ async def receive():
+ await started.wait()
+ return {"type": "http.disconnect"}
+
+ async def send(message):
+ pass
+
+ async with asyncio.timeout(3):
+ await response(request.scope, receive, send)
+ await cancelled.wait()
+
+
+async def test_cancellation_scope_isolates_turns_and_rejects_late_tools():
+ from hudi_agent_gateway.cancellation import ToolCancellationScope
+
+ first, second = ToolCancellationScope(), ToolCancellationScope()
+ started = asyncio.Event()
+
+ async def wait():
+ started.set()
+ await asyncio.Event().wait()
+
+ first_task, second_task = asyncio.create_task(wait()),
asyncio.create_task(wait())
+ first.add(first_task)
+ second.add(second_task)
+ await started.wait()
+ first.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await first_task
+ assert not second_task.done()
+ with pytest.raises(asyncio.CancelledError):
+ first.add(asyncio.current_task())
+ second.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await second_task
diff --git a/hudi-agent-gateway/tests/test_guardrails.py
b/hudi-agent-gateway/tests/test_guardrails.py
index d7c610c511b1..ff4b04c1a1c3 100644
--- a/hudi-agent-gateway/tests/test_guardrails.py
+++ b/hudi-agent-gateway/tests/test_guardrails.py
@@ -19,6 +19,8 @@
from __future__ import annotations
import pytest
+import sqlglot
+from sqlglot import exp
from hudi_agent_gateway.tools.guardrails import enforce_guardrails
from hudi_agent_gateway.tools.registry import ToolInputError
@@ -41,6 +43,56 @@ def test_oversized_limit_is_lowered() -> None:
assert f"LIMIT {CAP}" in out and "100000" not in out
[email protected]("dialect", ["trino", "spark"])
[email protected]("count", [0, 1, CAP, CAP + 1])
+def test_limit_boundaries(dialect: str, count: int) -> None:
+ out = enforce_guardrails(f"SELECT a FROM t LIMIT {count}", CAP,
dialect=dialect)
+ stmt = sqlglot.parse_one(out, read=dialect)
+ assert int(stmt.args["limit"].expression.this) == min(count, CAP)
+
+
[email protected]("clause, expected", [
+ ("FETCH FIRST 10 ROWS ONLY", 10),
+ ("FETCH NEXT 1000 ROWS ONLY", CAP),
+ ("FETCH FIRST ROW ONLY", 1),
+ ("FETCH FIRST 0 ROWS ONLY", 0),
+])
+def test_fetch_is_normalized_and_capped(clause: str, expected: int) -> None:
+ out = enforce_guardrails(f"SELECT a FROM t ORDER BY a OFFSET 2 ROWS
{clause}", CAP)
+ stmt = sqlglot.parse_one(out, read="trino")
+ assert isinstance(stmt.args["limit"], exp.Limit)
+ assert int(stmt.args["limit"].expression.this) == expected
+ assert stmt.args["offset"].expression.this == "2"
+ assert stmt.args["order"] is not None
+
+
[email protected]("count, expected", [("10", 10), ("1000", 1000), ("",
1)])
+def test_fetch_with_ties_keeps_inner_selection_and_has_outer_cap(count: str,
expected: int) -> None:
+ sql = f"SELECT a FROM t ORDER BY a OFFSET 2 ROWS FETCH FIRST {count} ROWS
WITH TIES"
+ out = enforce_guardrails(sql, CAP)
+ stmt = sqlglot.parse_one(out, read="trino")
+ assert int(stmt.args["limit"].expression.this) == CAP
+ inner = stmt.find(exp.Subquery).this
+ assert isinstance(inner.args["limit"], exp.Fetch)
+ assert int(inner.args["limit"].args["count"].this) == expected
+ assert inner.args["limit"].args["limit_options"].args["with_ties"]
+ assert inner.args["offset"].expression.this == "2"
+ assert inner.args["order"] is not None
+
+
+def test_union_fetch_with_ties_caps_the_whole_result() -> None:
+ sql = "SELECT a FROM t UNION ALL SELECT a FROM u ORDER BY a FETCH FIRST 10
ROWS WITH TIES"
+ stmt = sqlglot.parse_one(enforce_guardrails(sql, CAP), read="trino")
+ assert int(stmt.args["limit"].expression.this) == CAP
+ assert isinstance(stmt.find(exp.Subquery).this, exp.Union)
+
+
[email protected]("dialect", ["trino", "spark"])
+def test_negative_limit_is_rejected(dialect: str) -> None:
+ with pytest.raises(ToolInputError, match="non-negative"):
+ enforce_guardrails("SELECT a FROM t LIMIT -1", CAP, dialect=dialect)
+
+
def test_cte_select_is_allowed() -> None:
out = enforce_guardrails("WITH x AS (SELECT a FROM t) SELECT * FROM x",
row_cap=CAP)
assert out.upper().startswith("WITH") and f"LIMIT {CAP}" in out
@@ -102,3 +154,165 @@ def test_string_literal_containing_dml_is_fine() -> None:
def test_quoted_identifiers_round_trip() -> None:
out = enforce_guardrails('SELECT "weird col" FROM "my table"', row_cap=CAP)
assert '"weird col"' in out
+
+
[email protected]("sql", [
+ "SELECT TRANSFORM(a) USING 'sh -c id' AS (x) FROM t",
+ "SELECT * FROM text.`file:///proc/self/environ`",
+ "SELECT * FROM json.`/var/run/secrets/token`",
+ "SELECT * FROM parquet.`s3a://other-bucket/path`",
+ "WITH x AS (SELECT * FROM text.`file:///etc/passwd`) SELECT * FROM x",
+ "SELECT * FROM (SELECT * FROM parquet.`s3a://bucket/path`) x",
+ "SELECT * FROM hudi_table_changes('s3a://other-bucket/tbl',
'latest_state', 'earliest')",
+ "SELECT * FROM hudi_query_timeline('s3a://other-bucket/tbl', 'false')",
+ "SELECT * FROM IDENTIFIER('text.`/etc/passwd`')",
+])
+def test_spark_script_and_path_reads_are_rejected(sql: str) -> None:
+ with pytest.raises(ToolInputError):
+ enforce_guardrails(sql, row_cap=CAP, dialect="spark")
+
+
[email protected]("sql", [
+ "SELECT * FROM `2024`",
+ "SELECT * FROM `schema with space`.`my-table`",
+ "SELECT * FROM text.orders",
+ "SELECT * FROM `json`.`orders`",
+ "SELECT * FROM other.text.orders",
+])
+def test_spark_catalog_identifiers_are_allowed(sql: str) -> None:
+ out = enforce_guardrails(sql, row_cap=CAP, dialect="spark",
spark_catalog="my_catalog")
+ assert "LIMIT 100" in out
+ if sql in ("SELECT * FROM text.orders", "SELECT * FROM `json`.`orders`",
+ "SELECT * FROM `schema with space`.`my-table`"):
+ assert "`my_catalog`." in out
+
+
+def test_spark_relative_datasource_path_cannot_use_file_fallback() -> None:
+ out = enforce_guardrails("SELECT * FROM text.`relative-file`", CAP,
dialect="spark")
+ assert "`spark_catalog`.text." in out
+
+
[email protected]("sql", [
+ "SELECT * FROM xml.`secrets`",
+ "SELECT * FROM libsvm.`secrets`",
+ "SELECT * FROM
`org.apache.spark.sql.execution.datasources.json`.`secrets`",
+ "SELECT * FROM custom_provider.`*`",
+ "WITH x AS (SELECT * FROM xml.`secrets`) SELECT * FROM x",
+ "SELECT * FROM (SELECT * FROM custom_provider.`relative`) x",
+])
+def test_spark_all_two_part_tables_are_catalog_qualified(sql: str) -> None:
+ out = enforce_guardrails(sql, CAP, dialect="spark",
spark_catalog="my_catalog")
+ tables = list(sqlglot.parse_one(out, read="spark").find_all(exp.Table))
+ assert any(table.catalog == "my_catalog" for table in tables)
+ assert all(not table.db or table.catalog for table in tables)
+
+
+def test_spark_qualified_tables_and_allowed_table_functions_are_preserved() ->
None:
+ out = enforce_guardrails(
+ "SELECT * FROM other.text.orders JOIN range(1) ON true", CAP,
dialect="spark"
+ )
+ assert "other.text.orders" in out
+ assert "RANGE(1)" in out
+
+
+_HUDI_FUNCTION_CALLS = [
+ "hudi_query({table}, 'snapshot')",
+ "hudi_table_changes({table}, 'latest_state', 'earliest')",
+ "hudi_query_timeline({table}, 'false')",
+ "hudi_filesystem_view({table})",
+ "hudi_metadata({table})",
+]
+
+
[email protected]("call", _HUDI_FUNCTION_CALLS)
[email protected]("table", [
+ "'s3a://other-bucket/tbl'", "'/etc/passwd'", "'../tbl'", "'file:tbl'",
+ r"'C:\\data\\tbl'", "'`db`.`/tmp/tbl`'", "'db.tbl extra'", "'range(1)'",
+ "concat('s3a:', '//other-bucket/tbl')", "'s3a:' || '//other-bucket/tbl'",
+ "concat(chr(47), 'tmp', chr(47), 'tbl')", "CAST('/tmp/tbl' AS STRING)",
+ "(SELECT '/tmp/tbl')", "NULL", "table_name", "''",
+])
+def test_spark_hudi_table_arguments_are_checked(call: str, table: str) -> None:
+ with pytest.raises(ToolInputError, match="literal catalog table
identifier"):
+ enforce_guardrails(f"SELECT * FROM {call.format(table=table)}", CAP,
dialect="spark")
+
+
[email protected]("call", _HUDI_FUNCTION_CALLS)
[email protected]("table", ["orders", "default.orders", "`my
db`.`my``table`"])
+def test_spark_hudi_catalog_tables_are_allowed(call: str, table: str) -> None:
+ sql = "SELECT * FROM " + call.format(table=f"'{table}'")
+ out = enforce_guardrails(sql, CAP, dialect="spark")
+ assert table in out and "LIMIT 100" in out
+
+
[email protected]("call", [
+ "HuDi_Query_Timeline('/tmp/tbl', 'false')",
+ "`hudi_metadata`('/tmp/tbl')",
+ "hudi_filesystem_view('default.orders', '../other-table')",
+ "hudi_filesystem_view('default.orders', concat('..', '/other-table'))",
+ "hudi_metadata()",
+ "IdEnTiFiEr(concat('text.', '`/etc/passwd`'))",
+ "`identifier`('text.`/etc/passwd`')",
+ "custom_reader(concat('s3a:', '//other-bucket/tbl'))",
+ "schema.hudi_metadata('/tmp/tbl')",
+ "schema.range(1)",
+])
[email protected]("query", [
+ "SELECT * FROM {call}",
+ "WITH x AS (SELECT * FROM {call}) SELECT * FROM x",
+ "SELECT * FROM (SELECT * FROM {call}) x",
+ "SELECT * FROM t JOIN {call} ON true",
+ "SELECT * FROM t, LATERAL {call}",
+ "SELECT * FROM t LATERAL VIEW {call} x AS y",
+])
+def test_spark_unsafe_table_functions_are_rejected_everywhere(call: str,
query: str) -> None:
+ with pytest.raises(ToolInputError):
+ enforce_guardrails(query.format(call=call), CAP, dialect="spark")
+
+
[email protected]("sql", [
+ "SELECT IDENTIFIER('col') FROM t",
+ "SELECT builtin.identifier('col') FROM t",
+ "SELECT * FROM t WHERE IDENTIFIER('col') = 1",
+])
+def test_spark_dynamic_identifiers_are_rejected_outside_from(sql: str) -> None:
+ with pytest.raises(ToolInputError, match="Spark function identifier"):
+ enforce_guardrails(sql, CAP, dialect="spark")
+
+
[email protected]("sql", [
+ "SELECT '/etc/passwd', 's3a://bucket/tbl', 'IDENTIFIER(x)', `identifier`
FROM t",
+ "SELECT * FROM range(10)",
+ "SELECT * FROM explode(array('a/b', 'file:/tmp'))",
+ "SELECT * FROM explode_outer(array(1))",
+ "SELECT * FROM posexplode(array(1))",
+ "SELECT * FROM posexplode_outer(array(1))",
+ "SELECT * FROM inline(array(struct(1)))",
+ "SELECT * FROM inline_outer(array(struct(1)))",
+ "SELECT * FROM stack(1, 2)",
+ "SELECT * FROM json_tuple('{\"path\":\"a/b\"}', 'path')",
+ "SELECT * FROM t LATERAL VIEW explode(array('a/b')) x AS y",
+ "SELECT * FROM t, LATERAL hudi_query_timeline('default.orders', 'false')",
+ "SELECT * FROM t, LATERAL (SELECT * FROM range(1)) x",
+])
+def test_spark_safe_functions_and_path_literals_are_allowed(sql: str) -> None:
+ assert "LIMIT 100" in enforce_guardrails(sql, CAP, dialect="spark")
+
+
+def test_trino_table_functions_are_unchanged() -> None:
+ out = enforce_guardrails("SELECT * FROM
schema.custom_reader('s3a://bucket/tbl')", CAP)
+ assert "CUSTOM_READER" in out and "LIMIT 100" in out
+
+
[email protected]("sql, name, hint", [
+ ("SELECT * FROM custom_reader('x')", "custom_reader", "unqualified table
function"),
+ ("SELECT * FROM schema.range(1)", "range", "unqualified table function"),
+ ("SELECT * FROM hudi_metadata(concat('db.', 't'))", "Hudi",
"hudi_metadata"),
+ ("SELECT * FROM hudi_filesystem_view('db.t', '../x')", "subpaths",
"without a subpath"),
+])
+def test_spark_function_errors_explain_how_to_correct_the_query(
+ sql: str, name: str, hint: str,
+) -> None:
+ with pytest.raises(ToolInputError, match=name) as error:
+ enforce_guardrails(sql, CAP, dialect="spark")
+ assert hint in error.value.hint
diff --git a/hudi-agent-gateway/tests/test_mcp.py
b/hudi-agent-gateway/tests/test_mcp.py
index f1455acc8b83..b7f81c609bb1 100644
--- a/hudi-agent-gateway/tests/test_mcp.py
+++ b/hudi-agent-gateway/tests/test_mcp.py
@@ -46,7 +46,7 @@ async def test_mounted_mcp_serves_over_http(settings,
fake_trino) -> None:
async with httpx.AsyncClient(
transport=httpx.ASGITransport(app=app), base_url="http://test";
) as c, app.router.lifespan_context(app):
- app.state.trino_client = fake_trino
+ app.state.lakehouse_client = fake_trino
resp = await c.post(
"/mcp/",
json={
diff --git a/hudi-agent-gateway/tests/test_spark_connect.py
b/hudi-agent-gateway/tests/test_spark_connect.py
new file mode 100644
index 000000000000..c9df3aa1d52e
--- /dev/null
+++ b/hudi-agent-gateway/tests/test_spark_connect.py
@@ -0,0 +1,498 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""PySpark Connect lifecycle, bounded results, guardrails, and shared
surfaces."""
+
+import asyncio
+import gc
+import json
+import threading
+import weakref
+from contextlib import contextmanager
+from unittest.mock import Mock, patch
+
+import httpx
+import pytest
+from fastmcp import Client
+from pydantic import ValidationError
+from pyspark.errors import PySparkException
+from pyspark.sql.types import IntegerType, StringType, StructField, StructType
+
+from hudi_agent_gateway.agent import build_system_prompt
+from hudi_agent_gateway.app import create_app
+from hudi_agent_gateway.config import GatewaySettings
+from hudi_agent_gateway.mcp_server import build_mcp
+from hudi_agent_gateway.tools import build_registry
+from hudi_agent_gateway.tools.guardrails import enforce_guardrails
+from hudi_agent_gateway.tools.lakehouse_client import QueryError, QueryResult,
QueryTimeoutError
+from hudi_agent_gateway.tools.registry import ToolInputError
+from hudi_agent_gateway.tools.spark_connect_client import SparkConnectClient,
_RequestSparkSession
+
+
[email protected]("closed", [True, False])
+def test_session_gc_does_not_close_an_already_stopped_client(closed):
+ session = _RequestSparkSession.__new__(_RequestSparkSession)
+ client = Mock(is_closed=closed)
+ session._client = client
+ # Spark's cached Catalog retains its session, so finalization happens
during
+ # cyclic GC, potentially in an unrelated thread while locks are held.
+ session.catalog = Mock(_sparkSession=session)
+ ref = weakref.ref(session)
+ del session
+ gc.collect()
+ assert ref() is None
+ assert client.close.call_count == (0 if closed else 1)
+
+
+class DataFrame:
+ def __init__(self, session):
+ self.session = session
+ self.columns = ["id"]
+ self.schema = StructType(
+ [StructField("id", IntegerType(), False), StructField("region",
StringType())]
+ )
+ self.limits = []
+ self.collected = False
+
+ def limit(self, rows):
+ self.limits.append(rows)
+ return self
+
+ def collect(self):
+ self.collected = True
+ self.session.started.set()
+ if self.session.release is not None:
+ self.session.release.wait(3)
+ if self.session.error:
+ raise self.session.error
+ return [(i,) for i in range(self.limits[-1])] if self.columns else []
+
+
+class Session:
+ def __init__(self):
+ self.catalog = self
+ self.catalog_name = None
+ self.schema_name = None
+ self.catalog_error = None
+ self.executed = []
+ self.closed = threading.Event()
+ self.cancelled = threading.Event()
+ self.started = threading.Event()
+ self.release = None
+ self.error = None
+ self.cancel_error = None
+ self.df = DataFrame(self)
+
+ def setCurrentCatalog(self, name):
+ self.catalog_name = name
+ if self.catalog_error:
+ raise self.catalog_error
+
+ def setCurrentDatabase(self, name):
+ self.schema_name = name
+
+ def sql(self, sql):
+ self.executed.append(sql)
+ return self.df
+
+ def interruptAll(self):
+ self.cancelled.set()
+ if self.cancel_error:
+ raise self.cancel_error
+ return ["operation"]
+
+ def stop(self):
+ self.closed.set()
+
+
+@contextmanager
+def pyspark_client(settings, *sessions):
+ client = SparkConnectClient(settings)
+ with patch(
+ "hudi_agent_gateway.tools.spark_connect_client._RequestSparkSession",
side_effect=sessions
+ ) as connect:
+ try:
+ yield client, connect
+ finally:
+ for session in sessions:
+ if session.release is not None:
+ session.release.set()
+ client.close()
+
+
+def spark_settings(settings):
+ return GatewaySettings(
+ **{
+ **settings.model_dump(),
+ "sql_backend": "spark-connect",
+ "spark_connect_uri": "sc://localhost:15002",
+ }
+ )
+
+
+async def test_pyspark_collect_is_bounded_and_session_closes(settings):
+ s = spark_settings(settings)
+ session = Session()
+ with pyspark_client(s, session) as (client, connect):
+ out = await client.execute("SELECT id FROM t LIMIT 3", timeout=1,
max_rows=3)
+ assert out.rows == [[0], [1], [2]]
+ assert session.df.limits == [3]
+ assert session.closed.is_set()
+ assert connect.call_args.kwargs == {"connection": s.spark_connect_uri}
+ assert session.catalog_name == "spark_catalog" and session.schema_name
== "default"
+
+
[email protected]("failure", ["query", "catalog"])
+async def test_pyspark_failures_close_session(settings, failure):
+ session = Session()
+ if failure == "query":
+ session.error = RuntimeError("gRPC call failed")
+ else:
+ session.catalog_error = PySparkException("Catalog not found")
+ with pyspark_client(spark_settings(settings), session) as (client, _):
+ with pytest.raises(QueryError):
+ await client.execute("SELECT 1", timeout=1, max_rows=1)
+ assert session.closed.is_set()
+
+
+async def test_pyspark_query_error_is_actionable_and_redacted(settings):
+ s = spark_settings(settings).model_copy(
+ update={"spark_connect_uri":
"sc://localhost:15002/;token=secret%20token"}
+ )
+ session = Session()
+ session.error = PySparkException(
+ f"TABLE_NOT_FOUND at {s.spark_connect_uri}: secret token"
+ )
+ with pyspark_client(s, session) as (client, _):
+ with pytest.raises(QueryError, match="TABLE_NOT_FOUND") as caught:
+ await client.execute("SELECT * FROM absent", timeout=1, max_rows=1)
+ assert "secret" not in str(caught.value)
+
+
+async def test_pyspark_describe_uses_schema_without_collecting(settings):
+ session = Session()
+ with pyspark_client(spark_settings(settings), session) as (client, _):
+ out = await client.describe("SELECT * FROM t LIMIT 0", timeout=1,
max_rows=2)
+ assert out.columns == ["column_name", "data_type", "is_nullable"]
+ assert out.rows == [["id", "int", "NO"], ["region", "string", "YES"]]
+ assert not session.df.collected
+ assert session.closed.is_set()
+
+
[email protected]("cancel_task", [False, True])
+async def
test_timeout_and_cancellation_cleanup_without_blocking_loop(settings,
cancel_task):
+ session = Session()
+ session.release = threading.Event()
+ with pyspark_client(spark_settings(settings), session) as (client, _):
+ task = asyncio.create_task(client.execute("SELECT 1", timeout=0.05,
max_rows=1))
+ assert await asyncio.to_thread(session.started.wait, 1)
+ if cancel_task:
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError if cancel_task else
QueryTimeoutError):
+ await task
+ assert await asyncio.to_thread(session.cancelled.wait, 1)
+ assert not session.closed.is_set() # Worker still owns the session.
+ session.release.set()
+ assert await asyncio.to_thread(session.closed.wait, 1)
+
+
+async def
test_cancellation_is_isolated_and_failure_does_not_prevent_cleanup(settings):
+ slow, fast = Session(), Session()
+ slow.release = threading.Event()
+ slow.cancel_error = RuntimeError("interrupt RPC failed")
+ with pyspark_client(spark_settings(settings), slow, fast) as (client,
connect):
+ task = asyncio.create_task(client.execute("SELECT 1", timeout=0.1,
max_rows=1))
+ assert await asyncio.to_thread(slow.started.wait, 1)
+ assert (await client.execute("SELECT 2", timeout=1, max_rows=1)).rows
== [[0]]
+ with pytest.raises(QueryTimeoutError):
+ await task
+ assert await asyncio.to_thread(slow.cancelled.wait, 1)
+ assert not fast.cancelled.is_set() and fast.closed.is_set()
+ assert connect.call_count == 2
+ slow.release.set()
+ assert await asyncio.to_thread(slow.closed.wait, 1)
+
+
+async def test_timed_out_workers_keep_capacity_slots(settings):
+ sessions = [Session() for _ in range(8)]
+ for session in sessions:
+ session.release = threading.Event()
+ with pyspark_client(spark_settings(settings), *sessions) as (client, _):
+ jobs = [client.execute("SELECT 1", timeout=0.1, max_rows=1) for _ in
range(8)]
+ assert all(
+ isinstance(e, QueryTimeoutError)
+ for e in await asyncio.gather(*jobs, return_exceptions=True)
+ )
+ with pytest.raises(QueryError, match="capacity exhausted"):
+ await client.execute("SELECT 1", timeout=1, max_rows=1)
+ for session in sessions:
+ session.release.set()
+ assert await asyncio.to_thread(session.closed.wait, 1)
+
+
+async def test_ping_cached_and_closed_client_rejected(settings):
+ session = Session()
+ with pyspark_client(spark_settings(settings), session) as (client, _):
+ assert await client.ping()
+ assert await client.ping()
+ assert session.executed == ["SELECT 1"]
+ client.close()
+ with pytest.raises(QueryError, match="closed"):
+ await client.execute("SELECT 1", timeout=1, max_rows=1)
+
+
+async def test_ping_works_when_query_capacity_is_exhausted(settings):
+ sessions = [Session() for _ in range(9)]
+ for session in sessions[:8]:
+ session.release = threading.Event()
+ with pyspark_client(spark_settings(settings), *sessions) as (client, _):
+ tasks = [asyncio.create_task(client.execute("SELECT 1", timeout=5,
max_rows=1))
+ for _ in range(8)]
+ for session in sessions[:8]:
+ assert await asyncio.to_thread(session.started.wait, 1)
+ assert await client.ping(timeout=1)
+ for session in sessions[:8]:
+ session.release.set()
+ await asyncio.gather(*tasks)
+
+
+async def
test_concurrent_pings_share_probe_even_if_a_waiter_is_cancelled(settings):
+ session = Session()
+ session.release = threading.Event()
+ with pyspark_client(spark_settings(settings), session) as (client,
connect):
+ first = asyncio.create_task(client.ping())
+ assert await asyncio.to_thread(session.started.wait, 1)
+ second = asyncio.create_task(client.ping())
+ await asyncio.sleep(0)
+ first.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await first
+ assert not second.done()
+ session.release.set()
+ assert await second
+ assert await client.ping()
+ assert connect.call_count == 1
+
+
+async def test_close_interrupts_active_operations(settings):
+ session = Session()
+ session.release = threading.Event()
+ with pyspark_client(spark_settings(settings), session) as (client, _):
+ task = asyncio.create_task(client.execute("SELECT 1", timeout=5,
max_rows=1))
+ assert await asyncio.to_thread(session.started.wait, 1)
+ client.close()
+ assert await asyncio.to_thread(session.cancelled.wait, 1)
+ session.release.set()
+ await task
+
+
[email protected](
+ "sql",
+ [
+ "INSERT INTO t VALUES (1)",
+ "UPDATE t SET a=1",
+ "DELETE FROM t",
+ "DROP TABLE t",
+ "WITH x AS (INSERT INTO t VALUES (1)) SELECT * FROM x",
+ "SELECT 1; DELETE FROM t",
+ "INSERT OVERWRITE DIRECTORY '/tmp/x' SELECT 1",
+ "SELECT * INTO t FROM u",
+ "CACHE TABLE t",
+ "SET spark.sql.shuffle.partitions=1",
+ ],
+)
+def test_spark_writes_rejected(sql):
+ with pytest.raises(ToolInputError):
+ enforce_guardrails(sql, 100, dialect="spark")
+
+
[email protected]("function", ["java_method", "ReFlEcT", "TRY_REFLECT"])
[email protected](
+ "query",
+ [
+ "SELECT {call}",
+ "WITH x AS (SELECT {call} AS value) SELECT value FROM x",
+ "SELECT value FROM (SELECT {call} AS value) x",
+ "SELECT builtin.{call}",
+ "SELECT 1 UNION ALL SELECT {call}",
+ ],
+)
+async def test_spark_reflection_rejected_before_client(settings, fake_trino,
function, query):
+ call = f"{function}('java.lang.System', 'setProperty', 'k', 'v')"
+ sql = query.format(call=call)
+ with pytest.raises(ToolInputError, match="read-only: Spark function"):
+ enforce_guardrails(sql, 100, dialect="spark")
+ registry = build_registry(spark_settings(settings), fake_trino)
+ result = json.loads(await registry.get("query_lakehouse").handler(sql=sql))
+ assert "read-only: Spark function" in result["error"]
+ async with Client(build_mcp(registry)) as mcp:
+ result = await mcp.call_tool("query_lakehouse", {"sql": sql})
+ assert "read-only: Spark function" in
json.loads(result.content[0].text)["error"]
+ assert not fake_trino.executed
+
+
+def test_spark_sql_quoting_and_functions():
+ out = enforce_guardrails(
+ "SELECT `weird col`, date_trunc('DAY', ts) FROM `t`", 100,
dialect="spark"
+ )
+ assert "`weird col`" in out and "LIMIT 100" in out
+
+
+async def
test_query_qualifies_format_named_schema_with_configured_catalog(settings,
fake_trino):
+ s = spark_settings(settings).model_copy(update={"spark_connect_catalog":
"custom_catalog"})
+ registry = build_registry(s, fake_trino)
+ await registry.get("query_lakehouse").handler(sql="SELECT * FROM
text.orders")
+ assert "`custom_catalog`.text.orders" in fake_trino.executed[-1]
+
+
[email protected]("sql", [
+ "SELECT * FROM hudi_table_changes('s3a://other-bucket/tbl',
'latest_state', 'earliest')",
+ "SELECT * FROM hudi_query_timeline(concat('s3a:', '//other-bucket/tbl'),
'false')",
+ "SELECT * FROM hudi_filesystem_view('default.orders', '../other-table')",
+ "SELECT * FROM IDENTIFIER('text.`/etc/passwd`')",
+])
+async def test_spark_table_function_bypasses_rejected_before_client(settings,
fake_trino, sql):
+ registry = build_registry(spark_settings(settings), fake_trino)
+ result = json.loads(await registry.get("query_lakehouse").handler(sql=sql))
+ assert "read-only:" in result["error"]
+ async with Client(build_mcp(registry)) as mcp:
+ result = await mcp.call_tool("query_lakehouse", {"sql": sql})
+ assert "read-only:" in json.loads(result.content[0].text)["error"]
+ assert not fake_trino.executed
+
+
[email protected](
+ "sql",
+ [
+ "SELECT 'java_method', 'reflect', 'try_reflect' AS function_names",
+ "SELECT `java_method`, `reflect`, `try_reflect` FROM t",
+ "SELECT 1 /* java_method('java.lang.System', 'setProperty', 'k', 'v')
*/",
+ ],
+)
+def test_spark_reflection_names_without_calls_are_allowed(sql):
+ assert "LIMIT 100" in enforce_guardrails(sql, 100, dialect="spark")
+
+
+async def test_spark_tools_metadata_mcp_and_readonly(settings, fake_trino):
+ s = spark_settings(settings)
+ registry = build_registry(s, fake_trino)
+ fake_trino.result = QueryResult(
+ ["namespace", "tableName", "isTemporary"], [["default", "orders",
False]]
+ )
+ out = json.loads(await registry.get("list_tables").handler())
+ assert out["columns"] == ["table_schema", "table_name", "table_type"]
+ assert out["rows"] == [["default", "orders", "TABLE"]]
+ assert fake_trino.executed[-1] == "SHOW TABLES IN
`spark_catalog`.`default`"
+ fake_trino.result = QueryResult(
+ ["column_name", "data_type", "is_nullable"], [["id", "int32", "YES"]]
+ )
+ await registry.get("describe_table").handler(table="orders")
+ assert fake_trino.executed[-1] == "SELECT * FROM
`spark_catalog`.`default`.`orders` LIMIT 0"
+ async with Client(build_mcp(registry)) as mcp:
+ assert {t.name for t in await mcp.list_tools()} == {
+ "list_tables",
+ "describe_table",
+ "query_lakehouse",
+ }
+ before = len(fake_trino.executed)
+ result = await mcp.call_tool("query_lakehouse", {"sql": "DELETE FROM
orders"})
+ assert "read-only" in json.loads(result.content[0].text)["error"]
+ assert len(fake_trino.executed) == before
+ await mcp.call_tool("query_lakehouse", {"sql": "SELECT `id` FROM
orders"})
+ assert "LIMIT 100" in fake_trino.executed[-1]
+ assert "Spark Connect" in build_system_prompt(s)
+ assert "spark_catalog" in build_system_prompt(s)
+
+
[email protected]("port, expected_port", [("", 15002), (":1", 1),
(":65535", 65535)])
+async def test_spark_ready_info_and_shutdown(
+ settings, fake_trino, monkeypatch, port, expected_port
+):
+ s = GatewaySettings(
+ **{
+ **spark_settings(settings).model_dump(),
+ "spark_connect_uri": f"sc://localhost{port}",
+ }
+ )
+ monkeypatch.setattr("hudi_agent_gateway.app.build_client", lambda _:
fake_trino)
+ app = create_app(s.model_copy(update={"mcp_enabled": False}))
+ async with (
+ app.router.lifespan_context(app),
+ httpx.AsyncClient(
+ transport=httpx.ASGITransport(app=app),
+ base_url="http://test";,
+ ) as http,
+ ):
+ info = (await http.get("/v1/info")).json()
+ assert info["sql_backend"] == "spark-connect"
+ assert info["catalog"] == "spark_catalog" and not info["trino_url"]
+ assert info["lakehouse_url"] == f"sc://localhost:{expected_port}"
+ fake_trino.fail_with("unreachable")
+ ready = await http.get("/ready")
+ assert ready.status_code == 503
+ assert "spark-connect" in ready.json()["checks"]
+ assert "trino" not in ready.json()["checks"]
+
+
[email protected](
+ "uri, endpoint",
+ [
+ ("sc://localhost", "sc://localhost:15002"),
+ ("sc://localhost/;use_ssl=true;token=secret%20token",
"sc://localhost:15002"),
+ ("sc://[::1]:15002/;token=secret", "sc://[::1]:15002"),
+ ],
+)
+async def test_connect_uri_parameters_are_accepted_and_info_is_redacted(
+ settings, fake_trino, monkeypatch, uri, endpoint
+):
+ from pyspark.sql.connect.client import DefaultChannelBuilder
+
+ s = GatewaySettings(**{**spark_settings(settings).model_dump(),
"spark_connect_uri": uri})
+ assert DefaultChannelBuilder(uri).endpoint ==
endpoint.removeprefix("sc://")
+ monkeypatch.setattr("hudi_agent_gateway.app.build_client", lambda _:
fake_trino)
+ app = create_app(s.model_copy(update={"mcp_enabled": False}))
+ async with (
+ app.router.lifespan_context(app),
+ httpx.AsyncClient(transport=httpx.ASGITransport(app=app),
base_url="http://test";) as http,
+ ):
+ response = await http.get("/v1/info")
+ assert response.json()["lakehouse_url"] == endpoint
+ assert "secret" not in response.text
+
+
[email protected](
+ "uri",
+ [
+ "spark://localhost:15002?api=connect",
+ "http://localhost:15002";,
+ "sc://localhost?api=connect",
+ "sc://localhost:not-a-port",
+ "sc://localhost:65536",
+ "sc://localhost:-1",
+ "sc://localhost:0",
+ "sc://[::1]:not-a-port",
+ "sc://[::1]:65536",
+ "sc:///",
+ "sc://user:secret@localhost:15002",
+ "sc://localhost/database",
+ "sc://localhost/;token",
+ "sc://localhost/;token=a;",
+ "sc://localhost#fragment",
+ ],
+)
+def test_bad_connect_uri_fails_fast(uri):
+ with pytest.raises(ValidationError, match="GATEWAY_SPARK_CONNECT_URI"):
+ GatewaySettings(sql_backend="spark-connect", spark_connect_uri=uri)
diff --git a/hudi-lakehouse/README.md b/hudi-lakehouse/README.md
index 29bcca8f3b6c..d362eec1f2ca 100644
--- a/hudi-lakehouse/README.md
+++ b/hudi-lakehouse/README.md
@@ -83,7 +83,7 @@ catalog, and query engines are independent. This directory
keeps the
|---|---|---|
| `hudi-trino` | Trino (server 472) with the Hudi connector built from this
repository | your Hive Metastore or AWS Glue; your S3/GCS |
| `hudi-spark-connect` | optional Spark Connect 4.1.3 with Hudi 1.2.0; one
shared local-mode driver | your Hive Metastore and S3-compatible storage |
- | `hudi-agent-gateway` | the [Hudi AI gateway](../hudi-agent-gateway):
agent chat API (sessions + SSE), MCP server, and chat UI over guarded lakehouse
tools | your `hudi-trino`; your LLM (see below) |
+ | `hudi-agent-gateway` | the [Hudi AI gateway](../hudi-agent-gateway):
agent chat API (sessions + SSE), MCP server, and chat UI over guarded lakehouse
tools | your Trino or Spark Connect endpoint; your LLM (see below) |
| `vllm` | optional: vLLM serving one open-weight model (default
`Qwen/Qwen3-8B`, ungated) behind an OpenAI-compatible API; GPU required | —
(the gateway points at it) |
**LLM out-of-box experience**: the gateway defaults to **Ollama**
@@ -170,7 +170,12 @@ TLS; use it on a trusted network. See the [chart
installation guide](charts/hudi
for storage, credentials, and resource settings, and the [ADBC client
guide](charts/hudi-spark-connect/README.md#connect-and-validate-with-adbc)
for connection setup and CRUD/restart validation.
-## BI clients on Spark Connect
+## Gateway and BI clients on Spark Connect
+
+The gateway selects the official PySpark Connect client with
+`sqlBackend: spark-connect`; its `list_tables`, `describe_table`, and
+`query_lakehouse` tools remain read-only. See the
+[gateway setup](charts/hudi-spark-connect/README.md#agent-gateway) for the
Helm overlay.
See the [BI client
notes](charts/hudi-spark-connect/README.md#spark-connect-client-tutorials)
for Superset and Zeppelin connection requirements and ADBC limitations.
diff --git a/hudi-lakehouse/charts/hudi-agent-gateway/templates/NOTES.txt
b/hudi-lakehouse/charts/hudi-agent-gateway/templates/NOTES.txt
index 1670610e5ec4..01578a317bd5 100644
--- a/hudi-lakehouse/charts/hudi-agent-gateway/templates/NOTES.txt
+++ b/hudi-lakehouse/charts/hudi-agent-gateway/templates/NOTES.txt
@@ -18,7 +18,12 @@ limitations under the License.
hudi-agent-gateway {{ .Chart.Version }} deployed.
provider : {{ .Values.gateway.provider }} ({{ .Values.gateway.model }})
+ backend : {{ .Values.sqlBackend }}
+{{- if eq .Values.sqlBackend "trino" }}
trino : {{ .Values.trino.host }}:{{ .Values.trino.port }} (catalog {{
.Values.trino.catalog }})
+{{- else }}
+ catalog : {{ .Values.sparkConnect.catalog }} / {{
.Values.sparkConnect.schema }}
+{{- end }}
mcp : {{ if .Values.mcpEnabled }}enabled at /mcp{{ else }}disabled{{
end }}
Try it:
@@ -42,5 +47,5 @@ loopback, restart it reachable from the cluster:
{{- end }}
Readiness (kubectl -n {{ .Release.Namespace }} get pods) stays NotReady until
-both Trino and the LLM endpoint are reachable -- `curl localhost:8000/ready`
+both {{ .Values.sqlBackend }} and the LLM endpoint are reachable -- `curl
localhost:8000/ready`
shows per-dependency status.
diff --git a/hudi-lakehouse/charts/hudi-agent-gateway/templates/deployment.yaml
b/hudi-lakehouse/charts/hudi-agent-gateway/templates/deployment.yaml
index 2258d2f5466e..1e12acf4e901 100644
--- a/hudi-lakehouse/charts/hudi-agent-gateway/templates/deployment.yaml
+++ b/hudi-lakehouse/charts/hudi-agent-gateway/templates/deployment.yaml
@@ -14,6 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
{{- $g := .Values.gateway }}
+{{- if not (has .Values.sqlBackend (list "trino" "spark-connect")) }}
+{{- fail "sqlBackend must be trino or spark-connect" }}
+{{- end }}
apiVersion: apps/v1
kind: Deployment
metadata:
@@ -57,6 +60,8 @@ spec:
value: {{ $g.model | quote }}
- name: GATEWAY_LLM_TIMEOUT_SECONDS
value: {{ $g.llmTimeoutSeconds | quote }}
+ - name: GATEWAY_CHAT_TIMEOUT_SECONDS
+ value: {{ $g.chatTimeoutSeconds | quote }}
{{- with $g.ollama.baseUrl }}
- name: GATEWAY_OLLAMA_BASE_URL
value: {{ . | quote }}
@@ -69,6 +74,23 @@ spec:
- name: GATEWAY_SYSTEM_PROMPT_EXTRA
value: {{ . | quote }}
{{- end }}
+ - name: GATEWAY_SQL_BACKEND
+ value: {{ .Values.sqlBackend | quote }}
+ {{- if eq .Values.sqlBackend "spark-connect" }}
+ - name: GATEWAY_SPARK_CONNECT_URI
+ {{- if .Values.sparkConnect.existingSecret }}
+ valueFrom:
+ secretKeyRef:
+ name: {{ .Values.sparkConnect.existingSecret | quote }}
+ key: {{ .Values.sparkConnect.uriSecretKey | quote }}
+ {{- else }}
+ value: {{ .Values.sparkConnect.uri | quote }}
+ {{- end }}
+ - name: GATEWAY_SPARK_CONNECT_CATALOG
+ value: {{ .Values.sparkConnect.catalog | quote }}
+ - name: GATEWAY_SPARK_CONNECT_SCHEMA
+ value: {{ .Values.sparkConnect.schema | quote }}
+ {{- end }}
- name: GATEWAY_TRINO_HOST
value: {{ .Values.trino.host | quote }}
- name: GATEWAY_TRINO_PORT
@@ -109,7 +131,7 @@ spec:
path: /ready
port: http
periodSeconds: 10
- # /ready's own dependency checks budget up to ~8s (5s trino ping +
+ # /ready's own dependency checks budget up to ~8s (5s lakehouse
ping +
# 3s LLM probe) on a cache miss; the k8s default of 1s would flap.
timeoutSeconds: 9
failureThreshold: 3
diff --git a/hudi-lakehouse/charts/hudi-agent-gateway/values-local-dev.yaml
b/hudi-lakehouse/charts/hudi-agent-gateway/values-local-dev.yaml
index 143a24e4403c..a24e8c7eccf2 100644
--- a/hudi-lakehouse/charts/hudi-agent-gateway/values-local-dev.yaml
+++ b/hudi-lakehouse/charts/hudi-agent-gateway/values-local-dev.yaml
@@ -24,7 +24,6 @@ fullnameOverride: hudi-agent-gateway
gateway:
provider: ollama
# A small tool-capable model that fits a laptop; `ollama pull qwen3:4b`.
- # (1.7b also works but makes more reasoning mistakes on SQL tasks.)
model: qwen3:4b
ollama:
baseUrl: http://host.minikube.internal:11434
diff --git a/hudi-agent-gateway/src/hudi_agent_gateway/tools/__init__.py
b/hudi-lakehouse/charts/hudi-agent-gateway/values-spark-connect.yaml
similarity index 52%
copy from hudi-agent-gateway/src/hudi_agent_gateway/tools/__init__.py
copy to hudi-lakehouse/charts/hudi-agent-gateway/values-spark-connect.yaml
index 9e4db82cdc7b..ad42c0a59c38 100644
--- a/hudi-agent-gateway/src/hudi_agent_gateway/tools/__init__.py
+++ b/hudi-lakehouse/charts/hudi-agent-gateway/values-spark-connect.yaml
@@ -14,24 +14,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""Tool assembly: the single place where tool modules get registered.
-
-Adding a new tool module = write the module with a ``register(registry, ...)``
-function, then call it here.
-"""
-
-from __future__ import annotations
-
-from hudi_agent_gateway.config import GatewaySettings
-from hudi_agent_gateway.tools import trino_tools
-from hudi_agent_gateway.tools.registry import ToolRegistry
-from hudi_agent_gateway.tools.trino_client import TrinoClient
-
-
-def build_registry(
- settings: GatewaySettings, trino_client: TrinoClient | None = None
-) -> ToolRegistry:
- registry = ToolRegistry()
- client = trino_client or TrinoClient(settings)
- trino_tools.register(registry, client, settings)
- return registry
+# Overlay after values-local-dev.yaml to use the shared Connect endpoint.
+sqlBackend: spark-connect
+sparkConnect:
+ uri: "sc://hudi-spark-connect.hudi-lakehouse.svc:15002"
+ catalog: spark_catalog
+ schema: default
diff --git a/hudi-lakehouse/charts/hudi-agent-gateway/values.yaml
b/hudi-lakehouse/charts/hudi-agent-gateway/values.yaml
index e8880ba9ab60..d4d2f3a73699 100644
--- a/hudi-lakehouse/charts/hudi-agent-gateway/values.yaml
+++ b/hudi-lakehouse/charts/hudi-agent-gateway/values.yaml
@@ -54,6 +54,19 @@ gateway:
# Endpoint for provider=openai-compatible (vLLM, Together, Fireworks, ...).
baseUrl: ""
llmTimeoutSeconds: 120
+ # Total wall-clock limit for a chat turn, including all model/tool calls.
+ chatTimeoutSeconds: 300
+
+# trino | spark-connect
+sqlBackend: trino
+
+sparkConnect:
+ uri: "sc://hudi-spark-connect.hudi-lakehouse.svc:15002"
+ catalog: spark_catalog
+ schema: default
+ # Optional Secret key containing GATEWAY_SPARK_CONNECT_URI (e.g. auth
tokens).
+ existingSecret: ""
+ uriSecretKey: GATEWAY_SPARK_CONNECT_URI
trino:
host: hudi-trino.hudi-lakehouse.svc
diff --git a/hudi-lakehouse/charts/hudi-spark-connect/README.md
b/hudi-lakehouse/charts/hudi-spark-connect/README.md
index ba949104c08e..fd7c7d45c207 100644
--- a/hudi-lakehouse/charts/hudi-spark-connect/README.md
+++ b/hudi-lakehouse/charts/hudi-spark-connect/README.md
@@ -284,11 +284,37 @@ An existing session did not see an externally added
column until
recreation passed. Further end-to-end validation is tracked in
[#20233](https://github.com/apache/hudi/issues/20233).
+## Agent gateway
+
+See the [gateway setup and
limits](../../../hudi-agent-gateway/README.md#spark-connect-backend).
+The Agent and MCP interfaces share the same three read-only tools:
+`list_tables`, `describe_table`, and `query_lakehouse`.
+
+Build the gateway image, which includes the official PySpark Connect client:
+
+```bash
+docker build -t hudi-lakehouse-agent-gateway:0.1.0 hudi-agent-gateway
+```
+
+Load or push that image for your cluster, then install the gateway from the
+repository root:
+
+```bash
+helm upgrade --install hudi-agent-gateway
hudi-lakehouse/charts/hudi-agent-gateway \
+ -n hudi-lakehouse \
+ -f hudi-lakehouse/charts/hudi-agent-gateway/values-local-dev.yaml \
+ -f hudi-lakehouse/charts/hudi-agent-gateway/values-spark-connect.yaml
+```
+
+This retains the local-dev Ollama configuration. `/ready` checks Spark with
+`SELECT 1` and checks the configured LLM separately. MCP tool calls themselves
+do not require an LLM. Override the provider as described in the gateway
README.
+
<a id="spark-connect-client-tutorials"></a>
## BI clients
-Gateway integration and BI adapters/examples for Superset and other tools are
tracked
+BI adapters/examples for Superset and other tools are tracked
in [#20232](https://github.com/apache/hudi/issues/20232) and will be delivered
in follow-up PRs.
The initial SQLAlchemy adapter is planned under `hudi-lakehouse`; upstreaming
can
be evaluated separately after the integration stabilizes.
diff --git a/hudi-lakehouse/charts/hudi-spark-connect/templates/configmap.yaml
b/hudi-lakehouse/charts/hudi-spark-connect/templates/configmap.yaml
index 5fa146214ca8..11298a929992 100644
--- a/hudi-lakehouse/charts/hudi-spark-connect/templates/configmap.yaml
+++ b/hudi-lakehouse/charts/hudi-spark-connect/templates/configmap.yaml
@@ -22,6 +22,7 @@
"spark.connect.grpc.binding.address" "0.0.0.0"
"spark.connect.grpc.binding.port" "15002"
"spark.sql.catalogImplementation" "hive"
+ "spark.sql.runSQLOnFiles" "false"
"spark.hadoop.hive.metastore.uris" $uri
"spark.sql.warehouse.dir" $warehouse
"spark.sql.extensions"
"org.apache.spark.sql.hudi.HoodieSparkSessionExtension"
diff --git a/hudi-lakehouse/tests/test_gateway_chart.py
b/hudi-lakehouse/tests/test_gateway_chart.py
new file mode 100644
index 000000000000..5d9a77c59cf2
--- /dev/null
+++ b/hudi-lakehouse/tests/test_gateway_chart.py
@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Gateway backend and Secret chart contracts."""
+
+import pathlib
+import subprocess
+import unittest
+
+import yaml
+
+CHART = pathlib.Path(__file__).resolve().parents[1] /
"charts/hudi-agent-gateway"
+
+
+def render(*args):
+ result = subprocess.run(["helm", "template", "test", str(CHART), *args],
+ capture_output=True, text=True, check=True)
+ deployment = next(d for d in yaml.safe_load_all(result.stdout) if d and
d["kind"] == "Deployment")
+ container = deployment["spec"]["template"]["spec"]["containers"][0]
+ return {e["name"]: e for e in container["env"]}
+
+
+class GatewayChartTest(unittest.TestCase):
+ def test_default_trino(self):
+ env = render()
+ self.assertEqual(env["GATEWAY_SQL_BACKEND"]["value"], "trino")
+ self.assertNotIn("GATEWAY_SPARK_CONNECT_URI", env)
+
+ def test_connect_overlay(self):
+ env = render("-f", str(CHART / "values-spark-connect.yaml"))
+ self.assertEqual(env["GATEWAY_SQL_BACKEND"]["value"], "spark-connect")
+ self.assertEqual(env["GATEWAY_SPARK_CONNECT_URI"]["value"],
+ "sc://hudi-spark-connect.hudi-lakehouse.svc:15002")
+ self.assertEqual(env["GATEWAY_SPARK_CONNECT_CATALOG"]["value"],
"spark_catalog")
+ self.assertNotIn("GATEWAY_SPARK_CONNECT_DRIVER", env)
+
+ def test_connect_uri_secret(self):
+ env = render("--set", "sqlBackend=spark-connect", "--set",
"sparkConnect.existingSecret=connect-auth",
+ "--set", "sparkConnect.uriSecretKey=uri")
+ self.assertNotIn("value", env["GATEWAY_SPARK_CONNECT_URI"])
+
self.assertEqual(env["GATEWAY_SPARK_CONNECT_URI"]["valueFrom"]["secretKeyRef"],
+ {"name": "connect-auth", "key": "uri"})
+
+ def test_invalid_backend_rejected(self):
+ with self.assertRaises(subprocess.CalledProcessError):
+ render("--set", "sqlBackend=typo")
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/hudi-lakehouse/tests/test_spark_connect_chart.py
b/hudi-lakehouse/tests/test_spark_connect_chart.py
index 862586968422..b309aeb8206f 100644
--- a/hudi-lakehouse/tests/test_spark_connect_chart.py
+++ b/hudi-lakehouse/tests/test_spark_connect_chart.py
@@ -54,6 +54,7 @@ class ConnectChartTest(unittest.TestCase):
self.assertIn("spark.connect.grpc.arrow.maxBatchSize 1m", config)
self.assertIn("spark.hadoop.hive.metastore.uris
thrift://hive-metastore.", config)
self.assertIn("spark.sql.warehouse.dir s3a://warehouse/connect",
config)
+ self.assertIn("spark.sql.runSQLOnFiles false", config)
self.assertIn("org.apache.spark.sql.hudi.catalog.HoodieCatalog",
config)
self.assertIn("HoodieSparkSessionExtension", config)
self.assertEqual(docs["Deployment"]["spec"]["replicas"], 1)
