wangzhigang1999 commented on code in PR #7400: URL: https://github.com/apache/kyuubi/pull/7400#discussion_r3109372208
########## 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: - `conn.setReadOnly(true)` — Spark/Hive Thrift JDBC drivers throw or silently ignore this, and they're our main targets. Won't work. - Parsing past CTEs — Spark, Trino, Hive dialects break JSqlParser and Calcite all the time on valid SQL. Per-dialect grammars would be a maintenance nightmare. - Full-text DML/DDL token scan — `SELECT 'delete' AS action` or a column named `update` would get wrongly blocked. The SQL comes from an LLM already prompted to route mutations through `run_mutation_query`. Happy to revisit if any of this bites us. -- 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]
