From 721caccc58a9385c510f81b826f078dc334bf3b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=90=B4=E6=96=87=E5=B3=B0?= Date: Tue, 9 Jun 2026 19:37:26 +0800 Subject: [PATCH] up --- agents/SQL_query/sql_query.go | 763 +++++++++++++++++++++++++++++ agents/SQL_query/sql_query_test.go | 36 ++ go.mod | 13 +- go.sum | 55 ++- main.go | 164 ++++++- 5 files changed, 1025 insertions(+), 6 deletions(-) create mode 100644 agents/SQL_query/sql_query.go create mode 100644 agents/SQL_query/sql_query_test.go diff --git a/agents/SQL_query/sql_query.go b/agents/SQL_query/sql_query.go new file mode 100644 index 0000000..49d1d82 --- /dev/null +++ b/agents/SQL_query/sql_query.go @@ -0,0 +1,763 @@ +package sqlquery + +import ( + "context" + "database/sql" + "errors" + "fmt" + "os" + "path/filepath" + "regexp" + "sort" + "strings" + "sync" + "time" + "unicode/utf8" + + _ "github.com/go-sql-driver/mysql" + "gopkg.in/yaml.v3" + _ "modernc.org/sqlite" +) + +const ( + defaultActivationPrompt = `判断用户问题是否需要查询业务数据库。 +仅当用户询问数据库表、记录、字段、时间、状态、内容、统计、最近/最早/某时间范围内的数据时返回 activate=true。 +普通知识问答、代码问题、闲聊、联网搜索问题返回 activate=false。 +只返回 JSON: {"activate": true/false, "reason": "..."}` + defaultDatabaseName = "default" + defaultSQLiteDSN = "file:data/app.db?mode=ro" + defaultTimeout = 10 + defaultMaxRows = 50 + defaultMaxCellBytes = 4096 + defaultSchemaCacheSecond = 300 +) + +type Config struct { + Enabled bool `yaml:"enabled" json:"enabled"` + ActivationPrompt string `yaml:"activation_prompt" json:"activation_prompt"` + DefaultDatabase string `yaml:"default_database" json:"default_database"` + SchemaCacheSeconds int `yaml:"schema_cache_seconds" json:"schema_cache_seconds"` + Databases []DatabaseConfig `yaml:"databases" json:"databases"` +} + +type DatabaseConfig struct { + Name string `yaml:"name" json:"name"` + Active bool `yaml:"active" json:"active"` + Driver string `yaml:"driver" json:"driver"` + DSN string `yaml:"dsn" json:"-"` + Timeout int `yaml:"timeout" json:"timeout"` + MaxRows int `yaml:"max_rows" json:"max_rows"` + MaxCellBytes int `yaml:"max_cell_bytes" json:"max_cell_bytes"` + Schema SchemaConfig `yaml:"schema" json:"schema"` +} + +type SchemaConfig struct { + IncludeTables []string `yaml:"include_tables" json:"include_tables"` + ExcludeTables []string `yaml:"exclude_tables" json:"exclude_tables"` +} + +type State struct { + cfg *Config + dbs map[string]*database + order []string + cacheMu sync.Mutex + cacheText string + cacheAt time.Time +} + +type database struct { + cfg DatabaseConfig + db *sql.DB +} + +type QueryResult struct { + Database string `json:"database"` + SQL string `json:"sql"` + Columns []string `json:"columns"` + Rows [][]string `json:"rows"` + Truncated bool `json:"truncated"` + MaxRows int `json:"max_rows"` +} + +func LoadConfig(path string) (*Config, error) { + if _, err := os.Stat(path); err != nil { + if !os.IsNotExist(err) { + return nil, fmt.Errorf("检查 SQL 查询插件配置失败: %w", err) + } + cfg := defaultConfig() + if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil { + return nil, fmt.Errorf("创建 SQL 查询插件目录失败: %w", err) + } + data, err := yaml.Marshal(&cfg) + if err != nil { + return nil, fmt.Errorf("生成 SQL 查询插件配置失败: %w", err) + } + if err := os.WriteFile(path, data, 0644); err != nil { + return nil, fmt.Errorf("写入 SQL 查询插件配置失败: %w", err) + } + } + + data, err := os.ReadFile(path) + if err != nil { + return nil, fmt.Errorf("读取 SQL 查询插件配置失败: %w", err) + } + var cfg Config + if err := yaml.Unmarshal(data, &cfg); err != nil { + return nil, fmt.Errorf("解析 SQL 查询插件配置失败: %w", err) + } + if err := normalizeConfig(&cfg); err != nil { + return nil, err + } + return &cfg, nil +} + +func NewState(cfg *Config) (*State, error) { + state := &State{cfg: cfg, dbs: map[string]*database{}} + if cfg == nil || !cfg.Enabled { + return state, nil + } + + for _, item := range cfg.Databases { + db, err := sql.Open(driverName(item.Driver), item.DSN) + if err != nil { + return nil, fmt.Errorf("打开数据库 %s 失败: %w", item.Name, err) + } + if item.Driver == "sqlite" { + db.SetMaxOpenConns(1) + db.SetMaxIdleConns(1) + } else { + db.SetMaxOpenConns(5) + db.SetMaxIdleConns(2) + } + db.SetConnMaxLifetime(30 * time.Minute) + + ctx, cancel := context.WithTimeout(context.Background(), time.Duration(item.Timeout)*time.Second) + if err := db.PingContext(ctx); err != nil { + cancel() + db.Close() + return nil, fmt.Errorf("连接数据库 %s 失败: %w", item.Name, err) + } + if item.Driver == "sqlite" { + if _, err := db.ExecContext(ctx, "PRAGMA query_only = ON"); err != nil { + cancel() + db.Close() + return nil, fmt.Errorf("启用 SQLite 只读模式失败: %w", err) + } + } + cancel() + + state.dbs[item.Name] = &database{cfg: item, db: db} + state.order = append(state.order, item.Name) + } + return state, nil +} + +func (s *State) Close() error { + if s == nil { + return nil + } + var errs []string + for _, db := range s.dbs { + if err := db.db.Close(); err != nil { + errs = append(errs, err.Error()) + } + } + if len(errs) > 0 { + return errors.New(strings.Join(errs, "; ")) + } + return nil +} + +func (s *State) Enabled() bool { + return s != nil && s.cfg != nil && s.cfg.Enabled && len(s.dbs) > 0 +} + +func (s *State) ActivationPrompt() string { + if s == nil || s.cfg == nil || strings.TrimSpace(s.cfg.ActivationPrompt) == "" { + return defaultActivationPrompt + } + return strings.TrimSpace(s.cfg.ActivationPrompt) +} + +func (s *State) DefaultDatabase() string { + if s == nil || s.cfg == nil || strings.TrimSpace(s.cfg.DefaultDatabase) == "" { + return defaultDatabaseName + } + return strings.TrimSpace(s.cfg.DefaultDatabase) +} + +func (s *State) SchemaContext(ctx context.Context) (string, error) { + if !s.Enabled() { + return "", errors.New("SQL 查询插件未启用") + } + + s.cacheMu.Lock() + if s.cacheText != "" && time.Since(s.cacheAt) < time.Duration(s.cfg.SchemaCacheSeconds)*time.Second { + text := s.cacheText + s.cacheMu.Unlock() + return text, nil + } + s.cacheMu.Unlock() + + var b strings.Builder + fmt.Fprintf(&b, "可查询数据库列表（只能生成 SELECT/WITH 查询）：\n") + for _, name := range s.order { + handle := s.dbs[name] + fmt.Fprintf(&b, "\n数据库 %s，类型 %s，单次最多返回 %d 行：\n", handle.cfg.Name, handle.cfg.Driver, handle.cfg.MaxRows) + schema, err := handle.schemaContext(ctx) + if err != nil { + return "", err + } + b.WriteString(schema) + } + text := b.String() + + s.cacheMu.Lock() + s.cacheText = text + s.cacheAt = time.Now() + s.cacheMu.Unlock() + return text, nil +} + +func (s *State) ExecuteReadOnly(ctx context.Context, databaseName string, query string) (*QueryResult, error) { + if !s.Enabled() { + return nil, errors.New("SQL 查询插件未启用") + } + if err := ValidateReadOnlySQL(query); err != nil { + return nil, err + } + + handle := s.databaseByName(databaseName) + if handle == nil { + return nil, fmt.Errorf("数据库配置不存在: %s", databaseName) + } + if err := handle.rejectExcludedTables(query); err != nil { + return nil, err + } + + timeout := time.Duration(handle.cfg.Timeout) * time.Second + queryCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + if handle.cfg.Driver == "mysql" { + tx, err := handle.db.BeginTx(queryCtx, &sql.TxOptions{ReadOnly: true}) + if err != nil { + return nil, fmt.Errorf("开启只读事务失败: %w", err) + } + defer tx.Rollback() + rows, err := tx.QueryContext(queryCtx, query) + if err != nil { + return nil, fmt.Errorf("执行 SQL 查询失败: %w", err) + } + result, err := scanRows(rows, handle.cfg, query) + if err != nil { + return nil, err + } + if err := tx.Commit(); err != nil { + return nil, fmt.Errorf("提交只读事务失败: %w", err) + } + return result, nil + } + + rows, err := handle.db.QueryContext(queryCtx, query) + if err != nil { + return nil, fmt.Errorf("执行 SQL 查询失败: %w", err) + } + return scanRows(rows, handle.cfg, query) +} + +func BuildResultContext(userQuery string, generatedSQL string, result *QueryResult) string { + var b strings.Builder + fmt.Fprintf(&b, "用户问题需要查询本地数据库。请仅根据以下 SQL 查询结果回答；不要编造结果中不存在的记录。\n") + fmt.Fprintf(&b, "用户问题：%s\n", userQuery) + fmt.Fprintf(&b, "数据库：%s\n", result.Database) + fmt.Fprintf(&b, "已执行只读 SQL：%s\n", generatedSQL) + if len(result.Columns) == 0 { + b.WriteString("查询没有返回列。\n") + return b.String() + } + if len(result.Rows) == 0 { + b.WriteString("查询结果：没有匹配记录。\n") + return b.String() + } + b.WriteString("查询结果：\n") + b.WriteString(markdownTable(result.Columns, result.Rows)) + if result.Truncated { + fmt.Fprintf(&b, "\n结果已按配置截断，只展示前 %d 行。\n", result.MaxRows) + } + return b.String() +} + +func BuildErrorContext(userQuery string, err error) string { + return fmt.Sprintf("用户问题可能需要查询本地数据库，但 SQL 查询插件执行失败：%s\n用户问题：%s\n请向用户说明无法完成数据库查询，不要编造数据库记录。", err.Error(), userQuery) +} + +func ValidateReadOnlySQL(query string) error { + trimmed := strings.TrimSpace(query) + if trimmed == "" { + return errors.New("SQL 不能为空") + } + if strings.Contains(trimmed, "--") || strings.Contains(trimmed, "/*") || strings.Contains(trimmed, "*/") { + return errors.New("SQL 不允许包含注释") + } + body := strings.TrimSuffix(trimmed, ";") + if strings.Contains(body, ";") { + return errors.New("SQL 只允许单条语句") + } + upper := strings.ToUpper(body) + first := firstToken(upper) + if first != "SELECT" && first != "WITH" { + return fmt.Errorf("SQL 只允许 SELECT/WITH 查询，当前为 %s", first) + } + + stripped := stripSingleQuotedStrings(upper) + forbidden := []string{ + "INSERT", "UPDATE", "DELETE", "DROP", "ALTER", "CREATE", "TRUNCATE", "REPLACE", "MERGE", + "GRANT", "REVOKE", "VACUUM", "ANALYZE", "ATTACH", "DETACH", "LOAD", "CALL", "EXEC", + "SET", "USE", "LOCK", "UNLOCK", "BEGIN", "COMMIT", "ROLLBACK", + } + for _, word := range forbidden { + if hasSQLWord(stripped, word) { + return fmt.Errorf("SQL 包含禁止关键词: %s", word) + } + } + risky := []string{"INTO OUTFILE", "INTO DUMPFILE", "LOAD_FILE", "SLEEP", "BENCHMARK", "LOAD_EXTENSION"} + for _, phrase := range risky { + if strings.Contains(stripped, phrase) { + return fmt.Errorf("SQL 包含禁止函数或语法: %s", phrase) + } + } + return nil +} + +func defaultConfig() Config { + return Config{ + Enabled: false, + ActivationPrompt: defaultActivationPrompt, + DefaultDatabase: defaultDatabaseName, + SchemaCacheSeconds: defaultSchemaCacheSecond, + Databases: []DatabaseConfig{{ + Name: defaultDatabaseName, + Active: true, + Driver: "sqlite", + DSN: defaultSQLiteDSN, + Timeout: defaultTimeout, + MaxRows: defaultMaxRows, + MaxCellBytes: defaultMaxCellBytes, + }}, + } +} + +func normalizeConfig(cfg *Config) error { + if strings.TrimSpace(cfg.ActivationPrompt) == "" { + cfg.ActivationPrompt = defaultActivationPrompt + } + if strings.TrimSpace(cfg.DefaultDatabase) == "" { + cfg.DefaultDatabase = defaultDatabaseName + } + if cfg.SchemaCacheSeconds <= 0 { + cfg.SchemaCacheSeconds = defaultSchemaCacheSecond + } + if len(cfg.Databases) == 0 { + cfg.Databases = defaultConfig().Databases + } + + seen := map[string]bool{} + activeIndex := -1 + for i := range cfg.Databases { + item := &cfg.Databases[i] + item.Name = strings.TrimSpace(item.Name) + if item.Name == "" { + item.Name = fmt.Sprintf("database-%d", i+1) + } + if seen[item.Name] { + return fmt.Errorf("SQL 查询插件数据库名称重复: %s", item.Name) + } + seen[item.Name] = true + + item.Driver = strings.ToLower(strings.TrimSpace(item.Driver)) + if item.Driver == "" { + item.Driver = "sqlite" + } + if item.Driver != "sqlite" && item.Driver != "mysql" { + return fmt.Errorf("SQL 查询插件暂不支持数据库类型: %s", item.Driver) + } + if strings.TrimSpace(item.DSN) == "" { + return fmt.Errorf("数据库 %s 缺少 dsn", item.Name) + } + if item.Timeout <= 0 { + item.Timeout = defaultTimeout + } + if item.MaxRows <= 0 { + item.MaxRows = defaultMaxRows + } + if item.MaxCellBytes <= 0 { + item.MaxCellBytes = defaultMaxCellBytes + } + item.Schema.IncludeTables = cleanList(item.Schema.IncludeTables) + item.Schema.ExcludeTables = cleanList(item.Schema.ExcludeTables) + if item.Active { + if activeIndex == -1 { + activeIndex = i + } else { + item.Active = false + } + } + } + if activeIndex == -1 { + cfg.Databases[0].Active = true + } + if !seen[cfg.DefaultDatabase] { + for _, item := range cfg.Databases { + if item.Active { + cfg.DefaultDatabase = item.Name + break + } + } + } + return nil +} + +func driverName(driver string) string { + if driver == "sqlite" { + return "sqlite" + } + return driver +} + +func (s *State) databaseByName(name string) *database { + name = strings.TrimSpace(name) + if name == "" { + name = s.DefaultDatabase() + } + if db := s.dbs[name]; db != nil { + return db + } + return s.dbs[s.DefaultDatabase()] +} + +func (d *database) schemaContext(ctx context.Context) (string, error) { + timeout := time.Duration(d.cfg.Timeout) * time.Second + schemaCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + if d.cfg.Driver == "mysql" { + return d.mysqlSchemaContext(schemaCtx) + } + return d.sqliteSchemaContext(schemaCtx) +} + +func (d *database) sqliteSchemaContext(ctx context.Context) (string, error) { + rows, err := d.db.QueryContext(ctx, `SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name`) + if err != nil { + return "", fmt.Errorf("读取 SQLite 表列表失败: %w", err) + } + defer rows.Close() + + var tables []string + for rows.Next() { + var name string + if err := rows.Scan(&name); err != nil { + return "", err + } + if d.tableAllowed(name) { + tables = append(tables, name) + } + } + if err := rows.Err(); err != nil { + return "", err + } + + var b strings.Builder + for _, table := range tables { + fmt.Fprintf(&b, "- 表 %s\n", table) + colRows, err := d.db.QueryContext(ctx, "PRAGMA table_info("+quoteSQLiteString(table)+")") + if err != nil { + return "", fmt.Errorf("读取 SQLite 表 %s 字段失败: %w", table, err) + } + for colRows.Next() { + var cid int + var name, typ string + var notNull int + var defaultValue any + var pk int + if err := colRows.Scan(&cid, &name, &typ, ¬Null, &defaultValue, &pk); err != nil { + colRows.Close() + return "", err + } + extra := "" + if pk > 0 { + extra = " primary_key" + } + fmt.Fprintf(&b, " - %s %s%s\n", name, typ, extra) + } + if err := colRows.Close(); err != nil { + return "", err + } + } + if len(tables) == 0 { + b.WriteString("- 没有可查询表，或表被 include/exclude 规则过滤。\n") + } + return b.String(), nil +} + +func (d *database) mysqlSchemaContext(ctx context.Context) (string, error) { + rows, err := d.db.QueryContext(ctx, `SELECT table_name FROM information_schema.tables WHERE table_schema = DATABASE() AND table_type = 'BASE TABLE' ORDER BY table_name`) + if err != nil { + return "", fmt.Errorf("读取 MySQL 表列表失败: %w", err) + } + defer rows.Close() + + var tables []string + for rows.Next() { + var name string + if err := rows.Scan(&name); err != nil { + return "", err + } + if d.tableAllowed(name) { + tables = append(tables, name) + } + } + if err := rows.Err(); err != nil { + return "", err + } + + var b strings.Builder + for _, table := range tables { + fmt.Fprintf(&b, "- 表 %s\n", table) + colRows, err := d.db.QueryContext(ctx, `SELECT column_name, data_type, is_nullable, column_key FROM information_schema.columns WHERE table_schema = DATABASE() AND table_name = ? ORDER BY ordinal_position`, table) + if err != nil { + return "", fmt.Errorf("读取 MySQL 表 %s 字段失败: %w", table, err) + } + for colRows.Next() { + var name, typ, nullable, key string + if err := colRows.Scan(&name, &typ, &nullable, &key); err != nil { + colRows.Close() + return "", err + } + extra := "" + if key == "PRI" { + extra = " primary_key" + } + if nullable == "NO" { + extra += " not_null" + } + fmt.Fprintf(&b, " - %s %s%s\n", name, typ, extra) + } + if err := colRows.Close(); err != nil { + return "", err + } + } + if len(tables) == 0 { + b.WriteString("- 没有可查询表，或表被 include/exclude 规则过滤。\n") + } + return b.String(), nil +} + +func (d *database) tableAllowed(table string) bool { + name := strings.ToLower(strings.TrimSpace(table)) + include := lowerSet(d.cfg.Schema.IncludeTables) + if len(include) > 0 && !include[name] { + return false + } + exclude := lowerSet(d.cfg.Schema.ExcludeTables) + return !exclude[name] +} + +func (d *database) rejectExcludedTables(query string) error { + cleaned := strings.ToLower(stripSingleQuotedStrings(query)) + include := lowerSet(d.cfg.Schema.IncludeTables) + if len(include) > 0 { + matched := false + for table := range include { + if hasSQLWord(cleaned, table) { + matched = true + break + } + } + if !matched { + return errors.New("SQL 未访问 include_tables 中允许的表") + } + } + exclude := lowerSet(d.cfg.Schema.ExcludeTables) + for table := range exclude { + if hasSQLWord(cleaned, table) { + return fmt.Errorf("SQL 访问了被排除的表: %s", table) + } + } + return nil +} + +func scanRows(rows *sql.Rows, cfg DatabaseConfig, query string) (*QueryResult, error) { + defer rows.Close() + columns, err := rows.Columns() + if err != nil { + return nil, fmt.Errorf("读取查询列失败: %w", err) + } + result := &QueryResult{ + Database: cfg.Name, + SQL: query, + Columns: columns, + MaxRows: cfg.MaxRows, + } + for rows.Next() { + values := make([]any, len(columns)) + ptrs := make([]any, len(columns)) + for i := range values { + ptrs[i] = &values[i] + } + if err := rows.Scan(ptrs...); err != nil { + return nil, fmt.Errorf("读取查询结果失败: %w", err) + } + if len(result.Rows) >= cfg.MaxRows { + result.Truncated = true + break + } + row := make([]string, len(columns)) + for i, value := range values { + row[i] = formatCell(value, cfg.MaxCellBytes) + } + result.Rows = append(result.Rows, row) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("读取查询结果失败: %w", err) + } + return result, nil +} + +func formatCell(value any, maxBytes int) string { + if value == nil { + return "NULL" + } + switch v := value.(type) { + case time.Time: + return v.Format(time.RFC3339) + case []byte: + if !utf8.Valid(v) { + return fmt.Sprintf("", len(v)) + } + return truncateString(string(v), maxBytes) + case string: + return truncateString(v, maxBytes) + default: + return truncateString(fmt.Sprint(v), maxBytes) + } +} + +func truncateString(s string, maxBytes int) string { + if maxBytes <= 0 || len(s) <= maxBytes { + return s + } + cut := s[:maxBytes] + for !utf8.ValidString(cut) && len(cut) > 0 { + cut = cut[:len(cut)-1] + } + return cut + "..." +} + +func markdownTable(columns []string, rows [][]string) string { + var b strings.Builder + b.WriteString("| ") + for i, col := range columns { + if i > 0 { + b.WriteString(" | ") + } + b.WriteString(escapeMarkdownCell(col)) + } + b.WriteString(" |\n| ") + for i := range columns { + if i > 0 { + b.WriteString(" | ") + } + b.WriteString("---") + } + b.WriteString(" |\n") + for _, row := range rows { + b.WriteString("| ") + for i, cell := range row { + if i > 0 { + b.WriteString(" | ") + } + b.WriteString(escapeMarkdownCell(cell)) + } + b.WriteString(" |\n") + } + return b.String() +} + +func escapeMarkdownCell(s string) string { + s = strings.ReplaceAll(s, "|", "\\|") + s = strings.ReplaceAll(s, "\r\n", " ") + s = strings.ReplaceAll(s, "\n", " ") + s = strings.ReplaceAll(s, "\r", " ") + return s +} + +func firstToken(s string) string { + fields := strings.Fields(s) + if len(fields) == 0 { + return "" + } + return strings.Trim(fields[0], "();") +} + +func stripSingleQuotedStrings(s string) string { + var b strings.Builder + inString := false + for i := 0; i < len(s); i++ { + ch := s[i] + if ch == '\'' { + if inString && i+1 < len(s) && s[i+1] == '\'' { + i++ + continue + } + inString = !inString + b.WriteByte(' ') + continue + } + if inString { + b.WriteByte(' ') + } else { + b.WriteByte(ch) + } + } + return b.String() +} + +func hasSQLWord(s string, word string) bool { + pattern := `(?i)(^|[^a-zA-Z0-9_])` + regexp.QuoteMeta(word) + `([^a-zA-Z0-9_]|$)` + return regexp.MustCompile(pattern).FindStringIndex(s) != nil +} + +func quoteSQLiteString(s string) string { + return "'" + strings.ReplaceAll(s, "'", "''") + "'" +} + +func lowerSet(items []string) map[string]bool { + set := map[string]bool{} + for _, item := range items { + item = strings.ToLower(strings.TrimSpace(item)) + if item != "" { + set[item] = true + } + } + return set +} + +func cleanList(items []string) []string { + seen := map[string]bool{} + var cleaned []string + for _, item := range items { + item = strings.TrimSpace(item) + if item == "" { + continue + } + key := strings.ToLower(item) + if seen[key] { + continue + } + seen[key] = true + cleaned = append(cleaned, item) + } + sort.Strings(cleaned) + return cleaned +} diff --git a/agents/SQL_query/sql_query_test.go b/agents/SQL_query/sql_query_test.go new file mode 100644 index 0000000..d1f4a03 --- /dev/null +++ b/agents/SQL_query/sql_query_test.go @@ -0,0 +1,36 @@ +package sqlquery + +import "testing" + +func TestValidateReadOnlySQLAllowsSelectAndWith(t *testing.T) { + queries := []string{ + "SELECT * FROM events LIMIT 10", + "select id, created_at from events where content = 'delete keyword in text' limit 5;", + "WITH recent AS (SELECT * FROM events LIMIT 10) SELECT * FROM recent", + } + for _, query := range queries { + if err := ValidateReadOnlySQL(query); err != nil { + t.Fatalf("ValidateReadOnlySQL(%q) returned error: %v", query, err) + } + } +} + +func TestValidateReadOnlySQLRejectsUnsafeStatements(t *testing.T) { + queries := []string{ + "", + "DELETE FROM events", + "UPDATE events SET content='x'", + "DROP TABLE events", + "SELECT * FROM events; DELETE FROM events", + "SELECT * INTO OUTFILE '/tmp/x' FROM events", + "SELECT SLEEP(10)", + "ATTACH DATABASE 'x' AS y", + "VACUUM", + "SELECT * FROM events -- comment", + } + for _, query := range queries { + if err := ValidateReadOnlySQL(query); err == nil { + t.Fatalf("ValidateReadOnlySQL(%q) returned nil, want error", query) + } + } +} diff --git a/go.mod b/go.mod index bcaef50..3ad2bfd 100644 --- a/go.mod +++ b/go.mod @@ -4,15 +4,19 @@ go 1.25.0 require ( github.com/gin-gonic/gin v1.12.0 + github.com/go-sql-driver/mysql v1.10.0 github.com/volcengine/volcengine-go-sdk v1.2.28 gopkg.in/yaml.v3 v3.0.1 + modernc.org/sqlite v1.52.0 ) require ( + filippo.io/edwards25519 v1.2.0 // indirect github.com/bytedance/gopkg v0.1.3 // indirect github.com/bytedance/sonic v1.15.0 // indirect github.com/bytedance/sonic/loader v0.5.0 // indirect github.com/cloudwego/base64x v0.1.6 // indirect + github.com/dustin/go-humanize v1.0.1 // indirect github.com/gabriel-vasile/mimetype v1.4.12 // indirect github.com/gin-contrib/sse v1.1.0 // indirect github.com/go-playground/locales v0.14.1 // indirect @@ -20,7 +24,7 @@ require ( github.com/go-playground/validator/v10 v10.30.1 // indirect github.com/goccy/go-json v0.10.5 // indirect github.com/goccy/go-yaml v1.19.2 // indirect - github.com/google/uuid v1.3.0 // indirect + github.com/google/uuid v1.6.0 // indirect github.com/jmespath/go-jmespath v0.4.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect @@ -28,9 +32,11 @@ require ( github.com/mattn/go-isatty v0.0.20 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect github.com/modern-go/reflect2 v1.0.2 // indirect + github.com/ncruces/go-strftime v1.0.0 // indirect github.com/pelletier/go-toml/v2 v2.2.4 // indirect github.com/quic-go/qpack v0.6.0 // indirect github.com/quic-go/quic-go v0.59.0 // indirect + github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/ugorji/go/codec v1.3.1 // indirect github.com/volcengine/volc-sdk-golang v1.0.23 // indirect @@ -38,8 +44,11 @@ require ( golang.org/x/arch v0.22.0 // indirect golang.org/x/crypto v0.48.0 // indirect golang.org/x/net v0.51.0 // indirect - golang.org/x/sys v0.41.0 // indirect + golang.org/x/sys v0.42.0 // indirect golang.org/x/text v0.34.0 // indirect google.golang.org/protobuf v1.36.10 // indirect gopkg.in/yaml.v2 v2.2.8 // indirect + modernc.org/libc v1.72.3 // indirect + modernc.org/mathutil v1.7.1 // indirect + modernc.org/memory v1.11.0 // indirect ) diff --git a/go.sum b/go.sum index 0905cb7..4dfa0fe 100644 --- a/go.sum +++ b/go.sum @@ -1,4 +1,6 @@ cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= +filippo.io/edwards25519 v1.2.0 h1:crnVqOiS4jqYleHd9vaKZ+HKtHfllngJIiOpNpoJsjo= +filippo.io/edwards25519 v1.2.0/go.mod h1:xzAOLCNug/yB62zG1bQ8uziwrIqIuxhctzJT18Q77mc= github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/avast/retry-go v3.0.0+incompatible/go.mod h1:XtSnn+n/sHqQIpZ10K1qAevBhOOCWBLXXy3hyiqqBrY= github.com/bytedance/gopkg v0.1.3 h1:TPBSwH8RsouGCBcMBktLt1AymVo2TVsBVCY4b6TnZ/M= @@ -14,6 +16,8 @@ github.com/cloudwego/base64x v0.1.6/go.mod h1:OFcloc187FXDaYHvrNIjxSe8ncn0OOM8gE github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY= +github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto= github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= github.com/gabriel-vasile/mimetype v1.4.12 h1:e9hWvmLYvtp846tLHam2o++qitpguFiYCKbn0w9jyqw= @@ -30,6 +34,8 @@ github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJn github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY= github.com/go-playground/validator/v10 v10.30.1 h1:f3zDSN/zOma+w6+1Wswgd9fLkdwy06ntQJp0BBvFG0w= github.com/go-playground/validator/v10 v10.30.1/go.mod h1:oSuBIQzuJxL//3MelwSLD5hc2Tu889bF0Idm9Dg26cM= +github.com/go-sql-driver/mysql v1.10.0 h1:Q+1LV8DkHJvSYAdR83XzuhDaTykuDx0l6fkXxoWCWfw= +github.com/go-sql-driver/mysql v1.10.0/go.mod h1:M+cqaI7+xxXGG9swrdeUIoPG3Y3KCkF0pZej+SK+nWk= github.com/goccy/go-json v0.10.5 h1:Fq85nIqj+gXn/S5ahsiTlK3TmC85qgirsdTP/+DeaC4= github.com/goccy/go-json v0.10.5/go.mod h1:oq7eo15ShAhp70Anwd5lgX2pLfOS3QCiwU/PULtXL6M= github.com/goccy/go-yaml v1.19.2 h1:PmFC1S6h8ljIz6gMRBopkjP1TVT7xuwrButHID66PoM= @@ -55,8 +61,13 @@ github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/ github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg= -github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e h1:ijClszYn+mADRFY17kjQEVQ1XRhq2/JR1M3sGqeJoxs= +github.com/google/pprof v0.0.0-20250317173921-a4b03ec1a45e/go.mod h1:boTsfXsheKC2y+lKOCMpSfarhxDeIzfZG1jqGcPl3cA= github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= +github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg= github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo= github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8= @@ -81,6 +92,8 @@ github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd h1:TRLaZ9cD/w github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= github.com/modern-go/reflect2 v1.0.2 h1:xBagoLtFs94CBntxluKeaWgTMpvLxC4ur3nMaC9Gz0M= github.com/modern-go/reflect2 v1.0.2/go.mod h1:yWuevngMOJpCy52FWWMvUC8ws7m/LJsjYzDa0/r8luk= +github.com/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w= +github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls= github.com/pelletier/go-toml/v2 v2.2.4 h1:mye9XuhQ6gvn5h28+VilKrrPoQVanw5PMw/TB0t5Ec4= github.com/pelletier/go-toml/v2 v2.2.4/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= @@ -91,6 +104,8 @@ github.com/quic-go/qpack v0.6.0 h1:g7W+BMYynC1LbYLSqRt8PBg5Tgwxn214ZZR34VIOjz8= github.com/quic-go/qpack v0.6.0/go.mod h1:lUpLKChi8njB4ty2bFLX2x4gzDqXwUpaO1DP9qMDZII= github.com/quic-go/quic-go v0.59.0 h1:OLJkp1Mlm/aS7dpKgTc6cnpynnD2Xg7C1pwL6vy/SAw= github.com/quic-go/quic-go v0.59.0/go.mod h1:upnsH4Ju1YkqpLXC305eW3yDZ4NfnNbmQRCMWS58IKU= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= +github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ= github.com/rogpeppe/go-internal v1.10.0/go.mod h1:UQnix2H7Ngw/k4C5ijL5+65zddjncjaFoBhdsK/akog= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= @@ -126,6 +141,8 @@ golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU= golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= +golang.org/x/mod v0.33.0 h1:tHFzIWbBifEmbwtGz65eaWyGiGZatSrT9prnU8DbVL8= +golang.org/x/mod v0.33.0/go.mod h1:swjeQEj+6r7fODbD2cqrnje9PnziFuw4bmLbBZFrQ5w= golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= @@ -136,11 +153,13 @@ golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAG golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k= -golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/sys v0.42.0 h1:omrd2nAlyT5ESRdCLYdm3+fMfNFE/+Rf4bDIQImRJeo= +golang.org/x/sys v0.42.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.34.0 h1:oL/Qq0Kdaqxa1KbNeMKwQq0reLCCaFtqu2eNuSeNHbk= golang.org/x/text v0.34.0/go.mod h1:homfLqTYRFyVYemLBFl5GgL/DWEiH5wcsQ5gSh1yziA= @@ -148,6 +167,8 @@ golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGm golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY= golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q= +golang.org/x/tools v0.42.0 h1:uNgphsn75Tdz5Ji2q36v/nsFSfR/9BRFvqhGBaJGd5k= +golang.org/x/tools v0.42.0/go.mod h1:Ma6lCIwGZvHK6XtgbswSoWroEkhugApmsXyrUmBhfr0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM= google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= @@ -182,3 +203,31 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= +modernc.org/cc/v4 v4.28.2 h1:3tQ0lf2ADtoby2EtSP+J7IE2SHwEJdP8ioR59wx7XpY= +modernc.org/cc/v4 v4.28.2/go.mod h1:OnovgIhbbMXMu1aISnJ0wvVD1KnW+cAUJkIrAWh+kVI= +modernc.org/ccgo/v4 v4.34.0 h1:yRLPFZieg532OT4rp4JFNIVcquwalMX26G95WQDqwCQ= +modernc.org/ccgo/v4 v4.34.0/go.mod h1:AS5WYMyBakQ+fhsHhtP8mWB82KTGPkNNJDGfGQCe0/A= +modernc.org/fileutil v1.4.0 h1:j6ZzNTftVS054gi281TyLjHPp6CPHr2KCxEXjEbD6SM= +modernc.org/fileutil v1.4.0/go.mod h1:EqdKFDxiByqxLk8ozOxObDSfcVOv/54xDs/DUHdvCUU= +modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI= +modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito= +modernc.org/gc/v3 v3.1.2 h1:ZtDCnhonXSZexk/AYsegNRV1lJGgaNZJuKjJSWKyEqo= +modernc.org/gc/v3 v3.1.2/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY= +modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks= +modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI= +modernc.org/libc v1.72.3 h1:ZnDF4tXn4NBXFutMMQC4vtbTFSXhhKzR73fv0beZEAU= +modernc.org/libc v1.72.3/go.mod h1:dn0dZNnnn1clLyvRxLxYExxiKRZIRENOfqQ8XEeg4Qs= +modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU= +modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg= +modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI= +modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw= +modernc.org/opt v0.2.0 h1:tGyef5ApycA7FSEOMraay9SaTk5zmbx7Tu+cJs4QKZg= +modernc.org/opt v0.2.0/go.mod h1:03fq9lsNfvkYSfxrfUhZCWPk1lm4cq4N+Bh//bEtgns= +modernc.org/sortutil v1.2.1 h1:+xyoGf15mM3NMlPDnFqrteY07klSFxLElE2PVuWIJ7w= +modernc.org/sortutil v1.2.1/go.mod h1:7ZI3a3REbai7gzCLcotuw9AC4VZVpYMjDzETGsSMqJE= +modernc.org/sqlite v1.52.0 h1:p4dhYh2tXZCiyaqHwRVJDjIGKWyXayiQpThxgDzJaxo= +modernc.org/sqlite v1.52.0/go.mod h1:tcNzv5p84E0skkmJn038y+hWJbLQXQqEnQfeh5r2JLM= +modernc.org/strutil v1.2.1 h1:UneZBkQA+DX2Rp35KcM69cSsNES9ly8mQWD71HKlOA0= +modernc.org/strutil v1.2.1/go.mod h1:EHkiggD70koQxjVdSBM3JKM7k6L0FbGE5eymy9i3B9A= +modernc.org/token v1.1.0 h1:Xl7Ap9dKaEs5kLoOQeQmPWevfnk/DM5qcLcYlA8ys6Y= +modernc.org/token v1.1.0/go.mod h1:UGzOrNV1mAFSEB63lOFHIpNRUVMvYTc6yu1SMY/XTDM= diff --git a/main.go b/main.go index eead515..1aecd70 100644 --- a/main.go +++ b/main.go @@ -19,6 +19,8 @@ import ( "sync" "time" + sqlquery "aichat/agents/SQL_query" + "github.com/gin-gonic/gin" ark "github.com/volcengine/volcengine-go-sdk/service/arkruntime" "github.com/volcengine/volcengine-go-sdk/service/arkruntime/model" @@ -624,6 +626,7 @@ var ( cfg *Config aiState *OpenAIState searchState *SearchState + sqlState *sqlquery.State store *ConvStore ) @@ -737,8 +740,16 @@ func chatHandler(c *gin.Context) { } chatMessages := req.Messages + if sqlState != nil && sqlState.Enabled() { + withSQL, err := enrichMessagesWithSQL(c.Request.Context(), profile, chatMessages) + if err != nil { + fmt.Fprintln(os.Stderr, "SQL 查询插件调用失败:", err) + } else { + chatMessages = withSQL + } + } if req.WebSearch { - withSearch, err := enrichMessagesWithSearch(c.Request.Context(), req.Messages) + withSearch, err := enrichMessagesWithSearch(c.Request.Context(), chatMessages) if err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) return @@ -882,6 +893,146 @@ func latestUserQuery(messages []ChatMessage) string { return "" } +type sqlActivationDecision struct { + Activate bool `json:"activate"` + Reason string `json:"reason"` +} + +type sqlGenerationResult struct { + Database string `json:"database"` + SQL string `json:"sql"` + Reason string `json:"reason"` +} + +func enrichMessagesWithSQL(ctx context.Context, profile *OpenAIProfile, messages []ChatMessage) ([]ChatMessage, error) { + query := latestUserQuery(messages) + if query == "" { + return messages, nil + } + + activate, reason, err := classifySQLActivation(ctx, profile, messages) + if err != nil { + return messages, err + } + if !activate { + return messages, nil + } + + schemaContext, err := sqlState.SchemaContext(ctx) + if err != nil { + return prependSQLContext(messages, sqlquery.BuildErrorContext(query, err)), nil + } + + generated, err := generateSQLForUserQuery(ctx, profile, query, schemaContext) + if err != nil { + return prependSQLContext(messages, sqlquery.BuildErrorContext(query, err)), nil + } + generated.Database = strings.TrimSpace(generated.Database) + generated.SQL = strings.TrimSpace(generated.SQL) + if generated.SQL == "" { + err := fmt.Errorf("模型未生成可执行 SQL: %s", generated.Reason) + return prependSQLContext(messages, sqlquery.BuildErrorContext(query, err)), nil + } + + result, err := sqlState.ExecuteReadOnly(ctx, generated.Database, generated.SQL) + if err != nil { + return prependSQLContext(messages, sqlquery.BuildErrorContext(query, err)), nil + } + contextText := sqlquery.BuildResultContext(query, generated.SQL, result) + if strings.TrimSpace(reason) != "" { + contextText += "\n激活原因：" + reason + } + return prependSQLContext(messages, contextText), nil +} + +func prependSQLContext(messages []ChatMessage, content string) []ChatMessage { + withSQL := make([]ChatMessage, 0, len(messages)+1) + withSQL = append(withSQL, ChatMessage{Role: "system", Content: content, Hidden: true}) + withSQL = append(withSQL, messages...) + return withSQL +} + +func classifySQLActivation(ctx context.Context, profile *OpenAIProfile, messages []ChatMessage) (bool, string, error) { + query := latestUserQuery(messages) + prompt := fmt.Sprintf("%s\n\n最新用户问题：%s", sqlState.ActivationPrompt(), query) + text, err := completeText(ctx, profile, []ChatMessage{{Role: "system", Content: prompt}}, 512) + if err != nil { + return false, "", err + } + var decision sqlActivationDecision + if err := json.Unmarshal([]byte(extractJSONObject(text)), &decision); err != nil { + return false, "", fmt.Errorf("解析 SQL 查询激活结果失败: %w", err) + } + return decision.Activate, decision.Reason, nil +} + +func generateSQLForUserQuery(ctx context.Context, profile *OpenAIProfile, userQuery string, schemaContext string) (*sqlGenerationResult, error) { + prompt := fmt.Sprintf(`你是只读 SQL 生成器。请根据用户问题和数据库 schema 生成一条只读 SQL。 +要求： +- 只能返回 JSON，不要使用 Markdown。 +- JSON 格式：{"database":"数据库名称","sql":"SELECT ... LIMIT N","reason":"生成原因"} +- 只能生成 SELECT 或 WITH 查询，禁止 INSERT/UPDATE/DELETE/DROP/ALTER/CREATE 等任何修改语句。 +- 必须只使用 schema 中出现的数据库、表和字段。 +- 必须添加 LIMIT，且 LIMIT 不超过插件配置的 max_rows。 +- 如果无法根据 schema 回答，返回 {"database":"","sql":"","reason":"无法根据已知表结构生成查询"}。 + +%s + +用户问题：%s`, schemaContext, userQuery) + text, err := completeText(ctx, profile, []ChatMessage{{Role: "system", Content: prompt}}, 1024) + if err != nil { + return nil, err + } + var generated sqlGenerationResult + if err := json.Unmarshal([]byte(extractJSONObject(text)), &generated); err != nil { + return nil, fmt.Errorf("解析 SQL 生成结果失败: %w", err) + } + return &generated, nil +} + +func completeText(ctx context.Context, profile *OpenAIProfile, chatMessages []ChatMessage, maxTokens int) (string, error) { + messages, err := buildArkMessages(chatMessages) + if err != nil { + return "", err + } + timeout := time.Duration(profile.Config.Timeout) * time.Second + completionCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + stream, err := profile.Client.CreateChatCompletionStream(completionCtx, model.CreateChatCompletionRequest{ + Model: profile.Config.Model, + Messages: messages, + MaxTokens: intPtr(maxTokens), + }.WithStream(true)) + if err != nil { + return "", err + } + defer stream.Close() + + var b strings.Builder + for { + resp, err := stream.Recv() + if errors.Is(err, io.EOF) { + return b.String(), nil + } + if err != nil { + return "", err + } + if len(resp.Choices) > 0 { + b.WriteString(resp.Choices[0].Delta.Content) + } + } +} + +func extractJSONObject(text string) string { + text = strings.TrimSpace(text) + start := strings.Index(text, "{") + end := strings.LastIndex(text, "}") + if start >= 0 && end > start { + return text[start : end+1] + } + return text +} + func webSearch(ctx context.Context, config SearchConfig, query string) ([]searchResult, error) { switch strings.ToLower(config.Provider) { case "duckduckgo", "ddg": @@ -1363,6 +1514,17 @@ func main() { fmt.Fprintln(os.Stderr, "搜索配置初始化失败:", err) os.Exit(1) } + sqlConfig, err := sqlquery.LoadConfig("agents/SQL_query/config.yaml") + if err != nil { + fmt.Fprintln(os.Stderr, "SQL 查询插件配置加载失败:", err) + os.Exit(1) + } + sqlState, err = sqlquery.NewState(sqlConfig) + if err != nil { + fmt.Fprintln(os.Stderr, "SQL 查询插件初始化失败:", err) + os.Exit(1) + } + defer sqlState.Close() store = NewConvStore("conversations") // Gin 路由