diff --git a/apps/desktop/src/components/editor/QueryEditor.vue b/apps/desktop/src/components/editor/QueryEditor.vue index d51bcdfb57..ca3f83169b 100644 --- a/apps/desktop/src/components/editor/QueryEditor.vue +++ b/apps/desktop/src/components/editor/QueryEditor.vue @@ -30,6 +30,7 @@ import { getSqlCompletionContext, getSqlCompletionResultValidFor, isSqlLikeCompletionStatement, + recordCompletionSelection, shouldAutoOpenSqlCompletion, extractCteDefinitions, } from "@/lib/sqlCompletion"; @@ -887,28 +888,49 @@ function buildCompletionResult(items: QueryCompletionItem[], from: number, valid return { from, filter: false, - options: items.map((item) => - (item.type === "snippet" || item.type === "function") && item.apply - ? codeMirrorSnippetCompletion(item.apply, { - label: item.label, - type: item.type, - detail: item.detail, - info: item.info, - boost: item.boost, - }) - : { - label: item.label, - type: item.type, - detail: item.detail, - info: item.info, - apply: item.apply, - boost: item.boost, - }, - ), + options: items.map((item) => completionOptionForItem(item)), validFor, }; } +function completionOptionForItem(item: QueryCompletionItem) { + const record = () => { + recordCompletionSelection(item.label, item.type); + }; + if ((item.type === "snippet" || item.type === "function") && item.apply) { + const completion = codeMirrorSnippetCompletion(item.apply, { + label: item.label, + type: item.type, + detail: item.detail, + info: item.info, + boost: item.boost, + }); + const originalApply = completion.apply; + return { + ...completion, + apply(view: EditorViewType, completionItem: unknown, from: number, to: number) { + record(); + if (typeof originalApply === "function") { + originalApply(view, completionItem as never, from, to); + } else { + view.dispatch({ changes: { from, to, insert: String(originalApply ?? item.label) } }); + } + }, + }; + } + return { + label: item.label, + type: item.type, + detail: item.detail, + info: item.info, + boost: item.boost, + apply(view: EditorViewType, _completionItem: unknown, from: number, to: number) { + record(); + view.dispatch({ changes: { from, to, insert: item.apply ?? item.label } }); + }, + }; +} + async function provideElasticsearchCompletions( currentState: import("@codemirror/state").EditorState, position: number, @@ -1000,6 +1022,16 @@ async function provideSqlCompletions( ); } + const localResult = buildLocalSqlCompletionResult(completionContext, fullDoc, position); + if (localResult) { + scheduleCompletionMetadataRefresh(completionContext); + if (!explicit) return localResult; + } + if (!explicit) { + scheduleCompletionMetadataRefresh(completionContext); + return null; + } + // Cancel any pending debounced completion if (completionDebounceTimer) { clearTimeout(completionDebounceTimer); @@ -1029,6 +1061,208 @@ async function provideSqlCompletions( } } +function buildLocalSqlCompletionResult( + completionContext: ReturnType, + fullDoc: string, + position: number, +) { + if (!props.connectionId || props.database == null) return null; + const shouldLoadTables = + completionContext.suggestTables || + (!!completionContext.qualifier && !isReferencedTableQualifier(completionContext)); + const tableLookupSchema = + completionContext.qualifier && completionContext.suggestTables ? completionContext.qualifier : props.schema; + const tableLookupFilter = + completionContext.qualifier && completionContext.suggestTables + ? completionContext.prefix + : completionContext.qualifier || completionContext.prefix; + const tables = shouldLoadTables + ? connectionStore.lookupLocalCompletionTables( + props.connectionId, + props.database, + tableLookupFilter, + MAX_COMPLETION_TABLES, + tableLookupSchema, + ) + : cachedTables; + + const shouldLoadObjects = + completionContext.suggestRoutines || + completionContext.exclusiveRoutineSuggestions || + (!!completionContext.qualifier && !completionContext.exclusiveColumnSuggestions); + const completionObjects = shouldLoadObjects + ? connectionStore.lookupLocalCompletionObjects( + props.connectionId, + props.database, + completionContext.qualifier || completionContext.prefix, + MAX_COMPLETION_TABLES, + completionContext.qualifier && !completionContext.exclusiveColumnSuggestions + ? completionContext.qualifier + : props.schema, + ) + : cachedCompletionObjects; + + const schemaNames = + completionContext.suggestTables && !completionContext.qualifier && !completionContext.insertTable + ? connectionStore.lookupLocalCompletionSchemas( + props.connectionId, + props.database, + completionContext.prefix, + MAX_COMPLETION_TABLES, + ) + : []; + + const columnsByTable = new Map(); + if (completionContext.insertTable) { + const insertSchema = completionContext.insertSchema ?? props.schema; + const insertColumns = connectionStore.lookupLocalCompletionColumns( + props.connectionId, + props.database, + completionContext.insertTable, + insertSchema, + ); + if (insertColumns.length > 0) { + columnsByTable.set( + insertSchema ? `${insertSchema}.${completionContext.insertTable}` : completionContext.insertTable, + insertColumns, + ); + } + } + + const cteDefs = extractCteDefinitions(fullDoc); + for (const refTable of completionContext.referencedTables) { + const cteDef = cteDefs.find((c) => c.name.toLowerCase() === refTable.name.toLowerCase()); + if (cteDef) { + columnsByTable.set( + refTable.name, + cteDef.columns.map((name) => ({ name, table: refTable.name, dataType: undefined })), + ); + continue; + } + const cacheKey = refTable.schema ? `${refTable.schema}.${refTable.name}` : refTable.name; + const cached = cachedColumnsByTable.get(cacheKey); + if (cached) { + columnsByTable.set(cacheKey, cached); + continue; + } + const localColumns = connectionStore.lookupLocalCompletionColumns( + props.connectionId, + props.database, + refTable.name, + refTable.schema ?? props.schema, + ); + if (localColumns.length > 0) { + columnsByTable.set(cacheKey, localColumns); + } + } + + if ( + tables.length === 0 && + completionObjects.length === 0 && + schemaNames.length === 0 && + columnsByTable.size === 0 && + (completionContext.exclusiveTableSuggestions || + completionContext.exclusiveColumnSuggestions || + completionContext.exclusiveRoutineSuggestions) + ) { + return null; + } + + const items = buildSqlCompletionItemsFromContext(completionContext, { + tables, + objects: completionObjects, + columnsByTable, + foreignKeysByTable: cachedForeignKeysByTable, + schemas: schemaNames, + translations: completionTranslations.value, + snippets: settingsStore.editorSettings.snippets, + dialect: props.dialect, + databaseType: props.databaseType, + }); + + return buildCompletionResult( + items, + position - completionContext.prefix.length, + getSqlCompletionResultValidFor(fullDoc, position), + ); +} + +function scheduleCompletionMetadataRefresh(completionContext: ReturnType) { + if (!props.connectionId || props.database == null) return; + const connectionId = props.connectionId; + const database = props.database; + const schema = + completionContext.qualifier && completionContext.suggestTables ? completionContext.qualifier : props.schema; + if ( + completionContext.suggestTables || + (!!completionContext.qualifier && !isReferencedTableQualifier(completionContext)) + ) { + void connectionStore + .refreshCompletionTables( + connectionId, + database, + completionContext.qualifier && !schema ? completionContext.qualifier : completionContext.prefix, + MAX_COMPLETION_TABLES, + schema, + ) + .then((tables) => { + cachedTables = mergeCompletionTables(cachedTables, tables); + }) + .catch(() => {}); + } + if ( + completionContext.suggestRoutines || + completionContext.exclusiveRoutineSuggestions || + (!!completionContext.qualifier && !completionContext.exclusiveColumnSuggestions) + ) { + void connectionStore + .refreshCompletionObjects(connectionId, database, completionContext.prefix, MAX_COMPLETION_TABLES, props.schema) + .then((objects) => { + cachedCompletionObjects = mergeCompletionObjects(cachedCompletionObjects, objects); + }) + .catch(() => {}); + } + if (completionContext.suggestTables && !completionContext.qualifier && !completionContext.insertTable) { + void connectionStore.refreshCompletionSchemas(connectionId, database).catch(() => {}); + } + if (completionContext.insertTable) { + const insertTable = completionContext.insertTable; + void connectionStore + .refreshCompletionColumns(connectionId, database, insertTable, completionContext.insertSchema ?? props.schema) + .then((columns) => { + const insertSchema = completionContext.insertSchema ?? props.schema; + cachedColumnsByTable.set(insertSchema ? `${insertSchema}.${insertTable}` : insertTable, columns); + }) + .catch(() => {}); + } + for (const refTable of completionContext.referencedTables) { + if (refTable.columns && refTable.columns.length > 0) continue; + const cacheKey = refTable.schema ? `${refTable.schema}.${refTable.name}` : refTable.name; + if (cachedColumnsByTable.has(cacheKey)) continue; + void connectionStore + .refreshCompletionColumns(connectionId, database, refTable.name, refTable.schema ?? props.schema) + .then((columns) => { + if (columns.length > 0) cachedColumnsByTable.set(cacheKey, columns); + }) + .catch(() => {}); + } +} + +function mergeCompletionTables( + existing: Array<{ name: string; schema?: string; type?: "table" | "view" }>, + incoming: Array<{ name: string; schema?: string; type?: "table" | "view" }>, +) { + const merged = [...existing]; + const seen = new Set(existing.map((table) => `${table.schema ?? ""}.${table.name}`.toLowerCase())); + for (const table of incoming) { + const key = `${table.schema ?? ""}.${table.name}`.toLowerCase(); + if (seen.has(key)) continue; + seen.add(key); + merged.push(table); + } + return merged; +} + async function performAsyncCompletionWithResult( epoch: number, completionContext: ReturnType, diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index b9411081f3..52bde92071 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -125,6 +125,10 @@ export const useConnectionStore = defineStore("connection", () => { const completionColumnsCache = ref>({}); const elasticsearchCompletionIndicesCache = ref>({}); const schemaListCache = ref>({}); + const completionTableIndex = new Map(); + const completionObjectIndex = new Map(); + const completionColumnIndex = new Map(); + const completionInFlight = new Map>(); const transferSource = ref<{ connectionId: string; database: string } | null>(null); const schemaDiffSource = ref<{ connectionId: string; database: string; schema?: string } | null>(null); const dataCompareSource = ref<{ @@ -633,6 +637,18 @@ export const useConnectionStore = defineStore("connection", () => { for (const key of Object.keys(elasticsearchCompletionIndicesCache.value)) { if (key === exactCacheKey || key.startsWith(cachePrefix)) delete elasticsearchCompletionIndicesCache.value[key]; } + for (const key of [...completionTableIndex.keys()]) { + if (key.startsWith(cachePrefix)) completionTableIndex.delete(key); + } + for (const key of [...completionObjectIndex.keys()]) { + if (key.startsWith(cachePrefix)) completionObjectIndex.delete(key); + } + for (const key of [...completionColumnIndex.keys()]) { + if (key.startsWith(cachePrefix)) completionColumnIndex.delete(key); + } + for (const key of [...completionInFlight.keys()]) { + if (key.startsWith(cachePrefix)) completionInFlight.delete(key); + } } async function removeConnections(ids: Iterable) { @@ -1678,14 +1694,185 @@ export const useConnectionStore = defineStore("connection", () => { } } + function completionScopeKey(connectionId: string, database: string, schema?: string): string { + return `${connectionId}:${database}:${schema ?? ""}`; + } + + function completionColumnsKey(connectionId: string, database: string, table: string, schema?: string): string { + return `${completionScopeKey(connectionId, database, schema)}:${table.toLowerCase()}`; + } + + function touchCompletionIndex( + index: Map, + key: string, + value: T, + max = COMPLETION_CACHE_MAX, + ) { + index.set(key, { ...value, touched: Date.now() }); + if (index.size <= max) return; + const oldest = [...index.entries()].sort(([, a], [, b]) => a.touched - b.touched).slice(0, index.size - max); + for (const [oldKey] of oldest) index.delete(oldKey); + } + + function withCompletionInFlight(key: string, load: () => Promise): Promise { + const existing = completionInFlight.get(key) as Promise | undefined; + if (existing) return existing; + const promise = load().finally(() => { + if (completionInFlight.get(key) === promise) completionInFlight.delete(key); + }); + completionInFlight.set(key, promise); + return promise; + } + + function tableMatchScore(table: SqlCompletionTable, filter: string, preferredSchema?: string): number { + const text = table.name.toLowerCase(); + const schema = table.schema?.toLowerCase(); + const normalized = filter.trim().toLowerCase(); + let score = schema && preferredSchema && schema === preferredSchema.toLowerCase() ? 10_000 : 0; + if (!normalized) return score; + if (text === normalized) return score + 9_000 - text.length; + if (text.startsWith(normalized)) return score + 7_000 - text.length; + if (text.includes(normalized)) return score + 4_000 - text.length; + let index = 0; + for (const ch of normalized) { + index = text.indexOf(ch, index); + if (index < 0) return -1; + index++; + } + return score + 1_000 - text.length; + } + + function objectMatchScore(object: SqlCompletionObject, filter: string, preferredSchema?: string): number { + const tableLike: SqlCompletionTable = { name: object.name, schema: object.schema }; + return tableMatchScore(tableLike, filter, preferredSchema); + } + + function indexCompletionTables( + connectionId: string, + database: string, + schema: string | undefined, + tables: SqlCompletionTable[], + ) { + const groups = new Map(); + for (const table of tables) { + const tableSchema = table.schema ?? schema; + const key = completionScopeKey(connectionId, database, tableSchema); + const list = groups.get(key) ?? []; + list.push({ ...table, schema: tableSchema }); + groups.set(key, list); + } + for (const [key, group] of groups) { + const previous = completionTableIndex.get(key)?.tables ?? []; + touchCompletionIndex(completionTableIndex, key, { + tables: dedupeCompletionTables([...previous, ...group]), + }); + } + } + + function indexCompletionObjects( + connectionId: string, + database: string, + schema: string | undefined, + objects: SqlCompletionObject[], + ) { + const groups = new Map(); + for (const object of objects) { + const objectSchema = object.schema ?? schema; + const key = completionScopeKey(connectionId, database, objectSchema); + const list = groups.get(key) ?? []; + list.push({ ...object, schema: objectSchema }); + groups.set(key, list); + } + for (const [key, group] of groups) { + const previous = completionObjectIndex.get(key)?.objects ?? []; + touchCompletionIndex(completionObjectIndex, key, { + objects: dedupeCompletionObjects([...previous, ...group]), + }); + } + } + + function indexCompletionColumns( + connectionId: string, + database: string, + table: string, + schema: string | undefined, + columns: SqlCompletionColumn[], + ) { + touchCompletionIndex(completionColumnIndex, completionColumnsKey(connectionId, database, table, schema), { + columns, + }); + } + + function lookupLocalCompletionTables( + connectionId: string, + database: string, + filter = "", + limit?: number, + schema?: string, + ): SqlCompletionTable[] { + const allScopes = [...completionTableIndex.entries()] + .filter(([key]) => key.startsWith(`${connectionId}:${database}:`)) + .map(([, entry]) => entry); + const preferred = schema ? completionTableIndex.get(completionScopeKey(connectionId, database, schema)) : undefined; + const scopes = preferred ? [preferred, ...allScopes.filter((entry) => entry !== preferred)] : allScopes; + const ranked = scopes + .flatMap((entry) => entry?.tables ?? []) + .map((table) => ({ table, score: tableMatchScore(table, filter, schema) })) + .filter((entry) => entry.score >= 0) + .sort((a, b) => b.score - a.score || a.table.name.localeCompare(b.table.name)); + return dedupeCompletionTables(ranked.map((entry) => entry.table)).slice(0, limit ?? 200); + } + + function lookupLocalCompletionObjects( + connectionId: string, + database: string, + filter = "", + limit?: number, + schema?: string, + ): SqlCompletionObject[] { + const allScopes = [...completionObjectIndex.entries()] + .filter(([key]) => key.startsWith(`${connectionId}:${database}:`)) + .map(([, entry]) => entry); + const preferred = schema + ? completionObjectIndex.get(completionScopeKey(connectionId, database, schema)) + : undefined; + const scopes = preferred ? [preferred, ...allScopes.filter((entry) => entry !== preferred)] : allScopes; + const ranked = scopes + .flatMap((entry) => entry?.objects ?? []) + .map((object) => ({ object, score: objectMatchScore(object, filter, schema) })) + .filter((entry) => entry.score >= 0) + .sort((a, b) => b.score - a.score || a.object.name.localeCompare(b.object.name)); + return dedupeCompletionObjects(ranked.map((entry) => entry.object)).slice(0, limit ?? 200); + } + + function lookupLocalCompletionSchemas(connectionId: string, database: string, filter = "", limit = 50): string[] { + const schemas = schemaListCache.value[`${connectionId}:${database}`] ?? []; + const normalized = filter.trim().toLowerCase(); + return schemas + .filter((schema) => fuzzyTextMatch(schema, normalized)) + .sort((a, b) => tableMatchScore({ name: b }, normalized) - tableMatchScore({ name: a }, normalized)) + .slice(0, limit); + } + + function lookupLocalCompletionColumns( + connectionId: string, + database: string, + table: string, + schema?: string, + ): SqlCompletionColumn[] { + return completionColumnIndex.get(completionColumnsKey(connectionId, database, table, schema))?.columns ?? []; + } + async function listCompletionSchemas(connectionId: string, database: string): Promise { const cacheKey = `${connectionId}:${database}`; if (schemaListCache.value[cacheKey]) { return schemaListCache.value[cacheKey]; } - const schemas = await api.listSchemas(connectionId, database); - schemaListCache.value[cacheKey] = schemas; - return schemas; + return withCompletionInFlight(`${cacheKey}:schemas`, async () => { + const schemas = await api.listSchemas(connectionId, database); + schemaListCache.value[cacheKey] = schemas; + return schemas; + }); } async function listElasticsearchCompletionIndices(connectionId: string, database: string): Promise { @@ -1714,46 +1901,21 @@ export const useConnectionStore = defineStore("connection", () => { return completionTablesCache.value[cacheKey]; } - await ensureConnected(connectionId); + return withCompletionInFlight(`${cacheKey}:tables`, async () => { + await ensureConnected(connectionId); - if (isSchemaAwareDatabase(connectionId)) { - const schemas = schema ? [schema] : await listCompletionSchemas(connectionId, database); - if (normalizedFilter || limit) { - const batchSize = 5; - const results: SqlCompletionTable[] = []; - for (let i = 0; i < schemas.length && results.length < (limit ?? Infinity); i += batchSize) { - const batch = schemas.slice(i, i + batchSize); - const batchResults = await Promise.all( - batch.map(async (s) => { - try { - const tables = await api.listTables(connectionId, database, s, normalizedFilter, limit); - return tables.map((table) => ({ - name: table.name, - schema: s, - type: table.table_type === "VIEW" ? ("view" as const) : ("table" as const), - })) as SqlCompletionTable[]; - } catch { - return [] as SqlCompletionTable[]; - } - }), - ); - for (const group of batchResults) { - results.push(...group); - } - } - if (results.length === 0 && relaxedFilter) { - for (let i = 0; i < schemas.length && results.length < (limit ?? Infinity); i += batchSize) { + if (isSchemaAwareDatabase(connectionId)) { + const schemas = schema ? [schema] : await listCompletionSchemas(connectionId, database); + if (normalizedFilter || limit) { + const batchSize = 5; + const results: SqlCompletionTable[] = []; + const maxResults = limit ?? Infinity; + for (let i = 0; i < schemas.length && results.length < maxResults; i += batchSize) { const batch = schemas.slice(i, i + batchSize); const batchResults = await Promise.all( batch.map(async (s) => { try { - const tables = await api.listTables( - connectionId, - database, - s, - relaxedFilter, - expandedCompletionLimit(limit), - ); + const tables = await api.listTables(connectionId, database, s, normalizedFilter, limit); return tables.map((table) => ({ name: table.name, schema: s, @@ -1766,46 +1928,80 @@ export const useConnectionStore = defineStore("connection", () => { ); for (const group of batchResults) { results.push(...group); + indexCompletionTables(connectionId, database, undefined, group); } } + if (results.length === 0 && relaxedFilter) { + for (let i = 0; i < schemas.length && results.length < maxResults; i += batchSize) { + const batch = schemas.slice(i, i + batchSize); + const batchResults = await Promise.all( + batch.map(async (s) => { + try { + const tables = await api.listTables( + connectionId, + database, + s, + relaxedFilter, + expandedCompletionLimit(limit), + ); + return tables.map((table) => ({ + name: table.name, + schema: s, + type: table.table_type === "VIEW" ? ("view" as const) : ("table" as const), + })) as SqlCompletionTable[]; + } catch { + return [] as SqlCompletionTable[]; + } + }), + ); + for (const group of batchResults) { + results.push(...group); + indexCompletionTables(connectionId, database, undefined, group); + } + } + } + const limitedTables = limit ? dedupeCompletionTables(results).slice(0, limit) : results; + completionTablesCache.value[cacheKey] = limitedTables; + indexCompletionTables(connectionId, database, schema, limitedTables); + evictOldestCacheEntries(completionTablesCache.value, COMPLETION_CACHE_MAX); + return completionTablesCache.value[cacheKey]; } - const limitedTables = limit - ? dedupeCompletionTables(results).slice(0, expandedCompletionLimit(limit)) - : results; - completionTablesCache.value[cacheKey] = limitedTables; + + const tableGroups = await Promise.all( + schemas.map(async (schema) => { + try { + const tables = await api.listTables(connectionId, database, schema); + return tables.map((table) => ({ + name: table.name, + schema, + type: table.table_type === "VIEW" ? ("view" as const) : ("table" as const), + })); + } catch { + return []; + } + }), + ); + completionTablesCache.value[cacheKey] = tableGroups.flat(); + indexCompletionTables(connectionId, database, schema, completionTablesCache.value[cacheKey]); evictOldestCacheEntries(completionTablesCache.value, COMPLETION_CACHE_MAX); return completionTablesCache.value[cacheKey]; } - const tableGroups = await Promise.all( - schemas.map(async (schema) => { - try { - const tables = await api.listTables(connectionId, database, schema); - return tables.map((table) => ({ - name: table.name, - schema, - type: table.table_type === "VIEW" ? ("view" as const) : ("table" as const), - })); - } catch { - return []; - } - }), - ); - completionTablesCache.value[cacheKey] = tableGroups.flat(); + let tables = await api.listTables(connectionId, database, database, normalizedFilter, limit); + if (tables.length === 0 && relaxedFilter) { + tables = await api.listTables(connectionId, database, database, relaxedFilter, expandedCompletionLimit(limit)); + } + completionTablesCache.value[cacheKey] = tables.map((table) => ({ + name: table.name, + type: table.table_type === "VIEW" ? ("view" as const) : ("table" as const), + })); + completionTablesCache.value[cacheKey] = limit + ? completionTablesCache.value[cacheKey].slice(0, limit) + : completionTablesCache.value[cacheKey]; + indexCompletionTables(connectionId, database, schema, completionTablesCache.value[cacheKey]); evictOldestCacheEntries(completionTablesCache.value, COMPLETION_CACHE_MAX); return completionTablesCache.value[cacheKey]; - } - - let tables = await api.listTables(connectionId, database, database, normalizedFilter, limit); - if (tables.length === 0 && relaxedFilter) { - tables = await api.listTables(connectionId, database, database, relaxedFilter, expandedCompletionLimit(limit)); - } - completionTablesCache.value[cacheKey] = tables.map((table) => ({ - name: table.name, - type: table.table_type === "VIEW" ? ("view" as const) : ("table" as const), - })); - evictOldestCacheEntries(completionTablesCache.value, COMPLETION_CACHE_MAX); - return completionTablesCache.value[cacheKey]; + }); } function relaxedCompletionTableFilter(filter: string): string | undefined { @@ -1840,21 +2036,24 @@ export const useConnectionStore = defineStore("connection", () => { const normalizedFilter = filter.trim().toLowerCase(); const cacheKey = `${connectionId}:${database}:${schema ?? ""}`; if (!completionObjectsCache.value[cacheKey]) { - await ensureConnected(connectionId); - const objects = isSchemaAwareDatabase(connectionId) - ? await listSchemaAwareCompletionObjects(connectionId, database, schema) - : await api.listCompletionObjects(connectionId, database, schema || database); - completionObjectsCache.value[cacheKey] = dedupeCompletionObjects( - objects.map(toSqlCompletionObject).filter((object): object is SqlCompletionObject => object != null), - ); - evictOldestCacheEntries(completionObjectsCache.value, COMPLETION_CACHE_MAX); + await withCompletionInFlight(`${cacheKey}:objects`, async () => { + await ensureConnected(connectionId); + const objects = isSchemaAwareDatabase(connectionId) + ? await listSchemaAwareCompletionObjects(connectionId, database, schema) + : await api.listCompletionObjects(connectionId, database, schema || database); + completionObjectsCache.value[cacheKey] = dedupeCompletionObjects( + objects.map(toSqlCompletionObject).filter((object): object is SqlCompletionObject => object != null), + ); + indexCompletionObjects(connectionId, database, schema, completionObjectsCache.value[cacheKey]); + evictOldestCacheEntries(completionObjectsCache.value, COMPLETION_CACHE_MAX); + }); } const objects = completionObjectsCache.value[cacheKey]; const filtered = normalizedFilter ? objects.filter((object) => fuzzyCompletionObjectMatch(object, normalizedFilter)) : objects; - return typeof limit === "number" ? filtered.slice(0, expandedCompletionLimit(limit)) : filtered; + return typeof limit === "number" ? filtered.slice(0, limit) : filtered; } async function listSchemaAwareCompletionObjects( @@ -1944,13 +2143,15 @@ export const useConnectionStore = defineStore("connection", () => { } const cacheKey = `${connectionId}:${database}:${schema || ""}:${table}`; if (!completionColumnsCache.value[cacheKey]) { - await ensureConnected(connectionId); - const querySchema = metadataQuerySchema(connectionId, database, schema); - completionColumnsCache.value[cacheKey] = await api.getColumns(connectionId, database, querySchema, table); - evictOldestCacheEntries(completionColumnsCache.value, COMPLETION_CACHE_MAX); + await withCompletionInFlight(`${cacheKey}:columns`, async () => { + await ensureConnected(connectionId); + const querySchema = metadataQuerySchema(connectionId, database, schema); + completionColumnsCache.value[cacheKey] = await api.getColumns(connectionId, database, querySchema, table); + evictOldestCacheEntries(completionColumnsCache.value, COMPLETION_CACHE_MAX); + }); } - return completionColumnsCache.value[cacheKey].map((column) => ({ + const columns = completionColumnsCache.value[cacheKey].map((column) => ({ name: column.name, table, schema, @@ -1958,6 +2159,41 @@ export const useConnectionStore = defineStore("connection", () => { isNullable: column.is_nullable, comment: column.comment, })); + indexCompletionColumns(connectionId, database, table, schema, columns); + return columns; + } + + function refreshCompletionTables( + connectionId: string, + database: string, + filter = "", + limit?: number, + schema?: string, + ): Promise { + return listCompletionTables(connectionId, database, filter, limit, schema); + } + + function refreshCompletionObjects( + connectionId: string, + database: string, + filter = "", + limit?: number, + schema?: string, + ): Promise { + return listCompletionObjects(connectionId, database, filter, limit, schema); + } + + function refreshCompletionSchemas(connectionId: string, database: string): Promise { + return listCompletionSchemas(connectionId, database); + } + + function refreshCompletionColumns( + connectionId: string, + database: string, + table: string, + schema?: string, + ): Promise { + return listCompletionColumns(connectionId, database, table, schema); } function findNode(nodes: TreeNode[], id: string): TreeNode | null { @@ -2349,6 +2585,14 @@ export const useConnectionStore = defineStore("connection", () => { listCompletionObjects, listCompletionColumns, listCompletionSchemas, + lookupLocalCompletionTables, + lookupLocalCompletionObjects, + lookupLocalCompletionColumns, + lookupLocalCompletionSchemas, + refreshCompletionTables, + refreshCompletionObjects, + refreshCompletionColumns, + refreshCompletionSchemas, listElasticsearchCompletionIndices, exportConnectionsToFile, readImportFile, diff --git a/packages/app-tests/sqlCompletionPerformance.test.ts b/packages/app-tests/sqlCompletionPerformance.test.ts new file mode 100644 index 0000000000..bbb84d9f44 --- /dev/null +++ b/packages/app-tests/sqlCompletionPerformance.test.ts @@ -0,0 +1,142 @@ +import assert from "node:assert/strict"; +import { test } from "vitest"; +import { createPinia, setActivePinia } from "pinia"; +import { buildSqlCompletionItems, recordCompletionSelection } from "../../apps/desktop/src/lib/sqlCompletion.ts"; +import { useConnectionStore } from "../../apps/desktop/src/stores/connectionStore.ts"; +import type { ConnectionConfig, TableInfo } from "../../apps/desktop/src/types/database.ts"; + +function installMemoryStorage() { + const values = new Map(); + const original = Object.getOwnPropertyDescriptor(globalThis, "localStorage"); + Object.defineProperty(globalThis, "localStorage", { + configurable: true, + value: { + getItem: (key: string) => values.get(key) ?? null, + setItem: (key: string, value: string) => values.set(key, value), + removeItem: (key: string) => values.delete(key), + clear: () => values.clear(), + }, + }); + return { + restore() { + if (original) Object.defineProperty(globalThis, "localStorage", original); + else Reflect.deleteProperty(globalThis, "localStorage"); + }, + }; +} + +function postgresConnection(): ConnectionConfig { + return { + id: "conn-large", + name: "Large Postgres", + db_type: "postgres", + host: "127.0.0.1", + port: 5432, + username: "postgres", + password: "secret", + }; +} + +function largeCatalogFixture(schemaCount = 120, tablesPerSchema = 25) { + const schemas = Array.from({ length: schemaCount }, (_, index) => `schema_${String(index).padStart(3, "0")}`); + const tablesBySchema = new Map(); + for (const schema of schemas) { + tablesBySchema.set( + schema, + Array.from({ length: tablesPerSchema }, (_, index) => ({ + name: `customer_event_${schema}_${String(index).padStart(3, "0")}`, + table_type: index % 7 === 0 ? "VIEW" : "BASE TABLE", + })), + ); + } + return { schemas, tablesBySchema }; +} + +test("large table catalogs produce bounded SQL completion items", () => { + const tables = Array.from({ length: 2500 }, (_, index) => ({ + name: `customer_event_${String(index).padStart(4, "0")}`, + schema: `schema_${index % 100}`, + type: "table" as const, + })); + + const items = buildSqlCompletionItems("select * from customer", "select * from customer".length, { + tables, + columnsByTable: new Map(), + dialect: "postgres", + }); + + const tableItems = items.filter((item) => item.type === "table"); + assert.equal(tableItems.length, 200); + assert.equal(tableItems[0].label, "customer_event_0000"); +}); + +test("recently selected completion items receive a ranking boost", () => { + const tables = [ + { name: "customer_accounts", schema: "public", type: "table" as const }, + { name: "customer_accounts_archive", schema: "public", type: "table" as const }, + ]; + + const before = buildSqlCompletionItems("select * from customer_acc", "select * from customer_acc".length, { + tables, + columnsByTable: new Map(), + dialect: "postgres", + }).filter((item) => item.type === "table"); + assert.equal(before[0].label, "customer_accounts"); + + recordCompletionSelection("customer_accounts_archive", "table"); + const after = buildSqlCompletionItems("select * from customer_acc", "select * from customer_acc".length, { + tables, + columnsByTable: new Map(), + dialect: "postgres", + }).filter((item) => item.type === "table"); + assert.equal(after[0].label, "customer_accounts_archive"); +}); + +test("completion metadata refresh is deduplicated and local lookups rank the current schema first", async () => { + const originalFetch = globalThis.fetch; + const storage = installMemoryStorage(); + const { schemas, tablesBySchema } = largeCatalogFixture(); + let schemaCalls = 0; + let tableCalls = 0; + + globalThis.fetch = (async (input) => { + const url = new URL(String(input), "http://localhost"); + if (url.pathname === "/api/schema/schemas") { + schemaCalls++; + return Response.json(schemas); + } + if (url.pathname === "/api/schema/tables") { + tableCalls++; + const schema = url.searchParams.get("schema") ?? ""; + const filter = (url.searchParams.get("filter") ?? "").toLowerCase(); + const limit = Number(url.searchParams.get("limit") || "0") || undefined; + const tables = (tablesBySchema.get(schema) ?? []) + .filter((table) => !filter || table.name.toLowerCase().includes(filter)) + .slice(0, limit); + return Response.json(tables); + } + return Response.json(null); + }) as typeof fetch; + + try { + setActivePinia(createPinia()); + const store = useConnectionStore(); + store.addEphemeralConnection(postgresConnection()); + + await Promise.all([ + store.refreshCompletionTables("conn-large", "app", "customer", 20, "schema_010"), + store.refreshCompletionTables("conn-large", "app", "customer", 20, "schema_010"), + ]); + + assert.equal(tableCalls, 1); + assert.equal(schemaCalls, 0); + + await store.refreshCompletionTables("conn-large", "app", "customer", 20, "schema_011"); + const localTables = store.lookupLocalCompletionTables("conn-large", "app", "customer", 5, "schema_011"); + assert.equal(localTables.length, 5); + assert.equal(localTables[0].schema, "schema_011"); + } finally { + globalThis.fetch = originalFetch; + storage.restore(); + } +});