fix(vscode): allow override of pip executable

Author: benfdking

vscode/extension/src/auth/auth.ts

@@ -79,10 +79,11 @@ export class AuthenticationProviderTobikoCloud
     }
     const tcloudBinPath = tcloudBin.value
     const result = await execAsync(
-      tcloudBinPath,
+      tcloudBinPath.bin,
       ['auth', 'vscode', 'status'],
       {
         cwd: workspacePath.uri.fsPath,
+        env: tcloudBinPath.env,
       },
     )
     if (result.exitCode !== 0) {
@@ -162,8 +163,9 @@ export class AuthenticationProviderTobikoCloud
       throw new Error('Failed to get tcloud bin')
     }
     const tcloudBinPath = tcloudBin.value
-    const result = await execAsync(tcloudBinPath, ['auth', 'logout'], {
+    const result = await execAsync(tcloudBinPath.bin, ['auth', 'logout'], {
       cwd: workspacePath.uri.fsPath,
+      env: tcloudBinPath.env,
     })
     if (result.exitCode !== 0) {
       throw new Error('Failed to logout from tcloud')
@@ -187,7 +189,7 @@ export class AuthenticationProviderTobikoCloud
     }
     const tcloudBinPath = tcloudBin.value
     const result = await execAsync(
-      tcloudBinPath,
+      tcloudBinPath.bin,
       ['auth', 'vscode', 'login-url'],
       {
         cwd: workspacePath.uri.fsPath,
@@ -214,11 +216,12 @@ export class AuthenticationProviderTobikoCloud
         1000 * 60 * 5,
       )
       const backgroundServerForLogin = execAsync(
-        tcloudBinPath,
+        tcloudBinPath.bin,
         ['auth', 'vscode', 'start-server', urlCode.verifier_code],
         {
           cwd: workspacePath.uri.fsPath,
           signal: ac.signal,
+          env: tcloudBinPath.env,
         },
       )
 
@@ -283,10 +286,11 @@ export class AuthenticationProviderTobikoCloud
     }
     const tcloudBinPath = tcloudBin.value
     const result = await execAsync(
-      tcloudBinPath,
+      tcloudBinPath.bin,
       ['auth', 'vscode', 'device'],
       {
         cwd: workspacePath.uri.fsPath,
+        env: tcloudBinPath.env,
       },
     )
     if (result.exitCode !== 0) {
@@ -305,11 +309,12 @@ export class AuthenticationProviderTobikoCloud
         1000 * 60 * 5,
       )
       const waiting = execAsync(
-        tcloudBinPath,
+        tcloudBinPath.bin,
         ['auth', 'vscode', 'poll_device', deviceCodeResponse.device_code],
         {
           cwd: workspacePath.uri.fsPath,
           signal: ac.signal,
+          env: tcloudBinPath.env,
         },
       )
 

vscode/extension/src/utilities/sqlmesh/sqlmesh.ts

@@ -55,7 +55,7 @@ export const isTcloudProject = async (): Promise<Result<boolean, string>> => {
  *
  * @returns The tcloud executable for the current Python environment.
  */
-export const getTcloudBin = async (): Promise<Result<string, ErrorType>> => {
+export const getTcloudBin = async (): Promise<Result<SqlmeshExecInfo, ErrorType>> => {
   const tcloud = IS_WINDOWS ? 'tcloud.exe' : 'tcloud'
   const interpreterDetails = await getInterpreterDetails()
   if (!interpreterDetails.path) {
@@ -68,7 +68,25 @@ export const getTcloudBin = async (): Promise<Result<string, ErrorType>> => {
   if (!fs.existsSync(binPath)) {
     return err({type: 'tcloud_bin_not_found'})
   }
-  return ok(binPath)
+  const envVariables = await getPythonEnvVariables()
+  if (isErr(envVariables)) {
+    return err({
+      type: 'generic',
+      message: envVariables.error,
+    })
+  }
+  return ok({
+    bin: binPath,
+    workspacePath: interpreterDetails.resource?.fsPath ?? '',
+    env: {
+      ...process.env,
+      ...envVariables.value,
+      PYTHONPATH: interpreterDetails.path[0],
+      VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!),
+      PATH: interpreterDetails.binPath!,
+    },
+    args: [],
+  })
 }
 
 const isSqlmeshInstalledSchema = z.object({
@@ -96,8 +114,9 @@ export const isSqlmeshEnterpriseInstalled = async (): Promise<
       message: resolvedPath.error,
     })
   }
-  const called = await execAsync(tcloudBin.value, ['is_sqlmesh_installed'], {
+  const called = await execAsync(tcloudBin.value.bin, ['is_sqlmesh_installed'], {
     cwd: resolvedPath.value,
+    env: tcloudBin.value.env,
   })
   if (called.exitCode !== 0) {
     return err({
@@ -128,16 +147,24 @@ export const installSqlmeshEnterprise = async (
     return tcloudBin
   }
   const projectRoot = await getProjectRoot()
+  const envVariables = await getPythonEnvVariables()
+  if (isErr(envVariables)) {
+    return err({
+      type: 'generic',
+      message: envVariables.error,
+    })
+  }
   const resolvedPath = resolveProjectPath(projectRoot)
   if (isErr(resolvedPath)) {
     return err({
       type: 'generic',
       message: resolvedPath.error,
     })
   }
-  const called = await execAsync(tcloudBin.value, ['install_sqlmesh'], {
+  const called = await execAsync(tcloudBin.value.bin, ['install_sqlmesh'], {
     signal: abortController.signal,
     cwd: resolvedPath.value,
+    env: tcloudBin.value.env,
   })
   if (called.exitCode !== 0) {
     return err({
@@ -272,16 +299,10 @@ export const sqlmeshExec = async (): Promise<
         return ensured
       }
       return ok({
-        bin: `${tcloudBin.value} sqlmesh`,
+        bin: tcloudBin.value.bin,
         workspacePath,
-        env: {
-          ...process.env,
-          ...envVariables.value,
-          PYTHONPATH: interpreterDetails.path?.[0],
-          VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!),
-          PATH: interpreterDetails.binPath!,
-        },
-        args: [],
+        env: tcloudBin.value.env,
+        args: ["sqlmesh"],
       })
     }
     const binPath = path.join(interpreterDetails.binPath!, sqlmesh)
@@ -423,15 +444,9 @@ export const sqlmeshLspExec = async (): Promise<
       // TODO: Remove this once we have a stable version of tcloud that supports sqlmesh_lsp.
       if (isSemVerGreaterThanOrEqual(tcloudBinVersion.value, [2, 10, 1])) {
         return ok ({
-          bin: tcloudBin.value,
+          bin: tcloudBin.value.bin,
           workspacePath,
-          env: {
-            PYTHONPATH: interpreterDetails.path?.[0],
-            VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!),
-            PATH: interpreterDetails.binPath!,
-            ...process.env,
-            ...envVariables.value,
-          },
+          env: tcloudBin.value.env,
           args: ['sqlmesh_lsp'],
         })
       }
@@ -503,7 +518,9 @@ async function getTcloudBinVersion(): Promise<Result<[number, number, number], E
   if (isErr(tcloudBin)) {
     return tcloudBin
   }
-  const called = await execAsync(tcloudBin.value, ['--version'])
+  const called = await execAsync(tcloudBin.value.bin, ['--version'], {
+    env: tcloudBin.value.env,
+  })
   if (called.exitCode !== 0) {
     return err({
       type: 'generic',