nsivabalan commented on code in PR #20329: URL: https://github.com/apache/hudi/pull/20329#discussion_r4385843815
########## hudi-agent-gateway/src/hudi_agent_gateway/tools/spark_connect_client.py: ########## @@ -0,0 +1,184 @@ +# 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._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_cache: tuple[float, bool] | None = None + self._closed = False + + def close(self) -> None: + self._closed = True + self._executor.shutdown(wait=False, cancel_futures=True) + self._cancel_executor.shutdown(wait=False, cancel_futures=True) + + 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 + ) -> QueryResult: + if self._closed: + raise QueryError("Spark Connect client is closed") + if not self._slots.acquire(blocking=False): + raise QueryError( + "Spark Connect query capacity exhausted; retry after active queries finish" + ) + op = _Operation() + try: + work = self._executor.submit(self._execute_sync, sql, max_rows, op, describe) + except Exception: + self._slots.release() + raise + work.add_done_callback(lambda _: self._slots.release()) + 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: Review Comment: ⚠️ **IMPORTANT.** `ping()` goes through `execute()`, so it competes for the 8 slots. Slots are deliberately retained past the timeout, so 8 slow queries make `/ready` fail. After 3 probes the pod leaves the Service, even for chats that don't need Spark. Could ping use a reserved slot, or treat 'capacity exhausted' as reachable? ########## hudi-agent-gateway/src/hudi_agent_gateway/tools/guardrails.py: ########## @@ -42,42 +43,47 @@ exp.Command, # also covers CALL and EXPLAIN, which parse as generic commands exp.Use, ) +_SPARK_REFLECTION_FUNCTIONS = frozenset({"java_method", "reflect", "try_reflect"}) -_HINT = "Provide exactly one read-only SELECT statement (Trino SQL)." - -def enforce_guardrails(sql: str, row_cap: int) -> str: +def enforce_guardrails(sql: str, row_cap: int, *, dialect: str = "trino") -> 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 # (catches e.g. a CTE wrapping an INSERT). 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" and isinstance(node, exp.Func): Review Comment: 🚨 **BLOCKING.** the reflection check is good, but in the Spark dialect `SELECT TRANSFORM(a) USING 'sh -c id' AS (x) FROM t` and `SELECT * FROM text.\`file:///proc/self/environ\`` both pass `enforce_guardrails` (verified with sqlglot). The Connect chart uses the Hive catalog, so TRANSFORM spawns a process on the driver, and path reads expose pod files, secrets, and arbitrary buckets. Please reject `exp.QueryTransform` and path- or format-qualified tables (e.g. `text.`/`parquet.`/`json.` or URI-like names), and set `spark.sql.runSQLOnFiles=false` in the #20243 chart as defense in depth. ########## hudi-agent-gateway/pyproject.toml: ########## @@ -35,6 +35,8 @@ dependencies = [ "langchain-ollama>=1.0", "fastmcp>=2.10,<3", "trino>=0.330", + "pyspark-client==4.1.3", Review Comment: ⚠️ **IMPORTANT.** Since `spark_connect_client` is lazily imported, could `pyspark-client`/`pyarrow` be an optional extra (`[spark-connect]`) so Trino-only installs don't take exact pins? The exact `pyspark-client` pin is needed because `_RequestSparkSession.__del__` relies on the private `client.is_closed`; worth a comment saying so. ########## hudi-agent-gateway/src/hudi_agent_gateway/config.py: ########## @@ -44,7 +43,7 @@ class GatewaySettings(BaseSettings): # --- LLM --- llm_provider: LLMProvider = "ollama" - llm_model: str = "qwen3:8b" + llm_model: str = "qwen3:4b-instruct" Review Comment: ⚠️ **IMPORTANT.** This default model change looks unrelated to the Spark Connect backend. Can it go in a separate PR, or be called out in the description? ########## hudi-agent-gateway/src/hudi_agent_gateway/api/chat.py: ########## @@ -97,10 +100,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): Review Comment: ⚠️ **IMPORTANT.** If this deadline fires between the agent node checkpointing an AIMessage with tool_calls and the ToolMessage being written, the thread is left with a dangling tool call, and the next turn on this session_id gets a provider 400. The UI rotates the session, but API and MCP callers won't. Can we repair the checkpoint (synthetic error ToolMessages) or at least tell the caller to start a new session in the 504 detail? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
