Prompt for cluster start (#95)

kartikgupta-db · web-flow · commit adf1bf404ec4 · 2022-09-20T10:01:24.000Z
- Fixes #92 - Also propagates error message on API failures when executing a python file ## Stop a cluster https://user-images.githubusercontent.com/88345179/189954436-673bc4b4-150b-4279-8005-97af430f0a1b.mov ## Start a stopped cluster https://user-images.githubusercontent.com/88345179/189954448-a6ef20b1-72a3-4081-8ab1-20637d94a1c1.mov ## Prompts when running the jobs https://user-images.githubusercontent.com/88345179/189954497-7315a98b-7258-4a13-b4ee-c5a0c4538eb7.mov ## Starting cluster from prompt shows `Cluster starting...` in debug output https://user-images.githubusercontent.com/88345179/190226346-7dfcc018-4bd8-4239-bc70-6b72287fcd6c.mov ## A job run after waiting for cluster start <img width="1445" alt="Screenshot 2022-09-13 at 18 16 35" src="https://user-images.githubusercontent.com/88345179/189954685-d836ddfc-630d-46b5-be20-1907d1a89387.png">
diff --git a/packages/databricks-sdk-js/src/index.ts b/packages/databricks-sdk-js/src/index.ts
@@ -33,3 +33,8 @@ export * from "./auth/fromConfigFile";
 export * from "./auth/configFile";
 
 export * from "./types";
+
+export {ClusterFixture, TokenFixture} from "./test/fixtures";
+
+export {RetryConfigs, default as retry} from "./retries/retries";
+export {TimeUnits, default as Time} from "./retries/Time";
diff --git a/packages/databricks-sdk-js/src/services/Cluster.test.ts b/packages/databricks-sdk-js/src/services/Cluster.test.ts
@@ -5,9 +5,9 @@ import chai, {assert} from "chai";
 import {mock, when, instance, deepEqual, verify, anything} from "ts-mockito";
 import chaiAsPromised from "chai-as-promised";
 import Time, {TimeUnits} from "../retries/Time";
-import getMockTestCluster from "../test/fixtures/ClusterFixtures";
+import {getMockTestCluster} from "../test/fixtures/ClusterFixtures";
 import {ClusterInfo} from "../apis/clusters";
-import TokenFixture from "../test/fixtures/TokenFixture";
+import {TokenFixture} from "../test/fixtures/TokenFixtures";
 import {RetryConfigs} from "../retries/retries";
 
 chai.use(chaiAsPromised);
diff --git a/packages/databricks-sdk-js/src/services/Cluster.ts b/packages/databricks-sdk-js/src/services/Cluster.ts
@@ -80,7 +80,10 @@ export class Cluster {
         });
     }
 
-    async start(token?: CancellationToken) {
+    async start(
+        token?: CancellationToken,
+        onProgress: (state: ClusterInfoState) => void = (state) => {}
+    ) {
         await this.refresh();
         if (this.state === "RUNNING") {
             return;
@@ -104,6 +107,7 @@ export class Cluster {
                 }
 
                 await this.refresh();
+                onProgress(this.state);
 
                 switch (this.state) {
                     case "RUNNING":
diff --git a/packages/databricks-sdk-js/src/services/ExecutionContext.integ.ts b/packages/databricks-sdk-js/src/services/ExecutionContext.integ.ts
@@ -5,7 +5,7 @@ import assert from "assert";
 
 import {IntegrationTestSetup} from "../test/IntegrationTestSetup";
 import {mock, when, instance} from "ts-mockito";
-import TokenFixture from "../test/fixtures/TokenFixture";
+import {TokenFixture} from "../test/fixtures/TokenFixtures";
 
 describe(__filename, function () {
     let integSetup: IntegrationTestSetup;
diff --git a/packages/databricks-sdk-js/src/test/fixtures/ClusterFixtures.ts b/packages/databricks-sdk-js/src/test/fixtures/ClusterFixtures.ts
@@ -9,7 +9,7 @@ const testClusterDetails: ClusterInfo = {
     cluster_name: "testClusterName",
 };
 
-export default async function getMockTestCluster() {
+export async function getMockTestCluster() {
     const mockedClient = mock(ApiClient);
     when(
         mockedClient.request(
diff --git a/packages/databricks-sdk-js/src/test/fixtures/TokenFixtures.ts b/packages/databricks-sdk-js/src/test/fixtures/TokenFixtures.ts
@@ -1,6 +1,6 @@
 import {CancellationToken} from "../../types";
 
-export default class TokenFixture implements CancellationToken {
+export class TokenFixture implements CancellationToken {
     get isCancellationRequested() {
         return false;
     }
diff --git a/packages/databricks-sdk-js/src/test/fixtures/index.ts b/packages/databricks-sdk-js/src/test/fixtures/index.ts
@@ -0,0 +1,2 @@
+export * as ClusterFixture from "./ClusterFixtures";
+export * as TokenFixture from "./TokenFixtures";
diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json
@@ -162,6 +162,15 @@
             {
                 "command": "databricks.quickstart.open",
                 "title": "Databricks: Open Quickstart"
+            },
+            {  "command": "databricks.cluster.start",
+                "title": "Start Cluster",
+                "icon": "$(debug-start)"
+            },
+            {
+                "command": "databricks.cluster.stop",
+                "title": "Stop Cluster",
+                "icon": "$(debug-stop)"
             }
         ],
         "viewsContainers": {
@@ -229,12 +238,12 @@
                 {
                     "command": "databricks.connection.attachClusterQuickPick",
                     "when": "view == configurationView && viewItem == clusterDetached",
-                    "group": "inline@0"
+                    "group": "inline@1"
                 },
                 {
                     "command": "databricks.connection.detachCluster",
-                    "when": "view == configurationView && viewItem == clusterAttached",
-                    "group": "inline@0"
+                    "when": "view == configurationView && viewItem == clusterRunning || view == configurationView && viewItem == clusterStopped || view == configurationView && viewItem == clusterPending",
+                    "group": "inline@1"
                 },
                 {
                     "command": "databricks.connection.attachSyncDestination",
@@ -255,6 +264,16 @@
                     "command": "databricks.sync.stop",
                     "when": "view == configurationView && viewItem == syncRunning",
                     "group": "inline@0"
+                },
+                {
+                    "command": "databricks.cluster.stop",
+                    "when": "view == configurationView && viewItem == clusterRunning || view == configurationView && viewItem == clusterPending",
+                    "group": "inline@0"
+                },
+                {
+                    "command": "databricks.cluster.start",
+                    "when": "view == configurationView && viewItem == clusterStopped",
+                    "group": "inline@0"
                 }
             ],
             "editor/title/run": [
diff --git a/packages/databricks-vscode/src/cluster/ClusterCommands.ts b/packages/databricks-vscode/src/cluster/ClusterCommands.ts
@@ -1,10 +1,16 @@
+import {Cluster} from "@databricks/databricks-sdk";
+import {TreeItem} from "vscode";
+import {ConnectionManager} from "../configuration/ConnectionManager";
 import {ClusterFilter, ClusterModel} from "./ClusterModel";
 
 /**
  * Cluster related commands
  */
 export class ClusterCommands {
-    constructor(private clusterModel: ClusterModel) {}
+    constructor(
+        private clusterModel: ClusterModel,
+        readonly connectionManager: ConnectionManager
+    ) {}
 
     /**
      * Refresh cluster tree view by reloading them throug the API
@@ -23,4 +29,20 @@ export class ClusterCommands {
             this.clusterModel.filter = filter;
         };
     }
+
+    /**
+     * Command to start a cluster
+     * @param cluster cluster to be started
+     */
+    async startClusterCommand() {
+        await this.connectionManager.startCluster();
+    }
+
+    /**
+     * Command to stop a cluster
+     * @param cluster cluster to be stopped
+     */
+    async stopClusterCommand() {
+        await this.connectionManager.stopCluster();
+    }
 }
diff --git a/packages/databricks-vscode/src/cluster/ClusterListDataProvider.test.ts b/packages/databricks-vscode/src/cluster/ClusterListDataProvider.test.ts
@@ -149,7 +149,6 @@ describe(__filename, () => {
         let item = ClusterListDataProvider.clusterNodeToTreeItem(cluster);
         assert.deepEqual(item, {
             collapsibleState: 1,
-            contextValue: "clusterStopped",
             iconPath: {
                 color: undefined,
                 id: "debug-stop",
@@ -166,7 +165,6 @@ describe(__filename, () => {
         item = ClusterListDataProvider.clusterNodeToTreeItem(cluster);
         assert.deepEqual(item, {
             collapsibleState: 1,
-            contextValue: "clusterRunning",
             iconPath: {
                 color: undefined,
                 id: "debug-start",
diff --git a/packages/databricks-vscode/src/cluster/ClusterListDataProvider.ts b/packages/databricks-vscode/src/cluster/ClusterListDataProvider.ts
@@ -112,10 +112,6 @@ export class ClusterListDataProvider
             iconPath: icon,
             id: element.id,
             collapsibleState: TreeItemCollapsibleState.Collapsed,
-            contextValue:
-                element.state === "RUNNING"
-                    ? "clusterRunning"
-                    : "clusterStopped",
         };
     }
 
diff --git a/packages/databricks-vscode/src/cluster/ClusterManager.test.ts b/packages/databricks-vscode/src/cluster/ClusterManager.test.ts
@@ -0,0 +1,91 @@
+import {describe} from "mocha";
+import {ClusterManager} from "./ClusterManager";
+import {
+    ApiClient,
+    cluster,
+    Cluster,
+    ClusterFixture,
+    RetryConfigs,
+    Time,
+    TimeUnits,
+} from "@databricks/databricks-sdk";
+import {deepEqual, instance, mock, resetCalls, verify, when} from "ts-mockito";
+import assert from "assert";
+
+describe(__filename, async () => {
+    let mockedCluster: Cluster;
+    let mockedClient: ApiClient;
+    let testClusterDetails: cluster.ClusterInfo;
+
+    beforeEach(async () => {
+        ({testClusterDetails} = await ClusterFixture.getMockTestCluster());
+        mockedClient = mock(ApiClient);
+        when(
+            mockedClient.request(
+                "/api/2.0/clusters/get",
+                "GET",
+                deepEqual({
+                    cluster_id: testClusterDetails.cluster_id,
+                })
+            )
+        ).thenResolve({
+            ...testClusterDetails,
+            state: "RUNNING",
+        });
+        mockedCluster = await Cluster.fromClusterId(
+            instance(mockedClient),
+            testClusterDetails.cluster_id!
+        );
+
+        resetCalls(mockedClient);
+
+        RetryConfigs.waitTime = (attempt) => {
+            return new Time(0, TimeUnits.milliseconds);
+        };
+    });
+
+    it("should start a cluster with progress", async () => {
+        when(
+            mockedClient.request(
+                "/api/2.0/clusters/get",
+                "GET",
+                deepEqual({
+                    cluster_id: testClusterDetails.cluster_id,
+                })
+            )
+        ).thenResolve(
+            {
+                ...testClusterDetails,
+                state: "TERMINATED",
+            },
+            {
+                ...testClusterDetails,
+                state: "PENDING",
+            },
+            {
+                ...testClusterDetails,
+                state: "PENDING",
+            },
+            {
+                ...testClusterDetails,
+                state: "RUNNING",
+            }
+        );
+
+        await mockedCluster.refresh();
+        assert.equal(mockedCluster.state, "TERMINATED");
+        interface OnProgContainer {
+            onProgress: (state: cluster.ClusterInfoState) => void;
+        }
+        const mockOnProgContainer = mock<OnProgContainer>();
+        await new ClusterManager(mockedCluster).start(
+            instance(mockOnProgContainer).onProgress
+        );
+
+        verify(mockOnProgContainer.onProgress("RUNNING")).calledAfter(
+            mockOnProgContainer.onProgress("PENDING")
+        );
+
+        assert.equal(mockedCluster.state, "RUNNING");
+    });
+});
diff --git a/packages/databricks-vscode/src/cluster/ClusterManager.ts b/packages/databricks-vscode/src/cluster/ClusterManager.ts
@@ -0,0 +1,40 @@
+import {cluster, Cluster} from "@databricks/databricks-sdk";
+import {CancellationToken, CancellationTokenSource, Disposable} from "vscode";
+import {ConnectionManager} from "../configuration/ConnectionManager";
+
+export class ClusterManager implements Disposable {
+    private cancellationTokenSource?: CancellationTokenSource;
+    // cancellationToken?: CancellationToken
+
+    constructor(readonly cluster: Cluster) {}
+
+    dispose() {
+        this.cancellationTokenSource?.cancel();
+        this.cancellationTokenSource?.dispose();
+    }
+
+    async start(
+        onProgress: (state: cluster.ClusterInfoState) => void = (state) => {}
+    ) {
+        this.cancellationTokenSource?.cancel();
+        this.cancellationTokenSource = new CancellationTokenSource();
+
+        await this.cluster.start(
+            this.cancellationTokenSource.token,
+            onProgress
+        );
+
+        onProgress(this.cluster.state);
+    }
+
+    async stop(
+        onProgress: (state: cluster.ClusterInfoState) => void = (state) => {}
+    ) {
+        this.cancellationTokenSource?.cancel();
+        this.cancellationTokenSource = new CancellationTokenSource();
+
+        //TODO: add cancellation and onProgress cb after adding these to API generator
+        await this.cluster.stop();
+        onProgress(this.cluster.state);
+    }
+}
diff --git a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.test.ts b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.test.ts
@@ -130,7 +130,7 @@ describe(__filename, () => {
             },
             {
                 collapsibleState: 2,
-                contextValue: "clusterAttached",
+                contextValue: "clusterStopped",
                 iconPath: {
                     color: undefined,
                     id: "server",
diff --git a/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts b/packages/databricks-vscode/src/configuration/ConfigurationDataProvider.ts
@@ -74,11 +74,16 @@ export class ConfigurationDataProvider
 
                 if (cluster) {
                     children.push({
-                        label: `Cluster`,
+                        label: "Cluster",
                         iconPath: new ThemeIcon("server"),
                         id: "CLUSTER",
                         collapsibleState: TreeItemCollapsibleState.Expanded,
-                        contextValue: "clusterAttached",
+                        contextValue:
+                            cluster.state === "RUNNING"
+                                ? "clusterRunning"
+                                : cluster.state === "PENDING"
+                                ? "clusterPending"
+                                : "clusterStopped",
                     });
                 } else {
                     children.push({
@@ -124,7 +129,7 @@ export class ConfigurationDataProvider
                 ];
             }
 
-            if (element.id === "CLUSTER" && cluster) {
+            if (element.id?.startsWith("CLUSTER") && cluster) {
                 let clusterItem =
                     ClusterListDataProvider.clusterNodeToTreeItem(cluster);
 
diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts
diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts
diff --git a/packages/databricks-vscode/src/run/DabaricksRuntime.test.ts b/packages/databricks-vscode/src/run/DabaricksRuntime.test.ts
diff --git a/packages/databricks-vscode/src/run/DabaricksRuntime.ts b/packages/databricks-vscode/src/run/DabaricksRuntime.ts
diff --git a/packages/databricks-vscode/src/run/DabaricksWorkflowDebugAdapter.ts b/packages/databricks-vscode/src/run/DabaricksWorkflowDebugAdapter.ts
diff --git a/packages/databricks-vscode/src/ui/prompts.ts b/packages/databricks-vscode/src/ui/prompts.ts

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`import {CancellationToken} from "../../types";`
`2`	`2`
`3`		`-export default class TokenFixture implements CancellationToken {`
	`3`	`+export class TokenFixture implements CancellationToken {`
`4`	`4`	`get isCancellationRequested() {`
`5`	`5`	`return false;`
`6`	`6`	`}`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+export * as ClusterFixture from "./ClusterFixtures";`
	`2`	`+export * as TokenFixture from "./TokenFixtures";`