From f5cdbb8662e22c8fb4c6a11e9b6b8a43ceb5b787 Mon Sep 17 00:00:00 2001 From: zipg Date: Mon, 20 Jul 2026 01:55:52 +0800 Subject: [PATCH] feat(completion): improve SQL function suggestions --- .../src/components/editor/QueryEditor.vue | 61 ++++++++++-- .../__tests__/sql/semantic/completion.spec.ts | 31 ++++++ .../sql/sqlCompletion.context.spec.ts | 2 +- .../sql/sqlCompletionLookupTarget.spec.ts | 22 ++++- .../src/lib/sql/semantic/completion.ts | 4 +- apps/desktop/src/lib/sql/sqlCompletion.ts | 40 +++++--- .../src/lib/sql/sqlCompletionLookupTarget.ts | 17 ++++ .../connectionStore.completion.spec.ts | 97 ++++++++++++++++++- apps/desktop/src/stores/connectionStore.ts | 59 ++++++++--- crates/dbx-core/src/db/postgres.rs | 6 +- crates/dbx-core/src/db/sqlserver.rs | 32 +++++- packages/app-tests/sqlCompletion.test.ts | 81 ++++++++++++++++ 12 files changed, 412 insertions(+), 40 deletions(-) diff --git a/apps/desktop/src/components/editor/QueryEditor.vue b/apps/desktop/src/components/editor/QueryEditor.vue index 1ae4fdea7..5739d619e 100644 --- a/apps/desktop/src/components/editor/QueryEditor.vue +++ b/apps/desktop/src/components/editor/QueryEditor.vue @@ -41,7 +41,7 @@ import { buildSqlSemanticModel } from "@/lib/sql/semantic/model"; import { mergeSqlSemanticReferenceAnalysis, resolveSqlSemanticNavigationTarget } from "@/lib/sql/semantic/references"; import { buildElasticsearchCompletionItemsFromContext, getElasticsearchCompletionContext, getElasticsearchCompletionResultValidFor, shouldAutoOpenElasticsearchCompletion, type ElasticsearchCompletionItem } from "@/lib/elasticsearch/elasticsearchCompletion"; import { buildMongoCompletionItemsFromContext, getMongoCompletionContext, getMongoCompletionResultValidFor, mongoCompletionNeedsCollections, mongoCompletionNeedsFields, shouldAutoOpenMongoCompletion, type MongoCompletionItem } from "@/lib/mongo/mongoCompletion"; -import { resolveSqlCompletionTableLookupTarget } from "@/lib/sql/sqlCompletionLookupTarget"; +import { resolveSqlCompletionRoutineLookupTarget, resolveSqlCompletionTableLookupTarget } from "@/lib/sql/sqlCompletionLookupTarget"; import { extractIdentifierDetailsAt, isSqlKeyword, matchTable, mergeSqlObjectNavigationType, splitQualifiedIdentifier, sqlObjectHoverDetail, sqlObjectNavigationTarget, type SqlObjectNavigationTarget } from "@/lib/sql/sqlNavigation"; import { lineColumnToOffset, parseSqlErrorLocation } from "@/lib/sql/sqlDiagnostics"; import { @@ -74,7 +74,7 @@ import { areSqlSemanticDiagnosticsEqual, buildSqlParserErrorDiagnostic, buildSql import { buildRedisSyntaxDiagnostics, shouldRunRedisDiagnostics } from "@/lib/redis/redisSyntaxDiagnostics"; import { buildRedisCompletionItemsFromContext, getRedisCompletionContext, getRedisCompletionResultValidFor, shouldAutoOpenRedisCompletion, takesKeyArgument, type RedisCompletionItem } from "@/lib/redis/redisCompletion"; import type { SqlCompletionColumn, SqlCompletionForeignKey, SqlCompletionItem, SqlCompletionObject, SqlCompletionReferencedTable, SqlCompletionTable } from "@/lib/sql/sqlCompletion"; -import type { DatabaseType, SqlReferenceAnalysis, SqlTableReference, SqlTextSpan } from "@/types/database"; +import type { CompletionAssistantObjectKind, DatabaseType, SqlReferenceAnalysis, SqlTableReference, SqlTextSpan } from "@/types/database"; const props = defineProps<{ modelValue: string; @@ -2278,12 +2278,12 @@ async function provideSqlCompletions(context: CompletionContext) { const shouldResolveAsyncCompletion = tableNameCompletion || shouldResolveColumnCompletion; const localResult = buildLocalSqlCompletionResult(completionContext, fullDoc, position); if (localResult) { - scheduleCompletionMetadataRefresh(completionContext); + scheduleCompletionMetadataRefresh(completionContext, fullDoc, position); const hasLocalColumnResult = localResult.options.some((option) => option.type === "column"); if ((!explicit || typedActivation) && (!shouldResolveColumnCompletion || hasLocalColumnResult)) return localResult; } if ((!explicit || typedActivation) && !shouldResolveAsyncCompletion) { - scheduleCompletionMetadataRefresh(completionContext); + scheduleCompletionMetadataRefresh(completionContext, fullDoc, position); return null; } @@ -2465,7 +2465,7 @@ function buildLocalSqlCompletionResult(completionContext: ReturnType) { +function scheduleCompletionMetadataRefresh(completionContext: ReturnType, fullDoc: string, position: number) { if (!props.connectionId || props.database == null) return; const localOnlyMetadata = usesLocalOnlyCompletionMetadata(); const onDemandOnlyColumns = usesOnDemandOnlyCompletionColumns(); @@ -2494,7 +2494,10 @@ function scheduleCompletionMetadataRefresh(completionContext: ReturnType { - cachedCompletionObjects = mergeCompletionObjects(cachedCompletionObjects, objects); + const merged = mergeCompletionObjects(cachedCompletionObjects, objects); + const changed = completionObjectsDiffer(cachedCompletionObjects, merged); + cachedCompletionObjects = merged; + if (changed) refreshActiveSqlCompletion(fullDoc, position, completionContext); }) .catch(() => {}); } @@ -2609,18 +2612,29 @@ function lookupLocalCompletionObjectsForContext(completionContext: ReturnType): Promise { if (!props.connectionId || props.database == null) return []; + const objectKinds = completionObjectKindsForContext(completionContext); if (props.databaseType !== "oracle") { - return connectionStore.listCompletionObjects(props.connectionId, props.database, completionContext.qualifier || completionContext.prefix, MAX_COMPLETION_TABLES, props.schema); + const target = resolveSqlCompletionRoutineLookupTarget({ currentSchema: props.schema, completionContext }); + return connectionStore.listCompletionObjects(props.connectionId, props.database, target.mask, MAX_COMPLETION_TABLES, target.schema, undefined, false, props.schema, objectKinds); } - const groups = await Promise.all(oracleRoutineCompletionTargets(completionContext).map((target) => connectionStore.listCompletionObjects(props.connectionId!, props.database!, completionContext.prefix, MAX_COMPLETION_TABLES, target.schema, target.parentName, target.globalSearch, props.schema))); + const groups = await Promise.all( + oracleRoutineCompletionTargets(completionContext).map((target) => connectionStore.listCompletionObjects(props.connectionId!, props.database!, completionContext.prefix, MAX_COMPLETION_TABLES, target.schema, target.parentName, target.globalSearch, props.schema, objectKinds)), + ); return groups.reduce((objects, group) => mergeCompletionObjects(objects, group), [] as SqlCompletionObject[]); } +function completionObjectKindsForContext(completionContext: ReturnType): CompletionAssistantObjectKind[] { + if (completionContext.contextKind === "exec") return ["procedure"]; + if (completionContext.suggestColumns && completionContext.referencedTables.length > 0 && !completionContext.qualifier) return ["function"]; + return ["routine"]; +} + async function performAsyncCompletionWithResult(epoch: number, completionContext: ReturnType, fullDoc: string, position: number) { const localOnlyMetadata = usesLocalOnlyCompletionMetadata(); const onDemandOnlyColumns = usesOnDemandOnlyCompletionColumns(); @@ -2876,6 +2890,35 @@ function mergeCompletionObjects(existing: SqlCompletionObject[], incoming: SqlCo return merged; } +function completionObjectsDiffer(existing: SqlCompletionObject[], incoming: SqlCompletionObject[]): boolean { + if (existing.length !== incoming.length) return true; + return existing.some((object, index) => { + const other = incoming[index]; + return ( + !other || + object.name !== other.name || + object.schema !== other.schema || + object.type !== other.type || + object.parentSchema !== other.parentSchema || + object.parentName !== other.parentName || + object.dataType !== other.dataType || + object.signature !== other.signature || + object.comment !== other.comment || + object.applyName !== other.applyName || + object.boost !== other.boost + ); + }); +} + +function refreshActiveSqlCompletion(fullDoc: string, position: number, completionContext: ReturnType) { + const currentView = view.value; + if (!currentView || codeMirrorCompletionStatus?.(currentView.state) !== "active") return; + if (currentView.state.doc.toString() !== fullDoc || currentView.state.selection.main.head !== position) return; + const currentContext = getSqlCompletionContext(fullDoc, position); + if (currentContext.prefix !== completionContext.prefix || currentContext.contextKind !== completionContext.contextKind) return; + scheduleSqlCompletionStart(currentView); +} + function refreshCompletionCache() { cachedTables = []; cachedCompletionObjects = []; 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 0cec221b2..377f76930 100644 --- a/apps/desktop/src/lib/__tests__/sql/semantic/completion.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/semantic/completion.spec.ts @@ -35,6 +35,37 @@ function semanticCompletion(markedSql: string, input: Partial { + it("keeps matching functions available in column expressions", () => { + const columnsByTable = new Map([ + [ + "routes", + [ + { name: "start_sid", table: "routes" }, + { name: "start_dept", table: "routes" }, + ], + ], + ]); + + const { context, items } = semanticCompletion( + "SELECT * FROM routes WHERE st_|", + { + columnsByTable, + objects: [ + { name: "st_area", type: "function", dataType: "double precision" }, + { name: "st_refresh", type: "procedure" }, + ], + }, + { databaseType: "postgres", dialect: "postgres" }, + ); + + expect(context.contextKind).toBe("column"); + expect(context.suggestColumns).toBe(true); + expect(context.suggestRoutines).toBe(true); + expect(context.exclusiveRoutineSuggestions).toBe(false); + expect(items.some((item) => item.label === "st_area" && item.type === "function")).toBe(true); + expect(items.some((item) => item.label === "st_refresh")).toBe(false); + }); + it("keeps alias-qualified column completion scoped to one row source", () => { const columnsByTable = new Map([ ["users", ["id", "name", "email"].map((name) => ({ name, table: "users" }))], diff --git a/apps/desktop/src/lib/__tests__/sql/sqlCompletion.context.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlCompletion.context.spec.ts index b460b6473..28d004c40 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlCompletion.context.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/sqlCompletion.context.spec.ts @@ -285,7 +285,7 @@ describe("sqlCompletion scoped context classification", () => { expect(context.prefix).toBe("userc"); expect(context.referencedTables).toEqual(expect.arrayContaining([expect.objectContaining({ name: "A1User" })])); expect(context.suggestColumns).toBe(true); - expect(context.suggestRoutines).toBe(false); + expect(context.suggestRoutines).toBe(true); }); it("auto-opens column completion after WHERE whitespace before LIMIT", () => { diff --git a/apps/desktop/src/lib/__tests__/sql/sqlCompletionLookupTarget.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlCompletionLookupTarget.spec.ts index eac09297c..5e12df62e 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlCompletionLookupTarget.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/sqlCompletionLookupTarget.spec.ts @@ -1,5 +1,6 @@ import { describe, expect, it } from "vitest"; -import { resolveSqlCompletionTableLookupTarget } from "@/lib/sql/sqlCompletionLookupTarget"; +import { getSqlCompletionContext } from "@/lib/sql/sqlCompletion"; +import { resolveSqlCompletionRoutineLookupTarget, resolveSqlCompletionTableLookupTarget } from "@/lib/sql/sqlCompletionLookupTarget"; describe("sqlCompletionLookupTarget", () => { it("treats qualified table completion as a database lookup for MySQL-compatible engines", () => { @@ -75,4 +76,23 @@ describe("sqlCompletionLookupTarget", () => { filter: "ord", }); }); + + it.each([ + ["SELECT dbo.fn_", "dbo", "fn_"], + ["SELECT public.st_", "public", "st_"], + ])("separates a routine schema from its name mask for %s", (sql, schema, mask) => { + const completionContext = getSqlCompletionContext(sql, sql.length); + + expect(resolveSqlCompletionRoutineLookupTarget({ currentSchema: "fallback", completionContext })).toEqual({ schema, mask }); + }); + + it("uses the current schema for an unqualified routine mask", () => { + const sql = "SELECT st_"; + const completionContext = getSqlCompletionContext(sql, sql.length); + + expect(resolveSqlCompletionRoutineLookupTarget({ currentSchema: "public", completionContext })).toEqual({ + schema: "public", + mask: "st_", + }); + }); }); diff --git a/apps/desktop/src/lib/sql/semantic/completion.ts b/apps/desktop/src/lib/sql/semantic/completion.ts index 2d633bbb4..6d86e9d5d 100644 --- a/apps/desktop/src/lib/sql/semantic/completion.ts +++ b/apps/desktop/src/lib/sql/semantic/completion.ts @@ -183,7 +183,7 @@ export function sqlCompletionContextFromSemantic(model: SqlSemanticModel, base: const mutationSchema = mutationTarget?.qualifierParts[mutationTarget.qualifierParts.length - 1]; const suggestTables = scope.kind === "table" || scope.kind === "schema" || scope.kind === "catalog"; const suggestColumns = scope.kind === "columns"; - const suggestRoutines = scope.kind === "routine"; + const suggestRoutines = scope.kind === "routine" || (suggestColumns && base.suggestRoutines && !base.exclusiveColumnSuggestions); const projectionAliases = sqlSemanticProjectionAliasColumns(model).map((column) => column.name); return { @@ -198,7 +198,7 @@ export function sqlCompletionContextFromSemantic(model: SqlSemanticModel, base: suggestJoinConditions: model.cursorIntent.kind === "join_condition", exclusiveTableSuggestions: suggestTables, exclusiveColumnSuggestions: model.cursorIntent.kind === "alias_column" || model.cursorIntent.kind === "insert_column" || model.cursorIntent.kind === "update_column", - exclusiveRoutineSuggestions: suggestRoutines, + exclusiveRoutineSuggestions: scope.kind === "routine", prioritizeSelectAliases: base.prioritizeSelectAliases || projectionAliases.length > 0, selectAliases: projectionAliases.length > 0 ? projectionAliases : base.selectAliases, referencedTables: referencedTables.length > 0 ? referencedTables : base.referencedTables, diff --git a/apps/desktop/src/lib/sql/sqlCompletion.ts b/apps/desktop/src/lib/sql/sqlCompletion.ts index f74c0fc0f..8e3452d6a 100644 --- a/apps/desktop/src/lib/sql/sqlCompletion.ts +++ b/apps/desktop/src/lib/sql/sqlCompletion.ts @@ -1171,6 +1171,9 @@ export interface SqlCompletionObject { type: "procedure" | "function" | "trigger" | "package"; parentSchema?: string; parentName?: string; + dataType?: string; + signature?: string; + comment?: string | null; applyName?: string; boost?: number; } @@ -1338,10 +1341,15 @@ class SqlCompletionProvider { } const preferReferencedColumns = hasMatchingReferencedColumnPrefix(context, this.input.columnsByTable); - if (!pendingJoinKeyword && !preferReferencedColumns && !context.exclusiveTableSuggestions && !context.exclusiveColumnSuggestions && !context.exclusiveRoutineSuggestions) { + if (!pendingJoinKeyword && !context.exclusiveTableSuggestions && !context.exclusiveColumnSuggestions && !context.exclusiveRoutineSuggestions) { const snippets = this.databaseType === "manticoresearch" ? [...(this.input.snippets ?? DEFAULT_SQL_SNIPPETS), ...MANTICORESEARCH_SQL_SNIPPETS] : (this.input.snippets ?? DEFAULT_SQL_SNIPPETS); - this.items.push(...buildSnippetItems(context.prefix, snippets, this.input.keywordCase)); - this.items.push(...buildFunctionSnippetItems(context.prefix, getFunctionDescriptions(this.t), this.databaseType)); + if (!preferReferencedColumns) { + this.items.push(...buildSnippetItems(context.prefix, snippets, this.input.keywordCase)); + } + if (!preferReferencedColumns || context.suggestRoutines) { + const functionItems = buildFunctionSnippetItems(context.prefix, getFunctionDescriptions(this.t), this.databaseType); + this.items.push(...(preferReferencedColumns ? functionItems.filter((item) => item.label.toLowerCase().startsWith(context.prefix.toLowerCase())) : functionItems)); + } } if (this.databaseType === "manticoresearch" && context.exclusiveRoutineSuggestions) { @@ -1818,10 +1826,9 @@ export function getSqlCompletionContext(sql: string, cursor: number): SqlComplet const inJoinConditionContext = isInJoinConditionContext(beforeCursor); const prioritizeSelectAliases = isInOrderOrGroupByContext(beforeCursor); const inCallRoutineContext = isCallRoutineContext(beforeCursor); - const inPotentialPackageMemberContext = !!qualifier && !exclusiveTableSuggestions && !insertInfo && !oracleTableFunctionContext; + const inPotentialPackageMemberContext = !!qualifier && !exclusiveTableSuggestions && !insertInfo && !updateInfo?.inSetClause && !oracleTableFunctionContext; const suggestColumns = !!qualifier || !!updateInfo?.inSetClause || !!insertInfo || (inColumnContext && referencedTables.length > 0); - const preferColumnsOverGlobalRoutines = suggestColumns && referencedTables.length > 0 && !qualifier; - const suggestRoutines = inCallRoutineContext || oracleTableFunctionContext || inPotentialPackageMemberContext || (!preferColumnsOverGlobalRoutines && !exclusiveTableSuggestions && !exclusiveColumnSuggestions && !insertInfo && prefix.length >= 2); + const suggestRoutines = inCallRoutineContext || oracleTableFunctionContext || inPotentialPackageMemberContext || (!exclusiveTableSuggestions && !exclusiveColumnSuggestions && !insertInfo && !updateInfo?.inSetClause && prefix.length >= 2); const statementKind = detectStatementKind(beforeCursor || fullStatement); const preferredKeywords = qualifier ? [] : preferredKeywordsForCompletion(beforeCursor, beforeToken, selectListColumnContext, exclusiveTableSuggestions, updateInfo, deleteInfo); @@ -1903,12 +1910,12 @@ function detectCompletionContextKind(options: { if (options.updateInfo?.inSetClause) return "column"; if (options.inCallRoutineContext) return "exec"; if (options.qualifier && options.exclusiveColumnSuggestions) return "alias_column"; + if (options.suggestColumns) return options.qualifier ? "alias_column" : "column"; if (options.oracleTableFunctionContext || options.suggestRoutines) return "routine"; if (options.exclusiveTableSuggestions || options.afterTableTrigger) { if (options.statementKind === "insert" && options.lastWord === "into") return "insert_target"; return options.lastWord === "join" ? "join" : "table"; } - if (options.suggestColumns) return options.qualifier ? "alias_column" : "column"; return "keyword"; } @@ -2190,7 +2197,7 @@ function detectInsertColumnListContext(beforeCursor: string): { table: string; s function detectUpdateCompletionContext(beforeCursor: string): { target: { table: string; schema?: string }; afterTarget: boolean; inSetClause: boolean; afterSetAssignments: boolean } | null { const cleaned = beforeCursor.replace(/'[^']*'/g, "''").replace(/"[^"]*"/g, '""'); - const match = /^\s*update\s+((?:"[^"]+"|`[^`]+`|[A-Za-z_][\w$]*)(?:\.(?:"[^"]+"|`[^`]+`|[A-Za-z_][\w$]*))?)(?:\s+(?:as\s+)?([A-Za-z_][\w$]*))?/i.exec(cleaned); + const match = /^\s*update\s+((?:"[^"]+"|`[^`]+`|[A-Za-z_][\w$]*)(?:\.(?:"[^"]+"|`[^`]+`|[A-Za-z_][\w$]*))?)(?:\s+(?:as\s+)?((?!set\b|where\b)[A-Za-z_][\w$]*))?/i.exec(cleaned); if (!match) return null; const [first, second] = splitQualifiedName(match[1] ?? ""); if (!first) return null; @@ -2975,9 +2982,10 @@ function buildSchemaItems(prefix: string, schemas: string[], dialect?: "mysql" | function buildObjectItems(context: SqlCompletionContext, objects: SqlCompletionObject[], dialect?: "mysql" | "postgres" | "sqlserver", databaseType?: DatabaseType, currentSchema?: string): SqlCompletionItem[] { if (completionQualifierIsReferencedTable(context)) return []; const onlyProcedures = context.contextKind === "exec"; + const onlyFunctions = context.suggestColumns && context.referencedTables.length > 0 && !context.qualifier; const prioritizeOracleFunctions = databaseType === "oracle" && context.statementKind === "select"; return objects - .filter((object) => (!onlyProcedures || object.type === "procedure") && objectMatchesCompletionContext(object, context)) + .filter((object) => (!onlyProcedures || object.type === "procedure") && (!onlyFunctions || (object.type === "function" && object.name.toLowerCase().startsWith(context.prefix.toLowerCase()))) && objectMatchesCompletionContext(object, context)) .map((object) => { const qualifiedByContext = objectIsQualifiedByContext(object, context); const objectInCurrentSchema = !!currentSchema && !!object.schema && normalizeIdentifierPart(object.schema) === normalizeIdentifierPart(currentSchema); @@ -2985,13 +2993,17 @@ function buildObjectItems(context: SqlCompletionContext, objects: SqlCompletionO qualifiedByContext || (context.qualifier && object.schema?.toLowerCase() === context.qualifier.toLowerCase()) ? quoteSqlIdentifier(object.name, dialect) : (object.applyName ?? (object.schema && !objectInCurrentSchema ? `${quoteSqlIdentifier(object.schema, dialect)}.${quoteSqlIdentifier(object.name, dialect)}` : quoteSqlIdentifier(object.name, dialect))); - const detail = object.type === "trigger" && object.parentName ? `trigger on ${object.parentName}` : object.parentName ? `${object.type} in ${object.parentName}` : object.schema ? `${object.type} in ${object.schema}` : object.type; + const locationDetail = object.type === "trigger" && object.parentName ? `trigger on ${object.parentName}` : object.parentName ? `${object.type} in ${object.parentName}` : object.schema ? `${object.type} in ${object.schema}` : object.type; + const detail = object.dataType ? `${locationDetail} [${object.dataType}]` : locationDetail; + const schemaBoost = onlyFunctions ? Math.min(object.boost ?? 0, 1000) : (object.boost ?? 0); + const typeBoost = routineTypeBoost(object.type, prioritizeOracleFunctions && !onlyFunctions); return { label: object.name, type: "function" as const, detail, + info: buildRoutineInfo(object), apply: object.type === "trigger" || object.type === "package" ? applyName : `${applyName}()`, - boost: computeBoost(object.name, context.prefix) + routineTypeBoost(object.type, prioritizeOracleFunctions) + (object.boost ?? 0), + boost: computeBoost(object.name, context.prefix) + typeBoost + schemaBoost, dedupeKey: object.applyName || (databaseType === "oracle" && object.schema) ? applyName : undefined, // Preserve exact routine matches before the capped candidate list is truncated. exactMatch: !!context.prefix && object.name.toLowerCase() === context.prefix.toLowerCase(), @@ -3001,6 +3013,12 @@ function buildObjectItems(context: SqlCompletionContext, objects: SqlCompletionO .slice(0, MAX_TABLE_COMPLETION_ITEMS); } +function buildRoutineInfo(object: SqlCompletionObject): string | undefined { + const qualifiedName = object.parentName ? [object.parentSchema ?? object.schema, object.parentName, object.name].filter(Boolean).join(".") : [object.schema, object.name].filter(Boolean).join("."); + const parts = [qualifiedName || object.name, object.signature?.trim(), object.comment?.trim()].filter((part): part is string => !!part); + return parts.length > 1 ? parts.join("\n") : undefined; +} + function routineTypeBoost(type: SqlCompletionObject["type"], prioritizeFunctions: boolean): number { if (type === "package") return 1600; if (type === "function") return prioritizeFunctions ? 1800 : 900; diff --git a/apps/desktop/src/lib/sql/sqlCompletionLookupTarget.ts b/apps/desktop/src/lib/sql/sqlCompletionLookupTarget.ts index 6d7cd4da5..70ec178be 100644 --- a/apps/desktop/src/lib/sql/sqlCompletionLookupTarget.ts +++ b/apps/desktop/src/lib/sql/sqlCompletionLookupTarget.ts @@ -7,6 +7,11 @@ export interface SqlCompletionTableLookupTarget { qualifierDatabase?: string; } +export interface SqlCompletionRoutineLookupTarget { + schema?: string; + mask: string; +} + function findExactName(names: readonly string[] | undefined, value: string): string | undefined { return names?.find((name) => name.toLowerCase() === value.toLowerCase()); } @@ -40,3 +45,15 @@ export function resolveSqlCompletionTableLookupTarget(options: { filter: qualifier && completionContext.suggestTables ? completionContext.prefix : qualifier || completionContext.prefix, }; } + +export function resolveSqlCompletionRoutineLookupTarget(options: { currentSchema?: string; completionContext: Pick }): SqlCompletionRoutineLookupTarget { + const qualifierParts = options.completionContext.qualifierParts?.filter(Boolean); + const schema = qualifierParts?.[qualifierParts.length - 1] ?? options.completionContext.qualifier?.trim() ?? options.currentSchema; + + // A qualified routine uses the qualifier as metadata scope; only the final + // identifier fragment is the function/procedure name mask. + return { + schema: schema || undefined, + mask: options.completionContext.prefix, + }; +} diff --git a/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts index b18ec7bf5..e72e5ba40 100644 --- a/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts +++ b/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts @@ -37,6 +37,18 @@ function oracleConnection(): ConnectionConfig { } as ConnectionConfig; } +function sqlServerConnection(): ConnectionConfig { + return { + ...postgresConnection(), + id: "sqlserver-1", + name: "SQL Server", + db_type: "sqlserver", + port: 1433, + username: "sa", + database: "app", + } as ConnectionConfig; +} + function deferred() { let resolve!: (value: T) => void; const promise = new Promise((res) => { @@ -242,7 +254,90 @@ describe("connectionStore completion assistant", () => { const objects = await store.listCompletionObjects("oracle-1", "ORCL", "CALC", 20, "HR", "PAYROLL", false, "APP"); expect(completionAssistantSearch).toHaveBeenCalledWith(expect.objectContaining({ object_kinds: ["routine"], mask: "CALC", schema: "APP", parent_schema: "HR", parent_name: "PAYROLL", global_search: false })); - expect(objects).toEqual([expect.objectContaining({ name: "CALCULATE_BONUS", schema: "HR", type: "function", parentSchema: "HR", parentName: "PAYROLL", applyName: "HR.CALCULATE_BONUS", boost: 0 })]); + expect(objects).toEqual([expect.objectContaining({ name: "CALCULATE_BONUS", schema: "HR", type: "function", parentSchema: "HR", parentName: "PAYROLL", dataType: undefined, applyName: "HR.CALCULATE_BONUS", boost: 0 })]); + }); + + it("loads PostgreSQL routines by prefix and preserves return metadata", async () => { + const completionAssistantSearch = vi.fn().mockResolvedValue({ + candidates: [{ name: "st_area", kind: "function", schema: "public", data_type: "double precision", comment: "Returns an area" }], + incomplete: false, + fallback_used: false, + }); + + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => ({ + checkConnectionHealth: vi.fn().mockResolvedValue(undefined), + completionAssistantSearch, + listCompletionObjects: vi.fn().mockResolvedValue([]), + })); + + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [postgresConnection()]; + store.connectedIds.add("pg-1"); + + const objects = await store.listCompletionObjects("pg-1", "app", "st_", 20, "public", undefined, false, "public", ["function"]); + + expect(completionAssistantSearch).toHaveBeenCalledWith( + expect.objectContaining({ + object_kinds: ["function"], + mask: "st_", + schema: "public", + parent_schema: "public", + match_mode: "prefix", + }), + ); + expect(objects).toEqual([ + expect.objectContaining({ + name: "st_area", + schema: "public", + type: "function", + dataType: "double precision", + comment: "Returns an area", + applyName: "st_area", + boost: 1000, + }), + ]); + }); + + it("searches default SQL Server schemas without treating the username as a schema", async () => { + const completionAssistantSearch = vi.fn().mockResolvedValue({ + candidates: [{ name: "st_area", kind: "function", schema: "dbo", data_type: "float" }], + incomplete: false, + fallback_used: false, + }); + + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => ({ + checkConnectionHealth: vi.fn().mockResolvedValue(undefined), + completionAssistantSearch, + listCompletionObjects: vi.fn().mockResolvedValue([]), + })); + + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [sqlServerConnection()]; + store.connectedIds.add("sqlserver-1"); + + const objects = await store.listCompletionObjects("sqlserver-1", "app", "st_", 20); + + expect(completionAssistantSearch).toHaveBeenCalledWith( + expect.objectContaining({ + schema: null, + parent_schema: null, + mask: "st_", + }), + ); + expect(objects).toEqual([ + expect.objectContaining({ + name: "st_area", + schema: "dbo", + type: "function", + dataType: "float", + applyName: "dbo.st_area", + boost: 1000, + }), + ]); }); it("limits concurrent completion column metadata requests per connection database", async () => { diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index 8fd437a02..1b6ec4e7a 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -7,6 +7,7 @@ import type { CompletionAssistantObjectKind, CompletionAssistantRequest, ConnectionConfig, + DatabaseType, DatabaseConnectionInfo, CatalogInfo, ForeignKeyInfo, @@ -4125,6 +4126,7 @@ export const useConnectionStore = defineStore("connection", () => { } const ORACLE_SYSTEM_COMPLETION_SCHEMAS = new Set(["SYS", "SYSTEM", "SYSMAN", "DBSNMP", "OUTLN", "XDB", "MDSYS", "CTXSYS", "WMSYS"]); + const FILTERED_ROUTINE_COMPLETION_DATABASES = new Set(["mysql", "postgres", "sqlserver", "oracle"]); function completionPreferredSchema(connectionId: string, preferredSchema?: string): string | undefined { return preferredSchema?.trim() || getConfig(connectionId)?.username?.trim() || undefined; @@ -4161,25 +4163,34 @@ export const useConnectionStore = defineStore("connection", () => { }); } - function completionAssistantObjects(candidates: CompletionAssistantCandidate[], preferredSchema?: string): SqlCompletionObject[] { + function completionAssistantObjects(candidates: CompletionAssistantCandidate[], preferredSchema?: string, oracleMetadata = false): SqlCompletionObject[] { return candidates .map((candidate): SqlCompletionObject | null => { const candidateType = candidate.data_type?.toUpperCase(); const type = candidate.kind === "procedure" ? "procedure" : candidate.kind === "function" ? "function" : candidate.kind === "object" && candidateType === "PACKAGE" ? "package" : null; if (!type) return null; + const dataType = candidate.data_type && !["FUNCTION", "PROCEDURE", "PACKAGE"].includes(candidateType ?? "") ? candidate.data_type : undefined; return { name: candidate.name, schema: candidate.schema ?? undefined, type, parentSchema: candidate.parent_schema ?? undefined, parentName: candidate.parent_name ?? undefined, + dataType, + comment: candidate.comment ?? null, applyName: completionCandidateApplyName(candidate.name, candidate.schema, preferredSchema), - boost: completionCandidateSchemaBoost(candidate.schema, preferredSchema), + boost: oracleMetadata ? completionCandidateSchemaBoost(candidate.schema, preferredSchema) : completionRoutineSchemaBoost(candidate.schema, preferredSchema), }; }) .filter((object): object is SqlCompletionObject => object != null); } + function completionRoutineSchemaBoost(schema: string | null | undefined, preferredSchema?: string): number { + if (schema && preferredSchema && schema.toLowerCase() === preferredSchema.toLowerCase()) return 1000; + if (schema?.toUpperCase() === "PUBLIC") return 600; + return 0; + } + function completionAssistantColumns(candidates: CompletionAssistantCandidate[], table: string, schema?: string): SqlCompletionColumn[] { return candidates .filter((candidate) => candidate.kind === "column") @@ -4212,13 +4223,26 @@ export const useConnectionStore = defineStore("connection", () => { return tables; } - async function listCompletionAssistantObjects(connectionId: string, database: string, filter: string, limit: number | undefined, schema: string | undefined, parentName: string | undefined, globalSearch: boolean, currentSchema?: string): Promise { - const preferredSchema = completionPreferredSchema(connectionId, currentSchema); + async function listCompletionAssistantObjects( + connectionId: string, + database: string, + filter: string, + limit: number | undefined, + schema: string | undefined, + parentName: string | undefined, + globalSearch: boolean, + currentSchema: string | undefined, + objectKinds: CompletionAssistantObjectKind[], + ): Promise { + const databaseType = getConfig(connectionId)?.db_type; + const oracleAssistant = databaseType === "oracle"; + const requestedSchema = currentSchema?.trim() || schema?.trim() || undefined; + const preferredSchema = oracleAssistant ? completionPreferredSchema(connectionId, currentSchema) : requestedSchema || (databaseType === "sqlserver" ? "dbo" : databaseType === "postgres" ? "public" : databaseType === "mysql" ? database : undefined); const response = await completionAssistantSearch({ connection_id: connectionId, database, - schema: preferredSchema ?? null, - object_kinds: ["routine"], + schema: oracleAssistant ? (preferredSchema ?? null) : (requestedSchema ?? null), + object_kinds: objectKinds, mask: filter.trim(), max_results: limit ?? 200, global_search: globalSearch, @@ -4226,7 +4250,10 @@ export const useConnectionStore = defineStore("connection", () => { parent_name: parentName ?? null, match_mode: "prefix", }); - const objects = completionAssistantObjects(response.candidates, preferredSchema); + const objects = completionAssistantObjects(response.candidates, preferredSchema, oracleAssistant).map((object) => ({ + ...object, + applyName: databaseType === "sqlserver" && object.schema ? `${object.schema}.${object.name}` : object.applyName, + })); indexCompletionObjects(connectionId, database, schema, objects); return objects; } @@ -4655,17 +4682,23 @@ export const useConnectionStore = defineStore("connection", () => { return deduped; } - async function listCompletionObjects(connectionId: string, database: string, filter = "", limit?: number, schema?: string, parentName?: string, globalSearch = false, currentSchema?: string): Promise { + async function listCompletionObjects(connectionId: string, database: string, filter = "", limit?: number, schema?: string, parentName?: string, globalSearch = false, currentSchema?: string, objectKinds: CompletionAssistantObjectKind[] = ["routine"]): Promise { const normalizedFilter = filter.trim().toLowerCase(); - const oracleAssistant = getConfig(connectionId)?.db_type === "oracle" && (!!normalizedFilter || typeof limit === "number" || !!parentName || globalSearch); - const cacheKey = oracleAssistant ? `${connectionId}:${database}:${schema ?? ""}:${parentName ?? ""}:${normalizedFilter}:${limit ?? ""}:${globalSearch ? "global" : "scoped"}:${currentSchema ?? ""}` : `${connectionId}:${database}:${schema ?? ""}`; + const databaseType = getConfig(connectionId)?.db_type; + const filteredRoutineAssistant = !!databaseType && FILTERED_ROUTINE_COMPLETION_DATABASES.has(databaseType) && (!!normalizedFilter || typeof limit === "number" || !!parentName || globalSearch); + const cacheKey = filteredRoutineAssistant ? `${connectionId}:${database}:${schema ?? ""}:${parentName ?? ""}:${normalizedFilter}:${limit ?? ""}:${globalSearch ? "global" : "scoped"}:${currentSchema ?? ""}:${[...objectKinds].sort().join(",")}` : `${connectionId}:${database}:${schema ?? ""}`; if (!completionObjectsCache.value[cacheKey]) { await withCompletionInFlight( `${cacheKey}:objects`, async () => { await ensureConnected(connectionId); - if (oracleAssistant) { - completionObjectsCache.value[cacheKey] = dedupeCompletionObjects(await listCompletionAssistantObjects(connectionId, database, filter, limit, schema, parentName, globalSearch, currentSchema)); + if (filteredRoutineAssistant) { + try { + completionObjectsCache.value[cacheKey] = dedupeCompletionObjects(await listCompletionAssistantObjects(connectionId, database, filter, limit, schema, parentName, globalSearch, currentSchema, objectKinds)); + } catch { + const objects = isSchemaAwareDatabase(connectionId) ? await listSchemaAwareCompletionObjects(connectionId, database, schema) : await api.listCompletionObjects(connectionId, database, schema || database); + completionObjectsCache.value[cacheKey] = dedupeCompletionObjects(objects.map(toSqlCompletionObject).filter((object): object is SqlCompletionObject => object != null)); + } } else { const objects = isSchemaAwareDatabase(connectionId) ? await listSchemaAwareCompletionObjects(connectionId, database, schema) : await api.listCompletionObjects(connectionId, database, schema || database); completionObjectsCache.value[cacheKey] = dedupeCompletionObjects(objects.map(toSqlCompletionObject).filter((object): object is SqlCompletionObject => object != null)); @@ -4712,6 +4745,8 @@ export const useConnectionStore = defineStore("connection", () => { type, parentSchema: object.parent_schema ?? undefined, parentName: object.parent_name ?? undefined, + signature: object.signature ?? undefined, + comment: object.comment ?? null, }; } diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 15202fd73..69430edfa 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -1748,7 +1748,7 @@ fn postgres_completion_tables_sql() -> &'static str { LEFT JOIN pg_catalog.pg_namespace pn ON pn.oid = pc.relnamespace \ WHERE ($1::text IS NOT NULL AND n.nspname = $1 \ OR $1::text IS NULL AND pg_catalog.pg_table_is_visible(c.oid)) \ - AND c.relkind = ANY($3) \ + AND c.relkind::text = ANY($3::text[]) \ AND ($2 = '%%' OR c.relname ILIKE $2 ESCAPE '~') \ ORDER BY c.relname LIMIT $4" } @@ -1758,7 +1758,7 @@ fn postgres_completion_routines_sql() -> &'static str { obj_description(p.oid) AS routine_comment, COALESCE(pg_get_function_result(p.oid), '') AS data_type \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ - WHERE n.nspname = $1 AND p.prokind = ANY($3) \ + WHERE n.nspname = $1 AND p.prokind::text = ANY($3::text[]) \ AND ($2 = '%%' OR p.proname ILIKE $2 ESCAPE '~') \ ORDER BY p.proname LIMIT $4" } @@ -4523,8 +4523,10 @@ mod tests { fn postgres_completion_sql_filters_before_limit() { assert!(postgres_completion_tables_sql().contains("c.relname ILIKE $2 ESCAPE '~'")); assert!(postgres_completion_tables_sql().contains("pg_catalog.pg_table_is_visible(c.oid)")); + assert!(postgres_completion_tables_sql().contains("c.relkind::text = ANY($3::text[])")); assert!(postgres_completion_tables_sql().contains("ORDER BY c.relname LIMIT $4")); assert!(postgres_completion_routines_sql().contains("p.proname ILIKE $2 ESCAPE '~'")); + assert!(postgres_completion_routines_sql().contains("p.prokind::text = ANY($3::text[])")); assert!(postgres_completion_routines_sql().contains("ORDER BY p.proname LIMIT $4")); assert!(postgres_completion_columns_sql().contains("a.attname ILIKE $3 ESCAPE '~'")); assert!(postgres_visible_table_schema_sql().contains("pg_catalog.pg_table_is_visible(c.oid)")); diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index ba6739c44..72fd98f15 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -1270,10 +1270,15 @@ fn sqlserver_completion_assistant_sql(request: &crate::types::CompletionAssistan } let object_like = sqlserver_completion_object_search_clause(request, &like_pattern); let object_visibility = sqlserver_visible_object_predicate(); + let data_type = if object_kinds.iter().any(crate::types::CompletionAssistantObjectKind::is_routine_like) { + "CASE WHEN o.type IN ('IF','TF','FT') THEN 'table' ELSE (SELECT TOP (1) TYPE_NAME(p.user_type_id) FROM sys.parameters p WHERE p.object_id = o.object_id AND p.parameter_id = 0) END" + } else { + "CAST(NULL AS NVARCHAR(128))" + }; queries.push(format!( "SELECT TOP ({limit}) o.name, s.name AS schema_name, \ CASE o.type WHEN 'U' THEN 'TABLE' WHEN 'V' THEN 'VIEW' WHEN 'P' THEN 'PROCEDURE' WHEN 'FN' THEN 'FUNCTION' WHEN 'IF' THEN 'FUNCTION' WHEN 'TF' THEN 'FUNCTION' WHEN 'FS' THEN 'FUNCTION' WHEN 'FT' THEN 'FUNCTION' ELSE o.type_desc END AS object_type, \ - CAST(NULL AS NVARCHAR(128)) AS parent_schema, CAST(NULL AS NVARCHAR(128)) AS parent_name, ep.value AS object_comment, CAST(NULL AS NVARCHAR(128)) AS data_type \ + CAST(NULL AS NVARCHAR(128)) AS parent_schema, CAST(NULL AS NVARCHAR(128)) AS parent_name, ep.value AS object_comment, {data_type} AS data_type \ FROM sys.objects o \ JOIN sys.schemas s ON s.schema_id = o.schema_id \ OUTER APPLY (SELECT CAST(ep.value AS NVARCHAR(MAX)) AS value FROM sys.extended_properties ep WHERE ep.major_id = o.object_id AND ep.minor_id = 0 AND ep.name = N'MS_Description') ep \ @@ -2578,6 +2583,31 @@ mod tests { assert!(sql.contains("CAST(NULL AS NVARCHAR(MAX)) AS object_comment")); } + #[test] + fn sqlserver_completion_assistant_returns_function_result_types() { + let request = CompletionAssistantRequest { + connection_id: "c1".to_string(), + database: "app".to_string(), + schema: Some("dbo".to_string()), + object_kinds: vec![CompletionAssistantObjectKind::Routine], + mask: "fn_".to_string(), + case_sensitive: false, + global_search: false, + max_results: Some(50), + search_in_comments: false, + search_in_definitions: false, + parent_schema: Some("dbo".to_string()), + parent_name: None, + match_mode: Some(CompletionAssistantMatchMode::Prefix), + }; + + let sql = sqlserver_completion_assistant_sql(&request, 50); + + assert!(sql.contains("WHEN o.type IN ('IF','TF','FT') THEN 'table'")); + assert!(sql.contains("p.parameter_id = 0")); + assert!(sql.contains("TYPE_NAME(p.user_type_id)")); + } + #[test] fn sqlserver_completion_assistant_searches_tempdb_for_temp_table_masks() { let request = CompletionAssistantRequest { diff --git a/packages/app-tests/sqlCompletion.test.ts b/packages/app-tests/sqlCompletion.test.ts index 2062116aa..d6a3d90c6 100644 --- a/packages/app-tests/sqlCompletion.test.ts +++ b/packages/app-tests/sqlCompletion.test.ts @@ -444,6 +444,87 @@ test("does not mix routines into an explicit table alias column completion", () assert.equal(items.some((item) => item.label === "name_formatter"), false); }); +test("suggests matching database functions alongside referenced columns", () => { + const sql = "SELECT st_ FROM public.routes"; + const items = buildSqlCompletionItems(sql, "SELECT st_".length, { + tables: [{ name: "routes", schema: "public", type: "table" }], + columnsByTable: new Map([ + [ + "public.routes", + [ + { name: "start_sid", table: "routes", schema: "public", dataType: "integer" }, + { name: "start_dept", table: "routes", schema: "public", dataType: "text" }, + ], + ], + ]), + objects: [ + { name: "st_area", schema: "public", type: "function", dataType: "double precision", comment: "Returns an area" }, + { name: "st_x", schema: "public", type: "function", dataType: "double precision" }, + { name: "st_refresh", schema: "public", type: "procedure" }, + ], + databaseType: "postgres", + currentSchema: "public", + }); + + const area = items.find((item) => item.label === "st_area" && item.type === "function"); + assert.ok(area); + assert.equal(area.detail, "function in public [double precision]"); + assert.equal(area.info, "public.st_area\nReturns an area"); + assert.ok(items.some((item) => item.label === "start_sid" && item.type === "column")); + assert.equal( + items.some((item) => item.label === "st_refresh"), + false, + ); + assert.equal( + items.some((item) => item.label === "LAST_VALUE" || item.label === "FIRST_VALUE"), + false, + ); + assert.ok(items.findIndex((item) => item.label === "st_area" && item.type === "function") < items.findIndex((item) => item.label === "start_sid" && item.type === "column")); +}); + +test("keeps an exact referenced column above an exact database function", () => { + const sql = "SELECT st_area FROM public.routes"; + const items = buildSqlCompletionItems(sql, "SELECT st_area".length, { + tables: [{ name: "routes", schema: "public", type: "table" }], + columnsByTable: new Map([["public.routes", [{ name: "st_area", table: "routes", schema: "public", dataType: "numeric" }]]]), + objects: [{ name: "st_area", schema: "public", type: "function", dataType: "double precision", boost: 2400 }], + databaseType: "oracle", + currentSchema: "public", + }); + + const exactMatches = items.filter((item) => item.label === "st_area"); + assert.deepEqual( + exactMatches.map((item) => item.type), + ["column", "function"], + ); +}); + +test("does not suggest database functions in exclusive table or column contexts", () => { + const input = { + tables: [{ name: "routes", schema: "public", type: "table" as const }], + columnsByTable: new Map([["public.routes", [{ name: "start_sid", table: "routes", schema: "public" }]]]), + objects: [{ name: "st_area", schema: "public", type: "function" as const }], + databaseType: "postgres" as const, + }; + + const tableSql = "SELECT * FROM st_"; + assert.equal( + buildSqlCompletionItems(tableSql, tableSql.length, input).some((item) => item.label === "st_area"), + false, + ); + + const updateSql = "UPDATE public.routes SET st_"; + const updateContext = getSqlCompletionContext(updateSql, updateSql.length); + assert.equal(updateContext.qualifier, undefined); + assert.equal(updateContext.contextKind, "column"); + assert.equal(updateContext.exclusiveColumnSuggestions, true); + assert.equal(updateContext.suggestRoutines, false); + assert.equal( + buildSqlCompletionItems(updateSql, updateSql.length, input).some((item) => item.label === "st_area"), + false, + ); +}); + test("keeps explicit alias column suggestions scoped to the alias table", () => { const sql = "select * from public.users u join public.orders o on u.id = o.user_id where o.st"; const items = buildSqlCompletionItems(sql, sql.length, {