From cc74fabad7a6656f5010a9e4f170966b8fdeae08 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Wed, 3 Jun 2026 11:27:31 +0800 Subject: [PATCH] feat(completion): suggest database routines --- .../src/components/editor/QueryEditor.vue | 55 ++++++++- apps/desktop/src/lib/api.ts | 1 + apps/desktop/src/lib/http.ts | 8 ++ apps/desktop/src/lib/sqlCompletion.ts | 96 +++++++++++++++- apps/desktop/src/lib/tauri.ts | 8 ++ apps/desktop/src/stores/connectionStore.ts | 108 +++++++++++++++++- crates/dbx-core/src/db/mysql.rs | 54 ++++++++- crates/dbx-core/src/schema.rs | 74 ++++++++++++ crates/dbx-web/src/main.rs | 1 + crates/dbx-web/src/routes/schema.rs | 12 ++ packages/app-tests/sqlCompletion.test.ts | 46 ++++++++ src-tauri/src/commands/schema.rs | 10 ++ src-tauri/src/lib.rs | 1 + 13 files changed, 463 insertions(+), 11 deletions(-) diff --git a/apps/desktop/src/components/editor/QueryEditor.vue b/apps/desktop/src/components/editor/QueryEditor.vue index c63df89d3..0cc39206a 100644 --- a/apps/desktop/src/components/editor/QueryEditor.vue +++ b/apps/desktop/src/components/editor/QueryEditor.vue @@ -57,7 +57,7 @@ import { shouldRunSqlSemanticDiagnostics, type SqlSemanticDiagnostic, } from "@/lib/sqlSemanticDiagnostics"; -import type { SqlCompletionColumn, SqlCompletionForeignKey } from "@/lib/sqlCompletion"; +import type { SqlCompletionColumn, SqlCompletionForeignKey, SqlCompletionObject } from "@/lib/sqlCompletion"; import type { DatabaseType, ForeignKeyInfo, @@ -201,6 +201,7 @@ function editorThemeAppearance() { // Completion cache let cachedTables: Array<{ name: string; schema?: string; type?: "table" | "view" }> = []; +let cachedCompletionObjects: SqlCompletionObject[] = []; // Persistent column cache keyed by "schema.table" or "table" const cachedColumnsByTable = new Map(); const cachedForeignKeysByTable = new Map(); @@ -892,6 +893,7 @@ async function provideSqlCompletions( if (!hasDatabase) { const items = buildSqlCompletionItemsFromContext(completionContext, { tables: [], + objects: [], columnsByTable: new Map(), schemas: [], translations: completionTranslations.value, @@ -903,6 +905,8 @@ async function provideSqlCompletions( const needsAsyncData = completionContext.suggestTables || + completionContext.suggestRoutines || + completionContext.exclusiveRoutineSuggestions || !!completionContext.qualifier || !!completionContext.insertTable || completionContext.exclusiveColumnSuggestions || @@ -911,6 +915,7 @@ async function provideSqlCompletions( if (!needsAsyncData) { const items = buildSqlCompletionItemsFromContext(completionContext, { tables: [], + objects: [], columnsByTable: new Map(), schemas: [], translations: completionTranslations.value, @@ -990,6 +995,35 @@ async function performAsyncCompletionWithResult( : cachedTables; if (epoch !== completionEpoch) return null; + const shouldLoadObjects = + completionContext.suggestRoutines || + completionContext.exclusiveRoutineSuggestions || + (!!completionContext.qualifier && !completionContext.exclusiveColumnSuggestions); + let completionObjects = shouldLoadObjects + ? await connectionStore.listCompletionObjects( + props.connectionId!, + props.database!, + completionContext.qualifier || completionContext.prefix, + MAX_COMPLETION_TABLES, + ) + : cachedCompletionObjects; + if (epoch !== completionEpoch) return null; + + if (completionContext.qualifier && completionObjects.length === 0) { + const schemaObjects = await connectionStore.listCompletionObjects( + props.connectionId!, + props.database!, + completionContext.prefix, + MAX_COMPLETION_TABLES, + completionContext.qualifier, + ); + if (schemaObjects.length > 0) { + completionObjects = schemaObjects; + } + if (epoch !== completionEpoch) return null; + } + cachedCompletionObjects = mergeCompletionObjects(cachedCompletionObjects, completionObjects); + // Fetch schemas for schema completion let schemaNames: string[] = []; if (completionContext.suggestTables && !completionContext.qualifier && !completionContext.insertTable) { @@ -1142,6 +1176,7 @@ async function performAsyncCompletionWithResult( const items = buildSqlCompletionItemsFromContext(effectiveContext, { tables, + objects: completionObjects, columnsByTable, foreignKeysByTable, schemas: schemaNames, @@ -1161,8 +1196,25 @@ function isReferencedTableQualifier(completionContext: ReturnType + `${object.type}:${object.schema ?? ""}:${object.name}:${object.parentName ?? ""}`.toLowerCase(), + ), + ); + for (const object of incoming) { + const key = `${object.type}:${object.schema ?? ""}:${object.name}:${object.parentName ?? ""}`.toLowerCase(); + if (seen.has(key)) continue; + seen.add(key); + merged.push(object); + } + return merged; +} + async function refreshCompletionCache() { cachedTables = []; + cachedCompletionObjects = []; cachedColumnsByTable.clear(); cachedForeignKeysByTable.clear(); } @@ -1566,6 +1618,7 @@ onMounted(async () => { registerTableReferenceDropListener(); cachedTables = []; + cachedCompletionObjects = []; scheduleSemanticDiagnostics(); }); diff --git a/apps/desktop/src/lib/api.ts b/apps/desktop/src/lib/api.ts index ff84bb3dd..a77715a8a 100644 --- a/apps/desktop/src/lib/api.ts +++ b/apps/desktop/src/lib/api.ts @@ -74,6 +74,7 @@ export const deleteSchemaCachePrefix = forward("deleteSchemaCachePrefix"); export const listSchemas = forward("listSchemas"); export const listTables = forward("listTables"); export const listObjects = forward("listObjects"); +export const listCompletionObjects = forward("listCompletionObjects"); export const getObjectSource = forward("getObjectSource"); export const getColumns = forward("getColumns"); export const listIndexes = forward("listIndexes"); diff --git a/apps/desktop/src/lib/http.ts b/apps/desktop/src/lib/http.ts index d8d42559c..2a94707b9 100644 --- a/apps/desktop/src/lib/http.ts +++ b/apps/desktop/src/lib/http.ts @@ -372,6 +372,14 @@ export async function listObjects(connectionId: string, database: string, schema return get(`/api/schema/objects?${qs({ connection_id: connectionId, database, schema })}`); } +export async function listCompletionObjects( + connectionId: string, + database: string, + schema: string, +): Promise { + return get(`/api/schema/completion-objects?${qs({ connection_id: connectionId, database, schema })}`); +} + export async function getObjectSource( connectionId: string, database: string, diff --git a/apps/desktop/src/lib/sqlCompletion.ts b/apps/desktop/src/lib/sqlCompletion.ts index 5b1234a16..81c77250c 100644 --- a/apps/desktop/src/lib/sqlCompletion.ts +++ b/apps/desktop/src/lib/sqlCompletion.ts @@ -600,6 +600,14 @@ export interface SqlCompletionTable { type?: "table" | "view"; } +export interface SqlCompletionObject { + name: string; + schema?: string; + type: "procedure" | "function" | "trigger"; + parentSchema?: string; + parentName?: string; +} + export interface SqlCompletionColumn { name: string; table: string; @@ -639,9 +647,11 @@ export interface SqlCompletionContext { suggestTables: boolean; suggestColumns: boolean; suggestKeywords: boolean; + suggestRoutines: boolean; suggestJoinConditions: boolean; exclusiveTableSuggestions: boolean; exclusiveColumnSuggestions: boolean; + exclusiveRoutineSuggestions: boolean; prioritizeSelectAliases: boolean; selectAliases: string[]; referencedTables: SqlCompletionReferencedTable[]; @@ -678,6 +688,7 @@ export function buildSqlCompletionItems( cursor: number, input: { tables: SqlCompletionTable[]; + objects?: SqlCompletionObject[]; columnsByTable: Map; foreignKeysByTable?: Map; schemas?: string[]; @@ -693,6 +704,7 @@ export function buildSqlCompletionItemsFromContext( context: SqlCompletionContext, input: { tables: SqlCompletionTable[]; + objects?: SqlCompletionObject[]; columnsByTable: Map; foreignKeysByTable?: Map; schemas?: string[]; @@ -705,29 +717,44 @@ export function buildSqlCompletionItemsFromContext( const t = input.translations; const dialect = input.dialect; - if (!context.exclusiveTableSuggestions && !context.exclusiveColumnSuggestions) { + if ( + !context.exclusiveTableSuggestions && + !context.exclusiveColumnSuggestions && + !context.exclusiveRoutineSuggestions + ) { items.push(...buildSnippetItems(context.prefix, input.snippets ?? DEFAULT_SQL_SNIPPETS)); items.push(...buildFunctionSnippetItems(context.prefix, getFunctionDescriptions(t))); } - if (!context.exclusiveTableSuggestions && !context.exclusiveColumnSuggestions && context.prioritizeSelectAliases) { + if ( + !context.exclusiveTableSuggestions && + !context.exclusiveColumnSuggestions && + !context.exclusiveRoutineSuggestions && + context.prioritizeSelectAliases + ) { items.push(...buildSelectAliasItems(context)); } if ( !context.exclusiveTableSuggestions && !context.exclusiveColumnSuggestions && + !context.exclusiveRoutineSuggestions && context.isGroupBy && context.nonAggregatedSelectColumns.length > 0 ) { items.push(...buildNonAggregatedColumnItems(context, input.columnsByTable, dialect)); } - if (!context.exclusiveTableSuggestions && !context.exclusiveColumnSuggestions && context.suggestJoinConditions) { + if ( + !context.exclusiveTableSuggestions && + !context.exclusiveColumnSuggestions && + !context.exclusiveRoutineSuggestions && + context.suggestJoinConditions + ) { items.push(...buildJoinConditionItems(context, input.columnsByTable, input.foreignKeysByTable, dialect)); } - if (context.suggestKeywords) { + if (context.suggestKeywords && !context.exclusiveRoutineSuggestions) { items.push(...buildKeywordItems(context.prefix, context)); } @@ -747,6 +774,10 @@ export function buildSqlCompletionItemsFromContext( } } + if (context.suggestRoutines || context.exclusiveRoutineSuggestions) { + items.push(...buildObjectItems(context, input.objects ?? [], dialect)); + } + // Type-aware value hints after comparison operator if (context.comparisonLeftColumn && context.suggestKeywords) { items.push(...buildComparisonValueItems(context, input.columnsByTable, t)); @@ -765,9 +796,17 @@ export function shouldAutoOpenSqlCompletion(sql: string, cursor: number): boolea const previousChar = sql[cursor - 1]; if (!previousChar) return false; if (/\bon\s+$/i.test(sql.slice(0, cursor))) return true; + if (/\bcall\s+(?:[A-Za-z_][\w$]*\.)?$/i.test(sql.slice(0, cursor))) return true; if (/[,;()[\]]/.test(previousChar)) return false; const context = getSqlCompletionContext(sql, cursor); - if (context.exclusiveTableSuggestions || context.exclusiveColumnSuggestions || context.suggestTables) return true; + if ( + context.exclusiveTableSuggestions || + context.exclusiveColumnSuggestions || + context.exclusiveRoutineSuggestions || + context.suggestTables + ) { + return true; + } return /[\w$.]/.test(previousChar); } @@ -859,6 +898,13 @@ function detectStatementKind(previousStatements: string): SqlStatementKind { return kindMap[firstWord] ?? "unknown"; } +function isCallRoutineContext(beforeToken: string): boolean { + return ( + /\bcall\s+(?:[A-Za-z_][\w$]*\.)?$/i.test(beforeToken) || + /\bcall\s+(?:[A-Za-z_][\w$]*\.)?[A-Za-z_][\w$]*$/i.test(beforeToken) + ); +} + export function getSqlCompletionContext(sql: string, cursor: number): SqlCompletionContext { // Extract the full statement at cursor position for referenced tables const fullStatement = extractStatementAt(sql, cursor); @@ -917,6 +963,7 @@ export function getSqlCompletionContext(sql: string, cursor: number): SqlComplet const inColumnContext = isInColumnContext(beforeCursor) || !!insertInfo; const inJoinConditionContext = isInJoinConditionContext(beforeCursor); const prioritizeSelectAliases = isInOrderOrGroupByContext(beforeCursor); + const inCallRoutineContext = isCallRoutineContext(beforeCursor); const statementKind = detectStatementKind(beforeCursor || fullStatement); @@ -925,10 +972,14 @@ export function getSqlCompletionContext(sql: string, cursor: number): SqlComplet qualifier: insertInfo ? undefined : qualifier, suggestTables: insertInfo ? false : afterTableTrigger, suggestColumns: !!qualifier || (inColumnContext && referencedTables.length > 0), - suggestKeywords: !exclusiveTableSuggestions && !exclusiveColumnSuggestions && !insertInfo, + suggestKeywords: !exclusiveTableSuggestions && !exclusiveColumnSuggestions && !insertInfo && !inCallRoutineContext, + suggestRoutines: + inCallRoutineContext || + (!exclusiveTableSuggestions && !exclusiveColumnSuggestions && !insertInfo && prefix.length >= 2), suggestJoinConditions: insertInfo ? false : inJoinConditionContext && referencedTables.length >= 2, exclusiveTableSuggestions: insertInfo ? false : exclusiveTableSuggestions, exclusiveColumnSuggestions: exclusiveColumnSuggestions || !!insertInfo, + exclusiveRoutineSuggestions: inCallRoutineContext, prioritizeSelectAliases: insertInfo ? false : prioritizeSelectAliases, selectAliases: prioritizeSelectAliases ? extractSelectAliases(fullStatement) : [], referencedTables, @@ -1547,6 +1598,39 @@ function buildSchemaItems( })); } +function buildObjectItems( + context: SqlCompletionContext, + objects: SqlCompletionObject[], + dialect?: "mysql" | "postgres" | "sqlserver", +): SqlCompletionItem[] { + const onlyProcedures = context.exclusiveRoutineSuggestions; + return objects + .filter((object) => (!onlyProcedures || object.type === "procedure") && matchesPrefix(object.name, context.prefix)) + .map((object) => { + const applyName = + context.qualifier && object.schema?.toLowerCase() === context.qualifier.toLowerCase() + ? quoteSqlIdentifier(object.name, dialect) + : object.schema + ? `${quoteSqlIdentifier(object.schema, dialect)}.${quoteSqlIdentifier(object.name, dialect)}` + : quoteSqlIdentifier(object.name, dialect); + const detail = + object.type === "trigger" && object.parentName + ? `trigger on ${object.parentName}` + : object.schema + ? `${object.type} in ${object.schema}` + : object.type; + return { + label: object.name, + type: "function" as const, + detail, + apply: object.type === "trigger" ? applyName : `${applyName}()`, + boost: computeBoost(object.name, context.prefix) + (object.type === "procedure" ? 1800 : 900), + }; + }) + .sort(compareCompletionItems) + .slice(0, MAX_TABLE_COMPLETION_ITEMS); +} + function buildStarExpansionItem( columnsByTable: Map, t?: SqlCompletionTranslations, diff --git a/apps/desktop/src/lib/tauri.ts b/apps/desktop/src/lib/tauri.ts index 31ab08ffa..9769d3a5f 100644 --- a/apps/desktop/src/lib/tauri.ts +++ b/apps/desktop/src/lib/tauri.ts @@ -407,6 +407,14 @@ export async function listObjects(connectionId: string, database: string, schema return invoke("list_objects", { connectionId, database, schema }); } +export async function listCompletionObjects( + connectionId: string, + database: string, + schema: string, +): Promise { + return invoke("list_completion_objects", { connectionId, database, schema }); +} + export async function getObjectSource( connectionId: string, database: string, diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index a9cba9774..7d459814c 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -1,7 +1,7 @@ import { defineStore } from "pinia"; import { uuid } from "@/lib/utils"; import { ref, computed, watch } from "vue"; -import type { ColumnInfo, ConnectionConfig, SidebarLayout, TreeNode } from "@/types/database"; +import type { ColumnInfo, ConnectionConfig, ObjectInfo, SidebarLayout, TreeNode } from "@/types/database"; import { applyPinnedTreeNodeState, orderPinnedFirst } from "@/lib/pinnedItems"; import { reconcileLayout, @@ -17,7 +17,7 @@ import { reorderEntry as reorderEntryOp, type DropPosition, } from "@/lib/sidebarLayout"; -import type { SqlCompletionColumn, SqlCompletionTable } from "@/lib/sqlCompletion"; +import type { SqlCompletionColumn, SqlCompletionObject, SqlCompletionTable } from "@/lib/sqlCompletion"; import * as api from "@/lib/api"; import { isTauriRuntime } from "@/lib/tauriRuntime"; import { isSchemaAware, usesTreeSchemaMode } from "@/lib/databaseCapabilities"; @@ -98,6 +98,7 @@ export const useConnectionStore = defineStore("connection", () => { const editingConnectionId = ref(null); const newConnectionGroupId = ref(null); const completionTablesCache = ref>({}); + const completionObjectsCache = ref>({}); const completionColumnsCache = ref>({}); const schemaListCache = ref>({}); const transferSource = ref<{ connectionId: string; database: string } | null>(null); @@ -552,6 +553,9 @@ export const useConnectionStore = defineStore("connection", () => { for (const key of Object.keys(completionTablesCache.value)) { if (key === exactCacheKey || key.startsWith(cachePrefix)) delete completionTablesCache.value[key]; } + for (const key of Object.keys(completionObjectsCache.value)) { + if (key === exactCacheKey || key.startsWith(cachePrefix)) delete completionObjectsCache.value[key]; + } for (const key of Object.keys(completionColumnsCache.value)) { if (key === exactCacheKey || key.startsWith(cachePrefix)) delete completionColumnsCache.value[key]; } @@ -1579,6 +1583,105 @@ export const useConnectionStore = defineStore("connection", () => { return deduped; } + async function listCompletionObjects( + connectionId: string, + database: string, + filter = "", + limit?: number, + schema?: string, + ): Promise { + const normalizedFilter = filter.trim().toLowerCase(); + const cacheKey = `${connectionId}:${database}:${schema ?? ""}`; + if (!completionObjectsCache.value[cacheKey]) { + await ensureConnected(connectionId); + 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), + ); + evictOldestCacheEntries(completionObjectsCache.value, COMPLETION_CACHE_MAX); + } + + const objects = completionObjectsCache.value[cacheKey]; + const filtered = normalizedFilter + ? objects.filter((object) => fuzzyCompletionObjectMatch(object, normalizedFilter)) + : objects; + return typeof limit === "number" ? filtered.slice(0, expandedCompletionLimit(limit)) : filtered; + } + + async function listSchemaAwareCompletionObjects( + connectionId: string, + database: string, + schema?: string, + ): Promise { + const schemas = schema ? [schema] : await listCompletionSchemas(connectionId, database); + const batchSize = 5; + const results: ObjectInfo[] = []; + for (let i = 0; i < schemas.length; i += batchSize) { + const batch = schemas.slice(i, i + batchSize); + const groups = await Promise.all( + batch.map(async (s) => { + try { + return await api.listCompletionObjects(connectionId, database, s); + } catch { + return [] as ObjectInfo[]; + } + }), + ); + for (const group of groups) results.push(...group); + } + return results; + } + + function toSqlCompletionObject(object: ObjectInfo): SqlCompletionObject | null { + const objectType = object.object_type.toUpperCase(); + const type = objectType.includes("PROCEDURE") + ? "procedure" + : objectType.includes("FUNCTION") + ? "function" + : objectType.includes("TRIGGER") + ? "trigger" + : null; + if (!type) return null; + return { + name: object.name, + schema: object.schema ?? undefined, + type, + parentSchema: object.parent_schema ?? undefined, + parentName: object.parent_name ?? undefined, + }; + } + + function fuzzyCompletionObjectMatch(object: SqlCompletionObject, filter: string): boolean { + return fuzzyTextMatch(object.name, filter) || (!!object.schema && fuzzyTextMatch(object.schema, filter)); + } + + function fuzzyTextMatch(value: string, filter: string): boolean { + if (!filter) return true; + const text = value.toLowerCase(); + if (text.includes(filter)) return true; + let index = 0; + for (const ch of filter) { + index = text.indexOf(ch, index); + if (index < 0) return false; + index++; + } + return true; + } + + function dedupeCompletionObjects(objects: SqlCompletionObject[]): SqlCompletionObject[] { + const seen = new Set(); + const deduped: SqlCompletionObject[] = []; + for (const object of objects) { + const key = `${object.type}:${object.schema ?? ""}:${object.name}:${object.parentName ?? ""}`.toLowerCase(); + if (seen.has(key)) continue; + seen.add(key); + deduped.push(object); + } + return deduped; + } + async function listCompletionColumns( connectionId: string, database: string, @@ -1987,6 +2090,7 @@ export const useConnectionStore = defineStore("connection", () => { loadForeignKeys, loadTriggers, listCompletionTables, + listCompletionObjects, listCompletionColumns, listCompletionSchemas, exportConnectionsToFile, diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index ae8981660..2a17413fe 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -744,6 +744,7 @@ fn list_tables_objects_sql(database: &str) -> String { TABLE_COMMENT AS object_comment, \ CREATE_TIME AS created_at, \ UPDATE_TIME AS updated_at, \ + NULL AS parent_schema, NULL AS parent_name, \ CASE WHEN TABLE_TYPE = 'VIEW' THEN 1 ELSE 0 END AS sort_order \ FROM information_schema.TABLES \ WHERE TABLE_SCHEMA = {db} \ @@ -756,6 +757,7 @@ fn list_routines_sql(database: &str) -> String { format!( "SELECT ROUTINE_NAME AS object_name, ROUTINE_TYPE AS object_type, NULL AS object_comment, \ NULL AS created_at, NULL AS updated_at, \ + NULL AS parent_schema, NULL AS parent_name, \ CASE WHEN ROUTINE_TYPE = 'PROCEDURE' THEN 2 ELSE 3 END AS sort_order \ FROM information_schema.ROUTINES \ WHERE ROUTINE_SCHEMA = {db} AND ROUTINE_TYPE IN ('PROCEDURE', 'FUNCTION') \ @@ -764,6 +766,19 @@ fn list_routines_sql(database: &str) -> String { ) } +fn list_completion_triggers_sql(database: &str) -> String { + format!( + "SELECT TRIGGER_NAME AS object_name, 'TRIGGER' AS object_type, NULL AS object_comment, \ + CREATED AS created_at, NULL AS updated_at, \ + TRIGGER_SCHEMA AS parent_schema, EVENT_OBJECT_TABLE AS parent_name, \ + 4 AS sort_order \ + FROM information_schema.TRIGGERS \ + WHERE TRIGGER_SCHEMA = {db} \ + ORDER BY object_name", + db = quote_value(database), + ) +} + fn row_to_object(row: &mysql_async::Row, database: &str) -> ObjectInfo { ObjectInfo { name: get_str_by_name(row, "object_name"), @@ -772,8 +787,8 @@ fn row_to_object(row: &mysql_async::Row, database: &str) -> ObjectInfo { comment: get_opt_str(row, "object_comment").filter(|s| !s.is_empty()), created_at: get_opt_str(row, "created_at"), updated_at: get_opt_str(row, "updated_at"), - parent_schema: None, - parent_name: None, + parent_schema: get_opt_str(row, "parent_schema"), + parent_name: get_opt_str(row, "parent_name"), } } @@ -806,6 +821,31 @@ pub async fn list_objects(pool: &MySqlPool, database: &str) -> Result Result, String> { + let mut conn = pool.get_conn().await.map_err(|e| e.to_string())?; + let mut objects = Vec::new(); + + let routines_sql = list_routines_sql(database); + match conn.query_iter(&routines_sql).await { + Ok(result) => match result.collect_and_drop::().await { + Ok(rows) => objects.extend(rows.iter().map(|row| row_to_object(row, database))), + Err(e) => log::warn!("Skipping routines for completion in database `{}`: {}", database, e), + }, + Err(e) => log::warn!("Skipping routines for completion in database `{}`: {}", database, e), + } + + let triggers_sql = list_completion_triggers_sql(database); + match conn.query_iter(&triggers_sql).await { + Ok(result) => match result.collect_and_drop::().await { + Ok(rows) => objects.extend(rows.iter().map(|row| row_to_object(row, database))), + Err(e) => log::warn!("Skipping triggers for completion in database `{}`: {}", database, e), + }, + Err(e) => log::warn!("Skipping triggers for completion in database `{}`: {}", database, e), + } + + Ok(objects) +} + fn columns_sql(database: &str, table: &str) -> String { format!( "SELECT c.COLUMN_NAME, c.COLUMN_TYPE, c.IS_NULLABLE, c.COLUMN_DEFAULT, c.EXTRA, c.COLUMN_COMMENT, \ @@ -1289,6 +1329,16 @@ mod tests { assert!(!sql.contains("CREATED AS created_at")); } + #[test] + fn mysql_completion_triggers_sql_lists_database_triggers() { + let sql = list_completion_triggers_sql("app"); + + assert!(sql.contains("information_schema.TRIGGERS")); + assert!(sql.contains("'TRIGGER' AS object_type")); + assert!(sql.contains("EVENT_OBJECT_TABLE AS parent_name")); + assert!(sql.contains("TRIGGER_SCHEMA = 'app'")); + } + #[test] fn mysql_columns_sql_joins_key_column_usage_for_primary_keys() { let sql = columns_sql("app", "users"); diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 9c0c4fb25..ef3824875 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -519,6 +519,18 @@ pub async fn list_objects_core( .await } +pub async fn list_completion_objects_core( + state: &AppState, + connection_id: &str, + database: &str, + schema: &str, +) -> Result, String> { + retry_metadata_connection(state, connection_id, Some(database), || { + list_completion_objects_once(state, connection_id, database, schema) + }) + .await +} + async fn list_objects_once( state: &AppState, connection_id: &str, @@ -608,6 +620,68 @@ async fn list_objects_once( } } +async fn list_completion_objects_once( + state: &AppState, + connection_id: &str, + database: &str, + schema: &str, +) -> Result, String> { + let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; + let db_config = connection_config(state, connection_id).await; + + let connections = state.connections.read().await; + if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { + let config = config.clone(); + let session = session.clone(); + drop(connections); + return session + .invoke::>( + "listObjects", + serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema }), + ) + .await + .map(filter_completion_objects); + } + if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { + let is_oracle = db_config.as_ref().is_some_and(|config| config.db_type == DatabaseType::Oracle); + drop(connections); + let objects = if is_oracle { + oracle_agent_list_objects(client, database, schema).await? + } else { + let mut client = client.lock().await; + client.list_objects(database, schema).await? + }; + return Ok(filter_completion_objects(objects)); + } + + let pool = connections.get(&pool_key).ok_or("Pool not found")?; + match pool { + PoolKind::Mysql(p, mode) if *mode != MysqlMode::OceanBaseOracle => { + db::mysql::list_completion_objects(p, database).await + } + PoolKind::Mysql(p, mode) if *mode == MysqlMode::OceanBaseOracle => { + db::ob_oracle::list_objects(p, schema).await.map(filter_completion_objects) + } + PoolKind::Postgres(p) => db::postgres::list_objects(p, schema).await.map(filter_completion_objects), + PoolKind::SqlServer(_) => { + drop(connections); + let objects = list_objects_once(state, connection_id, database, schema).await?; + Ok(filter_completion_objects(objects)) + } + _ => Ok(Vec::new()), + } +} + +fn filter_completion_objects(objects: Vec) -> Vec { + objects + .into_iter() + .filter(|object| { + let object_type = object.object_type.to_ascii_uppercase(); + object_type.contains("PROCEDURE") || object_type.contains("FUNCTION") || object_type.contains("TRIGGER") + }) + .collect() +} + async fn retry_metadata_connection( state: &AppState, connection_id: &str, diff --git a/crates/dbx-web/src/main.rs b/crates/dbx-web/src/main.rs index 866f3a54d..338da4426 100644 --- a/crates/dbx-web/src/main.rs +++ b/crates/dbx-web/src/main.rs @@ -120,6 +120,7 @@ async fn main() { .route("/schema/schemas", get(routes::schema::list_schemas)) .route("/schema/tables", get(routes::schema::list_tables)) .route("/schema/objects", get(routes::schema::list_objects)) + .route("/schema/completion-objects", get(routes::schema::list_completion_objects)) .route("/schema/object-source", get(routes::schema::get_object_source)) .route("/schema/columns", get(routes::schema::list_columns)) .route("/schema/indexes", get(routes::schema::list_indexes)) diff --git a/crates/dbx-web/src/routes/schema.rs b/crates/dbx-web/src/routes/schema.rs index d4fabf094..6e165a96d 100644 --- a/crates/dbx-web/src/routes/schema.rs +++ b/crates/dbx-web/src/routes/schema.rs @@ -65,6 +65,18 @@ pub async fn list_objects( Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) } +pub async fn list_completion_objects( + State(state): State>, + Query(q): Query, +) -> Result, AppError> { + let database = q.database.as_deref().unwrap_or(""); + let schema = q.schema.as_deref().unwrap_or(""); + let result = dbx_core::schema::list_completion_objects_core(&state.app, &q.connection_id, database, schema) + .await + .map_err(AppError)?; + Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) +} + pub async fn get_object_source( State(state): State>, Query(q): Query, diff --git a/packages/app-tests/sqlCompletion.test.ts b/packages/app-tests/sqlCompletion.test.ts index 19382cd30..b7e74a97b 100644 --- a/packages/app-tests/sqlCompletion.test.ts +++ b/packages/app-tests/sqlCompletion.test.ts @@ -10,6 +10,7 @@ import { recordCompletionSelection, type SqlCompletionColumn, type SqlCompletionForeignKey, + type SqlCompletionObject, type SqlCompletionTable, } from "../../apps/desktop/src/lib/sqlCompletion.ts"; @@ -39,6 +40,12 @@ const columnsByTable = new Map([ ], ]); +const completionObjects: SqlCompletionObject[] = [ + { name: "refresh_user_stats", schema: "app", type: "procedure" }, + { name: "format_user_name", schema: "app", type: "function" }, + { name: "trg_users_audit", schema: "app", type: "trigger", parentName: "users" }, +]; + const postgresQuotedTables: SqlCompletionTable[] = [ { name: "article", schema: "public", type: "table" }, { name: "order_lines", schema: "public", type: "table" }, @@ -515,6 +522,45 @@ test("suggests DATE_FORMAT as parameter snippet", () => { assert.equal(snippet.apply, "DATE_FORMAT(${date}, ${format})"); }); +test("suggests stored procedures after CALL", () => { + const sql = "CALL rfs"; + const items = buildSqlCompletionItems(sql, sql.length, { + tables, + objects: completionObjects, + columnsByTable, + dialect: "mysql", + }); + + const procedure = items.find((item) => item.label === "refresh_user_stats"); + assert.ok(procedure); + assert.equal(procedure.type, "function"); + assert.equal(procedure.apply, "app.refresh_user_stats()"); + assert.equal( + items.some((item) => item.label === "format_user_name"), + false, + ); +}); + +test("suggests user functions and triggers with fuzzy matching", () => { + const sql = "select fun"; + const items = buildSqlCompletionItems(sql, sql.length, { + tables, + objects: completionObjects, + columnsByTable, + dialect: "mysql", + }); + + assert.ok(items.some((item) => item.label === "format_user_name" && item.detail === "function in app")); + + const triggerItems = buildSqlCompletionItems("drop trigger tua", "drop trigger tua".length, { + tables, + objects: completionObjects, + columnsByTable, + dialect: "mysql", + }); + assert.ok(triggerItems.some((item) => item.label === "trg_users_audit" && item.detail === "trigger on users")); +}); + test("matches alias qualifier case-insensitively", () => { const sql = "select O. from public.orders o"; const cursor = "select O.".length; diff --git a/src-tauri/src/commands/schema.rs b/src-tauri/src/commands/schema.rs index 8b90bfb36..312651b33 100644 --- a/src-tauri/src/commands/schema.rs +++ b/src-tauri/src/commands/schema.rs @@ -43,6 +43,16 @@ pub async fn list_objects( dbx_core::schema::list_objects_core(&state, &connection_id, &database, &schema).await } +#[tauri::command] +pub async fn list_completion_objects( + state: State<'_, Arc>, + connection_id: String, + database: String, + schema: String, +) -> Result, String> { + dbx_core::schema::list_completion_objects_core(&state, &connection_id, &database, &schema).await +} + #[tauri::command] pub async fn get_object_source( state: State<'_, Arc>, diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 138336fbb..5855c200b 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -380,6 +380,7 @@ pub fn run() { commands::schema::list_databases, commands::schema::list_tables, commands::schema::list_objects, + commands::schema::list_completion_objects, commands::schema::get_object_source, commands::schema::list_schemas, commands::schema::get_columns,