diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index a9be7f887..d95ada886 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -738,6 +738,16 @@ pub async fn get_object_source_core( db::sqlserver::execute_query(&mut client, &sqlserver_object_source_sql(schema, name, &object_type)) .await?, )? + } else if let Some(client) = extract_agent(&connections, &pool_key) { + drop(connections); + let mut client = client.lock().await; + let result: db::ObjectSource = client + .call( + "get_object_source", + serde_json::json!({"schema": schema, "name": name, "object_type": object_type}), + ) + .await?; + return Ok(result); } else { match connections.get(&pool_key).ok_or("Pool not found")? { PoolKind::Mysql(pool, _) => mysql_object_source(pool, name, &object_type).await?, diff --git a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java index 82fb2286e..9e94d64c4 100644 --- a/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java +++ b/plugins/jdbc/src/main/java/app/dbx/jdbc/DbxJdbcPlugin.java @@ -104,6 +104,7 @@ public final class DbxJdbcPlugin { case "listDatabases" -> listDatabases(connection); case "listSchemas" -> listSchemas(connection, optionalText(params, "database")); case "listTables" -> listTables(connection, optionalText(params, "database"), optionalText(params, "schema")); + case "list_objects" -> listObjects(connection, optionalText(params, "database"), optionalText(params, "schema")); case "getColumns" -> getColumns( connection, optionalText(params, "database"), @@ -285,6 +286,61 @@ public final class DbxJdbcPlugin { return result; } + private static JsonNode listObjects(JsonNode connection, String database, String schema) throws SQLException { + ArrayNode result = MAPPER.createArrayNode(); + Connection conn = openConnection(connection); + DatabaseMetaData meta = conn.getMetaData(); + String catalog = emptyToNull(database); + String schemaPattern = emptyToNull(schema); + + String[] tableTypes = new String[] {"TABLE", "VIEW", "MATERIALIZED VIEW", "SYSTEM TABLE", "SYSTEM VIEW"}; + try (ResultSet rs = meta.getTables(catalog, schemaPattern, "%", tableTypes)) { + while (rs.next()) { + ObjectNode item = MAPPER.createObjectNode(); + item.put("name", rs.getString("TABLE_NAME")); + item.put("object_type", rs.getString("TABLE_TYPE")); + putNullable(item, "schema", schema); + putNullable(item, "comment", rs.getString("REMARKS")); + result.add(item); + } + } + + try (ResultSet rs = meta.getProcedures(catalog, schemaPattern, "%")) { + while (rs.next()) { + ObjectNode item = MAPPER.createObjectNode(); + item.put("name", rs.getString("PROCEDURE_NAME")); + item.put("object_type", "PROCEDURE"); + putNullable(item, "schema", schema); + putNullable(item, "comment", rs.getString("REMARKS")); + result.add(item); + } + } catch (SQLException ignored) { + } + + Set procedureNames = new HashSet<>(); + for (JsonNode node : result) { + if ("PROCEDURE".equals(node.path("object_type").asText())) { + procedureNames.add(node.path("name").asText()); + } + } + try (ResultSet rs = meta.getFunctions(catalog, schemaPattern, "%")) { + while (rs.next()) { + String name = rs.getString("FUNCTION_NAME"); + if (!procedureNames.contains(name)) { + ObjectNode item = MAPPER.createObjectNode(); + item.put("name", name); + item.put("object_type", "FUNCTION"); + putNullable(item, "schema", schema); + putNullable(item, "comment", rs.getString("REMARKS")); + result.add(item); + } + } + } catch (SQLException ignored) { + } + + return result; + } + private static JsonNode getColumns(JsonNode connection, String database, String schema, String table) throws SQLException { ArrayNode result = MAPPER.createArrayNode(); Connection conn = openConnection(connection); diff --git a/src/components/sidebar/TreeItem.vue b/src/components/sidebar/TreeItem.vue index caeb4e0e4..f0c23dd19 100644 --- a/src/components/sidebar/TreeItem.vue +++ b/src/components/sidebar/TreeItem.vue @@ -322,6 +322,8 @@ function runRowClickAction() { const action = treeNodeRowAction(node.type, canExpand.value); if (action === "open-data") { openData(); + } else if (node.type === "procedure" || node.type === "function") { + void viewObjectSource(); } else if (node.type === "saved-sql-file") { openSavedSqlFile(); } else if (action === "toggle") { @@ -575,7 +577,23 @@ function buildDropObjectSql(): string { } function viewObjectSource() { - void openObjectBrowser(); + const node = props.node; + if (!node.connectionId || !node.database) return; + const objectType = node.type === "procedure" ? "PROCEDURE" : "FUNCTION"; + const schema = node.schema || node.database; + connectionStore + .ensureConnected(node.connectionId) + .then(() => { + connectionStore.activeConnectionId = node.connectionId!; + return api.getObjectSource(node.connectionId!, node.database!, schema, node.label, objectType as any); + }) + .then((result) => { + const tabId = queryStore.createTab(node.connectionId!, node.database!, node.label); + queryStore.updateSql(tabId, result.source); + }) + .catch((e: any) => { + toast(e?.message || String(e), 5000); + }); } function requestDropObject() { diff --git a/src/lib/sqlServerTree.ts b/src/lib/sqlServerTree.ts index c6afc226f..9024ecf43 100644 --- a/src/lib/sqlServerTree.ts +++ b/src/lib/sqlServerTree.ts @@ -1,4 +1,5 @@ -import type { TableInfo, TreeNode } from "@/types/database"; +import type { ObjectInfo, TreeNode } from "@/types/database"; +import { buildGroupedObjectTreeNodes } from "@/lib/tableTree"; export const SQLSERVER_DEFAULT_SCHEMA = "dbo"; @@ -10,20 +11,18 @@ export function buildSqlServerDatabaseTreeNodes( connectionId: string, database: string, schemas: string[], - defaultSchemaTables: TableInfo[], + defaultSchemaObjects: ObjectInfo[], ): 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", + + const defaultObjectNodes = buildGroupedObjectTreeNodes({ + nodeId: databaseNodeId, connectionId, database, schema: defaultSchema, - isExpanded: false, - children: [], - })); + objects: defaultSchemaObjects, + }); const schemaNodes = schemas .filter((schema) => !isDefaultSchema(schema)) @@ -38,5 +37,5 @@ export function buildSqlServerDatabaseTreeNodes( children: [], })); - return [...defaultTableNodes, ...schemaNodes]; + return [...defaultObjectNodes, ...schemaNodes]; } diff --git a/src/lib/tableTree.ts b/src/lib/tableTree.ts index bf9d3db4d..1ec12253d 100644 --- a/src/lib/tableTree.ts +++ b/src/lib/tableTree.ts @@ -108,7 +108,7 @@ export function buildGroupedObjectTreeNodes({ objectCount: items.length, isExpanded: false, children: items.map((obj) => ({ - id: `${nodeId}:${obj.name}`, + id: `${nodeId}:${def.key}:${obj.name}`, label: obj.name, type: def.childType, connectionId, diff --git a/src/stores/connectionStore.ts b/src/stores/connectionStore.ts index 96337a853..280f196bc 100644 --- a/src/stores/connectionStore.ts +++ b/src/stores/connectionStore.ts @@ -736,11 +736,11 @@ export const useConnectionStore = defineStore("connection", () => { const cacheKey = schemaCacheKey(connectionId, database, "sqlserver-objects"); if (!options?.force && (await loadPersistedTreeChildren(node, cacheKey))) return; - const [schemas, defaultSchemaTables] = await Promise.all([ + const [schemas, defaultSchemaObjects] = await Promise.all([ api.listSchemas(connectionId, database), - api.listTables(connectionId, database, SQLSERVER_DEFAULT_SCHEMA), + api.listObjects(connectionId, database, SQLSERVER_DEFAULT_SCHEMA), ]); - const children = buildSqlServerDatabaseTreeNodes(connectionId, database, schemas, defaultSchemaTables); + const children = buildSqlServerDatabaseTreeNodes(connectionId, database, schemas, defaultSchemaObjects); setChildren(node, children); await savePersistedTreeChildren(cacheKey, children); node.isExpanded = true; @@ -779,8 +779,15 @@ export const useConnectionStore = defineStore("connection", () => { } } - async function loadTableGroups(connectionId: string, database: string, table: string, schema?: string) { - const parentId = schema ? `${connectionId}:${database}:${schema}:${table}` : `${connectionId}:${database}:${table}`; + async function loadTableGroups( + connectionId: string, + database: string, + table: string, + schema?: string, + nodeId?: string, + ) { + const parentId = + nodeId ?? (schema ? `${connectionId}:${database}:${schema}:${table}` : `${connectionId}:${database}:${table}`); const node = findNode(treeNodes.value, parentId); if (!node) return; @@ -982,7 +989,7 @@ export const useConnectionStore = defineStore("connection", () => { } else if (node.type === "schema" && node.connectionId && node.database && node.schema) { await loadTables(node.connectionId, node.database, node.schema, options); } else if ((node.type === "table" || node.type === "view") && node.connectionId && node.database) { - await loadTableGroups(node.connectionId, node.database, node.label, node.schema); + await loadTableGroups(node.connectionId, node.database, node.label, node.schema, node.id); } else if (node.type === "group-columns" && node.connectionId && node.database && node.tableName) { await loadColumns(node.connectionId, node.database, node.tableName, node.schema); } else if (node.type === "group-indexes" && node.connectionId && node.database && node.tableName) { diff --git a/tests/sqlServerTree.test.ts b/tests/sqlServerTree.test.ts index beeb34d68..cab2d3715 100644 --- a/tests/sqlServerTree.test.ts +++ b/tests/sqlServerTree.test.ts @@ -1,38 +1,63 @@ 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"; +import type { ObjectInfo } from "../src/types/database.ts"; -function table(name: string, tableType = "BASE TABLE"): TableInfo { +function obj(name: string, objectType = "TABLE"): ObjectInfo { return { name, - table_type: tableType, + object_type: objectType, }; } -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")], +test("SQL Server database tree groups dbo objects and keeps non-default schemas", () => { + const nodes = buildSqlServerDatabaseTreeNodes("conn", "app", ["dbo", "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, [ + { id: "conn:app:__tables", label: "tree.tables", type: "group-tables" }, + { id: "conn:app:__views", label: "tree.views", type: "group-views" }, + { id: "conn:app:__functions", label: "tree.functions", type: "group-functions" }, + { id: "conn:app:sales", label: "sales", type: "schema" }, + ]); + + const tableGroup = nodes.find((n) => n.type === "group-tables"); + assert.deepEqual( + tableGroup?.children?.map((c) => c.label), + ["customers"], ); + const functionGroup = nodes.find((n) => n.type === "group-functions"); 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" }, - ], + functionGroup?.children?.map((c) => c.label), + ["get_total"], ); }); -test("SQL Server database tree hides default schema node when dbo has no tables", () => { +test("SQL Server database tree shows only schemas when dbo has no objects", () => { 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" }], + nodes.map((node) => ({ id: node.id, label: node.label, type: node.type })), + [{ id: "conn:app:archive", label: "archive", type: "schema" }], + ); +}); + +test("SQL Server database tree groups procedures alongside tables", () => { + const nodes = buildSqlServerDatabaseTreeNodes("conn", "app", ["dbo"], [ + obj("orders"), + obj("sp_refresh", "PROCEDURE"), + ]); + + assert.deepEqual( + nodes.map((n) => ({ label: n.label, type: n.type })), + [ + { label: "tree.tables", type: "group-tables" }, + { label: "tree.procedures", type: "group-procedures" }, + ], ); });