diff --git a/apps/desktop/src/components/editor/QueryEditor.vue b/apps/desktop/src/components/editor/QueryEditor.vue index 109980f0be..ac71af608f 100644 --- a/apps/desktop/src/components/editor/QueryEditor.vue +++ b/apps/desktop/src/components/editor/QueryEditor.vue @@ -57,8 +57,14 @@ import { shouldRunSqlSemanticDiagnostics, type SqlSemanticDiagnostic, } from "@/lib/sqlSemanticDiagnostics"; -import type { SqlCompletionColumn } from "@/lib/sqlCompletion"; -import type { DatabaseType, SqlReferenceAnalysis, SqlTableReference, SqlTextSpan } from "@/types/database"; +import type { SqlCompletionColumn, SqlCompletionForeignKey } from "@/lib/sqlCompletion"; +import type { + DatabaseType, + ForeignKeyInfo, + SqlReferenceAnalysis, + SqlTableReference, + SqlTextSpan, +} from "@/types/database"; import { vscodeSelectionLayer } from "@/lib/codemirrorVscodeSelectionLayer"; const props = defineProps<{ @@ -195,6 +201,7 @@ function editorThemeAppearance() { let cachedTables: Array<{ name: string; schema?: string; type?: "table" | "view" }> = []; // Persistent column cache keyed by "schema.table" or "table" const cachedColumnsByTable = new Map(); +const cachedForeignKeysByTable = new Map(); const zoomCommitScheduler = createEditorZoomCommitScheduler((fontSize) => { if (settingsStore.editorSettings.fontSize === fontSize) return; @@ -429,6 +436,28 @@ async function ensureColumnsForTable(table: { name: string; schema?: string | nu cachedColumnsByTable.set(cacheKey, columns); } +async function ensureForeignKeysForTable(table: { name: string; schema?: string | null }) { + const cacheKey = completionCacheKey(table); + if (cachedForeignKeysByTable.has(cacheKey) || !props.connectionId || props.database == null) return; + const querySchema = table.schema ?? props.database; + try { + const foreignKeys = await api.listForeignKeys(props.connectionId, props.database, querySchema, table.name); + cachedForeignKeysByTable.set( + cacheKey, + foreignKeys.map((foreignKey: ForeignKeyInfo) => ({ + name: foreignKey.name, + column: foreignKey.column, + ref_schema: foreignKey.ref_schema, + ref_table: foreignKey.ref_table, + ref_column: foreignKey.ref_column, + })), + ); + } catch (e) { + console.warn(`[DBX] Failed to load foreign keys for ${cacheKey}:`, e); + cachedForeignKeysByTable.set(cacheKey, []); + } +} + function createHoverDom(title: string, detail: string, rows: string[] = []) { const dom = document.createElement("div"); dom.className = "rounded-md border bg-popover px-3 py-2 text-xs text-popover-foreground shadow-md"; @@ -1027,8 +1056,19 @@ async function performAsyncCompletionWithResult( ); if (epoch !== completionEpoch) return null; + if (completionContext.suggestJoinConditions) { + await Promise.all( + refs.map(async (refTable) => { + if (refTable.columns && refTable.columns.length > 0) return; + await ensureForeignKeysForTable(refTable); + }), + ); + if (epoch !== completionEpoch) return null; + } + // Build columnsByTable — from cache or CTE definitions const columnsByTable = new Map(); + const foreignKeysByTable = new Map(); if (insertColumnsByTable.size > 0) { for (const [key, cols] of insertColumnsByTable.entries()) { columnsByTable.set(key, cols); @@ -1052,6 +1092,10 @@ async function performAsyncCompletionWithResult( if (cached) { columnsByTable.set(cacheKey, cached); } + const cachedForeignKeys = cachedForeignKeysByTable.get(cacheKey); + if (cachedForeignKeys) { + foreignKeysByTable.set(cacheKey, cachedForeignKeys); + } } } @@ -1068,6 +1112,7 @@ async function performAsyncCompletionWithResult( const items = buildSqlCompletionItemsFromContext(effectiveContext, { tables, columnsByTable, + foreignKeysByTable, schemas: schemaNames, translations: completionTranslations.value, snippets: settingsStore.editorSettings.snippets, @@ -1088,6 +1133,7 @@ function isReferencedTableQualifier(completionContext: ReturnType { diff --git a/apps/desktop/src/lib/sqlCompletion.ts b/apps/desktop/src/lib/sqlCompletion.ts index d0574c0254..5b1234a168 100644 --- a/apps/desktop/src/lib/sqlCompletion.ts +++ b/apps/desktop/src/lib/sqlCompletion.ts @@ -608,6 +608,14 @@ export interface SqlCompletionColumn { isNullable?: boolean; } +export interface SqlCompletionForeignKey { + name: string; + column: string; + ref_schema?: string | null; + ref_table: string; + ref_column: string; +} + export interface SqlCompletionItem { label: string; type: "keyword" | "table" | "column" | "snippet" | "function" | "schema"; @@ -671,6 +679,7 @@ export function buildSqlCompletionItems( input: { tables: SqlCompletionTable[]; columnsByTable: Map; + foreignKeysByTable?: Map; schemas?: string[]; translations?: SqlCompletionTranslations; dialect?: "mysql" | "postgres" | "sqlserver"; @@ -685,6 +694,7 @@ export function buildSqlCompletionItemsFromContext( input: { tables: SqlCompletionTable[]; columnsByTable: Map; + foreignKeysByTable?: Map; schemas?: string[]; translations?: SqlCompletionTranslations; snippets?: SqlSnippet[]; @@ -714,7 +724,7 @@ export function buildSqlCompletionItemsFromContext( } if (!context.exclusiveTableSuggestions && !context.exclusiveColumnSuggestions && context.suggestJoinConditions) { - items.push(...buildJoinConditionItems(context, input.columnsByTable, dialect)); + items.push(...buildJoinConditionItems(context, input.columnsByTable, input.foreignKeysByTable, dialect)); } if (context.suggestKeywords) { @@ -1836,6 +1846,7 @@ function buildColumnDetail(column: SqlCompletionColumn): string { function buildJoinConditionItems( context: SqlCompletionContext, columnsByTable: Map, + foreignKeysByTable?: Map, dialect?: "mysql" | "postgres" | "sqlserver", ): SqlCompletionItem[] { const refs = context.referencedTables; @@ -1849,6 +1860,7 @@ function buildJoinConditionItems( const previousColumns = columnsForReferencedTable(previous, columnsByTable); const latestColumns = columnsForReferencedTable(latest, columnsByTable); items.push( + ...buildForeignKeyJoinConditionItemsForPair(previous, latest, foreignKeysByTable, context.prefix, dialect), ...buildJoinConditionItemsForPair(previous, previousColumns, latest, latestColumns, context.prefix, dialect), ); } @@ -1868,6 +1880,128 @@ function columnsForReferencedTable( return []; } +function foreignKeysForReferencedTable( + table: SqlCompletionReferencedTable, + foreignKeysByTable?: Map, +): SqlCompletionForeignKey[] { + if (!foreignKeysByTable) return []; + const keys = table.schema ? [`${table.schema}.${table.name}`, table.name] : [table.name]; + for (const key of keys) { + const foreignKeys = foreignKeysByTable.get(key); + if (foreignKeys) return foreignKeys; + } + return []; +} + +function buildForeignKeyJoinConditionItemsForPair( + left: SqlCompletionReferencedTable, + right: SqlCompletionReferencedTable, + foreignKeysByTable?: Map, + prefix = "", + dialect?: "mysql" | "postgres" | "sqlserver", +): SqlCompletionItem[] { + if (!foreignKeysByTable) return []; + return [ + ...buildDirectionalForeignKeyJoinConditionItems( + left, + right, + foreignKeysForReferencedTable(left, foreignKeysByTable), + prefix, + dialect, + ), + ...buildDirectionalForeignKeyJoinConditionItems( + right, + left, + foreignKeysForReferencedTable(right, foreignKeysByTable), + prefix, + dialect, + ), + ]; +} + +function buildDirectionalForeignKeyJoinConditionItems( + owner: SqlCompletionReferencedTable, + referenced: SqlCompletionReferencedTable, + foreignKeys: SqlCompletionForeignKey[], + prefix: string, + dialect?: "mysql" | "postgres" | "sqlserver", +): SqlCompletionItem[] { + const matchingForeignKeys = foreignKeys.filter((foreignKey) => + referencedTableMatchesName(referenced, foreignKey.ref_table, foreignKey.ref_schema), + ); + const groups = groupForeignKeysByConstraint(matchingForeignKeys); + const items: SqlCompletionItem[] = []; + + for (const group of groups) { + const parts = group.map((foreignKey) => + buildJoinConditionPart(owner, foreignKey.column, referenced, foreignKey.ref_column, dialect), + ); + const label = parts.map((part) => part.label).join(" AND "); + if (!label || (prefix && !matchesPrefix(label, prefix))) continue; + const apply = parts.map((part) => part.apply).join(" AND "); + items.push({ + label, + type: "snippet", + detail: group.length > 1 ? "JOIN condition from composite foreign key" : "JOIN condition from foreign key", + apply, + boost: 3200 + group.length, + }); + } + + return items; +} + +function buildJoinConditionPart( + owner: SqlCompletionReferencedTable, + ownerColumn: string, + referenced: SqlCompletionReferencedTable, + referencedColumn: string, + dialect?: "mysql" | "postgres" | "sqlserver", +): { label: string; apply: string } { + const ownerRef = owner.alias || owner.name; + const referencedRef = referenced.alias || referenced.name; + const ownerApplyRef = owner.alias ? owner.alias : quoteSqlIdentifier(owner.name, dialect); + const referencedApplyRef = referenced.alias ? referenced.alias : quoteSqlIdentifier(referenced.name, dialect); + return { + label: `${ownerRef}.${ownerColumn} = ${referencedRef}.${referencedColumn}`, + apply: `${ownerApplyRef}.${quoteSqlIdentifier(ownerColumn, dialect)} = ${referencedApplyRef}.${quoteSqlIdentifier(referencedColumn, dialect)}`, + }; +} + +function groupForeignKeysByConstraint(foreignKeys: SqlCompletionForeignKey[]): SqlCompletionForeignKey[][] { + const groups = new Map(); + for (const foreignKey of foreignKeys) { + const key = `${foreignKey.name || `${foreignKey.column}->${foreignKey.ref_table}.${foreignKey.ref_column}`}:${foreignKey.ref_table}`; + if (!groups.has(key)) groups.set(key, []); + groups.get(key)!.push(foreignKey); + } + return [...groups.values()]; +} + +function referencedTableMatchesName( + table: SqlCompletionReferencedTable, + candidate: string, + candidateSchema?: string | null, +): boolean { + const normalizedCandidate = normalizeTableName(candidate); + if (normalizeTableName(table.name) !== normalizedCandidate) return false; + if (!candidateSchema || !table.schema) return true; + return normalizeIdentifierPart(table.schema) === normalizeIdentifierPart(candidateSchema); +} + +function normalizeTableName(name: string): string { + return name + .split(".") + .filter(Boolean) + .pop()! + .replace(/^["`[]|["`\]]$/g, "") + .toLowerCase(); +} + +function normalizeIdentifierPart(name: string): string { + return name.replace(/^["`[]|["`\]]$/g, "").toLowerCase(); +} + function buildJoinConditionItemsForPair( left: SqlCompletionReferencedTable, leftColumns: SqlCompletionColumn[], @@ -1884,72 +2018,222 @@ function buildJoinConditionItemsForPair( const leftTableKey = singularTableName(left.name); const rightTableKey = singularTableName(right.name); + const leftByName = indexColumnsByLowerName(leftColumns); + const rightByName = indexColumnsByLowerName(rightColumns); + const emittedPairs = new Set(); + + const addPair = ( + leftColumn: SqlCompletionColumn | undefined, + rightColumn: SqlCompletionColumn | undefined, + boost: number, + ) => { + if (!leftColumn || !rightColumn || !areJoinColumnTypesCompatible(leftColumn, rightColumn)) return; + const key = `${leftColumn.name.toLowerCase()}:${rightColumn.name.toLowerCase()}`; + if (emittedPairs.has(key)) return; + emittedPairs.add(key); + const label = `${leftRef}.${leftColumn.name} = ${rightRef}.${rightColumn.name}`; + if (prefix && !matchesPrefix(label, prefix)) return; + const apply = `${leftApplyRef}.${quoteSqlIdentifier(leftColumn.name, dialect)} = ${rightApplyRef}.${quoteSqlIdentifier(rightColumn.name, dialect)}`; + items.push({ + label, + type: "snippet", + detail: "JOIN condition", + apply, + boost, + }); + }; + + const leftId = leftByName.get("id")?.[0]; + const rightId = rightByName.get("id")?.[0]; + + // Pattern 1: a.id = b.{singular_a}_id (e.g., users.id = orders.user_id) + addPair(leftId, rightByName.get(`${leftTableKey}_id`)?.[0], 2300); + // Pattern 2: a.{singular_b}_id = b.id (e.g., orders.user_id = users.id) + addPair(leftByName.get(`${rightTableKey}_id`)?.[0], rightId, 2300); + + // Pattern 3/4: same-name columns, with FK-looking names above generic shared columns. + for (const [name, leftMatches] of leftByName.entries()) { + if (name === "id") continue; + const rightMatches = rightByName.get(name); + if (!rightMatches?.length) continue; + addPair(leftMatches[0], rightMatches[0], name.endsWith("_id") ? 2000 : 1700); + } + + // Pattern 5: parent_id -> id (self-referencing / hierarchical) + if (leftTableKey === rightTableKey) { + addPair(leftByName.get("parent_id")?.[0], rightId, 2100); + addPair(leftId, rightByName.get("parent_id")?.[0], 2100); + } + + // Pattern 6: created_by / modified_by / owned_by -> users.id + for (const auditColumnName of ["created_by", "modified_by", "owned_by"]) { + addPair(leftId, rightByName.get(auditColumnName)?.[0], 1800); + addPair(leftByName.get(auditColumnName)?.[0], rightId, 1800); + } + + // Pattern 7: Generic FK column -> id when table names do not reveal the relationship. for (const leftColumn of leftColumns) { - for (const rightColumn of rightColumns) { - const leftName = leftColumn.name.toLowerCase(); - const rightName = rightColumn.name.toLowerCase(); - const leftLabel = `${leftRef}.${leftColumn.name}`; - const rightLabel = `${rightRef}.${rightColumn.name}`; - let boost = 0; + const leftName = leftColumn.name.toLowerCase(); + if (leftName !== "id" && leftName.endsWith("_id")) addPair(leftColumn, rightId, 1650); + } + for (const rightColumn of rightColumns) { + const rightName = rightColumn.name.toLowerCase(); + if (rightName !== "id" && rightName.endsWith("_id")) addPair(leftId, rightColumn, 1650); + } - // Pattern 1: a.id = b.{singular_a}_id (e.g., users.id = orders.user_id) - if (leftName === "id" && rightName === `${leftTableKey}_id`) { - boost = 2300; - } - // Pattern 2: a.{singular_b}_id = b.id (e.g., orders.user_id = users.id) - else if (rightName === "id" && leftName === `${rightTableKey}_id`) { - boost = 2300; - } - // Pattern 3: Same FK column name in both tables (e.g., both have user_id) - else if (leftName === rightName && leftName.endsWith("_id")) { - boost = 2000; - } - // Pattern 4: Same column name that isn't "id" (e.g., both have "code", "email") - else if (leftName !== "id" && leftName === rightName) { - boost = 1700; - } - // Pattern 5: parent_id → id (self-referencing / hierarchical) - else if (leftName === "parent_id" && rightName === "id" && leftTableKey === rightTableKey) { - boost = 2100; - } else if (rightName === "parent_id" && leftName === "id" && leftTableKey === rightTableKey) { - boost = 2100; - } - // Pattern 6: created_by / modified_by / owned_by → users.id - else if ( - leftName === "id" && - (rightName === "created_by" || rightName === "modified_by" || rightName === "owned_by") - ) { - boost = 1800; - } else if ( - rightName === "id" && - (leftName === "created_by" || leftName === "modified_by" || leftName === "owned_by") - ) { - boost = 1800; - } - // Pattern 7: Generic FK column → id when table names do not reveal the relationship. - else if (leftName !== "id" && leftName.endsWith("_id") && rightName === "id") { - boost = 1650; - } else if (rightName !== "id" && rightName.endsWith("_id") && leftName === "id") { - boost = 1650; - } + items.push( + ...buildCompositeHeuristicJoinConditionItems(left, leftColumns, right, leftByName, rightByName, prefix, dialect), + ); - if (!boost) continue; - const label = `${leftLabel} = ${rightLabel}`; - if (prefix && !matchesPrefix(label, prefix)) continue; - const apply = `${leftApplyRef}.${quoteSqlIdentifier(leftColumn.name, dialect)} = ${rightApplyRef}.${quoteSqlIdentifier(rightColumn.name, dialect)}`; - items.push({ - label, - type: "snippet", - detail: "JOIN condition", - apply, - boost, - }); + return items; +} + +function indexColumnsByLowerName(columns: SqlCompletionColumn[]): Map { + const index = new Map(); + for (const column of columns) { + const key = column.name.toLowerCase(); + const existing = index.get(key); + if (existing) existing.push(column); + else index.set(key, [column]); + } + return index; +} + +function buildCompositeHeuristicJoinConditionItems( + left: SqlCompletionReferencedTable, + leftColumns: SqlCompletionColumn[], + right: SqlCompletionReferencedTable, + leftByName: Map, + rightByName: Map, + prefix: string, + dialect?: "mysql" | "postgres" | "sqlserver", +): SqlCompletionItem[] { + const leftId = leftByName.get("id")?.[0]; + const rightId = rightByName.get("id")?.[0]; + const leftTableKey = singularTableName(left.name); + const rightTableKey = singularTableName(right.name); + const candidates: Array<{ parent: "left" | "right"; parentId: SqlCompletionColumn; childFk: SqlCompletionColumn }> = + []; + const rightNamedFk = rightByName.get(`${leftTableKey}_id`)?.[0]; + const leftNamedFk = leftByName.get(`${rightTableKey}_id`)?.[0]; + if (leftId && rightNamedFk && areJoinColumnTypesCompatible(leftId, rightNamedFk)) { + candidates.push({ parent: "left", parentId: leftId, childFk: rightNamedFk }); + } + if (rightId && leftNamedFk && areJoinColumnTypesCompatible(leftNamedFk, rightId)) { + candidates.push({ parent: "right", parentId: rightId, childFk: leftNamedFk }); + } + if (candidates.length === 0) return []; + + const sharedScopeColumns = leftColumns + .map((leftColumn) => { + const name = leftColumn.name.toLowerCase(); + const rightColumn = rightByName.get(name)?.[0]; + if (!rightColumn || !isLikelyScopeColumnName(name) || !areJoinColumnTypesCompatible(leftColumn, rightColumn)) { + return null; + } + return { leftColumn, rightColumn }; + }) + .filter((value): value is { leftColumn: SqlCompletionColumn; rightColumn: SqlCompletionColumn } => !!value) + .slice(0, 2); + if (sharedScopeColumns.length === 0) return []; + + const leftRef = left.alias || left.name; + const rightRef = right.alias || right.name; + const leftApplyRef = left.alias ? left.alias : quoteSqlIdentifier(left.name, dialect); + const rightApplyRef = right.alias ? right.alias : quoteSqlIdentifier(right.name, dialect); + const items: SqlCompletionItem[] = []; + + for (const candidate of candidates.slice(0, 2)) { + const parts = sharedScopeColumns.map(({ leftColumn, rightColumn }) => + buildHeuristicJoinConditionPart(leftRef, leftApplyRef, leftColumn, rightRef, rightApplyRef, rightColumn, dialect), + ); + if (candidate.parent === "left") { + parts.push( + buildHeuristicJoinConditionPart( + leftRef, + leftApplyRef, + candidate.parentId, + rightRef, + rightApplyRef, + candidate.childFk, + dialect, + ), + ); + } else { + parts.push( + buildHeuristicJoinConditionPart( + leftRef, + leftApplyRef, + candidate.childFk, + rightRef, + rightApplyRef, + candidate.parentId, + dialect, + ), + ); } + const label = parts.map((part) => part.label).join(" AND "); + if (prefix && !matchesPrefix(label, prefix)) continue; + items.push({ + label, + type: "snippet", + detail: "Likely composite JOIN condition", + apply: parts.map((part) => part.apply).join(" AND "), + boost: 2400 + parts.length, + }); } return items; } +function buildHeuristicJoinConditionPart( + leftRef: string, + leftApplyRef: string, + leftColumn: SqlCompletionColumn, + rightRef: string, + rightApplyRef: string, + rightColumn: SqlCompletionColumn, + dialect?: "mysql" | "postgres" | "sqlserver", +): { label: string; apply: string } { + return { + label: `${leftRef}.${leftColumn.name} = ${rightRef}.${rightColumn.name}`, + apply: `${leftApplyRef}.${quoteSqlIdentifier(leftColumn.name, dialect)} = ${rightApplyRef}.${quoteSqlIdentifier(rightColumn.name, dialect)}`, + }; +} + +function isLikelyScopeColumnName(name: string): boolean { + return ( + name !== "id" && + (name.endsWith("_id") || + name === "tenant" || + name === "tenant_id" || + name === "account_id" || + name === "workspace_id" || + name === "organization_id" || + name === "org_id") + ); +} + +function areJoinColumnTypesCompatible(left: SqlCompletionColumn, right: SqlCompletionColumn): boolean { + const leftType = normalizeJoinType(left.dataType); + const rightType = normalizeJoinType(right.dataType); + if (!leftType || !rightType) return true; + return leftType === rightType; +} + +function normalizeJoinType(dataType?: string): string | null { + if (!dataType) return null; + const type = dataType.toLowerCase(); + if (/\b(uuid|uniqueidentifier)\b/.test(type)) return "uuid"; + if (/\b(bigint|int8|integer|int|int4|smallint|int2|tinyint|serial|bigserial|number|numeric|decimal)\b/.test(type)) { + return "number"; + } + if (/\b(char|text|clob|string|varchar|nvarchar|nchar|uuid)\b/.test(type)) return "text"; + if (/\b(bool|boolean|bit)\b/.test(type)) return "boolean"; + if (/\b(date|time|timestamp|datetime)\b/.test(type)) return "temporal"; + return type.replace(/\(.+\)/, "").trim() || null; +} + function singularTableName(name: string): string { const lower = name.toLowerCase(); // Irregular plurals diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts index a2e3da3da0..e56d9d223a 100644 --- a/apps/desktop/src/types/database.ts +++ b/apps/desktop/src/types/database.ts @@ -221,6 +221,7 @@ export interface IndexInfo { export interface ForeignKeyInfo { name: string; column: string; + ref_schema?: string | null; ref_table: string; ref_column: string; } diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 24c5525930..ae8981660b 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -1140,11 +1140,11 @@ pub async fn list_indexes(pool: &MySqlPool, database: &str, table: &str) -> Resu pub async fn list_foreign_keys(pool: &MySqlPool, database: &str, table: &str) -> Result, String> { let sql = format!( "SELECT kcu.CONSTRAINT_NAME, kcu.COLUMN_NAME, \ - kcu.REFERENCED_TABLE_NAME, kcu.REFERENCED_COLUMN_NAME \ + kcu.REFERENCED_TABLE_SCHEMA, kcu.REFERENCED_TABLE_NAME, kcu.REFERENCED_COLUMN_NAME \ FROM information_schema.KEY_COLUMN_USAGE kcu \ WHERE kcu.TABLE_SCHEMA = {} AND kcu.TABLE_NAME = {} \ AND kcu.REFERENCED_TABLE_NAME IS NOT NULL \ - ORDER BY kcu.CONSTRAINT_NAME", + ORDER BY kcu.CONSTRAINT_NAME, kcu.ORDINAL_POSITION", quote_value(database), quote_value(table), ); @@ -1157,6 +1157,7 @@ pub async fn list_foreign_keys(pool: &MySqlPool, database: &str, table: &str) -> .map(|row| ForeignKeyInfo { name: get_str_by_name(row, "CONSTRAINT_NAME"), column: get_str_by_name(row, "COLUMN_NAME"), + ref_schema: Some(get_str_by_name(row, "REFERENCED_TABLE_SCHEMA")), ref_table: get_str_by_name(row, "REFERENCED_TABLE_NAME"), ref_column: get_str_by_name(row, "REFERENCED_COLUMN_NAME"), }) diff --git a/crates/dbx-core/src/db/ob_oracle.rs b/crates/dbx-core/src/db/ob_oracle.rs index ebd3fd89a8..b74f05d336 100644 --- a/crates/dbx-core/src/db/ob_oracle.rs +++ b/crates/dbx-core/src/db/ob_oracle.rs @@ -214,7 +214,7 @@ pub async fn list_foreign_keys( ) -> Result, String> { let sql = format!( "SELECT ac.CONSTRAINT_NAME, acc.COLUMN_NAME, \ - ac2.TABLE_NAME AS R_TABLE, acc2.COLUMN_NAME AS R_COLUMN \ + ac2.OWNER AS R_OWNER, ac2.TABLE_NAME AS R_TABLE, acc2.COLUMN_NAME AS R_COLUMN \ FROM ALL_CONSTRAINTS ac \ JOIN ALL_CONS_COLUMNS acc ON ac.CONSTRAINT_NAME = acc.CONSTRAINT_NAME AND ac.OWNER = acc.OWNER \ JOIN ALL_CONSTRAINTS ac2 ON ac.R_CONSTRAINT_NAME = ac2.CONSTRAINT_NAME AND ac.R_OWNER = ac2.OWNER \ @@ -234,8 +234,9 @@ pub async fn list_foreign_keys( .map(|row| ForeignKeyInfo { name: get_str(row, 0), column: get_str(row, 1), - ref_table: get_str(row, 2), - ref_column: get_str(row, 3), + ref_schema: Some(get_str(row, 2)), + ref_table: get_str(row, 3), + ref_column: get_str(row, 4), }) .collect()) } diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 707837d250..8781fb7dcb 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -1099,17 +1099,24 @@ pub async fn list_foreign_keys(pool: &Pool, schema: &str, table: &str) -> Result let client = pool.get().await.map_err(|e| e.to_string())?; let stmt = client .prepare_cached( - "SELECT kcu.constraint_name, kcu.column_name, \ - ccu.table_name AS ref_table, ccu.column_name AS ref_column \ - FROM information_schema.key_column_usage kcu \ + "SELECT fk.constraint_name, fk.column_name, \ + pk.table_schema AS ref_schema, pk.table_name AS ref_table, pk.column_name AS ref_column \ + FROM information_schema.table_constraints tc \ + JOIN information_schema.key_column_usage fk \ + ON fk.constraint_name = tc.constraint_name \ + AND fk.constraint_schema = tc.constraint_schema \ + AND fk.table_schema = tc.table_schema \ + AND fk.table_name = tc.table_name \ JOIN information_schema.referential_constraints rc \ - ON kcu.constraint_name = rc.constraint_name \ - AND kcu.constraint_schema = rc.constraint_schema \ - JOIN information_schema.constraint_column_usage ccu \ - ON rc.unique_constraint_name = ccu.constraint_name \ - AND rc.unique_constraint_schema = ccu.constraint_schema \ - WHERE kcu.table_schema = $1 AND kcu.table_name = $2 \ - ORDER BY kcu.constraint_name", + ON rc.constraint_name = tc.constraint_name \ + AND rc.constraint_schema = tc.constraint_schema \ + JOIN information_schema.key_column_usage pk \ + ON pk.constraint_name = rc.unique_constraint_name \ + AND pk.constraint_schema = rc.unique_constraint_schema \ + AND pk.ordinal_position = fk.position_in_unique_constraint \ + WHERE tc.constraint_type = 'FOREIGN KEY' \ + AND fk.table_schema = $1 AND fk.table_name = $2 \ + ORDER BY fk.constraint_name, fk.ordinal_position", ) .await .map_err(|e| e.to_string())?; @@ -1120,8 +1127,9 @@ pub async fn list_foreign_keys(pool: &Pool, schema: &str, table: &str) -> Result .map(|row| ForeignKeyInfo { name: row.get::<_, String>(0), column: row.get::<_, String>(1), - ref_table: row.get::<_, String>(2), - ref_column: row.get::<_, String>(3), + ref_schema: Some(row.get::<_, String>(2)), + ref_table: row.get::<_, String>(3), + ref_column: row.get::<_, String>(4), }) .collect()) } diff --git a/crates/dbx-core/src/db/sqlite.rs b/crates/dbx-core/src/db/sqlite.rs index f81fe172c5..3cc2f672b3 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -451,6 +451,7 @@ pub async fn list_foreign_keys(pool: &SqliteHandle, _schema: &str, table: &str) Ok(ForeignKeyInfo { name: format!("fk_{}", row.get::<_, i32>("id")?), column: row.get("from")?, + ref_schema: None, ref_table: row.get("table")?, ref_column: row.get("to")?, }) diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index ad3990cf59..dc18fc0c35 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -528,14 +528,14 @@ pub async fn list_foreign_keys( table: &str, ) -> Result, String> { let sql = format!( - "SELECT fk.name, c.name, rt.name, rc.name \ + "SELECT fk.name, c.name, SCHEMA_NAME(rt.schema_id), rt.name, rc.name \ FROM sys.foreign_keys fk \ JOIN sys.foreign_key_columns fkc ON fk.object_id = fkc.constraint_object_id \ JOIN sys.columns c ON fkc.parent_object_id = c.object_id AND fkc.parent_column_id = c.column_id \ JOIN sys.tables rt ON fkc.referenced_object_id = rt.object_id \ JOIN sys.columns rc ON fkc.referenced_object_id = rc.object_id AND fkc.referenced_column_id = rc.column_id \ WHERE fk.parent_object_id = OBJECT_ID('{s}.{t}') \ - ORDER BY fk.name", + ORDER BY fk.name, fkc.constraint_column_id", s = schema.replace('\'', "''"), t = table.replace('\'', "''") ); @@ -546,8 +546,9 @@ pub async fn list_foreign_keys( .map(|row| ForeignKeyInfo { name: row.get::<&str, _>(0).unwrap_or("").to_string(), column: row.get::<&str, _>(1).unwrap_or("").to_string(), - ref_table: row.get::<&str, _>(2).unwrap_or("").to_string(), - ref_column: row.get::<&str, _>(3).unwrap_or("").to_string(), + ref_schema: Some(row.get::<&str, _>(2).unwrap_or("").to_string()), + ref_table: row.get::<&str, _>(3).unwrap_or("").to_string(), + ref_column: row.get::<&str, _>(4).unwrap_or("").to_string(), }) .collect()) } diff --git a/crates/dbx-core/src/schema_diff.rs b/crates/dbx-core/src/schema_diff.rs index d61b7d5f27..d688687fb7 100644 --- a/crates/dbx-core/src/schema_diff.rs +++ b/crates/dbx-core/src/schema_diff.rs @@ -462,6 +462,13 @@ pub fn diff_foreign_keys(source: &[ForeignKeyInfo], target: &[ForeignKeyInfo]) - if source_fk.ref_table != target_fk.ref_table { changes.push(format!("ref table: {} → {}", target_fk.ref_table, source_fk.ref_table)); } + if source_fk.ref_schema != target_fk.ref_schema { + changes.push(format!( + "ref schema: {} → {}", + target_fk.ref_schema.as_deref().unwrap_or(""), + source_fk.ref_schema.as_deref().unwrap_or("") + )); + } if source_fk.ref_column != target_fk.ref_column { changes.push(format!("ref column: {} → {}", target_fk.ref_column, source_fk.ref_column)); } @@ -872,6 +879,7 @@ mod tests { ForeignKeyInfo { name: if overrides.name.is_empty() { "orders_user_id_fk".to_string() } else { overrides.name }, column: if overrides.column.is_empty() { "user_id".to_string() } else { overrides.column }, + ref_schema: overrides.ref_schema, ref_table: if overrides.ref_table.is_empty() { "users".to_string() } else { overrides.ref_table }, ref_column: if overrides.ref_column.is_empty() { "id".to_string() } else { overrides.ref_column }, } @@ -929,12 +937,14 @@ mod tests { foreign_key(ForeignKeyInfo { name: "orders_user_id_fk".to_string(), column: String::new(), + ref_schema: None, ref_table: String::new(), ref_column: String::new(), }), foreign_key(ForeignKeyInfo { name: "orders_account_id_fk".to_string(), column: "account_id".to_string(), + ref_schema: None, ref_table: "accounts".to_string(), ref_column: String::new(), }), @@ -943,12 +953,14 @@ mod tests { foreign_key(ForeignKeyInfo { name: "orders_user_id_fk".to_string(), column: String::new(), + ref_schema: None, ref_table: "members".to_string(), ref_column: String::new(), }), foreign_key(ForeignKeyInfo { name: "orders_region_id_fk".to_string(), column: "region_id".to_string(), + ref_schema: None, ref_table: "regions".to_string(), ref_column: String::new(), }), @@ -995,6 +1007,7 @@ mod tests { source: Some(foreign_key(ForeignKeyInfo { name: "orders_user_id_fk".to_string(), column: String::new(), + ref_schema: None, ref_table: "users".to_string(), ref_column: String::new(), })), diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index a9fb75033d..52579b50f1 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -271,7 +271,8 @@ fn generate_postgres_foreign_key_ddl(foreign_keys: &[db::ForeignKeyInfo], table: .map(|foreign_key| quote_identifier(&foreign_key.ref_column, &DatabaseType::Postgres)) .collect::>() .join(", "); - let referenced_table = qualified_table(&group[0].ref_table, schema, &DatabaseType::Postgres); + let referenced_schema = group[0].ref_schema.as_deref().unwrap_or(schema); + let referenced_table = qualified_table(&group[0].ref_table, referenced_schema, &DatabaseType::Postgres); statements.push(format!( "ALTER TABLE {full_table} ADD CONSTRAINT {} FOREIGN KEY ({columns}) REFERENCES {referenced_table} ({ref_columns})", quote_identifier(name, &DatabaseType::Postgres) @@ -2761,12 +2762,14 @@ mod tests { db::ForeignKeyInfo { name: "orders_user_id_fkey".to_string(), column: "user_id".to_string(), + ref_schema: None, ref_table: "users".to_string(), ref_column: "id".to_string(), }, db::ForeignKeyInfo { name: "orders_user_id_fkey".to_string(), column: "tenant_id".to_string(), + ref_schema: None, ref_table: "users".to_string(), ref_column: "tenant_id".to_string(), }, diff --git a/crates/dbx-core/src/types.rs b/crates/dbx-core/src/types.rs index 63bdb8df03..b6accfbf79 100644 --- a/crates/dbx-core/src/types.rs +++ b/crates/dbx-core/src/types.rs @@ -88,6 +88,8 @@ pub struct IndexInfo { pub struct ForeignKeyInfo { pub name: String, pub column: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ref_schema: Option, pub ref_table: String, pub ref_column: String, } diff --git a/packages/app-tests/sqlCompletion.test.ts b/packages/app-tests/sqlCompletion.test.ts index 52a656d5b3..19382cd303 100644 --- a/packages/app-tests/sqlCompletion.test.ts +++ b/packages/app-tests/sqlCompletion.test.ts @@ -9,6 +9,7 @@ import { getSqlCompletionContext, recordCompletionSelection, type SqlCompletionColumn, + type SqlCompletionForeignKey, type SqlCompletionTable, } from "../../apps/desktop/src/lib/sqlCompletion.ts"; @@ -957,6 +958,194 @@ test("getSqlCompletionContext returns nonAggregatedSelectColumns", () => { // --- Better FK join inference --- +test("prefers explicit foreign-key join condition with table aliases", () => { + const foreignKeysByTable = new Map([ + [ + "public.orders", + [{ name: "orders_customer_id_fkey", column: "customer_id", ref_table: "customers", ref_column: "id" }], + ], + ]); + const sql = "select * from public.orders o join public.customers c on "; + const items = buildSqlCompletionItems(sql, sql.length, { + tables: [ + { name: "orders", schema: "public", type: "table" }, + { name: "customers", schema: "public", type: "table" }, + ], + columnsByTable, + foreignKeysByTable, + }); + + assert.deepEqual(items[0], { + label: "o.customer_id = c.id", + type: "snippet", + detail: "JOIN condition from foreign key", + apply: "o.customer_id = c.id", + boost: 3201, + }); +}); + +test("suggests explicit foreign-key join when the joined table owns the key", () => { + const foreignKeysByTable = new Map([ + [ + "public.orders", + [{ name: "orders_customer_id_fkey", column: "customer_id", ref_table: "customers", ref_column: "id" }], + ], + ]); + const sql = "select * from public.customers c join public.orders o on "; + const items = buildSqlCompletionItems(sql, sql.length, { + tables: [ + { name: "customers", schema: "public", type: "table" }, + { name: "orders", schema: "public", type: "table" }, + ], + columnsByTable, + foreignKeysByTable, + }); + + const fkJoin = items.find((item) => item.label === "o.customer_id = c.id"); + assert.ok(fkJoin, "should suggest FK join when the right side owns the foreign key"); +}); + +test("suggests composite explicit foreign-key join conditions", () => { + const colsWithCompositeFk = new Map([ + [ + "public.products", + [ + { name: "tenant_id", table: "products", schema: "public", dataType: "text" }, + { name: "id", table: "products", schema: "public", dataType: "text" }, + ], + ], + [ + "public.order_lines", + [ + { name: "tenant_id", table: "order_lines", schema: "public", dataType: "text" }, + { name: "product_id", table: "order_lines", schema: "public", dataType: "text" }, + ], + ], + ]); + const foreignKeysByTable = new Map([ + [ + "public.order_lines", + [ + { name: "order_lines_product_fkey", column: "tenant_id", ref_table: "products", ref_column: "tenant_id" }, + { name: "order_lines_product_fkey", column: "product_id", ref_table: "products", ref_column: "id" }, + ], + ], + ]); + const sql = "select * from public.order_lines ol join public.products p on "; + const items = buildSqlCompletionItems(sql, sql.length, { + tables: [ + { name: "order_lines", schema: "public", type: "table" }, + { name: "products", schema: "public", type: "table" }, + ], + columnsByTable: colsWithCompositeFk, + foreignKeysByTable, + }); + + const fkJoin = items.find((item) => item.label === "ol.tenant_id = p.tenant_id AND ol.product_id = p.id"); + assert.ok(fkJoin, "should suggest full composite FK join"); + assert.equal(fkJoin?.detail, "JOIN condition from composite foreign key"); +}); + +test("uses referenced schema to disambiguate explicit foreign-key joins", () => { + const foreignKeysByTable = new Map([ + [ + "sales.orders", + [ + { + name: "orders_customer_id_fkey", + column: "customer_id", + ref_schema: "crm", + ref_table: "customers", + ref_column: "id", + }, + ], + ], + ]); + const sql = "select * from sales.orders o join public.customers pc on "; + const items = buildSqlCompletionItems(sql, sql.length, { + tables: [ + { name: "orders", schema: "sales", type: "table" }, + { name: "customers", schema: "public", type: "table" }, + { name: "customers", schema: "crm", type: "table" }, + ], + columnsByTable, + foreignKeysByTable, + }); + + assert.equal( + items.some((item) => item.label === "o.customer_id = pc.id" && item.detail === "JOIN condition from foreign key"), + false, + ); +}); + +test("suggests likely composite joins from shared scope columns and id naming", () => { + const colsWithTenantRelationship = new Map([ + [ + "public.orders", + [ + { name: "tenant_id", table: "orders", schema: "public", dataType: "text" }, + { name: "id", table: "orders", schema: "public", dataType: "uuid" }, + ], + ], + [ + "public.order_lines", + [ + { name: "tenant_id", table: "order_lines", schema: "public", dataType: "text" }, + { name: "order_id", table: "order_lines", schema: "public", dataType: "uuid" }, + { name: "article_id", table: "order_lines", schema: "public", dataType: "text" }, + ], + ], + ]); + const sql = "select * from public.orders o join public.order_lines ol on "; + const items = buildSqlCompletionItems(sql, sql.length, { + tables: [ + { name: "orders", schema: "public", type: "table" }, + { name: "order_lines", schema: "public", type: "table" }, + ], + columnsByTable: colsWithTenantRelationship, + }); + + const compositeJoin = items.find((item) => item.label === "o.tenant_id = ol.tenant_id AND o.id = ol.order_id"); + assert.ok(compositeJoin, "should combine tenant scope with id/FK naming"); + assert.equal(compositeJoin?.detail, "Likely composite JOIN condition"); +}); + +test("does not suggest heuristic joins for incompatible column types", () => { + const colsWithIncompatibleIds = new Map([ + [ + "public.users", + [ + { name: "id", table: "users", schema: "public", dataType: "uuid" }, + { name: "tenant_id", table: "users", schema: "public", dataType: "text" }, + ], + ], + [ + "public.orders", + [ + { name: "user_id", table: "orders", schema: "public", dataType: "bigint" }, + { name: "tenant_id", table: "orders", schema: "public", dataType: "text" }, + ], + ], + ]); + const sql = "select * from public.users u join public.orders o on "; + const items = buildSqlCompletionItems(sql, sql.length, { + tables: [ + { name: "users", schema: "public", type: "table" }, + { name: "orders", schema: "public", type: "table" }, + ], + columnsByTable: colsWithIncompatibleIds, + }); + + assert.equal( + items.some((item) => item.label === "u.id = o.user_id"), + false, + ); + assert.equal( + items.some((item) => item.label === "u.tenant_id = o.tenant_id AND u.id = o.user_id"), + false, + ); +}); + test("suggests join condition for same FK column in both tables", () => { const colsWithFk = new Map([ [