diff --git a/vscode/extension/src/utilities/common/python.ts b/vscode/extension/src/utilities/common/python.ts index bffae16a01..b30e2e91f7 100644 --- a/vscode/extension/src/utilities/common/python.ts +++ b/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() + export const onDidChangePythonInterpreter: Event = onDidChangePythonInterpreterEvent.event let _api: PythonExtension | undefined + async function getPythonExtensionAPI(): Promise { 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, 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 = {} + 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) +} diff --git a/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts b/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts index 37eb3c0b5c..a07336be21 100644 --- a/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts +++ b/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)}`) @@ -268,6 +275,8 @@ export const sqlmeshExec = async (): Promise< bin: `${tcloudBin.value} sqlmesh`, workspacePath, env: { + ...process.env, + ...envVariables.value, PYTHONPATH: interpreterDetails.path?.[0], VIRTUAL_ENV: path.dirname(interpreterDetails.binPath!), PATH: interpreterDetails.binPath!, @@ -281,6 +290,8 @@ export const sqlmeshExec = async (): Promise< bin: binPath, workspacePath, env: { + ...process.env, + ...envVariables.value, PYTHONPATH: interpreterDetails.path?.[0], VIRTUAL_ENV: path.dirname(path.dirname(interpreterDetails.binPath!)), // binPath now points to bin dir PATH: interpreterDetails.binPath!, @@ -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: [], }) } diff --git a/vscode/extension/tests/python_env.spec.ts b/vscode/extension/tests/python_env.spec.ts new file mode 100644 index 0000000000..9fae408696 --- /dev/null +++ b/vscode/extension/tests/python_env.spec.ts @@ -0,0 +1,139 @@ +import { test } from '@playwright/test' +import fs from 'fs-extra' +import { + createVirtualEnvironment, + openLineageView, + pipInstall, + PythonEnvironment, + REPO_ROOT, + startVSCode, + SUSHI_SOURCE_PATH, +} from './utils' +import os from 'os' +import path from 'path' +import { setTcloudVersion, setupAuthenticatedState } from './tcloud_utils' + +function writeEnvironmentConfig(sushiPath: string) { + const configPath = path.join(sushiPath, 'config.py') + const originalConfig = fs.readFileSync(configPath, 'utf8') + + const newConfig = + ` +import os + +test_var = os.getenv("TEST_VAR") +if test_var is None or test_var == "": + raise Exception("TEST_VAR is not set") +` + originalConfig + + fs.writeFileSync(configPath, newConfig) +} + +async function setupEnvironment(): Promise<[string, PythonEnvironment]> { + 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, pythonDetails] +} + +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() + } + }) +}) + +async function setupTcloudProject( + tempDir: string, + pythonDetails: PythonEnvironment, +) { + // Install the mock tcloud package + const mockTcloudPath = path.join(__dirname, 'tcloud') + await pipInstall(pythonDetails, [mockTcloudPath]) + + // Create a tcloud.yaml to mark this as a tcloud project + const tcloudConfig = { + url: 'https://mock.tobikodata.com', + org: 'test-org', + project: 'test-project', + } + await fs.writeFile( + path.join(tempDir, 'tcloud.yaml'), + `url: ${tcloudConfig.url}\norg: ${tcloudConfig.org}\nproject: ${tcloudConfig.project}\n`, + ) + // Write mock ".tcloud_auth_state.json" file + await setupAuthenticatedState(tempDir) + // Set tcloud version to 2.10.1 + await setTcloudVersion(tempDir, '2.10.1') +} + +test.describe('tcloud version', () => { + test('normal setup - error ', async () => { + const [tempDir, pythonDetails] = await setupEnvironment() + await setupTcloudProject(tempDir, pythonDetails) + 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, pythonDetails] = await setupEnvironment() + await setupTcloudProject(tempDir, pythonDetails) + 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() + } + }) +}) diff --git a/vscode/extension/tests/tcloud.spec.ts b/vscode/extension/tests/tcloud.spec.ts index fe3a61f4dc..f01dfc1a33 100644 --- a/vscode/extension/tests/tcloud.spec.ts +++ b/vscode/extension/tests/tcloud.spec.ts @@ -9,6 +9,7 @@ import { startVSCode, SUSHI_SOURCE_PATH, } from './utils' +import { setTcloudVersion, setupAuthenticatedState } from './tcloud_utils' /** * Helper function to create and set up a Python virtual environment @@ -33,38 +34,6 @@ async function setupPythonEnvironment(envDir: string): Promise { return pythonDetails.pythonPath } -/** - * Helper function to set up a pre-authenticated tcloud state - */ -async function setupAuthenticatedState(tempDir: string): Promise { - const authStateFile = path.join(tempDir, '.tcloud_auth_state.json') - const authState = { - is_logged_in: true, - id_token: { - iss: 'https://mock.tobikodata.com', - aud: 'mock-audience', - sub: 'user-123', - scope: 'openid email profile', - iat: Math.floor(Date.now() / 1000), - exp: Math.floor(Date.now() / 1000) + 3600, // Valid for 1 hour - email: 'test@example.com', - name: 'Test User', - }, - } - await fs.writeJson(authStateFile, authState) -} - -/** - * Helper function to set the tcloud version for testing - */ -async function setTcloudVersion( - tempDir: string, - version: string, -): Promise { - const versionStateFile = path.join(tempDir, '.tcloud_version_state.json') - await fs.writeJson(versionStateFile, { version }) -} - test.describe('Tcloud', () => { test('not signed in, shows sign in window', async ({}, testInfo) => { testInfo.setTimeout(120_000) // 2 minutes for venv creation and package installation diff --git a/vscode/extension/tests/tcloud_utils.ts b/vscode/extension/tests/tcloud_utils.ts new file mode 100644 index 0000000000..5334b69ef6 --- /dev/null +++ b/vscode/extension/tests/tcloud_utils.ts @@ -0,0 +1,34 @@ +import path from 'path' +import fs from 'fs-extra' + +/** + * Helper function to set up a pre-authenticated tcloud state + */ +export async function setupAuthenticatedState(tempDir: string): Promise { + const authStateFile = path.join(tempDir, '.tcloud_auth_state.json') + const authState = { + is_logged_in: true, + id_token: { + iss: 'https://mock.tobikodata.com', + aud: 'mock-audience', + sub: 'user-123', + scope: 'openid email profile', + iat: Math.floor(Date.now() / 1000), + exp: Math.floor(Date.now() / 1000) + 3600, // Valid for 1 hour + email: 'test@example.com', + name: 'Test User', + }, + } + await fs.writeJson(authStateFile, authState) +} + +/** + * Helper function to set the tcloud version for testing + */ +export async function setTcloudVersion( + tempDir: string, + version: string, +): Promise { + const versionStateFile = path.join(tempDir, '.tcloud_version_state.json') + await fs.writeJson(versionStateFile, { version }) +}