Add currently selected python environment to path for wheel builds

databricks · Mar 26, 2024 · d85c73e · d85c73e
1 parent 2cff68e
commit d85c73e
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 1 deletion.
diff --git 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
@@ -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,15 +413,25 @@ 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: {
                 ...EnvVarGenerators.getEnvVarsForCli(configfilePath),
                 ...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
@@ -232,6 +232,7 @@ export async function activate(
     const bundleRemoteStateModel = new BundleRemoteStateModel(
         cli,
         workspaceUri,
+        pythonExtensionWrapper,
         workspaceConfigs
     );
     const configModel = new ConfigModel(