xiaoxue_admin
/
dbx
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(completion): include PostgreSQL function parameters

This commit is contained in:
zipg 2026-07-28 21:25:30 +08:00 committed by GitHub
parent a0d4a31467
commit 89c69ebe65
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 166 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
apps/desktop/src/components/editor/QueryEditor.vue
View File

@ -3187,9 +3187,9 @@ function isTableNameCompletionContext(completionContext: ReturnType<typeof getSq
function mergeCompletionObjects(existing: SqlCompletionObject[], incoming: SqlCompletionObject[]) {
const merged = [...existing];
const indexes = new Map(existing.map((object, index) => [`${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) => {

53
apps/desktop/src/lib/sql/sqlCompletion.ts
View File

@ -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);

3
apps/desktop/src/stores/connectionStore.ts
View File

@ -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<string>();
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);

1
apps/desktop/src/types/database.ts
View File

@ -106,6 +106,7 @@ export interface CompletionAssistantCandidate {
parent_name?: string | null;
comment?: string | null;
data_type?: string | null;
signature?: string | null;
}
export interface CompletionAssistantResponse {

4
crates/dbx-core/src/db/mysql.rs
View File

@ -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,
});
}
}

8
crates/dbx-core/src/db/postgres.rs
View File

@ -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<String>>(5).ok().flatten(),
comment: row.try_get::<_, Option<String>>(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<String>>(3).ok().flatten(),
data_type: row.try_get::<_, Option<String>>(4).ok().flatten(),
signature: row.try_get::<_, Option<String>>(5).ok().flatten(),
});
}
}
@ -1962,6 +1965,7 @@ pub async fn completion_assistant_search(
parent_name: Some(table.to_string()),
comment: row.try_get::<_, Option<String>>(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)"));

3
crates/dbx-core/src/db/sqlite.rs
View File

@ -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;

1
crates/dbx-core/src/db/sqlserver.rs
View File

@ -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::<Vec<_>>();

6
crates/dbx-core/src/schema.rs
View File

@ -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 {

1
crates/dbx-core/src/types.rs
View File

@ -215,6 +215,7 @@ pub struct CompletionAssistantCandidate {
pub parent_name: Option<String>,
pub comment: Option<String>,
pub data_type: Option<String>,
pub signature: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]

85
packages/app-tests/sqlCompletionRoutineParameters.test.ts vendored Normal file
View File

@ -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()");
});