diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index bd1a49dfa4..c82c353fdf 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -74,7 +74,8 @@ pub async fn export_database_sql_core( // 3. List tables let all_tables = - crate::schema::list_tables_core(state, &request.connection_id, &request.database, &request.schema).await?; + crate::schema::list_tables_core(state, &request.connection_id, &request.database, &request.schema, None, None) + .await?; // 4. Create file let mut file = std::fs::File::create(&request.file_path).map_err(|e| format!("Failed to write file: {e}"))?; diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index e683afecaa..2ee93ac8f6 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -114,14 +114,25 @@ pub async fn list_schemas(client: &mut SqlServerClient) -> Result, S Ok(rows.iter().map(|row| row.get::<&str, _>(0).unwrap_or("").to_string()).collect()) } -pub async fn list_tables(client: &mut SqlServerClient, schema: &str) -> Result, String> { +pub async fn list_tables( + client: &mut SqlServerClient, + schema: &str, + filter: Option<&str>, + limit: Option, +) -> Result, String> { + let top = limit.map(|value| format!("TOP ({}) ", value.min(1000))).unwrap_or_default(); + let filter_clause = filter + .filter(|value| !value.trim().is_empty()) + .map(|value| format!(" AND o.name LIKE '%{}%' ESCAPE '\\' ", escape_like_literal(value.trim()))) + .unwrap_or_default(); let sql = format!( - "SELECT o.name, CASE WHEN o.type = 'V' THEN 'VIEW' ELSE 'BASE TABLE' END \ + "SELECT {top}o.name, CASE WHEN o.type = 'V' THEN 'VIEW' ELSE 'BASE TABLE' END \ FROM sys.objects o \ JOIN sys.schemas s ON s.schema_id = o.schema_id \ WHERE s.name = '{}' \ AND o.type IN ('U','V') \ AND o.is_ms_shipped = 0 \ + {filter_clause}\ ORDER BY o.name", schema.replace('\'', "''") ); @@ -137,6 +148,10 @@ pub async fn list_tables(client: &mut SqlServerClient, schema: &str) -> Result String { + value.replace('\\', "\\\\").replace('\'', "''").replace('%', "\\%").replace('_', "\\_").replace('[', "\\[") +} + pub async fn list_objects(client: &mut SqlServerClient, schema: &str) -> Result, String> { let sql = format!( "SELECT o.name, \ diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 3d0e4a8b4c..ce66576308 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -199,6 +199,8 @@ pub async fn list_tables_core( connection_id: &str, database: &str, schema: &str, + filter: Option<&str>, + limit: Option, ) -> Result, String> { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; @@ -237,7 +239,7 @@ pub async fn list_tables_core( if let Some(client) = extract_sqlserver(&connections, &pool_key) { drop(connections); let mut client = client.lock().await; - return db::sqlserver::list_tables(&mut client, schema).await; + return db::sqlserver::list_tables(&mut client, schema, filter, limit).await; } if let Some(client) = extract_agent(&connections, &pool_key) { drop(connections); @@ -252,17 +254,31 @@ pub async fn list_tables_core( match pool { PoolKind::Mysql(p, mode) => { if *mode == MysqlMode::OceanBaseOracle { - db::ob_oracle::list_tables(p, schema).await + db::ob_oracle::list_tables(p, schema).await.map(|tables| filter_table_infos(tables, filter, limit)) } else { - db::mysql::list_tables(p, schema).await + db::mysql::list_tables(p, schema).await.map(|tables| filter_table_infos(tables, filter, limit)) } } - PoolKind::Postgres(p) => db::postgres::list_tables(p, schema).await, - PoolKind::Sqlite(p) => db::sqlite::list_tables(p, schema).await, + PoolKind::Postgres(p) => { + db::postgres::list_tables(p, schema).await.map(|tables| filter_table_infos(tables, filter, limit)) + } + PoolKind::Sqlite(p) => { + db::sqlite::list_tables(p, schema).await.map(|tables| filter_table_infos(tables, filter, limit)) + } _ => Ok(vec![]), } } +fn filter_table_infos(tables: Vec, filter: Option<&str>, limit: Option) -> Vec { + let filter = filter.unwrap_or("").to_lowercase(); + let limit = limit.unwrap_or(usize::MAX); + tables + .into_iter() + .filter(|table| filter.is_empty() || table.name.to_lowercase().contains(&filter)) + .take(limit) + .collect() +} + pub async fn list_objects_core( state: &AppState, connection_id: &str, @@ -298,7 +314,7 @@ pub async fn list_objects_core( } } - Ok(list_tables_core(state, connection_id, database, schema) + Ok(list_tables_core(state, connection_id, database, schema, None, None) .await? .into_iter() .map(|table| db::ObjectInfo { diff --git a/src-tauri/src/commands/schema.rs b/src-tauri/src/commands/schema.rs index ee8514481f..8b90bfb369 100644 --- a/src-tauri/src/commands/schema.rs +++ b/src-tauri/src/commands/schema.rs @@ -27,8 +27,10 @@ pub async fn list_tables( connection_id: String, database: String, schema: String, + filter: Option, + limit: Option, ) -> Result, String> { - dbx_core::schema::list_tables_core(&state, &connection_id, &database, &schema).await + dbx_core::schema::list_tables_core(&state, &connection_id, &database, &schema, filter.as_deref(), limit).await } #[tauri::command] diff --git a/src-web/src/routes/schema.rs b/src-web/src/routes/schema.rs index 12a2db4d8c..d4fabf094a 100644 --- a/src-web/src/routes/schema.rs +++ b/src-web/src/routes/schema.rs @@ -13,6 +13,8 @@ pub struct SchemaQuery { pub database: Option, pub schema: Option, pub table: Option, + pub filter: Option, + pub limit: Option, pub object_type: Option, } @@ -39,8 +41,16 @@ pub async fn list_tables( ) -> Result, AppError> { let database = q.database.as_deref().unwrap_or(""); let schema = q.schema.as_deref().unwrap_or(""); - let result = - dbx_core::schema::list_tables_core(&state.app, &q.connection_id, database, schema).await.map_err(AppError)?; + let result = dbx_core::schema::list_tables_core( + &state.app, + &q.connection_id, + database, + schema, + q.filter.as_deref(), + q.limit, + ) + .await + .map_err(AppError)?; Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) } diff --git a/src/components/editor/QueryEditor.vue b/src/components/editor/QueryEditor.vue index fc658666b3..359f57bfdf 100644 --- a/src/components/editor/QueryEditor.vue +++ b/src/components/editor/QueryEditor.vue @@ -39,6 +39,7 @@ const connectionStore = useConnectionStore(); const settingsStore = useSettingsStore(); const MIN_FONT_SIZE = 10; const MAX_FONT_SIZE = 24; +const MAX_COMPLETION_TABLES = 200; let editorViewModule: typeof import("@codemirror/view") | null = null; let fontThemeComp: import("@codemirror/state").Compartment | null = null; let codeMirrorTheme: import("@codemirror/state").Compartment | null = null; @@ -126,9 +127,17 @@ async function provideSqlCompletions(currentState: import("@codemirror/state").E try { const fullDoc = currentState.doc.toString(); const completionContext = getSqlCompletionContext(fullDoc, position); - const tables = await connectionStore.listCompletionTables(props.connectionId, props.database); + const shouldLoadTables = completionContext.suggestTables || !!completionContext.qualifier; + const tables = shouldLoadTables + ? await connectionStore.listCompletionTables( + props.connectionId, + props.database, + completionContext.qualifier || completionContext.prefix, + MAX_COMPLETION_TABLES, + ) + : cachedTables; - // Collect referenced tables — enrich with schema from cached tables list + // Collect referenced tables — enrich with schema from filtered table lookup let refs = completionContext.referencedTables.map((rt) => { // If no schema, look it up in the cached tables if (!rt.schema) { @@ -139,6 +148,20 @@ async function provideSqlCompletions(currentState: import("@codemirror/state").E } return rt; }); + const unresolvedRefs = refs.filter((rt) => !rt.schema); + if (unresolvedRefs.length > 0) { + const lookupGroups = await Promise.all( + unresolvedRefs.map((rt) => + connectionStore.listCompletionTables(props.connectionId!, props.database!, rt.name, 20), + ), + ); + const lookupTables = lookupGroups.flat(); + refs = refs.map((rt) => { + if (rt.schema) return rt; + const matched = lookupTables.find((table) => table.name.toLowerCase() === rt.name.toLowerCase()); + return matched?.schema ? { ...rt, schema: matched.schema } : rt; + }); + } // If no referenced tables but qualifier exists, infer table from tables list if (refs.length === 0 && completionContext.qualifier) { @@ -200,13 +223,7 @@ async function provideSqlCompletions(currentState: import("@codemirror/state").E } async function refreshCompletionCache() { - if (!props.connectionId || !props.database) return; - try { - cachedTables = await connectionStore.listCompletionTables(props.connectionId, props.database); - } catch { - cachedTables = []; - } - // Clear column cache when connection/database changes + cachedTables = []; cachedColumnsByTable.clear(); } @@ -370,7 +387,12 @@ onMounted(async () => { try { // Ensure table cache is populated if (cachedTables.length === 0) { - cachedTables = await connectionStore.listCompletionTables(props.connectionId!, props.database!); + cachedTables = await connectionStore.listCompletionTables( + props.connectionId!, + props.database!, + identifier, + MAX_COMPLETION_TABLES, + ); } // 1. Check if it's a table name @@ -463,8 +485,7 @@ onMounted(async () => { view.value = new EditorView({ state, parent: editorRef.value }); - // Fetch completion data - refreshCompletionCache(); + cachedTables = []; }); watch( diff --git a/src/lib/http.ts b/src/lib/http.ts index 2a85f2daa7..fa9212076b 100644 --- a/src/lib/http.ts +++ b/src/lib/http.ts @@ -172,8 +172,14 @@ export async function listSchemas(connectionId: string, database: string): Promi return get(`/api/schema/schemas?${qs({ connection_id: connectionId, database })}`); } -export async function listTables(connectionId: string, database: string, schema: string): Promise { - return get(`/api/schema/tables?${qs({ connection_id: connectionId, database, schema })}`); +export async function listTables( + connectionId: string, + database: string, + schema: string, + filter?: string, + limit?: number, +): Promise { + return get(`/api/schema/tables?${qs({ connection_id: connectionId, database, schema, filter, limit })}`); } export async function listObjects(connectionId: string, database: string, schema: string): Promise { diff --git a/src/lib/sqlCompletion.ts b/src/lib/sqlCompletion.ts index 421bd94135..a017248201 100644 --- a/src/lib/sqlCompletion.ts +++ b/src/lib/sqlCompletion.ts @@ -133,6 +133,7 @@ const SQL_KEYWORDS = [ const TABLE_TRIGGER_KEYWORDS = new Set(["from", "join", "update", "into", "table", "describe", "explain", "apply"]); const JOIN_MODIFIERS = new Set(["left", "right", "inner", "outer", "cross", "full", "natural"]); +const MAX_TABLE_COMPLETION_ITEMS = 200; export interface SqlCompletionTable { name: string; @@ -463,6 +464,7 @@ function unquoteIdentifier(value: string): string { function buildTableItems(prefix: string, tables: SqlCompletionTable[]): SqlCompletionItem[] { return tables .filter((table) => matchesPrefix(table.name, prefix)) + .slice(0, MAX_TABLE_COMPLETION_ITEMS) .map((table) => ({ label: table.name, type: "table" as const, diff --git a/src/lib/tauri.ts b/src/lib/tauri.ts index 07c220a089..370e753245 100644 --- a/src/lib/tauri.ts +++ b/src/lib/tauri.ts @@ -138,8 +138,14 @@ export async function deleteSchemaCachePrefix(prefix: string): Promise { return invoke("delete_schema_cache_prefix", { prefix }); } -export async function listTables(connectionId: string, database: string, schema: string): Promise { - return invoke("list_tables", { connectionId, database, schema }); +export async function listTables( + connectionId: string, + database: string, + schema: string, + filter?: string, + limit?: number, +): Promise { + return invoke("list_tables", { connectionId, database, schema, filter, limit }); } export async function listObjects(connectionId: string, database: string, schema: string): Promise { diff --git a/src/stores/connectionStore.ts b/src/stores/connectionStore.ts index 4d7dfb2c49..d4163d0461 100644 --- a/src/stores/connectionStore.ts +++ b/src/stores/connectionStore.ts @@ -990,8 +990,14 @@ export const useConnectionStore = defineStore("connection", () => { } } - async function listCompletionTables(connectionId: string, database: string): Promise { - const cacheKey = `${connectionId}:${database}`; + async function listCompletionTables( + connectionId: string, + database: string, + filter = "", + limit?: number, + ): Promise { + const normalizedFilter = filter.trim().toLowerCase(); + const cacheKey = `${connectionId}:${database}:${normalizedFilter}:${limit ?? ""}`; if (completionTablesCache.value[cacheKey]) { return completionTablesCache.value[cacheKey]; } @@ -1000,6 +1006,29 @@ export const useConnectionStore = defineStore("connection", () => { if (isSchemaAwareDatabase(connectionId)) { const schemas = await api.listSchemas(connectionId, database); + if (normalizedFilter || limit) { + const limitedTables: SqlCompletionTable[] = []; + for (const schema of schemas) { + try { + const remaining = limit ? Math.max(limit - limitedTables.length, 0) : undefined; + if (remaining === 0) break; + const tables = await api.listTables(connectionId, database, schema, normalizedFilter, remaining); + limitedTables.push( + ...tables.map((table) => ({ + name: table.name, + schema, + type: table.table_type === "VIEW" ? ("view" as const) : ("table" as const), + })), + ); + } catch { + /* ignore schema failures */ + } + } + completionTablesCache.value[cacheKey] = limitedTables; + evictOldestCacheEntries(completionTablesCache.value, COMPLETION_CACHE_MAX); + return completionTablesCache.value[cacheKey]; + } + const tableGroups = await Promise.all( schemas.map(async (schema) => { try { @@ -1019,7 +1048,7 @@ export const useConnectionStore = defineStore("connection", () => { return completionTablesCache.value[cacheKey]; } - const tables = await api.listTables(connectionId, database, database); + const tables = await api.listTables(connectionId, database, database, normalizedFilter, limit); completionTablesCache.value[cacheKey] = tables.map((table) => ({ name: table.name, type: table.table_type === "VIEW" ? ("view" as const) : ("table" as const), diff --git a/tests/sqlCompletion.test.ts b/tests/sqlCompletion.test.ts index 9f183f0eb8..41e187dd61 100644 --- a/tests/sqlCompletion.test.ts +++ b/tests/sqlCompletion.test.ts @@ -13,16 +13,22 @@ const tables: SqlCompletionTable[] = [ ]; const columnsByTable = new Map([ - ["public.users", [ - { name: "id", table: "users", schema: "public", dataType: "bigint" }, - { name: "name", table: "users", schema: "public", dataType: "varchar" }, - { name: "email", table: "users", schema: "public", dataType: "varchar" }, - ]], - ["public.orders", [ - { name: "id", table: "orders", schema: "public", dataType: "bigint" }, - { name: "user_id", table: "orders", schema: "public", dataType: "bigint" }, - { name: "status", table: "orders", schema: "public", dataType: "varchar" }, - ]], + [ + "public.users", + [ + { name: "id", table: "users", schema: "public", dataType: "bigint" }, + { name: "name", table: "users", schema: "public", dataType: "varchar" }, + { name: "email", table: "users", schema: "public", dataType: "varchar" }, + ], + ], + [ + "public.orders", + [ + { name: "id", table: "orders", schema: "public", dataType: "bigint" }, + { name: "user_id", table: "orders", schema: "public", dataType: "bigint" }, + { name: "status", table: "orders", schema: "public", dataType: "varchar" }, + ], + ], ]); test("suggests SQL keywords for generic keyword input", () => { @@ -116,3 +122,22 @@ test("always includes keywords alongside table suggestions", () => { assert.ok(items.some((item) => item.type === "table")); assert.ok(items.some((item) => item.type === "keyword" && item.label === "USING")); }); + +test("limits table suggestions for large schemas after filtering by prefix", () => { + const largeTables: SqlCompletionTable[] = Array.from({ length: 500 }, (_, index) => ({ + name: `erp_invoice_${String(index).padStart(4, "0")}`, + schema: "dbo", + type: "table", + })); + + const sql = "select * from erp_invoice_"; + const items = buildSqlCompletionItems(sql, sql.length, { + tables: largeTables, + columnsByTable, + }); + + const tableItems = items.filter((item) => item.type === "table"); + assert.equal(tableItems.length, 200); + assert.equal(tableItems[0]?.label, "erp_invoice_0000"); + assert.equal(tableItems.at(-1)?.label, "erp_invoice_0199"); +});