Add run as job functionality (#39)

fjakobs · web-flow · commit 5ef0919c9c48 · 2022-08-02T10:40:51.000+02:00
diff --git a/.vscode/extensions.json b/.vscode/extensions.json
@@ -4,6 +4,7 @@
     "recommendations": [
         "dbaeumer.vscode-eslint",
         "amodio.tsl-problem-matcher",
-        "esbenp.prettier-vscode"
+        "esbenp.prettier-vscode",
+        "Tobermory.es6-string-html"
     ]
 }
diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json
@@ -46,6 +46,8 @@
         "onCommand:databricks.cluster.filterByAll",
         "onCommand:databricks.cluster.filterByMe",
         "onCommand:databricks.cluster.filterByRunning",
+        "onCommand:databricks.run.getProgramName",
+        "onCommand:databricks.run.runEditorContentsAsWorkflow",
         "onView:clusterManager",
         "onView:clusterList"
     ],
@@ -95,6 +97,13 @@
                 "command": "databricks.cluster.refresh",
                 "icon": "$(refresh)",
                 "title": "Refresh"
+            },
+            {
+                "command": "databricks.run.runEditorContentsAsWorkflow",
+                "title": "Run File as Workflow on Databricks",
+                "category": "Databricks",
+                "enablement": "!inDebugMode",
+                "icon": "$(play)"
             }
         ],
         "viewsContainers": {
@@ -169,6 +178,19 @@
                     "when": "view == configurationView && viewItem == clusterStopped",
                     "group": "inline@0"
                 }
+            ],
+            "editor/title/run": [
+                {
+                    "command": "databricks.run.runEditorContentsAsWorkflow",
+                    "when": "resourceLangId == python",
+                    "group": "navigation@1"
+                }
+            ],
+            "commandPalette": [
+                {
+                    "command": "databricks.run.runEditorContentsAsWorkflow",
+                    "when": "resourceLangId == python"
+                }
             ]
         },
         "submenus": [
@@ -193,11 +215,11 @@
         "test:lint": "eslint src --ext ts && prettier . -c",
         "test:compile": "tsc -p . --outDir out",
         "test:watch": "tsc -p . -w --outDir out",
-        "test:unit": "node ./out/test/runTest.js",
+        "test:unit": "yarn run test:compile && node ./out/test/runTest.js",
         "test:integ:prepare": "yarn run vscode:package && extest get-vscode --type ${VSCODE_TEST_VERSION:-stable} --storage /tmp/vscode-test-databricks && extest get-chromedriver --storage /tmp/vscode-test-databricks && extest install-vsix -f databricks-0.0.1.vsix --storage /tmp/vscode-test-databricks",
         "test:integ:run": "yarn run test:compile && ts-node src/test/e2e/scripts/e2e.ts --storage /tmp/vscode-test-databricks --code_settings src/test/e2e/settings.json 'out/**/*.e2e.js'",
         "test:integ": "yarn run test:integ:prepare && yarn run test:integ:run",
-        "test": "yarn run test:compile && yarn run compile && yarn run test:lint && yarn run test:unit",
+        "test": "yarn run compile && yarn run test:lint && yarn run test:unit",
         "build": "yarn run vscode:prepublish",
         "clean": "rm -rf node_modules out dist .vscode-test"
     },
@@ -227,4 +249,4 @@
         "webpack": "^5.73.0",
         "webpack-cli": "^4.10.0"
     }
-}
+}
diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts
@@ -1,7 +1,10 @@
-import {execFile, ExecFileException, spawn} from "child_process";
+import {spawn} from "child_process";
 
 /**
  * Entrypoint for all wrapped CLI commands
+ *
+ * Righ now this is a placeholder for a future implementation
+ * of the bricks CLI
  */
 export class CliWrapper {
     constructor() {}
diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.test.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.test.ts
@@ -0,0 +1,60 @@
+/* eslint-disable @typescript-eslint/naming-convention */
+
+import {mock, instance} from "ts-mockito";
+import {cluster} from "@databricks/databricks-sdk";
+import {ConnectionManager} from "./ConnectionManager";
+import {Disposable} from "vscode";
+import {CliWrapper} from "../cli/CliWrapper";
+
+const me = "user-1";
+const mockListClustersResponse: cluster.ListClustersResponse = {
+    clusters: [
+        {
+            cluster_id: "cluster-id-2",
+            cluster_name: "cluster-name-2",
+            cluster_source: "UI",
+            creator_user_name: "user-2",
+            state: "TERMINATED",
+        },
+        {
+            cluster_id: "cluster-id-1",
+            cluster_name: "cluster-name-1",
+            cluster_source: "UI",
+            creator_user_name: me,
+            state: "RUNNING",
+        },
+        {
+            cluster_id: "cluster-id-3",
+            cluster_name: "cluster-name-3",
+            cluster_source: "JOB",
+            creator_user_name: "user-3",
+            state: "RUNNING",
+        },
+    ],
+};
+
+describe(__filename, () => {
+    let mockedCliWrapper: CliWrapper;
+    let disposables: Array<Disposable>;
+
+    beforeEach(() => {
+        disposables = [];
+        mockedCliWrapper = mock(CliWrapper);
+    });
+
+    afterEach(() => {
+        disposables.forEach((d) => d.dispose());
+    });
+
+    it("should create an instance", async () => {
+        let manager = new ConnectionManager(instance(mockedCliWrapper));
+    });
+
+    // login
+    // logout
+    // configure
+    // attach cluster
+    // detach cluster
+    // attach workspace
+    // detach workspace
+});
diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts
@@ -4,8 +4,15 @@ import {
     fromConfigFile,
     ScimService,
 } from "@databricks/databricks-sdk";
-import {commands, EventEmitter, window, workspace} from "vscode";
+import {
+    commands,
+    EventEmitter,
+    Uri,
+    window,
+    workspace as vscodeWorkspace,
+} from "vscode";
 import {CliWrapper} from "../cli/CliWrapper";
+import {PathMapper} from "./PathMapper";
 import {ProjectConfigFile} from "./ProjectConfigFile";
 import {selectProfile} from "./selectProfileWizard";
 
@@ -21,6 +28,7 @@ export class ConnectionManager {
     private _state: ConnectionState = "DISCONNECTED";
     private _cluster?: Cluster;
     private _apiClient?: ApiClient;
+    private _pathMapper?: PathMapper;
     private _projectConfigFile?: ProjectConfigFile;
     private _me?: string;
     private _profile?: string;
@@ -47,6 +55,10 @@ export class ConnectionManager {
         return this._state;
     }
 
+    get pathMapper(): PathMapper | undefined {
+        return this._pathMapper;
+    }
+
     /**
      * Get a pre-configured APIClient. Do not hold on to references to this class as
      * it might be invalidated as the configuration changes. If you have to store a reference
@@ -64,14 +76,14 @@ export class ConnectionManager {
         let profile;
 
         try {
-            if (!workspace.rootPath) {
+            if (!vscodeWorkspace.rootPath) {
                 throw new Error(
                     "Can't login to Databricks: Not in a VSCode workspace"
                 );
             }
 
             projectConfigFile = await ProjectConfigFile.load(
-                workspace.rootPath
+                vscodeWorkspace.rootPath
             );
 
             profile = projectConfigFile.config.profile;
@@ -106,6 +118,12 @@ export class ConnectionManager {
         if (projectConfigFile.config.clusterId) {
             await this.attachCluster(projectConfigFile.config.clusterId);
         }
+
+        if (projectConfigFile.config.workspacePath) {
+            await this.attachWorkspace(
+                Uri.file(projectConfigFile.config.workspacePath)
+            );
+        }
     }
 
     async logout() {
@@ -167,17 +185,20 @@ export class ConnectionManager {
     }
 
     private async writeConfigFile(profile: string) {
-        if (!workspace.rootPath) {
+        if (!vscodeWorkspace.rootPath) {
             throw new Error("Not in a VSCode workspace");
         }
 
         let projectConfigFile;
         try {
             projectConfigFile = await ProjectConfigFile.load(
-                workspace.rootPath
+                vscodeWorkspace.rootPath
             );
         } catch (e) {
-            projectConfigFile = new ProjectConfigFile({}, workspace.rootPath);
+            projectConfigFile = new ProjectConfigFile(
+                {},
+                vscodeWorkspace.rootPath
+            );
         }
 
         projectConfigFile.profile = profile;
@@ -224,6 +245,19 @@ export class ConnectionManager {
         this.updateCluster(undefined);
     }
 
+    async attachWorkspace(workspacePath: Uri): Promise<void> {
+        if (
+            !vscodeWorkspace.workspaceFolders ||
+            !vscodeWorkspace.workspaceFolders.length
+        ) {
+            // TODO how do we handle this?
+            return;
+        }
+
+        const wsUri = vscodeWorkspace.workspaceFolders[0].uri;
+        this._pathMapper = new PathMapper(workspacePath, wsUri);
+    }
+
     private async getMe(apiClient: ApiClient): Promise<string> {
         let scimApi = new ScimService(apiClient);
         let response = await scimApi.me({});
diff --git a/packages/databricks-vscode/src/configuration/PathMapper.test.ts b/packages/databricks-vscode/src/configuration/PathMapper.test.ts
@@ -0,0 +1,41 @@
+import assert from "assert";
+import {Uri} from "vscode";
+import {PathMapper} from "./PathMapper";
+
+describe(__filename, () => {
+    it("should map a file", async () => {
+        let mapper = new PathMapper(
+            Uri.file(
+                "/Workspace/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(
+            Uri.file(
+                "/Workspace/Repos/fabian.jakobs@databricks.com/notebook-best-practices"
+            ),
+            Uri.file("/Users/fabian.jakobs/Desktop/notebook-best-practices")
+        );
+
+        assert.equal(
+            mapper.localToRemoteDir(
+                Uri.file(
+                    "/Users/fabian.jakobs/Desktop/notebook-best-practices/jobs/hello.py"
+                )
+            ),
+            "/Workspace/Repos/fabian.jakobs@databricks.com/notebook-best-practices/jobs"
+        );
+    });
+});
diff --git a/packages/databricks-vscode/src/configuration/PathMapper.ts b/packages/databricks-vscode/src/configuration/PathMapper.ts
@@ -0,0 +1,23 @@
+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 {
+    readonly repo: string;
+
+    constructor(readonly repoPath: Uri, readonly workspacePath: Uri) {
+        this.repo = path.basename(repoPath.path);
+    }
+
+    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;
+    }
+}
diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts
@@ -6,6 +6,7 @@ import {ClusterListDataProvider} from "./cluster/ClusterListDataProvider";
 import {ClusterModel} from "./cluster/ClusterModel";
 import {ClusterCommands} from "./cluster/ClusterCommands";
 import {ConfigurationDataProvider} from "./configuration/ConfigurationDataProvider";
+import {WorkflowCommands} from "./workflow/WorkflowCommands";
 
 export function activate(context: ExtensionContext) {
     let cli = new CliWrapper();
@@ -58,6 +59,16 @@ export function activate(context: ExtensionContext) {
         )
     );
 
+    // Workflow group
+    const workflowCommands = new WorkflowCommands(connectionManager, context);
+    context.subscriptions.push(
+        commands.registerCommand(
+            "databricks.run.runEditorContentsAsWorkflow",
+            workflowCommands.runEditorContentsAsWorkflow(),
+            workflowCommands
+        )
+    );
+
     // Cluster group
     const clusterModel = new ClusterModel(connectionManager);
     const clusterTreeDataProvider = new ClusterListDataProvider(clusterModel);
diff --git a/packages/databricks-vscode/src/workflow/WorkflowCommands.ts b/packages/databricks-vscode/src/workflow/WorkflowCommands.ts
diff --git a/packages/databricks-vscode/src/workflow/WorkflowOutputPanel.ts b/packages/databricks-vscode/src/workflow/WorkflowOutputPanel.ts

Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,7 @@`
`4`	`4`	`"recommendations": [`
`5`	`5`	`"dbaeumer.vscode-eslint",`
`6`	`6`	`"amodio.tsl-problem-matcher",`
`7`		`- "esbenp.prettier-vscode"`
	`7`	`+ "esbenp.prettier-vscode",`
	`8`	`+ "Tobermory.es6-string-html"`
`8`	`9`	`]`
`9`	`10`	`}`