diff --git a/apps/desktop/src/components/sidebar/TreeItem.vue b/apps/desktop/src/components/sidebar/TreeItem.vue index 498b84bc90..60257096cd 100644 --- a/apps/desktop/src/components/sidebar/TreeItem.vue +++ b/apps/desktop/src/components/sidebar/TreeItem.vue @@ -350,14 +350,19 @@ async function toggle() { return; } - if ( + if (node.type === "group-partitions") { + node.isExpanded = !node.isExpanded; + emit("node-toggled", node, wasExpanded); + return; + } + + const databaseObjectGroup = node.type === "group-tables" || node.type === "group-views" || node.type === "group-procedures" || node.type === "group-functions" || - node.type === "group-packages" || - node.type === "group-partitions" - ) { + node.type === "group-packages"; + if (databaseObjectGroup && connectionStore.isTreeNodeChildrenLoaded(node.id)) { node.isExpanded = !node.isExpanded; emit("node-toggled", node, wasExpanded); return; @@ -428,6 +433,8 @@ async function toggle() { node.tableName ) { await connectionStore.loadTriggers(node.connectionId, node.database, node.tableName, node.schema, node.id); + } else if (databaseObjectGroup) { + await connectionStore.loadObjectGroupChildren(node); } emit("node-toggled", node, wasExpanded); } catch (e: any) { diff --git a/apps/desktop/src/lib/databaseCapabilities.ts b/apps/desktop/src/lib/databaseCapabilities.ts index 9b4e979dae..7d1f40fb84 100644 --- a/apps/desktop/src/lib/databaseCapabilities.ts +++ b/apps/desktop/src/lib/databaseCapabilities.ts @@ -1,3 +1,4 @@ export * from "./databaseCapabilitySets"; export * from "./databaseFeatureSupport"; +export * from "./databaseObjectCapabilities"; export * from "./databaseTableDataCapabilities"; diff --git a/apps/desktop/src/lib/databaseObjectCapabilities.ts b/apps/desktop/src/lib/databaseObjectCapabilities.ts new file mode 100644 index 0000000000..95456b0e03 --- /dev/null +++ b/apps/desktop/src/lib/databaseObjectCapabilities.ts @@ -0,0 +1,59 @@ +import type { DatabaseType } from "@/types/database"; + +export type SidebarObjectKind = "TABLE" | "VIEW" | "PROCEDURE" | "FUNCTION" | "PACKAGE" | "PACKAGE_BODY"; + +export interface DatabaseObjectCapabilities { + sidebarObjects: SidebarObjectKind[]; + sourceReadable: SidebarObjectKind[]; + executable: SidebarObjectKind[]; +} + +const TABLE_VIEW_OBJECTS: SidebarObjectKind[] = ["TABLE", "VIEW"]; +const ROUTINE_OBJECTS: SidebarObjectKind[] = ["TABLE", "VIEW", "PROCEDURE", "FUNCTION"]; +const ORACLE_OBJECTS: SidebarObjectKind[] = ["TABLE", "VIEW", "PROCEDURE", "FUNCTION", "PACKAGE", "PACKAGE_BODY"]; + +const TABLE_VIEW_ONLY_TYPES = new Set([ + "sqlite", + "rqlite", + "duckdb", + "clickhouse", + "doris", + "starrocks", + "databend", + "hive", + "trino", + "cassandra", + "bigquery", + "kylin", + "tdengine", + "iotdb", + "neo4j", +]); + +const ORACLE_PACKAGE_TYPES = new Set(["oracle", "oceanbase-oracle"]); + +export function databaseObjectCapabilities(dbType?: DatabaseType): DatabaseObjectCapabilities { + const sidebarObjects = sidebarObjectKindsForDatabase(dbType); + return { + sidebarObjects, + sourceReadable: sidebarObjects.filter((kind) => kind !== "TABLE"), + executable: sidebarObjects.filter((kind) => kind === "PROCEDURE"), + }; +} + +export function sidebarObjectKindsForDatabase(dbType?: DatabaseType): SidebarObjectKind[] { + if (!dbType) return [...TABLE_VIEW_OBJECTS]; + if (ORACLE_PACKAGE_TYPES.has(dbType)) return [...ORACLE_OBJECTS]; + if (TABLE_VIEW_ONLY_TYPES.has(dbType)) return [...TABLE_VIEW_OBJECTS]; + return [...ROUTINE_OBJECTS]; +} + +export function normalizeSidebarObjectKind(type: string): SidebarObjectKind { + const value = type.toUpperCase(); + if (value.includes("PACKAGE BODY") || value.includes("PACKAGE_BODY")) return "PACKAGE_BODY"; + if (value.includes("PACKAGE")) return "PACKAGE"; + if (value.includes("VIEW")) return "VIEW"; + if (value.includes("PROC")) return "PROCEDURE"; + if (value.includes("FUNC")) return "FUNCTION"; + return "TABLE"; +} diff --git a/apps/desktop/src/lib/http.ts b/apps/desktop/src/lib/http.ts index fc525ae333..047282a6b4 100644 --- a/apps/desktop/src/lib/http.ts +++ b/apps/desktop/src/lib/http.ts @@ -20,6 +20,7 @@ import type { SavedSqlFolder, SavedSqlLibrary, } from "@/types/database"; +import type { SidebarObjectKind } from "@/lib/databaseObjectCapabilities"; import type { AiConfig } from "@/stores/settingsStore"; import type { AgentDriverInfo, @@ -394,8 +395,20 @@ export async function listTables( return get(`/api/schema/tables?${qs({ connection_id: connectionId, database, schema, filter, limit })}`); } -export async function listObjects(connectionId: string, database: string, schema: string): Promise { - return get(`/api/schema/objects?${qs({ connection_id: connectionId, database, schema })}`); +export async function listObjects( + connectionId: string, + database: string, + schema: string, + objectTypes?: SidebarObjectKind[], +): Promise { + return get( + `/api/schema/objects?${qs({ + connection_id: connectionId, + database, + schema, + object_types: objectTypes?.join(","), + })}`, + ); } export async function listCompletionObjects( diff --git a/apps/desktop/src/lib/sqlServerTree.ts b/apps/desktop/src/lib/sqlServerTree.ts index adbb4df361..d66234ff25 100644 --- a/apps/desktop/src/lib/sqlServerTree.ts +++ b/apps/desktop/src/lib/sqlServerTree.ts @@ -1,6 +1,11 @@ import type { ObjectInfo, TreeNode } from "@/types/database"; import { sortSidebarNames } from "@/lib/databaseTree"; -import { buildGroupedObjectTreeNodes, buildSimpleObjectTreeNodes } from "@/lib/tableTree"; +import { + buildGroupedObjectTreeNodes, + buildObjectGroupPlaceholderNodes, + buildSimpleObjectTreeNodes, + type DatabaseObjectTreeKind, +} from "@/lib/tableTree"; export const SQLSERVER_DEFAULT_SCHEMA = "dbo"; @@ -13,26 +18,34 @@ export function buildSqlServerDatabaseTreeNodes( database: string, schemas: string[], defaultSchemaObjects: ObjectInfo[], - options: { simpleObjectDisplay?: boolean } = {}, + options: { lazyObjectTypes?: DatabaseObjectTreeKind[]; simpleObjectDisplay?: boolean } = {}, ): TreeNode[] { const databaseNodeId = `${connectionId}:${database}`; const defaultSchema = schemas.find(isDefaultSchema) || SQLSERVER_DEFAULT_SCHEMA; - const defaultObjectNodes = options.simpleObjectDisplay - ? buildSimpleObjectTreeNodes({ + const defaultObjectNodes = options.lazyObjectTypes + ? buildObjectGroupPlaceholderNodes({ nodeId: databaseNodeId, connectionId, database, schema: defaultSchema, - objects: defaultSchemaObjects, + objectTypes: options.lazyObjectTypes, }) - : buildGroupedObjectTreeNodes({ - nodeId: databaseNodeId, - connectionId, - database, - schema: defaultSchema, - objects: defaultSchemaObjects, - }); + : options.simpleObjectDisplay + ? buildSimpleObjectTreeNodes({ + nodeId: databaseNodeId, + connectionId, + database, + schema: defaultSchema, + objects: defaultSchemaObjects, + }) + : buildGroupedObjectTreeNodes({ + nodeId: databaseNodeId, + connectionId, + database, + schema: defaultSchema, + objects: defaultSchemaObjects, + }); const schemaNodes = sortSidebarNames(schemas.filter((schema) => !isDefaultSchema(schema))).map((schema) => ({ id: `${databaseNodeId}:${schema}`, diff --git a/apps/desktop/src/lib/tableTree.ts b/apps/desktop/src/lib/tableTree.ts index a6c54b474d..46b82ce1fd 100644 --- a/apps/desktop/src/lib/tableTree.ts +++ b/apps/desktop/src/lib/tableTree.ts @@ -1,4 +1,5 @@ import type { ObjectInfo, TableInfo, TreeNode, TreeNodeType } from "@/types/database"; +import { normalizeSidebarObjectKind, type SidebarObjectKind } from "@/lib/databaseObjectCapabilities"; const databaseObjectNameCollator = new Intl.Collator(undefined, { numeric: true, sensitivity: "base" }); @@ -315,7 +316,7 @@ export function hasTablePartitionGroups(node: TreeNode): boolean { return partitionGroupChildren(node).length > 0; } -type DatabaseObjectTreeKind = "TABLE" | "VIEW" | "PROCEDURE" | "FUNCTION" | "PACKAGE" | "PACKAGE_BODY"; +export type DatabaseObjectTreeKind = SidebarObjectKind; function buildObjectTreeEntries({ nodeId, @@ -369,12 +370,11 @@ export function buildSimpleObjectTreeNodes({ }): TreeNode[] { const seen = new Set(); const tableEntries: TableTreeEntry[] = []; - const viewNodes: TreeNode[] = []; + const objectNodes: TreeNode[] = []; for (const obj of objects) { const objectType = normalizeObjectType(obj.object_type); - if (objectType !== "TABLE" && objectType !== "VIEW" && objectType !== "PACKAGE" && objectType !== "PACKAGE_BODY") - continue; + if (!["TABLE", "VIEW", "PROCEDURE", "FUNCTION", "PACKAGE", "PACKAGE_BODY"].includes(objectType)) continue; const name = normalizeDatabaseObjectName(obj.name); if (!name) continue; @@ -397,13 +397,15 @@ export function buildSimpleObjectTreeNodes({ }); if (objectType === "TABLE") { tableEntries.push(entry); - } else if (objectType === "VIEW") { - viewNodes.push(entry.node); } else { - viewNodes.push({ - id: `${nodeId}:${childSchema ? `${childSchema}:` : ""}${name}:${objectType}`, + const simpleNodeType = simpleObjectNodeType(objectType); + objectNodes.push({ + id: + objectType === "VIEW" + ? entry.node.id + : `${nodeId}:${childSchema ? `${childSchema}:` : ""}${name}:${objectType}`, label: name, - type: objectType === "PACKAGE_BODY" ? "package-body" : "package", + type: simpleNodeType, comment: obj.comment, connectionId, database, @@ -416,18 +418,21 @@ export function buildSimpleObjectTreeNodes({ return [ ...buildPartitionTree(tableEntries, connectionId, database), - ...sortDatabaseObjectsByName(viewNodes, (node) => node.label), + ...sortDatabaseObjectsByName(objectNodes, (node) => node.label), ]; } +function simpleObjectNodeType(objectType: DatabaseObjectTreeKind): TreeNodeType { + if (objectType === "VIEW") return "view"; + if (objectType === "PROCEDURE") return "procedure"; + if (objectType === "FUNCTION") return "function"; + if (objectType === "PACKAGE_BODY") return "package-body"; + if (objectType === "PACKAGE") return "package"; + return "table"; +} + function normalizeObjectType(type: string): DatabaseObjectTreeKind { - const v = type.toUpperCase(); - if (v.includes("PACKAGE BODY") || v.includes("PACKAGE_BODY")) return "PACKAGE_BODY"; - if (v.includes("PACKAGE")) return "PACKAGE"; - if (v.includes("VIEW")) return "VIEW"; - if (v.includes("PROC")) return "PROCEDURE"; - if (v.includes("FUNC")) return "FUNCTION"; - return "TABLE"; + return normalizeSidebarObjectKind(type); } const groupDefs: Array<{ @@ -470,6 +475,34 @@ const objectGroupNodeTypes = new Set([ "group-packages", ]); +export function buildObjectGroupPlaceholderNodes({ + nodeId, + connectionId, + database, + schema, + objectTypes, +}: { + nodeId: string; + connectionId: string; + database: string; + schema?: string; + objectTypes: DatabaseObjectTreeKind[]; +}): TreeNode[] { + const supported = new Set(objectTypes); + return groupDefs + .filter((def) => def.objectTypes.some((objectType) => supported.has(objectType))) + .map((def) => ({ + id: `${nodeId}:${def.key}`, + label: def.label, + type: def.nodeType, + connectionId, + database, + schema, + isExpanded: false, + children: [], + })); +} + export function objectGroupRefreshParentId(node: TreeNode): string | null { if (!objectGroupNodeTypes.has(node.type)) return null; const suffixStart = node.id.lastIndexOf(":__"); @@ -477,6 +510,10 @@ export function objectGroupRefreshParentId(node: TreeNode): string | null { return node.id.slice(0, suffixStart); } +export function objectTypesForGroupNode(type: TreeNodeType): DatabaseObjectTreeKind[] | null { + return groupDefs.find((def) => def.nodeType === type)?.objectTypes ?? null; +} + export function buildGroupedObjectTreeNodes({ nodeId, connectionId, diff --git a/apps/desktop/src/lib/tauri.ts b/apps/desktop/src/lib/tauri.ts index a26e0ec079..ca68c22678 100644 --- a/apps/desktop/src/lib/tauri.ts +++ b/apps/desktop/src/lib/tauri.ts @@ -21,6 +21,7 @@ import type { SavedSqlFolder, SavedSqlLibrary, } from "@/types/database"; +import type { SidebarObjectKind } from "@/lib/databaseObjectCapabilities"; import type { AiConfig } from "@/stores/settingsStore"; import type { QueryEditability } from "@/lib/sqlAnalysis"; import type { @@ -452,8 +453,13 @@ export async function listTables( return invoke("list_tables", { connectionId, database, schema, filter, limit }); } -export async function listObjects(connectionId: string, database: string, schema: string): Promise { - return invoke("list_objects", { connectionId, database, schema }); +export async function listObjects( + connectionId: string, + database: string, + schema: string, + objectTypes?: SidebarObjectKind[], +): Promise { + return invoke("list_objects", { connectionId, database, schema, objectTypes }); } export async function listCompletionObjects( diff --git a/apps/desktop/src/stores/connectionStore.ts b/apps/desktop/src/stores/connectionStore.ts index b21ea9fe60..943c303bfb 100644 --- a/apps/desktop/src/stores/connectionStore.ts +++ b/apps/desktop/src/stores/connectionStore.ts @@ -20,11 +20,17 @@ import { 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"; +import { + isSchemaAware, + normalizeSidebarObjectKind, + sidebarObjectKindsForDatabase, + usesTreeSchemaMode, +} from "@/lib/databaseCapabilities"; import { connectionObjectTreeNodeSchema, connectionObjectTreeQuerySchema, connectionUsesDatabaseObjectTreeMode, + effectiveDatabaseTypeForConnection, } from "@/lib/jdbcDialect"; import { buildDatabaseTreeNodes, @@ -42,12 +48,16 @@ import { normalizeVisibleDatabaseSelection, } from "@/lib/visibleDatabases"; import { + buildObjectGroupPlaceholderNodes, buildGroupedObjectTreeNodes, + buildSimpleObjectTreeNodes, buildTableTreeNodes, expandCachedObjectBrowserNodes, mergeTableInfosIntoObjects, objectGroupRefreshParentId, + objectTypesForGroupNode, tablePartitionGroups, + type DatabaseObjectTreeKind, } from "@/lib/tableTree"; import { hasTreeNodeDatabaseContext, @@ -439,6 +449,11 @@ export const useConnectionStore = defineStore("connection", () => { return parts.map((part) => encodeURIComponent(part)).join(":"); } + function supportedSidebarObjectTypes(config?: ConnectionConfig): DatabaseObjectTreeKind[] { + const dbType = effectiveDatabaseTypeForConnection(config); + return sidebarObjectKindsForDatabase(dbType); + } + function refreshStaleTreeNode(node: TreeNode) { if (staleTreeRefreshIds.has(node.id)) return; staleTreeRefreshIds.add(node.id); @@ -1070,7 +1085,12 @@ export const useConnectionStore = defineStore("connection", () => { try { await ensureConnected(connectionId); if (useCachedChildren(node, options)) return; - const cacheKey = schemaCacheKey(connectionId, database, "sqlserver-objects"); + const simpleObjectDisplay = useSettingsStore().editorSettings.sidebarObjectDisplay === "simple"; + const cacheKey = schemaCacheKey( + connectionId, + database, + simpleObjectDisplay ? "sqlserver-objects-simple-v2" : "sqlserver-objects-grouped-v2", + ); if (!options?.force) { const cached = await loadPersistedTreeChildren(node, cacheKey); if (cached.hit) { @@ -1079,12 +1099,14 @@ export const useConnectionStore = defineStore("connection", () => { } } - const [schemas, defaultSchemaObjects] = await Promise.all([ - api.listSchemas(connectionId, database), - api.listObjects(connectionId, database, SQLSERVER_DEFAULT_SCHEMA), - ]); + const config = getConfig(connectionId); + const schemas = await api.listSchemas(connectionId, database); + const defaultSchemaObjects = simpleObjectDisplay + ? await api.listObjects(connectionId, database, SQLSERVER_DEFAULT_SCHEMA) + : []; const children = buildSqlServerDatabaseTreeNodes(connectionId, database, schemas, defaultSchemaObjects, { - simpleObjectDisplay: useSettingsStore().editorSettings.sidebarObjectDisplay === "simple", + lazyObjectTypes: simpleObjectDisplay ? undefined : supportedSidebarObjectTypes(config), + simpleObjectDisplay, }); setChildren(node, children); await savePersistedTreeChildren(cacheKey, children); @@ -1105,7 +1127,13 @@ export const useConnectionStore = defineStore("connection", () => { try { await ensureConnected(connectionId); if (useCachedChildren(node, options)) return; - const cacheKey = schemaCacheKey(connectionId, database, schema || "", "objects"); + const simpleObjectDisplay = useSettingsStore().editorSettings.sidebarObjectDisplay === "simple"; + const cacheKey = schemaCacheKey( + connectionId, + database, + schema || "", + simpleObjectDisplay ? "objects-simple-v2" : "objects-grouped-v2", + ); if (!options?.force) { const cached = await loadPersistedTreeChildren(node, cacheKey); if (cached.hit) { @@ -1117,18 +1145,14 @@ export const useConnectionStore = defineStore("connection", () => { const config = getConfig(connectionId); const querySchema = connectionObjectTreeQuerySchema(config, database, schema); const effectiveSchema = connectionObjectTreeNodeSchema(config, database, schema); - const simpleObjectDisplay = useSettingsStore().editorSettings.sidebarObjectDisplay === "simple"; let children: TreeNode[]; if (simpleObjectDisplay) { - const tables = await api.listTables(connectionId, database, querySchema); - children = buildTableTreeNodes({ nodeId, connectionId, database, schema: effectiveSchema, tables }); - } else { try { const [objects, tables] = await Promise.all([ api.listObjects(connectionId, database, querySchema), api.listTables(connectionId, database, querySchema), ]); - children = buildGroupedObjectTreeNodes({ + children = buildSimpleObjectTreeNodes({ nodeId, connectionId, database, @@ -1139,6 +1163,14 @@ export const useConnectionStore = defineStore("connection", () => { const tables = await api.listTables(connectionId, database, querySchema); children = buildTableTreeNodes({ nodeId, connectionId, database, schema: effectiveSchema, tables }); } + } else { + children = buildObjectGroupPlaceholderNodes({ + nodeId, + connectionId, + database, + schema: effectiveSchema, + objectTypes: supportedSidebarObjectTypes(config), + }); } setChildren(node, children); await savePersistedTreeChildren(cacheKey, children); @@ -1151,6 +1183,61 @@ export const useConnectionStore = defineStore("connection", () => { } } + async function loadObjectGroupChildren(node: TreeNode, options?: LoadTreeOptions) { + if (!node.connectionId || !hasTreeNodeDatabaseContext(node)) return; + node.isLoading = true; + try { + await ensureConnected(node.connectionId); + if (useCachedChildren(node, options)) return; + const objectTypes = objectTypesForGroupNode(node.type); + const parentNodeId = objectGroupRefreshParentId(node); + if (!objectTypes || !parentNodeId) return; + + 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"); + if (!options?.force) { + const cached = await loadPersistedTreeChildren(node, cacheKey); + if (cached.hit) { + if (cached.isStale) refreshStaleTreeNode(node); + return; + } + } + + 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; + setChildren(node, children); + await savePersistedTreeChildren(cacheKey, children); + node.isExpanded = true; + } catch (e) { + recordMetadataLoadError(node.connectionId, e); + throw e; + } finally { + node.isLoading = false; + } + } + + function normalizedObjectTreeKind(type: string): DatabaseObjectTreeKind { + return normalizeSidebarObjectKind(type); + } + async function loadTableGroups( connectionId: string, database: string, @@ -1431,9 +1518,10 @@ export const useConnectionStore = defineStore("connection", () => { node.type === "group-views" || node.type === "group-procedures" || node.type === "group-functions" || - node.type === "group-packages" || - node.type === "group-partitions" + node.type === "group-packages" ) { + await loadObjectGroupChildren(node, options); + } else if (node.type === "group-partitions") { node.isExpanded = true; } } @@ -1448,6 +1536,12 @@ export const useConnectionStore = defineStore("connection", () => { } async function refreshTreeNode(node: TreeNode) { + if (objectTypesForGroupNode(node.type)) { + clearLoadedChildrenCache(node.id); + await loadObjectGroupChildren(node, { force: true }); + return; + } + const parentId = objectGroupRefreshParentId(node); const parentNode = parentId ? findNode(treeNodes.value, parentId) : null; if (parentNode) { @@ -2166,6 +2260,7 @@ export const useConnectionStore = defineStore("connection", () => { loadSchemas, loadSqlServerDatabaseObjects, loadTables, + loadObjectGroupChildren, loadTableGroups, loadColumns, loadIndexes, diff --git a/crates/dbx-core/src/database_export.rs b/crates/dbx-core/src/database_export.rs index c43732ef88..5a1569bc9f 100644 --- a/crates/dbx-core/src/database_export.rs +++ b/crates/dbx-core/src/database_export.rs @@ -378,7 +378,8 @@ pub async fn export_database_sql_core( if request.include_objects && request.selected_tables.is_empty() { if let Ok(objects) = - crate::schema::list_objects_core(state, &request.connection_id, &request.database, &request.schema).await + crate::schema::list_objects_core(state, &request.connection_id, &request.database, &request.schema, None) + .await { for obj in &objects { let ot = obj.object_type.to_uppercase(); diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 2bc07a579e..4d650e13b2 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -2,6 +2,7 @@ use crate::connection::{connection_url_for_endpoint, database_connection_config, use crate::db; use crate::models::connection::{ConnectionConfig, DatabaseType}; use crate::query::{agent_execute_query_params, QueryExecutionOptions}; +use std::collections::HashSet; use std::future::Future; use std::sync::Arc; use std::time::Duration; @@ -555,7 +556,7 @@ fn filter_table_infos(tables: Vec, filter: Option<&str>, limit: O mod tests { use super::{ clickhouse_metadata_database, deduplicate_column_infos, duckdb_attach_database, duckdb_list_databases, - duckdb_query_tables_in_database, is_agent_postgres_metadata_fallback_config, + duckdb_query_tables_in_database, filter_objects_by_types, is_agent_postgres_metadata_fallback_config, }; use crate::models::connection::{ConnectionConfig, DatabaseType}; @@ -682,6 +683,49 @@ mod tests { assert!(!is_agent_postgres_metadata_fallback_config(&test_connection_config(DatabaseType::Postgres))); assert!(!is_agent_postgres_metadata_fallback_config(&test_connection_config(DatabaseType::Mysql))); } + + #[test] + fn filters_list_objects_by_normalized_object_types() { + let objects = vec![ + super::db::ObjectInfo { + name: "orders".to_string(), + object_type: "BASE TABLE".to_string(), + schema: None, + comment: None, + created_at: None, + updated_at: None, + parent_schema: None, + parent_name: None, + }, + super::db::ObjectInfo { + name: "active_orders".to_string(), + object_type: "MATERIALIZED VIEW".to_string(), + schema: None, + comment: None, + created_at: None, + updated_at: None, + parent_schema: None, + parent_name: None, + }, + super::db::ObjectInfo { + name: "payroll".to_string(), + object_type: "PACKAGE BODY".to_string(), + schema: None, + comment: None, + created_at: None, + updated_at: None, + parent_schema: None, + parent_name: None, + }, + ]; + + let filtered = filter_objects_by_types(objects, Some(&["VIEW".to_string(), "PACKAGE_BODY".to_string()])); + + assert_eq!( + filtered.iter().map(|object| object.name.as_str()).collect::>(), + ["active_orders", "payroll"] + ); + } } pub async fn list_objects_core( @@ -689,9 +733,10 @@ pub async fn list_objects_core( connection_id: &str, database: &str, schema: &str, + object_types: Option>, ) -> Result, String> { retry_metadata_connection(state, connection_id, Some(database), || { - list_objects_once(state, connection_id, database, schema) + list_objects_once(state, connection_id, database, schema, object_types.as_deref()) }) .await } @@ -713,6 +758,18 @@ async fn list_objects_once( connection_id: &str, database: &str, schema: &str, + object_types: Option<&[String]>, +) -> Result, String> { + list_objects_once_unfiltered(state, connection_id, database, schema) + .await + .map(|objects| filter_objects_by_types(objects, object_types)) +} + +async fn list_objects_once_unfiltered( + 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; @@ -832,6 +889,35 @@ async fn list_objects_once( } } +fn filter_objects_by_types(objects: Vec, object_types: Option<&[String]>) -> Vec { + let Some(object_types) = object_types else { + return objects; + }; + if object_types.is_empty() { + return objects; + } + let wanted: HashSet = + object_types.iter().map(|object_type| normalize_object_info_type(object_type)).collect(); + objects.into_iter().filter(|object| wanted.contains(&normalize_object_info_type(&object.object_type))).collect() +} + +fn normalize_object_info_type(object_type: &str) -> String { + let value = object_type.to_ascii_uppercase().replace(' ', "_"); + if value.contains("PACKAGE_BODY") { + "PACKAGE_BODY".to_string() + } else if value.contains("PACKAGE") { + "PACKAGE".to_string() + } else if value.contains("VIEW") { + "VIEW".to_string() + } else if value.contains("PROC") { + "PROCEDURE".to_string() + } else if value.contains("FUNC") { + "FUNCTION".to_string() + } else { + "TABLE".to_string() + } +} + async fn list_completion_objects_once( state: &AppState, connection_id: &str, @@ -919,7 +1005,7 @@ async fn list_completion_objects_once( 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?; + let objects = list_objects_once(state, connection_id, database, schema, None).await?; Ok(filter_completion_objects(objects)) } _ => Ok(Vec::new()), diff --git a/crates/dbx-web/src/routes/schema.rs b/crates/dbx-web/src/routes/schema.rs index 6e165a96dd..49c5647611 100644 --- a/crates/dbx-web/src/routes/schema.rs +++ b/crates/dbx-web/src/routes/schema.rs @@ -16,6 +16,7 @@ pub struct SchemaQuery { pub filter: Option, pub limit: Option, pub object_type: Option, + pub object_types: Option, } pub async fn list_databases( @@ -60,11 +61,17 @@ pub async fn list_objects( ) -> Result, AppError> { let database = q.database.as_deref().unwrap_or(""); let schema = q.schema.as_deref().unwrap_or(""); - let result = - dbx_core::schema::list_objects_core(&state.app, &q.connection_id, database, schema).await.map_err(AppError)?; + let object_types = q.object_types.as_deref().map(parse_object_types).filter(|types| !types.is_empty()); + let result = dbx_core::schema::list_objects_core(&state.app, &q.connection_id, database, schema, object_types) + .await + .map_err(AppError)?; Ok(Json(serde_json::to_value(result).map_err(|e| AppError(e.to_string()))?)) } +fn parse_object_types(value: &str) -> Vec { + value.split(',').map(str::trim).filter(|value| !value.is_empty()).map(ToString::to_string).collect() +} + pub async fn list_completion_objects( State(state): State>, Query(q): Query, diff --git a/packages/app-tests/databaseCapabilities.test.ts b/packages/app-tests/databaseCapabilities.test.ts index 4d89112d9d..4eb4fc2068 100644 --- a/packages/app-tests/databaseCapabilities.test.ts +++ b/packages/app-tests/databaseCapabilities.test.ts @@ -6,6 +6,7 @@ import { databaseObjectTreeNodeSchema, databaseObjectTreeQuerySchema, getDatabaseCapability, + sidebarObjectKindsForDatabase, supportsDatabaseCreation, supportsDatabaseSearch, supportsDriverManagement, @@ -266,3 +267,16 @@ test("object browser entry follows database tree shape", () => { assert.equal(supportsObjectBrowserTreeNode("jdbc", "database"), true); assert.equal(supportsObjectBrowserTreeNode("mongodb", "database"), false); }); + +test("sidebar object capability registry describes object groups by database type", () => { + assert.deepEqual(sidebarObjectKindsForDatabase("databend"), ["TABLE", "VIEW"]); + assert.deepEqual(sidebarObjectKindsForDatabase("postgres"), ["TABLE", "VIEW", "PROCEDURE", "FUNCTION"]); + assert.deepEqual(sidebarObjectKindsForDatabase("oracle"), [ + "TABLE", + "VIEW", + "PROCEDURE", + "FUNCTION", + "PACKAGE", + "PACKAGE_BODY", + ]); +}); diff --git a/packages/app-tests/sqlServerTree.test.ts b/packages/app-tests/sqlServerTree.test.ts index bd4416bbdc..71e8a56b05 100644 --- a/packages/app-tests/sqlServerTree.test.ts +++ b/packages/app-tests/sqlServerTree.test.ts @@ -11,11 +11,12 @@ function obj(name: string, objectType = "TABLE"): ObjectInfo { } test("SQL Server database tree groups dbo objects and keeps non-default schemas", () => { - const nodes = buildSqlServerDatabaseTreeNodes("conn", "app", ["dbo", "zeta", "sales"], [ - obj("customers"), - obj("customer_view", "VIEW"), - obj("get_total", "FUNCTION"), - ]); + const nodes = buildSqlServerDatabaseTreeNodes( + "conn", + "app", + ["dbo", "zeta", "sales"], + [obj("customers"), obj("customer_view", "VIEW"), obj("get_total", "FUNCTION")], + ); const topLevel = nodes.map((n) => ({ id: n.id, label: n.label, type: n.type })); assert.deepEqual(topLevel, [ @@ -52,10 +53,12 @@ test("SQL Server database tree shows only schemas when dbo has no objects", () = }); test("SQL Server database tree groups procedures alongside tables", () => { - const nodes = buildSqlServerDatabaseTreeNodes("conn", "app", ["dbo"], [ - obj("orders"), - obj("sp_refresh", "PROCEDURE"), - ]); + const nodes = buildSqlServerDatabaseTreeNodes( + "conn", + "app", + ["dbo"], + [obj("orders"), obj("sp_refresh", "PROCEDURE")], + ); assert.deepEqual( nodes.map((n) => ({ label: n.label, type: n.type })), @@ -65,3 +68,25 @@ test("SQL Server database tree groups procedures alongside tables", () => { ], ); }); + +test("SQL Server database tree can render lazy object groups before metadata is loaded", () => { + const nodes = buildSqlServerDatabaseTreeNodes("conn", "app", ["dbo", "sales"], [], { + lazyObjectTypes: ["TABLE", "VIEW", "PROCEDURE", "FUNCTION"], + }); + + assert.deepEqual( + nodes.map((n) => ({ label: n.label, type: n.type, count: n.objectCount, children: n.children })), + [ + { label: "tree.tables", type: "group-tables", count: undefined, children: [] }, + { label: "tree.views", type: "group-views", count: undefined, children: [] }, + { label: "tree.procedures", type: "group-procedures", count: undefined, children: [] }, + { label: "tree.functions", type: "group-functions", count: undefined, children: [] }, + { + label: "sales", + type: "schema", + count: undefined, + children: [], + }, + ], + ); +}); diff --git a/packages/app-tests/tableTree.test.ts b/packages/app-tests/tableTree.test.ts index 193fbcf5d3..cbe27a80f8 100644 --- a/packages/app-tests/tableTree.test.ts +++ b/packages/app-tests/tableTree.test.ts @@ -2,6 +2,8 @@ import test from "node:test"; import assert from "node:assert/strict"; import { buildGroupedObjectTreeNodes, + buildObjectGroupPlaceholderNodes, + buildSimpleObjectTreeNodes, buildTableTreeNodes, mergeTableInfosIntoObjects, } from "../../apps/desktop/src/lib/tableTree.ts"; @@ -169,6 +171,55 @@ test("buildGroupedObjectTreeNodes groups Oracle packages and package bodies", () ); }); +test("buildObjectGroupPlaceholderNodes creates capability-driven lazy sidebar groups", () => { + const groups = buildObjectGroupPlaceholderNodes({ + nodeId: "conn:app:HR", + connectionId: "conn", + database: "app", + schema: "HR", + objectTypes: ["TABLE", "VIEW", "PROCEDURE", "FUNCTION"], + }); + + assert.deepEqual( + groups.map((node) => ({ label: node.label, type: node.type, count: node.objectCount, children: node.children })), + [ + { label: "tree.tables", type: "group-tables", count: undefined, children: [] }, + { label: "tree.views", type: "group-views", count: undefined, children: [] }, + { label: "tree.procedures", type: "group-procedures", count: undefined, children: [] }, + { label: "tree.functions", type: "group-functions", count: undefined, children: [] }, + ], + ); +}); + +test("buildSimpleObjectTreeNodes keeps routines and packages visible in flat sidebar mode", () => { + const nodes = buildSimpleObjectTreeNodes({ + nodeId: "conn:app:HR", + connectionId: "conn", + database: "app", + schema: "HR", + objects: [ + { name: "ORDERS", object_type: "TABLE", schema: "HR" }, + { name: "ACTIVE_ORDERS", object_type: "VIEW", schema: "HR" }, + { name: "REFRESH_STATS", object_type: "PROCEDURE", schema: "HR" }, + { name: "TOTAL_DUE", object_type: "FUNCTION", schema: "HR" }, + { name: "PAYROLL", object_type: "PACKAGE", schema: "HR" }, + { name: "PAYROLL", object_type: "PACKAGE_BODY", schema: "HR" }, + ], + }); + + assert.deepEqual( + nodes.map((node) => ({ label: node.label, type: node.type })), + [ + { label: "ORDERS", type: "table" }, + { label: "ACTIVE_ORDERS", type: "view" }, + { label: "PAYROLL", type: "package" }, + { label: "PAYROLL", type: "package-body" }, + { label: "REFRESH_STATS", type: "procedure" }, + { label: "TOTAL_DUE", type: "function" }, + ], + ); +}); + test("mergeTableInfosIntoObjects restores views missing from object metadata", () => { const merged = mergeTableInfosIntoObjects( [object("orders")], diff --git a/src-tauri/src/commands/schema.rs b/src-tauri/src/commands/schema.rs index 312651b33e..9b4c197e69 100644 --- a/src-tauri/src/commands/schema.rs +++ b/src-tauri/src/commands/schema.rs @@ -39,8 +39,9 @@ pub async fn list_objects( connection_id: String, database: String, schema: String, + object_types: Option>, ) -> Result, String> { - dbx_core::schema::list_objects_core(&state, &connection_id, &database, &schema).await + dbx_core::schema::list_objects_core(&state, &connection_id, &database, &schema, object_types).await } #[tauri::command]