Prevent running resources when cli already launched but no run status (#1072)

## Changes
* When we launch a run, there is a lag between the CLI process execution
and when the IDE is able to poll for the first status update from the
API. We should not allow a second run of the resource during this time.
This PR blocks second runs by considering the "unknown" status also as a
running state for the viewitem context. This allows vscode to hide the
run button, even when we, we do not have a status update. Errors and
cancellations are still handled gracefully, because the status reaches
completed state.

* Also, clear the run statuses map on target change. 
## Tests
<!-- How is this tested? -->

Author: kartikgupta-db

packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts

@@ -30,6 +30,17 @@ export type BundleRemoteState = BundleTarget & {
 };
 
 /* eslint-enable @typescript-eslint/naming-convention */
+export function getResource(
+    key: string,
+    resources?: BundleRemoteState["resources"]
+) {
+    return key.split(".").reduce((prev: any, k) => {
+        if (prev === undefined) {
+            return undefined;
+        }
+        return prev[k];
+    }, resources ?? {});
+}
 
 export class BundleRemoteStateModel extends BaseModelWithStateCache<BundleRemoteState> {
     public target: string | undefined;

packages/databricks-vscode/src/bundle/run/BundleRunStatus.ts

@@ -2,14 +2,22 @@ import {EventEmitter} from "vscode";
 import {RunState} from "./types";
 import {ResourceKey} from "../types";
 import {BundleRemoteState} from "../models/BundleRemoteStateModel";
-
+import {Time, TimeUnits} from "@databricks/databricks-sdk";
 export abstract class BundleRunStatus {
     abstract readonly type: ResourceKey<BundleRemoteState>;
     protected readonly onDidChangeEmitter = new EventEmitter<void>();
     readonly onDidChange = this.onDidChangeEmitter.event;
     runId: string | undefined;
     data: any;
 
+    constructor() {
+        //  Timeout in 60 seconds if we don't have a runId till then.
+        setTimeout(() => {
+            if (this.runState === "unknown") {
+                this.runState = "timeout";
+            }
+        }, new Time(60, TimeUnits.seconds).toMillSeconds().value);
+    }
     protected _runState: RunState = "unknown";
     public get runState(): RunState {
         return this._runState;

packages/databricks-vscode/src/bundle/run/BundleRunStatusManager.ts

@@ -1,14 +1,12 @@
 import {Disposable, EventEmitter} from "vscode";
-import {
-    BundleRemoteState,
-    BundleRemoteStateModel,
-} from "../models/BundleRemoteStateModel";
+import {BundleRemoteState, getResource} from "../models/BundleRemoteStateModel";
 import {BundleRunTerminalManager} from "./BundleRunTerminalManager";
 import {JobRunStatus} from "./JobRunStatus";
 import {AuthProvider} from "../../configuration/auth/AuthProvider";
 import {BundleRunStatus} from "./BundleRunStatus";
 import {PipelineRunStatus} from "./PipelineRunStatus";
 import {Resource, ResourceKey} from "../types";
+import {ConfigModel} from "../../configuration/models/ConfigModel";
 /**
  * This class monitors the cli bundle run output and record ids for runs. It also polls for status of the these runs.
  */
@@ -20,9 +18,16 @@ export class BundleRunStatusManager implements Disposable {
     readonly onDidChange = this.onDidChangeEmitter.event;
 
     constructor(
-        private readonly bundleRemoteStateModel: BundleRemoteStateModel,
+        private readonly configModel: ConfigModel,
         private readonly bundleRunTerminalManager: BundleRunTerminalManager
-    ) {}
+    ) {
+        this.disposables.push(
+            this.configModel.onDidChangeTarget(() => {
+                this.runStatuses.clear();
+                this.onDidChangeEmitter.fire();
+            })
+        );
+    }
 
     getRunStatusMonitor(
         resourceKey: string,
@@ -54,10 +59,12 @@ export class BundleRunStatusManager implements Disposable {
         resourceKey: string,
         resourceType: ResourceKey<BundleRemoteState>
     ) {
-        const target = this.bundleRemoteStateModel.target;
-        const authProvider = this.bundleRemoteStateModel.authProvider;
-        const resource =
-            await this.bundleRemoteStateModel.getResource(resourceKey);
+        const target = this.configModel.target;
+        const authProvider = this.configModel.authProvider;
+        const resource = getResource(
+            resourceKey,
+            (await this.configModel.get("remoteStateConfig"))?.resources
+        );
 
         if (target === undefined) {
             throw new Error(`Cannot run ${resourceKey}, Target is undefined`);
@@ -85,6 +92,7 @@ export class BundleRunStatusManager implements Disposable {
                 this.onDidChangeEmitter.fire();
             })
         );
+        this.onDidChangeEmitter.fire();
         await this.bundleRunTerminalManager.run(resourceKey, (data) => {
             remoteRunStatus.parseId(data);
         });

packages/databricks-vscode/src/bundle/run/BundleRunTerminalManager.ts

@@ -99,7 +99,7 @@ export class BundleRunTerminalManager implements Disposable {
                 }, disposables);
                 window.onDidCloseTerminal((e) => {
                     // Resolve when the process is closed by human action
-                    e.name === terminal.terminal.name && reject(resolve());
+                    e.name === terminal.terminal.name && resolve();
                 }, disposables);
             });
         } finally {
@@ -135,6 +135,8 @@ export class BundleRunTerminalManager implements Disposable {
         }
 
         const terminalName = this.getTerminalName(target, resourceKey);
+        window.terminals.find((i) => i.name === terminalName)?.show();
+
         this.cancellationTokenSources.get(terminalName)?.cancel();
         this.cancellationTokenSources.get(terminalName)?.dispose();
         this.cancellationTokenSources.delete(terminalName);

packages/databricks-vscode/src/bundle/run/CustomOutputTerminal.ts

@@ -62,6 +62,20 @@ export class CustomOutputTerminal implements Pseudoterminal {
         this.process.stderr.on("data", handleOutput);
 
         this.process.on("close", (exitCode) => {
+            if (exitCode === 0) {
+                this.writeEmitter.fire(
+                    "\x1b[32mProcess completed successfully\x1b[0m\r\n"
+                );
+            }
+
+            if (exitCode !== 0) {
+                this.writeEmitter.fire(
+                    "\x1b[31mProcess exited with code " +
+                        exitCode +
+                        "\x1b[0m\r\n"
+                );
+            }
+
             this.onDidCloseProcessEmitter.fire(exitCode);
             this._process = undefined;
         });

packages/databricks-vscode/src/bundle/run/JobRunStatus.ts

@@ -66,14 +66,14 @@ export class JobRunStatus extends BundleRunStatus {
 
     async cancel(): Promise<void> {
         if (this.runId === undefined || this.runState !== "running") {
-            this.runState = "completed";
+            this.runState = "cancelled";
             return;
         }
 
         const client = this.authProvider.getWorkspaceClient();
         await (
             await client.jobs.cancelRun({run_id: parseInt(this.runId)})
         ).wait();
-        this.runState = "completed";
+        this.runState = "cancelled";
     }
 }

packages/databricks-vscode/src/bundle/run/PipelineRunStatus.ts

@@ -85,9 +85,17 @@ export class PipelineRunStatus extends BundleRunStatus {
         this.runState = "completed";
     }
 
+    private markCancelled() {
+        if (this.interval !== undefined) {
+            clearInterval(this.interval);
+            this.interval = undefined;
+        }
+        this.runState = "cancelled";
+    }
+
     async cancel() {
         if (this.runState !== "running" || this.runId === undefined) {
-            this.markCompleted();
+            this.markCancelled();
             return;
         }
 
@@ -109,6 +117,6 @@ export class PipelineRunStatus extends BundleRunStatus {
             pipeline_id: this.pipelineId,
             update_id: this.runId,
         });
-        this.markCompleted();
+        this.markCancelled();
     }
 }

packages/databricks-vscode/src/bundle/run/types.ts

@@ -1 +1,7 @@
-export type RunState = "running" | "completed" | "unknown" | "error";
+export type RunState =
+    | "running"
+    | "completed"
+    | "unknown"
+    | "error"
+    | "timeout"
+    | "cancelled";

packages/databricks-vscode/src/extension.ts

@@ -534,7 +534,7 @@ export async function activate(
         bundleRemoteStateModel
     );
     const bundleRunStatusManager = new BundleRunStatusManager(
-        bundleRemoteStateModel,
+        configModel,
         bundleRunTerminalManager
     );
     const bundleResourceExplorerTreeDataProvider =

packages/databricks-vscode/src/ui/bundle-resource-explorer/JobRunStatusTreeNode.ts

@@ -68,6 +68,17 @@ export class JobRunStatusTreeNode implements BundleResourceExplorerTreeNode {
     }
 
     getTreeItem(): BundleResourceExplorerTreeItem {
+        const runMonitorRunStateTreeItem =
+            RunStateUtils.getTreeItemFromRunMonitorStatus(
+                this.type,
+                this.url,
+                this.runMonitor
+            );
+
+        if (runMonitorRunStateTreeItem) {
+            return runMonitorRunStateTreeItem;
+        }
+
         if (this.runDetails === undefined) {
             return {
                 label: "Run Status",
@@ -76,7 +87,7 @@ export class JobRunStatusTreeNode implements BundleResourceExplorerTreeNode {
                 contextValue: ContextUtils.getContextString({
                     nodeType: this.type,
                 }),
-                collapsibleState: TreeItemCollapsibleState.Collapsed,
+                collapsibleState: TreeItemCollapsibleState.None,
             };
         }