diff --git a/apps/desktop/src/lib/__tests__/sql/semantic/completion.spec.ts b/apps/desktop/src/lib/__tests__/sql/semantic/completion.spec.ts index 678eec421..0cec221b2 100644 --- a/apps/desktop/src/lib/__tests__/sql/semantic/completion.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/semantic/completion.spec.ts @@ -46,6 +46,86 @@ describe("semantic SQL completion candidates", () => { expect(items.filter((item) => item.type === "column").map((item) => item.label)).toEqual(["id", "name", "email"]); }); + it.each([ + ["PostgreSQL", "postgres", "postgres"], + ["SQL Server", "sqlserver", "sqlserver"], + ] as const)("uses row-source aliases for %s self-join column collisions", (_label, databaseType, dialect) => { + const columnsByTable = new Map([["users", ["id", "name"].map((name) => ({ name, table: "users" }))]]); + + const { items } = semanticCompletion("SELECT * FROM users u JOIN users v ON u.id = v.id WHERE |", { columnsByTable }, { databaseType, dialect }); + const columns = items.filter((item) => item.type === "column"); + + expect(columns.map((item) => item.label)).toEqual(expect.arrayContaining(["u.id", "u.name", "v.id", "v.name"])); + expect(columns.find((item) => item.label === "u.id")?.apply).toBe("u.id"); + expect(columns.find((item) => item.label === "v.id")?.apply).toBe("v.id"); + }); + + it("completes columns for aliases in comma-separated table lists", () => { + const columnsByTable = new Map([ + ["table_a", ["id", "name"].map((name) => ({ name, table: "table_a" }))], + ["table_b", ["id", "status"].map((name) => ({ name, table: "table_b" }))], + ]); + + const { context, items } = semanticCompletion("SELECT * FROM table_a a, table_b b WHERE a.id = b.|", { columnsByTable }); + + expect(context.referencedTables).toEqual(expect.arrayContaining([expect.objectContaining({ name: "table_b", alias: "b" })])); + expect(items.filter((item) => item.type === "column").map((item) => item.label)).toEqual(["id", "status"]); + }); + + it("completes correlation columns for generic PostgreSQL table functions", () => { + const { context, items } = semanticCompletion("SELECT * FROM generate_series(1, 3) g(value) WHERE g.|", {}, { databaseType: "postgres", dialect: "postgres" }); + + expect(context.referencedTables).toEqual(expect.arrayContaining([expect.objectContaining({ name: "g", alias: "g" })])); + expect(items.filter((item) => item.type === "column").map((item) => item.label)).toEqual(["value"]); + }); + + it("completes correlation columns after PostgreSQL WITH ORDINALITY", () => { + const { items } = semanticCompletion("SELECT * FROM generate_series(1, 3) WITH ORDINALITY AS g(value, ord), orders o WHERE g.|", {}, { databaseType: "postgres", dialect: "postgres" }); + + expect(items.filter((item) => item.type === "column").map((item) => item.label)).toEqual(["value", "ord"]); + }); + + it("completes later comma-separated sources after a joined table", () => { + const columnsByTable = new Map([["audit_log", ["event_id", "action"].map((name) => ({ name, table: "audit_log" }))]]); + const { context, items } = semanticCompletion("SELECT * FROM users u JOIN orders o ON o.user_id = u.id, audit_log a WHERE a.|", { columnsByTable }, { databaseType: "postgres", dialect: "postgres" }); + + expect(context.referencedTables).toEqual(expect.arrayContaining([expect.objectContaining({ name: "audit_log", alias: "a" })])); + expect(items.filter((item) => item.type === "column").map((item) => item.label)).toEqual(["event_id", "action"]); + }); + + it("completes correlation columns for aliased table sources", () => { + const columnsByTable = new Map([["table_a", ["source_id", "source_label"].map((name) => ({ name, table: "table_a" }))]]); + const { items } = semanticCompletion("SELECT * FROM table_a a(id, label), table_b b WHERE a.|", { columnsByTable }, { databaseType: "postgres", dialect: "postgres" }); + + expect(items.filter((item) => item.type === "column").map((item) => item.label)).toEqual(["id", "label"]); + }); + + it("loads real SQL Server columns after aliased table hints", () => { + const columnsByTable = new Map([["users", ["id", "name", "email"].map((name) => ({ name, table: "users" }))]]); + + const { items } = semanticCompletion("SELECT * FROM users u (NOLOCK) WHERE u.|", { columnsByTable }, { databaseType: "sqlserver", dialect: "sqlserver" }); + + expect(items.filter((item) => item.type === "column").map((item) => item.label)).toEqual(["id", "name", "email"]); + }); + + it("merges partial PostgreSQL correlation names with metadata positionally", () => { + const columnsByTable = new Map([["users", ["id", "name", "email"].map((name) => ({ name, table: "users" }))]]); + + const { context, items } = semanticCompletion("SELECT * FROM users u(user_id) WHERE u.|", { columnsByTable }, { databaseType: "postgres", dialect: "postgres" }); + + expect(context.referencedTables).toEqual(expect.arrayContaining([expect.objectContaining({ name: "users", alias: "u", columns: undefined, columnAliases: ["user_id"] })])); + expect(items.filter((item) => item.type === "column").map((item) => item.label)).toEqual(["user_id", "name", "email"]); + }); + + it("completes an unquoted SQL Server table named lateral", () => { + const columnsByTable = new Map([["lateral", ["id", "value"].map((name) => ({ name, table: "lateral" }))]]); + + const { context, items } = semanticCompletion("SELECT * FROM lateral l WHERE l.|", { columnsByTable }, { databaseType: "sqlserver", dialect: "sqlserver" }); + + expect(context.referencedTables).toEqual(expect.arrayContaining([expect.objectContaining({ name: "lateral", alias: "l" })])); + expect(items.filter((item) => item.type === "column").map((item) => item.label)).toEqual(["id", "value"]); + }); + it("uses CTE projected columns without remote metadata", () => { const { items, context } = semanticCompletion("WITH recent_orders(id, total) AS (SELECT id, total FROM orders) SELECT * FROM recent_orders ro WHERE ro.|"); diff --git a/apps/desktop/src/lib/__tests__/sql/semantic/model.spec.ts b/apps/desktop/src/lib/__tests__/sql/semantic/model.spec.ts index a8f73c159..bbe381b21 100644 --- a/apps/desktop/src/lib/__tests__/sql/semantic/model.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/semantic/model.spec.ts @@ -72,6 +72,108 @@ describe("sqlSemanticModel baseline fixtures", () => { expect(scope.kind).toBe("table"); }); + it("extracts aliases from completed comma-separated table lists", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM table_a a, table_b b WHERE a.id = b.|"); + const model = buildSqlSemanticModel(sql, cursor); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ name: "table_a", alias: "a" }), expect.objectContaining({ name: "table_b", alias: "b" })])); + expect(model.cursorIntent).toEqual(expect.objectContaining({ kind: "alias_column", qualifierParts: ["b"] })); + }); + + it("consumes correlation column lists before parsing later comma-separated sources", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM table_a a(id), table_b b, table_c c WHERE c.|"); + const model = buildSqlSemanticModel(sql, cursor, { databaseType: "postgres" }); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ name: "table_a", alias: "a", columns: undefined, columnAliases: ["id"] }), expect.objectContaining({ name: "table_b", alias: "b" }), expect.objectContaining({ name: "table_c", alias: "c" })])); + expect(model.cursorIntent).toEqual(expect.objectContaining({ kind: "alias_column", qualifierParts: ["c"] })); + }); + + it("parses generic PostgreSQL table functions and their correlation columns", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM generate_series(1, 3) g(value) WHERE g.|"); + const model = buildSqlSemanticModel(sql, cursor, { databaseType: "postgres" }); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ kind: "table_function", name: "g", alias: "g", columns: ["value"] })])); + expect(model.cursorIntent).toEqual(expect.objectContaining({ kind: "alias_column", qualifierParts: ["g"] })); + }); + + it("consumes LATERAL functions before parsing later comma-separated sources", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM users u, LATERAL generate_series(1, 3) g(value), orders o WHERE o.|"); + const model = buildSqlSemanticModel(sql, cursor, { databaseType: "postgres" }); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ name: "users", alias: "u" }), expect.objectContaining({ kind: "table_function", name: "g", alias: "g", columns: ["value"] }), expect.objectContaining({ name: "orders", alias: "o" })])); + expect(model.cursorIntent).toEqual(expect.objectContaining({ kind: "alias_column", qualifierParts: ["o"] })); + }); + + it("consumes WITH ORDINALITY before function aliases and later sources", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM generate_series(1, 3) WITH ORDINALITY AS g(value, ord), orders o WHERE o.|"); + const model = buildSqlSemanticModel(sql, cursor, { databaseType: "postgres" }); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ kind: "table_function", name: "g", alias: "g", columns: ["value", "ord"] }), expect.objectContaining({ name: "orders", alias: "o" })])); + expect(model.cursorIntent).toEqual(expect.objectContaining({ kind: "alias_column", qualifierParts: ["o"] })); + }); + + it("parses comma-separated sources after a complete joined table", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM users u JOIN orders o ON o.user_id = u.id, audit_log a WHERE a.|"); + const model = buildSqlSemanticModel(sql, cursor, { databaseType: "postgres" }); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ name: "users", alias: "u" }), expect.objectContaining({ name: "orders", alias: "o" }), expect.objectContaining({ name: "audit_log", alias: "a" })])); + expect(model.cursorIntent).toEqual(expect.objectContaining({ kind: "alias_column", qualifierParts: ["a"] })); + }); + + it("does not parse commas in later SELECT clauses as row sources", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM users u WINDOW w1 AS (PARTITION BY u.id), w2 AS (PARTITION BY u.id)|"); + const model = buildSqlSemanticModel(sql, cursor, { databaseType: "postgres" }); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ name: "users", alias: "u" })])); + expect(model.rowSources.some((source) => source.name === "w2")).toBe(false); + }); + + it("keeps LATERAL subqueries available as row sources", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM users u, LATERAL (SELECT u.id AS user_id) s WHERE s.|"); + const model = buildSqlSemanticModel(sql, cursor, { databaseType: "postgres" }); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ kind: "subquery", name: "s", alias: "s", columns: ["user_id"] })])); + expect(model.cursorIntent).toEqual(expect.objectContaining({ kind: "alias_column", qualifierParts: ["s"] })); + }); + + it("does not classify SQL Server table hints as generic table functions", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM users (NOLOCK) WHERE users.|"); + const model = buildSqlSemanticModel(sql, cursor, { databaseType: "sqlserver" }); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ kind: "table", name: "users" })])); + expect(model.rowSources.some((source) => source.kind === "table_function")).toBe(false); + }); + + it("keeps aliased SQL Server table hints separate from correlation columns", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM users u (NOLOCK), orders o WHERE u.|"); + const model = buildSqlSemanticModel(sql, cursor, { databaseType: "sqlserver" }); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ kind: "table", name: "users", alias: "u", columns: undefined, columnAliases: undefined }), expect.objectContaining({ kind: "table", name: "orders", alias: "o" })])); + expect(model.cursorIntent).toEqual(expect.objectContaining({ kind: "alias_column", qualifierParts: ["u"] })); + }); + + it("consumes SQL Server WITH table hints without treating WITH as an alias", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM users WITH (NOLOCK), orders o WHERE users.|"); + const model = buildSqlSemanticModel(sql, cursor, { databaseType: "sqlserver" }); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ kind: "table", name: "users", alias: undefined, columns: undefined }), expect.objectContaining({ kind: "table", name: "orders", alias: "o" })])); + }); + + it("keeps partial PostgreSQL correlation names separate from the source schema", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM users u(user_id) WHERE u.|"); + const model = buildSqlSemanticModel(sql, cursor, { databaseType: "postgres" }); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ kind: "table", name: "users", alias: "u", columns: undefined, columnAliases: ["user_id"], metadataTarget: { table: "users" } })])); + }); + + it("treats LATERAL as a regular SQL Server table name", () => { + const { sql, cursor } = sqlFixtureCursor("SELECT * FROM lateral l WHERE l.|"); + const model = buildSqlSemanticModel(sql, cursor, { databaseType: "sqlserver" }); + + expect(model.rowSources).toEqual(expect.arrayContaining([expect.objectContaining({ kind: "table", name: "lateral", alias: "l" })])); + expect(model.cursorIntent).toEqual(expect.objectContaining({ kind: "alias_column", qualifierParts: ["l"] })); + }); + it("classifies alias-qualified star with replacement range", () => { const { sql, cursor } = sqlFixtureCursor("SELECT u.*| FROM users u"); const model = buildSqlSemanticModel(sql, cursor); diff --git a/apps/desktop/src/lib/sql/semantic/completion.ts b/apps/desktop/src/lib/sql/semantic/completion.ts index 793a6d844..2d633bbb4 100644 --- a/apps/desktop/src/lib/sql/semantic/completion.ts +++ b/apps/desktop/src/lib/sql/semantic/completion.ts @@ -21,6 +21,7 @@ export function sqlSemanticReferencedTables(model: SqlSemanticModel): SqlComplet schema: source.qualifierParts[source.qualifierParts.length - 1], alias: source.alias, columns: source.columns, + columnAliases: source.columnAliases, })); } diff --git a/apps/desktop/src/lib/sql/semantic/model.ts b/apps/desktop/src/lib/sql/semantic/model.ts index 7b1e13d3f..14ea65b32 100644 --- a/apps/desktop/src/lib/sql/semantic/model.ts +++ b/apps/desktop/src/lib/sql/semantic/model.ts @@ -17,9 +17,11 @@ import type { } from "@/lib/sql/semantic/types"; const TABLE_INTRODUCERS = new Set(["from", "join", "straight_join", "update", "into", "using", "apply"]); +const TABLE_FUNCTION_NAMES = new Set(["table", "xmltable", "json_table", "the", "read_csv", "read_parquet", "read_json", "unnest"]); const JOIN_MODIFIERS = new Set(["left", "right", "inner", "outer", "cross", "full", "natural"]); const CLAUSE_BOUNDARIES = new Set(["where", "group", "having", "order", "limit", "offset", "union", "intersect", "except", "on", "set", "values", "returning"]); -const ALIAS_BLACKLIST = new Set([...CLAUSE_BOUNDARIES, "join", "straight_join", "left", "right", "inner", "outer", "cross", "full", "natural", "as", "select", "from"]); +const FROM_CLAUSE_BOUNDARIES = new Set([...CLAUSE_BOUNDARIES, "window", "qualify", "fetch", "for", "connect", "start", "model"].filter((item) => item !== "on")); +const ALIAS_BLACKLIST = new Set([...CLAUSE_BOUNDARIES, "join", "straight_join", "left", "right", "inner", "outer", "cross", "full", "natural", "as", "select", "from", "with"]); interface ParseState { dialect: SqlSemanticDialectAdapter; @@ -218,26 +220,55 @@ function parseCteSources(state: ParseState): SqlSemanticRowSource[] { return sources; } -function aliasAfter(tokens: readonly SqlSemanticToken[], index: number, dialect: SqlSemanticDialectAdapter): { alias?: string; aliasSpan?: SqlSemanticSpan; nextIndex: number } { +function correlationColumnsAfter(tokens: readonly SqlSemanticToken[], index: number, dialect: SqlSemanticDialectAdapter): { columns: string[]; nextIndex: number } | null { + if (tokens[index]?.text !== "(") return null; + const close = findMatchingParenToken(tokens, index); + if (close < 0) return null; + const columns = splitTopLevelByComma(tokens.slice(index + 1, close)) + .map((group) => group.find(tokenIsIdentifier)) + .filter((item): item is SqlSemanticToken => item != null) + .map((item) => identifierPart(item, dialect).name); + return { columns, nextIndex: close + 1 }; +} + +function aliasAfter(tokens: readonly SqlSemanticToken[], index: number, dialect: SqlSemanticDialectAdapter, options: { allowCorrelationColumns?: boolean } = {}): { alias?: string; aliasSpan?: SqlSemanticSpan; columns?: string[]; nextIndex: number } { let cursor = index; if (tokens[cursor]?.kind === "word" && tokens[cursor]?.normalized === "as") cursor += 1; const aliasToken = tokens[cursor]; if (tokenIsIdentifier(aliasToken)) { const alias = identifierPart(aliasToken, dialect).name; if (!ALIAS_BLACKLIST.has(alias.toLowerCase())) { - return { alias, aliasSpan: aliasToken.span, nextIndex: cursor + 1 }; + const columns = options.allowCorrelationColumns === false ? null : correlationColumnsAfter(tokens, cursor + 1, dialect); + return { alias, aliasSpan: aliasToken.span, columns: columns?.columns, nextIndex: columns?.nextIndex ?? cursor + 1 }; } } return { nextIndex: index }; } +function mergeColumnAliases(columns: readonly string[], aliases: readonly string[] | undefined): string[] { + if (!aliases?.length) return [...columns]; + if (columns.length === 0) return [...aliases]; + return columns.map((column, index) => aliases[index] ?? column); +} + +function consumeSqlServerTableHint(tokens: readonly SqlSemanticToken[], index: number, dialect: SqlSemanticDialectAdapter): number { + if (dialect.id !== "sqlserver") return index; + const openIndex = tokens[index]?.normalized === "with" && tokens[index + 1]?.text === "(" ? index + 1 : index; + if (tokens[openIndex]?.text !== "(") return index; + const close = findMatchingParenToken(tokens, openIndex); + return close < 0 ? index : close + 1; +} + function parseSubquerySource(state: ParseState, openIndex: number, introducer: string, sourceIndex: number): { source: SqlSemanticRowSource; nextIndex: number } | null { const close = findMatchingParenToken(state.tokens, openIndex); if (close < 0) return null; const alias = aliasAfter(state.tokens, close + 1, state.dialect); if (!alias.alias) return null; const bodyTokens = state.tokens.slice(openIndex + 1, close); - const columns = parseSelectProjections(bodyTokens, state.dialect).map((projection) => projection.name); + const columns = mergeColumnAliases( + parseSelectProjections(bodyTokens, state.dialect).map((projection) => projection.name), + alias.columns, + ); return { source: { id: `${introducer}:subquery:${sourceIndex}`, @@ -246,7 +277,7 @@ function parseSubquerySource(state: ParseState, openIndex: number, introducer: s qualifierParts: [], alias: alias.alias, aliasSpan: alias.aliasSpan, - sourceSpan: { start: state.tokens[openIndex]?.span.start ?? 0, end: alias.aliasSpan?.end ?? state.tokens[close]?.span.end ?? 0 }, + sourceSpan: { start: state.tokens[openIndex]?.span.start ?? 0, end: state.tokens[alias.nextIndex - 1]?.span.end ?? alias.aliasSpan?.end ?? state.tokens[close]?.span.end ?? 0 }, columns, }, nextIndex: alias.nextIndex, @@ -254,22 +285,29 @@ function parseSubquerySource(state: ParseState, openIndex: number, introducer: s } function parseTableFunctionSource(state: ParseState, nameIndex: number, introducer: string, sourceIndex: number): { source: SqlSemanticRowSource; nextIndex: number } | null { - const nameToken = state.tokens[nameIndex]; - if (!nameToken || nameToken.kind !== "word" || !["table", "xmltable", "json_table", "the", "read_csv", "read_parquet", "read_json", "unnest"].includes(nameToken.normalized)) return null; - if (state.tokens[nameIndex + 1]?.text !== "(") return null; - const close = findMatchingParenToken(state.tokens, nameIndex + 1); - const safeClose = close < 0 ? nameIndex + 1 : close; - const alias = aliasAfter(state.tokens, safeClose + 1, state.dialect); - const sourceName = alias.alias ?? nameToken.normalized; + const isMutationTarget = introducer === "update" || introducer === "into" || (state.statement.kind === "delete" && introducer === "from"); + if (isMutationTarget) return null; + const qualified = readQualifiedName(state.tokens, nameIndex, state.dialect); + if (!qualified || state.tokens[qualified.nextIndex]?.text !== "(") return null; + const { name, qualifierParts } = sourceNameFromQualifiedName(qualified.name); + if (state.dialect.id !== "postgres" && !TABLE_FUNCTION_NAMES.has(name.toLowerCase())) return null; + const close = findMatchingParenToken(state.tokens, qualified.nextIndex); + const safeClose = close < 0 ? qualified.nextIndex : close; + let aliasIndex = safeClose + 1; + if (state.dialect.id === "postgres" && state.tokens[aliasIndex]?.normalized === "with" && state.tokens[aliasIndex + 1]?.normalized === "ordinality") aliasIndex += 2; + const alias = aliasAfter(state.tokens, aliasIndex, state.dialect); + const sourceName = alias.alias ?? name; return { source: { id: `${introducer}:table_function:${sourceIndex}`, kind: "table_function", name: sourceName, - qualifierParts: [], + qualifiedName: qualified.name, + qualifierParts, alias: alias.alias, aliasSpan: alias.aliasSpan, - sourceSpan: { start: nameToken.span.start, end: alias.aliasSpan?.end ?? state.tokens[safeClose]?.span.end ?? nameToken.span.end }, + sourceSpan: { start: qualified.name.span.start, end: state.tokens[alias.nextIndex - 1]?.span.end ?? alias.aliasSpan?.end ?? state.tokens[safeClose]?.span.end ?? qualified.name.span.end }, + columns: alias.columns, unresolved: close < 0, }, nextIndex: alias.nextIndex, @@ -280,7 +318,8 @@ function parseTableSource(state: ParseState, nameIndex: number, introducer: stri const qualified = readQualifiedName(state.tokens, nameIndex, state.dialect); if (!qualified) return null; const { name, qualifierParts } = sourceNameFromQualifiedName(qualified.name); - const alias = aliasAfter(state.tokens, qualified.nextIndex, state.dialect); + const alias = aliasAfter(state.tokens, qualified.nextIndex, state.dialect, { allowCorrelationColumns: state.dialect.id !== "sqlserver" }); + const nextIndex = consumeSqlServerTableHint(state.tokens, alias.nextIndex, state.dialect); const cte = state.cteSources.find((source) => source.name.toLowerCase() === name.toLowerCase()); const kind = cte ? "cte" : introducer === "update" || introducer === "into" || (state.statement.kind === "delete" && introducer === "from") ? "mutation_target" : "table"; const source: SqlSemanticRowSource = { @@ -291,47 +330,67 @@ function parseTableSource(state: ParseState, nameIndex: number, introducer: stri qualifierParts, alias: alias.alias, aliasSpan: alias.aliasSpan, - sourceSpan: { start: qualified.name.span.start, end: alias.aliasSpan?.end ?? qualified.name.span.end }, - columns: cte?.columns, + sourceSpan: { start: qualified.name.span.start, end: state.tokens[nextIndex - 1]?.span.end ?? alias.aliasSpan?.end ?? qualified.name.span.end }, + columns: cte?.columns ? mergeColumnAliases(cte.columns, alias.columns) : undefined, + columnAliases: alias.columns, metadataTarget: { schema: qualifierParts[qualifierParts.length - 1], table: name, }, }; - return { source, nextIndex: alias.nextIndex }; + return { source, nextIndex }; +} + +function isPostgresLateralSource(state: ParseState, target: number, introducer: string): boolean { + if (state.dialect.id !== "postgres" || state.tokens[target]?.normalized !== "lateral" || (introducer !== "from" && introducer !== "join")) return false; + const sourceIndex = target + 1; + if (state.tokens[sourceIndex]?.text === "(") return true; + const qualified = readQualifiedName(state.tokens, sourceIndex, state.dialect); + return !!qualified && state.tokens[qualified.nextIndex]?.text === "("; +} + +function parseRowSource(state: ParseState, target: number, introducer: string, sourceIndex: number): { source: SqlSemanticRowSource; nextIndex: number } | null { + if (isPostgresLateralSource(state, target, introducer)) target += 1; + if (state.tokens[target]?.text === "(") return parseSubquerySource(state, target, introducer, sourceIndex); + return parseTableFunctionSource(state, target, introducer, sourceIndex) ?? parseTableSource(state, target, introducer, sourceIndex); } function parseRowSources(state: ParseState): SqlSemanticRowSource[] { const sources: SqlSemanticRowSource[] = [...state.cteSources]; const rootDepth = state.tokens.reduce((min, item) => Math.min(min, item.depth), Number.POSITIVE_INFINITY); const sourceDepth = Number.isFinite(rootDepth) ? rootDepth : 0; + let inSelectFromClause = false; for (let index = 0; index < state.tokens.length; index += 1) { const item = state.tokens[index]; - if (!item || item.kind !== "word") continue; + if (!item) continue; if (item.depth !== sourceDepth) continue; + if (item.kind === "word") { + if (state.statement.kind === "select" && item.normalized === "from") inSelectFromClause = true; + else if (inSelectFromClause && FROM_CLAUSE_BOUNDARIES.has(item.normalized)) inSelectFromClause = false; + } + if (inSelectFromClause && item.text === ",") { + const parsed = parseRowSource(state, index + 1, "from", sources.length); + if (parsed) { + sources.push(parsed.source); + index = parsed.nextIndex - 1; + } + continue; + } + if (item.kind !== "word") continue; const normalized = item.normalized; if (!TABLE_INTRODUCERS.has(normalized)) continue; if (JOIN_MODIFIERS.has(normalized)) continue; let target = index + 1; while (JOIN_MODIFIERS.has(state.tokens[target]?.normalized ?? "")) target += 1; - if (state.tokens[target]?.text === "(") { - const subquery = parseSubquerySource(state, target, normalized, sources.length); - if (subquery) { - sources.push(subquery.source); - index = subquery.nextIndex - 1; - } - continue; - } - const tableFunction = parseTableFunctionSource(state, target, normalized, sources.length); - if (tableFunction) { - sources.push(tableFunction.source); - index = tableFunction.nextIndex - 1; - continue; - } - const table = parseTableSource(state, target, normalized, sources.length); - if (table) { - sources.push(table.source); - index = table.nextIndex - 1; + for (;;) { + const parsed = parseRowSource(state, target, normalized, sources.length); + if (!parsed) break; + sources.push(parsed.source); + index = parsed.nextIndex - 1; + + const separator = state.tokens[parsed.nextIndex]; + if (normalized !== "from" || separator?.text !== "," || separator.depth !== sourceDepth) break; + target = parsed.nextIndex + 1; } } return dedupeSources(sources); diff --git a/apps/desktop/src/lib/sql/semantic/types.ts b/apps/desktop/src/lib/sql/semantic/types.ts index 9c6d46cc9..c5a1e0c72 100644 --- a/apps/desktop/src/lib/sql/semantic/types.ts +++ b/apps/desktop/src/lib/sql/semantic/types.ts @@ -59,6 +59,7 @@ export interface SqlSemanticRowSource { aliasSpan?: SqlSemanticSpan; sourceSpan: SqlSemanticSpan; columns?: string[]; + columnAliases?: string[]; metadataTarget?: SqlSemanticMetadataTarget; unresolved?: boolean; } diff --git a/apps/desktop/src/lib/sql/sqlCompletion.ts b/apps/desktop/src/lib/sql/sqlCompletion.ts index 7ea3815ec..ad23749e9 100644 --- a/apps/desktop/src/lib/sql/sqlCompletion.ts +++ b/apps/desktop/src/lib/sql/sqlCompletion.ts @@ -1175,6 +1175,7 @@ export interface SqlCompletionObject { export interface SqlCompletionColumn { name: string; table: string; + sourceAlias?: string; schema?: string; dataType?: string; isNullable?: boolean; @@ -1211,6 +1212,7 @@ export interface SqlCompletionReferencedTable { schema?: string; alias?: string; columns?: string[]; + columnAliases?: string[]; } export type SqlStatementKind = "select" | "insert" | "update" | "delete" | "create" | "alter" | "drop" | "unknown"; @@ -3352,9 +3354,10 @@ function buildColumnItems(context: SqlCompletionContext, columnsByTable: Map referencedTableMatchesColumnQualifier(table, q, qLower, qualifiedTarget)); - relevantCols = allColumns.filter((column) => relatedTables.some((table) => columnMatchesReferencedTable(column, table)) || (!!qualifiedTarget && columnMatchesQualifiedTable(column, qualifiedTarget))); + relevantCols = relatedTables.flatMap((table) => completionColumnsForReferencedTable(table, allColumns)); + if (relatedTables.length === 0 && qualifiedTarget) relevantCols = allColumns.filter((column) => columnMatchesQualifiedTable(column, qualifiedTarget)); } else if (context.referencedTables.length > 0) { - relevantCols = allColumns.filter((column) => context.referencedTables.some((table) => columnMatchesReferencedTable(column, table))); + relevantCols = context.referencedTables.flatMap((table) => completionColumnsForReferencedTable(table, allColumns)); } // Count name frequencies to detect duplicates across tables @@ -3369,10 +3372,11 @@ function buildColumnItems(context: SqlCompletionContext, columnsByTable: Map 1) { - const qualifiedKey = `${c.table}.${c.name}`; + const qualifier = c.sourceAlias ?? c.table; + const qualifiedKey = `${qualifier}.${c.name}`; if (seen.has(qualifiedKey)) continue; seen.add(qualifiedKey); - uniqueColumns.push({ ...c, key: c.key, displayLabel: `${c.table}.${c.name}` }); + uniqueColumns.push({ ...c, key: c.key, displayLabel: `${qualifier}.${c.name}` }); } else { if (seen.has(c.name)) continue; seen.add(c.name); @@ -3402,19 +3406,24 @@ function buildColumnItems(context: SqlCompletionContext, columnsByTable: Map): boolean { - if (!context.suggestColumns || !context.prefix || context.referencedTables.length === 0) return false; +function completionColumnsForReferencedTable(table: SqlCompletionReferencedTable, columns: readonly T[]): T[] { + const matched = columns.filter((column) => columnMatchesReferencedTable(column, table)); + const aliasedColumns = applyReferencedColumnAliases(table, matched); + if (!table.alias) return aliasedColumns; + return aliasedColumns.map((column) => ({ ...column, sourceAlias: table.alias })); +} - for (const [key, cols] of columnsByTable.entries()) { - for (const column of cols) { - if (!matchesPrefix(column.name, context.prefix)) continue; - if (context.referencedTables.some((table) => columnMatchesReferencedTable({ ...column, key }, table))) { - return true; - } - } - } +function applyReferencedColumnAliases(table: SqlCompletionReferencedTable, columns: readonly T[]): T[] { + if (!table.columnAliases?.length) return [...columns]; + return columns.map((column, index) => { + const alias = table.columnAliases?.[index]; + return alias ? { ...column, name: alias } : column; + }); +} - return false; +function hasMatchingReferencedColumnPrefix(context: SqlCompletionContext, columnsByTable: Map): boolean { + if (!context.suggestColumns || !context.prefix || context.referencedTables.length === 0) return false; + return context.referencedTables.some((table) => columnsForReferencedTable(table, columnsByTable).some((column) => matchesPrefix(column.name, context.prefix))); } function qualifiedTableTargetFromContext(context: SqlCompletionContext): { schema: string; table: string } | null { @@ -3458,7 +3467,7 @@ function buildColumnApply(column: SqlCompletionColumn & { displayLabel: string } if (context.qualifier || column.displayLabel === column.name || !column.displayLabel.includes(".")) { return quoteSqlIdentifier(column.name, dialect); } - return `${quoteSqlIdentifier(column.table, dialect)}.${quoteSqlIdentifier(column.name, dialect)}`; + return `${quoteSqlIdentifier(column.sourceAlias ?? column.table, dialect)}.${quoteSqlIdentifier(column.name, dialect)}`; } function isKeyColumn(name: string): boolean { @@ -3510,7 +3519,7 @@ function columnsForReferencedTable(table: SqlCompletionReferencedTable, columnsB const keys = table.schema ? [`${table.schema}.${table.name}`, table.name] : [table.name]; for (const key of keys) { const columns = columnsByTable.get(key); - if (columns) return columns; + if (columns) return applyReferencedColumnAliases(table, columns); } return []; } diff --git a/packages/app-tests/sqlCompletion.test.ts b/packages/app-tests/sqlCompletion.test.ts index 69d70ae51..588c3d17b 100644 --- a/packages/app-tests/sqlCompletion.test.ts +++ b/packages/app-tests/sqlCompletion.test.ts @@ -2091,7 +2091,7 @@ test("filters data type keywords out of SELECT context", () => { // --- Qualified column names for duplicates --- -test("shows qualified column names when multiple tables share column name", () => { +test("uses row-source aliases when multiple tables share column names", () => { const sql = "select from public.users u join public.orders o on u.id = o.user_id"; const items = buildSqlCompletionItems(sql, "select ".length, { tables, @@ -2099,12 +2099,12 @@ test("shows qualified column names when multiple tables share column name", () = }); const columns = items.filter((item) => item.type === "column"); assert.ok( - columns.some((item) => item.label === "users.id"), - "should show users.id", + columns.some((item) => item.label === "u.id" && item.apply === "u.id"), + "should show u.id", ); assert.ok( - columns.some((item) => item.label === "orders.id"), - "should show orders.id", + columns.some((item) => item.label === "o.id" && item.apply === "o.id"), + "should show o.id", ); assert.ok( columns.some((item) => item.label === "name"),