vinothchandar commented on code in PR #19265: URL: https://github.com/apache/hudi/pull/19265#discussion_r3567425120
########## 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: Fixed in 8adefc01 — catalog/schema names are now validated against a strict identifier pattern before any SQL is built; invalid values return a tool error with a hint and never reach Trino. Covered by a test asserting the client receives nothing on quoted-breakout attempts. ########## 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: Fixed in 8adefc01 — all three segments of the split table name go through the same identifier validation as list_tables before interpolation. ########## 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: Fixed in 8adefc01 — shape_result now degrades in two bounded steps once rows are exhausted: a metadata-only payload (truncated sql, column count instead of names), then a bare error payload, so the size contract holds even for tiny caps. Tested at 400- and 80-byte limits. ########## 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: Fixed in 8adefc01 — README now documents the actual 120s default. -- 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]
