Add deployAndRunFromInput command

databricks · Nov 26, 2024 · ace04ef · ace04ef
1 parent 19b7522
commit ace04ef
Show file tree

Hide file tree

Showing 4 changed files with 91 additions and 11 deletions.
diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json
@@ -217,6 +217,13 @@
                 "enablement": "databricks.context.activated && databricks.context.bundle.isTargetSet && databricks.context.bundle.deploymentState == idle",
                 "category": "Databricks"
             },
+            {
+                "command": "databricks.bundle.deployAndRunFromInput",
+                "icon": "$(run)",
+                "title": "Deploy the bundle and run a resource",
+                "enablement": "databricks.context.activated && databricks.context.bundle.isTargetSet && databricks.context.bundle.deploymentState == idle",
+                "category": "Databricks"
+            },
             {
                 "command": "databricks.bundle.deployAndRunJob",
                 "icon": "$(run)",

diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts
@@ -683,6 +683,11 @@ export async function activate(
             bundleCommands.forceDestroyCommand,
             bundleCommands
         ),
+        telemetry.registerCommand(
+            "databricks.bundle.deployAndRunFromInput",
+            bundleCommands.deployAndRunFromInput,
+            bundleCommands
+        ),
         telemetry.registerCommand(
             "databricks.bundle.deployAndRunJob",
             bundleCommands.deployAndRun,

diff --git a/packages/databricks-vscode/src/telemetry/constants.ts b/packages/databricks-vscode/src/telemetry/constants.ts
@@ -30,7 +30,11 @@ export type ManualLoginSource =
     | "command"
     | "api";
 export type BundleRunResourceType = "pipelines" | "jobs";
-export type BundleRunType = "run" | "validate" | "partial-refresh";
+export type BundleRunType =
+    | "run"
+    | "validate"
+    | "partial-refresh"
+    | "manual-input";
 
 /** Documentation about all of the properties and metrics of the event. */
 type EventDescription<T> = {[K in keyof T]?: {comment?: string}};

diff --git a/packages/databricks-vscode/src/ui/bundle-resource-explorer/BundleCommands.ts b/packages/databricks-vscode/src/ui/bundle-resource-explorer/BundleCommands.ts
@@ -21,10 +21,9 @@ export const RUNNABLE_BUNDLE_RESOURCES = [
 ] satisfies BundleResourceExplorerTreeNode["type"][];
 
 type RunnableTreeNodes = PipelineTreeNode | JobTreeNode;
+type ResourceData = {type: string; resourceKey?: string};
 
-function isRunnable(
-    treeNode: BundleResourceExplorerTreeNode
-): treeNode is RunnableTreeNodes {
+function isRunnable(treeNode: ResourceData): treeNode is RunnableTreeNodes {
     return (RUNNABLE_BUNDLE_RESOURCES as string[]).includes(treeNode.type);
 }
 
@@ -178,34 +177,99 @@ export class BundleCommands implements Disposable {
 
     @onError({popup: {prefix: "Error running resource."}})
     async deployAndRun(
-        treeNode: BundleResourceExplorerTreeNode,
+        resourceData: ResourceData,
         additionalArgs: string[] = [],
         runType: BundleRunType = "run"
     ) {
-        if (!isRunnable(treeNode)) {
-            throw new Error(`Cannot run resource of type ${treeNode.type}`);
+        if (!isRunnable(resourceData)) {
+            throw new Error(`Cannot run resource of type ${resourceData.type}`);
         }
         const recordEvent = this.telemetry.start(Events.BUNDLE_RUN);
         try {
             // TODO: Don't deploy if there is no diff between local and remote state
             await this.deploy();
             const result = await this.bundleRunStatusManager.run(
-                treeNode.resourceKey,
-                treeNode.type,
+                resourceData.resourceKey,
+                resourceData.type,
                 additionalArgs
             );
             recordEvent({
                 success: true,
                 runType,
-                resourceType: treeNode.type,
+                resourceType: resourceData.type,
                 cancelled: result.cancelled,
             });
         } catch (e) {
-            recordEvent({success: false, resourceType: treeNode.type, runType});
+            recordEvent({
+                success: false,
+                resourceType: resourceData.type,
+                runType,
+            });
             throw e;
         }
     }
 
+    async deployAndRunFromInput(
+        options: {
+            resourceType?: string;
+            resourceKey?: string;
+            args?: string;
+        } = {}
+    ) {
+        let {resourceType, resourceKey} = options;
+        if (!resourceType || !resourceKey) {
+            const remoteState = await this.configModel.get("remoteStateConfig");
+            const pipelines = remoteState?.resources?.pipelines ?? {};
+            const pipelineItems = Object.keys(pipelines).map((key) => ({
+                label: key,
+                description: "Pipeline",
+                type: "pipelines",
+                key: `pipelines.${key}`,
+            }));
+            const workflows = remoteState?.resources?.jobs ?? {};
+            const workflowItems = Object.keys(workflows).map((key) => ({
+                label: key,
+                description: "Workflow",
+                type: "jobs",
+            }));
+            if (pipelineItems.length === 0 && workflowItems.length === 0) {
+                window.showErrorMessage(
+                    "No pipelines or workflows found in the bundle."
+                );
+                return;
+            }
+            const pick = await window.showQuickPick(
+                [...pipelineItems, ...workflowItems],
+                {
+                    placeHolder: "Select a pipeline or workflow to run",
+                }
+            );
+            if (!pick) {
+                return;
+            }
+            resourceType = pick.type;
+            resourceKey = pick.label;
+        }
+        const resourceData = {
+            type: resourceType,
+            resourceKey: `${resourceType}.${resourceKey}`,
+        };
+        const existingRun = this.bundleRunStatusManager.runStatuses.get(
+            resourceData.resourceKey
+        );
+        if (
+            existingRun?.runState === "running" ||
+            existingRun?.runState === "unknown"
+        ) {
+            window.showErrorMessage(
+                `A run for the "${resourceData.resourceKey}" is already in progress.`
+            );
+        } else {
+            const args = options.args?.split(" ").filter(Boolean);
+            return this.deployAndRun(resourceData, args, "manual-input");
+        }
+    }
+
     async deployAndValidate(treeNode: BundleResourceExplorerTreeNode) {
         return this.deployAndRun(treeNode, ["--validate-only"], "validate");
     }