From c972dab3d0244d095ad365555ef79e7fe612deaf Mon Sep 17 00:00:00 2001 From: zipg Date: Wed, 5 Aug 2026 14:58:03 +0800 Subject: [PATCH] fix(quick-open): search unloaded tables after cold start --- .../__tests__/useQuickOpen.spec.ts | 124 +++++++++++++++++- apps/desktop/src/composables/useQuickOpen.ts | 67 +++++++--- .../database/defaultDatabase.spec.ts | 5 + .../src/lib/database/defaultDatabase.ts | 6 +- .../connectionStore.completion.spec.ts | 29 ++++ apps/desktop/src/stores/connectionStore.ts | 9 +- 6 files changed, 211 insertions(+), 29 deletions(-) diff --git a/apps/desktop/src/composables/__tests__/useQuickOpen.spec.ts b/apps/desktop/src/composables/__tests__/useQuickOpen.spec.ts index 8b8f20983..d0f168db9 100644 --- a/apps/desktop/src/composables/__tests__/useQuickOpen.spec.ts +++ b/apps/desktop/src/composables/__tests__/useQuickOpen.spec.ts @@ -780,7 +780,7 @@ describe("useQuickOpen", () => { setQuery("ord"); await runDebouncedSearch(); - expect(mockStore.listCompletionTables).toHaveBeenCalledWith("conn1", "app", "ord", 25, undefined, true); + expect(mockStore.listCompletionTables).toHaveBeenCalledWith("conn1", "app", "ord", 25, undefined, true, undefined, undefined, { activateConnection: false }); expect(filteredItems.value).toEqual(expect.arrayContaining([expect.objectContaining({ label: "orders", type: "table", database: "app" })])); }); @@ -851,7 +851,7 @@ describe("useQuickOpen", () => { expect(mockStore.listCompletionTables).not.toHaveBeenCalled(); }); - it("does not request metadata from disconnected contexts", async () => { + it("searches disconnected contexts through lazy connection", async () => { const mockStore = remoteSearchStore({ connectedIds: new Set() }); vi.mocked(useConnectionStore).mockReturnValue(mockStore as any); @@ -859,7 +859,125 @@ describe("useQuickOpen", () => { setQuery("users"); await runDebouncedSearch(); - expect(mockStore.listCompletionTables).not.toHaveBeenCalled(); + expect(mockStore.listCompletionTables).toHaveBeenCalledWith("conn1", "app", "users", 25, undefined, true, undefined, undefined, { activateConnection: false }); + }); + + it("prioritizes the active connection before applying the remote request cap", async () => { + const connections = Array.from({ length: 9 }, (_, index) => ({ id: `conn${index}`, name: `Connection ${index}`, db_type: "mysql", database: `db${index}` })); + const mockStore = remoteSearchStore({ + connections, + activeConnectionId: "conn8", + connectedIds: new Set(["conn8"]), + treeNodes: [], + }); + vi.mocked(useConnectionStore).mockReturnValue(mockStore as any); + + const { setQuery } = useQuickOpen(); + setQuery("users"); + await runDebouncedSearch(); + + const requestedConnections = mockStore.listCompletionTables.mock.calls.map(([connectionId]) => connectionId); + expect(requestedConnections).toContain("conn8"); + expect(requestedConnections).not.toContain("conn7"); + expect(mockStore.listCompletionTables).toHaveBeenCalledTimes(8); + }); + + it("publishes active connection results without waiting for slow cold connections", async () => { + const slowSearch = deferred>(); + const mockStore = remoteSearchStore({ + connections: [ + { id: "cold", name: "Cold", db_type: "mysql", database: "cold_db" }, + { id: "active", name: "Active", db_type: "mysql", database: "active_db" }, + ], + activeConnectionId: "active", + connectedIds: new Set(["active"]), + treeNodes: [], + listCompletionTables: vi.fn((connectionId) => (connectionId === "active" ? Promise.resolve([{ name: "active_users", type: "table" as const }]) : slowSearch.promise)), + }); + vi.mocked(useConnectionStore).mockReturnValue(mockStore as any); + + const { filteredItems, setQuery } = useQuickOpen(); + setQuery("users"); + await runDebouncedSearch(); + + expect(filteredItems.value.map((item) => item.label)).toContain("active_users"); + slowSearch.resolve([]); + await flushAsyncWork(); + }); + + it("drops stale queued contexts before scheduling a newer query", async () => { + const alphaRequests: Array>>> = []; + const listCompletionTables = vi.fn((_connectionId, _database, query) => { + if (query === "alpha") { + const request = deferred>(); + alphaRequests.push(request); + return request.promise; + } + return Promise.resolve([{ name: "beta_table", type: "table" as const }]); + }); + const mockStore = remoteSearchStore({ + treeNodes: Array.from({ length: 4 }, (_, index) => ({ + id: `conn1:db${index}`, + connectionId: "conn1", + type: "database", + database: `db${index}`, + label: `db${index}`, + })), + listCompletionTables, + }); + vi.mocked(useConnectionStore).mockReturnValue(mockStore as any); + + const { filteredItems, setQuery } = useQuickOpen(); + setQuery("alpha"); + await runDebouncedSearch(); + expect(listCompletionTables).toHaveBeenCalledTimes(2); + + setQuery("beta"); + await runDebouncedSearch(); + expect(listCompletionTables).toHaveBeenCalledTimes(2); + + alphaRequests[0]!.resolve([]); + await flushAsyncWork(); + expect(listCompletionTables.mock.calls[2]?.[2]).toBe("beta"); + expect(listCompletionTables.mock.calls.filter(([, , query]) => query === "alpha")).toHaveLength(2); + + alphaRequests[1]!.resolve([]); + await flushAsyncWork(); + expect(filteredItems.value.map((item) => item.label)).toContain("beta_table"); + }); + + it("derives the SQLite main database before its tree is expanded", async () => { + const mockStore = remoteSearchStore({ + connections: [{ id: "sqlite-1", name: "SQLite", db_type: "sqlite", host: "/tmp/app.sqlite" }], + connectedIds: new Set(), + treeNodes: [], + listCompletionTables: vi.fn().mockResolvedValue([{ name: "scroll_test", type: "table" }]), + }); + vi.mocked(useConnectionStore).mockReturnValue(mockStore as any); + + const { filteredItems, setQuery } = useQuickOpen(); + setQuery("scroll_test"); + await runDebouncedSearch(); + + expect(mockStore.listCompletionTables).toHaveBeenCalledWith("sqlite-1", "main", "scroll_test", 25, undefined, true, undefined, undefined, { activateConnection: false }); + expect(filteredItems.value).toEqual(expect.arrayContaining([expect.objectContaining({ label: "scroll_test", type: "table", database: "main" })])); + }); + + it("searches the PostgreSQL backend default database before its tree is expanded", async () => { + const mockStore = remoteSearchStore({ + connections: [{ id: "pg-1", name: "PostgreSQL", db_type: "postgres", database: "" }], + connectedIds: new Set(), + treeNodes: [], + listCompletionTables: vi.fn().mockResolvedValue([{ name: "cold_start_table", type: "table" }]), + }); + vi.mocked(useConnectionStore).mockReturnValue(mockStore as any); + + const { filteredItems, setQuery } = useQuickOpen(); + setQuery("cold_start_table"); + await runDebouncedSearch(); + + expect(mockStore.listCompletionTables).toHaveBeenCalledWith("pg-1", "postgres", "cold_start_table", 25, undefined, true, undefined, undefined, { activateConnection: false }); + expect(filteredItems.value).toEqual(expect.arrayContaining([expect.objectContaining({ label: "cold_start_table", type: "table", database: "postgres" })])); }); it("keeps local results when remote metadata search fails", async () => { diff --git a/apps/desktop/src/composables/useQuickOpen.ts b/apps/desktop/src/composables/useQuickOpen.ts index 686473659..623486904 100644 --- a/apps/desktop/src/composables/useQuickOpen.ts +++ b/apps/desktop/src/composables/useQuickOpen.ts @@ -1,6 +1,7 @@ import { computed, ref, watch } from "vue"; import type { ConnectionConfig } from "@/types/database"; import type { SqlCompletionTable } from "@/lib/sql/sqlCompletion"; +import { resolveDefaultDatabase } from "@/lib/database/defaultDatabase"; import { useConnectionStore } from "@/stores/connectionStore"; import { useSavedSqlStore } from "@/stores/savedSqlStore"; import * as api from "@/lib/backend/api"; @@ -111,7 +112,7 @@ export function useQuickOpen() { let remoteSearchGeneration = 0; let remoteSearchTimer: ReturnType | undefined; let activeRemoteRequests = 0; - const remoteRequestWaiters: Array<() => void> = []; + const remoteRequestWaiters: Array<{ generation: number; resolve: (acquired: boolean) => void }> = []; let sqlFilesLoaded = false; let sqlFilesLoadingPromise: Promise | null = null; let sqlFilesLoadGeneration = 0; @@ -467,12 +468,20 @@ export function useQuickOpen() { function remoteSearchContexts(): Array<{ conn: ConnectionConfig; database: string }> { if (typeof connectionStore.listCompletionTables !== "function") return []; - const connectedIds = connectionStore.connectedIds; - if (!(connectedIds instanceof Set)) return []; const databasesByConnection: Array<{ conn: ConnectionConfig; databases: string[] }> = []; - for (const conn of connectionStore.connections) { - if (!connectedIds.has(conn.id) || REMOTE_SEARCH_UNSUPPORTED_TYPES.has(conn.db_type)) continue; + const orderedConnections = [...connectionStore.connections].sort((left, right) => { + const priority = (conn: ConnectionConfig) => { + if (conn.id === connectionStore.activeConnectionId) return 0; + if (connectionStore.connectedIds.has(conn.id)) return 1; + return 2; + }; + return priority(left) - priority(right); + }); + for (const conn of orderedConnections) { + // listCompletionTables connects on demand. Keeping disconnected connections + // out here makes quick-open blind to unloaded tables after a cold start. + if (REMOTE_SEARCH_UNSUPPORTED_TYPES.has(conn.db_type)) continue; const databases = new Set(); collectConnectionDatabases(connectionStore.treeNodes, conn.id, databases); if (conn.database?.trim()) databases.add(conn.database.trim()); @@ -482,6 +491,8 @@ export function useQuickOpen() { for (const database of conn.attached_databases ?? []) { if (database.name.trim()) databases.add(database.name.trim()); } + const defaultDatabase = resolveDefaultDatabase(conn, [...databases]); + if (defaultDatabase) databases.add(defaultDatabase); if (databases.size > 0) databasesByConnection.push({ conn, databases: [...databases] }); } @@ -500,39 +511,54 @@ export function useQuickOpen() { return contexts; } - async function acquireRemoteRequestSlot(): Promise { + async function acquireRemoteRequestSlot(generation: number): Promise { + if (generation !== remoteSearchGeneration) return false; if (activeRemoteRequests < REMOTE_SEARCH_CONCURRENCY) { activeRemoteRequests++; - return; + return true; + } + return new Promise((resolve) => remoteRequestWaiters.push({ generation, resolve })); + } + + function cancelStaleRemoteRequestWaiters(generation: number): void { + for (let index = remoteRequestWaiters.length - 1; index >= 0; index -= 1) { + const waiter = remoteRequestWaiters[index]!; + if (waiter.generation === generation) continue; + remoteRequestWaiters.splice(index, 1); + waiter.resolve(false); } - await new Promise((resolve) => remoteRequestWaiters.push(resolve)); } function releaseRemoteRequestSlot(): void { - const next = remoteRequestWaiters.shift(); - if (next) next(); + let next = remoteRequestWaiters.shift(); + while (next && next.generation !== remoteSearchGeneration) { + next.resolve(false); + next = remoteRequestWaiters.shift(); + } + if (next) next.resolve(true); else activeRemoteRequests--; } async function runRemoteSearch(query: string, generation: number, contexts: Array<{ conn: ConnectionConfig; database: string }>): Promise { - const groups = await Promise.all( - contexts.map(async ({ conn, database }) => { - await acquireRemoteRequestSlot(); + const groups = contexts.map(() => [] as QuickOpenItem[]); + await Promise.all( + contexts.map(async ({ conn, database }, index) => { + const acquired = await acquireRemoteRequestSlot(generation); + if (!acquired) return; try { // A newer query may supersede queued work before it reaches the metadata API. - if (generation !== remoteSearchGeneration) return []; - const tables = await connectionStore.listCompletionTables(conn.id, database, query, REMOTE_SEARCH_RESULTS_PER_REQUEST, undefined, true); - return tables.slice(0, REMOTE_SEARCH_RESULTS_PER_REQUEST).map((table) => remoteTableItem(table, conn, database)); + if (generation !== remoteSearchGeneration) return; + const tables = await connectionStore.listCompletionTables(conn.id, database, query, REMOTE_SEARCH_RESULTS_PER_REQUEST, undefined, true, undefined, undefined, { activateConnection: false }); + if (generation !== remoteSearchGeneration) return; + groups[index] = tables.slice(0, REMOTE_SEARCH_RESULTS_PER_REQUEST).map((table) => remoteTableItem(table, conn, database)); + remoteItems.value = groups.flat().slice(0, REMOTE_SEARCH_MAX_RESULTS); } catch { - return []; + return; } finally { releaseRemoteRequestSlot(); } }), ); - - if (generation !== remoteSearchGeneration) return; - remoteItems.value = groups.flat().slice(0, REMOTE_SEARCH_MAX_RESULTS); } /** @@ -550,6 +576,7 @@ export function useQuickOpen() { searchQuery, (query) => { const generation = ++remoteSearchGeneration; + cancelStaleRemoteRequestWaiters(generation); if (remoteSearchTimer) clearTimeout(remoteSearchTimer); remoteItems.value = []; diff --git a/apps/desktop/src/lib/__tests__/database/defaultDatabase.spec.ts b/apps/desktop/src/lib/__tests__/database/defaultDatabase.spec.ts index 1b388af73..23914b10c 100644 --- a/apps/desktop/src/lib/__tests__/database/defaultDatabase.spec.ts +++ b/apps/desktop/src/lib/__tests__/database/defaultDatabase.spec.ts @@ -32,4 +32,9 @@ describe("defaultDatabase selectable values", () => { expect(isDefaultDatabase({ db_type: "sqlite", database: "analytics" }, "analytics")).toBe(true); expect(resolveDefaultDatabase({ db_type: "sqlite", host: "primary.db", database: undefined }, ["analytics.db"])).toBe("analytics.db"); }); + + it("matches PostgreSQL backend defaults when the configured database is empty", () => { + expect(resolveDefaultDatabase({ db_type: "postgres", database: "" }, [])).toBe("postgres"); + expect(resolveDefaultDatabase({ db_type: "postgres", driver_profile: "cockroachdb", database: " " }, [])).toBe("defaultdb"); + }); }); diff --git a/apps/desktop/src/lib/database/defaultDatabase.ts b/apps/desktop/src/lib/database/defaultDatabase.ts index f78e6fa50..5b0be32d5 100644 --- a/apps/desktop/src/lib/database/defaultDatabase.ts +++ b/apps/desktop/src/lib/database/defaultDatabase.ts @@ -5,10 +5,12 @@ import { normalizeSqliteNamespace } from "@/lib/database/sqliteNamespace"; export const TREE_SCHEMA_DEFAULT_DATABASE_SELECT_VALUE = "__dbx_tree_schema_default_database__"; export const EMPTY_DATABASE_SELECT_VALUE = "__dbx_empty_database__"; -export function resolveDefaultDatabase(connection: Pick & Partial>, options: string[]): string { +export function resolveDefaultDatabase(connection: Pick & Partial>, options: string[]): string { if (connection.db_type === "cloudflare-d1") return "main"; if (connection.db_type === "sqlite") return normalizeSqliteNamespace(connection.database || options[0], connection); - return connection.database || options[0] || ""; + if (connection.database?.trim()) return connection.database; + if (connection.db_type === "postgres") return connection.driver_profile === "cockroachdb" ? "defaultdb" : "postgres"; + return options[0] || ""; } export function isTreeSchemaDefaultDatabase(dbType: DatabaseType | undefined, database: string): boolean { diff --git a/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts index 5092a2c7a..8c6412beb 100644 --- a/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts +++ b/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts @@ -112,6 +112,35 @@ describe("connectionStore completion assistant", () => { setActivePinia(createPinia()); }); + it("does not replace the active connection during a cold metadata search", async () => { + const connectDb = vi.fn().mockResolvedValue("pg-1"); + const completionAssistantSearch = vi.fn().mockResolvedValue({ + candidates: [{ name: "users", kind: "table", schema: "public" }], + incomplete: false, + fallback_used: false, + }); + + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => ({ + connectDb, + connectionDatabaseInfo: vi.fn().mockResolvedValue(null), + connectionIdentifierQuote: vi.fn().mockResolvedValue('"'), + completionAssistantSearch, + })); + + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [postgresConnection()]; + store.activeConnectionId = "already-active"; + + const tables = await store.listCompletionTables("pg-1", "app", "users", 20, undefined, true, undefined, undefined, { activateConnection: false }); + + expect(connectDb).toHaveBeenCalledOnce(); + expect(store.connectedIds.has("pg-1")).toBe(true); + expect(store.activeConnectionId).toBe("already-active"); + expect(tables).toEqual([{ name: "users", schema: "public", type: "table" }]); + }); + it("deduplicates in-flight assistant table requests", async () => { const completionAssistantSearch = vi.fn().mockResolvedValue({ candidates: [{ name: "accounts", kind: "table", schema: "public" }], diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index 34570c633..21854c388 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -2813,7 +2813,7 @@ export const useConnectionStore = defineStore("connection", () => { invalidateObjectBrowserRowsCache({ connectionId, database }); } - async function ensureConnected(connectionId: string) { + async function ensureConnected(connectionId: string, options: { activate?: boolean } = {}) { if (connectedIds.value.has(connectionId)) { if (hasRecentConnectionHealthCheck(connectionId)) return; // Optimistic: verify backend pool is actually healthy @@ -2842,6 +2842,7 @@ export const useConnectionStore = defineStore("connection", () => { const existingConnect = connectInFlight.get(connectionId); if (existingConnect) { await existingConnect; + if (options.activate !== false) activeConnectionId.value = connectionId; return; } const localAttempt = beginLocalConnectionAttempt(connectionId); @@ -2860,12 +2861,12 @@ export const useConnectionStore = defineStore("connection", () => { await refreshConnectionIdentifierQuote(connectionId, config); markSuccessfulLocalConnectionAttempt(connectionId, localAttempt); markConnectionHealthChecked(connectionId); - activeConnectionId.value = connectionId; clearConnectionError(connectionId); })(); connectInFlight.set(connectionId, connectPromise); try { await connectPromise; + if (options.activate !== false) activeConnectionId.value = connectionId; } catch (e) { if (isCancelledLocalConnectionAttempt(connectionId, localAttempt)) { clearConnectionError(connectionId); @@ -5924,7 +5925,7 @@ export const useConnectionStore = defineStore("connection", () => { return api.listTables(connectionId, database, schema, filter, limit); } - async function listCompletionTables(connectionId: string, database: string, filter = "", limit?: number, schema?: string, globalSearch = false, currentSchema?: string, catalog?: string): Promise { + async function listCompletionTables(connectionId: string, database: string, filter = "", limit?: number, schema?: string, globalSearch = false, currentSchema?: string, catalog?: string, options: { activateConnection?: boolean } = {}): Promise { const trimmedFilter = filter.trim(); const normalizedFilter = trimmedFilter.toLowerCase(); // Remote queries (Dameng/Oracle) are case-sensitive, so the cache key must @@ -5940,7 +5941,7 @@ export const useConnectionStore = defineStore("connection", () => { return withCompletionInFlight( `${cacheKey}:tables`, async () => { - await ensureConnected(connectionId); + await ensureConnected(connectionId, { activate: options.activateConnection !== false }); if (isSchemaAwareDatabase(connectionId)) { if (normalizedFilter || limit) {