diff --git a/apps/desktop/src/lib/sqlCompletion.ts b/apps/desktop/src/lib/sqlCompletion.ts index b2873b9d8..ae5e30b94 100644 --- a/apps/desktop/src/lib/sqlCompletion.ts +++ b/apps/desktop/src/lib/sqlCompletion.ts @@ -1088,6 +1088,7 @@ export interface SqlCompletionContext { nonAggregatedSelectColumns: string[]; comparisonLeftColumn?: string; onStar: boolean; + selectListColumnContext: boolean; preferredKeywords: string[]; updateTarget?: { table: string; schema?: string }; deleteTarget?: { table: string; schema?: string }; @@ -1213,6 +1214,7 @@ class SqlCompletionProvider { if (!context.exclusiveTableSuggestions && context.suggestColumns) { this.items.push(...buildColumnItems(context, this.input.columnsByTable, this.dialect)); + this.items.push(...buildSelectAllColumnItems(context, this.input.columnsByTable, this.t, this.dialect)); } if (context.referencedTables.length > 0 && !context.suggestColumns && !context.insertTable) { @@ -1550,7 +1552,8 @@ export function getSqlCompletionContext(sql: string, cursor: number): SqlComplet const exclusiveColumnSuggestions = !!qualifier && !exclusiveTableSuggestions && !insertInfo; // Check if we're in a context where columns are expected - const inColumnContext = isInColumnContext(beforeCursor) || !!insertInfo; + const selectListColumnContext = isInSelectListContext(beforeCursor); + const inColumnContext = selectListColumnContext || isInColumnContext(beforeCursor) || !!insertInfo; const inJoinConditionContext = isInJoinConditionContext(beforeCursor); const prioritizeSelectAliases = isInOrderOrGroupByContext(beforeCursor); const inCallRoutineContext = isCallRoutineContext(beforeCursor); @@ -1598,6 +1601,7 @@ export function getSqlCompletionContext(sql: string, cursor: number): SqlComplet nonAggregatedSelectColumns: extractNonAggregatedSelectColumns(fullStatement), comparisonLeftColumn: detectComparisonLeftColumn(beforeCursor), onStar: detectOnStar(beforeCursor), + selectListColumnContext, preferredKeywords, updateTarget: updateInfo?.target, deleteTarget: deleteInfo?.target, @@ -2569,6 +2573,92 @@ function buildStarExpansionItem(columnsByTable: Map, t?: SqlCompletionTranslations, dialect?: "mysql" | "postgres" | "sqlserver"): SqlCompletionItem[] { + if (!context.selectListColumnContext || context.statementKind !== "select" || context.onStar || context.referencedTables.length === 0) { + return []; + } + + const items: SqlCompletionItem[] = []; + const emittedRefs = new Set(); + const targetRefs = referencedTablesForSelectAllColumns(context); + const shouldQualify = !!context.qualifier || context.referencedTables.length > 1; + + for (const ref of targetRefs) { + const displayRef = context.qualifier || ref.alias || ref.name; + const refKey = `${displayRef}.${ref.schema ?? ""}.${ref.name}`.toLowerCase(); + if (emittedRefs.has(refKey)) continue; + emittedRefs.add(refKey); + + const columns = uniqueColumnsByName(columnsForSelectAllReferencedTable(ref, columnsByTable)); + if (columns.length === 0) continue; + + const label = `${displayRef}.*`; + if (!selectAllColumnItemMatchesPrefix(label, ref, columns, context.prefix)) continue; + + const qualifier = context.qualifier || ref.alias || (shouldQualify ? quoteSqlIdentifier(ref.name, dialect) : undefined); + const expansion = buildSelectAllColumnExpansion(columns, qualifier, !!context.qualifier, dialect); + const countText = (t?.starExpansionColumns ?? "{count} columns").replace("{count}", String(columns.length)); + items.push({ + label, + type: "snippet" as const, + detail: `${countText}: ${expansion.length > 60 ? expansion.slice(0, 57) + "..." : expansion}`, + apply: expansion, + boost: 2400 + selectAllColumnItemPrefixBoost(label, ref, columns, context.prefix) - items.length, + }); + } + + return items; +} + +function referencedTablesForSelectAllColumns(context: SqlCompletionContext): SqlCompletionReferencedTable[] { + if (!context.qualifier) return context.referencedTables; + const qualifier = context.qualifier; + const qualifierLower = qualifier.toLowerCase(); + const qualifiedTarget = qualifiedTableTargetFromContext(context); + return context.referencedTables.filter((table) => referencedTableMatchesColumnQualifier(table, qualifier, qualifierLower, qualifiedTarget)); +} + +function buildSelectAllColumnExpansion(columns: SqlCompletionColumn[], qualifier: string | undefined, qualifierAlreadyTyped: boolean, dialect?: "mysql" | "postgres" | "sqlserver"): string { + return columns + .map((column, index) => { + const columnName = quoteSqlIdentifier(column.name, dialect); + if (!qualifier || (qualifierAlreadyTyped && index === 0)) return columnName; + return `${qualifier}.${columnName}`; + }) + .join(", "); +} + +function columnsForSelectAllReferencedTable(table: SqlCompletionReferencedTable, columnsByTable: Map): SqlCompletionColumn[] { + const columns = columnsForReferencedTable(table, columnsByTable); + if (columns.length > 0) return columns; + if (!table.columns || table.columns.length === 0) return []; + return table.columns.map((name) => ({ name, table: table.name, schema: table.schema })); +} + +function uniqueColumnsByName(columns: SqlCompletionColumn[]): SqlCompletionColumn[] { + const seen = new Set(); + const unique: SqlCompletionColumn[] = []; + for (const column of columns) { + const key = normalizeIdentifierPart(column.name); + if (seen.has(key)) continue; + seen.add(key); + unique.push(column); + } + return unique; +} + +function selectAllColumnItemMatchesPrefix(label: string, ref: SqlCompletionReferencedTable, columns: SqlCompletionColumn[], prefix: string): boolean { + if (!prefix) return true; + if (matchesPrefix(label, prefix) || matchesPrefix(ref.name, prefix) || (!!ref.alias && matchesPrefix(ref.alias, prefix))) return true; + return columns.some((column) => matchesPrefix(column.name, prefix)); +} + +function selectAllColumnItemPrefixBoost(label: string, ref: SqlCompletionReferencedTable, columns: SqlCompletionColumn[], prefix: string): number { + if (!prefix) return 0; + const scores = [computeBoost(label, prefix), computeBoost(ref.name, prefix), ref.alias ? computeBoost(ref.alias, prefix) : -1, ...columns.map((column) => computeBoost(column.name, prefix))]; + return Math.min(Math.max(...scores, 0), 1000); +} + function buildComparisonValueItems(context: SqlCompletionContext, columnsByTable: Map, t?: SqlCompletionTranslations, keywordCase?: SqlKeywordCase): SqlCompletionItem[] { const colName = context.comparisonLeftColumn!; const parts = colName.split("."); diff --git a/packages/app-tests/sqlCompletion.test.ts b/packages/app-tests/sqlCompletion.test.ts index 2d3b8b16a..de7dddc9c 100644 --- a/packages/app-tests/sqlCompletion.test.ts +++ b/packages/app-tests/sqlCompletion.test.ts @@ -310,9 +310,10 @@ test("suggests only matching columns for an explicit alias qualifier prefix", () tables, columnsByTable, }); + const columnItems = items.filter((item) => item.type === "column"); assert.deepEqual( - items.map((item) => [item.label, item.type, item.detail]), + columnItems.map((item) => [item.label, item.type, item.detail]), [["name", "column", "public.users [varchar]"]], ); }); @@ -450,6 +451,87 @@ test("suggests columns from referenced tables in select list", () => { assert.equal(items[0]?.type, "column"); }); +test("suggests all columns expansion in select list when typing a column prefix", () => { + const sql = "select id from public.users"; + const cursor = "select id".length; + const items = buildSqlCompletionItems(sql, cursor, { + tables, + columnsByTable, + }); + + const allColumns = items.find((item) => item.type === "snippet" && item.label === "users.*"); + assert.ok(allColumns); + assert.equal(allColumns.apply, "id, name, email"); +}); + +test("qualifies all columns expansion with table aliases", () => { + const sql = "select id from public.users u"; + const cursor = "select id".length; + const items = buildSqlCompletionItems(sql, cursor, { + tables, + columnsByTable, + }); + + const allColumns = items.find((item) => item.type === "snippet" && item.label === "u.*"); + assert.ok(allColumns); + assert.equal(allColumns.apply, "u.id, u.name, u.email"); +}); + +test("suggests all columns expansion for each joined table", () => { + const sql = "select id from public.users u join public.orders o on u.id = o.user_id"; + const cursor = "select id".length; + const items = buildSqlCompletionItems(sql, cursor, { + tables, + columnsByTable, + }); + + assert.equal(items.find((item) => item.type === "snippet" && item.label === "u.*")?.apply, "u.id, u.name, u.email"); + assert.equal(items.find((item) => item.type === "snippet" && item.label === "o.*")?.apply, "o.id, o.user_id, o.status"); +}); + +test("suggests all columns expansion after an alias qualifier in select list", () => { + const sql = "select u. from public.users u join public.orders o on u.id = o.user_id"; + const cursor = "select u.".length; + const items = buildSqlCompletionItems(sql, cursor, { + tables, + columnsByTable, + }); + + const allColumns = items.find((item) => item.type === "snippet" && item.label === "u.*"); + assert.ok(allColumns); + assert.equal(allColumns.apply, "id, u.name, u.email"); + assert.equal( + items.some((item) => item.type === "snippet" && item.label === "o.*"), + false, + ); +}); + +test("keeps all columns expansion available after an alias-qualified column prefix", () => { + const sql = "select u.i from public.users u join public.orders o on u.id = o.user_id"; + const cursor = "select u.i".length; + const items = buildSqlCompletionItems(sql, cursor, { + tables, + columnsByTable, + }); + + const allColumns = items.find((item) => item.type === "snippet" && item.label === "u.*"); + assert.ok(allColumns); + assert.equal(allColumns.apply, "id, u.name, u.email"); +}); + +test("does not suggest all columns expansion outside select list", () => { + const sql = "select * from public.users u where id"; + const items = buildSqlCompletionItems(sql, sql.length, { + tables, + columnsByTable, + }); + + assert.equal( + items.some((item) => item.type === "snippet" && item.label === "u.*"), + false, + ); +}); + test("suggests tables after LEFT JOIN", () => { const sql = "select * from users left join us"; const items = buildSqlCompletionItems(sql, sql.length, {