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? -->

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,
             };
         }
 

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

@@ -34,21 +34,21 @@ export class JobTreeNode implements BundleResourceExplorerTreeNode {
         public parent?: BundleResourceExplorerTreeNode
     ) {}
 
-    isRunning(resourceKey: string) {
-        const runner = this.bundleRunStatusManager.runStatuses.get(resourceKey);
-        return runner?.runState === "running";
+    get isRunning() {
+        const runner = this.bundleRunStatusManager.runStatuses.get(
+            this.resourceKey
+        );
+        return runner?.runState === "running" || runner?.runState === "unknown";
     }
 
     getTreeItem(): BundleResourceExplorerTreeItem {
-        const isRunning = this.isRunning(this.resourceKey);
-
         return {
             label: this.data.name,
             contextValue: ContextUtils.getContextString({
                 resourceType: this.type,
-                running: isRunning,
+                running: this.isRunning,
                 hasUrl: this.url !== undefined,
-                cancellable: isRunning,
+                cancellable: this.isRunning,
                 nodeType: this.type,
                 modifiedStatus: this.data.modified_status,
             }),
@@ -57,7 +57,7 @@ export class JobTreeNode implements BundleResourceExplorerTreeNode {
                 this.data.modified_status
             ),
             collapsibleState: DecorationUtils.getCollapsibleState(
-                isRunning,
+                this.isRunning,
                 this.data.modified_status
             ),
         };
@@ -73,7 +73,7 @@ export class JobTreeNode implements BundleResourceExplorerTreeNode {
             this.resourceKey
         ) as JobRunStatus | undefined;
 
-        if (runMonitor?.runId !== undefined) {
+        if (runMonitor) {
             children.push(new JobRunStatusTreeNode(this, runMonitor));
         }
         children.push(

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

@@ -108,6 +108,17 @@ export class PipelineRunStatusTreeNode
     }
 
     getTreeItem(): BundleResourceExplorerTreeItem {
+        const runMonitorRunStateTreeItem =
+            RunStateUtils.getTreeItemFromRunMonitorStatus(
+                this.type,
+                this.url,
+                this.runMonitor
+            );
+
+        if (runMonitorRunStateTreeItem) {
+            return runMonitorRunStateTreeItem;
+        }
+
         if (this.update === undefined) {
             return {
                 label: "Run Status",
@@ -116,7 +127,7 @@ export class PipelineRunStatusTreeNode
                 contextValue: ContextUtils.getContextString({
                     nodeType: this.type,
                 }),
-                collapsibleState: TreeItemCollapsibleState.Collapsed,
+                collapsibleState: TreeItemCollapsibleState.None,
             };
         }
 

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

@@ -33,7 +33,7 @@ export class PipelineTreeNode implements BundleResourceExplorerTreeNode {
 
     isRunning(resourceKey: string) {
         const runner = this.bundleRunStatusManager.runStatuses.get(resourceKey);
-        return runner?.runState === "running";
+        return runner?.runState === "running" || runner?.runState === "unknown";
     }
 
     getTreeItem(): BundleResourceExplorerTreeItem {
@@ -65,7 +65,7 @@ export class PipelineTreeNode implements BundleResourceExplorerTreeNode {
         const runMonitor = this.bundleRunStatusManager.runStatuses.get(
             this.resourceKey
         ) as PipelineRunStatus | undefined;
-        if (runMonitor?.data?.update?.update_id !== undefined) {
+        if (runMonitor) {
             children.push(
                 new PipelineRunStatusTreeNode(
                     this.connectionManager,