diff --git a/crates/dbx-core/src/db/clickhouse_driver.rs b/crates/dbx-core/src/db/clickhouse_driver.rs index 0045ff1d72..7052a15f4e 100644 --- a/crates/dbx-core/src/db/clickhouse_driver.rs +++ b/crates/dbx-core/src/db/clickhouse_driver.rs @@ -80,7 +80,14 @@ pub async fn test_connection(client: &ChClient) -> Result<(), String> { } pub async fn list_databases(client: &ChClient) -> Result, String> { - let result = ch_query(client, "SELECT name FROM system.databases ORDER BY name", None).await?; + let result = ch_query( + client, + "SELECT name FROM system.databases \ + WHERE name NOT IN ('system', 'INFORMATION_SCHEMA', 'information_schema') \ + ORDER BY name", + None, + ) + .await?; Ok(result.data.iter().map(|row| DatabaseInfo { name: row[0].as_str().unwrap_or("").to_string() }).collect()) } diff --git a/crates/dbx-core/src/db/gaussdb_driver.rs b/crates/dbx-core/src/db/gaussdb_driver.rs index e473fc2d01..13e0112b18 100644 --- a/crates/dbx-core/src/db/gaussdb_driver.rs +++ b/crates/dbx-core/src/db/gaussdb_driver.rs @@ -53,7 +53,11 @@ fn normalize_database(database: &str) -> &str { pub async fn list_databases(client: &mut GaussdbClient) -> Result, String> { let rows = client - .query_single_column("SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname") + .query_single_column( + "SELECT datname FROM pg_database \ + WHERE datistemplate = false AND datallowconn = true \ + ORDER BY datname", + ) .await?; Ok(rows.into_iter().map(|name| DatabaseInfo { name }).collect()) } diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index b47f777082..e58a80e60e 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -183,10 +183,14 @@ pub async fn connect_bare(url: &str) -> Result { } pub async fn list_databases(pool: &MySqlPool) -> Result, String> { - let rows: Vec = sqlx::raw_sql("SELECT SCHEMA_NAME FROM information_schema.SCHEMATA ORDER BY SCHEMA_NAME") - .fetch_all(pool) - .await - .map_err(|e| e.to_string())?; + let rows: Vec = sqlx::raw_sql( + "SELECT SCHEMA_NAME FROM information_schema.SCHEMATA \ + WHERE SCHEMA_NAME NOT IN ('information_schema', 'mysql', 'performance_schema', 'sys') \ + ORDER BY SCHEMA_NAME", + ) + .fetch_all(pool) + .await + .map_err(|e| e.to_string())?; Ok(rows.iter().map(|row| DatabaseInfo { name: get_str(row, 0) }).collect()) } diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 2484db1b82..bacf8ddf04 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -153,10 +153,14 @@ fn validate_postgres_ssl_paths(url: &str) -> Result<(), String> { } pub async fn list_databases(pool: &PgPool) -> Result, String> { - let rows: Vec = sqlx::query("SELECT datname FROM pg_database WHERE datistemplate = false ORDER BY datname") - .fetch_all(pool) - .await - .map_err(|e| e.to_string())?; + let rows: Vec = sqlx::query( + "SELECT datname FROM pg_database \ + WHERE datistemplate = false AND datallowconn = true \ + ORDER BY datname", + ) + .fetch_all(pool) + .await + .map_err(|e| e.to_string())?; Ok(rows.iter().map(|row| DatabaseInfo { name: row.get::("datname") }).collect()) } diff --git a/crates/dbx-core/src/db/sqlserver.rs b/crates/dbx-core/src/db/sqlserver.rs index 9a7366d7ba..f46a5725d4 100644 --- a/crates/dbx-core/src/db/sqlserver.rs +++ b/crates/dbx-core/src/db/sqlserver.rs @@ -75,7 +75,17 @@ fn row_to_json(row: &tiberius::Row) -> Vec { } pub async fn list_databases(client: &mut SqlServerClient) -> Result, String> { - let stream = client.query("SELECT name FROM sys.databases ORDER BY name", &[]).await.map_err(|e| e.to_string())?; + let stream = client + .query( + "SELECT name \ + FROM sys.databases \ + WHERE database_id > 4 \ + AND state = 0 \ + ORDER BY name", + &[], + ) + .await + .map_err(|e| e.to_string())?; let rows = stream.into_first_result().await.map_err(|e| e.to_string())?; Ok(rows.iter().map(|row| DatabaseInfo { name: row.get::<&str, _>(0).unwrap_or("").to_string() }).collect()) } @@ -83,9 +93,16 @@ pub async fn list_databases(client: &mut SqlServerClient) -> Result Result, String> { let stream = client .query( - "SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA \ - WHERE SCHEMA_NAME NOT IN ('guest','INFORMATION_SCHEMA','sys') \ - ORDER BY SCHEMA_NAME", + "SELECT s.name \ + FROM sys.schemas s \ + WHERE s.name NOT IN ('guest','INFORMATION_SCHEMA','sys') \ + AND EXISTS ( \ + SELECT 1 FROM sys.objects o \ + WHERE o.schema_id = s.schema_id \ + AND o.type IN ('U','V') \ + AND o.is_ms_shipped = 0 \ + ) \ + ORDER BY CASE WHEN s.name = 'dbo' THEN 0 ELSE 1 END, s.name", &[], ) .await @@ -96,7 +113,13 @@ pub async fn list_schemas(client: &mut SqlServerClient) -> Result, S pub async fn list_tables(client: &mut SqlServerClient, schema: &str) -> Result, String> { let sql = format!( - "SELECT TABLE_NAME, TABLE_TYPE FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = '{}' ORDER BY TABLE_NAME", + "SELECT 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 \ + ORDER BY o.name", schema.replace('\'', "''") ); let stream = client.query(&*sql, &[]).await.map_err(|e| e.to_string())?; diff --git a/src/components/sidebar/TreeItem.vue b/src/components/sidebar/TreeItem.vue index 5002d464c7..627e80277c 100644 --- a/src/components/sidebar/TreeItem.vue +++ b/src/components/sidebar/TreeItem.vue @@ -220,7 +220,9 @@ async function toggle() { queryStore.updateSql(tab, node.label); } else if (node.type === "database" && node.connectionId && node.database) { const config = connectionStore.getConfig(node.connectionId); - if (config?.db_type && TREE_SCHEMA_TYPES.has(config.db_type)) { + if (config?.db_type === "sqlserver") { + await connectionStore.loadSqlServerDatabaseObjects(node.connectionId, node.database); + } else if (config?.db_type && TREE_SCHEMA_TYPES.has(config.db_type)) { await connectionStore.loadSchemas(node.connectionId, node.database); } else { await connectionStore.loadTables(node.connectionId, node.database); @@ -443,6 +445,18 @@ function dropTable() { showDropTableConfirm.value = true; } +async function refreshTableList(node: TreeNode) { + if (!node.connectionId || !node.database) return; + const config = connectionStore.getConfig(node.connectionId); + if (config?.db_type === "sqlserver" && node.schema?.toLowerCase() === "dbo") { + await connectionStore.loadSqlServerDatabaseObjects(node.connectionId, node.database, { force: true }); + } else if (node.schema) { + await connectionStore.loadTables(node.connectionId, node.database, node.schema, { force: true }); + } else { + await connectionStore.loadTables(node.connectionId, node.database, undefined, { force: true }); + } +} + async function confirmDropTable() { const node = props.node; if (!node.connectionId || !node.database) return; @@ -450,11 +464,7 @@ async function confirmDropTable() { await connectionStore.ensureConnected(node.connectionId); await api.executeQuery(node.connectionId, node.database, buildDropTableSql(), node.schema); toast(t("contextMenu.dropTableSuccess", { name: node.label }), 3000); - if (node.schema) { - await connectionStore.loadTables(node.connectionId, node.database, node.schema, { force: true }); - } else { - await connectionStore.loadTables(node.connectionId, node.database, undefined, { force: true }); - } + await refreshTableList(node); } catch (e: any) { toast(t("contextMenu.tableOperationFailed", { message: e?.message || String(e) }), 5000); } @@ -556,7 +566,12 @@ async function confirmCreateSchema() { const sql = `CREATE SCHEMA ${quoteIdent(name)};`; await api.executeQuery(node.connectionId, node.database, sql); toast(t("contextMenu.createSchemaSuccess", { name }), 3000); - await connectionStore.loadSchemas(node.connectionId, node.database, { force: true }); + const config = connectionStore.getConfig(node.connectionId); + if (config?.db_type === "sqlserver") { + await connectionStore.loadSqlServerDatabaseObjects(node.connectionId, node.database, { force: true }); + } else { + await connectionStore.loadSchemas(node.connectionId, node.database, { force: true }); + } } catch (e: any) { toast(t("contextMenu.tableOperationFailed", { message: e?.message || String(e) }), 5000); } @@ -573,7 +588,12 @@ async function confirmDropSchema() { await connectionStore.ensureConnected(node.connectionId); await api.executeQuery(node.connectionId, node.database, buildDropSchemaSql()); toast(t("contextMenu.dropSchemaSuccess", { name: node.label }), 3000); - await connectionStore.loadSchemas(node.connectionId, node.database, { force: true }); + const config = connectionStore.getConfig(node.connectionId); + if (config?.db_type === "sqlserver") { + await connectionStore.loadSqlServerDatabaseObjects(node.connectionId, node.database, { force: true }); + } else { + await connectionStore.loadSchemas(node.connectionId, node.database, { force: true }); + } } catch (e: any) { toast(t("contextMenu.tableOperationFailed", { message: e?.message || String(e) }), 5000); } @@ -608,11 +628,7 @@ async function confirmDuplicateStructure() { } await api.executeQuery(node.connectionId, node.database, sql, node.schema); toast(t("contextMenu.duplicateStructureSuccess", { name: newName }), 3000); - if (node.schema) { - await connectionStore.loadTables(node.connectionId, node.database, node.schema, { force: true }); - } else { - await connectionStore.loadTables(node.connectionId, node.database, undefined, { force: true }); - } + await refreshTableList(node); } catch (e: any) { toast(t("contextMenu.tableOperationFailed", { message: e?.message || String(e) }), 5000); } diff --git a/src/lib/sqlServerTree.ts b/src/lib/sqlServerTree.ts new file mode 100644 index 0000000000..c6afc226f0 --- /dev/null +++ b/src/lib/sqlServerTree.ts @@ -0,0 +1,42 @@ +import type { TableInfo, TreeNode } from "@/types/database"; + +export const SQLSERVER_DEFAULT_SCHEMA = "dbo"; + +function isDefaultSchema(schema: string): boolean { + return schema.toLowerCase() === SQLSERVER_DEFAULT_SCHEMA; +} + +export function buildSqlServerDatabaseTreeNodes( + connectionId: string, + database: string, + schemas: string[], + defaultSchemaTables: TableInfo[], +): TreeNode[] { + const databaseNodeId = `${connectionId}:${database}`; + const defaultSchema = schemas.find(isDefaultSchema) || SQLSERVER_DEFAULT_SCHEMA; + const defaultTableNodes = defaultSchemaTables.map((table) => ({ + id: `${databaseNodeId}:${defaultSchema}:${table.name}`, + label: table.name, + type: (table.table_type === "VIEW" ? "view" : "table") as "view" | "table", + connectionId, + database, + schema: defaultSchema, + isExpanded: false, + children: [], + })); + + const schemaNodes = schemas + .filter((schema) => !isDefaultSchema(schema)) + .map((schema) => ({ + id: `${databaseNodeId}:${schema}`, + label: schema, + type: "schema" as const, + connectionId, + database, + schema, + isExpanded: false, + children: [], + })); + + return [...defaultTableNodes, ...schemaNodes]; +} diff --git a/src/stores/connectionStore.ts b/src/stores/connectionStore.ts index a8822f2ccd..91aecb5a2a 100644 --- a/src/stores/connectionStore.ts +++ b/src/stores/connectionStore.ts @@ -22,6 +22,7 @@ import * as api from "@/lib/api"; import { isTauriRuntime } from "@/lib/tauriRuntime"; import { isSchemaAware, TREE_SCHEMA_TYPES } from "@/lib/databaseCapabilities"; import { buildDatabaseTreeNodes } from "@/lib/databaseTree"; +import { buildSqlServerDatabaseTreeNodes, SQLSERVER_DEFAULT_SCHEMA } from "@/lib/sqlServerTree"; const PINNED_TREE_NODES_STORAGE_KEY = "dbx-pinned-tree-nodes"; @@ -536,6 +537,30 @@ export const useConnectionStore = defineStore("connection", () => { } } + async function loadSqlServerDatabaseObjects(connectionId: string, database: string, options?: LoadTreeOptions) { + const nodeId = `${connectionId}:${database}`; + const node = findNode(treeNodes.value, nodeId); + if (!node) return; + if (useCachedChildren(node, options)) return; + const cacheKey = schemaCacheKey(connectionId, database, "sqlserver-objects"); + if (!options?.force && (await loadPersistedTreeChildren(node, cacheKey))) return; + + node.isLoading = true; + try { + await ensureConnected(connectionId); + const [schemas, defaultSchemaTables] = await Promise.all([ + api.listSchemas(connectionId, database), + api.listTables(connectionId, database, SQLSERVER_DEFAULT_SCHEMA), + ]); + const children = buildSqlServerDatabaseTreeNodes(connectionId, database, schemas, defaultSchemaTables); + setChildren(node, children); + await savePersistedTreeChildren(cacheKey, children); + node.isExpanded = true; + } finally { + node.isLoading = false; + } + } + async function loadTables(connectionId: string, database: string, schema?: string, options?: LoadTreeOptions) { const nodeId = schema ? `${connectionId}:${database}:${schema}` : `${connectionId}:${database}`; const node = findNode(treeNodes.value, nodeId); @@ -762,7 +787,9 @@ export const useConnectionStore = defineStore("connection", () => { await loadMongoCollections(node.connectionId, node.database); } else if (node.type === "database" && node.connectionId && node.database) { const config = getConfig(node.connectionId); - if (config?.db_type && TREE_SCHEMA_TYPES.has(config.db_type)) { + if (config?.db_type === "sqlserver") { + await loadSqlServerDatabaseObjects(node.connectionId, node.database, options); + } else if (config?.db_type && TREE_SCHEMA_TYPES.has(config.db_type)) { await loadSchemas(node.connectionId, node.database, options); } else { await loadTables(node.connectionId, node.database, undefined, options); @@ -1112,6 +1139,7 @@ export const useConnectionStore = defineStore("connection", () => { loadMongoDatabases, loadMongoCollections, loadSchemas, + loadSqlServerDatabaseObjects, loadTables, loadTableGroups, loadColumns, diff --git a/tests/sqlServerTree.test.ts b/tests/sqlServerTree.test.ts new file mode 100644 index 0000000000..beeb34d68b --- /dev/null +++ b/tests/sqlServerTree.test.ts @@ -0,0 +1,38 @@ +import test from "node:test"; +import assert from "node:assert/strict"; +import { buildSqlServerDatabaseTreeNodes } from "../src/lib/sqlServerTree.ts"; +import type { TableInfo } from "../src/types/database.ts"; + +function table(name: string, tableType = "BASE TABLE"): TableInfo { + return { + name, + table_type: tableType, + }; +} + +test("SQL Server database tree flattens dbo tables and keeps non-default schemas", () => { + const nodes = buildSqlServerDatabaseTreeNodes( + "conn", + "app", + ["dbo", "sales"], + [table("customers"), table("customer_view", "VIEW")], + ); + + assert.deepEqual( + nodes.map((node) => ({ id: node.id, label: node.label, type: node.type, schema: node.schema })), + [ + { id: "conn:app:dbo:customers", label: "customers", type: "table", schema: "dbo" }, + { id: "conn:app:dbo:customer_view", label: "customer_view", type: "view", schema: "dbo" }, + { id: "conn:app:sales", label: "sales", type: "schema", schema: "sales" }, + ], + ); +}); + +test("SQL Server database tree hides default schema node when dbo has no tables", () => { + const nodes = buildSqlServerDatabaseTreeNodes("conn", "app", ["dbo", "archive"], []); + + assert.deepEqual( + nodes.map((node) => ({ id: node.id, label: node.label, type: node.type, schema: node.schema })), + [{ id: "conn:app:archive", label: "archive", type: "schema", schema: "archive" }], + ); +});