Copilot commented on code in PR #7400: URL: https://github.com/apache/kyuubi/pull/7400#discussion_r3108086956
########## externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/tool/sql/SqlExecutor.java: ########## @@ -0,0 +1,137 @@ +/* + * 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. + */ + +package org.apache.kyuubi.engine.dataagent.tool.sql; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.Statement; +import javax.sql.DataSource; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Shared JDBC execution + result formatting for SQL tools. Owned by {@link RunSelectQueryTool} and + * {@link RunMutationQueryTool}; not part of the public tool API. + */ +final class SqlExecutor { + + private static final Logger LOG = LoggerFactory.getLogger(SqlExecutor.class); + + private SqlExecutor() {} + + /** + * Execute the given SQL and return a markdown-formatted result string. Errors are caught and + * returned as a one-line "Error: ..." message so the LLM can self-correct. + * + * <p><b>No client-side row cap.</b> The tool deliberately does not call {@link + * Statement#setMaxRows(int)}. Result-size discipline is delegated to the LLM via the system + * prompt, which requires every read query to include an explicit {@code LIMIT} (or equivalent + * pagination). Adding a hidden client-side cap here would silently truncate the LLM's view of the + * data and create magic numbers that are impossible to tune across deployments. + * + * <p><b>Two-layer timeout model.</b> The {@code timeoutSeconds} arg here is the JDBC inner + * timeout, set via {@link Statement#setQueryTimeout(int)} and sourced from {@code + * kyuubi.engine.data.agent.query.timeout}. This is the cooperative path: the driver tells the + * server (Spark/Trino/...) to cancel the query, and cluster resources are released gracefully. + * The outer hard wall-clock cap on the whole tool call is enforced separately by the agent + * runtime via {@code kyuubi.engine.data.agent.tool.call.timeout} — that one fires even if the + * JDBC driver ignores {@code setQueryTimeout} (some legacy Hive/Spark Thrift drivers do), and is + * the ungraceful "kill the thread" backstop. + */ + static String execute(DataSource dataSource, String sql, int timeoutSeconds) { + if (sql == null || sql.trim().isEmpty()) { + return "Error: 'sql' parameter is required."; + } + + try (Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement()) { + stmt.setQueryTimeout(timeoutSeconds); + boolean hasResultSet = stmt.execute(sql); + if (hasResultSet) { + try (ResultSet rs = stmt.getResultSet()) { + return formatResult(rs); + } + } else { + return "[Statement executed successfully. " + stmt.getUpdateCount() + " row(s) affected]"; + } + } catch (Exception e) { + LOG.warn("SQL execution error", e); + return "Error: SQL execution failed. " + extractRootCause(e); + } + } + + private static String formatResult(ResultSet rs) throws Exception { + ResultSetMetaData meta = rs.getMetaData(); + int colCount = meta.getColumnCount(); + StringBuilder sb = new StringBuilder(); + + sb.append("| "); + for (int i = 1; i <= colCount; i++) { + if (i > 1) sb.append(" | "); + sb.append(meta.getColumnName(i)); Review Comment: Result headers use `ResultSetMetaData.getColumnName(i)`, which can ignore SQL aliases and produce confusing column names for the agent (e.g., `SELECT a AS b`). Prefer `getColumnLabel(i)` so the rendered markdown table reflects the query output labels. ```suggestion sb.append(meta.getColumnLabel(i)); ``` ########## externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/tool/sql/SqlReadOnlyChecker.java: ########## @@ -0,0 +1,137 @@ +/* + * 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. + */ + +package org.apache.kyuubi.engine.dataagent.tool.sql; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.Locale; +import java.util.Set; + +/** + * Lightweight read-only check for the {@code run_select_query} tool. + * + * <p><b>Design note — we deliberately do NOT use a SQL parser.</b> Big-data engines (Spark, Trino, + * Hive, Flink, Impala, ClickHouse, ...) all extend ANSI SQL with non-standard syntax. Any general + * parser (JSqlParser, Calcite, ...) will either fail on legitimate dialect-specific queries (false + * rejection) or require a per-engine grammar (heavy dependency footprint, version drift). For the + * data-agent use case, the SQL is generated by an LLM that has been told via the system prompt + * which tool to use for which kind of statement, and modern LLMs essentially never make the mistake + * of feeding an INSERT to a tool whose description says "read-only". A keyword whitelist on the + * first significant token is a sufficient guard against accidents and clearly communicates intent + * to the LLM via the rejection message — without dragging in a parser dependency we'd have to + * maintain forever. + * + * <p>The whitelist is intentionally broad to cover common big-data read-only entry points. If a + * legitimate read-only keyword is missing, add it here rather than working around it elsewhere. + */ +final class SqlReadOnlyChecker { + + /** + * Read-only keywords accepted as the first token of a SQL statement. + * + * <ul> + * <li>{@code SELECT} — standard query + * <li>{@code WITH} — CTE-prefixed query + * <li>{@code VALUES} — standalone VALUES expression (Trino/Postgres) + * <li>{@code TABLE} — Spark/Postgres {@code TABLE x} shorthand for {@code SELECT * FROM x} + * <li>{@code FROM} — Hive {@code FROM t SELECT ...} variant + * <li>{@code SHOW} — SHOW TABLES / DATABASES / CREATE / PARTITIONS / FUNCTIONS / COLUMNS / + * VIEWS / TBLPROPERTIES / INDEXES / GRANT / ROLE / STATS / LOCKS / ... + * <li>{@code DESCRIBE} / {@code DESC} — table / function / formatted / extended + * <li>{@code EXPLAIN} — query plan inspection (note: in Postgres {@code EXPLAIN ANALYZE} + * actually executes; we trust the LLM not to wrap a mutation in EXPLAIN ANALYZE for the + * read-only tool) + * <li>{@code USE} — switch session catalog/database; not data-mutating + * <li>{@code LIST} — Spark {@code LIST FILE} / {@code LIST JAR} inspection + * <li>{@code HELP} — some engines expose interactive help + * </ul> + */ + private static final Set<String> READ_ONLY_KEYWORDS = + Collections.unmodifiableSet( + new HashSet<>( + Arrays.asList( + "SELECT", + "WITH", + "VALUES", + "TABLE", + "FROM", + "SHOW", + "DESCRIBE", + "DESC", + "EXPLAIN", + "USE", + "LIST", + "HELP"))); + + private SqlReadOnlyChecker() {} + + /** + * Returns true if the SQL's first significant token is a known read-only keyword. Strips leading + * whitespace and {@code --} / {@code /* *}{@code /} comments before testing. + */ + static boolean isReadOnly(String sql) { + if (sql == null) { + return false; + } + String stripped = stripLeadingNoise(sql); + if (stripped.isEmpty()) { + return false; + } + int end = 0; + while (end < stripped.length() && Character.isLetter(stripped.charAt(end))) { + end++; + } + if (end == 0) { + return false; + } + String firstToken = stripped.substring(0, end).toUpperCase(Locale.ROOT); + return READ_ONLY_KEYWORDS.contains(firstToken); + } Review Comment: `SqlReadOnlyChecker` only whitelists the first token, which allows common mutation patterns to bypass `run_select_query` (e.g., `WITH x AS (DELETE ... RETURNING ...) SELECT ...`, or `EXPLAIN ANALYZE INSERT/UPDATE/DELETE ...` on dialects that execute the statement). This undermines the approval gate for destructive operations. Please harden read-only enforcement (e.g., explicitly detect/forbid DML/DDL tokens beyond the first keyword, parse/validate the main statement keyword after a CTE, and/or set the JDBC connection/transaction to read-only so the server enforces it). ########## externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/tool/sql/SqlExecutor.java: ########## @@ -0,0 +1,137 @@ +/* + * 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. + */ + +package org.apache.kyuubi.engine.dataagent.tool.sql; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.Statement; +import javax.sql.DataSource; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Shared JDBC execution + result formatting for SQL tools. Owned by {@link RunSelectQueryTool} and + * {@link RunMutationQueryTool}; not part of the public tool API. + */ +final class SqlExecutor { + + private static final Logger LOG = LoggerFactory.getLogger(SqlExecutor.class); + + private SqlExecutor() {} + + /** + * Execute the given SQL and return a markdown-formatted result string. Errors are caught and + * returned as a one-line "Error: ..." message so the LLM can self-correct. + * + * <p><b>No client-side row cap.</b> The tool deliberately does not call {@link + * Statement#setMaxRows(int)}. Result-size discipline is delegated to the LLM via the system + * prompt, which requires every read query to include an explicit {@code LIMIT} (or equivalent + * pagination). Adding a hidden client-side cap here would silently truncate the LLM's view of the + * data and create magic numbers that are impossible to tune across deployments. + * Review Comment: This code/documentation explicitly disables any client-side row/output limiting, but the PR description mentions maxRows enforcement and output truncation. Returning entire result sets can easily explode response sizes and context windows if the LLM/user forgets a LIMIT. Please either implement a hard cap/truncation in `SqlExecutor` (while clearly indicating truncation + total rows), or update the PR description to match the actual behavior. ########## externals/kyuubi-data-agent-engine/src/test/java/org/apache/kyuubi/engine/dataagent/prompt/SystemPromptBuilderTest.java: ########## @@ -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. + */ + +package org.apache.kyuubi.engine.dataagent.prompt; + +import static org.junit.Assert.*; + +import java.time.LocalDate; +import org.junit.Test; + +public class SystemPromptBuilderTest { + + @Test + public void testDefaultBuildContainsBaseAndDate() { + String prompt = SystemPromptBuilder.create().build(); + assertTrue(prompt.contains("data analysis agent")); + assertTrue(prompt.contains(LocalDate.now().toString())); + } Review Comment: These assertions depend on `LocalDate.now()` being the same when the test runs and when `SystemPromptBuilder.build()` runs, which can be flaky around midnight in the test environment. Consider injecting a `Clock`/date supplier into `SystemPromptBuilder` (or letting tests pass a fixed date) so tests can assert deterministically. ########## externals/kyuubi-data-agent-engine/src/test/java/org/apache/kyuubi/engine/dataagent/mysql/WithMySQLContainer.java: ########## @@ -0,0 +1,125 @@ +/* + * 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. + */ + +package org.apache.kyuubi.engine.dataagent.mysql; + +import com.zaxxer.hikari.HikariConfig; +import com.zaxxer.hikari.HikariDataSource; +import java.sql.Connection; +import java.sql.Statement; +import org.apache.kyuubi.engine.dataagent.tool.ToolRegistry; +import org.apache.kyuubi.engine.dataagent.tool.sql.RunMutationQueryTool; +import org.apache.kyuubi.engine.dataagent.tool.sql.RunSelectQueryTool; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.testcontainers.mysql.MySQLContainer; +import org.testcontainers.utility.DockerImageName; + +/** + * Shared base class for MySQL integration tests. Starts a MySQL 8.0.32 container once per test + * class and provides a pooled DataSource. + */ +public abstract class WithMySQLContainer { + + private static final String MYSQL_IMAGE = + System.getProperty("kyuubi.test.mysql.image", "mysql:8.0.32"); + + private static final long TOOL_CALL_TIMEOUT_SECONDS = 30; + private static final int QUERY_TIMEOUT_SECONDS = 10; + + protected static MySQLContainer mysql; + protected static HikariDataSource dataSource; + protected static ToolRegistry registry; + + @BeforeClass + public static void startContainer() { + DockerImageName imageName = + DockerImageName.parse(MYSQL_IMAGE).asCompatibleSubstituteFor("mysql"); + mysql = + new MySQLContainer(imageName) + .withUsername("root") + .withPassword("kyuubi") + .withDatabaseName("test_db") + .withEnv("MYSQL_ROOT_PASSWORD", "kyuubi"); + mysql.start(); + + HikariConfig config = new HikariConfig(); + config.setJdbcUrl(mysql.getJdbcUrl()); + config.setUsername(mysql.getUsername()); + config.setPassword(mysql.getPassword()); + config.setMaximumPoolSize(5); + config.setMinimumIdle(1); + config.setPoolName("kyuubi-mysql-it"); + dataSource = new HikariDataSource(config); + + registry = new ToolRegistry(TOOL_CALL_TIMEOUT_SECONDS); + registry.register(new RunSelectQueryTool(dataSource, QUERY_TIMEOUT_SECONDS)); + registry.register(new RunMutationQueryTool(dataSource, QUERY_TIMEOUT_SECONDS)); + } + + @AfterClass + public static void stopContainer() { + if (dataSource != null) { + dataSource.close(); + } + if (mysql != null) { + mysql.stop(); + } + } + + /** Execute a SQL statement that does not return a result set. */ + protected static void exec(String sql) { + try (Connection conn = dataSource.getConnection(); + Statement stmt = conn.createStatement()) { + stmt.execute(sql); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + /** + * Invoke a tool through the ToolRegistry, same as the LLM runtime does. Builds the JSON args + * envelope automatically. + */ + protected static String callTool(String toolName, String sql) { + return registry.executeTool(toolName, "{\"sql\": \"" + sql.replace("\"", "\\\"") + "\"}"); + } Review Comment: `callTool` builds JSON via string concatenation and only escapes double-quotes. If `sql` contains backslashes, newlines, tabs, or other control characters this can produce invalid JSON and flaky tests. Consider constructing the JSON with an ObjectMapper (or a minimal JSON escaper) instead of manual concatenation. ########## externals/kyuubi-data-agent-engine/src/main/java/org/apache/kyuubi/engine/dataagent/tool/ToolRegistry.java: ########## @@ -0,0 +1,183 @@ +/* + * 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. + */ + +package org.apache.kyuubi.engine.dataagent.tool; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.openai.core.JsonValue; +import com.openai.models.FunctionDefinition; +import com.openai.models.FunctionParameters; +import com.openai.models.chat.completions.ChatCompletionCreateParams; +import com.openai.models.chat.completions.ChatCompletionFunctionTool; +import com.openai.models.chat.completions.ChatCompletionTool; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.concurrent.Callable; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Registry for agent tools. Provides tool lookup, schema generation for LLM requests, and + * deserialization + execution of tool calls. + * + * <p>Tool metadata (name, description) comes from {@link AgentTool}. Parameter schemas are + * generated from the args class via {@link ToolSchemaGenerator}. The SDK's {@code + * function.arguments(Class)} is used for deserializing LLM responses. + * + * <p><b>Concurrency model: single-writer at startup.</b> {@link #register} is expected to be called + * only during engine initialization, before any LLM request is dispatched. After that the registry + * is read-only. {@link #addToolsTo} therefore does not synchronize on the {@code tools} map, which + * is safe under this single-writer-at-startup assumption — do NOT "fix" this by adding a lock + * around the read path. If we ever need true runtime registration, the right move is to swap {@code + * tools} for an immutable snapshot per refresh, not to grab a mutex on the hot path. + */ +public class ToolRegistry { + + private static final Logger LOG = LoggerFactory.getLogger(ToolRegistry.class); + private static final ObjectMapper JSON = new ObjectMapper(); + + private final Map<String, AgentTool<?>> tools = new LinkedHashMap<>(); + private volatile Map<String, ChatCompletionTool> cachedSpecs; + private final long toolCallTimeoutSeconds; + private final ExecutorService executor; + + /** + * @param toolCallTimeoutSeconds wall-clock cap on every {@link #executeTool} call, sourced from + * {@code kyuubi.engine.data.agent.tool.call.timeout}. When the timeout fires, the thread is + * interrupted and a descriptive error is returned to the LLM. + */ + public ToolRegistry(long toolCallTimeoutSeconds) { + this.toolCallTimeoutSeconds = toolCallTimeoutSeconds; + this.executor = + Executors.newCachedThreadPool( + r -> { + Thread t = new Thread(r, "tool-call-worker"); + t.setDaemon(true); + return t; + }); + } Review Comment: `ToolRegistry` creates an unbounded `newCachedThreadPool` and never shuts it down. In a long-running engine this can leak threads across lifecycle events, and under load it can create an unbounded number of workers. Consider using a bounded executor (size configurable) and exposing a `close()`/`shutdown()` hook so the engine can terminate the pool cleanly on shutdown. ########## kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala: ########## @@ -3824,7 +3824,24 @@ object KyuubiConf { val ENGINE_DATA_AGENT_QUERY_TIMEOUT: ConfigEntry[Long] = buildConf("kyuubi.engine.data.agent.query.timeout") - .doc("The query execution timeout for the Data Agent SQL tool.") + .doc("The JDBC query execution timeout for the Data Agent SQL tools. " + + "Passed to <code>Statement.setQueryTimeout</code> so the server (Spark/Trino/...) " + + "can cooperatively cancel long-running queries and release cluster resources. " + + "Should be set lower than " + + "<code>kyuubi.engine.data.agent.tool.call.timeout</code> so server-side cancellation " + + "has time to react before the outer wall-clock cap fires.") + .version("1.12.0") + .timeConf + .createWithDefaultString("PT3M") + + val ENGINE_DATA_AGENT_TOOL_CALL_TIMEOUT: ConfigEntry[Long] = + buildConf("kyuubi.engine.data.agent.tool.call.timeout") + .doc("The maximum wall-clock execution time for any tool call in the Data Agent " + + "engine. Acts as the outer safety net enforced by the agent runtime via " + + "<code>Future.cancel()</code>, applied uniformly to every tool. " + + "For SQL tools the inner JDBC-level timeout is controlled separately by " + + "<code>kyuubi.engine.data.agent.query.timeout</code>, which should be set lower " + + "so server-side cancellation has time to react before this hard cap fires.") Review Comment: The config entry is defined as `timeConf` (stored as milliseconds), but the doc says it is passed to `Statement.setQueryTimeout`, which expects **seconds**. Please clarify the units/conversion here (and for `kyuubi.engine.data.agent.tool.call.timeout`) to avoid misconfiguration or accidentally applying millisecond values as seconds in the runtime integration. -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
