Added functionality to deploy the current bundle (#1010)

kartikgupta-db · web-flow · commit 936fc497bc91 · 2024-01-28T20:46:30.000Z
## Changes
* Allow deploying the current bundle.
* Users can do this with a command or by clicking the "deploy" button in
the dabs resource explorer view.
* The outputs of the deployment are shown in an output channel. 

## Tests
&lt;!-- How is this tested? --&gt;
diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json
@@ -231,6 +231,20 @@
                 "enablement": "databricks.context.activated",
                 "category": "Databricks"
             },
+            {
+                "command": "databricks.bundle.refreshRemoteState",
+                "icon": "$(refresh)",
+                "title": "Refresh remote state",
+                "enablement": "databricks.context.activated && databricks.context.bundle.isTargetSet",
+                "category": "Databricks"
+            },
+            {
+                "command": "databricks.bundle.deploy",
+                "icon": "$(cloud-upload)",
+                "title": "Deploy bundle",
+                "enablement": "databricks.context.activated && databricks.context.bundle.isTargetSet",
+                "category": "Databricks"
+            },
             {
                 "command": "databricks.bundle.initNewProject",
                 "title": "Initialize new project",
@@ -322,6 +336,16 @@
                     "command": "databricks.wsfs.createFolder",
                     "when": "view == workspaceFsView",
                     "group": "navigation@1"
+                },
+                {
+                    "command": "databricks.bundle.refreshRemoteState",
+                    "when": "view == dabsResourceExplorerView",
+                    "group": "navigation@1"
+                },
+                {
+                    "command": "databricks.bundle.deploy",
+                    "when": "view == dabsResourceExplorerView",
+                    "group": "navigation@1"
                 }
             ],
             "databricks.cluster.filter": [
diff --git a/packages/databricks-vscode/src/bundle/BundleCommands.ts b/packages/databricks-vscode/src/bundle/BundleCommands.ts
@@ -0,0 +1,40 @@
+import {Disposable, window} from "vscode";
+import {BundleRemoteStateModel} from "./models/BundleRemoteStateModel";
+import {onError} from "../utils/onErrorDecorator";
+
+export class BundleCommands implements Disposable {
+    private disposables: Disposable[] = [];
+    private outputChannel = window.createOutputChannel(
+        "Databricks Asset Bundles"
+    );
+
+    constructor(private bundleRemoteStateModel: BundleRemoteStateModel) {}
+
+    @onError({popup: {prefix: "Error refreshing remote state."}})
+    async refreshRemoteState() {
+        await this.bundleRemoteStateModel.refresh();
+    }
+
+    @onError({popup: {prefix: "Error deploying the bundle."}})
+    async deploy() {
+        this.outputChannel.show(true);
+        this.outputChannel.appendLine("");
+
+        const writeToChannel = (data: string) => {
+            this.outputChannel.append(data);
+        };
+        await window.withProgress(
+            {location: {viewId: "dabsResourceExplorerView"}},
+            async () => {
+                await this.bundleRemoteStateModel.deploy(
+                    writeToChannel,
+                    writeToChannel
+                );
+            }
+        );
+    }
+
+    dispose() {
+        this.disposables.forEach((i) => i.dispose());
+    }
+}
diff --git a/packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts b/packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts
@@ -41,6 +41,34 @@ export class BundleRemoteStateModel extends BaseModelWithStateCache<BundleRemote
         super();
     }
 
+    public async refresh() {
+        return await this.stateCache.refresh();
+    }
+
+    @Mutex.synchronise("mutex")
+    public async deploy(
+        onStdOut?: (data: string) => void,
+        onStdErr?: (data: string) => void
+    ) {
+        if (this.target === undefined) {
+            throw new Error("Target is undefined");
+        }
+        if (this.authProvider === undefined) {
+            throw new Error("No authentication method is set");
+        }
+
+        await this.cli.bundleDeploy(
+            this.target,
+            this.authProvider,
+            this.workspaceFolder,
+            this.workspaceConfigs.databrickscfgLocation,
+            onStdOut,
+            onStdErr
+        );
+
+        await this.refresh();
+    }
+
     @withLogContext(Loggers.Extension)
     public init(@context ctx?: Context) {
         this.refreshInterval = setInterval(async () => {
diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts
@@ -1,4 +1,8 @@
-import {execFile as execFileCb} from "child_process";
+import {
+    ChildProcessWithoutNullStreams,
+    execFile as execFileCb,
+    spawn,
+} from "child_process";
 import {ExtensionContext, window, commands, Uri} from "vscode";
 import {SyncDestinationMapper} from "../sync/SyncDestination";
 import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs";
@@ -29,18 +33,58 @@ export interface ConfigEntry {
 
 export type SyncType = "full" | "incremental";
 
+async function waitForProcess(
+    p: ChildProcessWithoutNullStreams,
+    onStdOut?: (data: string) => void,
+    onStdError?: (data: string) => void
+) {
+    const output: string[] = [];
+    p.stdout.on("data", (data) => {
+        output.push(data.toString());
+        if (onStdOut) {
+            onStdOut(data.toString());
+        }
+    });
+
+    const stderr: string[] = [];
+    p.stderr.on("data", (data) => {
+        stderr.push(data.toString());
+        output.push(data.toString());
+        if (onStdError) {
+            onStdError(data.toString());
+        }
+    });
+
+    await new Promise((resolve, reject) => {
+        p.on("close", (code) => {
+            if (code === 0) {
+                resolve(output.join(""));
+            } else {
+                reject(stderr.join(""));
+            }
+        });
+        p.on("error", reject);
+    });
+
+    return output.join("");
+}
 /**
  * Entrypoint for all wrapped CLI commands
  *
  * Righ now this is a placeholder for a future implementation
  * of the databricks CLI
  */
 export class CliWrapper {
+    private clusterId?: string;
     constructor(
         private extensionContext: ExtensionContext,
         private logFilePath?: string
     ) {}
 
+    public setClusterId(clusterId?: string) {
+        this.clusterId = clusterId;
+    }
+
     get cliPath(): string {
         return this.extensionContext.asAbsolutePath("./bin/databricks");
     }
@@ -192,7 +236,7 @@ export class CliWrapper {
         workspaceFolder: Uri,
         configfilePath?: string
     ) {
-        const {stdout, stderr} = await execFile(
+        const {stdout} = await execFile(
             this.cliPath,
             ["bundle", "validate", "--target", target],
             {
@@ -201,14 +245,13 @@ export class CliWrapper {
                     ...EnvVarGenerators.getEnvVarsForCli(configfilePath),
                     ...EnvVarGenerators.getProxyEnvVars(),
                     ...authProvider.toEnv(),
+                    // eslint-disable-next-line @typescript-eslint/naming-convention
+                    DATABRICKS_CLUSTER_ID: this.clusterId,
                 },
                 shell: true,
             }
         );
 
-        if (stderr !== "") {
-            throw new Error(stderr);
-        }
         return stdout;
     }
 
@@ -227,6 +270,8 @@ export class CliWrapper {
                     ...EnvVarGenerators.getEnvVarsForCli(configfilePath),
                     ...EnvVarGenerators.getProxyEnvVars(),
                     ...authProvider.toEnv(),
+                    // eslint-disable-next-line @typescript-eslint/naming-convention
+                    DATABRICKS_CLUSTER_ID: this.clusterId,
                 },
                 shell: true,
             }
@@ -237,4 +282,37 @@ export class CliWrapper {
         }
         return stdout;
     }
+    async bundleDeploy(
+        target: string,
+        authProvider: AuthProvider,
+        workspaceFolder: Uri,
+        configfilePath?: string,
+        onStdOut?: (data: string) => void,
+        onStdError?: (data: string) => void
+    ) {
+        if (onStdError) {
+            onStdError(`Deploying the bundle for target ${target}...\n\n`);
+            onStdError(`${this.cliPath} bundle deploy --target ${target}\n`);
+            if (this.clusterId) {
+                onStdError(`DATABRICKS_CLUSTER_ID=${this.clusterId}\n\n`);
+            }
+        }
+        const p = spawn(
+            this.cliPath,
+            ["bundle", "deploy", "--target", target],
+            {
+                cwd: workspaceFolder.fsPath,
+                env: {
+                    ...EnvVarGenerators.getEnvVarsForCli(configfilePath),
+                    ...EnvVarGenerators.getProxyEnvVars(),
+                    ...authProvider.toEnv(),
+                    // eslint-disable-next-line @typescript-eslint/naming-convention
+                    DATABRICKS_CLUSTER_ID: this.clusterId,
+                },
+                shell: true,
+            }
+        );
+
+        return await waitForProcess(p, onStdOut, onStdError);
+    }
 }
diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts
@@ -113,6 +113,8 @@ export class ConnectionManager implements Disposable {
                       }
                   )
                 : undefined;
+
+        this.cli.setClusterId(clusterId);
         this.onDidChangeClusterEmitter.fire(this.cluster);
     }
 
diff --git a/packages/databricks-vscode/src/configuration/models/ConfigModel.ts b/packages/databricks-vscode/src/configuration/models/ConfigModel.ts
@@ -208,6 +208,7 @@ export class ConfigModel implements Disposable {
                     this.bundlePreValidateModel.setTarget(target),
                     this.bundleValidateModel.setTarget(target),
                     this.overrideableConfigModel.setTarget(target),
+                    this.bundleRemoteStateModel.setTarget(target),
                 ]);
             });
             this.onDidChangeTargetEmitter.fire();
@@ -222,6 +223,7 @@ export class ConfigModel implements Disposable {
         this._authProvider = authProvider;
         await this.readStateMutex.synchronise(async () => {
             await this.bundleValidateModel.setAuthProvider(authProvider);
+            await this.bundleRemoteStateModel.setAuthProvider(authProvider);
         });
         this.onDidChangeAuthProviderEmitter.fire();
     }
diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts
@@ -58,6 +58,7 @@ import {ConfigModel} from "./configuration/models/ConfigModel";
 import {OverrideableConfigModel} from "./configuration/models/OverrideableConfigModel";
 import {BundlePreValidateModel} from "./bundle/models/BundlePreValidateModel";
 import {BundleRemoteStateModel} from "./bundle/models/BundleRemoteStateModel";
+import {BundleCommands} from "./bundle/BundleCommands";
 import {BundleProjectManager} from "./bundle/BundleProjectManager";
 
 const customWhenContext = new CustomWhenContext();
@@ -597,6 +598,21 @@ export async function activate(
             }
         })
     );
+
+    // Bundle
+    const bundleCommands = new BundleCommands(bundleRemoteStateModel);
+    context.subscriptions.push(
+        telemetry.registerCommand(
+            "databricks.bundle.refreshRemoteState",
+            bundleCommands.refreshRemoteState,
+            bundleCommands
+        ),
+        telemetry.registerCommand(
+            "databricks.bundle.deploy",
+            bundleCommands.deploy,
+            bundleCommands
+        )
+    );
     // generate a json schema for bundle root and load a custom provider into
     // redhat.vscode-yaml extension to validate bundle config files with this schema
     registerBundleAutocompleteProvider(
diff --git a/packages/databricks-vscode/src/utils/envVarGenerators.test.ts b/packages/databricks-vscode/src/utils/envVarGenerators.test.ts
@@ -58,7 +58,6 @@ describe(__filename, () => {
             DATABRICKS_AUTH_TYPE: "metadata-service",
             DATABRICKS_METADATA_SERVICE_URL:
                 "http://example.com/metadata-service",
-            DATABRICKS_CLUSTER_ID: mockClusterId,
         });
     });
 
diff --git a/packages/databricks-vscode/src/utils/envVarGenerators.ts b/packages/databricks-vscode/src/utils/envVarGenerators.ts
@@ -47,7 +47,6 @@ function getUserAgent(connectionManager: ConnectionManager) {
 }
 
 export function getAuthEnvVars(connectionManager: ConnectionManager) {
-    const cluster = connectionManager.cluster;
     const host = connectionManager.databricksWorkspace?.host.toString();
     if (!host || !connectionManager.metadataServiceUrl) {
         return;
@@ -58,18 +57,19 @@ export function getAuthEnvVars(connectionManager: ConnectionManager) {
         DATABRICKS_HOST: host,
         DATABRICKS_AUTH_TYPE: "metadata-service",
         DATABRICKS_METADATA_SERVICE_URL: connectionManager.metadataServiceUrl,
-        DATABRICKS_CLUSTER_ID: cluster?.id,
     };
     /* eslint-enable @typescript-eslint/naming-convention */
 }
 
 export function getCommonDatabricksEnvVars(
     connectionManager: ConnectionManager
 ) {
+    const cluster = connectionManager.cluster;
     /* eslint-disable @typescript-eslint/naming-convention */
     return {
         ...(getAuthEnvVars(connectionManager) || {}),
         ...(getProxyEnvVars() || {}),
+        DATABRICKS_CLUSTER_ID: cluster?.id,
     };
     /* eslint-enable @typescript-eslint/naming-convention */
 }

Original file line number	Diff line number	Diff line change
`@@ -113,6 +113,8 @@ export class ConnectionManager implements Disposable {`
`113`	`113`	`}`
`114`	`114`	`)`
`115`	`115`	`: undefined;`
	`116`	`+`
	`117`	`+ this.cli.setClusterId(clusterId);`
`116`	`118`	`this.onDidChangeClusterEmitter.fire(this.cluster);`
`117`	`119`	`}`
`118`	`120`