diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 48129b382..e7540c520 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -41,7 +41,10 @@ "onCommand:databricks.connection.configureProject", "onCommand:databricks.connection.openDatabricksConfigFile", "onCommand:databricks.connection.attachCluster", + "onCommand:databricks.connection.attachClusterQuickPick", "onCommand:databricks.connection.detachCluster", + "onCommand:databricks.connection.attachSyncDestination", + "onCommand:databricks.connection.detachSyncDestination", "onCommand:databricks.cli.startSync", "onCommand:databricks.cli.startSyncFull", "onCommand:databricks.cli.stopSync", @@ -79,6 +82,11 @@ "title": "Attach cluster", "icon": "$(plug)" }, + { + "command": "databricks.connection.attachClusterQuickPick", + "title": "Attach cluster", + "icon": "$(plug)" + }, { "command": "databricks.connection.detachCluster", "title": "Detach cluster", @@ -113,6 +121,16 @@ "icon": "$(refresh)", "title": "Refresh" }, + { + "command": "databricks.connection.attachSyncDestination", + "title": "Attach workspace", + "icon": "$(plug)" + }, + { + "command": "databricks.connection.detachSyncDestination", + "title": "Detach workspace", + "icon": "$(debug-disconnect)" + }, { "command": "databricks.run.runEditorContentsAsWorkflow", "title": "Run File as Workflow on Databricks", @@ -184,13 +202,23 @@ "group": "inline@0" }, { - "command": "databricks.connection.detachCluster", - "when": "view == configurationView && viewItem == clusterRunning", + "command": "databricks.connection.attachClusterQuickPick", + "when": "view == configurationView && viewItem == clusterDetached", "group": "inline@0" }, { "command": "databricks.connection.detachCluster", - "when": "view == configurationView && viewItem == clusterStopped", + "when": "view == configurationView && viewItem == clusterAttached", + "group": "inline@0" + }, + { + "command": "databricks.connection.attachSyncDestination", + "when": "view == configurationView && viewItem == syncDestinationDetached", + "group": "inline@0" + }, + { + "command": "databricks.connection.detachSyncDestination", + "when": "view == configurationView && viewItem == syncDestinationAttached", "group": "inline@0" } ], diff --git a/packages/databricks-vscode/src/cli/CliCommands.ts b/packages/databricks-vscode/src/cli/CliCommands.ts index 1c6e65d02..a6512dd71 100644 --- a/packages/databricks-vscode/src/cli/CliCommands.ts +++ b/packages/databricks-vscode/src/cli/CliCommands.ts @@ -31,7 +31,7 @@ export class CliCommands { return async () => { const workspacePath = workspace.rootPath; const me = this.connection.me; - const pathMapper = this.connection.pathMapper; + const syncDestination = this.connection.syncDestination; const profile = this.connection.profile; if (!workspacePath) { @@ -47,7 +47,7 @@ export class CliCommands { ); return; } - if (!pathMapper) { + if (!syncDestination) { window.showErrorMessage( "Can't start sync: Databricks synchronization destination not configured!" ); @@ -61,7 +61,7 @@ export class CliCommands { const {command, args} = this.cli.getSyncCommand( profile, me, - pathMapper, + syncDestination, syncType ); diff --git a/packages/databricks-vscode/src/cli/CliWrapper.test.ts b/packages/databricks-vscode/src/cli/CliWrapper.test.ts index 9872e435b..8d804d42b 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.test.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.test.ts @@ -1,13 +1,13 @@ import * as assert from "assert"; import {Uri} from "vscode"; -import {PathMapper} from "../configuration/PathMapper"; +import {SyncDestination} from "../configuration/SyncDestination"; import {CliWrapper} from "./CliWrapper"; describe(__filename, () => { it("should create sync command", () => { const cli = new CliWrapper(); - const mapper = new PathMapper( + const mapper = new SyncDestination( Uri.file( "/Workspace/Repos/fabian.jakobs@databricks.com/notebook-best-practices" ), @@ -29,7 +29,7 @@ describe(__filename, () => { it("should create full sync command", () => { const cli = new CliWrapper(); - const mapper = new PathMapper( + const mapper = new SyncDestination( Uri.file( "/Workspace/Repos/fabian.jakobs@databricks.com/notebook-best-practices" ), diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts index aac5a0c11..8df7aed1f 100644 --- a/packages/databricks-vscode/src/cli/CliWrapper.ts +++ b/packages/databricks-vscode/src/cli/CliWrapper.ts @@ -1,5 +1,5 @@ import {spawn} from "child_process"; -import {PathMapper} from "../configuration/PathMapper"; +import {SyncDestination} from "../configuration/SyncDestination"; interface Command { command: string; @@ -21,7 +21,7 @@ export class CliWrapper { getSyncCommand( profile: string, me: string, - pathMapper: PathMapper, + syncDestination: SyncDestination, syncType: "full" | "incremental" ): Command { const command = "dbx"; @@ -33,7 +33,7 @@ export class CliWrapper { "--user", me, "--dest-repo", - pathMapper.remoteWorkspaceName, + syncDestination.name, ]; if (syncType === "full") { diff --git a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.test.ts b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.test.ts index 19d25b97a..d2918b449 100644 --- a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.test.ts +++ b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.test.ts @@ -7,29 +7,46 @@ import {ConfigurationDataProvider} from "./ConfigurationDataProvider"; import {ApiClient, Cluster} from "@databricks/databricks-sdk"; import {ConnectionManager} from "./ConnectionManager"; import {resolveProviderResult} from "../test/utils"; +import {SyncDestination} from "./SyncDestination"; describe(__filename, () => { let mockedConnectionManager: ConnectionManager; let disposables: Array; let onChangeClusterListener: (e: Cluster) => void; + let onChangeSyncDestinationListener: (e: SyncDestination) => void; beforeEach(() => { disposables = []; mockedConnectionManager = mock(ConnectionManager); onChangeClusterListener = () => {}; + onChangeSyncDestinationListener = () => {}; + + when(mockedConnectionManager.onChangeState).thenReturn((_handler) => { + return { + dispose() {}, + }; + }); when(mockedConnectionManager.onChangeCluster).thenReturn((_handler) => { onChangeClusterListener = _handler; return { dispose() {}, }; }); + when(mockedConnectionManager.onChangeSyncDestination).thenReturn( + (_handler) => { + onChangeSyncDestinationListener = _handler; + return { + dispose() {}, + }; + } + ); }); afterEach(() => { disposables.forEach((d) => d.dispose()); }); - it("should reload tree on model change", async () => { + it("should reload tree on cluster change", async () => { let connectionManager = instance(mockedConnectionManager); let provider = new ConfigurationDataProvider(connectionManager); disposables.push(provider); @@ -46,6 +63,23 @@ describe(__filename, () => { assert(called); }); + it("should reload tree on sync destination change", async () => { + let connectionManager = instance(mockedConnectionManager); + let provider = new ConfigurationDataProvider(connectionManager); + disposables.push(provider); + + let called = false; + disposables.push( + provider.onDidChangeTreeData(() => { + called = true; + }) + ); + + assert(!called); + onChangeSyncDestinationListener(instance(mock(SyncDestination))); + assert(called); + }); + it("should get empty roots", async () => { let connectionManager = instance(mockedConnectionManager); let provider = new ConfigurationDataProvider(connectionManager); @@ -76,24 +110,33 @@ describe(__filename, () => { let children = await resolveProviderResult(provider.getChildren()); assert.deepEqual(children, [ { - collapsibleState: 0, - id: "CONNECTION", - label: "Profile: null", + collapsibleState: 2, + iconPath: { + color: undefined, + id: "tools", + }, + id: "PROFILE", + label: "Profile", }, { collapsibleState: 2, - contextValue: "clusterStopped", + contextValue: "clusterAttached", iconPath: { color: undefined, - id: "debug-stop", + id: "server", }, id: "CLUSTER", - label: "Cluster: cluster-name-2", + label: "Cluster", }, { collapsibleState: 2, + contextValue: "syncDestinationDetached", + iconPath: { + color: undefined, + id: "repo", + }, id: "WORKSPACE", - label: "Workspace", + label: 'Workspace - "None attached"', }, ]); }); diff --git a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts index 315538745..12ea800fb 100644 --- a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts +++ b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts @@ -3,6 +3,7 @@ import { Event, EventEmitter, ProviderResult, + ThemeIcon, TreeDataProvider, TreeItem, TreeItemCollapsibleState, @@ -25,17 +26,21 @@ export class ConfigurationDataProvider constructor(private connectionManager: ConnectionManager) { this.disposables.push( - this.connectionManager.onChangeState((state) => { + this.connectionManager.onChangeState(() => { this._onDidChangeTreeData.fire(); }), - this.connectionManager.onChangeCluster((cluster) => { - // console.log("change cluster event", cluster); + this.connectionManager.onChangeCluster(() => { + this._onDidChangeTreeData.fire(); + }), + this.connectionManager.onChangeSyncDestination(() => { this._onDidChangeTreeData.fire(); }) ); } - dispose() {} + dispose() { + this.disposables.forEach((d) => d.dispose()); + } getTreeItem(element: TreeItem): TreeItem | Thenable { return element; @@ -48,49 +53,93 @@ export class ConfigurationDataProvider return []; } return (async () => { + let cluster = this.connectionManager.cluster; + let syncDestination = this.connectionManager.syncDestination; + if (!element) { - let cluster = this.connectionManager.cluster; let children: Array = []; children.push({ - label: `Profile: ${this.connectionManager.profile}`, - id: "CONNECTION", - collapsibleState: TreeItemCollapsibleState.None, + label: `Profile`, + iconPath: new ThemeIcon("tools"), + id: "PROFILE", + collapsibleState: TreeItemCollapsibleState.Expanded, }); if (cluster) { - let clusterItem = - ClusterListDataProvider.clusterNodeToTreeItem(cluster); - children.push({ - ...clusterItem, - label: `Cluster: ${cluster.name}`, + label: `Cluster`, + iconPath: new ThemeIcon("server"), id: "CLUSTER", collapsibleState: TreeItemCollapsibleState.Expanded, + contextValue: "clusterAttached", }); } else { children.push({ - label: `Cluster: "None attached"`, + label: `Cluster - "None attached"`, + iconPath: new ThemeIcon("server"), id: "CLUSTER", collapsibleState: TreeItemCollapsibleState.Expanded, + contextValue: "clusterDetached", }); } - children.push({ - label: "Workspace", - id: "WORKSPACE", - collapsibleState: TreeItemCollapsibleState.Expanded, - }); + if (syncDestination) { + children.push({ + label: `Workspace`, + iconPath: new ThemeIcon("repo"), + id: "WORKSPACE", + collapsibleState: TreeItemCollapsibleState.Expanded, + contextValue: "syncDestinationAttached", + }); + } else { + children.push({ + label: `Workspace - "None attached"`, + iconPath: new ThemeIcon("repo"), + id: "WORKSPACE", + collapsibleState: TreeItemCollapsibleState.Expanded, + contextValue: "syncDestinationDetached", + }); + } return children; } - if (element.id === "CLUSTER") { - let cluster = this.connectionManager.cluster; - if (cluster) { - return ClusterListDataProvider.clusterNodeToTreeItems( - cluster - ); - } + if (element.id === "PROFILE" && this.connectionManager.profile) { + return [ + { + label: "Name", + description: this.connectionManager.profile, + collapsibleState: TreeItemCollapsibleState.None, + }, + ]; + } + + if (element.id === "CLUSTER" && cluster) { + let clusterItem = + ClusterListDataProvider.clusterNodeToTreeItem(cluster); + + return [ + { + label: "Name", + description: cluster.name, + iconPath: clusterItem.iconPath, + collapsibleState: TreeItemCollapsibleState.None, + }, + ...ClusterListDataProvider.clusterNodeToTreeItems(cluster), + ]; + } + + if (element.id === "WORKSPACE" && syncDestination) { + return [ + { + label: `Name: ${syncDestination.name}`, + collapsibleState: TreeItemCollapsibleState.None, + }, + { + label: `Path: ${syncDestination.path.path}`, + collapsibleState: TreeItemCollapsibleState.None, + }, + ]; } })(); } diff --git a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts index 5e205956e..5f94aedad 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionCommands.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionCommands.ts @@ -1,8 +1,18 @@ -import {Cluster} from "@databricks/databricks-sdk"; +import {Cluster, Repos} from "@databricks/databricks-sdk"; import {homedir} from "node:os"; -import {Uri, window, workspace} from "vscode"; +import {QuickPickItem, ThemeIcon, Uri, window, workspace} from "vscode"; +import {ClusterListDataProvider} from "../cluster/ClusterListDataProvider"; import {ConnectionManager} from "./ConnectionManager"; +export interface WorkspaceItem extends QuickPickItem { + id: number; + path: string; +} + +export interface ClusterItem extends QuickPickItem { + cluster: Cluster; +} + export class ConnectionCommands { constructor(private connectionManager: ConnectionManager) {} /** @@ -34,10 +44,6 @@ export class ConnectionCommands { openDatabricksConfigFileCommand() { return async () => { - console.log( - "really opening file", - Uri.joinPath(Uri.file(homedir()), ".databrickscfg").path - ); const doc = await workspace.openTextDocument( Uri.joinPath(Uri.file(homedir()), ".databrickscfg") ); @@ -56,6 +62,68 @@ export class ConnectionCommands { }; } + attachClusterQuickPickCommand() { + return async () => { + const apiClient = this.connectionManager.apiClient; + const me = this.connectionManager.me; + if (!apiClient || !me) { + // TODO + return; + } + + const quickPick = window.createQuickPick(); + + quickPick.busy = true; + quickPick.show(); + + let clusters = await Cluster.list(apiClient); + + function formatSize(sizeInMB: number): string { + if (sizeInMB > 1024) { + return Math.round(sizeInMB / 1024).toString() + " GB"; + } else { + return `${sizeInMB} MB`; + } + } + + function formatDetails(cluster: Cluster) { + let details = []; + if (cluster.memoryMb) { + details.push(formatSize(cluster.memoryMb)); + } + + if (cluster.cores) { + details.push(`${cluster.cores} Cores`); + } + + details.push(cluster.sparkVersion); + details.push(cluster.creator); + + return details.join(" | "); + } + + quickPick.items = clusters.map((c) => { + let treeItem = ClusterListDataProvider.clusterNodeToTreeItem(c); + return { + label: `$(${ + (treeItem.iconPath as ThemeIcon).id + }) ${c.name!} (${c.id})`, + detail: formatDetails(c), + cluster: c, + }; + }); + quickPick.busy = false; + + quickPick.onDidAccept(async () => { + const cluster = quickPick.selectedItems[0].cluster; + await this.connectionManager.attachCluster(cluster); + quickPick.dispose(); + }); + + quickPick.onDidHide(() => quickPick.dispose()); + }; + } + /** * Set cluster to undefined and remove cluster ID from settings file */ @@ -65,18 +133,55 @@ export class ConnectionCommands { }; } - /** - * Attach to a workspace from settings. If attach fails or no workspace is configured - * then show dialog to select (or create) one. Selected workspaces is saved in settings. - */ - attachWorkspaceCommand() { - return () => {}; + attachSyncDestinationCommand() { + return async () => { + const apiClient = this.connectionManager.apiClient; + const me = this.connectionManager.me; + if (!apiClient || !me) { + // TODO + return; + } + + const reposApi = new Repos(apiClient); + const quickPick = window.createQuickPick(); + + quickPick.busy = true; + quickPick.canSelectMany = false; + quickPick.show(); + + let repos = ( + await reposApi.getRepos({ + // eslint-disable-next-line @typescript-eslint/naming-convention + path_prefix: `/Repos/${me}`, + }) + ).repos; + + quickPick.items = repos.map((r) => ({ + label: r.path.split("/").pop() || "", + detail: r.path, + path: r.path, + id: r.id, + })); + quickPick.busy = false; + + quickPick.onDidAccept(async () => { + const repoPath = quickPick.selectedItems[0].path; + await this.connectionManager.attachSyncDestination( + Uri.file(repoPath) + ); + quickPick.dispose(); + }); + + quickPick.onDidHide(() => quickPick.dispose()); + }; } /** * Set workspace to undefined and remove workspace path from settings file. */ detachWorkspaceCommand() { - return () => {}; + return () => { + this.connectionManager.detachSyncDestination(); + }; } } diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts index 27b6c5e56..238c68931 100644 --- a/packages/databricks-vscode/src/configuration/ConnectionManager.ts +++ b/packages/databricks-vscode/src/configuration/ConnectionManager.ts @@ -13,7 +13,7 @@ import { workspace as vscodeWorkspace, } from "vscode"; import {CliWrapper} from "../cli/CliWrapper"; -import {PathMapper} from "./PathMapper"; +import {SyncDestination} from "./SyncDestination"; import {ProjectConfigFile} from "./ProjectConfigFile"; import {selectProfile} from "./selectProfileWizard"; @@ -31,7 +31,7 @@ export class ConnectionManager { private _state: ConnectionState = "DISCONNECTED"; private _cluster?: Cluster; private _apiClient?: ApiClient; - private _pathMapper?: PathMapper; + private _syncDestination?: SyncDestination; private _projectConfigFile?: ProjectConfigFile; private _me?: string; private _profile?: string; @@ -40,9 +40,14 @@ export class ConnectionManager { new EventEmitter(); private readonly onChangeClusterEmitter: EventEmitter = new EventEmitter(); + private readonly onChangeSyncDestinationEmitter: EventEmitter< + SyncDestination | undefined + > = new EventEmitter(); public readonly onChangeState = this.onChangeStateEmitter.event; public readonly onChangeCluster = this.onChangeClusterEmitter.event; + public readonly onChangeSyncDestination = + this.onChangeSyncDestinationEmitter.event; constructor(private cli: CliWrapper) {} @@ -58,8 +63,16 @@ export class ConnectionManager { return this._state; } - get pathMapper(): PathMapper | undefined { - return this._pathMapper; + get cluster(): Cluster | undefined { + if (this.state === "DISCONNECTED") { + return; + } + + return this._cluster; + } + + get syncDestination(): SyncDestination | undefined { + return this._syncDestination; } /** @@ -123,13 +136,18 @@ export class ConnectionManager { this.updateState("CONNECTED"); if (projectConfigFile.config.clusterId) { - await this.attachCluster(projectConfigFile.config.clusterId); + await this.attachCluster(projectConfigFile.config.clusterId, false); + } else { + this.updateCluster(undefined); } if (projectConfigFile.config.workspacePath) { - await this.attachWorkspace( - Uri.file(projectConfigFile.config.workspacePath) + await this.attachSyncDestination( + Uri.file(projectConfigFile.config.workspacePath), + false ); + } else { + this.updateSyncDestination(undefined); } } @@ -196,31 +214,20 @@ export class ConnectionManager { throw new Error("Not in a VSCode workspace"); } - let projectConfigFile; - try { - projectConfigFile = await ProjectConfigFile.load( - vscodeWorkspace.rootPath - ); - } catch (e) { - projectConfigFile = new ProjectConfigFile( - {}, - vscodeWorkspace.rootPath - ); - } + const projectConfigFile = new ProjectConfigFile( + {}, + vscodeWorkspace.rootPath + ); projectConfigFile.profile = profile; - await projectConfigFile.write(); - } - - get cluster(): Cluster | undefined { - if (this.state === "DISCONNECTED") { - return; - } - return this._cluster; + await projectConfigFile.write(); } - async attachCluster(cluster: Cluster | string): Promise { + async attachCluster( + cluster: Cluster | string, + skipWrite = false + ): Promise { if (this._cluster === cluster) { return; } @@ -233,8 +240,10 @@ export class ConnectionManager { cluster = await Cluster.fromClusterId(this._apiClient!, cluster); } - this._projectConfigFile!.clusterId = cluster.id; - await this._projectConfigFile!.write(); + if (!skipWrite) { + this._projectConfigFile!.clusterId = cluster.id; + await this._projectConfigFile!.write(); + } this.updateCluster(cluster); } @@ -252,7 +261,10 @@ export class ConnectionManager { this.updateCluster(undefined); } - async attachWorkspace(workspacePath: Uri): Promise { + async attachSyncDestination( + workspacePath: Uri, + skipWrite = false + ): Promise { if ( !vscodeWorkspace.workspaceFolders || !vscodeWorkspace.workspaceFolders.length @@ -261,8 +273,26 @@ export class ConnectionManager { return; } + if (!skipWrite) { + this._projectConfigFile!.workspacePath = workspacePath.path; + await this._projectConfigFile!.write(); + } + const wsUri = vscodeWorkspace.workspaceFolders[0].uri; - this._pathMapper = new PathMapper(workspacePath, wsUri); + this.updateSyncDestination(new SyncDestination(workspacePath, wsUri)); + } + + async detachSyncDestination(): Promise { + if (!this._syncDestination) { + return; + } + + if (this._projectConfigFile) { + this._projectConfigFile.workspacePath = undefined; + await this._projectConfigFile.write(); + } + + this.updateSyncDestination(undefined); } private async getMe(apiClient: ApiClient): Promise { @@ -287,6 +317,15 @@ export class ConnectionManager { } } + private updateSyncDestination( + newSyncDestination: SyncDestination | undefined + ) { + if (this._syncDestination !== newSyncDestination) { + this._syncDestination = newSyncDestination; + this.onChangeSyncDestinationEmitter.fire(this._syncDestination); + } + } + private async waitForConnect() { if (this._state === "CONNECTED") { return; diff --git a/packages/databricks-vscode/src/configuration/PathMapper.ts b/packages/databricks-vscode/src/configuration/PathMapper.ts deleted file mode 100644 index 1ee31ccfb..000000000 --- a/packages/databricks-vscode/src/configuration/PathMapper.ts +++ /dev/null @@ -1,23 +0,0 @@ -import path = require("path"); -import {Uri} from "vscode"; - -/** - * Class that maps paths between the local file system to the file systems - * on the Databricks driver - */ -export class PathMapper { - constructor(readonly repoPath: Uri, readonly workspacePath: Uri) {} - - localToRemoteDir(localPath: Uri): string { - return path.dirname(this.localToRemote(localPath)); - } - - localToRemote(localPath: Uri): string { - let relativePath = localPath.path.replace(this.workspacePath.path, ""); - return Uri.joinPath(this.repoPath, relativePath).path; - } - - get remoteWorkspaceName(): string { - return path.basename(this.repoPath.path); - } -} diff --git a/packages/databricks-vscode/src/configuration/PathMapper.test.ts b/packages/databricks-vscode/src/configuration/SyncDestination.test.ts similarity index 50% rename from packages/databricks-vscode/src/configuration/PathMapper.test.ts rename to packages/databricks-vscode/src/configuration/SyncDestination.test.ts index 8007e4081..f17132b12 100644 --- a/packages/databricks-vscode/src/configuration/PathMapper.test.ts +++ b/packages/databricks-vscode/src/configuration/SyncDestination.test.ts @@ -1,10 +1,10 @@ import assert from "assert"; import {Uri} from "vscode"; -import {PathMapper} from "./PathMapper"; +import {SyncDestination} from "./SyncDestination"; describe(__filename, () => { it("should map a file", async () => { - let mapper = new PathMapper( + let mapper = new SyncDestination( Uri.file( "/Workspace/Repos/fabian.jakobs@databricks.com/notebook-best-practices" ), @@ -21,8 +21,26 @@ describe(__filename, () => { ); }); + it("should prepend '/Workspace' if missing", async () => { + let mapper = new SyncDestination( + Uri.file( + "/Repos/fabian.jakobs@databricks.com/notebook-best-practices" + ), + Uri.file("/Users/fabian.jakobs/Desktop/notebook-best-practices") + ); + + assert.equal( + mapper.localToRemote( + Uri.file( + "/Users/fabian.jakobs/Desktop/notebook-best-practices/hello.py" + ) + ), + "/Workspace/Repos/fabian.jakobs@databricks.com/notebook-best-practices/hello.py" + ); + }); + it("should map a directory", async () => { - let mapper = new PathMapper( + let mapper = new SyncDestination( Uri.file( "/Workspace/Repos/fabian.jakobs@databricks.com/notebook-best-practices" ), @@ -40,13 +58,31 @@ describe(__filename, () => { }); it("should get repo name", async () => { - let mapper = new PathMapper( + let mapper = new SyncDestination( Uri.file( "/Workspace/Repos/fabian.jakobs@databricks.com/notebook-best-practices" ), Uri.file("/Users/fabian.jakobs/Desktop/notebook-best-practices") ); - assert.equal(mapper.remoteWorkspaceName, "notebook-best-practices"); + assert.equal(mapper.name, "notebook-best-practices"); + }); + + it("should map notebooks", async () => { + let mapper = new SyncDestination( + Uri.file( + "/Repos/fabian.jakobs@databricks.com/notebook-best-practices" + ), + Uri.file("/Users/fabian.jakobs/Desktop/notebook-best-practices") + ); + + assert.equal( + mapper.localToRemoteNotebook( + Uri.file( + "/Users/fabian.jakobs/Desktop/notebook-best-practices/notebooks/covid_eda.py" + ) + ), + "/Repos/fabian.jakobs@databricks.com/notebook-best-practices/notebooks/covid_eda" + ); }); }); diff --git a/packages/databricks-vscode/src/configuration/SyncDestination.ts b/packages/databricks-vscode/src/configuration/SyncDestination.ts new file mode 100644 index 000000000..4a6786964 --- /dev/null +++ b/packages/databricks-vscode/src/configuration/SyncDestination.ts @@ -0,0 +1,64 @@ +import path = require("path"); +import {Uri} from "vscode"; + +type SyncDestinationType = "workspace" | "repo"; + +/** + * Either Databricks repo or workspace that acts as a sync target for the current workspace. + */ +export class SyncDestination { + private repoPath: Uri; + + constructor(repoPath: Uri, readonly vscodeWorkspacePath: Uri) { + this.repoPath = repoPath; + + // Repo paths always start with "/Workspace" but the repos API strips this off. + if (!this.repoPath.path.startsWith("/Workspace/")) { + this.repoPath = Uri.file(`/Workspace${this.repoPath.path}`); + } + } + + get type(): SyncDestinationType { + return "repo"; + } + + get name(): string { + return path.basename(this.repoPath.path); + } + + get path(): Uri { + return this.repoPath; + } + + /** + * Maps a local notebook to notebook path used in workflow deifnitions. + */ + localToRemoteNotebook(localPath: Uri): string { + return this.localToRemote(localPath).replace( + /^\/Workspace(\/.*).py/g, + "$1" + ); + } + + /** + * Maps a local file path to the remote directory containing the file. + */ + localToRemoteDir(localPath: Uri): string { + return path.dirname(this.localToRemote(localPath)); + } + + /** + * Maps a local file path to the remote file path whre it gets synced to. + */ + localToRemote(localPath: Uri): string { + if (!localPath.path.startsWith(this.vscodeWorkspacePath.path)) { + throw new Error("local path is not within the workspace"); + } + + let relativePath = localPath.path.replace( + this.vscodeWorkspacePath.path, + "" + ); + return Uri.joinPath(this.repoPath, relativePath).path; + } +} diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 81fe6a6ce..d884bf2a8 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -53,10 +53,25 @@ export function activate(context: ExtensionContext) { connectionCommands.attachClusterCommand(), connectionCommands ), + commands.registerCommand( + "databricks.connection.attachClusterQuickPick", + connectionCommands.attachClusterQuickPickCommand(), + connectionCommands + ), commands.registerCommand( "databricks.connection.detachCluster", connectionCommands.detachClusterCommand(), connectionCommands + ), + commands.registerCommand( + "databricks.connection.attachSyncDestination", + connectionCommands.attachSyncDestinationCommand(), + connectionCommands + ), + commands.registerCommand( + "databricks.connection.detachSyncDestination", + connectionCommands.detachWorkspaceCommand(), + connectionCommands ) ); @@ -65,7 +80,7 @@ export function activate(context: ExtensionContext) { context.subscriptions.push( commands.registerCommand( "databricks.run.runEditorContentsAsWorkflow", - workflowCommands.runEditorContentsAsWorkflow(), + workflowCommands.runEditorContentsAsWorkflowCommand(), workflowCommands ) ); diff --git a/packages/databricks-vscode/src/test/e2e/configure.e2e.ts b/packages/databricks-vscode/src/test/e2e/configure.e2e.ts index 8f4032fd0..48d5e4bc0 100644 --- a/packages/databricks-vscode/src/test/e2e/configure.e2e.ts +++ b/packages/databricks-vscode/src/test/e2e/configure.e2e.ts @@ -29,6 +29,9 @@ describe("Configure Databricks Extension", function () { browser = VSBrowser.instance; driver = browser.driver; + assert(process.env.TEST_DEFAULT_CLUSTER_ID); + clusterId = process.env.TEST_DEFAULT_CLUSTER_ID; + ({path: projectDir, cleanup} = await tmp.dir()); await openFolder(browser, projectDir); }); @@ -77,10 +80,9 @@ describe("Configure Databricks Extension", function () { assert(labels.length > 0); }); + // test is skipped because context menus currently don't work in vscode-extension-tester + // https://github.com/redhat-developer/vscode-extension-tester/issues/444 it.skip("should filter clusters", async () => { - // test is skipped because context menus currently don't work in vscode-extension-tester - // https://github.com/redhat-developer/vscode-extension-tester/issues/444 - const section = await getViewSection("Clusters"); assert(section); const action = await section!.getAction("Filter clusters ..."); @@ -96,26 +98,6 @@ describe("Configure Databricks Extension", function () { }); it("should attach cluster", async () => { - const section = await getViewSection("Clusters"); - assert(section); - - const items = await section.getVisibleItems(); - assert(items.length > 0); - - // find top level cluster tree item - let item: TreeItem | undefined; - for (const i of items) { - if (await (i as TreeItem).hasChildren()) { - item = i as TreeItem; - break; - } - } - assert(item); - - const buttons = await (item as TreeItem).getActionButtons(); - await buttons[0].click(); - - // check if cluster is attached const config = await getViewSection("Configuration"); assert(config); const configTree = config as CustomTreeSection; @@ -125,16 +107,29 @@ describe("Configure Databricks Extension", function () { const configItems = await configTree.getVisibleItems(); let clusterConfigItem: TreeItem | undefined; - const itemLabel = await item.getLabel(); for (const i of configItems) { const label = await i.getLabel(); - if (label === `Cluster: ${itemLabel}`) { + if (label.startsWith("Cluster")) { clusterConfigItem = i; break; } } assert(clusterConfigItem); + const buttons = await ( + clusterConfigItem as TreeItem + ).getActionButtons(); + await buttons[0].click(); + + const input = await InputBox.create(); + while (await input.hasProgress()) { + await driver.sleep(200); + } + + await input.setText(clusterId); + await input.confirm(); + await input.selectQuickPick(0); + // get cluster ID const clusterPropsItems = await clusterConfigItem.getChildren(); const clusterProps: Record = {}; diff --git a/packages/databricks-vscode/src/workflow/WorkflowCommands.ts b/packages/databricks-vscode/src/workflow/WorkflowCommands.ts index 35d9baef8..f669beeb7 100644 --- a/packages/databricks-vscode/src/workflow/WorkflowCommands.ts +++ b/packages/databricks-vscode/src/workflow/WorkflowCommands.ts @@ -14,7 +14,7 @@ export class WorkflowCommands { /** * Run a Python file or notebook as a workflow on the connected cluster */ - runEditorContentsAsWorkflow() { + runEditorContentsAsWorkflowCommand() { return async (resource: Uri) => { let targetResource = resource; if (!targetResource && window.activeTextEditor) { @@ -31,8 +31,8 @@ export class WorkflowCommands { return; } - let pathMapper = this.connectionManager.pathMapper; - if (!pathMapper) { + let syncDestination = this.connectionManager.syncDestination; + if (!syncDestination) { window.showErrorMessage( "You must configure code synchronization to run a workflow" ); @@ -52,7 +52,7 @@ export class WorkflowCommands { await runNotebookAsWorkflow({ notebookUri: targetResource, cluster, - pathMapper, + syncDestination: syncDestination, context: this.context, }); } diff --git a/packages/databricks-vscode/src/workflow/WorkflowOutputPanel.ts b/packages/databricks-vscode/src/workflow/WorkflowOutputPanel.ts index fd178b4b6..8a668c29e 100644 --- a/packages/databricks-vscode/src/workflow/WorkflowOutputPanel.ts +++ b/packages/databricks-vscode/src/workflow/WorkflowOutputPanel.ts @@ -9,19 +9,19 @@ import { WebviewPanel, window, } from "vscode"; -import {PathMapper} from "../configuration/PathMapper"; +import {SyncDestination} from "../configuration/SyncDestination"; // TODO: add dispose, add persistence export async function runNotebookAsWorkflow({ notebookUri, cluster, - pathMapper, + syncDestination, context, }: { notebookUri: Uri; cluster: Cluster; - pathMapper: PathMapper; + syncDestination: SyncDestination; context: ExtensionContext; }) { const panel = new WorkflowOutputPanel( @@ -40,15 +40,11 @@ export async function runNotebookAsWorkflow({ const cancellation = new CancellationTokenSource(); panel.onDidDispose(() => cancellation.cancel()); - const clusterNotebookPath = pathMapper - .localToRemote(notebookUri) - .replace(/^\/Workspace(\/.*).py/g, "$1"); - try { let response = await cluster.runNotebookAndWait({ - path: clusterNotebookPath, + path: syncDestination.localToRemoteNotebook(notebookUri), onProgress: (state: jobs.RunLifeCycleState, run: WorkflowRun) => { - panel.updateState(state, run.runPageUrl); + panel.updateState(state, run); }, token: cancellation.token, }); @@ -79,12 +75,17 @@ export class WorkflowOutputPanel { this.panel.webview.html = htmlContent; } - updateState(state: jobs.RunLifeCycleState, pageUrl: string) { - this.panel.webview.postMessage({ - type: "status", - state, - pageUrl, - }); + updateState(state: jobs.RunLifeCycleState, run: WorkflowRun) { + if (state === "INTERNAL_ERROR") { + // TODO + this.showError(run.state!.state_message!); + } else { + this.panel.webview.postMessage({ + type: "status", + state, + pageUrl: run.runPageUrl, + }); + } } showError(error: string) { @@ -126,7 +127,7 @@ export class WorkflowOutputPanel { case "error": const pre = document.createElement("pre"); - pre.innterText = event.data.error; + pre.innerText = event.data.error; messageEl.appendChild(pre); clearInterval(interval);