xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(kingbase): preserve query schema context

This commit is contained in:
zipg 2026-07-29 18:43:53 +08:00 committed by GitHub
parent 51ab6fb9ec
commit 02aea73d28
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 306 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

93
agents/drivers/kingbase-go/main.go
View File

@ -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

98
agents/drivers/kingbase-go/main_test.go
View File

@ -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()

17
apps/desktop/src/App.vue
View File

@ -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 {

4
apps/desktop/src/lib/__tests__/database/jdbcDialect.spec.ts vendored
View File

@ -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();
});

19
apps/desktop/src/lib/history/historyRestoreTarget.ts Normal file
View File

@ -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 };
}

124
packages/app-tests/queryStore.test.ts vendored
View File

@ -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());