✨ feat(query): 支持多条SQL语句执行返回多结果集

- 新增 ResultSetData 结构体承载单个结果集数据 - 新增 MultiResultQuerier/MultiResultQuerierContext 可选接口 - 新增 scanMultiRows 函数利用 NextResultSet() 遍历所有结果集 - MySQL 驱动 DSN 开启 multiStatements=true 并实现多结果集接口 - 新增 DBQueryMulti Wails 方法，支持驱动原生多结果集及自动回退逐条执行 - 新增 Go 版 SQL 拆分函数 splitSQLStatements 及 10 个单元测试 - 前端 QueryEditor handleRun 改为一次性调用 DBQueryMulti - MongoDB 保持独立的逐条执行路径不受影响 - refs #235
2026-07-08 11:31:57 +08:00 · 2026-03-17 22:21:49 +08:00
parent 064cdc34be
commit 0ab10d2e80
10 changed files with 682 additions and 106 deletions
--- a/internal/app/methods_db.go
+++ b/internal/app/methods_db.go
@@ -487,6 +487,138 @@ func (a *App) DBQueryWithCancel(config connection.ConnectionConfig, dbName strin
 	}
 }

+// DBQueryMulti 执行可能包含多条 SQL 语句的查询，返回多个结果集。
+// 如果底层驱动支持 MultiResultQuerier，一次性执行所有语句；
+// 否则按分号拆分后逐条执行，模拟多结果集。
+func (a *App) DBQueryMulti(config connection.ConnectionConfig, dbName string, query string, queryID string) connection.QueryResult {
+	runConfig := normalizeRunConfig(config, dbName)
+
+	if queryID == "" {
+		queryID = generateQueryID()
+	}
+
+	dbInst, err := a.getDatabase(runConfig)
+	if err != nil {
+		logger.Error(err, "DBQueryMulti 获取连接失败：%s", formatConnSummary(runConfig))
+		return connection.QueryResult{Success: false, Message: err.Error(), QueryID: queryID}
+	}
+
+	query = sanitizeSQLForPgLike(runConfig.Type, query)
+	timeoutSeconds := runConfig.Timeout
+	if timeoutSeconds <= 0 {
+		timeoutSeconds = 30
+	}
+	ctx, cancel := utils.ContextWithTimeout(time.Duration(timeoutSeconds) * time.Second)
+	defer cancel()
+
+	a.queryMu.Lock()
+	a.runningQueries[queryID] = queryContext{
+		cancel:  cancel,
+		started: time.Now(),
+	}
+	a.queryMu.Unlock()
+	defer func() {
+		a.queryMu.Lock()
+		delete(a.runningQueries, queryID)
+		a.queryMu.Unlock()
+	}()
+
+	// 尝试使用驱动原生多结果集支持
+	runMultiQuery := func(inst db.Database) ([]connection.ResultSetData, error) {
+		if q, ok := inst.(db.MultiResultQuerierContext); ok {
+			return q.QueryMultiContext(ctx, query)
+		}
+		if q, ok := inst.(db.MultiResultQuerier); ok {
+			return q.QueryMulti(query)
+		}
+		return nil, nil // 返回 nil 表示不支持
+	}
+
+	results, err := runMultiQuery(dbInst)
+	if err != nil && shouldRefreshCachedConnection(err) {
+		if a.invalidateCachedDatabase(runConfig, err) {
+			retryInst, retryErr := a.getDatabaseForcePing(runConfig)
+			if retryErr != nil {
+				logger.Error(retryErr, "DBQueryMulti 重建连接失败：%s SQL片段=%q", formatConnSummary(runConfig), sqlSnippet(query))
+				return connection.QueryResult{Success: false, Message: retryErr.Error(), QueryID: queryID}
+			}
+			results, err = runMultiQuery(retryInst)
+		}
+	}
+	if err != nil {
+		logger.Error(err, "DBQueryMulti 执行失败：%s SQL片段=%q", formatConnSummary(runConfig), sqlSnippet(query))
+		return connection.QueryResult{Success: false, Message: err.Error(), QueryID: queryID}
+	}
+
+	// 驱动支持多结果集，直接返回
+	if results != nil {
+		return connection.QueryResult{Success: true, Data: results, QueryID: queryID}
+	}
+
+	// 驱动不支持多结果集，回退到逐条执行
+	statements := splitSQLStatements(query)
+	if len(statements) == 0 {
+		return connection.QueryResult{
+			Success: true,
+			Data:    []connection.ResultSetData{},
+			QueryID: queryID,
+		}
+	}
+
+	var resultSets []connection.ResultSetData
+	for _, stmt := range statements {
+		stmt = strings.TrimSpace(stmt)
+		if stmt == "" {
+			continue
+		}
+
+		if isReadOnlySQLQuery(runConfig.Type, stmt) {
+			var data []map[string]interface{}
+			var columns []string
+			if q, ok := dbInst.(interface {
+				QueryContext(context.Context, string) ([]map[string]interface{}, []string, error)
+			}); ok {
+				data, columns, err = q.QueryContext(ctx, stmt)
+			} else {
+				data, columns, err = dbInst.Query(stmt)
+			}
+			if err != nil {
+				logger.Error(err, "DBQueryMulti 逐条查询失败：%s SQL片段=%q", formatConnSummary(runConfig), sqlSnippet(stmt))
+				return connection.QueryResult{Success: false, Message: err.Error(), QueryID: queryID}
+			}
+			if data == nil {
+				data = make([]map[string]interface{}, 0)
+			}
+			if columns == nil {
+				columns = []string{}
+			}
+			resultSets = append(resultSets, connection.ResultSetData{Rows: data, Columns: columns})
+		} else {
+			var affected int64
+			if e, ok := dbInst.(interface {
+				ExecContext(context.Context, string) (int64, error)
+			}); ok {
+				affected, err = e.ExecContext(ctx, stmt)
+			} else {
+				affected, err = dbInst.Exec(stmt)
+			}
+			if err != nil {
+				logger.Error(err, "DBQueryMulti 逐条执行失败：%s SQL片段=%q", formatConnSummary(runConfig), sqlSnippet(stmt))
+				return connection.QueryResult{Success: false, Message: err.Error(), QueryID: queryID}
+			}
+			resultSets = append(resultSets, connection.ResultSetData{
+				Rows:    []map[string]interface{}{{"affectedRows": affected}},
+				Columns: []string{"affectedRows"},
+			})
+		}
+	}
+
+	if resultSets == nil {
+		resultSets = []connection.ResultSetData{}
+	}
+	return connection.QueryResult{Success: true, Data: resultSets, QueryID: queryID}
+}
+
 func (a *App) DBQueryIsolated(config connection.ConnectionConfig, dbName string, query string) connection.QueryResult {
 	runConfig := normalizeRunConfig(config, dbName)

--- a/internal/app/sql_split.go
+++ b/internal/app/sql_split.go
@@ -0,0 +1,167 @@
+package app
+
+import "strings"
+
+// splitSQLStatements 按分号拆分 SQL 文本为独立语句。
+// 正确处理单引号/双引号/反引号字符串、行注释（-- / #）、块注释（/* */）和
+// PostgreSQL/Kingbase 的 $$...$$ dollar-quoting，避免在这些上下文中错误拆分。
+func splitSQLStatements(sql string) []string {
+	text := strings.ReplaceAll(sql, "\r\n", "\n")
+	var statements []string
+
+	cur := ""
+	inSingle := false
+	inDouble := false
+	inBacktick := false
+	escaped := false
+	inLineComment := false
+	inBlockComment := false
+	var dollarTag string // postgres/kingbase: $$...$$ or $tag$...$tag$
+
+	push := func() {
+		s := strings.TrimSpace(cur)
+		if s != "" {
+			statements = append(statements, s)
+		}
+		cur = ""
+	}
+
+	for i := 0; i < len(text); i++ {
+		ch := text[i]
+		next := byte(0)
+		if i+1 < len(text) {
+			next = text[i+1]
+		}
+
+		// 行注释
+		if inLineComment {
+			if ch == '\n' {
+				inLineComment = false
+			}
+			cur += string(ch)
+			continue
+		}
+
+		// 块注释
+		if inBlockComment {
+			cur += string(ch)
+			if ch == '*' && next == '/' {
+				cur += "/"
+				i++
+				inBlockComment = false
+			}
+			continue
+		}
+
+		// Dollar-quoting
+		if dollarTag != "" {
+			if strings.HasPrefix(text[i:], dollarTag) {
+				cur += dollarTag
+				i += len(dollarTag) - 1
+				dollarTag = ""
+			} else {
+				cur += string(ch)
+			}
+			continue
+		}
+
+		// 转义字符
+		if escaped {
+			escaped = false
+			cur += string(ch)
+			continue
+		}
+		if (inSingle || inDouble) && ch == '\\' {
+			escaped = true
+			cur += string(ch)
+			continue
+		}
+
+		// 字符串开闭
+		if !inDouble && !inBacktick && ch == '\'' {
+			inSingle = !inSingle
+			cur += string(ch)
+			continue
+		}
+		if !inSingle && !inBacktick && ch == '"' {
+			inDouble = !inDouble
+			cur += string(ch)
+			continue
+		}
+		if !inSingle && !inDouble && ch == '`' {
+			inBacktick = !inBacktick
+			cur += string(ch)
+			continue
+		}
+
+		// 在引号/反引号内部不做任何判断
+		if inSingle || inDouble || inBacktick {
+			cur += string(ch)
+			continue
+		}
+
+		// 行注释开始
+		if ch == '-' && next == '-' {
+			inLineComment = true
+			cur += string(ch)
+			continue
+		}
+		if ch == '#' {
+			inLineComment = true
+			cur += string(ch)
+			continue
+		}
+
+		// 块注释开始
+		if ch == '/' && next == '*' {
+			inBlockComment = true
+			cur += "/*"
+			i++
+			continue
+		}
+
+		// Dollar-quoting 开始
+		if ch == '$' {
+			if tag := parseSQLDollarTag(text[i:]); tag != "" {
+				dollarTag = tag
+				cur += tag
+				i += len(tag) - 1
+				continue
+			}
+		}
+
+		// 分号分隔（支持全角分号"；"）
+		if ch == ';' {
+			push()
+			continue
+		}
+		// 全角分号 UTF-8 序列: 0xEF 0xBC 0x9B
+		if ch == 0xEF && i+2 < len(text) && text[i+1] == 0xBC && text[i+2] == 0x9B {
+			push()
+			i += 2
+			continue
+		}
+
+		cur += string(ch)
+	}
+
+	push()
+	return statements
+}
+
+// parseSQLDollarTag 解析 PostgreSQL/Kingbase 的 dollar-quoting 标签。
+func parseSQLDollarTag(s string) string {
+	if len(s) < 2 || s[0] != '$' {
+		return ""
+	}
+	for i := 1; i < len(s); i++ {
+		c := s[i]
+		if c == '$' {
+			return s[:i+1]
+		}
+		if !((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_') {
+			return ""
+		}
+	}
+	return ""
+}
--- a/internal/app/sql_split_test.go
+++ b/internal/app/sql_split_test.go
@@ -0,0 +1,94 @@
+package app
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestSplitSQLStatements_BasicSplit(t *testing.T) {
+	input := "SELECT 1; SELECT 2; SELECT 3"
+	got := splitSQLStatements(input)
+	want := []string{"SELECT 1", "SELECT 2", "SELECT 3"}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("splitSQLStatements(%q) = %v, want %v", input, got, want)
+	}
+}
+
+func TestSplitSQLStatements_QuotedSemicolon(t *testing.T) {
+	input := `SELECT 'hello;world'; SELECT 2`
+	got := splitSQLStatements(input)
+	want := []string{`SELECT 'hello;world'`, "SELECT 2"}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("splitSQLStatements(%q) = %v, want %v", input, got, want)
+	}
+}
+
+func TestSplitSQLStatements_LineComment(t *testing.T) {
+	input := "SELECT 1; -- this is a comment;\nSELECT 2"
+	got := splitSQLStatements(input)
+	want := []string{"SELECT 1", "-- this is a comment;\nSELECT 2"}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("splitSQLStatements(%q) = %v, want %v", input, got, want)
+	}
+}
+
+func TestSplitSQLStatements_BlockComment(t *testing.T) {
+	input := "SELECT /* ; */ 1; SELECT 2"
+	got := splitSQLStatements(input)
+	want := []string{"SELECT /* ; */ 1", "SELECT 2"}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("splitSQLStatements(%q) = %v, want %v", input, got, want)
+	}
+}
+
+func TestSplitSQLStatements_EmptyInput(t *testing.T) {
+	got := splitSQLStatements("")
+	if len(got) != 0 {
+		t.Errorf("splitSQLStatements(\"\") = %v, want empty slice", got)
+	}
+}
+
+func TestSplitSQLStatements_SingleStatement(t *testing.T) {
+	input := "SELECT * FROM users WHERE id = 1"
+	got := splitSQLStatements(input)
+	want := []string{"SELECT * FROM users WHERE id = 1"}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("splitSQLStatements(%q) = %v, want %v", input, got, want)
+	}
+}
+
+func TestSplitSQLStatements_DollarQuoting(t *testing.T) {
+	input := "SELECT $tag$hello;world$tag$; SELECT 2"
+	got := splitSQLStatements(input)
+	want := []string{"SELECT $tag$hello;world$tag$", "SELECT 2"}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("splitSQLStatements(%q) = %v, want %v", input, got, want)
+	}
+}
+
+func TestSplitSQLStatements_FullWidthSemicolon(t *testing.T) {
+	input := "SELECT 1；SELECT 2"
+	got := splitSQLStatements(input)
+	want := []string{"SELECT 1", "SELECT 2"}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("splitSQLStatements(%q) = %v, want %v", input, got, want)
+	}
+}
+
+func TestSplitSQLStatements_Backtick(t *testing.T) {
+	input := "SELECT `col;name` FROM t; SELECT 2"
+	got := splitSQLStatements(input)
+	want := []string{"SELECT `col;name` FROM t", "SELECT 2"}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("splitSQLStatements(%q) = %v, want %v", input, got, want)
+	}
+}
+
+func TestSplitSQLStatements_TrailingSemicolon(t *testing.T) {
+	input := "SELECT 1; SELECT 2;"
+	got := splitSQLStatements(input)
+	want := []string{"SELECT 1", "SELECT 2"}
+	if !reflect.DeepEqual(got, want) {
+		t.Errorf("splitSQLStatements(%q) = %v, want %v", input, got, want)
+	}
+}