vinothchandar commented on code in PR #19265: URL: https://github.com/apache/hudi/pull/19265#discussion_r3568121268
########## hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_tools.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. + +"""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 +import re +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, Review Comment: It's not in scope for the initial PR file. These are issues to file if you think it needs to be fixed. -- 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]
