Add DATABRICKS_CLI_UPSTREAM env var to CLI calls (#1112)

Also add DATABRICKS_CLI_UPSTREAM_VERSION

Author: ilia-db

packages/databricks-vscode/src/cli/CliWrapper.test.ts

@@ -18,10 +18,16 @@ import crypto from "node:crypto";
 import {Context} from "@databricks/databricks-sdk/dist/context";
 import {logging} from "@databricks/databricks-sdk";
 import {LoggerManager} from "../logger";
+import {ProfileAuthProvider} from "../configuration/auth/AuthProvider";
+import {isMatch} from "lodash";
+import {removeUndefinedKeys} from "../utils/envVarGenerators";
 
 const execFile = promisify(execFileCb);
 const cliPath = path.join(__dirname, "../../bin/databricks");
 
+// eslint-disable-next-line @typescript-eslint/no-var-requires
+const extensionVersion = require("../../package.json").version;
+
 function getTempLogFilePath() {
     return path.join(
         os.tmpdir(),
@@ -204,4 +210,50 @@ nothing = true
             assert.equal(profiles.length, 0);
         });
     });
+
+    it("should set required env vars to the bundle run CLI calls", async () => {
+        const logFilePath = getTempLogFilePath();
+        const cli = createCliWrapper(logFilePath);
+        const authProvider = new ProfileAuthProvider(
+            new URL("https://test.com"),
+            "PROFILE",
+            true
+        );
+        const workspaceFolder = Uri.file("/test/123");
+        const runCmd = cli.getBundleRunCommand(
+            "dev",
+            authProvider,
+            "resource-key",
+            workspaceFolder
+        );
+        const expected = {
+            args: ["bundle", "run", "--target", "dev", "resource-key"],
+            cmd: cli.cliPath,
+            options: {
+                cwd: workspaceFolder.fsPath,
+                env: removeUndefinedKeys({
+                    /* eslint-disable @typescript-eslint/naming-convention */
+                    DATABRICKS_CLI_UPSTREAM: "databricks-vscode",
+                    DATABRICKS_CLI_UPSTREAM_VERSION: extensionVersion,
+                    DATABRICKS_CONFIG_PROFILE: "PROFILE",
+                    DATABRICKS_HOST: "https://test.com/",
+                    DATABRICKS_LOG_FILE: logFilePath,
+                    DATABRICKS_LOG_FORMAT: "json",
+                    DATABRICKS_LOG_LEVEL: "debug",
+                    DATABRICKS_OUTPUT_FORMAT: "json",
+                    HOME: process.env.HOME,
+                    PATH: process.env.PATH,
+                    /* eslint-enable @typescript-eslint/naming-convention */
+                }),
+                shell: true,
+            },
+        };
+        try {
+            assert.ok(isMatch(runCmd, expected));
+        } catch (e) {
+            // Run this in the "catch" case to show better error messages
+            assert.deepStrictEqual(runCmd, expected);
+            throw e;
+        }
+    });
 });

packages/databricks-vscode/src/cli/CliWrapper.ts

@@ -87,6 +87,7 @@ async function waitForProcess(
  */
 export class CliWrapper {
     private clusterId?: string;
+
     constructor(
         private extensionContext: ExtensionContext,
         private loggerManager: LoggerManager,

packages/databricks-vscode/src/utils/envVarGenerators.ts

@@ -5,6 +5,9 @@ import {logging, Headers} from "@databricks/databricks-sdk";
 import {ConnectionManager} from "../configuration/ConnectionManager";
 import {ConfigModel} from "../configuration/models/ConfigModel";
 
+// eslint-disable-next-line @typescript-eslint/no-var-requires
+const extensionVersion = require("../../package.json").version;
+
 //Get env variables from user's .env file
 export async function getUserEnvVars(userEnvPath: Uri) {
     try {
@@ -137,6 +140,8 @@ export function getEnvVarsForCli(configfilePath?: string) {
         DATABRICKS_CONFIG_FILE:
             configfilePath ?? process.env.DATABRICKS_CONFIG_FILE,
         DATABRICKS_OUTPUT_FORMAT: "json",
+        DATABRICKS_CLI_UPSTREAM: "databricks-vscode",
+        DATABRICKS_CLI_UPSTREAM_VERSION: extensionVersion,
     };
     /* eslint-enable @typescript-eslint/naming-convention */
 }