feat(vscode): .env file into python env

Author: benfdking

vscode/extension/src/utilities/common/python.ts

@@ -5,6 +5,8 @@ import { commands, Disposable, Event, EventEmitter, Uri } from 'vscode'
 import { traceError, traceLog } from './log'
 import { PythonExtension, ResolvedEnvironment } from '@vscode/python-extension'
 import path from 'path'
+import { err, ok, Result } from '@bus/result'
+import * as vscode from 'vscode'
 
 export interface IInterpreterDetails {
   path?: string[]
@@ -15,10 +17,12 @@ export interface IInterpreterDetails {
 
 const onDidChangePythonInterpreterEvent =
   new EventEmitter<IInterpreterDetails>()
+
 export const onDidChangePythonInterpreter: Event<IInterpreterDetails> =
   onDidChangePythonInterpreterEvent.event
 
 let _api: PythonExtension | undefined
+
 async function getPythonExtensionAPI(): Promise<PythonExtension | undefined> {
   if (_api) {
     return _api
@@ -118,3 +122,34 @@ export function checkVersion(
   traceError('Supported versions are 3.8 and above.')
   return false
 }
+
+/**
+ * getPythonEnvVariables returns the environment variables for the current python interpreter.
+ *
+ * @returns The environment variables for the current python interpreter.
+ */
+export async function getPythonEnvVariables(): Promise<
+  Result<Record<string, string>, string>
+> {
+  const api = await getPythonExtensionAPI()
+  if (!api) {
+    return err('Python extension API not found')
+  }
+
+  const workspaces = vscode.workspace.workspaceFolders
+  if (!workspaces) {
+    return ok({})
+  }
+  const out: Record<string, string> = {}
+  for (const workspace of workspaces) {
+    const envVariables = api.environments.getEnvironmentVariables(workspace.uri)
+    if (envVariables) {
+      for (const [key, value] of Object.entries(envVariables)) {
+        if (value) {
+          out[key] = value
+        }
+      }
+    }
+  }
+  return ok(out)
+}

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

@@ -1,6 +1,6 @@
 import path from 'path'
 import { traceInfo, traceLog, traceVerbose } from '../common/log'
-import { getInterpreterDetails } from '../common/python'
+import { getInterpreterDetails, getPythonEnvVariables } from '../common/python'
 import { Result, err, isErr, ok } from '@bus/result'
 import { getProjectRoot } from '../common/utilities'
 import { isPythonModuleInstalled } from '../python'
@@ -230,6 +230,13 @@ export const sqlmeshExec = async (): Promise<
       message: resolvedPath.error,
     })
   }
+  const envVariables = await getPythonEnvVariables()
+  if (isErr(envVariables)) {
+    return err({
+      type: 'generic',
+      message: envVariables.error,
+    })
+  }
   const workspacePath = resolvedPath.value
   const interpreterDetails = await getInterpreterDetails()
   traceLog(`Interpreter details: ${JSON.stringify(interpreterDetails)}`)
@@ -271,6 +278,8 @@ export const sqlmeshExec = async (): Promise<
           PYTHONPATH: interpreterDetails.path?.[0],
           VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!),
           PATH: interpreterDetails.binPath!,
+          ...process.env,
+          ...envVariables.value,
         },
         args: [],
       })
@@ -284,6 +293,8 @@ export const sqlmeshExec = async (): Promise<
         PYTHONPATH: interpreterDetails.path?.[0],
         VIRTUAL_ENV: path.dirname(path.dirname(interpreterDetails.binPath!)), // binPath now points to bin dir
         PATH: interpreterDetails.binPath!,
+        ...process.env,
+        ...envVariables.value,
       },
       args: [],
     })
@@ -297,7 +308,10 @@ export const sqlmeshExec = async (): Promise<
     return ok({
       bin: sqlmesh,
       workspacePath,
-      env: {},
+      env: {
+        ...process.env,
+        ...envVariables.value,
+      },
       args: [],
     })
   }
@@ -353,6 +367,13 @@ export const sqlmeshLspExec = async (): Promise<
 > => {
   const sqlmeshLSP = IS_WINDOWS ? 'sqlmesh_lsp.exe' : 'sqlmesh_lsp'
   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({
@@ -408,6 +429,8 @@ export const sqlmeshLspExec = async (): Promise<
             PYTHONPATH: interpreterDetails.path?.[0],
             VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!),
             PATH: interpreterDetails.binPath!,
+            ...process.env,
+            ...envVariables.value,
           },
           args: ['sqlmesh_lsp'],
         })
@@ -431,6 +454,8 @@ export const sqlmeshLspExec = async (): Promise<
         PYTHONPATH: interpreterDetails.path?.[0],
         VIRTUAL_ENV: path.dirname(path.dirname(interpreterDetails.binPath!)), // binPath now points to bin dir
         PATH: interpreterDetails.binPath!, // binPath already points to the bin directory
+        ...process.env, 
+        ...envVariables.value,
       },
       args: [],
     })
@@ -444,7 +469,10 @@ export const sqlmeshLspExec = async (): Promise<
     return ok({
       bin: sqlmeshLSP,
       workspacePath,
-      env: {},
+      env: {
+        ...process.env,
+        ...envVariables.value,
+      },
       args: [],
     })
   }

vscode/extension/tests/python_env.spec.ts

@@ -0,0 +1,83 @@
+import { test } from '@playwright/test'
+import fs from 'fs-extra'
+import {
+  createVirtualEnvironment,
+  openLineageView,
+  pipInstall,
+  REPO_ROOT,
+  startVSCode,
+  SUSHI_SOURCE_PATH,
+} from './utils'
+import os from 'os'
+import path from 'path'
+
+function writeEnvironmentConfig(sushi_path: string) {
+  const config_path = path.join(sushi_path, 'config.py')
+  const original_config = fs.readFileSync(config_path, 'utf8')
+
+  const new_config =
+    `
+import os
+
+test_var = os.getenv("TEST_VAR")
+if test_var is None or test_var == "":
+    raise Exception("TEST_VAR is not set")
+` + original_config
+
+  fs.writeFileSync(config_path, new_config)
+}
+
+async function setupEnvironment(): Promise<string> {
+  const tempDir = await fs.mkdtemp(
+    path.join(os.tmpdir(), 'vscode-test-tcloud-'),
+  )
+  await fs.copy(SUSHI_SOURCE_PATH, tempDir)
+  const pythonEnvDir = path.join(tempDir, '.venv')
+  const pythonDetails = await createVirtualEnvironment(pythonEnvDir)
+  const custom_materializations = path.join(
+    REPO_ROOT,
+    'examples',
+    'custom_materializations',
+  )
+  const sqlmeshWithExtras = `${REPO_ROOT}[bigquery,lsp]`
+  await pipInstall(pythonDetails, [sqlmeshWithExtras, custom_materializations])
+
+  const settings = {
+    'python.defaultInterpreterPath': pythonDetails.pythonPath,
+    'sqlmesh.environmentPath': pythonEnvDir,
+  }
+  await fs.ensureDir(path.join(tempDir, '.vscode'))
+  await fs.writeJson(path.join(tempDir, '.vscode', 'settings.json'), settings, {
+    spaces: 2,
+  })
+
+  return tempDir
+}
+
+test.describe('python environment variable injection on sqlmesh_lsp', () => {
+  test('normal setup - error ', async () => {
+    const tempDir = await setupEnvironment()
+    writeEnvironmentConfig(tempDir)
+    const { window, close } = await startVSCode(tempDir)
+    try {
+      await openLineageView(window)
+      await window.waitForSelector('text=Error creating context')
+    } finally {
+      await close()
+    }
+  })
+
+  test('normal setup - set', async () => {
+    const tempDir = await setupEnvironment()
+    writeEnvironmentConfig(tempDir)
+    const env_file = path.join(tempDir, '.env')
+    fs.writeFileSync(env_file, 'TEST_VAR=test_value')
+    const { window, close } = await startVSCode(tempDir)
+    try {
+      await openLineageView(window)
+      await window.waitForSelector('text=Loaded SQLMesh context')
+    } finally {
+      await close()
+    }
+  })
+})