From 02aea73d28dbf77462f8248cfb14fe041e101062 Mon Sep 17 00:00:00 2001 From: zipg Date: Wed, 29 Jul 2026 18:43:53 +0800 Subject: [PATCH] fix(kingbase): preserve query schema context --- agents/drivers/kingbase-go/main.go | 93 ++++++++----- agents/drivers/kingbase-go/main_test.go | 98 +++++++++++++- apps/desktop/src/App.vue | 17 ++- .../__tests__/database/jdbcDialect.spec.ts | 4 + .../src/lib/history/historyRestoreTarget.ts | 19 +++ packages/app-tests/queryStore.test.ts | 124 ++++++++++++++++++ 6 files changed, 306 insertions(+), 49 deletions(-) create mode 100644 apps/desktop/src/lib/history/historyRestoreTarget.ts diff --git a/agents/drivers/kingbase-go/main.go b/agents/drivers/kingbase-go/main.go index 530b772ac..dbc4b9d5e 100644 --- a/agents/drivers/kingbase-go/main.go +++ b/agents/drivers/kingbase-go/main.go @@ -120,6 +120,7 @@ type queryPageResult struct { type querySession struct { rows *sql.Rows + conn *sql.Conn columns []string columnTypes []string pending []any @@ -572,17 +573,15 @@ func (s *server) cancelActiveQuery() { func (s *server) executeQuery(opts queryOptions) (queryResult, error) { start := time.Now() - if err := s.setSchema(opts.Schema); err != nil { - return queryResult{}, err - } sqlText := trimStatementSQL(opts.SQL) if isQuerySQL(sqlText) { - rows, cancel, err := s.queryRows(sqlText, opts.TimeoutSecs) + rows, conn, cancel, err := s.queryRows(sqlText, opts.Schema, opts.TimeoutSecs) if err != nil { return queryResult{}, err } defer func() { _ = rows.Close() + _ = conn.Close() s.endOperation(cancel) }() maxRows := opts.MaxRows @@ -593,13 +592,15 @@ func (s *server) executeQuery(opts queryOptions) (queryResult, error) { result.ExecutionTimeMS = time.Since(start).Milliseconds() return result, err } - db, err := s.requireDB() + conn, ctx, cancel, err := s.operationConn(opts.Schema, opts.TimeoutSecs) if err != nil { return queryResult{}, err } - ctx, cancel := s.beginOperation(opts.TimeoutSecs) - defer s.endOperation(cancel) - execResult, err := db.ExecContext(ctx, sqlText) + defer func() { + _ = conn.Close() + s.endOperation(cancel) + }() + execResult, err := conn.ExecContext(ctx, sqlText) if err != nil { return queryResult{}, err } @@ -607,37 +608,35 @@ func (s *server) executeQuery(opts queryOptions) (queryResult, error) { return queryResult{Columns: []string{}, ColumnTypes: []string{}, Rows: [][]any{}, AffectedRows: affected, ExecutionTimeMS: time.Since(start).Milliseconds()}, nil } -func (s *server) queryRows(sqlText string, timeoutSecs int) (*sql.Rows, context.CancelFunc, error) { - db, err := s.requireDB() +func (s *server) queryRows(sqlText string, schema string, timeoutSecs int) (*sql.Rows, *sql.Conn, context.CancelFunc, error) { + conn, ctx, cancel, err := s.operationConn(schema, timeoutSecs) if err != nil { - return nil, nil, err + return nil, nil, nil, err } - ctx, cancel := s.beginOperation(timeoutSecs) - rows, err := db.QueryContext(ctx, sqlText) + rows, err := conn.QueryContext(ctx, sqlText) if err != nil { + _ = conn.Close() s.endOperation(cancel) - return nil, nil, err + return nil, nil, nil, err } - return rows, cancel, nil + return rows, conn, cancel, nil } func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageResult, error) { start := time.Now() - if err := s.setSchema(opts.Schema); err != nil { - return queryPageResult{}, err - } sqlText := trimStatementSQL(opts.SQL) if !isQuerySQL(sqlText) { result, err := s.executeQuery(opts) return queryPageResult{Columns: result.Columns, ColumnTypes: result.ColumnTypes, Rows: result.Rows, AffectedRows: result.AffectedRows, ExecutionTimeMS: result.ExecutionTimeMS, Truncated: result.Truncated}, err } - rows, cancel, err := s.queryRows(sqlText, opts.TimeoutSecs) + rows, conn, cancel, err := s.queryRows(sqlText, opts.Schema, opts.TimeoutSecs) if err != nil { return queryPageResult{}, err } columns, err := rows.Columns() if err != nil { _ = rows.Close() + _ = conn.Close() s.endOperation(cancel) return queryPageResult{}, err } @@ -645,11 +644,12 @@ func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageRes if maxRows <= 0 { maxRows = defaultMaxRows } - session := &querySession{rows: rows, columns: columns, columnTypes: columnTypeNames(rows), remaining: maxRows, cancel: cancel} + session := &querySession{rows: rows, conn: conn, columns: columns, columnTypes: columnTypeNames(rows), remaining: maxRows, cancel: cancel} result, err := readQuerySessionPage(session, pageSize) result.ExecutionTimeMS = time.Since(start).Milliseconds() if err != nil { _ = rows.Close() + _ = conn.Close() s.endOperation(cancel) return queryPageResult{}, err } @@ -660,6 +660,7 @@ func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageRes result.SessionID = &id } else { _ = rows.Close() + _ = conn.Close() s.endOperation(cancel) } return result, nil @@ -689,6 +690,9 @@ func (s *server) closeQuerySession(id string) bool { return false } _ = session.rows.Close() + if session.conn != nil { + _ = session.conn.Close() + } if session.cancel != nil { s.endOperation(session.cancel) } @@ -789,22 +793,23 @@ func columnTypeNames(rows *sql.Rows) []string { } func (s *server) executeTransaction(params map[string]json.RawMessage) (queryResult, error) { - db, err := s.requireDB() + statements := stringSliceParam(params, "statements") + conn, ctx, cancel, err := s.operationConn(stringParam(params, "schema"), intParam(params, "timeoutSecs")) if err != nil { return queryResult{}, err } - statements := stringSliceParam(params, "statements") - if err := s.setSchema(stringParam(params, "schema")); err != nil { - return queryResult{}, err - } + defer func() { + _ = conn.Close() + s.endOperation(cancel) + }() start := time.Now() - tx, err := db.Begin() + tx, err := conn.BeginTx(ctx, nil) if err != nil { return queryResult{}, err } var affected int64 for _, statement := range statements { - result, execErr := tx.Exec(trimStatementSQL(statement)) + result, execErr := tx.ExecContext(ctx, trimStatementSQL(statement)) if execErr != nil { _ = tx.Rollback() return queryResult{}, execErr @@ -831,25 +836,41 @@ func (s *server) executeBatch(params map[string]json.RawMessage) (queryResult, e return queryResult{Columns: []string{}, ColumnTypes: []string{}, Rows: [][]any{}, AffectedRows: affected, ExecutionTimeMS: time.Since(start).Milliseconds()}, nil } -func (s *server) setSchema(schema string) error { - schema = strings.TrimSpace(schema) - if schema == "" && !s.schemaSet { - return nil - } - if schema != "" && s.schemaSet && schema == s.currentSchema { - return nil +func (s *server) operationConn(schema string, timeoutSecs int) (*sql.Conn, context.Context, context.CancelFunc, error) { + ctx, cancel := s.beginOperation(timeoutSecs) + conn, err := s.schemaConn(ctx, schema) + if err != nil { + s.endOperation(cancel) + return nil, nil, nil, err } + return conn, ctx, cancel, nil +} + +func (s *server) schemaConn(ctx context.Context, schema string) (*sql.Conn, error) { db, err := s.requireDB() if err != nil { - return err + return nil, err } + conn, err := db.Conn(ctx) + if err != nil { + return nil, err + } + if err := s.setSchema(ctx, conn, schema); err != nil { + _ = conn.Close() + return nil, err + } + return conn, nil +} + +func (s *server) setSchema(ctx context.Context, conn *sql.Conn, schema string) error { + schema = strings.TrimSpace(schema) statement := "RESET search_path" if schema != "" { // Kingbase implicitly prioritizes its system catalog when it is not // listed explicitly, matching the JDBC agent and DBeaver behavior. statement = "SET search_path TO " + quoteIdentifier(schema) } - if _, err = db.Exec(statement); err != nil { + if _, err := conn.ExecContext(ctx, statement); err != nil { return err } s.currentSchema = schema diff --git a/agents/drivers/kingbase-go/main_test.go b/agents/drivers/kingbase-go/main_test.go index 9ae68430c..783c746d4 100644 --- a/agents/drivers/kingbase-go/main_test.go +++ b/agents/drivers/kingbase-go/main_test.go @@ -26,14 +26,21 @@ var registerMetadataDriver sync.Once var metadataState atomic.Pointer[metadataDriverState] type fakeDriverState struct { - queryArgs int - queryCtx context.Context - rowCount int + mu sync.Mutex + nextConnID int + queryArgs int + queryCtx context.Context + queryConnID int + rowCount int + execStatements []string + execConnIDs []int } type fakeDriver struct{} -type fakeConn struct{} +type fakeConn struct { + id int +} type fakeRows struct { current int @@ -103,7 +110,13 @@ type valueRows struct { index int } -func (fakeDriver) Open(string) (driver.Conn, error) { return fakeConn{}, nil } +func (fakeDriver) Open(string) (driver.Conn, error) { + state := testDriverState.Load() + state.mu.Lock() + defer state.mu.Unlock() + state.nextConnID++ + return fakeConn{id: state.nextConnID}, nil +} func (fakeConn) Prepare(string) (driver.Stmt, error) { return nil, driver.ErrSkip } @@ -111,14 +124,22 @@ func (fakeConn) Close() error { return nil } func (fakeConn) Begin() (driver.Tx, error) { return nil, driver.ErrSkip } -func (fakeConn) QueryContext(ctx context.Context, _ string, args []driver.NamedValue) (driver.Rows, error) { +func (connection fakeConn) QueryContext(ctx context.Context, _ string, args []driver.NamedValue) (driver.Rows, error) { state := testDriverState.Load() + state.mu.Lock() + defer state.mu.Unlock() state.queryArgs = len(args) state.queryCtx = ctx + state.queryConnID = connection.id return &fakeRows{count: state.rowCount}, nil } -func (fakeConn) ExecContext(context.Context, string, []driver.NamedValue) (driver.Result, error) { +func (connection fakeConn) ExecContext(_ context.Context, query string, _ []driver.NamedValue) (driver.Result, error) { + state := testDriverState.Load() + state.mu.Lock() + defer state.mu.Unlock() + state.execStatements = append(state.execStatements, query) + state.execConnIDs = append(state.execConnIDs, connection.id) return driver.RowsAffected(1), nil } @@ -1237,6 +1258,69 @@ func TestExecuteQueryUsesSimpleProtocolAndReleasesContext(t *testing.T) { assertContextCanceled(t, state.queryCtx) } +func TestExecuteQueryReappliesSchemaForRepeatedRequests(t *testing.T) { + db, state := openFakeDB(t, 1) + server := newServer() + server.db = db + + for range 2 { + if _, err := server.executeQuery(queryOptions{SQL: "SELECT 1", Schema: "sdy_smartsite", MaxRows: 10}); err != nil { + t.Fatal(err) + } + } + + expected := []string{`SET search_path TO "sdy_smartsite"`, `SET search_path TO "sdy_smartsite"`} + if len(state.execStatements) != len(expected) { + t.Fatalf("expected repeated schema setup, got %v", state.execStatements) + } + for index, statement := range expected { + if state.execStatements[index] != statement { + t.Fatalf("unexpected schema statement at %d: %s", index, state.execStatements[index]) + } + } +} + +func TestExecuteQueryAppliesSchemaOnSamePoolConnection(t *testing.T) { + db, state := openFakeDB(t, 1) + db.SetMaxOpenConns(4) + server := newServer() + server.db = db + + if _, err := server.executeQuery(queryOptions{SQL: "SELECT 1", Schema: "sdy_smartsite", MaxRows: 10}); err != nil { + t.Fatal(err) + } + + state.mu.Lock() + defer state.mu.Unlock() + if len(state.execStatements) != 1 || state.execStatements[0] != `SET search_path TO "sdy_smartsite"` { + t.Fatalf("unexpected schema setup: %v", state.execStatements) + } + if len(state.execConnIDs) != 1 || state.execConnIDs[0] != state.queryConnID { + t.Fatalf("schema setup and query used different connections: exec=%v query=%d", state.execConnIDs, state.queryConnID) + } +} + +func TestExecuteStatementAppliesSchemaOnSamePoolConnection(t *testing.T) { + db, state := openFakeDB(t, 0) + db.SetMaxOpenConns(4) + server := newServer() + server.db = db + + if _, err := server.executeQuery(queryOptions{SQL: "UPDATE orders SET status = 1", Schema: "sdy_smartsite"}); err != nil { + t.Fatal(err) + } + + state.mu.Lock() + defer state.mu.Unlock() + expected := []string{`SET search_path TO "sdy_smartsite"`, "UPDATE orders SET status = 1"} + if strings.Join(state.execStatements, "\n") != strings.Join(expected, "\n") { + t.Fatalf("unexpected statements: %v", state.execStatements) + } + if len(state.execConnIDs) != 2 || state.execConnIDs[0] != state.execConnIDs[1] { + t.Fatalf("schema setup and statement used different connections: %v", state.execConnIDs) + } +} + func TestPagedQueryKeepsContextUntilSessionCloses(t *testing.T) { db, state := openFakeDB(t, 3) server := newServer() diff --git a/apps/desktop/src/App.vue b/apps/desktop/src/App.vue index 94d5dc507..525024553 100644 --- a/apps/desktop/src/App.vue +++ b/apps/desktop/src/App.vue @@ -45,6 +45,7 @@ import { findTreeNodeById, resolveNewQueryTarget, resolveNewQueryInitialSql } fr import { sqlObjectNavigationSourceKind, sqlObjectNavigationTableType, type SqlObjectNavigationTarget } from "@/lib/sql/sqlNavigation"; import { buildExecutableObjectSourceStatements, executeObjectSourceSave } from "@/lib/table/objectSourceEditor"; import { schemaAfterConnectionSwitch } from "@/lib/schema/connectionSchemaInitialization"; +import { resolveHistorySqlRestoreTarget } from "@/lib/history/historyRestoreTarget"; import { resolveExecutableSql, resolveExecutableSqlWithBackend, type SqlExecutionSnapshot } from "@/lib/sql/sqlExecutionTarget"; import { uuid } from "@/lib/common/utils"; import { isMacOS, isWindows } from "@/lib/backend/platform"; @@ -254,11 +255,14 @@ function restoreHistorySql(sql: string, entry: HistoryEntry) { return; } - const connectionId = entry.connection_id || tab?.connectionId || connectionStore.connections[0]?.id; - if (!connectionId) return; - const config = connectionStore.getConfig(connectionId); - const database = entry.database || tab?.database || (config ? resolveDefaultDatabase(config, []) : ""); - const tabId = queryStore.createTab(connectionId, database || "", t("tabs.sql")); + const target = resolveHistorySqlRestoreTarget({ + entry, + activeTab: tab, + firstConnectionId: connectionStore.connections[0]?.id, + getConfig: (connectionId) => connectionStore.getConfig(connectionId), + }); + if (!target) return; + const tabId = queryStore.createTab(target.connectionId, target.database, t("tabs.sql"), "query", target.schema); queryStore.updateSql(tabId, sql); } @@ -1554,7 +1558,8 @@ function ensureQueryTab(): string { if (tab && tab.mode === "query") return tab.id; const connId = connectionStore.activeConnectionId || connectionStore.connections[0]?.id || ""; const db = tab?.connectionId === connId ? tab.database : connectionStore.getConfig(connId)?.database || ""; - return queryStore.createTab(connId, db, undefined, "query"); + const schema = tab?.connectionId === connId ? tab.schema : undefined; + return queryStore.createTab(connId, db, undefined, "query", schema); } function routeAiRedisCommand(command: string, execute: boolean): boolean { diff --git a/apps/desktop/src/lib/__tests__/database/jdbcDialect.spec.ts b/apps/desktop/src/lib/__tests__/database/jdbcDialect.spec.ts index e78361d52..a24baeb94 100644 --- a/apps/desktop/src/lib/__tests__/database/jdbcDialect.spec.ts +++ b/apps/desktop/src/lib/__tests__/database/jdbcDialect.spec.ts @@ -85,6 +85,10 @@ describe("query execution schema", () => { expect(connectionQueryExecutionSchema({ db_type: "postgres" }, "app", "reporting", false)).toBe("reporting"); }); + it("prefers an explicit schema for Kingbase query execution", () => { + expect(connectionQueryExecutionSchema({ db_type: "kingbase" }, "qinzhou", "sdy_smartsite", false)).toBe("sdy_smartsite"); + }); + it("does not send a schema for MySQL database context", () => { expect(connectionQueryExecutionSchema({ db_type: "mysql" }, "app", undefined, false)).toBeUndefined(); }); diff --git a/apps/desktop/src/lib/history/historyRestoreTarget.ts b/apps/desktop/src/lib/history/historyRestoreTarget.ts new file mode 100644 index 000000000..36bfaadcf --- /dev/null +++ b/apps/desktop/src/lib/history/historyRestoreTarget.ts @@ -0,0 +1,19 @@ +import type { HistoryEntry } from "@/lib/backend/api"; +import { resolveDefaultDatabase } from "@/lib/database/defaultDatabase"; +import type { ConnectionConfig, QueryTab } from "@/types/database"; + +export interface HistorySqlRestoreTarget { + connectionId: string; + database: string; + schema?: string; +} + +export function resolveHistorySqlRestoreTarget(options: { entry: HistoryEntry; activeTab?: QueryTab; firstConnectionId?: string; getConfig: (connectionId: string) => ConnectionConfig | undefined }): HistorySqlRestoreTarget | null { + const { entry, activeTab, firstConnectionId, getConfig } = options; + const connectionId = entry.connection_id || activeTab?.connectionId || firstConnectionId; + if (!connectionId) return null; + const config = getConfig(connectionId); + const database = entry.database || activeTab?.database || (config ? resolveDefaultDatabase(config, []) : ""); + const schema = activeTab?.connectionId === connectionId && activeTab.database === database ? activeTab.schema : undefined; + return { connectionId, database: database || "", schema }; +} diff --git a/packages/app-tests/queryStore.test.ts b/packages/app-tests/queryStore.test.ts index a125ebddc..202d713dd 100644 --- a/packages/app-tests/queryStore.test.ts +++ b/packages/app-tests/queryStore.test.ts @@ -7,10 +7,12 @@ import { analyzeEditableQueryEditability } from "../../apps/desktop/src/lib/sql/ import { resultSqlForGrid } from "../../apps/desktop/src/lib/tabs/tabPresentation.ts"; import { parseMongoCommand } from "../../apps/desktop/src/lib/mongo/mongoShellCommand.ts"; import { useExportTracker } from "../../apps/desktop/src/composables/useExportTracker.ts"; +import { resolveHistorySqlRestoreTarget } from "../../apps/desktop/src/lib/history/historyRestoreTarget.ts"; import { useConnectionStore } from "../../apps/desktop/src/stores/connectionStore.ts"; import { useQueryStore } from "../../apps/desktop/src/stores/queryStore.ts"; import { useSettingsStore } from "../../apps/desktop/src/stores/settingsStore.ts"; import type { ConnectionConfig } from "../../apps/desktop/src/types/database.ts"; +import type { HistoryEntry } from "../../apps/desktop/src/lib/backend/tauri.ts"; import type { QueryResult } from "../../apps/desktop/src/types/database.ts"; afterEach(() => { @@ -88,6 +90,14 @@ function sparkConn(id: string): ConnectionConfig { }; } +function kingbaseConn(id: string): ConnectionConfig { + return { + ...conn(id), + db_type: "kingbase", + port: 54321, + }; +} + function withConnectionHealthMock(handler: typeof fetch): typeof fetch { return async (input, init) => { if (String(input) === "/api/connection/check-health") { @@ -4885,6 +4895,120 @@ test("Spark query execution applies the selected database as schema context", as } }); +test("Kingbase query execution sends the selected schema context", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + + connectionStore.addEphemeralConnection(kingbaseConn("kingbase-1")); + const tabId = store.createTab("kingbase-1", "qinzhou", "Query", "query", "sdy_smartsite"); + let executeBody: any; + + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + const url = String(input); + if (url === "/api/query/prepare-pagination-plan") { + return new Response(JSON.stringify({ sqlToExecute: "select * from busi_sea_trip_records", useAgentResultSession: false }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + if (url === "/api/query/execute-multi") { + executeBody = JSON.parse(String(init?.body ?? "{}")); + return new Response(JSON.stringify([{ columns: ["id"], rows: [[1]], affected_rows: 0, execution_time_ms: 1 }]), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + if (url === "/api/query/analyze-editability") { + return new Response(JSON.stringify({ editable: false, reason: "complex-source" }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + await store.executeTabSql(tabId, "select * from busi_sea_trip_records"); + + assert.equal(executeBody.database, "qinzhou"); + assert.equal(executeBody.schema, "sdy_smartsite"); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + +test("Kingbase history restore inherits schema when the history entry keeps its database", async () => { + const restoreStorage = installMemoryStorage(); + setActivePinia(createPinia()); + const connectionStore = useConnectionStore(); + const store = useQueryStore(); + const originalFetch = globalThis.fetch; + + connectionStore.addEphemeralConnection(kingbaseConn("kingbase-1")); + const currentTabId = store.createTab("kingbase-1", "qinzhou", "Current", "query", "sdy_smartsite"); + const currentTab = store.tabs.find((tab) => tab.id === currentTabId); + assert.ok(currentTab); + const entry: HistoryEntry = { + id: "history-1", + connection_id: "kingbase-1", + connection_name: "Kingbase", + database: "qinzhou", + sql: "select * from busi_sea_trip_records", + executed_at: "2026-07-29T08:00:00Z", + execution_time_ms: 12, + success: true, + }; + const target = resolveHistorySqlRestoreTarget({ + entry, + activeTab: currentTab, + firstConnectionId: connectionStore.connections[0]?.id, + getConfig: (connectionId) => connectionStore.getConfig(connectionId), + }); + assert.ok(target); + assert.deepEqual(target, { connectionId: "kingbase-1", database: "qinzhou", schema: "sdy_smartsite" }); + const restoredTabId = store.createTab(target.connectionId, target.database, "SQL", "query", target.schema); + store.updateSql(restoredTabId, entry.sql); + let executeBody: any; + + globalThis.fetch = withConnectionHealthMock(async (input, init) => { + const url = String(input); + if (url === "/api/query/prepare-pagination-plan") { + return new Response(JSON.stringify({ sqlToExecute: entry.sql, useAgentResultSession: false }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + if (url === "/api/query/execute-multi") { + executeBody = JSON.parse(String(init?.body ?? "{}")); + return new Response(JSON.stringify([{ columns: ["id"], rows: [[1]], affected_rows: 0, execution_time_ms: 1 }]), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + if (url === "/api/query/analyze-editability") { + return new Response(JSON.stringify({ editable: false, reason: "complex-source" }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + return new Response("unexpected request", { status: 500 }); + }); + + try { + await store.executeTabSql(restoredTabId, entry.sql); + + assert.equal(executeBody.database, "qinzhou"); + assert.equal(executeBody.schema, "sdy_smartsite"); + } finally { + globalThis.fetch = originalFetch; + restoreStorage(); + } +}); + test("data tab execution uses a tab-scoped client session", async () => { const restoreStorage = installMemoryStorage(); setActivePinia(createPinia());