From ea18ab3ba0c318059382e58d19281c05fa059715 Mon Sep 17 00:00:00 2001 From: Hz_ Date: Tue, 14 Jul 2026 15:02:28 +0800 Subject: [PATCH] fix(go-agent): align ExeSQL component inputs with Python (#16896) ## Summary - Resolve configured ExeSQL SQL references before tool invocation. - Align ExeSQL parameters, defaults, JSON tags, and tool schema with Python. - Preserve SQL string literals and restore Canvas output fields. ## Testing - `bash build.sh --test ./internal/agent/tool/...` - `bash build.sh --test ./internal/agent/component/...` - `bash build.sh --test ./internal/agent/runtime/...` image --- .../agent/component/exesql_component_test.go | 100 +++++ .../agent/component/universe_a_wrappers.go | 127 +++++- internal/agent/tool/exesql.go | 395 ++++++++++++------ internal/agent/tool/exesql_decoder_test.go | 60 +++ internal/agent/tool/exesql_test.go | 159 ++++++- internal/agent/tool/registry.go | 24 +- 6 files changed, 703 insertions(+), 162 deletions(-) create mode 100644 internal/agent/component/exesql_component_test.go diff --git a/internal/agent/component/exesql_component_test.go b/internal/agent/component/exesql_component_test.go new file mode 100644 index 0000000000..a069a863c2 --- /dev/null +++ b/internal/agent/component/exesql_component_test.go @@ -0,0 +1,100 @@ +package component + +import ( + "context" + "encoding/json" + "errors" + "strings" + "testing" + + einotool "github.com/cloudwego/eino/components/tool" + + "ragflow/internal/agent/runtime" +) + +type exesqlInvokerStub struct { + arguments string + result string + err error +} + +func (s *exesqlInvokerStub) InvokableRun(_ context.Context, arguments string, _ ...einotool.Option) (string, error) { + s.arguments = arguments + return s.result, s.err +} + +func TestExeSQLComponentResolvesConfiguredSQL(t *testing.T) { + stub := &exesqlInvokerStub{ + result: `{"columns":["id","status"],"rows":[{"id":1,"status":"Completed"}]}`, + } + state := runtime.NewCanvasState("run", "task") + state.SetVar("Agent:SparklyMooseDivide", "content", "SELECT id FROM orders WHERE status = 'Completed'") + c := &exesqlComponent{ + inner: stub, + sql: "{Agent:SparklyMooseDivide@content}", + } + + out, err := c.Invoke(runtime.WithState(context.Background(), state), map[string]any{ + "content": "this must not be used as SQL", + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + var arguments map[string]any + if err := json.Unmarshal([]byte(stub.arguments), &arguments); err != nil { + t.Fatalf("decode tool arguments: %v", err) + } + if got := arguments["sql"]; got != "SELECT id FROM orders WHERE status = 'Completed'" { + t.Fatalf("sql argument = %#v, want resolved SQL", got) + } + if _, exists := arguments["content"]; exists { + t.Fatalf("tool arguments unexpectedly contain upstream content: %s", stub.arguments) + } + formalized, ok := out["formalized_content"].(string) + if !ok || !strings.Contains(formalized, "Completed") { + t.Fatalf("formalized_content = %#v, want rendered row", out["formalized_content"]) + } + jsonResult, ok := out["json"].([]any) + if !ok || len(jsonResult) != 1 { + t.Fatalf("json = %#v, want one statement result", out["json"]) + } +} + +func TestNewExeSQLComponentKeepsConfiguredSQL(t *testing.T) { + component, err := newExeSQLComponent(map[string]any{ + "db_type": "mysql", + "database": "demo", + "username": "root", + "host": "db.example.com", + "port": 3306, + "password": "secret", + "max_records": 100, + "sql": "{Agent:SparklyMooseDivide@content}", + }) + if err != nil { + t.Fatalf("newExeSQLComponent: %v", err) + } + exeSQL, ok := component.(*exesqlComponent) + if !ok { + t.Fatalf("component type = %T, want *exesqlComponent", component) + } + if exeSQL.sql != "{Agent:SparklyMooseDivide@content}" { + t.Fatalf("configured SQL = %q", exeSQL.sql) + } +} + +func TestExeSQLComponentPreservesToolErrorAsCanvasOutput(t *testing.T) { + stub := &exesqlInvokerStub{ + result: `{"_ERROR":"exesql: empty sql"}`, + err: errors.New("exesql: empty sql"), + } + c := &exesqlComponent{inner: stub, sql: ""} + + out, err := c.Invoke(context.Background(), map[string]any{"sql": ""}) + if err != nil { + t.Fatalf("Invoke returned a hard error: %v", err) + } + if got := out["_ERROR"]; got != "exesql: empty sql" { + t.Fatalf("_ERROR = %#v, want tool error", got) + } +} diff --git a/internal/agent/component/universe_a_wrappers.go b/internal/agent/component/universe_a_wrappers.go index fa1123815e..56645d6baf 100644 --- a/internal/agent/component/universe_a_wrappers.go +++ b/internal/agent/component/universe_a_wrappers.go @@ -34,6 +34,7 @@ import ( "errors" "fmt" "regexp" + "sort" "strconv" "strings" @@ -1254,7 +1255,10 @@ func newExeSQLComponent(params map[string]any) (Component, error) { if err != nil { return nil, fmt.Errorf("canvas: ExeSQL: %w", err) } - return &exesqlComponent{inner: agenttool.NewExeSQLTool(conn)}, nil + return &exesqlComponent{ + inner: agenttool.NewExeSQLTool(conn), + sql: conn.SQL, + }, nil } // translateExeSQLParamsToToolShape adapts a v1 DSL ExeSQL params @@ -1316,16 +1320,20 @@ func translateExeSQLParamsToToolShape(v1Params map[string]any) map[string]any { return out } +type exeSQLInvoker interface { + InvokableRun(ctx context.Context, argsJSON string, opts ...einotool.Option) (string, error) +} + type exesqlComponent struct { - inner *agenttool.ExeSQLTool + inner exeSQLInvoker + sql string } func (c *exesqlComponent) Name() string { return "ExeSQL" } func (c *exesqlComponent) Inputs() map[string]string { return map[string]string{ - "sql": "SQL statement to execute (SELECT-only; DML/DDL rejected).", - "database": "Optional target database/schema (overrides the tool's configured DB).", + "sql": "SQL statement to execute (SELECT-only; DML/DDL rejected).", } } @@ -1340,19 +1348,114 @@ func (c *exesqlComponent) GetInputForm() map[string]any { func (c *exesqlComponent) Outputs() map[string]string { return map[string]string{ - "columns": "Result-set column names.", - "rows": "Result-set rows as column→value maps.", - "sql": "Resolved SQL string (after parameter substitution).", + "formalized_content": "SQL result rendered as Markdown.", + "json": "Raw SQL statement results.", } } func (c *exesqlComponent) Invoke(ctx context.Context, inputs map[string]any) (map[string]any, error) { - argsJSON, _ := json.Marshal(inputs) - out, err := c.inner.InvokableRun(ctx, string(argsJSON)) - if err != nil { - return nil, fmt.Errorf("canvas: ExeSQL: %w", err) + sqlText := c.sql + if value, ok := inputs["sql"].(string); ok && strings.TrimSpace(value) != "" { + sqlText = value } - return parseToolEnvelope(out), nil + if state, _, stateErr := runtime.GetStateFromContext[*runtime.CanvasState](ctx); stateErr == nil && state != nil { + resolved, resolveErr := runtime.ResolveTemplate(sqlText, state) + if resolveErr != nil { + return map[string]any{ + "formalized_content": "", + "json": []any{}, + "_ERROR": resolveErr.Error(), + }, nil + } + sqlText = resolved + } + + argsJSON, err := json.Marshal(map[string]any{"sql": sqlText}) + if err != nil { + return nil, fmt.Errorf("canvas: ExeSQL: encode SQL: %w", err) + } + out, invokeErr := c.inner.InvokableRun(ctx, string(argsJSON)) + decoded := parseToolEnvelope(out) + result := formatExeSQLCanvasOutput(decoded) + if invokeErr != nil { + if message := stringParam(decoded["_ERROR"]); message != "" { + result["_ERROR"] = message + } else { + result["_ERROR"] = invokeErr.Error() + } + } + return result, nil +} + +func formatExeSQLCanvasOutput(decoded map[string]any) map[string]any { + rows := anySlice(decoded["rows"]) + columns := anySlice(decoded["columns"]) + jsonResult := make([]any, 0, 1) + if len(rows) == 1 { + if row, ok := rows[0].(map[string]any); ok && len(row) == 1 { + if _, hasContent := row["content"]; hasContent { + jsonResult = append(jsonResult, row) + } + } + } + if len(jsonResult) == 0 && len(rows) > 0 { + jsonResult = append(jsonResult, rows) + } + result := map[string]any{ + "formalized_content": renderExeSQLMarkdown(columns, rows), + "json": jsonResult, + } + if message := stringParam(decoded["_ERROR"]); message != "" { + result["_ERROR"] = message + } + return result +} + +func renderExeSQLMarkdown(columns, rows []any) string { + if len(rows) == 0 { + return "" + } + for _, value := range rows { + if row, ok := value.(map[string]any); ok && len(row) == 1 { + if message, exists := row["content"]; exists { + return stringParam(message) + } + } + } + columnNames := make([]string, 0, len(columns)) + for _, column := range columns { + columnNames = append(columnNames, fmt.Sprint(column)) + } + if len(columnNames) == 0 { + if first, ok := rows[0].(map[string]any); ok { + for column := range first { + columnNames = append(columnNames, column) + } + sort.Strings(columnNames) + } + } + if len(columnNames) == 0 { + return "" + } + var builder strings.Builder + fmt.Fprintf(&builder, "| %s |\n", strings.Join(columnNames, " | ")) + separators := make([]string, len(columnNames)) + for i := range separators { + separators[i] = "---" + } + fmt.Fprintf(&builder, "| %s |\n", strings.Join(separators, " | ")) + for _, value := range rows { + row, ok := value.(map[string]any) + if !ok { + continue + } + cells := make([]string, len(columnNames)) + for i, column := range columnNames { + cells[i] = strings.ReplaceAll(strings.ReplaceAll(fmt.Sprint(row[column]), "|", "\\|"), "\n", "
") + } + fmt.Fprintf(&builder, "| %s |\n", strings.Join(cells, " | ")) + } + return strings.TrimSuffix(builder.String(), "\n") } func (c *exesqlComponent) Stream(_ context.Context, _ map[string]any) (<-chan map[string]any, error) { diff --git a/internal/agent/tool/exesql.go b/internal/agent/tool/exesql.go index c8b2f40b7b..851b3770a3 100644 --- a/internal/agent/tool/exesql.go +++ b/internal/agent/tool/exesql.go @@ -87,22 +87,34 @@ const ( exesqlToolName = "execute_sql" exesqlToolDescription = "This is a tool that can execute SQL." + exesqlDefaultSQL = "{sys.query}" + exesqlDefaultDBType = "mysql" + exesqlDefaultPort = 3306 exesqlDefaultMaxRecords = 1024 exesqlDefaultTimeout = 60 * time.Second ) -// exesqlConnParams captures the user-supplied DB connection details. -// These are tool-level config (set on the canvas node, not exposed -// to the LLM at function-call time), matching the Python ExeSQLParam -// fields. The LLM only sees `sql` and optional `database` in args. +// exesqlConnParams mirrors Python's ExeSQLParam fields. SQL is the only +// model-emitted runtime input; the remaining fields are Canvas node +// configuration and are not exposed from Info. type exesqlConnParams struct { - DBType string // mysql | postgres | mariadb | mssql | oceanbase - Database string - Username string - Host string - Port int - Password string - MaxRecords int + SQL string `json:"sql"` + DBType string `json:"db_type"` // mysql | postgres | mariadb | mssql | oceanbase + Database string `json:"database"` + Username string `json:"username"` + Host string `json:"host"` + Port int `json:"port"` + Password string `json:"password"` + MaxRecords int `json:"max_records"` +} + +func defaultExeSQLConnParams() exesqlConnParams { + return exesqlConnParams{ + SQL: exesqlDefaultSQL, + DBType: exesqlDefaultDBType, + Port: exesqlDefaultPort, + MaxRecords: exesqlDefaultMaxRecords, + } } // ExeSQLConnParams is the public alias of exesqlConnParams for @@ -112,14 +124,17 @@ type exesqlConnParams struct { type ExeSQLConnParams = exesqlConnParams // NewExeSQLConnParams decodes a canvas-node params map into an -// ExeSQLConnParams. Returns an error if any required field -// (db_type, host, database, username) is missing. +// ExeSQLConnParams. Python defaults are applied before node values; +// host, database, and username remain required configuration. // // Callers (e.g. the Universe A exesqlComponent wrapper) build the // params map from the canvas DSL; the tool-side decoding stays // in this package so the schema lives next to the type. func NewExeSQLConnParams(params map[string]any) (ExeSQLConnParams, error) { - conn := ExeSQLConnParams{} + conn := defaultExeSQLConnParams() + if v, ok := params["sql"].(string); ok { + conn.SQL = v + } if v, ok := params["db_type"].(string); ok { conn.DBType = v } @@ -132,13 +147,13 @@ func NewExeSQLConnParams(params map[string]any) (ExeSQLConnParams, error) { if v, ok := params["host"].(string); ok { conn.Host = v } - if v, ok := params["port"].(int); ok { + if v, ok := intParam(params, "port"); ok { conn.Port = v } if v, ok := params["password"].(string); ok { conn.Password = v } - if v, ok := params["max_records"].(int); ok { + if v, ok := intParam(params, "max_records"); ok { conn.MaxRecords = v } if conn.DBType == "" || conn.Host == "" || conn.Username == "" || conn.Database == "" { @@ -147,11 +162,10 @@ func NewExeSQLConnParams(params map[string]any) (ExeSQLConnParams, error) { return conn, nil } -// exesqlArgs is the JSON shape the model sends in. Matches the Python -// ExeSQLParam ToolMeta (sql is required, database is optional). +// exesqlArgs is the JSON shape the model sends in. It matches Python's +// ExeSQLParam ToolMeta: SQL is the only runtime parameter. type exesqlArgs struct { - SQL string `json:"sql"` - Database string `json:"database,omitempty"` + SQL string `json:"sql"` } // exesqlResult is the JSON envelope returned to the model. The shape @@ -188,7 +202,7 @@ func defaultExeSQLDialer(driver, dsn string) (*sql.DB, error) { } // ExeSQLTool is the ExeSQL tool. -// It validates SELECT-only at the parser level and executes the +// It validates read-only SQL before database access and executes the // statement against a user-configured external DB via `database/sql`. type ExeSQLTool struct { conn exesqlConnParams @@ -199,8 +213,18 @@ type ExeSQLTool struct { // params. The dialer defaults to `sql.Open`; tests can pass a // sqlmock-backed dialer via WithExeSQLDialer. func NewExeSQLTool(conn exesqlConnParams) *ExeSQLTool { - if conn.MaxRecords <= 0 { - conn.MaxRecords = exesqlDefaultMaxRecords + defaults := defaultExeSQLConnParams() + if conn.SQL == "" { + conn.SQL = defaults.SQL + } + if conn.DBType == "" { + conn.DBType = defaults.DBType + } + if conn.Port == 0 { + conn.Port = defaults.Port + } + if conn.MaxRecords == 0 { + conn.MaxRecords = defaults.MaxRecords } return &ExeSQLTool{ conn: conn, @@ -218,8 +242,8 @@ func (e *ExeSQLTool) WithExeSQLDialer(d exesqlDialer) *ExeSQLTool { } // Info returns the tool's metadata for the chat model. Mirrors the -// Python ExeSQLParam ToolMeta: only `sql` (and optional `database`) -// are visible to the LLM. Connection params are not exposed here — +// Python ExeSQLParam ToolMeta: only `sql` is visible to the LLM. +// Connection params are not exposed here — // they're set on the tool instance, matching the Python convention // where ExeSQLParam fields like `db_type` / `host` are tool // configuration, not function-call arguments. @@ -233,11 +257,6 @@ func (e *ExeSQLTool) Info(_ context.Context) (*schema.ToolInfo, error) { Desc: "The SQL statement to execute. Must be a SELECT (read-only).", Required: true, }, - "database": { - Type: schema.String, - Desc: "Optional target database / schema name. Overrides the tool's configured DB.", - Required: false, - }, }), }, nil } @@ -251,7 +270,7 @@ func (e *ExeSQLTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ if argumentsInJSON == "" { return exesqlErrorResult(errors.New("exesql: empty arguments")), errors.New("exesql: empty arguments") } - var args exesqlArgs + args := exesqlArgs{SQL: e.conn.SQL} if err := json.Unmarshal([]byte(argumentsInJSON), &args); err != nil { return exesqlErrorResult(fmt.Errorf("exesql: parse arguments: %w", err)), fmt.Errorf("exesql: parse arguments: %w", err) @@ -259,16 +278,11 @@ func (e *ExeSQLTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ if strings.TrimSpace(args.SQL) == "" { return exesqlErrorResult(errors.New("exesql: empty sql")), errors.New("exesql: empty sql") } - if !isSelectStatement(args.SQL) { + conn := e.conn + if err := validateExeSQLStatements(args.SQL, conn.DBType); err != nil { return exesqlErrorResult(ErrExeSQLNotSelect), ErrExeSQLNotSelect } - // Honor the per-call `database` override if the model supplied one; - // fall back to the tool's configured DB. - conn := e.conn - if args.Database != "" { - conn.Database = args.Database - } if err := conn.check(); err != nil { return exesqlErrorResult(err), err } @@ -309,20 +323,20 @@ func (e *ExeSQLTool) InvokableRun(ctx context.Context, argumentsInJSON string, _ fmt.Errorf("exesql: ping: %w", err) } - res, err := exesqlExecute(ctx, db, args.SQL, conn.MaxRecords) + res, err := exesqlExecute(ctx, db, args.SQL, conn.DBType, conn.MaxRecords) if err != nil { return exesqlErrorResult(err), err } return exesqlMarshalResult(res) } -// exesqlExecute splits the SQL on `;` (Python parity) and runs each -// statement independently. A failing statement is recorded as an +// exesqlExecute splits the SQL on statement-delimiting semicolons and runs +// each statement independently. A failing statement is recorded as an // error entry but does not abort subsequent statements — this is // the same isolation guarantee the Python tool provides so that // earlier results survive a bad statement later in the batch. -func exesqlExecute(ctx context.Context, db *sql.DB, sqlText string, maxRows int) (*exesqlResult, error) { - stmts := splitSQLStatements(sqlText) +func exesqlExecute(ctx context.Context, db *sql.DB, sqlText, dbType string, maxRows int) (*exesqlResult, error) { + stmts := splitSQLStatements(sqlText, dbType) res := &exesqlResult{} for _, stmt := range stmts { stmt = stripChunkIDMarkers(stmt) @@ -431,13 +445,46 @@ func isBadFloat(f float64) bool { return false } -// splitSQLStatements splits on `;`, ignoring semicolons inside string -// literals and line/block comments. This matches what the Python -// tool does with `sqls = sql.split(";")` — a naive split, but safe -// enough for read-only statements the LLM is expected to produce. -func splitSQLStatements(s string) []string { - cleaned := stripSQLStrings(stripSQLComments(s)) - return strings.Split(cleaned, ";") +// splitSQLStatements preserves the original SQL and splits only on delimiters +// outside strings, quoted identifiers, and comments. SQL quoting differs by +// database, so the scanner enables backslash escapes, dollar quotes, bracketed +// identifiers, and # comments only for dialects that support them. +func splitSQLStatements(s, dbType string) []string { + masked, _ := maskSQLLiteralsAndComments(s, dbType) + statements := make([]string, 0, strings.Count(masked, ";")+1) + start := 0 + for i := range masked { + if masked[i] != ';' { + continue + } + statements = append(statements, s[start:i]) + start = i + 1 + } + return append(statements, s[start:]) +} + +// validateExeSQLStatements rejects the entire batch before any database work. +// Each executable fragment is checked independently so a read-only first +// statement cannot hide a later write statement. +func validateExeSQLStatements(sqlText, dbType string) error { + if _, executableComment := maskSQLLiteralsAndComments(sqlText, dbType); executableComment { + return ErrExeSQLNotSelect + } + hasStatement := false + for _, stmt := range splitSQLStatements(sqlText, dbType) { + stmt = strings.TrimSpace(stripChunkIDMarkers(stmt)) + if stmt == "" { + continue + } + hasStatement = true + if !isReadOnlySQLStatement(stmt, dbType) { + return ErrExeSQLNotSelect + } + } + if !hasStatement { + return ErrExeSQLNotSelect + } + return nil } // stripChunkIDMarkers drops the [ID:123] tokens the RAGFlow chunker @@ -539,29 +586,21 @@ func (c exesqlConnParams) check() error { return nil } -// --------------------------------------------------------------------------- -// SELECT-only safety validator -// --------------------------------------------------------------------------- - -// leadingKeywordRe matches the first non-comment, non-whitespace keyword -// in a SQL statement. Comments (-- line, /* block */) and string literals -// are stripped before the match. -var leadingKeywordRe = regexp.MustCompile(`^[\s,;(]*([A-Za-z]+)`) - -// nonSelectKeywords lists DML/DDL/DCL verbs the parser rejects. These -// are the only top-level forms we refuse; everything else (WITH ... SELECT, -// SELECT INTO, SHOW, DESCRIBE, EXPLAIN) is allowed because they're -// read-only. -var nonSelectKeywords = map[string]struct{}{ +// writeCapableSQLKeywords are rejected anywhere in SELECT, WITH, and EXPLAIN +// statements. Checking every executable token prevents data-modifying CTEs +// such as `WITH t AS (DELETE ... RETURNING *) SELECT ...` from passing merely +// because their first keyword is WITH. +var writeCapableSQLKeywords = map[string]struct{}{ "INSERT": {}, "UPDATE": {}, "DELETE": {}, "REPLACE": {}, + "MERGE": {}, "TRUNCATE": {}, "CREATE": {}, "DROP": {}, "ALTER": {}, "RENAME": {}, "GRANT": {}, "REVOKE": {}, "LOCK": {}, "UNLOCK": {}, "CALL": {}, "EXEC": {}, "EXECUTE": {}, "COPY": {}, - "VACUUM": {}, "ANALYZE": {}, - "SET": {}, "RESET": {}, + "VACUUM": {}, + "SET": {}, "RESET": {}, "USE": {}, "KILL": {}, "LOAD": {}, @@ -570,90 +609,192 @@ var nonSelectKeywords = map[string]struct{}{ "SHUTDOWN": {}, } -// isSelectStatement returns true iff sql is a read-only statement. The -// heuristic is intentionally narrow: strip line + block comments and -// string literals, scan the first keyword, and reject if it's a -// DML/DDL/DCL verb. SQL parsers in Go stdlib don't exist; this matches -// the safety bar the Go shell needs. -func isSelectStatement(sql string) bool { - cleaned := stripSQLComments(sql) - cleaned = stripSQLStrings(cleaned) - m := leadingKeywordRe.FindStringSubmatch(cleaned) - if len(m) < 2 { +// isReadOnlySQLStatement validates executable tokens rather than only the +// leading keyword. Quoted text and comments are masked first, so write verbs +// in data values do not cause false rejections. +func isReadOnlySQLStatement(sql, dbType string) bool { + masked, executableComment := maskSQLLiteralsAndComments(sql, dbType) + if executableComment { return false } - kw := strings.ToUpper(m[1]) - if _, bad := nonSelectKeywords[kw]; bad { + words := sqlKeywordRe.FindAllString(masked, -1) + if len(words) == 0 { return false } - switch kw { - case "SELECT", "WITH", "SHOW", "DESCRIBE", "DESC", "EXPLAIN", "PRAGMA": + for i := range words { + words[i] = strings.ToUpper(words[i]) + } + + switch words[0] { + case "SHOW", "DESCRIBE", "DESC", "PRAGMA": + return true + case "SELECT", "WITH", "EXPLAIN": + for _, word := range words[1:] { + if _, bad := writeCapableSQLKeywords[word]; bad { + return false + } + switch word { + case "INTO", "OUTFILE", "DUMPFILE": + return false + } + } return true } - // Unknown verb → conservative reject. The Python tool would forward - // this; the Go shell declines to execute without a recognized form. return false } -// stripSQLComments removes -- line comments and /* ... */ block comments. -// We don't try to handle nested comments (MySQL/PG/SQLite differ) — this -// is a best-effort guard for the SELECT validator, not a SQL parser. -func stripSQLComments(s string) string { - var b strings.Builder - b.Grow(len(s)) - i := 0 - for i < len(s) { - if i+1 < len(s) && s[i] == '-' && s[i+1] == '-' { - for i < len(s) && s[i] != '\n' { - i++ +var sqlKeywordRe = regexp.MustCompile(`[A-Za-z_][A-Za-z0-9_$]*`) + +// maskSQLLiteralsAndComments replaces non-executable SQL regions with spaces +// without changing byte offsets. The returned flag identifies MySQL/MariaDB +// executable comments, which are rejected rather than mistaken for comments. +func maskSQLLiteralsAndComments(s, dbType string) (string, bool) { + masked := []byte(s) + dialect := strings.ToLower(dbType) + isMySQL := dialect == "mysql" || dialect == "mariadb" || dialect == "oceanbase" + isPostgres := dialect == "postgres" || dialect == "postgresql" + isMSSQL := dialect == "mssql" || dialect == "sqlserver" + executableComment := false + + mask := func(start, end int) { + for i := start; i < end; i++ { + masked[i] = ' ' + } + } + + for i := 0; i < len(s); { + if i+1 < len(s) && s[i] == '-' && s[i+1] == '-' && + (!isMySQL || i+2 == len(s) || s[i+2] <= ' ') { + end := i + 2 + for end < len(s) && s[end] != '\n' && s[end] != '\r' { + end++ } + mask(i, end) + i = end + continue + } + if isMySQL && s[i] == '#' { + end := i + 1 + for end < len(s) && s[end] != '\n' && s[end] != '\r' { + end++ + } + mask(i, end) + i = end continue } if i+1 < len(s) && s[i] == '/' && s[i+1] == '*' { - i += 2 - for i+1 < len(s) && !(s[i] == '*' && s[i+1] == '/') { - i++ + if isMySQL && (i+2 < len(s) && s[i+2] == '!' || + i+3 < len(s) && s[i+2] == 'M' && s[i+3] == '!') { + executableComment = true } - i += 2 + end := scanSQLBlockComment(s, i, isPostgres) + mask(i, end) + i = end continue } - b.WriteByte(s[i]) + if s[i] == '\'' || s[i] == '"' || s[i] == '`' { + backslashEscapes := isMySQL || isPostgres && s[i] == '\'' && isPostgresEscapeString(s, i) + end := scanSQLQuotedText(s, i, s[i], backslashEscapes) + mask(i, end) + i = end + continue + } + if isMSSQL && s[i] == '[' { + end := scanSQLQuotedText(s, i, ']', false) + mask(i, end) + i = end + continue + } + if isPostgres && s[i] == '$' { + if delimiter := postgresDollarQuoteDelimiter(s, i); delimiter != "" { + end := i + len(delimiter) + if closeAt := strings.Index(s[end:], delimiter); closeAt >= 0 { + end += closeAt + len(delimiter) + } else { + end = len(s) + } + mask(i, end) + i = end + continue + } + } i++ } - return b.String() + return string(masked), executableComment } -// stripSQLStrings removes single- and double-quoted string literals and -// replaces them with empty placeholders so that keywords inside string -// contents don't confuse the validator. -func stripSQLStrings(s string) string { - var b strings.Builder - b.Grow(len(s)) - i := 0 - inStr := byte(0) - for i < len(s) { - c := s[i] - if inStr != 0 { - if c == inStr { - if i+1 < len(s) && s[i+1] == inStr { - b.WriteByte(' ') - i += 2 - continue - } - inStr = 0 +func scanSQLBlockComment(s string, start int, nested bool) int { + depth := 1 + for i := start + 2; i < len(s); { + if nested && i+1 < len(s) && s[i] == '/' && s[i+1] == '*' { + depth++ + i += 2 + continue + } + if i+1 < len(s) && s[i] == '*' && s[i+1] == '/' { + depth-- + i += 2 + if depth == 0 { + return i } - b.WriteByte(' ') - i++ continue } - if c == '\'' || c == '"' || c == '`' { - inStr = c - b.WriteByte(' ') - i++ - continue - } - b.WriteByte(c) i++ } - return b.String() + return len(s) +} + +func scanSQLQuotedText(s string, start int, quote byte, backslashEscapes bool) int { + for i := start + 1; i < len(s); { + if backslashEscapes && s[i] == '\\' && i+1 < len(s) { + i += 2 + continue + } + if s[i] != quote { + i++ + continue + } + if i+1 < len(s) && s[i+1] == quote { + i += 2 + continue + } + return i + 1 + } + return len(s) +} + +func isPostgresEscapeString(s string, quoteAt int) bool { + if quoteAt == 0 || s[quoteAt-1] != 'E' && s[quoteAt-1] != 'e' { + return false + } + return quoteAt == 1 || !isSQLIdentifierByte(s[quoteAt-2]) +} + +func postgresDollarQuoteDelimiter(s string, start int) string { + if start > 0 && isSQLIdentifierByte(s[start-1]) { + return "" + } + for i := start + 1; i < len(s); i++ { + if s[i] == '$' { + return s[start : i+1] + } + if i == start+1 { + if !isSQLIdentifierStartByte(s[i]) { + return "" + } + continue + } + if !isSQLIdentifierByte(s[i]) { + return "" + } + } + return "" +} + +func isSQLIdentifierStartByte(b byte) bool { + return b == '_' || b >= 'A' && b <= 'Z' || b >= 'a' && b <= 'z' +} + +func isSQLIdentifierByte(b byte) bool { + return isSQLIdentifierStartByte(b) || b >= '0' && b <= '9' || b == '$' } diff --git a/internal/agent/tool/exesql_decoder_test.go b/internal/agent/tool/exesql_decoder_test.go index 2e678dc403..6fe20e83dc 100644 --- a/internal/agent/tool/exesql_decoder_test.go +++ b/internal/agent/tool/exesql_decoder_test.go @@ -17,6 +17,7 @@ package tool import ( + "encoding/json" "testing" ) @@ -46,6 +47,7 @@ func TestNewExeSQLConnParams_RequiredFields(t *testing.T) { // TestNewExeSQLConnParams_AllFields: full params map decodes correctly. func TestNewExeSQLConnParams_AllFields(t *testing.T) { conn, err := NewExeSQLConnParams(map[string]any{ + "sql": "SELECT 1", "db_type": "postgres", "host": "db.example.com", "port": 5432, @@ -57,6 +59,9 @@ func TestNewExeSQLConnParams_AllFields(t *testing.T) { if err != nil { t.Fatalf("NewExeSQLConnParams: %v", err) } + if conn.SQL != "SELECT 1" { + t.Errorf("SQL=%q, want SELECT 1", conn.SQL) + } if conn.DBType != "postgres" { t.Errorf("DBType=%q, want postgres", conn.DBType) } @@ -80,6 +85,61 @@ func TestNewExeSQLConnParams_AllFields(t *testing.T) { } } +func TestNewExeSQLConnParams_PythonDefaults(t *testing.T) { + conn, err := NewExeSQLConnParams(map[string]any{ + "host": "db.example.com", + "database": "demo", + "username": "ragflow", + }) + if err != nil { + t.Fatalf("NewExeSQLConnParams: %v", err) + } + if conn.SQL != "{sys.query}" { + t.Errorf("SQL=%q, want {sys.query}", conn.SQL) + } + if conn.DBType != "mysql" { + t.Errorf("DBType=%q, want mysql", conn.DBType) + } + if conn.Port != 3306 { + t.Errorf("Port=%d, want 3306", conn.Port) + } + if conn.MaxRecords != 1024 { + t.Errorf("MaxRecords=%d, want 1024", conn.MaxRecords) + } + if conn.Database != "demo" || conn.Username != "ragflow" || conn.Host != "db.example.com" || conn.Password != "" { + t.Errorf("unexpected connection fields: %+v", conn) + } +} + +func TestExeSQLConnParams_JSONTags(t *testing.T) { + raw, err := json.Marshal(exesqlConnParams{ + SQL: "SELECT 1", + DBType: "mysql", + Database: "demo", + Username: "ragflow", + Host: "db.example.com", + Port: 3306, + Password: "secret", + MaxRecords: 1024, + }) + if err != nil { + t.Fatalf("json.Marshal: %v", err) + } + var fields map[string]any + if err := json.Unmarshal(raw, &fields); err != nil { + t.Fatalf("json.Unmarshal: %v", err) + } + want := []string{"sql", "db_type", "database", "username", "host", "port", "password", "max_records"} + if len(fields) != len(want) { + t.Fatalf("JSON fields = %v", fields) + } + for _, field := range want { + if _, ok := fields[field]; !ok { + t.Errorf("JSON missing field %q: %s", field, raw) + } + } +} + // TestExeSQLConnParams_Alias: the public type alias ExeSQLConnParams // refers to the same underlying type as the lowercase exesqlConnParams. // The factory returns the public name, and existing in-package diff --git a/internal/agent/tool/exesql_test.go b/internal/agent/tool/exesql_test.go index 378dda5082..7890d86b38 100644 --- a/internal/agent/tool/exesql_test.go +++ b/internal/agent/tool/exesql_test.go @@ -97,12 +97,23 @@ func TestExeSQL_RejectsNonSelect(t *testing.T) { {"kill", `KILL 1234`}, {"use", `USE rag_flow`}, {"uppercase drop", `DROP DATABASE rag_flow`}, + {"select into", `SELECT * INTO archived_users FROM users`}, + {"select outfile", `SELECT * FROM users INTO OUTFILE '/tmp/users'`}, + {"select dumpfile", `SELECT payload FROM files INTO DUMPFILE '/tmp/payload'`}, + {"delete cte", `WITH deleted AS (DELETE FROM users RETURNING *) SELECT * FROM deleted`}, + {"update cte", `WITH changed AS (UPDATE users SET active = false RETURNING *) SELECT * FROM changed`}, + {"insert cte", `WITH added AS (INSERT INTO users(name) VALUES ('alice') RETURNING *) SELECT * FROM added`}, + {"merge cte", `WITH changed AS (MERGE INTO users USING incoming ON users.id = incoming.id WHEN MATCHED THEN UPDATE SET name = incoming.name RETURNING *) SELECT * FROM changed`}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { t.Parallel() - e := NewExeSQLTool(testConn()) + e := NewExeSQLTool(testConn()). + WithExeSQLDialer(func(_, _ string) (*sql.DB, error) { + t.Fatal("dialer called for rejected SQL") + return nil, nil + }) _, err := e.InvokableRun(context.Background(), `{"sql":`+jsonString(c.sql)+`}`) if !errors.Is(err, ErrExeSQLNotSelect) { @@ -112,6 +123,19 @@ func TestExeSQL_RejectsNonSelect(t *testing.T) { } } +func TestExeSQL_RejectsMixedBatchBeforeDatabaseAccess(t *testing.T) { + e := NewExeSQLTool(testConn()). + WithExeSQLDialer(func(_, _ string) (*sql.DB, error) { + t.Fatal("dialer called before every SQL statement was validated") + return nil, nil + }) + + _, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1; DROP TABLE users"}`) + if !errors.Is(err, ErrExeSQLNotSelect) { + t.Fatalf("err = %v, want ErrExeSQLNotSelect", err) + } +} + func TestExeSQL_AllowsSelect(t *testing.T) { t.Parallel() @@ -120,7 +144,7 @@ func TestExeSQL_AllowsSelect(t *testing.T) { `select * from t`, ` SELECT * FROM t WHERE a = 1`, `WITH cte AS (SELECT 1) SELECT * FROM cte`, - `SELECT * FROM t INTO OUTFILE '/tmp/x'`, + `WITH cte AS (SELECT 'DELETE; UPDATE' AS note) SELECT * FROM cte`, `SHOW TABLES`, `DESCRIBE t`, `EXPLAIN SELECT * FROM t`, @@ -189,6 +213,96 @@ func TestExeSQL_RejectsEmptyArgs(t *testing.T) { } } +func TestSplitSQLStatementsIgnoresQuotedAndCommentedSemicolons(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + dbType string + sql string + }{ + {"single quoted string", "mysql", `SELECT 'hello; world'; SELECT 2`}, + {"doubled single quote", "postgres", `SELECT 'it''s; intact'; SELECT 2`}, + {"mysql backslash escape", "mysql", `SELECT 'it\'; is intact'; SELECT 2`}, + {"double quoted identifier", "postgres", `SELECT "column;name" FROM t; SELECT 2`}, + {"backtick identifier", "mysql", "SELECT `column;name` FROM t; SELECT 2"}, + {"bracketed identifier", "mssql", `SELECT [column;name] FROM t; SELECT 2`}, + {"line comment", "postgres", "SELECT 1 -- ignored; delimiter\n; SELECT 2"}, + {"mysql hash comment", "mysql", "SELECT 1 # ignored; delimiter\n; SELECT 2"}, + {"block comment", "mysql", `SELECT 1 /* ignored; delimiter */; SELECT 2`}, + {"dollar quoted string", "postgres", `SELECT $$hello; world$$; SELECT 2`}, + {"tagged dollar quoted string", "postgres", `SELECT $body$hello; world$body$; SELECT 2`}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + statements := splitSQLStatements(tc.sql, tc.dbType) + if len(statements) != 2 { + t.Fatalf("splitSQLStatements(%q) = %#v, want 2 statements", tc.sql, statements) + } + if !strings.Contains(statements[0], ";") { + t.Fatalf("first statement = %q, want the quoted/commented semicolon preserved", statements[0]) + } + if statements[1] != " SELECT 2" { + t.Fatalf("second statement = %q, want %q", statements[1], " SELECT 2") + } + }) + } +} + +func TestExeSQL_ReadOnlyValidationIgnoresQuotedAndCommentedKeywords(t *testing.T) { + t.Parallel() + + cases := []struct { + dbType string + sql string + }{ + {"mysql", `SELECT 'DELETE; DROP TABLE users' AS note`}, + {"postgres", `WITH note AS (SELECT $$UPDATE users; DELETE FROM users$$ AS value) SELECT * FROM note`}, + {"postgres", "WITH note AS (SELECT 1 /* DELETE FROM users; */) SELECT * FROM note"}, + } + for _, tc := range cases { + if err := validateExeSQLStatements(tc.sql, tc.dbType); err != nil { + t.Errorf("validateExeSQLStatements(%q, %q) = %v, want nil", tc.sql, tc.dbType, err) + } + } +} + +func TestExeSQL_ExecutesStatementsWithQuotedSemicolonsIntact(t *testing.T) { + t.Parallel() + + dialer, mock, cleanup := sqlmockDialer(t) + defer cleanup() + mock.ExpectPing() + mock.ExpectQuery("SELECT 'hello; world'"). + WillReturnRows(sqlmock.NewRows([]string{"value"}).AddRow("hello; world")) + mock.ExpectQuery("SELECT 2"). + WillReturnRows(sqlmock.NewRows([]string{"value"}).AddRow(2)) + + e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer) + if _, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 'hello; world'; SELECT 2"}`); err != nil { + t.Fatalf("InvokableRun: %v", err) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatalf("SQL statements were not executed intact: %v", err) + } +} + +func TestExeSQL_RejectsMySQLExecutableComment(t *testing.T) { + t.Parallel() + + e := NewExeSQLTool(testConn()). + WithExeSQLDialer(func(_, _ string) (*sql.DB, error) { + t.Fatal("dialer called for an executable comment") + return nil, nil + }) + _, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1 /*!; DROP TABLE users */"}`) + if !errors.Is(err, ErrExeSQLNotSelect) { + t.Fatalf("err = %v, want ErrExeSQLNotSelect", err) + } +} + func TestExeSQL_Info(t *testing.T) { t.Parallel() @@ -200,6 +314,47 @@ func TestExeSQL_Info(t *testing.T) { if info.Name != "execute_sql" { t.Errorf("Name = %q, want execute_sql", info.Name) } + paramsSchema, err := info.ParamsOneOf.ToJSONSchema() + if err != nil { + t.Fatalf("ToJSONSchema: %v", err) + } + rawSchema, err := json.Marshal(paramsSchema) + if err != nil { + t.Fatalf("marshal params schema: %v", err) + } + params := string(rawSchema) + if !strings.Contains(params, `"sql"`) { + t.Fatalf("schema missing sql: %s", params) + } + if strings.Contains(params, `"database"`) { + t.Fatalf("schema leaked node-level database param: %s", params) + } + if !strings.Contains(params, `"required":["sql"]`) { + t.Fatalf("schema does not require sql: %s", params) + } +} + +func TestExeSQL_UsesConfiguredSQLDefault(t *testing.T) { + dialer, mock, cleanup := sqlmockDialer(t) + defer cleanup() + mock.ExpectPing() + mock.ExpectQuery("SELECT 1").WillReturnRows( + sqlmock.NewRows([]string{"value"}).AddRow(1), + ) + conn := testConn() + conn.SQL = "SELECT 1" + e := NewExeSQLTool(conn).WithExeSQLDialer(dialer) + + out, err := e.InvokableRun(context.Background(), `{}`) + if err != nil { + t.Fatalf("InvokableRun: %v", err) + } + if !strings.Contains(out, `"value":1`) { + t.Fatalf("output = %s, want configured SQL result", out) + } + if err := mock.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet SQL expectations: %v", err) + } } func TestExeSQL_ExecuteSelect_ReturnsRows(t *testing.T) { diff --git a/internal/agent/tool/registry.go b/internal/agent/tool/registry.go index c12485c66b..7ece5ff0ca 100644 --- a/internal/agent/tool/registry.go +++ b/internal/agent/tool/registry.go @@ -379,27 +379,9 @@ func decodeExeSQLConnParams(params map[string]any) (exesqlConnParams, error) { "(db_type/host/port/database/username/password)", ) } - conn := exesqlConnParams{} - if v, ok := stringParam(params, "db_type"); ok { - conn.DBType = v - } - if v, ok := stringParam(params, "database"); ok { - conn.Database = v - } - if v, ok := stringParam(params, "username"); ok { - conn.Username = v - } - if v, ok := stringParam(params, "host"); ok { - conn.Host = v - } - if v, ok := intParam(params, "port"); ok { - conn.Port = v - } - if v, ok := stringParam(params, "password"); ok { - conn.Password = v - } - if v, ok := intParam(params, "max_records"); ok { - conn.MaxRecords = v + conn, err := NewExeSQLConnParams(params) + if err != nil { + return exesqlConnParams{}, fmt.Errorf("agent tool: execute_sql config: %w", err) } if err := conn.check(); err != nil { return exesqlConnParams{}, fmt.Errorf("agent tool: execute_sql config: %w", err)