Copilot commented on code in PR #19265: URL: https://github.com/apache/hudi/pull/19265#discussion_r3567392604
########## hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_tools.py: ########## @@ -0,0 +1,174 @@ +# 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. + +"""Lakehouse tools backed by Trino. + +Handlers return JSON strings. Expected failures (guardrail rejections, query +errors, timeouts) are returned as ``{"error": ..., "hint": ...}`` payloads +rather than raised, so the agent can read the error and self-correct, and MCP +clients get a useful result instead of a protocol error. +""" + +from __future__ import annotations + +import json +from typing import Annotated, Any + +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 ( + QueryResult, + TrinoClient, + TrinoQueryError, + TrinoTimeoutError, +) + +_QUERY_DESC = ( + "Run a single read-only SELECT statement (Trino SQL) against the lakehouse " + "and return the result as JSON. A server-side row cap is enforced; results " + "may be truncated (indicated by `truncated: true`)." +) +_LIST_TABLES_DESC = ( + "List tables in the lakehouse. Optionally filter by catalog and schema; " + "defaults to the gateway's configured catalog and schema." +) +_DESCRIBE_DESC = ( + "Describe a table's columns and types. Accepts `table`, `schema.table`, or " + "`catalog.schema.table`." +) + + +def shape_result(result: QueryResult, *, max_bytes: int, sql: str) -> str: + payload: dict[str, Any] = { + "sql": sql, + "columns": result.columns, + "rows": result.rows, + "row_count": result.row_count, + "truncated": False, + } + text = json.dumps(payload, default=str) + while len(text.encode()) > max_bytes and payload["rows"]: + keep = max(1, len(payload["rows"]) // 2) + if keep == len(payload["rows"]): + keep -= 1 + payload["rows"] = payload["rows"][:keep] + payload["truncated"] = True + payload["notice"] = ( + f"Result truncated to {len(payload['rows'])} of {result.row_count} fetched rows " + "to fit the size limit. Narrow the query (filters, aggregation, fewer columns)." + ) + text = json.dumps(payload, default=str) + return text + + +def _error(message: str, hint: str = "") -> str: + payload = {"error": message} + if hint: + payload["hint"] = hint + return json.dumps(payload) + + +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] + if len(parts) == 2: + return settings.trino_catalog, parts[0], parts[1] + if len(parts) == 3: + return parts[0], parts[1], parts[2] + raise ToolInputError( + f"invalid table name {table!r}", hint="Use table, schema.table, or catalog.schema.table." + ) + + +def register(registry: ToolRegistry, client: TrinoClient, settings: GatewaySettings) -> None: + @registry.register("query_lakehouse", _QUERY_DESC) + async def query_lakehouse( + sql: Annotated[ + str, Field(description="A single read-only SELECT statement in Trino SQL.") + ], + ) -> str: + try: + safe_sql = enforce_guardrails(sql, row_cap=settings.sql_row_cap) + 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: + 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: + return _error(f"query failed: {e}", "Fix the SQL and try again.") + return shape_result(result, max_bytes=settings.tool_result_max_bytes, sql=safe_sql) + + @registry.register("list_tables", _LIST_TABLES_DESC) + async def list_tables( + catalog: Annotated[ + str, Field(description="Catalog to list from; empty for the default.") + ] = "", + schema_name: Annotated[ + str, Field(description="Schema to list from; empty for the default.") + ] = "", + ) -> str: + cat = catalog or settings.trino_catalog + sch = schema_name or settings.trino_schema + sql = ( + f'SELECT table_schema, table_name, table_type FROM "{cat}".information_schema.tables ' + f"WHERE table_schema = '{sch}' ORDER BY table_name" + ) Review Comment: `list_tables` interpolates `catalog` and `schema_name` directly into SQL without escaping/validation. Since these tool params are exposed via HTTP/MCP, a malicious value containing quotes could break out of the identifier/string literal and inject SQL against Trino’s information_schema. Escape/quote these values (at minimum: double quotes in catalog identifiers, and single quotes in string literals) or validate against a strict identifier regex before building SQL. ########## hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_tools.py: ########## @@ -0,0 +1,174 @@ +# 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. + +"""Lakehouse tools backed by Trino. + +Handlers return JSON strings. Expected failures (guardrail rejections, query +errors, timeouts) are returned as ``{"error": ..., "hint": ...}`` payloads +rather than raised, so the agent can read the error and self-correct, and MCP +clients get a useful result instead of a protocol error. +""" + +from __future__ import annotations + +import json +from typing import Annotated, Any + +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 ( + QueryResult, + TrinoClient, + TrinoQueryError, + TrinoTimeoutError, +) + +_QUERY_DESC = ( + "Run a single read-only SELECT statement (Trino SQL) against the lakehouse " + "and return the result as JSON. A server-side row cap is enforced; results " + "may be truncated (indicated by `truncated: true`)." +) +_LIST_TABLES_DESC = ( + "List tables in the lakehouse. Optionally filter by catalog and schema; " + "defaults to the gateway's configured catalog and schema." +) +_DESCRIBE_DESC = ( + "Describe a table's columns and types. Accepts `table`, `schema.table`, or " + "`catalog.schema.table`." +) + + +def shape_result(result: QueryResult, *, max_bytes: int, sql: str) -> str: + payload: dict[str, Any] = { + "sql": sql, + "columns": result.columns, + "rows": result.rows, + "row_count": result.row_count, + "truncated": False, + } + text = json.dumps(payload, default=str) + while len(text.encode()) > max_bytes and payload["rows"]: + keep = max(1, len(payload["rows"]) // 2) + if keep == len(payload["rows"]): + keep -= 1 + payload["rows"] = payload["rows"][:keep] + payload["truncated"] = True + payload["notice"] = ( + f"Result truncated to {len(payload['rows'])} of {result.row_count} fetched rows " + "to fit the size limit. Narrow the query (filters, aggregation, fewer columns)." + ) + text = json.dumps(payload, default=str) + return text + + +def _error(message: str, hint: str = "") -> str: + payload = {"error": message} + if hint: + payload["hint"] = hint + return json.dumps(payload) + + +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] + if len(parts) == 2: + return settings.trino_catalog, parts[0], parts[1] + if len(parts) == 3: + return parts[0], parts[1], parts[2] + raise ToolInputError( + f"invalid table name {table!r}", hint="Use table, schema.table, or catalog.schema.table." + ) + + +def register(registry: ToolRegistry, client: TrinoClient, settings: GatewaySettings) -> None: + @registry.register("query_lakehouse", _QUERY_DESC) + async def query_lakehouse( + sql: Annotated[ + str, Field(description="A single read-only SELECT statement in Trino SQL.") + ], + ) -> str: + try: + safe_sql = enforce_guardrails(sql, row_cap=settings.sql_row_cap) + 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: + 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: + return _error(f"query failed: {e}", "Fix the SQL and try again.") + return shape_result(result, max_bytes=settings.tool_result_max_bytes, sql=safe_sql) + + @registry.register("list_tables", _LIST_TABLES_DESC) + async def list_tables( + catalog: Annotated[ + str, Field(description="Catalog to list from; empty for the default.") + ] = "", + schema_name: Annotated[ + str, Field(description="Schema to list from; empty for the default.") + ] = "", + ) -> str: + cat = catalog or settings.trino_catalog + sch = schema_name or settings.trino_schema + sql = ( + f'SELECT table_schema, table_name, table_type 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: + return _error(f"listing tables failed: {e}") + return shape_result(result, max_bytes=settings.tool_result_max_bytes, sql=sql) + + @registry.register("describe_table", _DESCRIBE_DESC) + async def describe_table( + table: Annotated[ + str, Field(description="Table name: table, schema.table, or catalog.schema.table.") + ], + ) -> str: + try: + cat, sch, tbl = _split_table_name(table, settings) + 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" + ) Review Comment: `describe_table` builds a SQL string with unescaped `catalog/schema/table` values. Because `table` comes from tool input (HTTP/MCP), quotes in any segment can produce SQL injection against Trino’s information_schema. Please escape/quote the identifier and string-literal segments (or validate them as identifiers) before interpolation. ########## hudi-agent-gateway/README.md: ########## @@ -0,0 +1,129 @@ +<!-- +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. +--> +# hudi-agent-gateway + +**One deployable service that serves your Hudi lakehouse to AI.** A single +process hosts three surfaces over the same set of guarded lakehouse tools: + +| Surface | Where | What | +|---|---|---| +| Agent chat API | `POST /v1/chat` | prompt in → LangGraph agent loop (model ↔ tools) → grounded answer out; multi-turn sessions; optional SSE streaming | +| 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, +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. + +## Quickstart (local process) + +```bash +cd hudi-agent-gateway +python3.12 -m venv .venv && .venv/bin/pip install -e ".[dev]" + +# point at a Trino with the Hudi connector (e.g. the local-dev stack, port-forwarded): +GATEWAY_TRINO_HOST=localhost GATEWAY_TRINO_PORT=18080 \ +GATEWAY_LLM_PROVIDER=anthropic GATEWAY_LLM_MODEL=claude-haiku-4-5-20251001 \ + .venv/bin/hudi-agent-gateway serve + +curl -X POST localhost:8000/v1/chat -H 'Content-Type: application/json' \ + -d '{"message": "How many trips per city?", "session_id": "s1"}' +open http://localhost:8000/ui/ +``` + +`GET /v1/models` lists the models the configured provider offers (live: +Anthropic and OpenAI model APIs, Ollama's local tags, vLLM's served models), +and `POST /v1/chat` accepts an optional `"model"` to pick one per request — +the chat UI exposes this as a model picker. Sessions survive model switches. + +For a fully local model, install [Ollama](https://ollama.com), pull a +tool-capable model, and use the default provider: + +```bash +ollama pull qwen3:8b +GATEWAY_LLM_PROVIDER=ollama GATEWAY_LLM_MODEL=qwen3:8b hudi-agent-gateway serve +``` + +Connect an MCP client: + +```bash +claude mcp add --transport http hudi-lakehouse http://localhost:8000/mcp/ +``` + +## Deploying on Kubernetes + +See `hudi-lakehouse/charts/hudi-agent-gateway` — the product Helm chart — +and `hudi-lakehouse/local-dev/` for a complete laptop environment +(MinIO + Hive Metastore + Trino + this gateway) where the gateway is +installed alongside Trino by default. + +## Configuration + +Environment variables (prefix `GATEWAY_` except the standard key names): + +| Variable | Default | Purpose | +|---|---|---| +| `GATEWAY_LLM_PROVIDER` | `ollama` | `anthropic` \| `openai` \| `ollama` \| `openai-compatible` | +| `GATEWAY_LLM_MODEL` | `qwen3:8b` | model name for the provider | +| `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_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 | +| `GATEWAY_SQL_TIMEOUT_SECONDS` | `60` | per-query timeout | +| `GATEWAY_TOOL_RESULT_MAX_BYTES` | `50000` | tool results truncated beyond this (with notice) | +| `GATEWAY_AGENT_MAX_ITERATIONS` | `25` | agent loop recursion limit | Review Comment: README docs `GATEWAY_SQL_TIMEOUT_SECONDS` default as `60`, but the actual default in `GatewaySettings.sql_timeout_seconds` is `120.0`. This mismatch will mislead users troubleshooting readiness/timeouts. Please update the README default (or change the code default) so they match. ########## hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_tools.py: ########## @@ -0,0 +1,174 @@ +# 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. + +"""Lakehouse tools backed by Trino. + +Handlers return JSON strings. Expected failures (guardrail rejections, query +errors, timeouts) are returned as ``{"error": ..., "hint": ...}`` payloads +rather than raised, so the agent can read the error and self-correct, and MCP +clients get a useful result instead of a protocol error. +""" + +from __future__ import annotations + +import json +from typing import Annotated, Any + +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 ( + QueryResult, + TrinoClient, + TrinoQueryError, + TrinoTimeoutError, +) + +_QUERY_DESC = ( + "Run a single read-only SELECT statement (Trino SQL) against the lakehouse " + "and return the result as JSON. A server-side row cap is enforced; results " + "may be truncated (indicated by `truncated: true`)." +) +_LIST_TABLES_DESC = ( + "List tables in the lakehouse. Optionally filter by catalog and schema; " + "defaults to the gateway's configured catalog and schema." +) +_DESCRIBE_DESC = ( + "Describe a table's columns and types. Accepts `table`, `schema.table`, or " + "`catalog.schema.table`." +) + + +def shape_result(result: QueryResult, *, max_bytes: int, sql: str) -> str: + payload: dict[str, Any] = { + "sql": sql, + "columns": result.columns, + "rows": result.rows, + "row_count": result.row_count, + "truncated": False, + } + text = json.dumps(payload, default=str) + while len(text.encode()) > max_bytes and payload["rows"]: + keep = max(1, len(payload["rows"]) // 2) + if keep == len(payload["rows"]): + keep -= 1 + payload["rows"] = payload["rows"][:keep] + payload["truncated"] = True + payload["notice"] = ( + f"Result truncated to {len(payload['rows'])} of {result.row_count} fetched rows " + "to fit the size limit. Narrow the query (filters, aggregation, fewer columns)." + ) + text = json.dumps(payload, default=str) + return text Review Comment: `shape_result` stops truncating once `rows` becomes empty, but the JSON can still exceed `max_bytes` (e.g., very long `sql`, many columns, or a very small configured limit). That would violate the tool’s size contract and could break SSE/HTTP clients. Consider adding a final fail-safe: if the payload is still too large after dropping rows, return a minimal payload (or an error) that is guaranteed to fit. -- 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]
