Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions vscode/extension/src/utilities/common/python.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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[]
Expand All @@ -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
Expand Down Expand Up @@ -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)
}
34 changes: 31 additions & 3 deletions vscode/extension/src/utilities/sqlmesh/sqlmesh.ts
Original file line number Diff line number Diff line change
@@ -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'
Expand Down Expand Up @@ -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)}`)
Expand Down Expand Up @@ -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!,
Expand All @@ -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!,
Expand All @@ -297,7 +308,10 @@ export const sqlmeshExec = async (): Promise<
return ok({
bin: sqlmesh,
workspacePath,
env: {},
env: {
Comment thread
benfdking marked this conversation as resolved.
...process.env,
...envVariables.value,
},
args: [],
})
}
Expand Down Expand Up @@ -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({
Expand Down Expand Up @@ -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'],
})
Expand All @@ -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: [],
})
Expand All @@ -444,7 +469,10 @@ export const sqlmeshLspExec = async (): Promise<
return ok({
bin: sqlmeshLSP,
workspacePath,
env: {},
env: {
...process.env,
...envVariables.value,
},
args: [],
})
}
Expand Down
139 changes: 139 additions & 0 deletions vscode/extension/tests/python_env.spec.ts
Original file line number Diff line number Diff line change
@@ -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()
}
})
})
33 changes: 1 addition & 32 deletions vscode/extension/tests/tcloud.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -33,38 +34,6 @@ async function setupPythonEnvironment(envDir: string): Promise<string> {
return pythonDetails.pythonPath
}

/**
* Helper function to set up a pre-authenticated tcloud state
*/
async function setupAuthenticatedState(tempDir: string): Promise<void> {
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<void> {
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
Expand Down
34 changes: 34 additions & 0 deletions vscode/extension/tests/tcloud_utils.ts
Original file line number Diff line number Diff line change
@@ -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<void> {
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<void> {
const versionStateFile = path.join(tempDir, '.tcloud_version_state.json')
await fs.writeJson(versionStateFile, { version })
}