diff --git a/apps/desktop/src/components/sidebar/TreeItem.vue b/apps/desktop/src/components/sidebar/TreeItem.vue index ad4e07bbee..94ddb6ef2c 100644 --- a/apps/desktop/src/components/sidebar/TreeItem.vue +++ b/apps/desktop/src/components/sidebar/TreeItem.vue @@ -245,6 +245,8 @@ function getIconInfo(node: TreeNode): { icon: any; colorClass: string } | null { return { icon: Package, colorClass: "text-cyan-500" }; case "group-partitions": return { icon: node.isExpanded ? FolderOpen : FolderClosed, colorClass: "text-green-400" }; + case "load-more": + return { icon: Plus, colorClass: "text-primary" }; default: return { icon: Database, colorClass: "text-muted-foreground" }; } @@ -258,6 +260,7 @@ function isGroupLabel(node: TreeNode): boolean { } function displayLabel(node: TreeNode): string { + if (node.type === "load-more") return t(node.label); if (node.type === "object-browser") return t(node.label, { count: node.objectCount ?? 0 }); if (node.type === "user-admin") return t(node.label); if (node.label === "tree.defaultDatabase") return t(node.label); @@ -390,6 +393,10 @@ async function toggle() { function runRowClickAction() { const node = props.node; + if (node.type === "load-more") { + void loadMoreObjectGroupChildren(); + return; + } if (node.type === "object-browser") { void openObjectBrowser(); return; @@ -404,6 +411,14 @@ function runRowClickAction() { } } +async function loadMoreObjectGroupChildren() { + try { + await connectionStore.loadMoreObjectGroupChildren(props.node); + } catch (e: any) { + toast(t("connection.connectFailed", { message: translateBackendError(t, e?.message || String(e)) }), 5000); + } +} + function visibleTreeNodes(): TreeNode[] { if (sidebarTreeContext) return sidebarTreeContext.getVisibleNodes(); return flattenTree(connectionStore.treeNodes).map((item) => item.node); @@ -3068,6 +3083,7 @@ function treeItemMenuItems(): ContextMenuItem[] { + { - return get(`/api/schema/tables?${qs({ connection_id: connectionId, database, schema, filter, limit })}`); +export async function listTables(connectionId: string, database: string, schema: string, filter?: string, limit?: number, offset?: number): Promise { + return get(`/api/schema/tables?${qs({ connection_id: connectionId, database, schema, filter, limit, offset })}`); } export async function listObjects(connectionId: string, database: string, schema: string, objectTypes?: SidebarObjectKind[]): Promise { diff --git a/apps/desktop/src/lib/sidebarTreeItemLayout.ts b/apps/desktop/src/lib/sidebarTreeItemLayout.ts index 291babd0d4..94e57c0b9e 100644 --- a/apps/desktop/src/lib/sidebarTreeItemLayout.ts +++ b/apps/desktop/src/lib/sidebarTreeItemLayout.ts @@ -1,6 +1,6 @@ import type { TreeNodeType } from "@/types/database"; -const leafTypes: Set = new Set(["column", "index", "fkey", "trigger", "procedure", "function", "package", "package-body", "object-browser", "redis-db", "mongo-collection", "elasticsearch-index", "user-admin", "saved-sql-file"]); +const leafTypes: Set = new Set(["column", "index", "fkey", "trigger", "procedure", "function", "package", "package-body", "object-browser", "redis-db", "mongo-collection", "elasticsearch-index", "user-admin", "saved-sql-file", "load-more"]); const fullWidthLabelTypes: Set = new Set(["table", "view", "mongo-collection", "elasticsearch-index"]); diff --git a/apps/desktop/src/lib/tauri.ts b/apps/desktop/src/lib/tauri.ts index 65e4db5e4f..f01b2721f4 100644 --- a/apps/desktop/src/lib/tauri.ts +++ b/apps/desktop/src/lib/tauri.ts @@ -449,8 +449,8 @@ export async function deleteSchemaCachePrefix(prefix: string): Promise { return invoke("delete_schema_cache_prefix", { prefix }); } -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 listTables(connectionId: string, database: string, schema: string, filter?: string, limit?: number, offset?: number): Promise { + return invoke("list_tables", { connectionId, database, schema, filter, limit, offset }); } export async function listObjects(connectionId: string, database: string, schema: string, objectTypes?: SidebarObjectKind[]): Promise { diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index dc52be7cd2..ed9e276990 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -51,6 +51,7 @@ import { useSettingsStore } from "@/stores/settingsStore"; const PINNED_TREE_NODES_STORAGE_KEY = "dbx-pinned-tree-nodes"; const ACTIVE_CONNECTION_STORAGE_KEY = "dbx-active-connection"; +const SIDEBAR_OBJECT_GROUP_PAGE_SIZE = 500; type ImportSource = "dbx" | "navicat" | "dbeaver" | "datagrip"; // Temporary storage for DataGrip import payload (used to read Keychain passwords after import) @@ -395,6 +396,75 @@ export const useConnectionStore = defineStore("connection", () => { return sidebarObjectKindsForDatabase(dbType); } + function objectGroupCacheKey(node: TreeNode): string { + return schemaCacheKey(node.connectionId || "", node.database || "", node.schema || "", node.type, "objects-v2"); + } + + function buildLoadMoreNode(parent: TreeNode, offset: number, pageSize: number): TreeNode { + return { + id: `${parent.id}:__load_more:${offset}`, + label: "tree.loadMore", + type: "load-more", + connectionId: parent.connectionId, + database: parent.database, + schema: parent.schema, + isLoading: false, + loadMore: { + parentId: parent.id, + offset, + pageSize, + }, + }; + } + + function withoutLoadMoreNodes(children: TreeNode[] | undefined): TreeNode[] { + return (children || []).filter((child) => child.type !== "load-more"); + } + + function objectGroupChildrenFromObjects(options: { node: TreeNode; parentNodeId: string; effectiveSchema?: string; objectTypes: DatabaseObjectTreeKind[]; objects: ObjectInfo[] }): TreeNode[] { + const grouped = buildGroupedObjectTreeNodes({ + nodeId: options.parentNodeId, + connectionId: options.node.connectionId || "", + database: options.node.database || "", + schema: options.effectiveSchema, + objects: options.objects.filter((object) => options.objectTypes.includes(normalizedObjectTreeKind(object.object_type))), + }); + const refreshedGroup = grouped.find((group) => group.type === options.node.type); + return refreshedGroup?.children ?? []; + } + + async function loadPagedTableGroupChildren(options: { + node: TreeNode; + parentNodeId: string; + querySchema: string; + effectiveSchema?: string; + objectTypes: DatabaseObjectTreeKind[]; + offset: number; + pageSize: number; + }): Promise<{ children: TreeNode[]; objectCount: number; hasMore: boolean; nextOffset: number }> { + if (!options.node.connectionId || !options.node.database) { + return { children: [], objectCount: 0, hasMore: false, nextOffset: options.offset }; + } + const fetchLimit = options.pageSize + 1; + const tables = await api.listTables(options.node.connectionId, options.node.database, options.querySchema, undefined, fetchLimit, options.offset); + const hasMore = tables.length > options.pageSize; + const pageTables = hasMore ? tables.slice(0, options.pageSize) : tables; + const objects = mergeTableInfosIntoObjects([], pageTables, options.effectiveSchema); + const visibleObjectCount = objects.filter((object) => options.objectTypes.includes(normalizedObjectTreeKind(object.object_type))).length; + return { + children: objectGroupChildrenFromObjects({ + node: options.node, + parentNodeId: options.parentNodeId, + effectiveSchema: options.effectiveSchema, + objectTypes: options.objectTypes, + objects, + }), + objectCount: visibleObjectCount, + hasMore, + nextOffset: options.offset + pageTables.length, + }; + } + function refreshStaleTreeNode(node: TreeNode) { if (staleTreeRefreshIds.has(node.id)) return; staleTreeRefreshIds.add(node.id); @@ -1190,7 +1260,7 @@ export const useConnectionStore = defineStore("connection", () => { const config = getConfig(node.connectionId); const querySchema = connectionObjectTreeQuerySchema(config, node.database, node.schema); const effectiveSchema = connectionObjectTreeNodeSchema(config, node.database, node.schema); - const cacheKey = schemaCacheKey(node.connectionId, node.database, node.schema || "", node.type, "objects-v1"); + const cacheKey = objectGroupCacheKey(node); if (!options?.force) { const cached = await loadPersistedTreeChildren(node, cacheKey); if (cached.hit) { @@ -1200,17 +1270,30 @@ export const useConnectionStore = defineStore("connection", () => { } const wantsOnlyTablesOrViews = objectTypes.every((objectType) => objectType === "TABLE" || objectType === "VIEW"); - const objects = wantsOnlyTablesOrViews ? mergeTableInfosIntoObjects([], await api.listTables(node.connectionId, node.database, querySchema), effectiveSchema) : await api.listObjects(node.connectionId, node.database, querySchema, objectTypes); - const grouped = buildGroupedObjectTreeNodes({ - nodeId: parentNodeId, - connectionId: node.connectionId, - database: node.database, - schema: effectiveSchema, - objects: objects.filter((object) => objectTypes.includes(normalizedObjectTreeKind(object.object_type))), - }); - const refreshedGroup = grouped.find((group) => group.type === node.type); - const children = refreshedGroup?.children ?? []; - node.objectCount = refreshedGroup?.objectCount ?? children.length; + let children: TreeNode[]; + if (wantsOnlyTablesOrViews) { + const page = await loadPagedTableGroupChildren({ + node, + parentNodeId, + querySchema, + effectiveSchema, + objectTypes, + offset: 0, + pageSize: SIDEBAR_OBJECT_GROUP_PAGE_SIZE, + }); + children = page.hasMore ? [...page.children, buildLoadMoreNode(node, page.nextOffset, SIDEBAR_OBJECT_GROUP_PAGE_SIZE)] : page.children; + node.objectCount = page.objectCount; + } else { + const objects = await api.listObjects(node.connectionId, node.database, querySchema, objectTypes); + children = objectGroupChildrenFromObjects({ + node, + parentNodeId, + effectiveSchema, + objectTypes, + objects, + }); + node.objectCount = children.length; + } setChildren(node, children); await savePersistedTreeChildren(cacheKey, children); node.isExpanded = true; @@ -1222,6 +1305,44 @@ export const useConnectionStore = defineStore("connection", () => { } } + async function loadMoreObjectGroupChildren(node: TreeNode) { + if (node.type !== "load-more" || !node.loadMore) return; + const parent = findNode(treeNodes.value, node.loadMore.parentId); + if (!parent?.connectionId || !hasTreeNodeDatabaseContext(parent)) return; + node.isLoading = true; + try { + await ensureConnected(parent.connectionId); + const objectTypes = objectTypesForGroupNode(parent.type); + const parentNodeId = objectGroupRefreshParentId(parent); + if (!objectTypes || !parentNodeId) return; + if (!objectTypes.every((objectType) => objectType === "TABLE" || objectType === "VIEW")) return; + + const config = getConfig(parent.connectionId); + const querySchema = connectionObjectTreeQuerySchema(config, parent.database, parent.schema); + const effectiveSchema = connectionObjectTreeNodeSchema(config, parent.database, parent.schema); + const page = await loadPagedTableGroupChildren({ + node: parent, + parentNodeId, + querySchema, + effectiveSchema, + objectTypes, + offset: node.loadMore.offset, + pageSize: node.loadMore.pageSize, + }); + const currentChildren = withoutLoadMoreNodes(parent.children); + const nextChildren = page.hasMore ? [...currentChildren, ...page.children, buildLoadMoreNode(parent, page.nextOffset, node.loadMore.pageSize)] : [...currentChildren, ...page.children]; + parent.objectCount = (parent.objectCount ?? currentChildren.length) + page.objectCount; + setChildren(parent, nextChildren); + await savePersistedTreeChildren(objectGroupCacheKey(parent), nextChildren); + parent.isExpanded = true; + } catch (e) { + recordMetadataLoadError(parent.connectionId, e); + throw e; + } finally { + node.isLoading = false; + } + } + function normalizedObjectTreeKind(type: string): DatabaseObjectTreeKind { return normalizeSidebarObjectKind(type); } @@ -2575,6 +2696,7 @@ export const useConnectionStore = defineStore("connection", () => { loadSqlServerDatabaseObjects, loadTables, loadObjectGroupChildren, + loadMoreObjectGroupChildren, loadTableGroups, loadColumns, loadIndexes, diff --git a/apps/desktop/src/types/database.ts b/apps/desktop/src/types/database.ts index f71a164195..e96836f068 100644 --- a/apps/desktop/src/types/database.ts +++ b/apps/desktop/src/types/database.ts @@ -375,6 +375,7 @@ export type TreeNodeType = | "saved-sql-root" | "saved-sql-folder" | "saved-sql-file" + | "load-more" | "column" | "index" | "fkey" @@ -418,6 +419,11 @@ export interface TreeNode { savedSqlId?: string; savedSqlFolderId?: string; meta?: ColumnInfo | IndexInfo | ForeignKeyInfo | TriggerInfo; + loadMore?: { + parentId: string; + offset: number; + pageSize: number; + }; } export interface QueryTab { diff --git a/crates/dbx-core/src/agent_loop.rs b/crates/dbx-core/src/agent_loop.rs index d3154e0bac..74abda29e3 100644 --- a/crates/dbx-core/src/agent_loop.rs +++ b/crates/dbx-core/src/agent_loop.rs @@ -254,6 +254,7 @@ async fn build_schema_prompt(agent_ctx: &AgentLoopContext, system_prompt: &str) "", None, Some(50), // smaller limit for prompt injection + None, ) .await; diff --git a/crates/dbx-core/src/agent_tools.rs b/crates/dbx-core/src/agent_tools.rs index 4614f3143a..edbd730527 100644 --- a/crates/dbx-core/src/agent_tools.rs +++ b/crates/dbx-core/src/agent_tools.rs @@ -169,10 +169,17 @@ async fn execute_list_tables( let schema = tool_call.arguments.get("schema").and_then(|v| v.as_str()).unwrap_or("").to_string(); // Request one extra to detect whether more tables exist beyond the limit. - let tables = - crate::schema::list_tables_core(state, connection_id, database, &schema, None, Some(LIST_TABLES_LIMIT + 1)) - .await - .map_err(|e| format!("Failed to list tables: {e}"))?; + let tables = crate::schema::list_tables_core( + state, + connection_id, + database, + &schema, + None, + Some(LIST_TABLES_LIMIT + 1), + None, + ) + .await + .map_err(|e| format!("Failed to list tables: {e}"))?; let total = tables.len(); let truncated = total > LIST_TABLES_LIMIT; diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 2ebc59ae83..53fa034842 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -2286,7 +2286,8 @@ mod tests { let schemas = schema::list_schemas_core(&state, "kwdb-live", &database).await.unwrap(); assert!(schemas.iter().any(|schema| schema == test_schema)); - let tables = schema::list_tables_core(&state, "kwdb-live", &database, test_schema, None, None).await.unwrap(); + let tables = + schema::list_tables_core(&state, "kwdb-live", &database, test_schema, None, None, None).await.unwrap(); assert!(tables.iter().any(|table| table.name == "devices" && table.table_type == "BASE TABLE")); let columns = schema::get_columns_core(&state, "kwdb-live", &database, test_schema, "devices").await.unwrap(); let id_column = columns.iter().find(|column| column.name == "id").expect("id column should be listed"); diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index d81a96a35a..9302c03e03 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -352,9 +352,16 @@ pub async fn export_database_sql_core( let pool_key = state.get_or_create_pool(&request.connection_id, Some(&request.database)).await?; // 3. List tables - let all_tables = - crate::schema::list_tables_core(state, &request.connection_id, &request.database, &request.schema, None, None) - .await?; + let all_tables = crate::schema::list_tables_core( + state, + &request.connection_id, + &request.database, + &request.schema, + None, + None, + None, + ) + .await?; let all_tables = filter_selected_table_infos(all_tables, &request.selected_tables); // 4. Create file diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index c72d255ada..2320517fe3 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -1093,10 +1093,26 @@ pub async fn list_databases(pool: &Pool) -> Result, String> { } pub async fn list_tables(pool: &Pool, schema: &str) -> Result, String> { + list_tables_filtered(pool, schema, None, None, None).await +} + +pub async fn list_tables_filtered( + pool: &Pool, + schema: &str, + filter: Option<&str>, + limit: Option, + offset: Option, +) -> Result, String> { let schema = if schema.is_empty() { "public" } else { schema }; + let filter_pattern = like_contains_pattern(filter.unwrap_or("").trim()); + let limit_param = limit.and_then(|value| i64::try_from(value).ok()); + let offset_param = offset.and_then(|value| i64::try_from(value).ok()).unwrap_or(0); let client = pool.get().await.map_err(|e| e.to_string())?; let stmt = client.prepare_cached(postgres_tables_sql()).await.map_err(|e| e.to_string())?; - let rows = client.query(&stmt, &[&schema]).await.map_err(|e| e.to_string())?; + let rows = client + .query(&stmt, &[&schema, &filter_pattern, &limit_param, &offset_param]) + .await + .map_err(|e| e.to_string())?; Ok(rows .iter() @@ -1124,7 +1140,26 @@ fn postgres_tables_sql() -> &'static str { LEFT JOIN pg_catalog.pg_class pc ON pc.oid = i.inhparent \ LEFT JOIN pg_catalog.pg_namespace pn ON pn.oid = pc.relnamespace \ WHERE n.nspname = $1 AND c.relkind IN ('r','v','m','f','p') \ - ORDER BY c.relname" + AND ($2 = '%%' OR c.relname ILIKE $2 ESCAPE '\\') \ + ORDER BY c.relname \ + LIMIT $3 OFFSET $4" +} + +fn like_contains_pattern(value: &str) -> String { + if value.is_empty() { + return "%%".to_string(); + } + + let mut pattern = String::with_capacity(value.len() + 2); + pattern.push('%'); + for ch in value.chars() { + if ch == '\\' || ch == '%' || ch == '_' { + pattern.push('\\'); + } + pattern.push(ch); + } + pattern.push('%'); + pattern } fn list_objects_sql(include_timestamps: bool) -> &'static str { @@ -2365,4 +2400,11 @@ mod tests { fn url_with_uppercase_timezone_returns_true() { assert!(pg_url_has_timezone_setting("postgres://localhost/db?options=--TimeZone=UTC")); } + + #[test] + fn like_contains_pattern_escapes_wildcards() { + assert_eq!(like_contains_pattern(""), "%%"); + assert_eq!(like_contains_pattern("order_100%"), "%order\\_100\\%%"); + assert_eq!(like_contains_pattern(r"foo\bar"), r"%foo\\bar%"); + } } diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index 268d72ddf1..e48a0d48c4 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -633,15 +633,20 @@ pub async fn list_tables( schema: &str, filter: Option<&str>, limit: Option, + offset: Option, ) -> Result, String> { - let top = limit.map(|value| format!("TOP ({}) ", value.min(1000))).unwrap_or_default(); + let fetch_clause = limit + .map(|value| format!(" OFFSET {} ROWS FETCH NEXT {} ROWS ONLY", offset.unwrap_or(0), value.min(1000))) + .unwrap_or_else(|| { + offset.filter(|value| *value > 0).map(|value| format!(" OFFSET {value} ROWS")).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 schema_escaped = schema.replace('\'', "''"); let sql = format!( - "SELECT {top}o.name, CASE WHEN o.type = 'V' THEN 'VIEW' ELSE 'BASE TABLE' END, \ + "SELECT o.name, CASE WHEN o.type = 'V' THEN 'VIEW' ELSE 'BASE TABLE' END, \ ep.value AS TABLE_COMMENT \ FROM sys.objects o \ JOIN sys.schemas s ON s.schema_id = o.schema_id \ @@ -650,7 +655,7 @@ pub async fn list_tables( AND o.type IN ('U','V') \ AND o.is_ms_shipped = 0 \ {filter_clause}\ - ORDER BY o.name" + ORDER BY o.name{fetch_clause}" ); let stream = client.query(&*sql, &[]).await.map_err(|e| e.to_string())?; let rows = stream.into_first_result().await.map_err(|e| e.to_string())?; diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 92ef5cb39f..3ee5b9c9f7 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -416,9 +416,10 @@ pub async fn list_tables_core( schema: &str, filter: Option<&str>, limit: Option, + offset: Option, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || { - list_tables_once(state, connection_id, database, schema, filter, limit) + list_tables_once(state, connection_id, database, schema, filter, limit, offset) }) .await } @@ -430,6 +431,7 @@ async fn list_tables_once( schema: &str, filter: Option<&str>, limit: Option, + offset: Option, ) -> Result, String> { let pool_key = state.get_or_create_pool(connection_id, Some(database)).await?; #[cfg(feature = "duckdb-bundled")] @@ -447,7 +449,8 @@ async fn list_tables_once( duckdb_query_tables(&con) }) .await - .map_err(|e| e.to_string())?; + .map_err(|e| e.to_string())? + .map(|tables| filter_table_infos(tables, filter, limit, offset)); } if let Some(PoolKind::ExternalDriver { config, session, .. }) = connections.get(&pool_key) { let config = config.clone(); @@ -458,34 +461,36 @@ async fn list_tables_once( "listTables", serde_json::json!({ "connection": config.as_ref(), "database": database, "schema": schema }), ) - .await; + .await + .map(|tables| filter_table_infos(tables, filter, limit, offset)); } #[cfg(feature = "duckdb-bundled")] if let Some(con) = extract_pool!(&connections, &pool_key, DuckDb) { drop(connections); let con = con.lock().map_err(|e| e.to_string())?; - return duckdb_query_tables_in_database_with_attached(&con, database, schema, &duckdb_attached_names); + return duckdb_query_tables_in_database_with_attached(&con, database, schema, &duckdb_attached_names) + .map(|tables| filter_table_infos(tables, filter, limit, offset)); } if let Some(client) = extract_pool!(&connections, &pool_key, ClickHouse) { drop(connections); - return db::clickhouse_driver::list_tables(&client, clickhouse_metadata_database(database, schema)).await; + return db::clickhouse_driver::list_tables(&client, clickhouse_metadata_database(database, schema)) + .await + .map(|tables| filter_table_infos(tables, filter, limit, offset)); } - try_sqlserver!(connections, &pool_key, list_tables, schema, filter, limit); + try_sqlserver!(connections, &pool_key, list_tables, schema, filter, limit, offset); if let Some(client) = extract_pool!(&connections, &pool_key, Agent) { let fallback_config = db_config.clone(); drop(connections); let mut client = client.lock().await; match client.list_tables::>(database, schema).await { - Ok(tables) if !tables.is_empty() => return Ok(filter_table_infos(tables, filter, limit)), + Ok(tables) if !tables.is_empty() => return Ok(filter_table_infos(tables, filter, limit, offset)), Ok(tables) => { if let Some(config) = fallback_config.as_ref() { match native_postgres_metadata_pool(state, connection_id, database, config).await { Ok(Some(pool)) => { - return db::postgres::list_tables(&pool, schema) - .await - .map(|tables| filter_table_infos(tables, filter, limit)); + return db::postgres::list_tables_filtered(&pool, schema, filter, limit, offset).await; } - Ok(None) => return Ok(filter_table_infos(tables, filter, limit)), + Ok(None) => return Ok(filter_table_infos(tables, filter, limit, offset)), Err(error) => { log::warn!( "[schema][agent:list_tables:fallback-failed] connection_id={} database={} schema={} error={}", @@ -497,16 +502,15 @@ async fn list_tables_once( } } } - return Ok(filter_table_infos(tables, filter, limit)); + return Ok(filter_table_infos(tables, filter, limit, offset)); } Err(agent_error) => { if let Some(config) = fallback_config.as_ref() { if let Some(pool) = native_postgres_metadata_pool(state, connection_id, database, config).await? { - return db::postgres::list_tables(&pool, schema) + return db::postgres::list_tables_filtered(&pool, schema, filter, limit, offset) .await - .map(|tables| filter_table_infos(tables, filter, limit)) .map_err(|fallback_error| { format!( "{agent_error}\n\nNative PostgreSQL metadata fallback failed: {fallback_error}" @@ -525,29 +529,29 @@ async fn list_tables_once( match pool { PoolKind::Mysql(p, _) if db_config.as_ref().is_some_and(is_doris_family_config) => { - db::mysql::list_tables_show(p, database).await.map(|tables| filter_table_infos(tables, filter, limit)) + db::mysql::list_tables_show(p, database) + .await + .map(|tables| filter_table_infos(tables, filter, limit, offset)) } PoolKind::Mysql(p, mode) => { dispatch_mysql!(p, mode, db::mysql::list_tables, db::ob_oracle::list_tables, schema) - .map(|tables| filter_table_infos(tables, filter, limit)) - } - PoolKind::Postgres(p) => { - db::postgres::list_tables(p, schema).await.map(|tables| filter_table_infos(tables, filter, limit)) + .map(|tables| filter_table_infos(tables, filter, limit, offset)) } + PoolKind::Postgres(p) => db::postgres::list_tables_filtered(p, schema, filter, limit, offset).await, PoolKind::Sqlite(p) => { - db::sqlite::list_tables(p, schema).await.map(|tables| filter_table_infos(tables, filter, limit)) - } - PoolKind::Rqlite(client) => { - db::rqlite_driver::list_tables(client, schema).await.map(|tables| filter_table_infos(tables, filter, limit)) + db::sqlite::list_tables(p, schema).await.map(|tables| filter_table_infos(tables, filter, limit, offset)) } + PoolKind::Rqlite(client) => db::rqlite_driver::list_tables(client, schema) + .await + .map(|tables| filter_table_infos(tables, filter, limit, offset)), PoolKind::MongoDb(client) => db::mongo_driver::list_collections(client, database) .await .map(|names| collection_names_to_tables(names, "COLLECTION")) - .map(|tables| filter_table_infos(tables, filter, limit)), + .map(|tables| filter_table_infos(tables, filter, limit, offset)), PoolKind::Elasticsearch(client) => db::elasticsearch_driver::list_indices(client) .await .map(|names| collection_names_to_tables(names, "INDEX")) - .map(|tables| filter_table_infos(tables, filter, limit)), + .map(|tables| filter_table_infos(tables, filter, limit, offset)), _ => Ok(vec![]), } } @@ -565,19 +569,29 @@ fn collection_names_to_tables(names: Vec, table_type: &str) -> Vec, filter: Option<&str>, limit: Option) -> Vec { +fn filter_table_infos( + tables: Vec, + filter: Option<&str>, + limit: Option, + offset: Option, +) -> Vec { let filter = filter.unwrap_or("").to_lowercase(); let limit = limit.unwrap_or(usize::MAX); + let offset = offset.unwrap_or(0); tables .into_iter() .filter(|table| filter.is_empty() || table.name.to_lowercase().contains(&filter)) + .skip(offset) .take(limit) .collect() } #[cfg(test)] mod tests { - use super::{clickhouse_metadata_database, deduplicate_column_infos, is_agent_postgres_metadata_fallback_config}; + use super::{ + clickhouse_metadata_database, deduplicate_column_infos, filter_table_infos, + is_agent_postgres_metadata_fallback_config, + }; #[cfg(feature = "duckdb-bundled")] use super::{duckdb_attach_database, duckdb_list_databases, duckdb_query_tables_in_database}; use crate::models::connection::{ConnectionConfig, DatabaseType}; @@ -641,6 +655,31 @@ mod tests { } } + fn test_table_info(name: &str) -> super::db::TableInfo { + super::db::TableInfo { + name: name.to_string(), + table_type: "BASE TABLE".to_string(), + comment: None, + parent_schema: None, + parent_name: None, + } + } + + #[test] + fn filter_table_infos_applies_filter_offset_and_limit() { + let tables = vec![ + test_table_info("alpha"), + test_table_info("audit_log"), + test_table_info("audit_record"), + test_table_info("users"), + ]; + + let filtered = filter_table_infos(tables, Some("audit"), Some(1), Some(1)); + + assert_eq!(filtered.len(), 1); + assert_eq!(filtered[0].name, "audit_record"); + } + #[cfg(feature = "duckdb-bundled")] #[test] fn duckdb_list_databases_includes_attached_database() { @@ -846,7 +885,7 @@ async fn list_objects_once( PoolKind::Postgres(p) => db::postgres::list_objects(p, schema).await, _ => { drop(connections); - Ok(list_tables_core(state, connection_id, database, schema, None, None) + Ok(list_tables_core(state, connection_id, database, schema, None, None, None) .await? .into_iter() .map(|table| db::ObjectInfo { diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index 24af584aec..2ee02553fe 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -2726,6 +2726,7 @@ where &request.source_schema, Some(table), Some(1), + None, ) .await .unwrap_or_default() @@ -2740,6 +2741,7 @@ where &request.target_schema, Some(table), Some(1), + None, ) .await .map(|tables| !tables.is_empty()) diff --git a/crates/dbx-web/src/routes/schema.rs b/crates/dbx-web/src/routes/schema.rs index 2a54c93c7f..3b03337a72 100644 --- a/crates/dbx-web/src/routes/schema.rs +++ b/crates/dbx-web/src/routes/schema.rs @@ -15,6 +15,7 @@ pub struct SchemaQuery { pub table: Option, pub filter: Option, pub limit: Option, + pub offset: Option, pub object_type: Option, } @@ -48,6 +49,7 @@ pub async fn list_tables( schema, q.filter.as_deref(), q.limit, + q.offset, ) .await .map_err(AppError)?; diff --git a/src-tauri/src/commands/mcp_bridge.rs b/src-tauri/src/commands/mcp_bridge.rs index e559a3d858..b6ad60a937 100644 --- a/src-tauri/src/commands/mcp_bridge.rs +++ b/src-tauri/src/commands/mcp_bridge.rs @@ -327,7 +327,7 @@ async fn handle_list_tables_data(state: &Arc, body: &str, stream: &mut respond_error(stream, "403 Forbidden", &e).await; return; } - match dbx_core::schema::list_tables_core(state, &config.id, &database, &schema, None, None).await { + match dbx_core::schema::list_tables_core(state, &config.id, &database, &schema, None, None, None).await { Ok(tables) => respond_json(stream, &tables).await, Err(e) => respond_error(stream, "500 Internal Server Error", &e).await, } diff --git a/src-tauri/src/commands/schema.rs b/src-tauri/src/commands/schema.rs index 54e1df9966..fe10c4ea00 100644 --- a/src-tauri/src/commands/schema.rs +++ b/src-tauri/src/commands/schema.rs @@ -29,8 +29,10 @@ pub async fn list_tables( schema: String, filter: Option, limit: Option, + offset: Option, ) -> Result, String> { - dbx_core::schema::list_tables_core(&state, &connection_id, &database, &schema, filter.as_deref(), limit).await + dbx_core::schema::list_tables_core(&state, &connection_id, &database, &schema, filter.as_deref(), limit, offset) + .await } #[tauri::command]