diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 3e4793bec..ce979ccc7 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -187,6 +187,58 @@ "enablement": "databricks.context.activated && databricks.context.loggedIn && databricks.feature.views.workspace", "category": "Databricks" }, + { + "command": "databricks.unityCatalog.filter", + "title": "Filter", + "icon": "$(search)", + "category": "Databricks" + }, + { + "command": "databricks.unityCatalog.refresh", + "title": "Refresh Unity Catalog view", + "icon": "$(refresh)", + "enablement": "databricks.context.activated && databricks.context.loggedIn", + "category": "Databricks" + }, + { + "command": "databricks.unityCatalog.copyStorageLocation", + "title": "Copy storage location", + "category": "Databricks" + }, + { + "command": "databricks.unityCatalog.copyViewSql", + "title": "Copy view SQL", + "category": "Databricks" + }, + { + "command": "databricks.unityCatalog.copyName", + "title": "Copy", + "category": "Databricks" + }, + { + "command": "databricks.unityCatalog.refreshNode", + "title": "Refresh", + "icon": "$(refresh)", + "category": "Databricks" + }, + { + "command": "databricks.unityCatalog.pinSchema", + "title": "Pin Schema", + "icon": "$(star-empty)", + "category": "Databricks" + }, + { + "command": "databricks.unityCatalog.unpinSchema", + "title": "Unpin Schema", + "icon": "$(star-full)", + "category": "Databricks" + }, + { + "command": "databricks.unityCatalog.openExternal", + "title": "Open in Databricks", + "icon": "$(link-external)", + "category": "Databricks" + }, { "command": "databricks.call", "title": "Call", @@ -448,12 +500,107 @@ "name": "Workspace explorer", "when": "databricks.feature.views.workspace" }, + { + "id": "unityCatalogView", + "name": "Unity Catalog", + "when": "databricks.context.activated && databricks.context.loggedIn" + }, { "id": "databricksDocsView", "name": "Documentation" } ] }, + "colors": [ + { + "id": "databricks.unityCatalog.catalog", + "description": "Icon color for Unity Catalog catalog nodes", + "defaults": { + "dark": "#9B6CF7", + "light": "#6B2FD4", + "highContrast": "#C586C0", + "highContrastLight": "#6B2FD4" + } + }, + { + "id": "databricks.unityCatalog.schema", + "description": "Icon color for Unity Catalog schema nodes", + "defaults": { + "dark": "#0DB7C4", + "light": "#007A85", + "highContrast": "#4EC9B0", + "highContrastLight": "#007A85" + } + }, + { + "id": "databricks.unityCatalog.table", + "description": "Icon color for Unity Catalog table nodes", + "defaults": { + "dark": "#FF6B2C", + "light": "#C84B0A", + "highContrast": "#FF6B2C", + "highContrastLight": "#C84B0A" + } + }, + { + "id": "databricks.unityCatalog.volume", + "description": "Icon color for Unity Catalog volume nodes", + "defaults": { + "dark": "#4FC1E9", + "light": "#0E6FA0", + "highContrast": "#4FC1E9", + "highContrastLight": "#0E6FA0" + } + }, + { + "id": "databricks.unityCatalog.function", + "description": "Icon color for Unity Catalog function nodes", + "defaults": { + "dark": "#FFB347", + "light": "#A06000", + "highContrast": "#FFCA28", + "highContrastLight": "#A06000" + } + }, + { + "id": "databricks.unityCatalog.columnKey", + "description": "Icon color for Unity Catalog non-nullable (key) column nodes", + "defaults": { + "dark": "#F47C7C", + "light": "#C0392B", + "highContrast": "#F47C7C", + "highContrastLight": "#C0392B" + } + }, + { + "id": "databricks.unityCatalog.column", + "description": "Icon color for Unity Catalog nullable column nodes", + "defaults": { + "dark": "#8EAFC2", + "light": "#4A6B82", + "highContrast": "#8EAFC2", + "highContrastLight": "#4A6B82" + } + }, + { + "id": "databricks.unityCatalog.registeredModel", + "description": "Color for registered model nodes in the Unity Catalog view", + "defaults": { + "dark": "#C586C0", + "light": "#AF00DB", + "highContrast": "#C586C0" + } + }, + { + "id": "databricks.unityCatalog.modelVersion", + "description": "Color for model version nodes in the Unity Catalog view", + "defaults": { + "dark": "#B5CEA8", + "light": "#008000", + "highContrast": "#B5CEA8" + } + } + ], "viewsWelcome": [ { "view": "configurationView", @@ -531,6 +678,16 @@ "when": "view == workspaceFsView", "group": "navigation@1" }, + { + "command": "databricks.unityCatalog.filter", + "when": "view == unityCatalogView", + "group": "navigation@2" + }, + { + "command": "databricks.unityCatalog.refresh", + "when": "view == unityCatalogView", + "group": "navigation@1" + }, { "command": "databricks.bundle.refreshRemoteState", "when": "view == dabsResourceExplorerView && databricks.context.bundle.deploymentState == idle", @@ -587,6 +744,61 @@ "when": "viewItem =~ /^databricks.*\\.(has-url).*$/ && databricks.context.bundle.deploymentState == idle", "group": "navigation_2@0" }, + { + "command": "databricks.unityCatalog.openExternal", + "when": "view == unityCatalogView && viewItem =~ /\\.has-url/", + "group": "inline@1" + }, + { + "command": "databricks.unityCatalog.openExternal", + "when": "view == unityCatalogView && viewItem =~ /\\.has-url/", + "group": "navigation_2@0" + }, + { + "command": "databricks.unityCatalog.copyName", + "when": "view == unityCatalogView && viewItem =~ /unityCatalog/", + "group": "navigation_2@0" + }, + { + "command": "databricks.unityCatalog.copyStorageLocation", + "when": "view == unityCatalogView && viewItem =~ /\\.has-storage/", + "group": "navigation_2@1" + }, + { + "command": "databricks.unityCatalog.copyViewSql", + "when": "view == unityCatalogView && viewItem =~ /\\.is-view/", + "group": "navigation_2@2" + }, + { + "command": "databricks.unityCatalog.refreshNode", + "when": "view == unityCatalogView && viewItem =~ /^unityCatalog\\.(?!column)/", + "group": "navigation_2@3" + }, + { + "command": "databricks.unityCatalog.refreshNode", + "when": "view == unityCatalogView && viewItem =~ /^unityCatalog\\.(?!column)/", + "group": "inline@2" + }, + { + "command": "databricks.unityCatalog.pinSchema", + "when": "view == unityCatalogView && viewItem =~ /^unityCatalog\\.schema/ && !(viewItem =~ /\\.is-pinned/)", + "group": "inline@3" + }, + { + "command": "databricks.unityCatalog.unpinSchema", + "when": "view == unityCatalogView && viewItem =~ /^unityCatalog\\.schema.*\\.is-pinned/", + "group": "inline@3" + }, + { + "command": "databricks.unityCatalog.pinSchema", + "when": "view == unityCatalogView && viewItem =~ /^unityCatalog\\.schema/ && !(viewItem =~ /\\.is-pinned/)", + "group": "navigation_2@4" + }, + { + "command": "databricks.unityCatalog.unpinSchema", + "when": "view == unityCatalogView && viewItem =~ /^unityCatalog\\.schema.*\\.is-pinned/", + "group": "navigation_2@4" + }, { "command": "databricks.utils.goToDefinition", "when": "viewItem =~ /^databricks.*\\.(has-source-location).*$/", diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 4713b02eb..c75d6725a 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -1,6 +1,7 @@ import { commands, debug, + env, ExtensionContext, extensions, window, @@ -26,6 +27,7 @@ import { FileUtils, PackageJsonUtils, TerraformUtils, + UrlUtils, UtilsCommands, } from "./utils"; import {ConfigureAutocomplete} from "./language/ConfigureAutocomplete"; @@ -74,6 +76,10 @@ import {SyncCommands} from "./sync/SyncCommands"; import {CodeSynchronizer} from "./sync"; import {BundlePipelinesManager} from "./bundle/BundlePipelinesManager"; import {DocsViewTreeDataProvider} from "./ui/docs-view/DocsViewTreeDataProvider"; +import { + UnityCatalogTreeDataProvider, + UnityCatalogTreeNode, +} from "./ui/unity-catalog/UnityCatalogTreeDataProvider"; // eslint-disable-next-line @typescript-eslint/no-var-requires const packageJson = require("../package.json"); @@ -370,6 +376,100 @@ export async function activate( ) ); + const unityCatalogTreeDataProvider = new UnityCatalogTreeDataProvider( + connectionManager, + stateStorage + ); + const unityCatalogTreeView = window.createTreeView("unityCatalogView", { + treeDataProvider: unityCatalogTreeDataProvider, + // eslint-disable-next-line @typescript-eslint/no-explicit-any + filterOnType: true, + } as any); + context.subscriptions.push( + unityCatalogTreeDataProvider, + unityCatalogTreeView, + telemetry.registerCommand( + "databricks.unityCatalog.refresh", + unityCatalogTreeDataProvider.refresh, + unityCatalogTreeDataProvider + ), + telemetry.registerCommand( + "databricks.unityCatalog.refreshNode", + (node: UnityCatalogTreeNode) => + unityCatalogTreeDataProvider.refreshNode(node) + ), + telemetry.registerCommand( + "databricks.unityCatalog.copyStorageLocation", + async (node: UnityCatalogTreeNode) => { + if ( + (node.kind === "table" || node.kind === "volume") && + node.storageLocation + ) { + await env.clipboard.writeText(node.storageLocation); + } + } + ), + telemetry.registerCommand( + "databricks.unityCatalog.copyViewSql", + async (node: UnityCatalogTreeNode) => { + if (node.kind === "table" && node.viewDefinition) { + await env.clipboard.writeText(node.viewDefinition); + } + } + ), + telemetry.registerCommand( + "databricks.unityCatalog.copyName", + async (node: UnityCatalogTreeNode) => { + if (node.kind === "error" || node.kind === "empty") { + return; + } + const text = node.kind === "column" ? node.name : node.fullName; + await env.clipboard.writeText(text); + window.showInformationMessage("Copied to clipboard"); + } + ), + telemetry.registerCommand( + "databricks.unityCatalog.openExternal", + async (node: UnityCatalogTreeNode) => { + if (node.kind === "error" || node.kind === "column") { + return; + } + const url = + unityCatalogTreeDataProvider.getNodeExploreUrl(node); + if (!url) { + window.showErrorMessage( + "Databricks: Can't open external link. No URL found." + ); + return; + } + await UrlUtils.openExternal(url); + } + ), + commands.registerCommand( + "databricks.unityCatalog.filter", + async () => { + await commands.executeCommand("unityCatalogView.focus"); + await commands.executeCommand("list.find"); + } + ), + telemetry.registerCommand( + "databricks.unityCatalog.pinSchema", + (node: UnityCatalogTreeNode) => { + if (node.kind === "schema") { + return unityCatalogTreeDataProvider.pinSchema(node); + } + } + ), + telemetry.registerCommand( + "databricks.unityCatalog.unpinSchema", + (node: UnityCatalogTreeNode) => { + if (node.kind === "schema") { + return unityCatalogTreeDataProvider.unpinSchema(node); + } + } + ) + ); + const configureAutocomplete = new ConfigureAutocomplete( context, stateStorage, diff --git a/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.test.ts b/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.test.ts new file mode 100644 index 000000000..3ffc16a18 --- /dev/null +++ b/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.test.ts @@ -0,0 +1,750 @@ +/* eslint-disable @typescript-eslint/naming-convention */ + +import assert from "assert"; +import {anything, instance, mock, when} from "ts-mockito"; +import {Disposable} from "vscode"; +import {WorkspaceClient} from "@databricks/sdk-experimental"; +import { + CatalogsService, + FunctionsService, + ModelVersionsService, + RegisteredModelsService, + SchemasService, + TablesService, + VolumesService, +} from "@databricks/sdk-experimental/dist/apis/catalog/api"; +import { + ConnectionManager, + ConnectionState, +} from "../../configuration/ConnectionManager"; +import {resolveProviderResult} from "../../test/utils"; +import { + UnityCatalogTreeDataProvider, + UnityCatalogTreeItem, + UnityCatalogTreeNode, +} from "./UnityCatalogTreeDataProvider"; +import {StateStorage} from "../../vscode-objs/StateStorage"; + +describe(__filename, () => { + let disposables: Disposable[] = []; + let mockConnectionManager: ConnectionManager; + let stubStateStorage: StateStorage; + let mockWorkspaceClient: WorkspaceClient; + let mockCatalogs: CatalogsService; + let mockSchemas: SchemasService; + let mockTables: TablesService; + let mockVolumes: VolumesService; + let mockFunctions: FunctionsService; + let mockRegisteredModels: RegisteredModelsService; + let mockModelVersions: ModelVersionsService; + let onDidChangeStateHandler: (s: ConnectionState) => void; + + beforeEach(() => { + disposables = []; + onDidChangeStateHandler = () => {}; + stubStateStorage = { + get: () => [] as string[], + set: async () => {}, + onDidChange: () => ({dispose() {}}), + } as unknown as StateStorage; + + mockCatalogs = mock(CatalogsService); + when(mockCatalogs.list(anything())).thenCall(() => { + async function* impl() { + yield {name: "c_b", full_name: "c_b"}; + yield {name: "c_a", full_name: "c_a"}; + } + return impl(); + }); + + mockSchemas = mock(SchemasService); + when(mockSchemas.list(anything())).thenCall(() => { + async function* impl() { + yield {name: "s_b", full_name: "cat.s_b"}; + yield {name: "s_a", full_name: "cat.s_a"}; + } + return impl(); + }); + + mockTables = mock(TablesService); + when(mockTables.list(anything())).thenCall(() => { + async function* impl() { + yield { + name: "t1", + full_name: "cat.sch.t1", + table_type: "MANAGED", + data_source_format: "DELTA", + comment: "a test table", + owner: "alice", + columns: [ + { + name: "id", + type_text: "bigint", + nullable: false, + position: 0, + }, + { + name: "name", + type_text: "string", + nullable: true, + position: 1, + }, + ], + }; + } + return impl(); + }); + + mockVolumes = mock(VolumesService); + when(mockVolumes.list(anything())).thenCall(() => { + async function* impl() { + yield { + name: "v1", + full_name: "cat.sch.v1", + volume_type: "MANAGED", + }; + } + return impl(); + }); + + mockFunctions = mock(FunctionsService); + when(mockFunctions.list(anything())).thenCall(() => { + async function* impl() { + yield { + name: "f1", + catalog_name: "cat", + schema_name: "sch", + }; + } + return impl(); + }); + + mockRegisteredModels = mock(RegisteredModelsService); + when(mockRegisteredModels.list(anything())).thenCall(() => { + async function* impl() { + /* empty */ + } + return impl(); + }); + + mockModelVersions = mock(ModelVersionsService); + when(mockModelVersions.list(anything())).thenCall(() => { + async function* impl() { + /* empty */ + } + return impl(); + }); + + mockWorkspaceClient = mock(WorkspaceClient); + when(mockWorkspaceClient.catalogs).thenReturn(instance(mockCatalogs)); + when(mockWorkspaceClient.schemas).thenReturn(instance(mockSchemas)); + when(mockWorkspaceClient.tables).thenReturn(instance(mockTables)); + when(mockWorkspaceClient.volumes).thenReturn(instance(mockVolumes)); + when(mockWorkspaceClient.functions).thenReturn(instance(mockFunctions)); + when(mockWorkspaceClient.registeredModels).thenReturn( + instance(mockRegisteredModels) + ); + when(mockWorkspaceClient.modelVersions).thenReturn( + instance(mockModelVersions) + ); + + mockConnectionManager = mock(ConnectionManager); + when(mockConnectionManager.workspaceClient).thenReturn( + instance(mockWorkspaceClient) + ); + when(mockConnectionManager.onDidChangeState).thenReturn( + (cb: (s: ConnectionState) => void) => { + onDidChangeStateHandler = cb; + return {dispose() {}}; + } + ); + }); + + afterEach(() => { + disposables.forEach((d) => d.dispose()); + }); + + it("returns undefined when not connected", async () => { + when(mockConnectionManager.workspaceClient).thenReturn(undefined); + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const children = await resolveProviderResult(provider.getChildren()); + assert.strictEqual(children, undefined); + }); + + it("lists catalogs sorted by name", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const children = (await resolveProviderResult( + provider.getChildren() + )) as UnityCatalogTreeNode[]; + assert(children); + assert.strictEqual(children.length, 2); + const first = children[0]; + const second = children[1]; + assert.strictEqual(first.kind, "catalog"); + assert.strictEqual(second.kind, "catalog"); + if (first.kind !== "catalog" || second.kind !== "catalog") { + assert.fail("expected catalogs"); + } + assert.strictEqual(first.name, "c_a"); + assert.strictEqual(second.name, "c_b"); + }); + + it("lists schemas under a catalog", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const catalog: UnityCatalogTreeNode = { + kind: "catalog", + name: "cat", + fullName: "cat", + }; + const children = (await resolveProviderResult( + provider.getChildren(catalog) + )) as UnityCatalogTreeNode[]; + + assert(children); + assert.strictEqual(children.length, 2); + assert.strictEqual(children[0].kind, "schema"); + assert.strictEqual(children[0].name, "s_a"); + assert.strictEqual( + (children[0] as {catalogName: string}).catalogName, + "cat" + ); + }); + + it("lists tables, volumes, and functions under a schema", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const schema: UnityCatalogTreeNode = { + kind: "schema", + catalogName: "cat", + name: "sch", + fullName: "cat.sch", + }; + const children = (await resolveProviderResult( + provider.getChildren(schema) + )) as UnityCatalogTreeNode[]; + + assert(children); + assert.strictEqual(children.length, 3); + const kinds = children.map((c) => c.kind).sort(); + assert.deepStrictEqual(kinds, ["function", "table", "volume"]); + + const table = children.find((c) => c.kind === "table"); + assert(table && table.kind === "table"); + assert.strictEqual(table.name, "t1"); + + const volume = children.find((c) => c.kind === "volume"); + assert(volume && volume.kind === "volume"); + assert.strictEqual(volume.name, "v1"); + + const fn = children.find((c) => c.kind === "function"); + assert(fn && fn.kind === "function"); + assert.strictEqual(fn.name, "f1"); + assert.strictEqual(fn.fullName, "cat.sch.f1"); + }); + + it("fires onDidChangeTreeData when connection state changes", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + let count = 0; + disposables.push( + provider.onDidChangeTreeData(() => { + count += 1; + }) + ); + + assert.strictEqual(count, 0); + onDidChangeStateHandler("CONNECTED"); + assert.strictEqual(count, 1); + }); + + it("getTreeItem sets url when host is available", async () => { + const stubManager = { + onDidChangeState: () => ({dispose() {}}), + databricksWorkspace: { + host: new URL("https://adb-123.azuredatabricks.net/"), + }, + } as unknown as ConnectionManager; + + const provider = new UnityCatalogTreeDataProvider(stubManager, stubStateStorage); + disposables.push(provider); + + const catalog: UnityCatalogTreeNode = { + kind: "catalog", + name: "cat", + fullName: "cat", + }; + const item = provider.getTreeItem(catalog) as UnityCatalogTreeItem; + + assert(item.url, "url should be set"); + assert( + item.url!.includes("explore/data/cat"), + `url should contain explore/data/cat, got: ${item.url}` + ); + assert( + item.contextValue?.endsWith(".has-url"), + `contextValue should end with .has-url, got: ${item.contextValue}` + ); + assert.strictEqual(item.copyText, "cat"); + }); + + it("getTreeItem omits url when no host", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const catalog: UnityCatalogTreeNode = { + kind: "catalog", + name: "cat", + fullName: "cat", + }; + const item = provider.getTreeItem(catalog) as UnityCatalogTreeItem; + + assert.strictEqual(item.url, undefined); + assert.strictEqual(item.contextValue, "unityCatalog.catalog"); + }); + + it("getTreeItem for function node", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const fn: UnityCatalogTreeNode = { + kind: "function", + catalogName: "cat", + schemaName: "sch", + name: "f1", + fullName: "cat.sch.f1", + }; + const item = provider.getTreeItem(fn) as UnityCatalogTreeItem; + + assert.strictEqual(item.label, "f1"); + assert.strictEqual(item.copyText, "cat.sch.f1"); + assert( + item.contextValue === "unityCatalog.function" || + item.contextValue === "unityCatalog.function.has-url" + ); + }); + + it("table node carries enriched fields", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const schema: UnityCatalogTreeNode = { + kind: "schema", + catalogName: "cat", + name: "sch", + fullName: "cat.sch", + }; + const children = (await resolveProviderResult( + provider.getChildren(schema) + )) as UnityCatalogTreeNode[]; + + const table = children.find((c) => c.kind === "table"); + assert(table && table.kind === "table"); + assert.strictEqual(table.dataSourceFormat, "DELTA"); + assert.strictEqual(table.comment, "a test table"); + assert.strictEqual(table.owner, "alice"); + assert(table.columns && table.columns.length === 2); + assert.strictEqual(table.columns[0].name, "id"); + assert.strictEqual(table.columns[0].typeText, "bigint"); + assert.strictEqual(table.columns[0].nullable, false); + }); + + it("getChildren for table with columns returns sorted column nodes", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const tableNode: UnityCatalogTreeNode = { + kind: "table", + catalogName: "cat", + schemaName: "sch", + name: "t1", + fullName: "cat.sch.t1", + columns: [ + {name: "b_col", typeText: "string", position: 1}, + {name: "a_col", typeText: "bigint", position: 0}, + ], + }; + const children = (await resolveProviderResult( + provider.getChildren(tableNode) + )) as UnityCatalogTreeNode[]; + + assert(children); + assert.strictEqual(children.length, 2); + assert.strictEqual(children[0].kind, "column"); + if (children[0].kind === "column") { + assert.strictEqual(children[0].name, "a_col"); + } + assert.strictEqual(children[1].kind, "column"); + if (children[1].kind === "column") { + assert.strictEqual(children[1].name, "b_col"); + } + }); + + it("getChildren for table without columns returns undefined", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const tableNode: UnityCatalogTreeNode = { + kind: "table", + catalogName: "cat", + schemaName: "sch", + name: "t1", + fullName: "cat.sch.t1", + columns: [], + }; + const children = await resolveProviderResult( + provider.getChildren(tableNode) + ); + assert.strictEqual(children, undefined); + }); + + it("getTreeItem for non-nullable column uses symbol-key icon", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const col: UnityCatalogTreeNode = { + kind: "column", + tableFullName: "cat.sch.t1", + name: "id", + typeText: "bigint", + nullable: false, + }; + const item = provider.getTreeItem(col) as UnityCatalogTreeItem; + assert.strictEqual(item.label, "id"); + assert.strictEqual(item.description, "bigint"); + const icon = item.iconPath as {id: string}; + assert.strictEqual(icon.id, "symbol-key"); + }); + + it("getTreeItem for nullable column uses symbol-field icon", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const col: UnityCatalogTreeNode = { + kind: "column", + tableFullName: "cat.sch.t1", + name: "val", + typeText: "string", + nullable: true, + }; + const item = provider.getTreeItem(col) as UnityCatalogTreeItem; + const icon = item.iconPath as {id: string}; + assert.strictEqual(icon.id, "symbol-field"); + }); + + it("getTreeItem for EXTERNAL table with storage has has-storage in contextValue", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const tableNode: UnityCatalogTreeNode = { + kind: "table", + catalogName: "cat", + schemaName: "sch", + name: "ext", + fullName: "cat.sch.ext", + tableType: "EXTERNAL", + storageLocation: "s3://bucket/path", + }; + const item = provider.getTreeItem(tableNode) as UnityCatalogTreeItem; + assert( + item.contextValue?.includes("has-storage"), + `expected has-storage in contextValue, got: ${item.contextValue}` + ); + assert.strictEqual(item.storageLocation, "s3://bucket/path"); + }); + + it("getTreeItem for VIEW table with view_definition has is-view in contextValue", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const tableNode: UnityCatalogTreeNode = { + kind: "table", + catalogName: "cat", + schemaName: "sch", + name: "vw", + fullName: "cat.sch.vw", + tableType: "VIEW", + viewDefinition: "SELECT 1", + }; + const item = provider.getTreeItem(tableNode) as UnityCatalogTreeItem; + assert( + item.contextValue?.includes("is-view"), + `expected is-view in contextValue, got: ${item.contextValue}` + ); + assert.strictEqual(item.viewDefinition, "SELECT 1"); + }); + + it("volume node carries volumeType and shows EXTERNAL label suffix", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const volNode: UnityCatalogTreeNode = { + kind: "volume", + catalogName: "cat", + schemaName: "sch", + name: "ev", + fullName: "cat.sch.ev", + volumeType: "EXTERNAL", + storageLocation: "s3://bucket/vol", + }; + const item = provider.getTreeItem(volNode) as UnityCatalogTreeItem; + assert.strictEqual(item.label, "ev (EXTERNAL)"); + assert( + item.contextValue?.includes("has-storage"), + `expected has-storage in contextValue, got: ${item.contextValue}` + ); + }); + + it("catalog node carries comment", async () => { + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const catNode: UnityCatalogTreeNode = { + kind: "catalog", + name: "cat", + fullName: "cat", + comment: "my catalog", + }; + const item = provider.getTreeItem(catNode) as UnityCatalogTreeItem; + assert.strictEqual(item.label, "cat"); + }); + + it("returns error when functions API throws", async () => { + when(mockFunctions.list(anything())).thenCall(() => { + async function* impl(): AsyncGenerator { + throw new Error("functions API unavailable"); + // eslint-disable-next-line no-unreachable + yield undefined as never; + } + return impl(); + }); + + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + + const schema: UnityCatalogTreeNode = { + kind: "schema", + catalogName: "cat", + name: "sch", + fullName: "cat.sch", + }; + const children = (await resolveProviderResult( + provider.getChildren(schema) + )) as UnityCatalogTreeNode[]; + + assert(children); + // allSettled: tables (t1) and volumes (v1) still succeed; only functions errors + assert.strictEqual(children.length, 3); + assert.notStrictEqual(children[0].kind, "error"); + assert.strictEqual(children[children.length - 1].kind, "error"); + }); + + it("lists registered models under a schema", async () => { + when(mockRegisteredModels.list(anything())).thenCall(() => { + async function* impl() { + yield {name: "m1", full_name: "cat.sch.m1"}; + } + return impl(); + }); + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + const schema: UnityCatalogTreeNode = { + kind: "schema", + catalogName: "cat", + name: "sch", + fullName: "cat.sch", + }; + const children = (await resolveProviderResult( + provider.getChildren(schema) + )) as UnityCatalogTreeNode[]; + const model = children.find((c) => c.kind === "registeredModel"); + assert(model && model.kind === "registeredModel"); + assert.strictEqual(model.name, "m1"); + assert.strictEqual(model.fullName, "cat.sch.m1"); + }); + + it("lists model versions for a registered model, sorted descending", async () => { + when(mockModelVersions.list(anything())).thenCall(() => { + async function* impl() { + yield {version: 1}; + yield {version: 3}; + yield {version: 2}; + } + return impl(); + }); + const provider = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + stubStateStorage + ); + disposables.push(provider); + const modelNode: UnityCatalogTreeNode = { + kind: "registeredModel", + catalogName: "cat", + schemaName: "sch", + name: "m1", + fullName: "cat.sch.m1", + }; + const children = (await resolveProviderResult( + provider.getChildren(modelNode) + )) as UnityCatalogTreeNode[]; + assert(children); + assert.strictEqual(children.length, 3); + assert.strictEqual(children[0].kind, "modelVersion"); + if (children[0].kind === "modelVersion") { + assert.strictEqual(children[0].version, 3); + assert.strictEqual((children[2] as any).version, 1); + } + }); + + it("pinSchema adds fullName to stateStorage and fires tree change", async () => { + const stored: string[] = []; + const spyStorage = { + get: () => stored, + set: async (_key: string, val: string[]) => { + stored.length = 0; + stored.push(...val); + }, + onDidChange: () => ({dispose() {}}), + } as unknown as StateStorage; + const p = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + spyStorage + ); + disposables.push(p); + let fired = 0; + disposables.push(p.onDidChangeTreeData(() => {fired++;})); + const schema = { + kind: "schema" as const, + catalogName: "cat", + name: "sch", + fullName: "cat.sch", + }; + await p.pinSchema(schema); + assert(stored.includes("cat.sch")); + assert.strictEqual(fired, 1); + }); + + it("unpinSchema removes fullName from stateStorage and fires tree change", async () => { + const stored: string[] = ["cat.sch", "cat.other"]; + const spyStorage = { + get: () => [...stored], + set: async (_key: string, val: string[]) => { + stored.length = 0; + stored.push(...val); + }, + onDidChange: () => ({dispose() {}}), + } as unknown as StateStorage; + const p = new UnityCatalogTreeDataProvider( + instance(mockConnectionManager), + spyStorage + ); + disposables.push(p); + let fired = 0; + disposables.push(p.onDidChangeTreeData(() => {fired++;})); + const schema = { + kind: "schema" as const, + catalogName: "cat", + name: "sch", + fullName: "cat.sch", + }; + await p.unpinSchema(schema); + assert(!stored.includes("cat.sch")); + assert(stored.includes("cat.other")); + assert.strictEqual(fired, 1); + }); + + it("pinned schema sorts before owned, owned before unowned", async () => { + const pinnedStorage = { + get: () => ["cat.s_b"], + set: async () => {}, + onDidChange: () => ({dispose() {}}), + } as unknown as StateStorage; + when(mockSchemas.list(anything())).thenCall(() => { + async function* impl() { + yield {name: "s_c", full_name: "cat.s_c", owner: "carol"}; + yield {name: "s_b", full_name: "cat.s_b", owner: "bob"}; // pinned + yield {name: "s_a", full_name: "cat.s_a", owner: "alice"}; // owned + } + return impl(); + }); + const stubManager = { + onDidChangeState: () => ({dispose() {}}), + workspaceClient: instance(mockWorkspaceClient), + databricksWorkspace: {user: {userName: "alice"}}, + } as unknown as ConnectionManager; + const p = new UnityCatalogTreeDataProvider(stubManager, pinnedStorage); + disposables.push(p); + const catalog: UnityCatalogTreeNode = { + kind: "catalog", + name: "cat", + fullName: "cat", + }; + const children = (await resolveProviderResult( + p.getChildren(catalog) + )) as UnityCatalogTreeNode[]; + assert.strictEqual((children[0] as any).name, "s_b"); // pinned first + assert.strictEqual((children[1] as any).name, "s_a"); // owned second + assert.strictEqual((children[2] as any).name, "s_c"); // unowned last + }); +}); diff --git a/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.ts b/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.ts new file mode 100644 index 000000000..7ae7ec6e7 --- /dev/null +++ b/packages/databricks-vscode/src/ui/unity-catalog/UnityCatalogTreeDataProvider.ts @@ -0,0 +1,176 @@ +/* eslint-disable @typescript-eslint/naming-convention */ +import {Disposable, EventEmitter, TreeDataProvider} from "vscode"; +import {ConnectionManager} from "../../configuration/ConnectionManager"; +import {buildTreeItem} from "./nodeRenderer"; +import {UnityCatalogTreeItem, UnityCatalogTreeNode} from "./types"; +import {StateStorage} from "../../vscode-objs/StateStorage"; +import { + loadCatalogs, + loadSchemas, + loadSchemaChildren, + loadModelVersions, +} from "./loaders"; + +export type { + ColumnData, + UnityCatalogTreeItem, + UnityCatalogTreeNode, +} from "./types"; + +export class UnityCatalogTreeDataProvider + implements TreeDataProvider, Disposable +{ + private readonly _onDidChangeTreeData = new EventEmitter< + UnityCatalogTreeNode | undefined | void + >(); + readonly onDidChangeTreeData = this._onDidChangeTreeData.event; + private readonly disposables: Disposable[] = []; + + constructor( + private readonly connectionManager: ConnectionManager, + private readonly stateStorage: StateStorage + ) { + this.disposables.push( + this.connectionManager.onDidChangeState(() => { + this._onDidChangeTreeData.fire(undefined); + }) + ); + } + + private getExploreUrl(path: string): string | undefined { + const host = this.connectionManager.databricksWorkspace?.host; + if (!host) { + return undefined; + } + return `${host.toString()}explore/data/${path}`; + } + + getNodeExploreUrl(node: UnityCatalogTreeNode): string | undefined { + if ( + node.kind === "error" || + node.kind === "column" || + node.kind === "empty" + ) { + return undefined; + } + const fullNamePath = node.fullName.replaceAll(".", "/"); + let path = fullNamePath; + switch (node.kind) { + case "registeredModel": + path = `models/${fullNamePath}`; + break; + case "modelVersion": + path = `models/${fullNamePath}/version/${node.version}`; + break; + case "function": + path = `functions/${fullNamePath}`; + break; + } + return this.getExploreUrl(path); + } + + getTreeItem(element: UnityCatalogTreeNode): UnityCatalogTreeItem { + return buildTreeItem(element, this.getNodeExploreUrl(element)); + } + + async getChildren( + element?: UnityCatalogTreeNode + ): Promise { + const client = this.connectionManager.workspaceClient; + if (!client) { + return undefined; + } + + const currentUser = + this.connectionManager.databricksWorkspace?.user; + + if (!element) { + return loadCatalogs(client, currentUser); + } + + if (element.kind === "error") { + return undefined; + } + + if (element.kind === "catalog") { + const pinned = new Set( + this.stateStorage.get( + "databricks.unityCatalog.pinnedSchemas" + ) ?? [] + ); + return loadSchemas(client, element.name, currentUser, pinned); + } + + if (element.kind === "schema") { + return loadSchemaChildren( + client, + element.catalogName, + element.name + ); + } + + if (element.kind === "registeredModel") { + return loadModelVersions(client, element); + } + + if (element.kind === "table") { + if (!element.columns?.length) { + return undefined; + } + return [...element.columns] + .sort((a, b) => (a.position ?? 0) - (b.position ?? 0)) + .map((col) => ({ + kind: "column" as const, + tableFullName: element.fullName, + name: col.name, + typeName: col.typeName, + typeText: col.typeText, + comment: col.comment, + nullable: col.nullable, + position: col.position, + })); + } + + return undefined; + } + + async pinSchema( + node: Extract + ): Promise { + const current = + this.stateStorage.get("databricks.unityCatalog.pinnedSchemas") ?? + []; + if (!current.includes(node.fullName)) { + await this.stateStorage.set( + "databricks.unityCatalog.pinnedSchemas", + [...current, node.fullName] + ); + } + this._onDidChangeTreeData.fire(undefined); + } + + async unpinSchema( + node: Extract + ): Promise { + const current = + this.stateStorage.get("databricks.unityCatalog.pinnedSchemas") ?? + []; + await this.stateStorage.set( + "databricks.unityCatalog.pinnedSchemas", + current.filter((n) => n !== node.fullName) + ); + this._onDidChangeTreeData.fire(undefined); + } + + refresh(): void { + this._onDidChangeTreeData.fire(undefined); + } + + refreshNode(element: UnityCatalogTreeNode): void { + this._onDidChangeTreeData.fire(element); + } + + dispose(): void { + this.disposables.forEach((d) => d.dispose()); + } +} diff --git a/packages/databricks-vscode/src/ui/unity-catalog/loaders.ts b/packages/databricks-vscode/src/ui/unity-catalog/loaders.ts new file mode 100644 index 000000000..b1bb217ef --- /dev/null +++ b/packages/databricks-vscode/src/ui/unity-catalog/loaders.ts @@ -0,0 +1,300 @@ +/* eslint-disable @typescript-eslint/naming-convention */ +import {ApiError, logging, type iam} from "@databricks/sdk-experimental"; +import {ConnectionManager} from "../../configuration/ConnectionManager"; +import {Loggers} from "../../logger"; +import {UnityCatalogTreeNode} from "./types"; +import {drainAsyncIterable, isOwnedByUser} from "./utils"; + +const logger = logging.NamedLogger.getOrCreate(Loggers.Extension); + +type Client = NonNullable; + +function emptyNode(message: string): UnityCatalogTreeNode[] { + return [{kind: "empty", message}]; +} + +function errorNode(e: unknown, resource: string): UnityCatalogTreeNode[] { + const message = + e instanceof ApiError + ? `Failed to load ${resource}: ${e.message}` + : `Failed to load ${resource}`; + logger.error(`Unity Catalog: ${message}`, e); + return [{kind: "error", message}]; +} + +export async function loadCatalogs( + client: Client, + currentUser: iam.User | undefined +): Promise { + try { + const rows = await drainAsyncIterable(client.catalogs.list({})); + const result = rows + .filter((c) => c.name) + .map((c) => ({ + kind: "catalog" as const, + name: c.name!, + fullName: c.full_name ?? c.name!, + comment: c.comment, + owner: c.owner, + owned: isOwnedByUser(c.owner, currentUser), + })) + .sort((a, b) => { + if (a.owned && !b.owned) { + return -1; + } + if (!a.owned && b.owned) { + return 1; + } + return a.name.localeCompare(b.name); + }); + return result.length > 0 ? result : emptyNode("No catalogs found"); + } catch (e) { + return errorNode(e, "catalogs"); + } +} + +export async function loadSchemas( + client: Client, + catalogName: string, + currentUser: iam.User | undefined, + pinnedSchemas: Set +): Promise { + try { + const rows = await drainAsyncIterable( + client.schemas.list({catalog_name: catalogName}) + ); + const result = rows + .filter((s) => s.name) + .map((s) => { + const fullName = s.full_name ?? `${catalogName}.${s.name}`; + return { + kind: "schema" as const, + catalogName, + name: s.name!, + fullName, + comment: s.comment, + owner: s.owner, + pinned: pinnedSchemas.has(fullName), + owned: isOwnedByUser(s.owner, currentUser), + }; + }) + .sort((a, b) => { + const rank = (n: typeof a) => (n.pinned ? 0 : n.owned ? 1 : 2); + const r = rank(a) - rank(b); + if (r !== 0) { + return r; + } + return a.name.localeCompare(b.name); + }); + return result.length > 0 ? result : emptyNode("No schemas"); + } catch (e) { + return errorNode(e, "schemas"); + } +} + +export async function loadSchemaChildren( + client: Client, + catalogName: string, + schemaName: string +): Promise { + const [tablesResult, volumesResult, functionsResult, modelsResult] = + await Promise.allSettled([ + drainAsyncIterable( + client.tables.list({ + catalog_name: catalogName, + schema_name: schemaName, + }) + ), + drainAsyncIterable( + client.volumes.list({ + catalog_name: catalogName, + schema_name: schemaName, + }) + ), + drainAsyncIterable( + client.functions.list({ + catalog_name: catalogName, + schema_name: schemaName, + }) + ), + drainAsyncIterable( + client.registeredModels.list({ + catalog_name: catalogName, + schema_name: schemaName, + }) + ), + ]); + + const tableNodes: UnityCatalogTreeNode[] = + tablesResult.status === "fulfilled" + ? tablesResult.value + .filter((t) => t.name) + .map((t) => ({ + kind: "table" as const, + catalogName, + schemaName, + name: t.name!, + fullName: + t.full_name ?? + `${catalogName}.${schemaName}.${t.name}`, + tableType: t.table_type, + comment: t.comment, + dataSourceFormat: t.data_source_format, + storageLocation: t.storage_location, + viewDefinition: t.view_definition, + owner: t.owner, + createdBy: t.created_by, + createdAt: t.created_at, + updatedAt: t.updated_at, + columns: (t.columns ?? []).map((col) => ({ + name: col.name!, + typeName: col.type_name, + typeText: col.type_text, + comment: col.comment, + nullable: col.nullable, + position: col.position, + })), + })) + : []; + + const volumeNodes: UnityCatalogTreeNode[] = + volumesResult.status === "fulfilled" + ? volumesResult.value + .filter((v) => v.name) + .map((v) => ({ + kind: "volume" as const, + catalogName, + schemaName, + name: v.name!, + fullName: + v.full_name ?? + `${catalogName}.${schemaName}.${v.name}`, + volumeType: v.volume_type, + storageLocation: v.storage_location, + comment: v.comment, + owner: v.owner, + })) + : []; + + const functionNodes: UnityCatalogTreeNode[] = + functionsResult.status === "fulfilled" + ? functionsResult.value + .filter((f) => f.name) + .map((f) => ({ + kind: "function" as const, + catalogName, + schemaName, + name: f.name!, + fullName: `${catalogName}.${schemaName}.${f.name}`, + })) + : []; + + const modelNodes: UnityCatalogTreeNode[] = + modelsResult.status === "fulfilled" + ? modelsResult.value + .filter((m) => m.name) + .map((m) => ({ + kind: "registeredModel" as const, + catalogName, + schemaName, + name: m.name!, + fullName: + m.full_name ?? + `${catalogName}.${schemaName}.${m.name}`, + comment: m.comment, + owner: m.owner, + storageLocation: m.storage_location, + aliases: m.aliases?.map((a) => ({ + alias_name: a.alias_name, + version_num: a.version_num, + })), + createdAt: m.created_at, + updatedAt: m.updated_at, + })) + : []; + + const errNodes: UnityCatalogTreeNode[] = [ + ...(tablesResult.status === "rejected" + ? errorNode(tablesResult.reason, "tables") + : []), + ...(volumesResult.status === "rejected" + ? errorNode(volumesResult.reason, "volumes") + : []), + ...(functionsResult.status === "rejected" + ? errorNode(functionsResult.reason, "functions") + : []), + ...(modelsResult.status === "rejected" + ? errorNode(modelsResult.reason, "registered models") + : []), + ]; + + const kindOrder = { + table: 0, + volume: 1, + function: 2, + registeredModel: 3, + } as Record; + const contentNodes = [ + ...tableNodes, + ...volumeNodes, + ...functionNodes, + ...modelNodes, + ]; + if (contentNodes.length === 0 && errNodes.length === 0) { + return emptyNode("No items"); + } + return [ + ...contentNodes.sort((a, b) => { + const an = + a.kind === "table" || + a.kind === "volume" || + a.kind === "function" || + a.kind === "registeredModel" + ? a.name + : ""; + const bn = + b.kind === "table" || + b.kind === "volume" || + b.kind === "function" || + b.kind === "registeredModel" + ? b.name + : ""; + const c = an.localeCompare(bn); + if (c !== 0) { + return c; + } + return (kindOrder[a.kind] ?? 0) - (kindOrder[b.kind] ?? 0); + }), + ...errNodes, + ]; +} + +export async function loadModelVersions( + client: Client, + model: Extract +): Promise { + try { + const rows = await drainAsyncIterable( + client.modelVersions.list({full_name: model.fullName}) + ); + const nodes = rows + .filter((v) => v.version !== undefined) + .map((v) => ({ + kind: "modelVersion" as const, + catalogName: model.catalogName, + schemaName: model.schemaName, + modelName: model.name, + fullName: model.fullName, + version: v.version!, + comment: v.comment, + status: v.status, + storageLocation: v.storage_location, + createdAt: v.created_at, + createdBy: v.created_by, + })) + .sort((a, b) => b.version - a.version); + return nodes.length > 0 ? nodes : emptyNode("No versions"); + } catch (e) { + return errorNode(e, "model versions"); + } +} diff --git a/packages/databricks-vscode/src/ui/unity-catalog/nodeRenderer.ts b/packages/databricks-vscode/src/ui/unity-catalog/nodeRenderer.ts new file mode 100644 index 000000000..6d86f558a --- /dev/null +++ b/packages/databricks-vscode/src/ui/unity-catalog/nodeRenderer.ts @@ -0,0 +1,364 @@ +import { + MarkdownString, + ThemeColor, + ThemeIcon, + TreeItemCollapsibleState, +} from "vscode"; +import {UnityCatalogTreeNode, UnityCatalogTreeItem} from "./types"; +import {formatTs} from "./utils"; + +export function buildTreeItem( + node: UnityCatalogTreeNode, + exploreUrl: string | undefined +): UnityCatalogTreeItem { + switch (node.kind) { + case "error": + return renderError(node); + case "empty": + return renderEmpty(node); + case "catalog": + return renderCatalog(node, exploreUrl); + case "schema": + return renderSchema(node, exploreUrl); + case "table": + return renderTable(node, exploreUrl); + case "volume": + return renderVolume(node, exploreUrl); + case "function": + return renderFunction(node, exploreUrl); + case "registeredModel": + return renderRegisteredModel(node, exploreUrl); + case "modelVersion": + return renderModelVersion(node, exploreUrl); + case "column": + return renderColumn(node); + } +} + +function renderError( + node: Extract +): UnityCatalogTreeItem { + return { + label: node.message, + iconPath: new ThemeIcon( + "error", + new ThemeColor("notificationsErrorIcon.foreground") + ), + collapsibleState: TreeItemCollapsibleState.None, + }; +} + +function renderEmpty( + node: Extract +): UnityCatalogTreeItem { + return { + label: node.message, + iconPath: new ThemeIcon( + "info", + new ThemeColor("descriptionForeground") + ), + collapsibleState: TreeItemCollapsibleState.None, + }; +} + +function renderCatalog( + node: Extract, + exploreUrl: string | undefined +): UnityCatalogTreeItem { + const tt = new MarkdownString(`**${node.fullName}**`); + if (node.comment) { + tt.appendMarkdown(`\n\n${node.comment}`); + } + return { + label: node.name, + description: node.owned ? "yours" : undefined, + tooltip: tt, + iconPath: new ThemeIcon( + "library", + new ThemeColor("databricks.unityCatalog.catalog") + ), + contextValue: exploreUrl + ? "unityCatalog.catalog.has-url" + : "unityCatalog.catalog", + collapsibleState: TreeItemCollapsibleState.Collapsed, + url: exploreUrl, + copyText: node.fullName, + }; +} + +function renderSchema( + node: Extract, + exploreUrl: string | undefined +): UnityCatalogTreeItem { + const tt = new MarkdownString(`**${node.fullName}**`); + if (node.comment) { + tt.appendMarkdown(`\n\n${node.comment}`); + } + const baseContextValue = exploreUrl + ? "unityCatalog.schema.has-url" + : "unityCatalog.schema"; + let description: string | undefined; + if (node.pinned && node.owned) { + description = "★ · yours"; + } else if (node.pinned) { + description = "★"; + } else if (node.owned) { + description = "yours"; + } + return { + label: node.name, + description, + tooltip: tt, + iconPath: new ThemeIcon( + "folder-library", + new ThemeColor("databricks.unityCatalog.schema") + ), + contextValue: node.pinned + ? baseContextValue + ".is-pinned" + : baseContextValue, + collapsibleState: TreeItemCollapsibleState.Collapsed, + url: exploreUrl, + copyText: node.fullName, + }; +} + +function renderTable( + node: Extract, + exploreUrl: string | undefined +): UnityCatalogTreeItem { + const typeSuffix = + node.tableType && node.tableType !== "MANAGED" + ? ` (${node.tableType})` + : ""; + const flags = ["unityCatalog.table"]; + if (exploreUrl) { + flags.push("has-url"); + } + if (node.storageLocation) { + flags.push("has-storage"); + } + const isView = + node.tableType === "VIEW" || node.tableType === "MATERIALIZED_VIEW"; + if (isView && node.viewDefinition) { + flags.push("is-view"); + } + + const tt = new MarkdownString(`**${node.fullName}**`); + if (node.tableType) { + tt.appendMarkdown(`\n\n*Type:* ${node.tableType}`); + } + if (node.dataSourceFormat) { + tt.appendMarkdown(` · *Format:* ${node.dataSourceFormat}`); + } + if (node.owner) { + tt.appendMarkdown(`\n\n*Owner:* ${node.owner}`); + } + if (node.createdBy) { + tt.appendMarkdown(` · *Created by:* ${node.createdBy}`); + } + const cAt = formatTs(node.createdAt); + const uAt = formatTs(node.updatedAt); + if (cAt) { + tt.appendMarkdown(`\n\n*Created:* ${cAt}`); + } + if (uAt) { + tt.appendMarkdown(` *Updated:* ${uAt}`); + } + if (node.comment) { + tt.appendMarkdown(`\n\n${node.comment}`); + } + + const hasColumns = (node.columns?.length ?? 0) > 0; + return { + label: `${node.name}${typeSuffix}`, + description: node.dataSourceFormat, + tooltip: tt, + iconPath: new ThemeIcon( + "table", + new ThemeColor("databricks.unityCatalog.table") + ), + contextValue: flags.join("."), + collapsibleState: hasColumns + ? TreeItemCollapsibleState.Collapsed + : TreeItemCollapsibleState.None, + url: exploreUrl, + copyText: node.fullName, + storageLocation: node.storageLocation, + viewDefinition: node.viewDefinition, + }; +} + +function renderVolume( + node: Extract, + exploreUrl: string | undefined +): UnityCatalogTreeItem { + const isExternal = + node.volumeType !== undefined && node.volumeType !== "MANAGED"; + const label = isExternal ? `${node.name} (${node.volumeType})` : node.name; + const flags = ["unityCatalog.volume"]; + if (exploreUrl) { + flags.push("has-url"); + } + if (node.storageLocation) { + flags.push("has-storage"); + } + const tt = new MarkdownString(`**${node.fullName}**`); + if (node.volumeType) { + tt.appendMarkdown(`\n\n*Type:* ${node.volumeType}`); + } + if (node.owner) { + tt.appendMarkdown(`\n\n*Owner:* ${node.owner}`); + } + if (node.comment) { + tt.appendMarkdown(`\n\n${node.comment}`); + } + return { + label, + tooltip: tt, + iconPath: new ThemeIcon( + "package", + new ThemeColor("databricks.unityCatalog.volume") + ), + contextValue: flags.join("."), + collapsibleState: TreeItemCollapsibleState.None, + url: exploreUrl, + copyText: node.fullName, + storageLocation: node.storageLocation, + }; +} + +function renderFunction( + node: Extract, + exploreUrl: string | undefined +): UnityCatalogTreeItem { + return { + label: node.name, + tooltip: node.fullName, + iconPath: new ThemeIcon( + "symbol-function", + new ThemeColor("databricks.unityCatalog.function") + ), + contextValue: exploreUrl + ? "unityCatalog.function.has-url" + : "unityCatalog.function", + collapsibleState: TreeItemCollapsibleState.None, + url: exploreUrl, + copyText: node.fullName, + }; +} + +function renderRegisteredModel( + node: Extract, + exploreUrl: string | undefined +): UnityCatalogTreeItem { + const tt = new MarkdownString(`**${node.fullName}**`); + if (node.owner) { + tt.appendMarkdown(`\n\n*Owner:* ${node.owner}`); + } + if (node.comment) { + tt.appendMarkdown(`\n\n${node.comment}`); + } + if (node.aliases && node.aliases.length > 0) { + const aliasList = node.aliases + .filter((a) => a.alias_name) + .map((a) => + a.version_num !== undefined + ? `${a.alias_name} → v${a.version_num}` + : a.alias_name! + ) + .join(", "); + if (aliasList) { + tt.appendMarkdown(`\n\n*Aliases:* ${aliasList}`); + } + } + const cAt = formatTs(node.createdAt); + const uAt = formatTs(node.updatedAt); + if (cAt) { + tt.appendMarkdown(`\n\n*Created:* ${cAt}`); + } + if (uAt) { + tt.appendMarkdown(` *Updated:* ${uAt}`); + } + return { + label: node.name, + tooltip: tt, + iconPath: new ThemeIcon( + "beaker", + new ThemeColor("databricks.unityCatalog.registeredModel") + ), + contextValue: exploreUrl + ? "unityCatalog.registeredModel.has-url" + : "unityCatalog.registeredModel", + collapsibleState: TreeItemCollapsibleState.Collapsed, + url: exploreUrl, + copyText: node.fullName, + }; +} + +function renderModelVersion( + node: Extract, + exploreUrl: string | undefined +): UnityCatalogTreeItem { + const tt = new MarkdownString(`**v${node.version}**`); + if (node.status) { + tt.appendMarkdown(`\n\n*Status:* ${node.status}`); + } + if (node.comment) { + tt.appendMarkdown(`\n\n${node.comment}`); + } + if (node.createdBy) { + tt.appendMarkdown(`\n\n*Created by:* ${node.createdBy}`); + } + const cAt = formatTs(node.createdAt); + if (cAt) { + tt.appendMarkdown(`\n\n*Created:* ${cAt}`); + } + return { + label: `v${node.version}`, + description: + node.status && node.status !== "READY" ? node.status : undefined, + tooltip: tt, + iconPath: new ThemeIcon( + "tag", + new ThemeColor("databricks.unityCatalog.modelVersion") + ), + contextValue: exploreUrl + ? "unityCatalog.modelVersion.has-url" + : "unityCatalog.modelVersion", + collapsibleState: TreeItemCollapsibleState.None, + url: exploreUrl, + copyText: node.fullName, + }; +} + +function renderColumn( + node: Extract +): UnityCatalogTreeItem { + const icon = + node.nullable === false + ? new ThemeIcon( + "symbol-key", + new ThemeColor("databricks.unityCatalog.columnKey") + ) + : new ThemeIcon( + "symbol-field", + new ThemeColor("databricks.unityCatalog.column") + ); + const typeLabel = node.typeText ?? node.typeName ?? ""; + const tt = new MarkdownString(`**${node.name}** \`${typeLabel}\``); + if (node.nullable === false) { + tt.appendMarkdown(" *(not null)*"); + } + if (node.comment) { + tt.appendMarkdown(`\n\n${node.comment}`); + } + return { + label: node.name, + description: typeLabel, + tooltip: tt, + iconPath: icon, + contextValue: "unityCatalog.column", + collapsibleState: TreeItemCollapsibleState.None, + copyText: node.name, + }; +} diff --git a/packages/databricks-vscode/src/ui/unity-catalog/types.ts b/packages/databricks-vscode/src/ui/unity-catalog/types.ts new file mode 100644 index 000000000..0dd507a7b --- /dev/null +++ b/packages/databricks-vscode/src/ui/unity-catalog/types.ts @@ -0,0 +1,103 @@ +import {TreeItem} from "vscode"; + +export interface ColumnData { + name: string; + typeName?: string; + typeText?: string; + comment?: string; + nullable?: boolean; + position?: number; +} + +export type UnityCatalogTreeNode = + | {kind: "catalog"; name: string; fullName: string; comment?: string; owner?: string; owned?: boolean} + | { + kind: "schema"; + catalogName: string; + name: string; + fullName: string; + comment?: string; + pinned?: boolean; + owner?: string; + owned?: boolean; + } + | { + kind: "table"; + catalogName: string; + schemaName: string; + name: string; + fullName: string; + tableType?: string; + comment?: string; + dataSourceFormat?: string; + storageLocation?: string; + viewDefinition?: string; + owner?: string; + createdBy?: string; + createdAt?: number; + updatedAt?: number; + columns?: ColumnData[]; + } + | { + kind: "volume"; + catalogName: string; + schemaName: string; + name: string; + fullName: string; + volumeType?: string; + storageLocation?: string; + comment?: string; + owner?: string; + } + | { + kind: "function"; + catalogName: string; + schemaName: string; + name: string; + fullName: string; + } + | { + kind: "registeredModel"; + catalogName: string; + schemaName: string; + name: string; + fullName: string; + comment?: string; + owner?: string; + storageLocation?: string; + aliases?: Array<{alias_name?: string; version_num?: number}>; + createdAt?: number; + updatedAt?: number; + } + | { + kind: "modelVersion"; + catalogName: string; + schemaName: string; + modelName: string; + fullName: string; + version: number; + comment?: string; + status?: string; + storageLocation?: string; + createdAt?: number; + createdBy?: string; + } + | { + kind: "column"; + tableFullName: string; + name: string; + typeName?: string; + typeText?: string; + comment?: string; + nullable?: boolean; + position?: number; + } + | {kind: "error"; message: string} + | {kind: "empty"; message: string}; + +export interface UnityCatalogTreeItem extends TreeItem { + url?: string; + copyText?: string; + storageLocation?: string; + viewDefinition?: string; +} diff --git a/packages/databricks-vscode/src/ui/unity-catalog/utils.ts b/packages/databricks-vscode/src/ui/unity-catalog/utils.ts new file mode 100644 index 000000000..9cdc8160b --- /dev/null +++ b/packages/databricks-vscode/src/ui/unity-catalog/utils.ts @@ -0,0 +1,34 @@ +import {type iam} from "@databricks/sdk-experimental"; + +export async function drainAsyncIterable( + iter: AsyncIterable +): Promise { + const out: T[] = []; + for await (const item of iter) { + out.push(item); + } + return out; +} + +export function isOwnedByUser( + owner: string | undefined, + user: iam.User | undefined +): boolean { + if (!owner || !user) { + return false; + } + if (owner === user.userName) { + return true; + } + // TODO: Check if user is owner through group? like: return (user.groups ?? []).some((g) => g.display === owner); + return false; +} + +export function formatTs(ms: number | undefined): string | undefined { + if (ms === undefined) { + return undefined; + } + return ( + new Date(ms).toISOString().replace("T", " ").substring(0, 19) + " UTC" + ); +} diff --git a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts index 3c716599b..a32b72033 100644 --- a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts +++ b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts @@ -66,6 +66,11 @@ const StorageConfigurations = { location: "workspace", }), + "databricks.unityCatalog.pinnedSchemas": withType()({ + location: "workspace", + defaultValue: [], + }), + "databricks.lastInstalledExtensionVersion": withType()({ location: "global", defaultValue: "0.0.0",