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

fix(quick-open): search unloaded tables after cold start

This commit is contained in:
zipg 2026-08-05 14:58:03 +08:00 committed by GitHub
parent 605d20dd4f
commit c972dab3d0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 211 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

124
apps/desktop/src/composables/__tests__/useQuickOpen.spec.ts vendored
View File

@ -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<string>() });
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<Array<{ name: string; type: "table" }>>();
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<ReturnType<typeof deferred<Array<{ name: string; type: "table" }>>>> = [];
const listCompletionTables = vi.fn((_connectionId, _database, query) => {
if (query === "alpha") {
const request = deferred<Array<{ name: string; type: "table" }>>();
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<string>(),
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<string>(),
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 () => {

67
apps/desktop/src/composables/useQuickOpen.ts
View File

@ -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<typeof setTimeout> | undefined;
let activeRemoteRequests = 0;
const remoteRequestWaiters: Array<() => void> = [];
const remoteRequestWaiters: Array<{ generation: number; resolve: (acquired: boolean) => void }> = [];
let sqlFilesLoaded = false;
let sqlFilesLoadingPromise: Promise<void> | 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<string>();
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<void> {
async function acquireRemoteRequestSlot(generation: number): Promise<boolean> {
if (generation !== remoteSearchGeneration) return false;
if (activeRemoteRequests < REMOTE_SEARCH_CONCURRENCY) {
activeRemoteRequests++;
return;
return true;
}
return new Promise<boolean>((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<void>((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<void> {
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 = [];

5
apps/desktop/src/lib/__tests__/database/defaultDatabase.spec.ts vendored
View File

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

6
apps/desktop/src/lib/database/defaultDatabase.ts
View File

@ -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<ConnectionConfig, "database"> & Partial<Pick<ConnectionConfig, "db_type" | "host">>, options: string[]): string {
export function resolveDefaultDatabase(connection: Pick<ConnectionConfig, "database"> & Partial<Pick<ConnectionConfig, "db_type" | "driver_profile" | "host">>, 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 {

29
apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts vendored
View File

@ -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" }],

9
apps/desktop/src/stores/connectionStore.ts
View File

@ -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<SqlCompletionTable[]> {
async function listCompletionTables(connectionId: string, database: string, filter = "", limit?: number, schema?: string, globalSearch = false, currentSchema?: string, catalog?: string, options: { activateConnection?: boolean } = {}): Promise<SqlCompletionTable[]> {
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) {