diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java index 1733f0bd..9ac9b6b5 100644 --- a/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/base/BaseSQLActuator.java @@ -125,7 +125,7 @@ public int getAffectedRows(SQL sql) throws SQLException { @Override public List parseSQL(SQL sql) { - return SQLUtils.parseStatements(sql.getSql(), this.druidDbType).stream() + return SqlValidator.parse(this.druidDbType, sql.getSql()).stream() .map(stmt -> SQLUtils.toSQLString(stmt, this.druidDbType)) .toList(); } diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SqlValidationResult.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SqlValidationResult.java new file mode 100644 index 00000000..6e3da174 --- /dev/null +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SqlValidationResult.java @@ -0,0 +1,33 @@ +package org.jumpserver.chen.framework.datasource.sql; + +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +public record SqlValidationResult( + boolean parseable, + int statementCount, + String statementType, + List tables, + List columns, + List errors +) { + public SqlValidationResult { + tables = tables == null ? List.of() : List.copyOf(tables); + columns = columns == null ? List.of() : List.copyOf(columns); + errors = errors == null ? List.of() : List.copyOf(errors); + statementType = statementType == null ? "UNKNOWN" : statementType; + } + + public Map toAnalysisMap() { + Map result = new LinkedHashMap<>(); + result.put("valid", parseable); + result.put("parseable", parseable); + result.put("statementCount", statementCount); + result.put("statementType", statementType); + result.put("tables", tables); + result.put("columns", columns); + result.put("errors", errors); + return result; + } +} diff --git a/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SqlValidator.java b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SqlValidator.java new file mode 100644 index 00000000..3533a2bb --- /dev/null +++ b/backend/framework/src/main/java/org/jumpserver/chen/framework/datasource/sql/SqlValidator.java @@ -0,0 +1,106 @@ +package org.jumpserver.chen.framework.datasource.sql; + +import com.alibaba.druid.DbType; +import com.alibaba.druid.sql.SQLUtils; +import com.alibaba.druid.sql.ast.SQLStatement; +import com.alibaba.druid.sql.visitor.SchemaStatVisitor; +import com.alibaba.druid.stat.TableStat; +import org.apache.commons.lang3.StringUtils; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Set; + +public final class SqlValidator { + public static final String QUERY_UNSUPPORTED_MESSAGE = + "This SQL is not supported by Query because Chen cannot parse it with Druid. Use Console to draft and run it."; + + private static final int MAX_OBJECTS = 50; + private static final int MAX_ANALYSIS_COLUMNS = 512; + + private SqlValidator() { + } + + public static List parse(DbType dbType, String sql) { + return SQLUtils.parseStatements(sql, dbType); + } + + public static SqlValidationResult validate(DbType dbType, String sql) { + if (StringUtils.isBlank(sql)) { + return unparseable("SQL statement is empty"); + } + + List statements; + try { + statements = parse(dbType, sql); + if (statements == null || statements.isEmpty()) { + return unparseable("SQL statement is empty"); + } + } catch (RuntimeException e) { + return unparseable(safeError(e)); + } + + LinkedHashSet tables = new LinkedHashSet<>(); + LinkedHashSet columns = new LinkedHashSet<>(); + String statementType = statements.size() == 1 ? statementType(statements.get(0)) : "MULTI"; + for (SQLStatement statement : statements) { + try { + SchemaStatVisitor visitor = SQLUtils.createSchemaStatVisitor(dbType); + statement.accept(visitor); + for (TableStat.Name table : visitor.getTables().keySet()) { + addBounded(tables, table.toString(), MAX_OBJECTS); + } + Collection statementColumns = visitor.getColumns(); + for (TableStat.Column column : statementColumns) { + addBounded(columns, column.toString(), MAX_ANALYSIS_COLUMNS); + } + } catch (RuntimeException ignored) { + // Some vendor-specific statements are syntactically valid but do not support schema statistics. + } + } + return new SqlValidationResult( + true, + statements.size(), + statementType, + List.copyOf(tables), + List.copyOf(columns), + List.of() + ); + } + + private static SqlValidationResult unparseable(String error) { + return new SqlValidationResult( + false, + 0, + "UNKNOWN", + List.of(), + List.of(), + List.of(error) + ); + } + + private static String statementType(SQLStatement statement) { + String name = statement.getClass().getSimpleName().toUpperCase(Locale.ROOT); + if (name.startsWith("SQL")) { + name = name.substring(3); + } + if (name.endsWith("STATEMENT")) { + name = name.substring(0, name.length() - "STATEMENT".length()); + } + return name; + } + + private static String safeError(RuntimeException error) { + String message = StringUtils.defaultIfBlank(error.getMessage(), error.getClass().getSimpleName()); + return message.length() > 4096 ? message.substring(0, 4096) : message; + } + + private static void addBounded(Set values, String value, int max) { + if (values.size() < max && StringUtils.isNotBlank(value)) { + values.add(value); + } + } +} diff --git a/backend/modules/src/main/java/org.jumpserver.chen.modules/oracle/OracleActuator.java b/backend/modules/src/main/java/org.jumpserver.chen.modules/oracle/OracleActuator.java index 8f18190c..7ebcd6d3 100644 --- a/backend/modules/src/main/java/org.jumpserver.chen.modules/oracle/OracleActuator.java +++ b/backend/modules/src/main/java/org.jumpserver.chen.modules/oracle/OracleActuator.java @@ -8,6 +8,7 @@ import org.jumpserver.chen.framework.datasource.sql.SQLExecutePlan; import org.jumpserver.chen.framework.datasource.sql.SQLIdentifier; import org.jumpserver.chen.framework.datasource.sql.SQLQueryParams; +import org.jumpserver.chen.framework.datasource.sql.SqlValidator; import java.sql.Connection; import java.sql.SQLException; @@ -67,7 +68,7 @@ public SQLExecutePlan createPlan(SQL sql) throws SQLException { @Override public List parseSQL(SQL sql) { var dbType = this.getDbType(); - var statements = SQLUtils.parseStatements(sql.getSql(), dbType); + var statements = SqlValidator.parse(dbType, sql.getSql()); for (var i = 0; i < statements.size() - 1; i++) { if (!statements.get(i).isAfterSemi()) { throw new ParserException("Multiple SQL statements must be separated by semicolons"); diff --git a/backend/web/src/main/java/org/jumpserver/chen/web/ai/AgentWebSocketHandler.java b/backend/web/src/main/java/org/jumpserver/chen/web/ai/AgentWebSocketHandler.java index b08c7792..ca7ecc41 100644 --- a/backend/web/src/main/java/org/jumpserver/chen/web/ai/AgentWebSocketHandler.java +++ b/backend/web/src/main/java/org/jumpserver/chen/web/ai/AgentWebSocketHandler.java @@ -327,7 +327,8 @@ static List> toolDefinitions() { tools.add(tool( "validate_sql", "Validate SQL draft", - "Parse SQL locally in Chen and return statement count, type, referenced objects and risk. " + "Parse SQL locally in Chen with Druid and return statement count, type and referenced objects. " + + "valid=false means Chen cannot parse the SQL, not that it is illegal or unsafe. " + "This never executes SQL.", "{\"type\":\"object\",\"additionalProperties\":false,\"required\":[\"sql\"]," + "\"properties\":{\"sql\":{\"type\":\"string\",\"minLength\":1,\"maxLength\":131072," @@ -337,12 +338,14 @@ static List> toolDefinitions() { Map proposalTool = tool( "propose_sql", "Propose SQL draft", - "Validate exactly one SQL statement, prepare a draft for explicit user review, and wait for the " - + "user to apply or reject it. This never executes SQL.", + "Prepare SQL as a draft for explicit user review, and wait for the " + + "user to apply or reject it. Query only accepts SQL Chen can parse with Druid; " + + "unparseable vendor-native SQL must be proposed in Console, which inserts it with an explicit notice. " + + "This never executes SQL.", "{\"type\":\"object\",\"additionalProperties\":false," + "\"required\":[\"sql\",\"explanation\"],\"properties\":{" + "\"sql\":{\"type\":\"string\",\"minLength\":1,\"maxLength\":131072," - + "\"description\":\"Exactly one complete SQL statement in the verified dialect\"}," + + "\"description\":\"Complete SQL in the verified dialect\"}," + "\"explanation\":{\"type\":\"string\",\"maxLength\":4096," + "\"description\":\"Concise explanation for the user reviewing the draft\"}}}", false diff --git a/backend/web/src/main/java/org/jumpserver/chen/web/ai/SqlAgentToolService.java b/backend/web/src/main/java/org/jumpserver/chen/web/ai/SqlAgentToolService.java index ebce4a5c..3b5df396 100644 --- a/backend/web/src/main/java/org/jumpserver/chen/web/ai/SqlAgentToolService.java +++ b/backend/web/src/main/java/org/jumpserver/chen/web/ai/SqlAgentToolService.java @@ -1,10 +1,6 @@ package org.jumpserver.chen.web.ai; import com.alibaba.druid.DbType; -import com.alibaba.druid.sql.SQLUtils; -import com.alibaba.druid.sql.ast.SQLStatement; -import com.alibaba.druid.sql.visitor.SchemaStatVisitor; -import com.alibaba.druid.stat.TableStat; import com.google.gson.Gson; import com.google.gson.JsonElement; import com.google.gson.JsonObject; @@ -25,29 +21,30 @@ import org.jumpserver.chen.framework.datasource.metadata.RelationKind; import org.jumpserver.chen.framework.datasource.metadata.RelationMetadata; import org.jumpserver.chen.framework.datasource.metadata.RelationScope; +import org.jumpserver.chen.framework.datasource.sql.SqlValidationResult; +import org.jumpserver.chen.framework.datasource.sql.SqlValidator; import org.jumpserver.chen.framework.session.Session; import org.springframework.stereotype.Service; import java.sql.SQLException; import java.util.ArrayList; -import java.util.Collection; import java.util.LinkedHashMap; import java.util.LinkedHashSet; import java.util.List; import java.util.Locale; import java.util.Map; import java.util.Set; -import java.util.regex.Pattern; @Service public class SqlAgentToolService { + static final String CONSOLE_UNPARSEABLE_NOTICE = "Chen cannot validate this SQL with Druid.\n\n" + + "It may use database-native syntax.\n\n" + + "Please review it before execution."; private static final Gson GSON = new Gson(); private static final int MAX_CONTEXT_BYTES = 256 * 1024; private static final int MAX_SQL_BYTES = 128 * 1024; private static final int MAX_TOOL_ARGUMENT_BYTES = MAX_SQL_BYTES + 8 * 1024; private static final int MAX_IDENTIFIER_BYTES = 1024; - private static final int MAX_OBJECTS = 50; - private static final int MAX_ANALYSIS_COLUMNS = 512; static final int MAX_INSPECT_TABLES = 8; static final int MAX_DISCOVER_TABLES = 100; static final String DISCOVER_TABLES_QUERY = "*"; @@ -152,7 +149,6 @@ public AgentRequestContext resolveRequestContext(Session session, String context String targetSql = hasSelection ? selectedSql : documentSql; Map sqlAnalysis = StringUtils.isBlank(targetSql) ? null : validateSQL(datasource.getDruidDbType(), targetSql); - assertAllowedAiScope(dialect, database, schema, targetSql, sqlAnalysis); JsonObject sanitized = new JsonObject(); sanitized.addProperty("dialect", dialect); addNullableString(sanitized, "database", database); @@ -225,9 +221,9 @@ public String execute(Session session, AgentRequestContext context, String toolN ); case "validate_sql" -> Map.of( "kind", "validation", - "analysis", validateSQL( + "analysis", SqlValidator.validate( datasource.getDruidDbType(), boundedString(arguments, "sql", MAX_SQL_BYTES, true) - ) + ).toAnalysisMap() ); case "propose_sql" -> proposeSQL(datasource, context, arguments); default -> throw new IllegalArgumentException("Unsupported SQL assistant tool"); @@ -242,14 +238,16 @@ private static Map proposeSQL( ) { String sql = boundedString(arguments, "sql", MAX_SQL_BYTES, true).trim(); String explanation = boundedString(arguments, "explanation", 4 * 1024, false).trim(); - Map analysis = validateSQL(datasource.getDruidDbType(), sql); - assertAllowedAiScope(context.dialect(), context.database(), context.schema(), sql, analysis); - if (!Boolean.TRUE.equals(analysis.get("valid")) - || ((Number) analysis.getOrDefault("statementCount", 0)).intValue() != 1) { - throw new IllegalArgumentException("The SQL proposal must contain exactly one valid statement"); - } - JsonObject editor = JsonParser.parseString(context.sanitizedJson()).getAsJsonObject(); + boolean consoleWorkspace = "console".equals(boundedString(editor, "workspaceTabKind", 32, false)); + SqlValidationResult validation = SqlValidator.validate(datasource.getDruidDbType(), sql); + if (!validation.parseable()) { + if (!consoleWorkspace) { + throw new IllegalArgumentException(SqlValidator.QUERY_UNSUPPORTED_MESSAGE); + } + explanation = appendNotice(explanation, CONSOLE_UNPARSEABLE_NOTICE); + } + Map analysis = validation.parseable() ? validation.toAnalysisMap() : null; int selectionFrom = editor.get("selectionFrom").getAsInt(); int selectionTo = editor.get("selectionTo").getAsInt(); String tabId = editor.get("tabId").getAsString(); @@ -277,9 +275,22 @@ private static Map proposeSQL( proposal.put("sql", sql); proposal.put("originalSql", originalSQL); proposal.put("explanation", explanation); - proposal.put("analysis", analysis); + if (analysis != null) { + proposal.put("analysis", analysis); + } proposal.put("base", base); - return Map.of("kind", "proposal", "analysis", analysis, "proposal", proposal); + + Map result = new LinkedHashMap<>(); + result.put("kind", "proposal"); + if (analysis != null) { + result.put("analysis", analysis); + } + result.put("proposal", proposal); + return result; + } + + private static String appendNotice(String explanation, String notice) { + return explanation.isBlank() ? notice : explanation + "\n\n" + notice; } MetadataApprovalScope resolveMetadataApprovalScope( @@ -354,55 +365,7 @@ private static JsonObject parseToolArguments(String argumentsJson) { } static Map validateSQL(DbType dbType, String sql) { - Map result = new LinkedHashMap<>(); - LinkedHashSet tables = new LinkedHashSet<>(); - LinkedHashSet columns = new LinkedHashSet<>(); - List errors = new ArrayList<>(); - List statements; - try { - statements = SQLUtils.parseStatements(sql, dbType); - if (statements.isEmpty()) { - throw new IllegalArgumentException("SQL statement is empty"); - } - } catch (RuntimeException e) { - result.put("valid", false); - result.put("statementCount", 0); - result.put("statementType", "UNKNOWN"); - result.put("riskLevel", 0); - result.put("riskReason", "SQL syntax validation failed"); - result.put("tables", List.of()); - result.put("columns", List.of()); - result.put("errors", List.of(safeError(e))); - return result; - } - - int riskLevel = 0; - String statementType = statements.size() == 1 ? statementType(statements.get(0)) : "MULTI"; - for (SQLStatement statement : statements) { - riskLevel = Math.max(riskLevel, riskLevel(statementType(statement))); - try { - SchemaStatVisitor visitor = SQLUtils.createSchemaStatVisitor(dbType); - statement.accept(visitor); - for (TableStat.Name table : visitor.getTables().keySet()) { - addBounded(tables, table.toString(), MAX_OBJECTS); - } - Collection statementColumns = visitor.getColumns(); - for (TableStat.Column column : statementColumns) { - addBounded(columns, column.toString(), MAX_ANALYSIS_COLUMNS); - } - } catch (RuntimeException ignored) { - // Some vendor-specific statements are syntactically valid but do not support schema statistics. - } - } - result.put("valid", true); - result.put("statementCount", statements.size()); - result.put("statementType", statementType); - result.put("riskLevel", riskLevel); - result.put("riskReason", riskReason(riskLevel, statements.size())); - result.put("tables", List.copyOf(tables)); - result.put("columns", List.copyOf(columns)); - result.put("errors", errors); - return result; + return SqlValidator.validate(dbType, sql).toAnalysisMap(); } private Map inspectSchema( @@ -736,45 +699,6 @@ private static void assertRequestContextCurrent(Session session, AgentRequestCon } } - private static void assertAllowedAiScope( - String dialect, - String database, - String schema, - String sql, - Map sqlAnalysis - ) { - boolean sqlIsValid = sqlAnalysis == null || Boolean.TRUE.equals(sqlAnalysis.get("valid")); - if (isBlockedSystemScope(dialect, database, schema)) { - throw new IllegalArgumentException(blockedSystemScopeMessage(dialect, database, schema)); - } - if (!sqlIsValid && containsBlockedSystemQualifier(dialect, sql)) { - throw new IllegalArgumentException("The SQL references protected system database objects"); - } - if (sqlAnalysis == null) { - return; - } - Object tablesValue = sqlAnalysis.get("tables"); - if (!(tablesValue instanceof Collection tables)) { - return; - } - for (Object value : tables) { - List parts; - try { - parts = qualifiedIdentifierParts(String.valueOf(value)); - } catch (IllegalArgumentException ignored) { - continue; - } - if (parts.size() == 2 && (isBlockedSystemDatabase(dialect, parts.get(0)) - || isBlockedSystemSchema(dialect, parts.get(0)))) { - throw new IllegalArgumentException("The SQL references protected system database objects"); - } - if (parts.size() == 3 && (isBlockedSystemDatabase(dialect, parts.get(0)) - || isBlockedSystemSchema(dialect, parts.get(1)))) { - throw new IllegalArgumentException("The SQL references protected system database objects"); - } - } - } - private static String blockedSystemScopeMessage(String dialect, String database, String schema) { if (isBlockedSystemSchema(dialect, schema)) { return "The active schema '" + normalizePolicyIdentifier(schema) @@ -791,45 +715,6 @@ private static boolean isBlockedSystemScope(String dialect, String database, Str return isBlockedSystemDatabase(dialect, database) || isBlockedSystemSchema(dialect, schema); } - private static boolean containsBlockedSystemQualifier(String dialect, String sql) { - if (StringUtils.isBlank(sql)) { - return false; - } - String normalized = sql.toLowerCase(Locale.ROOT) - .replace("\"", "") - .replace("`", "") - .replace("[", "") - .replace("]", ""); - Set identifiers = new LinkedHashSet<>(); - identifiers.add("information_schema"); - String db = StringUtils.defaultString(dialect).toLowerCase(Locale.ROOT); - switch (db) { - case "mysql", "mariadb" -> identifiers.addAll(Set.of("mysql", "performance_schema", "sys")); - case "postgresql", "postgres" -> identifiers.addAll(Set.of("pg_catalog", "pg_toast")); - case "sqlserver" -> identifiers.addAll(Set.of("master", "model", "msdb", "tempdb", "sys")); - case "oracle" -> identifiers.addAll(Set.of("sys", "system", "xdb", "mdsys", "ctxsys", "audsys")); - case "db2" -> identifiers.addAll(Set.of( - "sysibm", "syscat", "sysstat", "sysfun", "sysproc", "systools" - )); - case "clickhouse" -> identifiers.add("system"); - case "dm", "dameng" -> identifiers.addAll(Set.of("sys", "system", "sysauditor")); - default -> { - } - } - for (String identifier : identifiers) { - Pattern qualifier = Pattern.compile( - "(? metadataScopeMap(String catalog, String schem return result; } - private static String statementType(SQLStatement statement) { - String name = statement.getClass().getSimpleName().toUpperCase(Locale.ROOT); - if (name.startsWith("SQL")) { - name = name.substring(3); - } - if (name.endsWith("STATEMENT")) { - name = name.substring(0, name.length() - "STATEMENT".length()); - } - return name; - } - - private static int riskLevel(String statementType) { - String type = statementType.toUpperCase(Locale.ROOT); - if (type.contains("SELECT") || type.contains("SHOW") || type.contains("DESC") - || type.contains("EXPLAIN") || type.contains("WITH")) { - return 1; - } - if (type.contains("INSERT") || type.contains("UPDATE") || type.contains("DELETE") - || type.contains("MERGE") || type.contains("REPLACE")) { - return 3; - } - if (type.contains("DROP") || type.contains("TRUNCATE") || type.contains("GRANT") - || type.contains("REVOKE")) { - return 4; - } - if (type.contains("CREATE") || type.contains("ALTER") || type.contains("RENAME")) { - return 3; - } - return 2; - } - - private static String riskReason(int riskLevel, int statementCount) { - String base = switch (riskLevel) { - case 1 -> "Read-only SQL statement"; - case 3 -> "SQL may change database data or schema"; - case 4 -> "SQL may remove data or change privileges"; - default -> "SQL statement requires manual review"; - }; - return statementCount > 1 ? base + "; contains multiple statements" : base; - } - - private static String safeError(RuntimeException error) { - String message = StringUtils.defaultIfBlank(error.getMessage(), error.getClass().getSimpleName()); - return message.length() > 4096 ? message.substring(0, 4096) : message; - } - - private static void addBounded(Set values, String value, int max) { - if (values.size() < max && StringUtils.isNotBlank(value)) { - values.add(value); - } - } - static JsonObject sanitizeLastError(JsonElement value) { if (value == null || value.isJsonNull()) { return null;