diff --git a/agents/common/src/main/java/com/dbx/agent/CompletionAssistantCandidateKind.java b/agents/common/src/main/java/com/dbx/agent/CompletionAssistantCandidateKind.java index 9e2a02d368..7cf29ed1d5 100644 --- a/agents/common/src/main/java/com/dbx/agent/CompletionAssistantCandidateKind.java +++ b/agents/common/src/main/java/com/dbx/agent/CompletionAssistantCandidateKind.java @@ -17,6 +17,8 @@ public enum CompletionAssistantCandidateKind { FUNCTION, @SerializedName("column") COLUMN, + @SerializedName("sequence") + SEQUENCE, @SerializedName("object") OBJECT } diff --git a/agents/common/src/main/java/com/dbx/agent/CompletionAssistantObjectKind.java b/agents/common/src/main/java/com/dbx/agent/CompletionAssistantObjectKind.java index ec7bbd3920..ad5ed7918f 100644 --- a/agents/common/src/main/java/com/dbx/agent/CompletionAssistantObjectKind.java +++ b/agents/common/src/main/java/com/dbx/agent/CompletionAssistantObjectKind.java @@ -18,5 +18,7 @@ public enum CompletionAssistantObjectKind { @SerializedName("function") FUNCTION, @SerializedName("column") - COLUMN + COLUMN, + @SerializedName("sequence") + SEQUENCE } diff --git a/agents/drivers/duckdb/src/wire.rs b/agents/drivers/duckdb/src/wire.rs index 995f724784..d276132a43 100644 --- a/agents/drivers/duckdb/src/wire.rs +++ b/agents/drivers/duckdb/src/wire.rs @@ -60,6 +60,7 @@ pub enum CompletionAssistantObjectKind { Procedure, Function, Column, + Sequence, } impl CompletionAssistantObjectKind { @@ -78,6 +79,7 @@ pub enum CompletionAssistantCandidateKind { Procedure, Function, Column, + Sequence, Object, } diff --git a/apps/desktop/src/components/editor/QueryEditor.vue b/apps/desktop/src/components/editor/QueryEditor.vue index c38e8230d1..0b87606efd 100644 --- a/apps/desktop/src/components/editor/QueryEditor.vue +++ b/apps/desktop/src/components/editor/QueryEditor.vue @@ -34,8 +34,10 @@ import { useToast } from "@/composables/useToast"; import { buildSelectStarExpansion, buildSqlCompletionItemsFromContext, + buildPostgresSequenceLiteralCompletionItems, getSqlFunctionSignatureHelp, getSqlCompletionContext, + getPostgresSequenceLiteralCompletionContext, getSqlCompletionResultValidFor, isSqlCompletionSuppressedContext, isSqlLikeCompletionStatement, @@ -3296,12 +3298,13 @@ async function provideSqlCompletions(context: CompletionContext) { } if (props.databaseType === "victoriametrics") return null; const hasDatabase = props.database != null; + const sequenceLiteralContext = getPostgresSequenceLiteralCompletionContext(fullDoc, position, props.databaseType); const epoch = ++completionEpoch; try { // 1. Suppressed context (comment / string literal) rejects everything, including explicit. - if (isSqlCompletionSuppressedContext(fullDoc, position)) return null; + if (isSqlCompletionSuppressedContext(fullDoc, position) && !sequenceLiteralContext) return null; // 2. Determine completion origin (session-level marker). activeCompletionOrigin = originForSqlCompletionProvider(activeCompletionOrigin, context.explicit); @@ -3324,12 +3327,12 @@ async function provideSqlCompletions(context: CompletionContext) { // require-prefix: only compute local facts (no positionalEligible). if (mode === "require-prefix") { - const ctx = getSqlCompletionContext(fullDoc, position, sqlCompletionDialectOptions()); + const ctx = sequenceLiteralContext ?? getSqlCompletionContext(fullDoc, position, sqlCompletionDialectOptions()); const prevChar = fullDoc[position - 1] ?? ""; const facts: SqlCompletionTriggerFacts = { origin, hasIdentifierPrefix: ctx.prefix.length > 0, - qualifierTriggered: prevChar === "." && ctx.qualifier != null, + qualifierTriggered: prevChar === "." && ("from" in ctx ? ctx.schema != null : ctx.qualifier != null), useDatabasePrefix, }; if (!shouldAllowSqlCompletionTrigger(mode, facts)) return null; @@ -3337,13 +3340,13 @@ async function provideSqlCompletions(context: CompletionContext) { // positional: compute positionalEligible (lazy). if (mode === "positional") { - const ctx = getSqlCompletionContext(fullDoc, position, sqlCompletionDialectOptions()); + const ctx = sequenceLiteralContext ?? getSqlCompletionContext(fullDoc, position, sqlCompletionDialectOptions()); const prevChar = fullDoc[position - 1] ?? ""; const positionalEligible = shouldAutoOpenSqlCompletion(fullDoc, position, sqlCompletionDialectOptions()); const facts: SqlCompletionTriggerFacts = { origin, hasIdentifierPrefix: ctx.prefix.length > 0, - qualifierTriggered: prevChar === "." && ctx.qualifier != null, + qualifierTriggered: prevChar === "." && ("from" in ctx ? ctx.schema != null : ctx.qualifier != null), useDatabasePrefix, positionalEligible, }; @@ -3379,6 +3382,13 @@ async function provideSqlCompletions(context: CompletionContext) { return buildCompletionResult(items, useDatabaseCompletion.from, undefined, useDatabaseCompletion.prefix); } + if (sequenceLiteralContext) { + if (!hasDatabase) return null; + const sequences = await connectionStore.listCompletionObjects(props.connectionId, props.database!, sequenceLiteralContext.prefix, MAX_COMPLETION_TABLES, sequenceLiteralContext.schema, undefined, false, undefined, ["sequence"], sequenceLiteralContext.nameQuoted); + if (epoch !== completionEpoch) return null; + return buildCompletionResult(buildPostgresSequenceLiteralCompletionItems(sequenceLiteralContext, sequences), sequenceLiteralContext.from, undefined, sequenceLiteralContext.prefix); + } + const legacyCompletionContext = getSqlCompletionContext(fullDoc, position, sqlCompletionDialectOptions()); const semanticModel = SEMANTIC_SQL_COMPLETION_ENABLED ? buildSqlSemanticModel(fullDoc, position, sqlCompletionDialectOptions()) : null; let completionContext = semanticModel ? sqlCompletionContextFromSemantic(semanticModel, legacyCompletionContext) : legacyCompletionContext; @@ -3546,7 +3556,8 @@ function flushImeComposition() { * Used by flushImeComposition and shouldStartSqlCompletionAfterInput. */ function shouldTriggerSqlCompletionForPosition(fullDoc: string, position: number): boolean { - if (isSqlCompletionSuppressedContext(fullDoc, position)) return false; + const sequenceLiteralContext = getPostgresSequenceLiteralCompletionContext(fullDoc, position, props.databaseType); + if (isSqlCompletionSuppressedContext(fullDoc, position) && !sequenceLiteralContext) return false; const mode = settingsStore.editorSettings.completionTriggerMode; if (mode === "manual") return false; @@ -3558,25 +3569,25 @@ function shouldTriggerSqlCompletionForPosition(fullDoc: string, position: number const useDatabasePrefix = useDatabaseCompletion?.prefix ?? null; if (mode === "require-prefix") { - const ctx = getSqlCompletionContext(fullDoc, position, sqlCompletionDialectOptions()); + const ctx = sequenceLiteralContext ?? getSqlCompletionContext(fullDoc, position, sqlCompletionDialectOptions()); const prevChar = fullDoc[position - 1] ?? ""; const facts: SqlCompletionTriggerFacts = { origin: "typing", hasIdentifierPrefix: ctx.prefix.length > 0, - qualifierTriggered: prevChar === "." && ctx.qualifier != null, + qualifierTriggered: prevChar === "." && ("from" in ctx ? ctx.schema != null : ctx.qualifier != null), useDatabasePrefix, }; return shouldAllowSqlCompletionTrigger(mode, facts); } // positional - const ctx = getSqlCompletionContext(fullDoc, position, sqlCompletionDialectOptions()); + const ctx = sequenceLiteralContext ?? getSqlCompletionContext(fullDoc, position, sqlCompletionDialectOptions()); const prevChar = fullDoc[position - 1] ?? ""; const positionalEligible = shouldAutoOpenSqlCompletion(fullDoc, position, sqlCompletionDialectOptions()); const facts: SqlCompletionTriggerFacts = { origin: "typing", hasIdentifierPrefix: ctx.prefix.length > 0, - qualifierTriggered: prevChar === "." && ctx.qualifier != null, + qualifierTriggered: prevChar === "." && ("from" in ctx ? ctx.schema != null : ctx.qualifier != null), useDatabasePrefix, positionalEligible, }; 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 d70dd5c9ea..f2e2ae10d9 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlCompletion.context.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/sqlCompletion.context.spec.ts @@ -1,5 +1,5 @@ import { describe, expect, it } from "vitest"; -import { buildSelectStarExpansion, buildSqlCompletionItems, getSqlCompletionContext, selectStarResultColumnsMatch, shouldAutoOpenSqlCompletion } from "@/lib/sql/sqlCompletion"; +import { buildPostgresSequenceLiteralCompletionItems, buildSelectStarExpansion, buildSqlCompletionItems, getPostgresSequenceLiteralCompletionContext, getSqlCompletionContext, selectStarResultColumnsMatch, shouldAutoOpenSqlCompletion } from "@/lib/sql/sqlCompletion"; import { sqlCompletionContextFromSemantic } from "@/lib/sql/semantic/completion"; import { buildSqlSemanticModel } from "@/lib/sql/semantic/model"; import { originForSqlCompletionProvider, originForTypedSqlCompletionStart, shouldAllowSqlCompletionTrigger, type SqlCompletionTriggerFacts } from "@/lib/sql/sqlCompletionTriggerPolicy"; @@ -17,6 +17,96 @@ describe("sqlCompletion keyword snippets", () => { }); }); +describe("PostgreSQL sequence literal completion", () => { + it.each(["nextval", "currval", "setval"])("recognizes the first %s regclass literal", (functionName) => { + const sql = `SELECT ${functionName}('public.order_`; + const context = getPostgresSequenceLiteralCompletionContext(sql, sql.length, "postgres"); + + expect(context).toEqual( + expect.objectContaining({ + prefix: "order_", + schema: "public", + schemaQuoted: false, + nameQuoted: false, + }), + ); + expect(context?.from).toBe(sql.lastIndexOf("order_")); + expect(shouldAutoOpenSqlCompletion(sql, sql.length, { databaseType: "postgres" })).toBe(true); + }); + + it("preserves quoted mixed-case schema and sequence identifiers", () => { + const sql = `SELECT pg_catalog.nextval('"App"."Order`; + const context = getPostgresSequenceLiteralCompletionContext(sql, sql.length, "postgres"); + + expect(context).toEqual( + expect.objectContaining({ + prefix: "Order", + schema: "App", + schemaQuoted: true, + nameQuoted: true, + nameQuoteClosed: false, + }), + ); + expect( + buildPostgresSequenceLiteralCompletionItems(context!, [ + { name: "OrderSequence", schema: "App", type: "sequence" }, + { name: "order_sequence", schema: "App", type: "sequence" }, + ]), + ).toEqual([ + expect.objectContaining({ + label: "OrderSequence", + apply: 'OrderSequence"', + replaceClosingQuote: '"', + detail: "sequence in App", + }), + ]); + }); + + it("keeps doubled apostrophes inside the sequence identifier and escapes insertion", () => { + const sql = `SELECT nextval('"customer''s_`; + const context = getPostgresSequenceLiteralCompletionContext(sql, sql.length, "postgres"); + + expect(context).toEqual(expect.objectContaining({ prefix: "customer's_", nameQuoted: true })); + expect(buildPostgresSequenceLiteralCompletionItems(context!, [{ name: "customer's_seq", schema: "public", type: "sequence" }])).toEqual([expect.objectContaining({ label: "customer's_seq", apply: `customer''s_seq"` })]); + }); + + it("quotes an accepted mixed-case identifier in an unquoted literal", () => { + const sql = "SELECT nextval('mix"; + const context = getPostgresSequenceLiteralCompletionContext(sql, sql.length, "postgres"); + + expect(buildPostgresSequenceLiteralCompletionItems(context!, [{ name: "MixedSequence", schema: "public", type: "sequence" }])).toEqual([expect.objectContaining({ label: "MixedSequence", filterText: "MixedSequence", apply: '"MixedSequence"' })]); + }); + + it("does not expose sequence metadata through ordinary SQL object completion", () => { + const sql = "SELECT order_"; + const items = buildSqlCompletionItems(sql, sql.length, { + tables: [], + objects: [{ name: "order_seq", schema: "public", type: "sequence" }], + columnsByTable: new Map(), + databaseType: "postgres", + dialect: "postgres", + }); + + expect(items.some((item) => item.label === "order_seq")).toBe(false); + }); + + it.each([ + ["SELECT 'order_", "postgres"], + ["SELECT nextval('order_", "mysql"], + ["SELECT nextval('order_') || 'suffix", "postgres"], + ["SELECT setval(42, 'order_", "postgres"], + ["SELECT app.nextval('order_", "postgres"], + ["SELECT app. nextval('order_", "postgres"], + ["SELECT app.pg_catalog.nextval('order_", "postgres"], + [`SELECT "PG_CATALOG".nextval('order_`, "postgres"], + ["SELECT $$ nextval('order_", "postgres"], + ["SELECT $body$ nextval('order_", "postgres"], + ] as const)("does not enable sequence completion for unrelated literals: %s (%s)", (sql, databaseType) => { + expect(getPostgresSequenceLiteralCompletionContext(sql, sql.length, databaseType)).toBeNull(); + expect(shouldAutoOpenSqlCompletion(sql, sql.length, { databaseType })).toBe(false); + }); +}); + describe("SELECT star expansion", () => { it("reuses completion column ordering for an unqualified star", () => { const sql = "SELECT * FROM apis"; diff --git a/apps/desktop/src/lib/sql/sqlCompletion.ts b/apps/desktop/src/lib/sql/sqlCompletion.ts index 57a3c56cf2..aa7ef37c2a 100644 --- a/apps/desktop/src/lib/sql/sqlCompletion.ts +++ b/apps/desktop/src/lib/sql/sqlCompletion.ts @@ -1172,7 +1172,7 @@ export interface SqlCompletionTable { export interface SqlCompletionObject { name: string; schema?: string; - type: "procedure" | "function" | "trigger" | "package"; + type: "procedure" | "function" | "trigger" | "package" | "sequence"; parentSchema?: string; parentName?: string; dataType?: string; @@ -1208,6 +1208,7 @@ export interface SqlCompletionItem { detail?: string; info?: string; apply?: string; + replaceClosingQuote?: '"' | "'"; boost: number; exactMatch?: boolean; dedupeKey?: string; @@ -1271,6 +1272,15 @@ export interface SqlCompletionContext { dataTypeContext: boolean; } +export interface PostgresSequenceLiteralCompletionContext { + from: number; + prefix: string; + schema?: string; + schemaQuoted: boolean; + nameQuoted: boolean; + nameQuoteClosed: boolean; +} + export interface SqlFunctionSignatureHelpOverload { signature: string; parameterGroups: string[][]; @@ -1478,6 +1488,7 @@ class SqlCompletionProvider { } export function shouldAutoOpenSqlCompletion(sql: string, cursor: number, options: SqlSemanticBuildOptions = {}): boolean { + if (getPostgresSequenceLiteralCompletionContext(sql, cursor, options.databaseType)) return true; if (isSqlCompletionSuppressedContext(sql, cursor)) return false; const previousChar = sql[cursor - 1]; if (!previousChar) return false; @@ -1511,6 +1522,209 @@ export function isSqlCompletionSuppressedContext(sql: string, cursor: number): b return context.inLineComment || context.inBlockComment || context.inStringLiteral; } +export function getPostgresSequenceLiteralCompletionContext(sql: string, cursor: number, databaseType?: DatabaseType): PostgresSequenceLiteralCompletionContext | null { + if (databaseType !== "postgres") return null; + const position = Math.max(0, Math.min(cursor, sql.length)); + const literalStart = activeSingleQuotedLiteralStart(sql, position); + if (literalStart == null) return null; + + const beforeLiteral = sql.slice(0, literalStart); + if (!isPostgresSequenceFunctionCall(beforeLiteral)) return null; + + const rawLiteral = sql.slice(literalStart + 1, position); + const parsed = parsePostgresRegclassPrefix(rawLiteral); + if (!parsed) return null; + return { + from: literalStart + 1 + parsed.nameStart, + prefix: parsed.name, + schema: parsed.schema, + schemaQuoted: parsed.schemaQuoted, + nameQuoted: parsed.nameQuoted, + nameQuoteClosed: parsed.nameQuoteClosed, + }; +} + +export function buildPostgresSequenceLiteralCompletionItems(context: PostgresSequenceLiteralCompletionContext, objects: SqlCompletionObject[]): SqlCompletionItem[] { + return objects + .filter((object) => object.type === "sequence") + .filter((object) => { + if (context.schema) { + if (!object.schema) return false; + const schemaMatches = context.schemaQuoted ? object.schema === context.schema : object.schema.toLowerCase() === context.schema.toLowerCase(); + if (!schemaMatches) return false; + } + return context.nameQuoted ? object.name.startsWith(context.prefix) : object.name.toLowerCase().startsWith(context.prefix.toLowerCase()); + }) + .map((object) => { + const identifier = context.nameQuoted ? `${escapePostgresSequenceLiteralQuotedIdentifier(object.name)}"` : escapePostgresSequenceLiteralIdentifier(quoteSqlIdentifier(object.name, "postgres")); + return { + label: object.name, + filterText: context.nameQuoted ? identifier.slice(0, -1) : escapePostgresSequenceLiteralIdentifier(object.name), + type: "variable" as const, + detail: object.schema ? `sequence in ${object.schema}` : "sequence", + info: object.comment?.trim() || undefined, + apply: identifier, + replaceClosingQuote: context.nameQuoted && !context.nameQuoteClosed ? ('"' as const) : undefined, + boost: computeBoost(object.name, context.prefix) + 1_200, + }; + }) + .sort(compareCompletionItems) + .slice(0, MAX_TABLE_COMPLETION_ITEMS); +} + +function escapePostgresSequenceLiteralIdentifier(value: string): string { + return value.replaceAll("'", "''"); +} + +function escapePostgresSequenceLiteralQuotedIdentifier(value: string): string { + return escapePostgresSequenceLiteralIdentifier(value.replaceAll('"', '""')); +} + +function isPostgresSequenceFunctionCall(beforeLiteral: string): boolean { + const call = /(?:nextval|currval|setval)\s*\(\s*$/i.exec(beforeLiteral); + if (!call) return false; + + const beforeFunction = beforeLiteral.slice(0, call.index); + const beforeFunctionTrimmed = beforeFunction.trimEnd(); + if (!beforeFunctionTrimmed.endsWith(".")) { + return !beforeFunction || /\s$/.test(beforeFunction) || !/[\w$."`]$/.test(beforeFunction); + } + + const beforeDot = beforeFunctionTrimmed.slice(0, -1).trimEnd(); + const qualifier = /(?:"((?:""|[^"])*)"|([A-Za-z_][\w$]*))$/.exec(beforeDot); + if (!qualifier) return false; + const quotedQualifier = qualifier[1]; + const qualifierName = (quotedQualifier ?? qualifier[2] ?? "").replaceAll('""', '"'); + if (quotedQualifier != null ? qualifierName !== "pg_catalog" : qualifierName.toLowerCase() !== "pg_catalog") return false; + const beforeQualifier = beforeDot.slice(0, qualifier.index); + return !beforeQualifier || /\s$/.test(beforeQualifier) || !/[\w$."`]$/.test(beforeQualifier); +} + +function activeSingleQuotedLiteralStart(sql: string, cursor: number): number | null { + let literalStart: number | null = null; + let dollarQuoteDelimiter: string | null = null; + let inDoubleQuote = false; + let inLineComment = false; + let inBlockComment = false; + for (let index = 0; index < cursor; index += 1) { + const char = sql[index] ?? ""; + const next = sql[index + 1] ?? ""; + if (dollarQuoteDelimiter) { + if (sql.startsWith(dollarQuoteDelimiter, index)) { + index += dollarQuoteDelimiter.length - 1; + dollarQuoteDelimiter = null; + } + continue; + } + if (inLineComment) { + if (char === "\n" || char === "\r") inLineComment = false; + continue; + } + if (inBlockComment) { + if (char === "*" && next === "/") { + inBlockComment = false; + index++; + } + continue; + } + if (literalStart != null) { + if (char === "'" && next === "'") { + index++; + } else if (char === "'") { + literalStart = null; + } + continue; + } + if (inDoubleQuote) { + if (char === '"' && next === '"') index++; + else if (char === '"') inDoubleQuote = false; + continue; + } + if (char === "-" && next === "-") { + inLineComment = true; + index++; + } else if (char === "/" && next === "*") { + inBlockComment = true; + index++; + } else if (char === '"') { + inDoubleQuote = true; + } else if (char === "$") { + const delimiter = /^\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$/.exec(sql.slice(index, cursor))?.[0]; + if (delimiter) { + dollarQuoteDelimiter = delimiter; + index += delimiter.length - 1; + } + } else if (char === "'") { + literalStart = index; + } + } + return literalStart; +} + +function parsePostgresRegclassPrefix(raw: string): { nameStart: number; name: string; schema?: string; schemaQuoted: boolean; nameQuoted: boolean; nameQuoteClosed: boolean } | null { + const parts: Array<{ start: number; value: string; quoted: boolean; quoteClosed: boolean }> = []; + let index = 0; + while (index < raw.length && /\s/.test(raw[index] ?? "")) index++; + while (index <= raw.length && parts.length < 2) { + const start = index; + let quoted = false; + let quoteClosed = false; + let value = ""; + if (raw[index] === '"') { + quoted = true; + index++; + const valueStart = index; + while (index < raw.length) { + const char = raw[index] ?? ""; + if (char === '"' && raw[index + 1] === '"') { + value += '"'; + index += 2; + } else if (char === '"') { + quoteClosed = true; + index++; + break; + } else if (char === "'" && raw[index + 1] === "'") { + value += "'"; + index += 2; + } else { + value += char; + index++; + } + } + parts.push({ start: valueStart, value, quoted, quoteClosed }); + } else { + while (index < raw.length && raw[index] !== ".") { + const char = raw[index] ?? ""; + if (char === "'" && raw[index + 1] === "'") { + value += "'"; + index += 2; + } else if (/\s/.test(char) || char === '"') { + return null; + } else { + value += char; + index++; + } + } + parts.push({ start, value, quoted, quoteClosed }); + } + if (index >= raw.length) break; + if (raw[index] !== "." || (!quoteClosed && quoted)) return null; + index++; + if (index >= raw.length) parts.push({ start: index, value: "", quoted: false, quoteClosed: false }); + } + if (index < raw.length || parts.length === 0 || parts.length > 2) return null; + const name = parts[parts.length - 1]!; + const schema = parts.length === 2 ? parts[0] : undefined; + return { + nameStart: name.start, + name: name.value, + schema: schema?.value ? (schema.quoted ? schema.value : schema.value.toLowerCase()) : undefined, + schemaQuoted: schema?.quoted ?? false, + nameQuoted: name.quoted, + nameQuoteClosed: name.quoteClosed, + }; +} + export function isSqlStringLiteralContext(sql: string, cursor: number): boolean { return getSqlLexicalContext(sql, cursor).inStringLiteral; } @@ -3057,7 +3271,7 @@ function buildObjectItems(context: SqlCompletionContext, objects: SqlCompletionO 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") && (!onlyFunctions || (object.type === "function" && object.name.toLowerCase().startsWith(context.prefix.toLowerCase()))) && objectMatchesCompletionContext(object, context)) + .filter((object) => object.type !== "sequence" && (!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); diff --git a/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts b/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts index ed062fc585..56adf118d8 100644 --- a/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts +++ b/apps/desktop/src/stores/__tests__/connectionStore.completion.spec.ts @@ -883,6 +883,86 @@ describe("connectionStore completion assistant", () => { ]); }); + it("loads PostgreSQL sequences without falling back to routine metadata", async () => { + const completionAssistantSearch = vi.fn().mockResolvedValue({ + candidates: [{ name: "OrderSequence", kind: "sequence", schema: "App" }], + incomplete: false, + fallback_used: false, + }); + const listCompletionObjects = vi.fn().mockResolvedValue([{ name: "not_a_sequence", object_type: "FUNCTION" }]); + + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => ({ + checkConnectionHealth: vi.fn().mockResolvedValue(undefined), + completionAssistantSearch, + listCompletionObjects, + })); + + 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", "Order", 20, "App", undefined, false, undefined, ["sequence"], true); + await store.listCompletionObjects("pg-1", "app", "Order", 20, undefined, undefined, false, undefined, ["sequence"]); + await store.listCompletionObjects("pg-1", "app", "order", 20, "App", undefined, false, undefined, ["sequence"], true); + + expect(completionAssistantSearch).toHaveBeenNthCalledWith( + 1, + expect.objectContaining({ + object_kinds: ["sequence"], + mask: "Order", + case_sensitive: true, + schema: "App", + parent_schema: null, + }), + ); + expect(completionAssistantSearch).toHaveBeenNthCalledWith( + 2, + expect.objectContaining({ + object_kinds: ["sequence"], + mask: "Order", + case_sensitive: false, + schema: null, + parent_schema: null, + }), + ); + expect(completionAssistantSearch).toHaveBeenNthCalledWith( + 3, + expect.objectContaining({ + object_kinds: ["sequence"], + mask: "order", + case_sensitive: true, + schema: "App", + parent_schema: null, + }), + ); + expect(listCompletionObjects).not.toHaveBeenCalled(); + expect(objects).toEqual([expect.objectContaining({ name: "OrderSequence", schema: "App", type: "sequence" })]); + }); + + it.each([ + ["empty", vi.fn().mockResolvedValue({ candidates: [], incomplete: false, fallback_used: false })], + ["permission error", vi.fn().mockRejectedValue(new Error("permission denied for sequence"))], + ])("returns an empty sequence fallback for %s metadata", async (_caseName, completionAssistantSearch) => { + const listCompletionObjects = vi.fn().mockResolvedValue([{ name: "not_a_sequence", object_type: "FUNCTION" }]); + + vi.doMock("@/lib/backend/tauriRuntime", () => ({ isTauriRuntime: () => false })); + vi.doMock("@/lib/backend/api", () => ({ + checkConnectionHealth: vi.fn().mockResolvedValue(undefined), + completionAssistantSearch, + listCompletionObjects, + })); + + const { useConnectionStore } = await import("@/stores/connectionStore"); + const store = useConnectionStore(); + store.connections = [postgresConnection()]; + store.connectedIds.add("pg-1"); + + await expect(store.listCompletionObjects("pg-1", "app", "missing", 20, undefined, undefined, false, undefined, ["sequence"])).resolves.toEqual([]); + expect(listCompletionObjects).not.toHaveBeenCalled(); + }); + it("searches the server-reported SQL Server default schema without treating the username as a schema", async () => { const completionAssistantSearch = vi.fn().mockResolvedValue({ candidates: [{ name: "st_area", kind: "function", schema: "app_user", data_type: "float" }], diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index 63fc198c86..813f658d43 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -5921,7 +5921,7 @@ export const useConnectionStore = defineStore("connection", () => { 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; + const type = candidate.kind === "procedure" ? "procedure" : candidate.kind === "function" ? "function" : candidate.kind === "sequence" ? "sequence" : 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 { @@ -6000,20 +6000,23 @@ export const useConnectionStore = defineStore("connection", () => { globalSearch: boolean, currentSchema: string | undefined, objectKinds: CompletionAssistantObjectKind[], + caseSensitive: boolean, ): Promise { const databaseType = getConfig(connectionId)?.db_type; const oracleAssistant = databaseType === "oracle"; const requestedSchema = schema?.trim() || currentSchema?.trim() || undefined; - const preferredSchema = oracleAssistant ? completionPreferredSchema(connectionId, currentSchema) : requestedSchema || (databaseType === "postgres" ? "public" : databaseType === "mysql" ? database : undefined); + const sequenceOnly = objectKinds.length === 1 && objectKinds[0] === "sequence"; + const preferredSchema = oracleAssistant ? completionPreferredSchema(connectionId, currentSchema) : requestedSchema || (!sequenceOnly && databaseType === "postgres" ? "public" : databaseType === "mysql" ? database : undefined); const response = await completionAssistantSearch({ connection_id: connectionId, database, schema: oracleAssistant ? (preferredSchema ?? null) : (requestedSchema ?? null), object_kinds: objectKinds, mask: filter.trim(), + case_sensitive: caseSensitive, max_results: limit ?? 200, global_search: globalSearch, - parent_schema: globalSearch ? null : (schema ?? null), + parent_schema: globalSearch || sequenceOnly ? null : (schema ?? null), parent_name: parentName ?? null, match_mode: "prefix", }); @@ -6528,11 +6531,25 @@ 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, objectKinds: CompletionAssistantObjectKind[] = ["routine"]): Promise { + async function listCompletionObjects( + connectionId: string, + database: string, + filter = "", + limit?: number, + schema?: string, + parentName?: string, + globalSearch = false, + currentSchema?: string, + objectKinds: CompletionAssistantObjectKind[] = ["routine"], + caseSensitive = false, + ): Promise { const normalizedFilter = filter.trim().toLowerCase(); + const cacheFilter = caseSensitive ? filter.trim() : normalizedFilter; 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 ?? ""}`; + const cacheKey = filteredRoutineAssistant + ? `${connectionId}:${database}:${schema ?? ""}:${parentName ?? ""}:${cacheFilter}:${limit ?? ""}:${globalSearch ? "global" : "scoped"}:${currentSchema ?? ""}:${[...objectKinds].sort().join(",")}:${caseSensitive ? "case-sensitive" : "case-insensitive"}` + : `${connectionId}:${database}:${schema ?? ""}`; if (!completionObjectsCache.value[cacheKey]) { await withCompletionInFlight( `${cacheKey}:objects`, @@ -6540,10 +6557,14 @@ export const useConnectionStore = defineStore("connection", () => { await ensureConnected(connectionId); if (filteredRoutineAssistant) { try { - completionObjectsCache.value[cacheKey] = dedupeCompletionObjects(await listCompletionAssistantObjects(connectionId, database, filter, limit, schema, parentName, globalSearch, currentSchema, objectKinds)); + completionObjectsCache.value[cacheKey] = dedupeCompletionObjects(await listCompletionAssistantObjects(connectionId, database, filter, limit, schema, parentName, globalSearch, currentSchema, objectKinds, caseSensitive)); } 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)); + if (objectKinds.length === 1 && objectKinds[0] === "sequence") { + completionObjectsCache.value[cacheKey] = []; + } 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)); + } } } else { const objects = isSchemaAwareDatabase(connectionId) ? await listSchemaAwareCompletionObjects(connectionId, database, schema) : await api.listCompletionObjects(connectionId, database, schema || database); diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts index 6d71f538a7..b866c5d6b3 100644 --- a/apps/desktop/src/types/database.ts +++ b/apps/desktop/src/types/database.ts @@ -91,9 +91,9 @@ export interface SqlSnippet { enabled?: boolean; } -export type CompletionAssistantObjectKind = "database" | "schema" | "table" | "view" | "routine" | "procedure" | "function" | "column"; +export type CompletionAssistantObjectKind = "database" | "schema" | "table" | "view" | "routine" | "procedure" | "function" | "column" | "sequence"; -export type CompletionAssistantCandidateKind = "database" | "schema" | "table" | "view" | "procedure" | "function" | "column" | "object"; +export type CompletionAssistantCandidateKind = "database" | "schema" | "table" | "view" | "procedure" | "function" | "column" | "sequence" | "object"; export type CompletionAssistantMatchMode = "prefix" | "contains"; diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 8b94a94cbf..8647b9d12a 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -2540,6 +2540,29 @@ pub async fn completion_assistant_search( } } + if candidates.len() < limit && kinds.iter().any(|kind| matches!(kind, CompletionAssistantObjectKind::Sequence)) { + let rows = postgres_query_cached( + &client, + postgres_completion_sequences_sql(), + &[&schema, &pattern, &request.case_sensitive, &((limit - candidates.len()) as i64)], + ) + .await + .map_err(|e| e.to_string())?; + for row in rows { + candidates.push(CompletionAssistantCandidate { + name: pg_row_try_string(&row, 0), + kind: CompletionAssistantCandidateKind::Sequence, + database: Some(request.database.clone()), + schema: Some(pg_row_try_string(&row, 1)), + parent_schema: None, + parent_name: None, + comment: row.try_get::<_, Option>(2).ok().flatten(), + data_type: None, + signature: None, + }); + } + } + if candidates.len() < limit && kinds.iter().any(|kind| matches!(kind, CompletionAssistantObjectKind::Column)) { let table = request.parent_name.as_deref().unwrap_or(""); if !table.is_empty() { @@ -2611,6 +2634,22 @@ fn postgres_completion_routines_sql() -> &'static str { ORDER BY p.proname LIMIT $4" } +fn postgres_completion_sequences_sql() -> &'static str { + "SELECT c.relname, n.nspname, obj_description(c.oid) AS sequence_comment \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.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 = 'S' \ + AND pg_catalog.has_schema_privilege(n.oid, 'USAGE') \ + AND CASE WHEN c.relkind = 'S' \ + THEN pg_catalog.has_sequence_privilege(c.oid, 'USAGE, SELECT, UPDATE') \ + ELSE FALSE END \ + AND ($2 = '%%' OR CASE WHEN $3 THEN c.relname LIKE $2 ESCAPE '~' \ + ELSE c.relname ILIKE $2 ESCAPE '~' END) \ + ORDER BY c.relname LIMIT $4" +} + fn postgres_completion_columns_sql() -> &'static str { "SELECT a.attname, pg_catalog.format_type(a.atttypid, a.atttypmod), col_description(c.oid, a.attnum) \ FROM pg_catalog.pg_attribute a \ @@ -9121,6 +9160,90 @@ mod tests { .expect("drop live table fixtures"); } + #[tokio::test] + #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] + async fn completion_assistant_searches_visible_and_qualified_sequences() { + let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); + let pool = connect_with_local_timezone(&url, Duration::from_secs(5), "UTC") + .await + .expect("connect PostgreSQL database"); + let client = pool.get().await.expect("checkout postgres"); + client + .batch_execute( + r#"DROP SCHEMA IF EXISTS dbx_issue_4005_visible CASCADE; + DROP SCHEMA IF EXISTS dbx_issue_4005_shadowed CASCADE; + CREATE SCHEMA dbx_issue_4005_visible; + CREATE SCHEMA dbx_issue_4005_shadowed; + CREATE SEQUENCE dbx_issue_4005_visible.order_seq; + CREATE SEQUENCE dbx_issue_4005_shadowed.order_seq; + CREATE SEQUENCE dbx_issue_4005_visible."OrderSequence"; + SET search_path = dbx_issue_4005_visible, dbx_issue_4005_shadowed, public;"#, + ) + .await + .expect("create live sequence fixtures"); + drop(client); + + let visible = completion_assistant_search( + &pool, + &CompletionAssistantRequest { + connection_id: "live".into(), + database: "postgres".into(), + schema: None, + object_kinds: vec![CompletionAssistantObjectKind::Sequence], + mask: "order_".into(), + case_sensitive: false, + global_search: false, + max_results: Some(20), + search_in_comments: false, + search_in_definitions: false, + parent_schema: None, + parent_name: None, + match_mode: Some(CompletionAssistantMatchMode::Prefix), + }, + ) + .await + .expect("complete visible sequence"); + assert_eq!( + visible + .candidates + .iter() + .map(|candidate| (candidate.name.as_str(), candidate.schema.as_deref())) + .collect::>(), + [("order_seq", Some("dbx_issue_4005_visible"))] + ); + + let qualified = completion_assistant_search( + &pool, + &CompletionAssistantRequest { + connection_id: "live".into(), + database: "postgres".into(), + schema: Some("dbx_issue_4005_visible".into()), + object_kinds: vec![CompletionAssistantObjectKind::Sequence], + mask: "OrderS".into(), + case_sensitive: true, + global_search: false, + max_results: Some(20), + search_in_comments: false, + search_in_definitions: false, + parent_schema: None, + parent_name: None, + match_mode: Some(CompletionAssistantMatchMode::Prefix), + }, + ) + .await + .expect("complete qualified mixed-case sequence"); + assert_eq!(qualified.candidates[0].name, "OrderSequence"); + + let client = pool.get().await.expect("checkout postgres for cleanup"); + client + .batch_execute( + r#"DROP SCHEMA IF EXISTS dbx_issue_4005_visible CASCADE; + DROP SCHEMA IF EXISTS dbx_issue_4005_shadowed CASCADE;"#, + ) + .await + .expect("drop live sequence fixtures"); + } + #[test] fn like_contains_pattern_escapes_wildcards() { assert_eq!(like_contains_pattern(""), "%%"); @@ -9174,6 +9297,15 @@ mod tests { assert!(postgres_completion_routines_sql().contains("p.prokind::text = ANY($3::text[])")); assert!(postgres_completion_routines_sql().contains("pg_get_function_identity_arguments(p.oid) AS signature")); assert!(postgres_completion_routines_sql().contains("ORDER BY p.proname LIMIT $4")); + assert!(postgres_completion_sequences_sql().contains("c.relkind = 'S'")); + assert!(postgres_completion_sequences_sql().contains("has_schema_privilege(n.oid, 'USAGE')")); + assert!(postgres_completion_sequences_sql().contains("pg_catalog.pg_table_is_visible(c.oid)")); + assert!(postgres_completion_sequences_sql().contains("CASE WHEN c.relkind = 'S'")); + assert!(postgres_completion_sequences_sql() + .contains("THEN pg_catalog.has_sequence_privilege(c.oid, 'USAGE, SELECT, UPDATE')")); + assert!(postgres_completion_sequences_sql().contains("CASE WHEN $3 THEN c.relname LIKE $2 ESCAPE '~'")); + assert!(postgres_completion_sequences_sql().contains("ELSE c.relname ILIKE $2 ESCAPE '~' END")); + assert!(postgres_completion_sequences_sql().contains("ORDER BY c.relname 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/types.rs b/crates/dbx-core/src/types.rs index 972404d7ac..ffb9de8fbc 100644 --- a/crates/dbx-core/src/types.rs +++ b/crates/dbx-core/src/types.rs @@ -181,6 +181,7 @@ pub enum CompletionAssistantObjectKind { Procedure, Function, Column, + Sequence, } impl CompletionAssistantObjectKind { @@ -203,6 +204,7 @@ pub enum CompletionAssistantCandidateKind { Procedure, Function, Column, + Sequence, Object, }