From d85c73e34ab44579e804fd5c3d2a41f95fcb6404 Mon Sep 17 00:00:00 2001
From: kartikgupta-db <kartik.gupta@databricks.com>
Date: Tue, 26 Mar 2024 10:34:39 +0100
Subject: [PATCH] Add currently selected python environment to path for wheel
 builds

---
 .../src/bundle/models/BundleRemoteStateModel.ts   |  3 +++
 packages/databricks-vscode/src/cli/CliWrapper.ts  | 15 ++++++++++++++-
 packages/databricks-vscode/src/extension.ts       |  1 +
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts b/packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts
index 9c5572ca2..b26645d1d 100644
--- a/packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts
+++ b/packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts
@@ -9,6 +9,7 @@ import lodash from "lodash";
 import {WorkspaceConfigs} from "../../vscode-objs/WorkspaceConfigs";
 import {logging} from "@databricks/databricks-sdk";
 import {Loggers} from "../../logger";
+import {MsPythonExtensionWrapper} from "../../language/MsPythonExtensionWrapper";
 
 /* eslint-disable @typescript-eslint/naming-convention */
 export type BundleResourceModifiedStatus = "created" | "deleted" | "updated";
@@ -49,6 +50,7 @@ export class BundleRemoteStateModel extends BaseModelWithStateCache<BundleRemote
     constructor(
         private readonly cli: CliWrapper,
         private readonly workspaceFolder: Uri,
+        private readonly pythonExtension: MsPythonExtensionWrapper,
         private readonly workspaceConfigs: WorkspaceConfigs
     ) {
         super();
@@ -72,6 +74,7 @@ export class BundleRemoteStateModel extends BaseModelWithStateCache<BundleRemote
             this.target,
             this.authProvider,
             this.workspaceFolder,
+            this.pythonExtension,
             this.workspaceConfigs.databrickscfgLocation,
             this.logger
         );
diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts
index be75fb46a..903aa892e 100644
--- a/packages/databricks-vscode/src/cli/CliWrapper.ts
+++ b/packages/databricks-vscode/src/cli/CliWrapper.ts
@@ -16,6 +16,8 @@ import {EnvVarGenerators, FileUtils, UrlUtils} from "../utils";
 import {AuthProvider} from "../configuration/auth/AuthProvider";
 import {removeUndefinedKeys} from "../utils/envVarGenerators";
 import {quote} from "shell-quote";
+import {MsPythonExtensionWrapper} from "../language/MsPythonExtensionWrapper";
+import path from "path";
 
 const withLogContext = logging.withLogContext;
 const execFile = promisify(execFileCb);
@@ -392,6 +394,7 @@ export class CliWrapper {
         target: string,
         authProvider: AuthProvider,
         workspaceFolder: Uri,
+        pythonExtension: MsPythonExtensionWrapper,
         configfilePath?: string,
         logger?: logging.NamedLogger
     ) {
@@ -410,6 +413,14 @@ export class CliWrapper {
             bundleOpName: "deploy",
         });
 
+        // Add python executable to PATH
+        const executable = await pythonExtension.getPythonExecutable();
+        let shellPath = process.env.PATH;
+        if (executable) {
+            shellPath = `${path.dirname(executable)}${path.delimiter}${
+                process.env.PATH
+            }`;
+        }
         const p = spawn(cmd[0], cmd.slice(1), {
             cwd: workspaceFolder.fsPath,
             env: {
@@ -417,8 +428,10 @@ export class CliWrapper {
                 ...EnvVarGenerators.getProxyEnvVars(),
                 ...authProvider.toEnv(),
                 ...this.getLogginEnvVars(),
-                // eslint-disable-next-line @typescript-eslint/naming-convention
+                /* eslint-disable @typescript-eslint/naming-convention */
+                PATH: shellPath,
                 DATABRICKS_CLUSTER_ID: this.clusterId,
+                /* eslint-enable @typescript-eslint/naming-convention */
             },
             shell: true,
         });
diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts
index c47d2722c..7e1b8f63f 100644
--- a/packages/databricks-vscode/src/extension.ts
+++ b/packages/databricks-vscode/src/extension.ts
@@ -232,6 +232,7 @@ export async function activate(
     const bundleRemoteStateModel = new BundleRemoteStateModel(
         cli,
         workspaceUri,
+        pythonExtensionWrapper,
         workspaceConfigs
     );
     const configModel = new ConfigModel(