diff --git a/apps/desktop/src/components/editor/QueryEditor.vue b/apps/desktop/src/components/editor/QueryEditor.vue index 160462e94..cf549dc73 100644 --- a/apps/desktop/src/components/editor/QueryEditor.vue +++ b/apps/desktop/src/components/editor/QueryEditor.vue @@ -3187,9 +3187,9 @@ function isTableNameCompletionContext(completionContext: ReturnType [`${object.type}:${object.schema ?? ""}:${object.name}:${object.parentName ?? ""}`.toLowerCase(), index])); + const indexes = new Map(existing.map((object, index) => [completionObjectIdentityKey(object), index])); for (const object of incoming) { - const key = `${object.type}:${object.schema ?? ""}:${object.name}:${object.parentName ?? ""}`.toLowerCase(); + const key = completionObjectIdentityKey(object); const index = indexes.get(key); if (index == null) { indexes.set(key, merged.length); @@ -3201,6 +3201,10 @@ function mergeCompletionObjects(existing: SqlCompletionObject[], incoming: SqlCo return merged; } +function completionObjectIdentityKey(object: SqlCompletionObject): string { + return `${object.type}:${object.schema ?? ""}:${object.name}:${object.parentName ?? ""}:${object.signature?.trim() ?? ""}`.toLowerCase(); +} + function completionObjectsDiffer(existing: SqlCompletionObject[], incoming: SqlCompletionObject[]): boolean { if (existing.length !== incoming.length) return true; return existing.some((object, index) => { diff --git a/apps/desktop/src/lib/sql/sqlCompletion.ts b/apps/desktop/src/lib/sql/sqlCompletion.ts index af56d0ea2..4f55f6915 100644 --- a/apps/desktop/src/lib/sql/sqlCompletion.ts +++ b/apps/desktop/src/lib/sql/sqlCompletion.ts @@ -3024,17 +3024,19 @@ function buildObjectItems(context: SqlCompletionContext, objects: SqlCompletionO ? quoteSqlIdentifier(object.name, dialect) : (object.applyName ?? (object.schema && !objectInCurrentSchema ? `${quoteSqlIdentifier(object.schema, dialect)}.${quoteSqlIdentifier(object.name, dialect)}` : quoteSqlIdentifier(object.name, dialect))); 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 signature = object.signature?.trim(); + const detail = [locationDetail, signature ? `(${signature})` : undefined, object.dataType ? `[${object.dataType}]` : undefined].filter(Boolean).join(" "); const schemaBoost = onlyFunctions ? Math.min(object.boost ?? 0, 1000) : (object.boost ?? 0); const typeBoost = routineTypeBoost(object.type, prioritizeOracleFunctions && !onlyFunctions); + const baseDedupeKey = object.applyName || (databaseType === "oracle" && object.schema) ? applyName : undefined; return { label: object.name, type: "function" as const, detail, info: buildRoutineInfo(object), - apply: object.type === "trigger" || object.type === "package" ? applyName : `${applyName}()`, + apply: object.type === "trigger" || object.type === "package" ? applyName : buildRoutineApply(applyName, object.signature), boost: computeBoost(object.name, context.prefix) + typeBoost + schemaBoost, - dedupeKey: object.applyName || (databaseType === "oracle" && object.schema) ? applyName : undefined, + dedupeKey: signature ? `${baseDedupeKey ?? object.name}(${signature})` : baseDedupeKey, // Preserve exact routine matches before the capped candidate list is truncated. exactMatch: !!context.prefix && object.name.toLowerCase() === context.prefix.toLowerCase(), }; @@ -3043,6 +3045,51 @@ function buildObjectItems(context: SqlCompletionContext, objects: SqlCompletionO .slice(0, MAX_TABLE_COMPLETION_ITEMS); } +function buildRoutineApply(applyName: string, signature?: string): string { + const parameters = splitRoutineSignatureParameters(signature?.trim() ?? ""); + if (parameters.length === 0) return `${applyName}()`; + return `${applyName}(${parameters.map((parameter, index) => `\${${index + 1}:${escapeSnippetFieldName(parameter)}}`).join(", ")})`; +} + +function splitRoutineSignatureParameters(signature: string): string[] { + if (!signature) return []; + const parameters: string[] = []; + let start = 0; + let parenthesisDepth = 0; + let bracketDepth = 0; + let quoted = false; + + for (let index = 0; index < signature.length; index++) { + const char = signature[index]; + if (char === '"') { + if (quoted && signature[index + 1] === '"') { + index++; + } else { + quoted = !quoted; + } + continue; + } + if (quoted) continue; + if (char === "(") parenthesisDepth++; + else if (char === ")" && parenthesisDepth > 0) parenthesisDepth--; + else if (char === "[") bracketDepth++; + else if (char === "]" && bracketDepth > 0) bracketDepth--; + else if (char === "," && parenthesisDepth === 0 && bracketDepth === 0) { + const parameter = signature.slice(start, index).trim(); + if (parameter) parameters.push(parameter); + start = index + 1; + } + } + + const parameter = signature.slice(start).trim(); + if (parameter) parameters.push(parameter); + return parameters; +} + +function escapeSnippetFieldName(value: string): string { + return value.replace(/[{}]/g, "\\$&"); +} + 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); diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index de677d269..87248c32d 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -5130,6 +5130,7 @@ export const useConnectionStore = defineStore("connection", () => { parentSchema: candidate.parent_schema ?? undefined, parentName: candidate.parent_name ?? undefined, dataType, + signature: candidate.signature ?? undefined, comment: candidate.comment ?? null, applyName: completionCandidateApplyName(candidate.name, candidate.schema, preferredSchema), boost: oracleMetadata ? completionCandidateSchemaBoost(candidate.schema, preferredSchema) : completionRoutineSchemaBoost(candidate.schema, preferredSchema), @@ -5736,7 +5737,7 @@ export const useConnectionStore = defineStore("connection", () => { const seen = new Set(); const deduped: SqlCompletionObject[] = []; for (const object of objects) { - const key = `${object.type}:${object.schema ?? ""}:${object.name}:${object.parentName ?? ""}`.toLowerCase(); + const key = `${object.type}:${object.schema ?? ""}:${object.name}:${object.parentName ?? ""}:${object.signature?.trim() ?? ""}`.toLowerCase(); if (seen.has(key)) continue; seen.add(key); deduped.push(object); diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts index 7075720af..3e8510efa 100644 --- a/apps/desktop/src/types/database.ts +++ b/apps/desktop/src/types/database.ts @@ -106,6 +106,7 @@ export interface CompletionAssistantCandidate { parent_name?: string | null; comment?: string | null; data_type?: string | null; + signature?: string | null; } export interface CompletionAssistantResponse { diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 3aa0c5b80..1050cc639 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -1919,6 +1919,7 @@ pub async fn completion_assistant_search( parent_name: None, comment: None, data_type: None, + signature: None, }); } } @@ -1944,6 +1945,7 @@ pub async fn completion_assistant_search( .map(|s| fix_potential_double_encoding(&s)) .filter(|s| !s.is_empty()), data_type: None, + signature: None, }); } } @@ -1969,6 +1971,7 @@ pub async fn completion_assistant_search( .map(|s| fix_potential_double_encoding(&s)) .filter(|s| !s.is_empty()), data_type: get_opt_str(&row, "data_type"), + signature: None, }); } } @@ -1990,6 +1993,7 @@ pub async fn completion_assistant_search( .map(|s| fix_potential_double_encoding(&s)) .filter(|s| !s.is_empty()), data_type: Some(get_str_by_name(&row, "data_type")), + signature: None, }); } } diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 8d3dbb754..0d9ba47b6 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -1869,6 +1869,7 @@ pub async fn completion_assistant_search( parent_name: None, comment: None, data_type: None, + signature: None, }); } } @@ -1897,6 +1898,7 @@ pub async fn completion_assistant_search( parent_name: row.try_get::<_, Option>(5).ok().flatten(), comment: row.try_get::<_, Option>(3).ok().flatten(), data_type: None, + signature: None, }); } } @@ -1925,6 +1927,7 @@ pub async fn completion_assistant_search( parent_name: None, comment: row.try_get::<_, Option>(3).ok().flatten(), data_type: row.try_get::<_, Option>(4).ok().flatten(), + signature: row.try_get::<_, Option>(5).ok().flatten(), }); } } @@ -1962,6 +1965,7 @@ pub async fn completion_assistant_search( parent_name: Some(table.to_string()), comment: row.try_get::<_, Option>(2).ok().flatten(), data_type: Some(pg_row_try_string(&row, 1)), + signature: None, }); } } @@ -1990,7 +1994,8 @@ fn postgres_completion_tables_sql() -> &'static str { fn postgres_completion_routines_sql() -> &'static str { "SELECT p.proname, n.nspname, CASE p.prokind WHEN 'p' THEN 'PROCEDURE' ELSE 'FUNCTION' END, \ - obj_description(p.oid) AS routine_comment, COALESCE(pg_get_function_result(p.oid), '') AS data_type \ + obj_description(p.oid) AS routine_comment, COALESCE(pg_get_function_result(p.oid), '') AS data_type, \ + pg_get_function_identity_arguments(p.oid) AS signature \ FROM pg_catalog.pg_proc p \ JOIN pg_catalog.pg_namespace n ON n.oid = p.pronamespace \ WHERE n.nspname = $1 AND p.prokind::text = ANY($3::text[]) \ @@ -5651,6 +5656,7 @@ mod tests { 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("pg_get_function_identity_arguments(p.oid) AS signature")); 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/sqlite.rs b/crates/dbx-core/src/db/sqlite.rs index 517009ba0..02f1c0034 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -1376,6 +1376,7 @@ fn sqlite_completion_schemas( parent_name: None, comment: None, data_type: None, + signature: None, }) .collect()) } @@ -1429,6 +1430,7 @@ fn sqlite_completion_tables( parent_name: None, comment: None, data_type: None, + signature: None, }) }) .map_err(|e| e.to_string())?; @@ -1467,6 +1469,7 @@ fn sqlite_completion_columns( parent_name: Some(table.to_string()), comment: None, data_type: Some(data_type), + signature: None, }); if candidates.len() >= limit { break; diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index ae00e0235..d2f486c50 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -1338,6 +1338,7 @@ pub async fn completion_assistant_search( parent_name: row.get::<&str, _>(4).map(str::to_string), comment: row.get::<&str, _>(5).filter(|s: &&str| !s.is_empty()).map(|s| (*s).to_string()), data_type: row.get::<&str, _>(6).map(str::to_string), + signature: None, } }) .collect::>(); diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 5ce12b490..6cf815fc8 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -591,6 +591,7 @@ fn duckdb_completion_schemas( parent_name: None, comment: None, data_type: None, + signature: None, }) }) .map_err(|e| e.to_string())?; @@ -641,6 +642,7 @@ fn duckdb_completion_tables( parent_name: None, comment: None, data_type: None, + signature: None, }) }) .map_err(|e| e.to_string())?; @@ -686,6 +688,7 @@ fn duckdb_completion_columns( parent_name: Some(table.to_string()), comment: None, data_type: Some(row.get(1)?), + signature: None, }) }) .map_err(|e| e.to_string())?; @@ -4327,6 +4330,7 @@ async fn completion_assistant_fallback_core( parent_name: None, comment: None, data_type: None, + signature: None, }); } if candidates.len() >= limit { @@ -4364,6 +4368,7 @@ async fn completion_assistant_fallback_core( parent_name: table.parent_name, comment: table.comment, data_type: None, + signature: None, }); if candidates.len() >= limit { return Ok(db::CompletionAssistantResponse { candidates, incomplete: true, fallback_used: true }); @@ -4385,6 +4390,7 @@ async fn completion_assistant_fallback_core( parent_name: Some(table.to_string()), comment: column.comment, data_type: Some(column.data_type), + signature: None, }); } if candidates.len() >= limit { diff --git a/crates/dbx-core/src/types.rs b/crates/dbx-core/src/types.rs index 296e82782..96c599612 100644 --- a/crates/dbx-core/src/types.rs +++ b/crates/dbx-core/src/types.rs @@ -215,6 +215,7 @@ pub struct CompletionAssistantCandidate { pub parent_name: Option, pub comment: Option, pub data_type: Option, + pub signature: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/packages/app-tests/sqlCompletionRoutineParameters.test.ts b/packages/app-tests/sqlCompletionRoutineParameters.test.ts new file mode 100644 index 000000000..107ebc049 --- /dev/null +++ b/packages/app-tests/sqlCompletionRoutineParameters.test.ts @@ -0,0 +1,85 @@ +import { strict as assert } from "node:assert"; +import { snippetCompletion } from "@codemirror/autocomplete"; +import { EditorState, type Transaction } from "@codemirror/state"; +import { test } from "vitest"; +import { buildSqlCompletionItems } from "../../apps/desktop/src/lib/sql/sqlCompletion.ts"; + +test("inserts PostgreSQL routine parameters and preserves overloaded functions", () => { + const sql = "SELECT st_astext"; + const items = buildSqlCompletionItems(sql, sql.length, { + tables: [], + columnsByTable: new Map(), + objects: [ + { name: "st_astext", schema: "public", type: "function", signature: "geometry", dataType: "text" }, + { name: "st_astext", schema: "public", type: "function", signature: "geography, integer", dataType: "text" }, + ], + databaseType: "postgres", + currentSchema: "public", + }); + + const functions = items.filter((item) => item.label === "st_astext" && item.type === "function"); + assert.deepEqual( + functions.map((item) => ({ apply: item.apply, detail: item.detail })), + [ + { apply: "st_astext(${1:geometry})", detail: "function in public (geometry) [text]" }, + { apply: "st_astext(${1:geography}, ${2:integer})", detail: "function in public (geography, integer) [text]" }, + ], + ); +}); + +test("keeps repeated routine parameter types as independent snippet fields", () => { + const sql = "SELECT add_pair"; + const item = buildSqlCompletionItems(sql, sql.length, { + tables: [], + columnsByTable: new Map(), + objects: [{ name: "add_pair", schema: "public", type: "function", signature: "integer, integer", dataType: "integer" }], + databaseType: "postgres", + currentSchema: "public", + }).find((candidate) => candidate.label === "add_pair"); + + assert.equal(item?.apply, "add_pair(${1:integer}, ${2:integer})"); + assert.equal(item?.detail, "function in public (integer, integer) [integer]"); + + let state = EditorState.create({ + doc: sql, + selection: { anchor: sql.length }, + extensions: [EditorState.allowMultipleSelections.of(true)], + }); + const editor = { + get state() { + return state; + }, + dispatch(transaction: Transaction) { + state = transaction.state; + }, + }; + const completion = snippetCompletion(item?.apply ?? "", { label: item?.label ?? "" }); + + assert.equal(typeof completion.apply, "function"); + if (typeof completion.apply !== "function") return; + completion.apply(editor as never, completion, "SELECT ".length, sql.length); + editor.dispatch(state.update(state.replaceSelection("first_value"))); + + assert.equal(state.doc.toString(), "SELECT add_pair(first_value, integer)"); +}); + +test("splits routine parameters only on top-level commas", () => { + const sql = "SELECT transform_value"; + const items = buildSqlCompletionItems(sql, sql.length, { + tables: [], + columnsByTable: new Map(), + objects: [{ name: "transform_value", schema: "public", type: "function", signature: 'numeric(10, 2), "custom,schema"."value,type"[], text[]' }], + databaseType: "postgres", + currentSchema: "public", + }); + + assert.equal(items.find((item) => item.label === "transform_value")?.apply, 'transform_value(${1:numeric(10, 2)}, ${2:"custom,schema"."value,type"[]}, ${3:text[]})'); +}); + +test("keeps empty and unavailable routine signatures as empty parentheses", () => { + const sql = "SELECT current_marker"; + const baseInput = { tables: [], columnsByTable: new Map(), databaseType: "postgres" as const, currentSchema: "public" }; + + assert.equal(buildSqlCompletionItems(sql, sql.length, { ...baseInput, objects: [{ name: "current_marker", schema: "public", type: "function", signature: "" }] }).find((item) => item.label === "current_marker")?.apply, "current_marker()"); + assert.equal(buildSqlCompletionItems(sql, sql.length, { ...baseInput, objects: [{ name: "current_marker", schema: "public", type: "function" }] }).find((item) => item.label === "current_marker")?.apply, "current_marker()"); +});