vinothchandar commented on code in PR #19265: URL: https://github.com/apache/hudi/pull/19265#discussion_r3634641470
########## hudi-agent-gateway/src/hudi_agent_gateway/tools/trino_client.py: ########## @@ -0,0 +1,96 @@ +# 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. + +"""Async wrapper over the (synchronous) Trino DB-API client.""" + +from __future__ import annotations + +import asyncio +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) + + +class TrinoClient: + """Thin async facade: each query runs the sync client in a worker thread + under a hard timeout, and fetches at most ``max_rows`` from the cursor + regardless of what the SQL says (defense in depth behind the guardrails). + """ + + def __init__(self, settings: GatewaySettings) -> None: + self._settings = settings + + def _connect(self) -> trino.dbapi.Connection: + s = self._settings + return trino.dbapi.connect( + host=s.trino_host, + port=s.trino_port, + user=s.trino_user, + catalog=s.trino_catalog, + schema=s.trino_schema, + http_scheme="http", + ) + + def _execute_sync(self, sql: str, max_rows: int) -> QueryResult: + with self._connect() as conn: + cursor = conn.cursor() + try: + cursor.execute(sql) + rows = cursor.fetchmany(max_rows) + columns = [d[0] for d in cursor.description or []] + return QueryResult(columns=columns, rows=[list(r) for r in rows]) + finally: + cursor.close() + + async def execute(self, sql: str, *, timeout: float, max_rows: int) -> QueryResult: + try: + return await asyncio.wait_for( Review Comment: folded it in rather than filing an issue — ba979e3 runs queries on a dedicated bounded pool and the timeout path calls `cursor.cancel()` so Trino kills the query server-side. ########## 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: fair, at 2 lines it's cheaper to fix than file. ba979e3: filling the row cap now sets `truncated` with a notice. ########## hudi-lakehouse/local-dev/example/spark-app.yaml: ########## @@ -0,0 +1,61 @@ +# 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. + +# SparkApplication for the Apache Spark Kubernetes Operator +# (https://github.com/apache/spark-kubernetes-operator), submitted by +# local-dev/scripts/run-example.sh after uploading hudi_table_writer.py to +# s3a://warehouse/jobs/. +# +# To run your own job: copy this file + your script, adjust pyFiles (or +# mainClass/jars for JVM jobs), the env block, and any spark conf. +apiVersion: spark.apache.org/v1alpha1 +kind: SparkApplication +metadata: + name: hudi-table-writer + namespace: hudi-lakehouse +spec: + pyFiles: "s3a://warehouse/jobs/hudi_table_writer.py" + runtimeVersions: + sparkVersion: "3.5.7" + scalaVersion: "2.12" + sparkConf: + # Service account created by the spark-kubernetes-operator helm chart + # (the namespace default SA cannot manage executor pods). + spark.kubernetes.authenticate.driver.serviceAccountName: "spark" + spark.kubernetes.container.image: "hudi-lakehouse-spark:3.5" Review Comment: dropped the flag everywhere in ba979e3 — the quickstart builds exactly the 3.5 line the example manifest pins. ########## hudi-lakehouse/scripts/build-jars.sh: ########## @@ -0,0 +1,62 @@ +#!/usr/bin/env bash +# 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. + +# Builds the two Maven artifacts the hudi-lakehouse images need: +# 1. the Hudi Spark bundle (packaging/hudi-spark-bundle) -- JDK 11/17 +# 2. the Hudi Trino plugin (hudi-trino-plugin) -- JDK 23+ +# +# Usage: build-jars.sh [--spark-version 3.5|4.1] [--skip-bundle] [--skip-plugin] + +set -euo pipefail +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +SPARK_VERSION="3.5" +BUILD_BUNDLE=1 +BUILD_PLUGIN=1 + +while [[ $# -gt 0 ]]; do + case "$1" in + --spark-version) SPARK_VERSION="$2"; shift 2 ;; + --skip-bundle) BUILD_BUNDLE=0; shift ;; + --skip-plugin) BUILD_PLUGIN=0; shift ;; + *) echo "unknown arg: $1" >&2; exit 1 ;; + esac +done + +MVN_ARGS=(-DskipTests -Dcheckstyle.skip=true -Dscalastyle.skip=true -Drat.skip=true -Dmaven.javadoc.skip=true -Dgpg.skip=true) + +if [[ "$BUILD_BUNDLE" == 1 ]]; then + # Profiles must be activated explicitly: activating any -D profile property + # disables Maven's activeByDefault profiles, so relying on defaults is fragile. + echo ">>> Building Hudi Spark bundle (spark${SPARK_VERSION})" + (cd "$REPO_ROOT" && mvn clean package -pl packaging/hudi-spark-bundle -am \ + "-Dspark${SPARK_VERSION}" -Dflink2.1 "${MVN_ARGS[@]}") +fi + +if [[ "$BUILD_PLUGIN" == 1 ]]; then + JAVA_MAJOR=$(java -version 2>&1 | sed -n 's/.*version "\([0-9]*\).*/\1/p' | head -1) Review Comment: fixed in ba979e3 — the guard probes `$JAVA_HOME`'s java, same as maven. ########## scripts/release/validate_source_copyright.sh: ########## @@ -50,7 +50,7 @@ echo "Performing custom Licensing Check " # Exclude the 'hudi-trino-plugin' directory. Its license checks are handled by airlift: # https://github.com/airlift/airbase/blob/823101482dbc60600d7862f0f5c93aded6190996/airbase/pom.xml#L1239 # --- -numfilesWithNoLicense=$(find . -path './hudi-trino-plugin' -prune -o -type f -iname '*' | grep -v './hudi-trino-plugin' | grep -v NOTICE | grep -v LICENSE | grep -v '.jpg' | grep -v '.json' | grep -v '.zip' | grep -v '.hfile' | grep -v '.data' | grep -v '.commit' | grep -v emptyFile | grep -v DISCLAIMER | grep -v '.sqltemplate' | grep -v KEYS | grep -v '.mailmap' | grep -v 'banner.txt' | grep -v '.txt' | grep -v "fixtures" | xargs grep -L "Licensed to the Apache Software Foundation (ASF)") +numfilesWithNoLicense=$(find . -path './hudi-trino-plugin' -prune -o -type f -iname '*' | grep -v './hudi-trino-plugin' | grep -v NOTICE | grep -v LICENSE | grep -v '.jpg' | grep -v '.png' | grep -v '.json' | grep -v '.zip' | grep -v '.hfile' | grep -v '.data' | grep -v '.commit' | grep -v emptyFile | grep -v DISCLAIMER | grep -v '.sqltemplate' | grep -v KEYS | grep -v '.mailmap' | grep -v 'banner.txt' | grep -v '.txt' | grep -v "fixtures" | xargs grep -L "Licensed to the Apache Software Foundation (ASF)") Review Comment: went a step simpler than the data-URI: the UI now hotlinks the logo from hudi.apache.org like the root README does (degrades to alt text offline). binary gone, exclusion reverted — that script is untouched by this PR now. PR description corrected too. -- 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]
