github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-21 10:28:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(completion): optimize qualified SQL completion queries

Browse Source
This commit is contained in:
miracle
2026-08-22 15:41:13 +08:00
committed by GitHub
parent fe269a28db
commit 3442acd52c
6 changed files with 393 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/desktop/src/components/editor/QueryEditor.vue
+118 -26
View File
@@ -334,7 +334,6 @@ const completionTranslations = computed(() => ({
const MAX_COMPLETION_TABLES = 200;
const PRESTO_ON_DEMAND_TABLE_COMPLETION_MIN_PREFIX = 2;
const PRESTO_ON_DEMAND_TABLE_COMPLETION_LIMIT = 20;
const MAX_JOIN_FK_PREFETCH_TABLES = 24;
const MAX_SEMANTIC_DIAGNOSTIC_COLUMN_TABLES = 4;
const liveFontSize = ref(settingsStore.editorSettings.fontSize);
const gestureStartFontSize = ref(settingsStore.editorSettings.fontSize);
@@ -617,6 +616,7 @@ let cachedTables: SqlCompletionTable[] = [];
const cachedCompletionObjectsByScope = new Map<string, SqlCompletionObject[]>();
// Persistent column cache keyed by "schema.table" or "table"
const cachedColumnsByTable = new Map<string, SqlCompletionColumn[]>();
const cachedPrefixColumnsByTable = new Map<string, SqlCompletionColumn[]>();
const cachedInsertValueHintColumnsByTable = new Map<string, string[]>();
const cachedForeignKeysByTable = new Map<string, SqlCompletionForeignKey[]>();
const loadedColumnsByTable = new Set<string>();
@@ -650,10 +650,19 @@ function completionColumnRequestContext(reference?: Pick<SqlCompletionReferenced
};
}
async function listCompletionColumnsForEditor(connectionId: string, database: string, table: string, schema?: string, catalog = props.catalog, reference?: Pick<SqlCompletionReferencedTable, "nameQuoted" | "schemaQuoted">) {
async function listCompletionColumnsForEditor(connectionId: string, database: string, table: string, schema?: string, catalog = props.catalog, reference?: Pick<SqlCompletionReferencedTable, "nameQuoted" | "schemaQuoted">, prefix?: string) {
const requestedVersion = props.completionContextVersion;
const sessionScoped = usesOracleSessionCompletionColumns(schema);
const columns = await connectionStore.listCompletionColumns(connectionId, database, table, schema, completionColumnRequestContext(reference), catalog);
let columns: SqlCompletionColumn[];
if (prefix && prefix.length >= 2 && (props.databaseType === "postgres" || props.databaseType === "mysql")) {
try {
columns = await connectionStore.listCompletionColumnsByPrefix(connectionId, database, table, schema, prefix, catalog, completionColumnRequestContext(reference));
} catch {
columns = await connectionStore.listCompletionColumns(connectionId, database, table, schema, completionColumnRequestContext(reference), catalog);
}
} else {
columns = await connectionStore.listCompletionColumns(connectionId, database, table, schema, completionColumnRequestContext(reference), catalog);
}
if (sessionScoped && requestedVersion !== props.completionContextVersion) throw new Error("Stale Oracle completion context");
return columns;
}
@@ -2156,11 +2165,35 @@ function identifierRangeAt(sql: string, pos: number): { from: number; to: number
type CompletionMetadataScope = Pick<SqlCompletionScope, "database" | "schema">;
function completionCacheKey(table: { name: string; catalog?: string | null; database?: string | null; schema?: string | null }, scope?: CompletionMetadataScope) {
function completionCacheKey(table: { name: string; catalog?: string | null; database?: string | null; schema?: string | null; nameQuoted?: boolean; schemaQuoted?: boolean }, scope?: CompletionMetadataScope) {
const schema = table.schema ?? scope?.schema ?? props.schema;
const scopedDatabase = scope && scope.database !== props.database ? scope.database : undefined;
const database = supportsDatabaseSchemaQualifierCompletion() ? (table.database ?? scopedDatabase) : undefined;
return schema ? `${database ? `${database}.` : ""}${schema}.${table.name}` : table.name;
const baseKey = schema ? `${database ? `${database}.` : ""}${schema}.${table.name}` : table.name;
if (props.databaseType !== "postgres" || (!table.nameQuoted && !table.schemaQuoted)) return baseKey;
return `${baseKey}:quoted:s=${table.schemaQuoted ? "1" : "0"}:t=${table.nameQuoted ? "1" : "0"}`;
}
function completionPrefixCacheKey(table: { name: string; catalog?: string | null; database?: string | null; schema?: string | null; nameQuoted?: boolean; schemaQuoted?: boolean }, scope: CompletionMetadataScope | undefined, prefix: string) {
return `${completionCacheKey(table, scope)}:prefix:${prefix.trim().toLowerCase()}`;
}
function lookupCachedPrefixColumns(table: { name: string; catalog?: string | null; database?: string | null; schema?: string | null; nameQuoted?: boolean; schemaQuoted?: boolean }, scope: CompletionMetadataScope | undefined, prefix: string): SqlCompletionColumn[] | undefined {
const normalizedPrefix = prefix.trim().toLowerCase();
const exactKey = completionPrefixCacheKey(table, scope, normalizedPrefix);
const exact = cachedPrefixColumnsByTable.get(exactKey);
if (exact) return exact;
const marker = `${completionCacheKey(table, scope)}:prefix:`;
let best: { prefix: string; columns: SqlCompletionColumn[] } | undefined;
for (const [key, columns] of cachedPrefixColumnsByTable) {
if (!key.startsWith(marker)) continue;
const cachedPrefix = key.slice(marker.length);
if (!normalizedPrefix.startsWith(cachedPrefix) || (best && cachedPrefix.length <= best.prefix.length)) continue;
best = { prefix: cachedPrefix, columns };
}
if (!best) return undefined;
return best.columns.filter((column) => column.name.toLowerCase().startsWith(normalizedPrefix));
}
const pendingInsertValueHintColumnLoads = new Set<string>();
@@ -3165,6 +3198,23 @@ function unregisterTableReferenceDropListener() {
}
let completionEpoch = 0;
let tableCompletionRefreshActive = false;
let latestTableCompletionRefresh: (() => Promise<void>) | null = null;
function queueTableCompletionRefresh(task: () => Promise<void>): void {
latestTableCompletionRefresh = task;
if (tableCompletionRefreshActive) return;
tableCompletionRefreshActive = true;
void (async () => {
while (latestTableCompletionRefresh) {
const next = latestTableCompletionRefresh;
latestTableCompletionRefresh = null;
await next();
}
})().finally(() => {
tableCompletionRefreshActive = false;
});
}
let completionDebounceTimer: ReturnType<typeof setTimeout> | null = null;
let typedCompletionActivationUntil = 0;
let suppressNextSqlCompletionAutoStartUntil = 0;
@@ -3837,13 +3887,25 @@ function buildLocalSqlCompletionResult(completionContext: ReturnType<typeof getS
const qualifiedColumnTarget = completionQualifiedTableTarget(completionContext);
if (qualifiedColumnTarget) {
const cacheKey = completionCacheKey(qualifiedColumnTarget, scope);
const cached = cachedColumnsByTable.get(cacheKey);
const reference = completionContext.referencedTables.find((table) => completionTablesMatch(table, qualifiedColumnTarget));
const qualifiedCacheTable = { ...qualifiedColumnTarget, nameQuoted: reference?.nameQuoted, schemaQuoted: reference?.schemaQuoted };
const cacheKey = completionCacheKey(qualifiedCacheTable, scope);
const cachedPrefix = completionContext.prefix.length >= 2 && (props.databaseType === "postgres" || props.databaseType === "mysql") ? lookupCachedPrefixColumns(qualifiedCacheTable, scope, completionContext.prefix) : undefined;
const cached = cachedPrefix ?? cachedColumnsByTable.get(cacheKey);
if (cached) {
columnsByTable.set(cacheKey, cached);
} else {
const target = completionMetadataTarget(qualifiedColumnTarget, scope);
const localColumns = target && !usesOracleSessionCompletionColumns(target.schema) ? connectionStore.lookupLocalCompletionColumns(props.connectionId, target.database, qualifiedColumnTarget.name, target.schema, target.catalog) : [];
const prefixColumns =
target && completionContext.prefix.length >= 2 && (props.databaseType === "postgres" || props.databaseType === "mysql")
? connectionStore.lookupLocalCompletionColumnsByPrefix(props.connectionId, target.database, qualifiedColumnTarget.name, target.schema, completionContext.prefix, target.catalog, completionColumnRequestContext(reference))
: [];
const localColumns =
prefixColumns.length > 0
? prefixColumns
: target && !usesOracleSessionCompletionColumns(target.schema)
? connectionStore.lookupLocalCompletionColumns(props.connectionId, target.database, qualifiedColumnTarget.name, target.schema, target.catalog, completionColumnRequestContext(reference))
: [];
if (localColumns.length > 0) {
columnsByTable.set(cacheKey, localColumns);
}
@@ -3866,13 +3928,22 @@ function buildLocalSqlCompletionResult(completionContext: ReturnType<typeof getS
continue;
}
const cacheKey = completionCacheKey(refTable, scope);
const cached = cachedColumnsByTable.get(cacheKey);
const prefixCompletion =
(props.databaseType === "postgres" || props.databaseType === "mysql") &&
completionContext.qualifier &&
completionContext.prefix.length >= 2 &&
isReferencedTableQualifier(completionContext) &&
(refTable.alias?.toLowerCase() === completionContext.qualifier.toLowerCase() || refTable.name.toLowerCase() === completionContext.qualifier.toLowerCase())
? completionContext.prefix
: undefined;
const cached = (prefixCompletion ? lookupCachedPrefixColumns(refTable, scope, prefixCompletion) : undefined) ?? cachedColumnsByTable.get(cacheKey);
if (cached) {
columnsByTable.set(cacheKey, cached);
continue;
}
const target = completionMetadataTarget(refTable, scope);
const localColumns = target && !usesOracleSessionCompletionColumns(target.schema) ? connectionStore.lookupLocalCompletionColumns(props.connectionId, target.database, refTable.name, target.schema, target.catalog, refTable) : [];
const prefixColumns = target && prefixCompletion ? connectionStore.lookupLocalCompletionColumnsByPrefix(props.connectionId, target.database, refTable.name, target.schema, prefixCompletion, target.catalog, refTable) : [];
const localColumns = prefixColumns.length > 0 ? prefixColumns : target && !usesOracleSessionCompletionColumns(target.schema) ? connectionStore.lookupLocalCompletionColumns(props.connectionId, target.database, refTable.name, target.schema, target.catalog, refTable) : [];
if (localColumns.length > 0) {
columnsByTable.set(cacheKey, localColumns);
}
@@ -3938,16 +4009,21 @@ function scheduleCompletionMetadataRefresh(completionContext: ReturnType<typeof
});
if (!localOnlyMetadata && !schemaLookupDatabase && (completionContext.suggestTables || (!!completionContext.qualifier && !isReferencedTableQualifier(completionContext)))) {
const globalOracleTableSearch = props.databaseType === "oracle" && completionContext.suggestTables && !completionContext.qualifier;
void connectionStore
.refreshCompletionTables(connectionId, tableLookupTarget.database, tableLookupTarget.filter, MAX_COMPLETION_TABLES, tableLookupTarget.schema, globalOracleTableSearch, scope.schema, props.catalog)
.then((tables) => {
const refreshEpoch = completionEpoch;
queueTableCompletionRefresh(async () => {
if (refreshEpoch !== completionEpoch) return;
try {
const tables = await connectionStore.refreshCompletionTables(connectionId, tableLookupTarget.database, tableLookupTarget.filter, MAX_COMPLETION_TABLES, tableLookupTarget.schema, globalOracleTableSearch, scope.schema, props.catalog);
if (refreshEpoch !== completionEpoch) return;
const scopedTables = tables.map((table) => ({ ...table, database: table.database ?? tableLookupTarget.database }));
cachedTables = mergeCompletionTables(cachedTables, scopedTables);
if (completionContext.suggestTables && completionContext.referencedTables.length > 0) {
void ensureForeignKeysForTables([...completionContext.referencedTables, ...scopedTables.slice(0, MAX_JOIN_FK_PREFETCH_TABLES)]);
if (completionContext.suggestJoinConditions && completionContext.referencedTables.length > 0) {
void ensureForeignKeysForTables(completionContext.referencedTables);
}
})
.catch(() => {});
} catch {
// Local candidates remain available when the remote refresh fails.
}
});
}
if (!localOnlyMetadata && shouldLoadCompletionObjects(completionContext)) {
const completionObjectScope = routineCompletionScopeForContext(completionContext, scope);
@@ -4009,7 +4085,7 @@ function scheduleCompletionMetadataRefresh(completionContext: ReturnType<typeof
.catch(() => {});
}
}
if (!tableNameCompletion && (completionContext.suggestTables || completionContext.suggestJoinConditions) && completionContext.referencedTables.length > 0) {
if (!tableNameCompletion && completionContext.suggestJoinConditions && completionContext.referencedTables.length > 0) {
void ensureForeignKeysForTables(completionContext.referencedTables);
}
}
@@ -4278,20 +4354,27 @@ async function performAsyncCompletionWithResult(epoch: number, completionContext
const tableNameCompletion = isTableNameCompletionContext(completionContext);
const shouldFetchColumnsForCompletion = !tableNameCompletion && (!onDemandOnlyColumns || completionContext.suggestColumns || completionContext.exclusiveColumnSuggestions || !!completionContext.insertTable);
const hasQualifiedColumnPrefix = (props.databaseType === "postgres" || props.databaseType === "mysql") && completionContext.qualifier && completionContext.prefix.length >= 2 && isReferencedTableQualifier(completionContext);
const columnRefs = hasQualifiedColumnPrefix
? refs.filter((refTable) => refTable.alias?.toLowerCase() === completionContext.qualifier!.toLowerCase() || refTable.name.toLowerCase() === completionContext.qualifier!.toLowerCase() || (!!qualifiedColumnTarget && completionTablesMatch(refTable, qualifiedColumnTarget)))
: refs.slice(0, 4);
if (shouldFetchColumnsForCompletion) {
await Promise.all(
refs.map(async (refTable) => {
columnRefs.map(async (refTable) => {
if (isVirtualCompletionTableReference(refTable)) return;
if (refTable.columns && refTable.columns.length > 0) return;
const cacheKey = completionCacheKey(refTable, scope);
if (cachedColumnsByTable.has(cacheKey)) return;
const prefixCompletion = hasQualifiedColumnPrefix ? completionContext.prefix : undefined;
const prefixCacheKey = prefixCompletion ? completionPrefixCacheKey(refTable, scope, prefixCompletion) : undefined;
if (prefixCompletion ? lookupCachedPrefixColumns(refTable, scope, prefixCompletion) : cachedColumnsByTable.has(cacheKey)) return;
try {
const target = completionMetadataTarget(refTable, scope);
if (!target) return;
const columns = await listCompletionColumnsForEditor(props.connectionId!, target.database, refTable.name, target.schema, target.catalog, refTable);
const columns = await listCompletionColumnsForEditor(props.connectionId!, target.database, refTable.name, target.schema, target.catalog, refTable, prefixCompletion);
if (epoch !== completionEpoch) return;
if (columns.length === 0) return;
cachedColumnsByTable.set(cacheKey, columns);
if (prefixCacheKey) cachedPrefixColumnsByTable.set(prefixCacheKey, columns);
else cachedColumnsByTable.set(cacheKey, columns);
} catch (e) {
console.error(`[DBX] Failed to load columns for ${cacheKey}:`, e);
}
@@ -4300,9 +4383,8 @@ async function performAsyncCompletionWithResult(epoch: number, completionContext
}
if (epoch !== completionEpoch) return null;
if (!tableNameCompletion && (completionContext.suggestTables || completionContext.suggestJoinConditions) && refs.length > 0) {
const fkPrefetchTables = completionContext.suggestTables ? [...refs, ...tables.slice(0, MAX_JOIN_FK_PREFETCH_TABLES)] : refs;
await ensureForeignKeysForTables(fkPrefetchTables.filter((table) => !("columns" in table) || !table.columns || table.columns.length === 0));
if (!tableNameCompletion && completionContext.suggestJoinConditions && refs.length > 0) {
await ensureForeignKeysForTables(refs.filter((table) => !("columns" in table) || !table.columns || table.columns.length === 0));
if (epoch !== completionEpoch) return null;
}
@@ -4315,6 +4397,7 @@ async function performAsyncCompletionWithResult(epoch: number, completionContext
}
} else {
for (const refTable of refs) {
if (hasQualifiedColumnPrefix && !columnRefs.includes(refTable)) continue;
if (refTable.columns && refTable.columns.length > 0) {
const key = refTable.name;
columnsByTable.set(
@@ -4328,7 +4411,15 @@ async function performAsyncCompletionWithResult(epoch: number, completionContext
continue;
}
const cacheKey = completionCacheKey(refTable, scope);
const cached = cachedColumnsByTable.get(cacheKey);
const prefixCompletion =
(props.databaseType === "postgres" || props.databaseType === "mysql") &&
completionContext.qualifier &&
completionContext.prefix.length >= 2 &&
isReferencedTableQualifier(completionContext) &&
(refTable.alias?.toLowerCase() === completionContext.qualifier.toLowerCase() || refTable.name.toLowerCase() === completionContext.qualifier.toLowerCase())
? completionContext.prefix
: undefined;
const cached = (prefixCompletion ? lookupCachedPrefixColumns(refTable, scope, prefixCompletion) : undefined) ?? cachedColumnsByTable.get(cacheKey);
if (cached) {
columnsByTable.set(cacheKey, cached);
}
@@ -4446,6 +4537,7 @@ function refreshCompletionCache() {
cachedTables = [];
cachedCompletionObjectsByScope.clear();
cachedColumnsByTable.clear();
cachedPrefixColumnsByTable.clear();
cachedInsertValueHintColumnsByTable.clear();
loadedColumnsByTable.clear();
cachedForeignKeysByTable.clear();
apps/desktop/src/lib/__tests__/editor/queryEditorColumnPrefixCache.spec.ts
Vendored
+13
View File
@@ -0,0 +1,13 @@
import { readFileSync } from "node:fs";
import { describe, expect, it } from "vitest";
const queryEditorSource = readFileSync(new URL("../../../components/editor/QueryEditor.vue", import.meta.url), "utf8");
describe("QueryEditor column prefix cache", () => {
it("hydrates qualified completions from the persistent prefix cache", () => {
const prefixLookups = queryEditorSource.match(/lookupLocalCompletionColumnsByPrefix/g) ?? [];
expect(prefixLookups.length).toBeGreaterThanOrEqual(2);
expect(queryEditorSource).toContain("completionColumnRequestContext(reference)");
});
});
apps/desktop/src/lib/sql/sqlCompletion.ts
+90 -32
View File
@@ -1,4 +1,5 @@
import { Cassandra, MariaSQL, MSSQL, MySQL, PLSQL, PostgreSQL, SQLite, StandardSQL } from "@codemirror/lang-sql";
import type { Completion, CompletionInfo } from "@codemirror/autocomplete";
import type { DatabaseType, SqlSnippet } from "@/types/database";
import { buildMongoCompletionItemsFromContext, type MongoCompletionItem } from "@/lib/mongo/mongoCompletion";
import { CLOUDFLARE_D1_COMMON_FUNCTION_NAMES } from "@/lib/sql/cloudflareD1";
@@ -1279,7 +1280,7 @@ export interface SqlCompletionItem {
filterText?: string;
type: "keyword" | "table" | "column" | "snippet" | "function" | "schema" | "variable" | "text";
detail?: string;
info?: string;
info?: string | ((completion: Completion) => CompletionInfo | Promise<CompletionInfo>);
apply?: string;
replaceClosingQuote?: '"' | "'";
boost: number;
@@ -3756,6 +3757,8 @@ function buildSelectAllColumnItems(context: SqlCompletionContext, columnsByTable
function buildInsertAllColumnItems(context: SqlCompletionContext, columnsByTable: Map<string, SqlCompletionColumn[]>, t?: SqlCompletionTranslations, dialect?: SqlCompletionApplyDialect, keywordCase?: SqlKeywordCase): SqlCompletionItem[] {
if (!context.insertTable) return [];
// The INSERT column prefix controls the replacement/filter range, but the
// all-column snippet must always expand the complete target table.
const columns = uniqueColumnsByName(columnsForInsertTarget(context, columnsByTable));
if (columns.length === 0) return [];
@@ -3844,7 +3847,7 @@ function buildComparisonValueItems(context: SqlCompletionContext, columnsByTable
// Find the column's data type
let dataType: string | undefined;
for (const [, cols] of columnsByTable) {
for (const col of cols) {
for (const col of completionColumnPrefixCandidates(cols, unqualified, 256)) {
if (col.name.toLowerCase() === unqualified.toLowerCase()) {
if (qualifier) {
const qualLower = qualifier.toLowerCase();
@@ -4090,10 +4093,63 @@ function activeQueryBlockSql(sql: string): string {
return activeSelectIndex == null ? sql : sql.slice(activeSelectIndex);
}
function collectCompletionColumns(columnsByTable: Map<string, SqlCompletionColumn[]>): Array<SqlCompletionColumn & { key: string }> {
interface CompletionColumnSearchIndex {
entries: Array<{ normalizedName: string; index: number }>;
}
const completionColumnSearchIndexes = new WeakMap<readonly SqlCompletionColumn[], CompletionColumnSearchIndex>();
function completionColumnSearchIndex(columns: readonly SqlCompletionColumn[]): CompletionColumnSearchIndex {
const cached = completionColumnSearchIndexes.get(columns);
if (cached) return cached;
const index: CompletionColumnSearchIndex = {
entries: columns.map((column, index) => ({ normalizedName: column.name.trim().toLowerCase(), index })).sort((left, right) => left.normalizedName.localeCompare(right.normalizedName) || left.index - right.index),
};
completionColumnSearchIndexes.set(columns, index);
return index;
}
function completionColumnPrefixCandidates(columns: readonly SqlCompletionColumn[], prefix: string, limit = 256): SqlCompletionColumn[] {
if (!prefix) return [...columns];
const normalizedPrefix = prefix.trim().toLowerCase();
if (!normalizedPrefix) return [...columns];
const index = completionColumnSearchIndex(columns).entries;
let low = 0;
let high = index.length;
while (low < high) {
const middle = (low + high) >>> 1;
if ((index[middle]?.normalizedName ?? "") < normalizedPrefix) low = middle + 1;
else high = middle;
}
const candidates: SqlCompletionColumn[] = [];
const seen = new Set<number>();
for (let position = low; position < index.length && candidates.length < limit; position += 1) {
const entry = index[position];
if (!entry || !entry.normalizedName.startsWith(normalizedPrefix)) break;
const column = columns[entry.index];
if (column) {
candidates.push(column);
seen.add(entry.index);
}
}
// Preserve fuzzy/substring matching when the prefix index has not filled
// the bounded candidate pool, while avoiding a second scan for the common
// case where a prefix already has enough results.
if (candidates.length < limit) {
for (let indexPosition = 0; indexPosition < columns.length && candidates.length < limit; indexPosition += 1) {
if (seen.has(indexPosition)) continue;
const column = columns[indexPosition];
if (!column || !matchesIdentifierSearch(column.name, prefix)) continue;
candidates.push(column);
}
}
return candidates;
}
function collectCompletionColumns(columnsByTable: Map<string, SqlCompletionColumn[]>, prefix = ""): Array<SqlCompletionColumn & { key: string }> {
const allColumns: Array<SqlCompletionColumn & { key: string }> = [];
for (const [key, cols] of columnsByTable.entries()) {
for (const col of cols) {
for (const col of completionColumnPrefixCandidates(cols, prefix)) {
allColumns.push({ ...col, key });
}
}
@@ -4115,8 +4171,10 @@ function columnsForInsertTarget(context: SqlCompletionContext, columnsByTable: M
}
function buildColumnItems(context: SqlCompletionContext, columnsByTable: Map<string, SqlCompletionColumn[]>, dialect?: SqlCompletionApplyDialect): SqlCompletionItem[] {
// Collect all columns from the map (all tables have been fetched)
const allColumns = collectCompletionColumns(columnsByTable);
// Build a bounded candidate pool before doing duplicate detection, ranking,
// and completion object materialization. Canonical column arrays remain the
// source of truth; the per-array prefix index is only a derived accelerator.
const allColumns = collectCompletionColumns(columnsByTable, context.prefix);
// Handle INSERT column list: filter to only the target table
let relevantCols = allColumns;
@@ -4168,22 +4226,28 @@ function buildColumnItems(context: SqlCompletionContext, columnsByTable: Map<str
// keywords so they rank at the top instead of being interleaved.
const relevanceBoost = context.referencedTables.length > 0 || !!context.qualifier || !!context.insertTable ? 2000 : 0;
return uniqueColumns
.filter((column) => matchesIdentifierSearch(column.name, context.prefix) || matchesIdentifierSearch(column.displayLabel, context.prefix))
.map((column) => {
const rankedColumns = uniqueColumns
.map((column, index) => ({ column, index }))
.filter(({ column }) => matchesIdentifierSearch(column.name, context.prefix) || matchesIdentifierSearch(column.displayLabel, context.prefix))
.map(({ column, index }) => {
const keyBoost = isKeyColumn(column.name) ? 500 : 0;
const matchScore = Math.max(identifierMatchScore(column.name, context.prefix), identifierMatchScore(column.displayLabel, context.prefix));
return {
label: column.displayLabel,
filterText: column.displayLabel === column.name ? undefined : column.name,
type: "column" as const,
detail: buildColumnDetail(column),
info: buildColumnInfo(column),
apply: buildColumnApply(column, context, dialect),
boost: matchScore + keyBoost + relevanceBoost,
};
return { column, index, boost: matchScore + keyBoost + relevanceBoost };
})
.sort(compareCompletionItems);
.sort((left, right) => right.boost - left.boost || left.index - right.index)
.slice(0, context.insertTable || !context.prefix ? 50 : context.qualifier ? 30 : 20);
return rankedColumns.map(({ column, boost }) => {
return {
label: column.displayLabel,
filterText: column.displayLabel === column.name ? undefined : column.name,
type: "column" as const,
detail: buildColumnDetail(column),
info: buildColumnInfo(column),
apply: buildColumnApply(column, context, dialect),
boost,
};
});
}
function completionColumnsForReferencedTable<T extends SqlCompletionColumn & { key: string }>(table: SqlCompletionReferencedTable, columns: readonly T[]): T[] {
@@ -4203,7 +4267,7 @@ function applyReferencedColumnAliases<T extends SqlCompletionColumn>(table: SqlC
function hasMatchingReferencedColumnPrefix(context: SqlCompletionContext, columnsByTable: Map<string, SqlCompletionColumn[]>): boolean {
if (!context.suggestColumns || !context.prefix || context.referencedTables.length === 0) return false;
return context.referencedTables.some((table) => columnsForReferencedTable(table, columnsByTable).some((column) => matchesIdentifierSearch(column.name, context.prefix)));
return context.referencedTables.some((table) => completionColumnPrefixCandidates(columnsForReferencedTable(table, columnsByTable), context.prefix, 1).length > 0);
}
function qualifiedTableTargetFromContext(context: SqlCompletionContext): { database?: string; schema: string; table: string } | null {
@@ -4265,21 +4329,15 @@ function buildColumnDetail(column: SqlCompletionColumn): string {
if (column.isNullable === false) {
detail += " NOT NULL";
}
const comment = column.comment?.trim();
if (comment) {
detail += ` -- ${comment}`;
}
return detail;
}
function buildColumnInfo(column: SqlCompletionColumn): string | undefined {
const parts = [
column.schema ? `${column.schema}.${column.table}.${column.name}` : `${column.table}.${column.name}`,
column.dataType ? `Type: ${column.dataType}` : undefined,
column.isNullable === false ? "Nullable: no" : column.isNullable === true ? "Nullable: yes" : undefined,
column.comment?.trim() ? `Comment: ${column.comment.trim()}` : undefined,
].filter((part): part is string => !!part);
return parts.length > 1 ? parts.join("\n") : undefined;
const hasDetails = !!column.dataType || column.isNullable !== undefined || !!column.comment?.trim();
if (!hasDetails) return undefined;
const title = column.schema ? `${column.schema}.${column.table}.${column.name}` : `${column.table}.${column.name}`;
const details = [column.dataType ? `Type: ${column.dataType}` : undefined, column.isNullable === false ? "Nullable: no" : column.isNullable === true ? "Nullable: yes" : undefined, column.comment?.trim() ? `Comment: ${column.comment.trim()}` : undefined].filter((part): part is string => !!part);
return [title, ...details].join("\n");
}
function buildJoinConditionItems(context: SqlCompletionContext, columnsByTable: Map<string, SqlCompletionColumn[]>, foreignKeysByTable?: Map<string, SqlCompletionForeignKey[]>, dialect?: SqlCompletionApplyDialect, keywordCase?: SqlKeywordCase): SqlCompletionItem[] {
@@ -4786,7 +4844,7 @@ function buildNonAggregatedColumnItems(context: SqlCompletionContext, columnsByT
const items: SqlCompletionItem[] = [];
for (const [, cols] of columnsByTable) {
for (const col of cols) {
for (const col of completionColumnPrefixCandidates(cols, context.prefix, 256)) {
const key = col.name.toLowerCase();
if (!nonAggSet.has(key) || seen.has(key)) continue;
if (context.prefix && !matchesIdentifierSearch(col.name, context.prefix)) continue;
apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts
Vendored
+72
View File
@@ -582,6 +582,78 @@ describe("connectionStore completion assistant", () => {
]);
});
it("keeps PostgreSQL quoted column caches separate from unquoted identifiers", async () => {
const completionAssistantSearch = vi.fn(async (request: { parent_name?: string | null; parent_schema?: string | null }) => ({
candidates: [
{
name: request.parent_name === "Orders" ? "QuotedId" : "lower_id",
kind: "column",
parent_schema: request.parent_schema,
parent_name: request.parent_name,
data_type: "integer",
},
],
incomplete: false,
fallback_used: false,
}));
vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false }));
vi.doMock("@/lib/backend/api", () => ({
checkConnectionHealth: vi.fn().mockResolvedValue(undefined),
completionAssistantSearch,
getColumns: vi.fn(),
}));
const { useConnectionStore } = await import("@/stores/connectionStore");
const store = useConnectionStore();
store.connections = [postgresConnection()];
store.connectedIds.add("pg-1");
const unquoted = await store.listCompletionColumns("pg-1", "app", "orders", "public", { tableQuoted: false, schemaQuoted: false });
const quoted = await store.listCompletionColumns("pg-1", "app", "Orders", "Public", { tableQuoted: true, schemaQuoted: true });
expect(completionAssistantSearch).toHaveBeenCalledTimes(2);
expect(unquoted.map((column) => column.name)).toEqual(["lower_id"]);
expect(quoted.map((column) => column.name)).toEqual(["QuotedId"]);
expect(store.lookupLocalCompletionColumns("pg-1", "app", "orders", "public", undefined, { tableQuoted: false, schemaQuoted: false }).map((column) => column.name)).toEqual(["lower_id"]);
expect(store.lookupLocalCompletionColumns("pg-1", "app", "Orders", "Public", undefined, { tableQuoted: true, schemaQuoted: true }).map((column) => column.name)).toEqual(["QuotedId"]);
});
it("reads PostgreSQL prefix columns back with the same quoted context", async () => {
const completionAssistantSearch = vi.fn(async (request: { parent_name?: string | null; parent_schema?: string | null }) => ({
candidates: (request.parent_name === "Orders" ? ["QuotedName"] : ["name", "native"]).map((name) => ({
name,
kind: "column",
parent_schema: request.parent_schema,
parent_name: request.parent_name,
data_type: "text",
})),
incomplete: false,
fallback_used: false,
}));
vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false }));
vi.doMock("@/lib/backend/api", () => ({
checkConnectionHealth: vi.fn().mockResolvedValue(undefined),
completionAssistantSearch,
}));
const { useConnectionStore } = await import("@/stores/connectionStore");
const store = useConnectionStore();
store.connections = [postgresConnection()];
store.connectedIds.add("pg-1");
await store.listCompletionColumnsByPrefix("pg-1", "app", "orders", "public", "na", undefined, { tableQuoted: false, schemaQuoted: false });
const narrowed = await store.listCompletionColumnsByPrefix("pg-1", "app", "orders", "public", "nam", undefined, { tableQuoted: false, schemaQuoted: false });
await store.listCompletionColumnsByPrefix("pg-1", "app", "Orders", "Public", "Na", undefined, { tableQuoted: true, schemaQuoted: true });
expect(completionAssistantSearch).toHaveBeenCalledTimes(2);
expect(narrowed.map((column) => column.name)).toEqual(["name"]);
expect(store.lookupLocalCompletionColumnsByPrefix("pg-1", "app", "orders", "public", "na", undefined, { tableQuoted: false, schemaQuoted: false }).map((column) => column.name)).toEqual(["name", "native"]);
expect(store.lookupLocalCompletionColumnsByPrefix("pg-1", "app", "orders", "public", "nam", undefined, { tableQuoted: false, schemaQuoted: false }).map((column) => column.name)).toEqual(["name"]);
expect(store.lookupLocalCompletionColumnsByPrefix("pg-1", "app", "Orders", "Public", "Na", undefined, { tableQuoted: true, schemaQuoted: true }).map((column) => column.name)).toEqual(["QuotedName"]);
});
it("normalizes unquoted Oracle aliases while preserving quoted case before catalog lookup", async () => {
const completionAssistantSearch = vi.fn();
const getColumns = vi.fn().mockResolvedValue([]);
apps/desktop/src/stores/connectionStore.ts
+99 -7
View File
@@ -459,6 +459,7 @@ export const useConnectionStore = defineStore("connection", () => {
const completionTableIndex = new Map<string, { touched: number; tables: SqlCompletionTable[] }>();
const completionObjectIndex = new Map<string, { touched: number; objects: SqlCompletionObject[] }>();
const completionColumnIndex = new Map<string, { touched: number; columns: SqlCompletionColumn[] }>();
const completionColumnPrefixIndex = new Map<string, { touched: number; columns: SqlCompletionColumn[]; complete: boolean }>();
const completionForeignKeyIndex = new Map<string, { touched: number; foreignKeys: SqlCompletionForeignKey[] }>();
const completionInFlight = new Map<string, Promise<unknown>>();
const completionCacheRevisions = ref<Record<string, number>>({});
@@ -2903,6 +2904,9 @@ export const useConnectionStore = defineStore("connection", () => {
for (const key of completionColumnIndex.keys()) {
if (key.startsWith(cachePrefix)) completionColumnIndex.delete(key);
}
for (const key of completionColumnPrefixIndex.keys()) {
if (key.startsWith(cachePrefix)) completionColumnPrefixIndex.delete(key);
}
for (const key of completionForeignKeyIndex.keys()) {
if (key.startsWith(cachePrefix)) completionForeignKeyIndex.delete(key);
}
@@ -6360,12 +6364,21 @@ export const useConnectionStore = defineStore("connection", () => {
}
function completionColumnsKey(connectionId: string, database: string, table: string, schema?: string, catalog?: string, context?: { tableQuoted?: boolean; schemaQuoted?: boolean }): string {
if (getConfig(connectionId)?.db_type === "oracle") {
const databaseType = getConfig(connectionId)?.db_type;
if (databaseType === "oracle") {
const normalizedTable = context?.tableQuoted === false ? table.toUpperCase() : table;
const normalizedSchema = schema && context?.schemaQuoted === false ? schema.toUpperCase() : (schema ?? "");
return `${connectionId}:${database}:${catalog ?? ""}:${normalizedSchema}:${normalizedTable}`;
}
return `${completionTableScopeKey(connectionId, database, schema, catalog)}:${table.toLowerCase()}`;
const baseKey = `${completionTableScopeKey(connectionId, database, schema, catalog)}:${table.toLowerCase()}`;
if (databaseType !== "postgres" || (!context?.tableQuoted && !context?.schemaQuoted)) return baseKey;
const quotedSchema = context.schemaQuoted ? encodeURIComponent(schema ?? "") : "";
const quotedTable = context.tableQuoted ? encodeURIComponent(table) : "";
return `${baseKey}:quoted:s=${quotedSchema}:t=${quotedTable}`;
}
function completionColumnPrefixKey(connectionId: string, database: string, table: string, schema: string | undefined, prefix: string, catalog?: string, context?: { tableQuoted?: boolean; schemaQuoted?: boolean }): string {
return `${completionColumnsKey(connectionId, database, table, schema, catalog, context)}:prefix:${prefix.trim().toLowerCase()}`;
}
function completionForeignKeysKey(connectionId: string, database: string, table: string, schema?: string): string {
@@ -6410,6 +6423,13 @@ export const useConnectionStore = defineStore("connection", () => {
removed++;
}
}
for (const key of completionColumnPrefixIndex.keys()) {
const prefixMarker = key.toLowerCase().lastIndexOf(":prefix:");
const baseKey = prefixMarker >= 0 ? key.slice(0, prefixMarker) : key;
if (!matches(baseKey)) continue;
completionColumnPrefixIndex.delete(key);
removed++;
}
return removed;
}
@@ -6774,7 +6794,7 @@ export const useConnectionStore = defineStore("connection", () => {
requestRevision,
);
const columns = completionAssistantColumns(response.candidates, table, schema, context);
if (columns.length > 0 && requestRevision === completionCacheRevision(connectionId, database)) indexCompletionColumns(connectionId, database, table, schema, columns);
if (columns.length > 0 && requestRevision === completionCacheRevision(connectionId, database)) indexCompletionColumns(connectionId, database, table, schema, columns, undefined, context);
return columns;
}
@@ -6872,12 +6892,16 @@ export const useConnectionStore = defineStore("connection", () => {
}
}
function indexCompletionColumns(connectionId: string, database: string, table: string, schema: string | undefined, columns: SqlCompletionColumn[], catalog?: string) {
touchCompletionIndex(completionColumnIndex, completionColumnsKey(connectionId, database, table, schema, catalog), {
function indexCompletionColumns(connectionId: string, database: string, table: string, schema: string | undefined, columns: SqlCompletionColumn[], catalog?: string, context?: { tableQuoted?: boolean; schemaQuoted?: boolean }) {
touchCompletionIndex(completionColumnIndex, completionColumnsKey(connectionId, database, table, schema, catalog, context), {
columns,
});
}
function indexCompletionColumnPrefix(connectionId: string, database: string, table: string, schema: string | undefined, prefix: string, columns: SqlCompletionColumn[], complete: boolean, catalog?: string, context?: { tableQuoted?: boolean; schemaQuoted?: boolean }) {
touchCompletionIndex(completionColumnPrefixIndex, completionColumnPrefixKey(connectionId, database, table, schema, prefix, catalog, context), { columns, complete });
}
function sqlCompletionForeignKeys(foreignKeys: ForeignKeyInfo[]): SqlCompletionForeignKey[] {
return foreignKeys.map((foreignKey) => ({
name: foreignKey.name,
@@ -6957,10 +6981,75 @@ export const useConnectionStore = defineStore("connection", () => {
return completionColumnIndex.get(completionColumnsKey(connectionId, database, table, schema, catalog, context))?.columns ?? [];
}
function completionColumnPrefixEntry(connectionId: string, database: string, table: string, schema: string | undefined, prefix: string, catalog?: string, context?: { tableQuoted?: boolean; schemaQuoted?: boolean }): { columns: SqlCompletionColumn[]; complete: boolean } | undefined {
const normalizedPrefix = prefix.trim().toLowerCase();
const exact = completionColumnPrefixIndex.get(completionColumnPrefixKey(connectionId, database, table, schema, normalizedPrefix, catalog, context));
if (exact) return exact;
const marker = `${completionColumnsKey(connectionId, database, table, schema, catalog, context)}:prefix:`;
let reusable: { prefix: string; columns: SqlCompletionColumn[] } | undefined;
for (const [key, entry] of completionColumnPrefixIndex) {
if (!entry.complete || !key.startsWith(marker)) continue;
const cachedPrefix = key.slice(marker.length);
if (!normalizedPrefix.startsWith(cachedPrefix) || (reusable && cachedPrefix.length <= reusable.prefix.length)) continue;
reusable = { prefix: cachedPrefix, columns: entry.columns };
}
if (!reusable) return undefined;
return {
columns: reusable.columns.filter((column) => column.name.toLowerCase().startsWith(normalizedPrefix)),
complete: true,
};
}
function lookupLocalCompletionColumnsByPrefix(connectionId: string, database: string, table: string, schema: string | undefined, prefix: string, catalog?: string, context?: { tableQuoted?: boolean; schemaQuoted?: boolean }): SqlCompletionColumn[] {
return completionColumnPrefixEntry(connectionId, database, table, schema, prefix, catalog, context)?.columns ?? [];
}
function lookupLocalCompletionForeignKeys(connectionId: string, database: string, table: string, schema?: string): SqlCompletionForeignKey[] {
return completionForeignKeyIndex.get(completionForeignKeysKey(connectionId, database, table, schema))?.foreignKeys ?? [];
}
async function listCompletionColumnsByPrefix(connectionId: string, database: string, table: string, schema: string | undefined, prefix: string, catalog?: string, context?: { tableQuoted?: boolean; schemaQuoted?: boolean }): Promise<SqlCompletionColumn[]> {
const normalizedPrefix = prefix.trim();
if (normalizedPrefix.length < 2) return [];
const databaseType = effectiveDatabaseTypeForConnection(getConfig(connectionId));
if (databaseType !== "postgres" && databaseType !== "mysql") return [];
const completionTable = table;
const completionSchema = schema?.trim() || (databaseType === "mysql" ? database : undefined);
const cacheKey = completionColumnPrefixKey(connectionId, database, completionTable, completionSchema, normalizedPrefix, catalog, context);
const cached = completionColumnPrefixEntry(connectionId, database, completionTable, completionSchema, normalizedPrefix, catalog, context);
if (cached) return cached.columns;
const requestRevision = completionCacheRevision(connectionId, database);
return withCompletionInFlight(
`column-prefix:${requestRevision}:${cacheKey}`,
async () => {
const existing = completionColumnPrefixIndex.get(cacheKey);
if (existing) return existing.columns;
await ensureConnected(connectionId);
const response = await completionAssistantSearch(
{
connection_id: connectionId,
database,
schema: completionSchema ?? null,
object_kinds: ["column"],
mask: normalizedPrefix,
max_results: 128,
parent_schema: completionSchema ?? null,
parent_name: completionTable,
match_mode: "prefix",
},
requestRevision,
);
const columns = completionAssistantColumns(response.candidates, completionTable, completionSchema, context);
if (requestRevision === completionCacheRevision(connectionId, database)) {
indexCompletionColumnPrefix(connectionId, database, completionTable, completionSchema, normalizedPrefix, columns, !response.incomplete, catalog, context);
}
return columns;
},
{ scope: completionLimiterScope(connectionId, database), kind: "column-prefix" },
);
}
function databaseNamesFromTree(connectionId: string): string[] {
const node = findConnectionNode(connectionId);
if (!node?.children) return [];
@@ -7402,13 +7491,14 @@ export const useConnectionStore = defineStore("connection", () => {
return [];
}
const sessionCacheScope = usesOracleCurrentSchema && context?.clientSessionId ? `:${context.clientSessionId}:${context.version ?? 0}` : "";
const cacheKey = `${connectionId}:${database}:${catalog ?? ""}:${completionSchema || ""}:${completionTable}${sessionCacheScope}`;
const cacheKey = `${completionColumnsKey(connectionId, database, completionTable, completionSchema, catalog, context)}${sessionCacheScope}`;
if (!completionColumnsCache.value[cacheKey]) {
const requestRevision = completionCacheRevision(connectionId, database);
await withCompletionInFlight(
`${cacheKey}:columns`,
async () => {
await ensureConnected(connectionId);
// Use assistant metadata opportunistically, then fall back to canonical metadata.
if (!usesOracleCurrentSchema && !catalog) {
try {
const assistantColumns = await listCompletionAssistantColumns(connectionId, database, completionTable, completionSchema, context, requestRevision);
@@ -7456,7 +7546,7 @@ export const useConnectionStore = defineStore("connection", () => {
isNullable: column.is_nullable,
comment: column.comment,
}));
if (!usesOracleCurrentSchema) indexCompletionColumns(connectionId, database, completionTable, completionSchema, columns, catalog);
if (!usesOracleCurrentSchema) indexCompletionColumns(connectionId, database, completionTable, completionSchema, columns, catalog, context);
return columns;
}
@@ -8396,6 +8486,7 @@ export const useConnectionStore = defineStore("connection", () => {
listCompletionTables,
listCompletionObjects,
listCompletionColumns,
listCompletionColumnsByPrefix,
listCompletionForeignKeys,
listCompletionSchemas,
listCompletionDatabases,
@@ -8403,6 +8494,7 @@ export const useConnectionStore = defineStore("connection", () => {
lookupLocalCompletionTables,
lookupLocalCompletionObjects,
lookupLocalCompletionColumns,
lookupLocalCompletionColumnsByPrefix,
lookupLocalCompletionForeignKeys,
lookupLocalCompletionSchemas,
lookupLocalCompletionDatabases,
packages/app-tests/sqlCompletion.test.ts
Vendored
+1 -1
View File
@@ -1224,7 +1224,7 @@ test("shows column comments in WHERE field completions", () => {
});
const column = items.find((item) => item.type === "column" && item.label === "status");
assert.equal(column?.detail, "public.orders [varchar] NOT NULL -- Order lifecycle state");
assert.equal(column?.detail, "public.orders [varchar] NOT NULL");
assert.equal(column?.info, "public.orders.status\nType: varchar\nNullable: no\nComment: Order lifecycle state");
});