From 946044678acef04315777677eb1dec4e4118b9b1 Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Sun, 27 Sep 2026 00:40:11 +0000 Subject: [PATCH] fix(vscode): honor global telemetry setting Signed-off-by: Vamsi-klu --- vscode/extension/src/auth/auth.test.ts | 82 ++++++++ vscode/extension/src/auth/auth.ts | 1 + vscode/extension/src/extension.ts | 56 +++-- .../common/configurationChange.test.ts | 77 ++++++- .../utilities/common/configurationChange.ts | 43 ++++ .../src/utilities/sqlmesh/sqlmesh.test.ts | 194 ++++++++++++++++++ .../src/utilities/sqlmesh/sqlmesh.ts | 32 ++- vscode/extension/tests/configuration.spec.ts | 19 +- vscode/extension/tests/utils_code_server.ts | 7 + vscode/extension/vitest.config.ts | 6 + 10 files changed, 490 insertions(+), 27 deletions(-) create mode 100644 vscode/extension/src/auth/auth.test.ts create mode 100644 vscode/extension/src/utilities/sqlmesh/sqlmesh.test.ts diff --git a/vscode/extension/src/auth/auth.test.ts b/vscode/extension/src/auth/auth.test.ts new file mode 100644 index 0000000000..6aaac40c43 --- /dev/null +++ b/vscode/extension/src/auth/auth.test.ts @@ -0,0 +1,82 @@ +// SPDX-License-Identifier: Apache-2.0 + +import { beforeEach, describe, expect, it, vi } from 'vitest' +import { ok } from '@bus/result' + +const mocks = vi.hoisted(() => ({ + execAsync: vi.fn(), + getProjectRoot: vi.fn(), + getTcloudBin: vi.fn(), + showInformationMessage: vi.fn(), +})) + +vi.mock('vscode', () => ({ + env: { openExternal: vi.fn() }, + Uri: { parse: vi.fn((value: string) => value) }, + EventEmitter: class { + event = vi.fn() + fire = vi.fn() + }, + window: { showInformationMessage: mocks.showInformationMessage }, +})) + +vi.mock('../utilities/exec', () => ({ execAsync: mocks.execAsync })) +vi.mock('../utilities/common/utilities', () => ({ + getProjectRoot: mocks.getProjectRoot, +})) +vi.mock('../utilities/sqlmesh/sqlmesh', () => ({ + getTcloudBin: mocks.getTcloudBin, +})) +vi.mock('../utilities/common/log', () => ({ traceError: vi.fn() })) + +import { AuthenticationProviderTobikoCloud } from './auth' + +describe('AuthenticationProviderTobikoCloud telemetry environment', () => { + beforeEach(() => { + vi.clearAllMocks() + mocks.getProjectRoot.mockResolvedValue({ uri: { fsPath: '/workspace' } }) + mocks.getTcloudBin.mockResolvedValue( + ok({ + bin: '/venv/bin/tcloud', + workspacePath: '/workspace', + env: { SQLMESH__DISABLE_ANONYMIZED_ANALYTICS: 'true' }, + args: [], + }), + ) + mocks.execAsync + .mockResolvedValueOnce({ + exitCode: 0, + stdout: JSON.stringify({ + url: 'https://example.com/login', + verifier_code: 'verifier', + }), + stderr: '', + }) + .mockResolvedValueOnce({ exitCode: 0, stdout: '', stderr: '' }) + // Dismissing the prompt exits without waiting for the mocked login server. + mocks.showInformationMessage.mockResolvedValue(undefined) + }) + + it('passes the telemetry-aware environment to every OAuth tcloud subprocess', async () => { + const provider = new AuthenticationProviderTobikoCloud() + + await provider.sign_in_oauth_flow() + + const expectedOptions = expect.objectContaining({ + cwd: '/workspace', + env: { SQLMESH__DISABLE_ANONYMIZED_ANALYTICS: 'true' }, + }) + expect(mocks.execAsync).toHaveBeenNthCalledWith( + 1, + '/venv/bin/tcloud', + ['auth', 'vscode', 'login-url'], + expectedOptions, + ) + expect(mocks.execAsync).toHaveBeenNthCalledWith( + 2, + '/venv/bin/tcloud', + ['auth', 'vscode', 'start-server', 'verifier'], + expectedOptions, + ) + }) +}) diff --git a/vscode/extension/src/auth/auth.ts b/vscode/extension/src/auth/auth.ts index 8d7908f06b..73daf058a4 100644 --- a/vscode/extension/src/auth/auth.ts +++ b/vscode/extension/src/auth/auth.ts @@ -193,6 +193,7 @@ export class AuthenticationProviderTobikoCloud ['auth', 'vscode', 'login-url'], { cwd: workspacePath.uri.fsPath, + env: tcloudBinPath.env, }, ) if (result.exitCode !== 0) { diff --git a/vscode/extension/src/extension.ts b/vscode/extension/src/extension.ts index db72414389..57d44f8eeb 100644 --- a/vscode/extension/src/extension.ts +++ b/vscode/extension/src/extension.ts @@ -24,7 +24,10 @@ import { traceError, } from './utilities/common/log' import { onDidChangePythonInterpreter } from './utilities/common/python' -import { requiresLspRestart } from './utilities/common/configurationChange' +import { + requiresLspRestart, + restartLspOnTelemetryChange, +} from './utilities/common/configurationChange' import { coalesceAsync } from './utilities/coalesceAsync' import { sleep } from './utilities/sleep' import { ErrorType, handleError } from './utilities/errors' @@ -128,6 +131,15 @@ export async function activate(context: vscode.ExtensionContext) { } } + // Subscribe before starting the first server so a consent change during a + // slow startup is not lost. The subscription defers that restart until the + // first client and its test controller are fully initialized. + const telemetryRestartSubscription = restartLspOnTelemetryChange( + vscode.env.onDidChangeTelemetryEnabled, + restartLsp, + ) + context.subscriptions.push(telemetryRestartSubscription) + // commands needing the restart helper context.subscriptions.push( vscode.commands.registerCommand( @@ -141,24 +153,34 @@ export async function activate(context: vscode.ExtensionContext) { vscode.commands.registerCommand('sqlmesh.signout', signOut(authProvider)), ) - // Instantiate the LSP client (once) - lspClient = new LSPClient() - const startResult = await lspClient.start() - if (isErr(startResult)) { - await handleError( - authProvider, - restartLsp, - startResult.error, - 'Failed to start LSP', - ) - return // abort activation – nothing else to do - } + // Instantiate the LSP client once. The telemetry subscription is completed + // in a finally block so even a failed initial start cannot leave every later + // telemetry change deferred forever. + const initialLspClient = new LSPClient() + lspClient = initialLspClient + let initialStartSucceeded = false + await telemetryRestartSubscription.runDuringInitialStart(async () => { + const startResult = await initialLspClient.start() + if (isErr(startResult)) { + await handleError( + authProvider, + restartLsp, + startResult.error, + 'Failed to start LSP', + ) + return + } - context.subscriptions.push(lspClient) + context.subscriptions.push(initialLspClient) - // Initialize the test controller - testControllerDisposable = setupTestController(lspClient) - context.subscriptions.push(testControllerDisposable, testController) + // Initialize the test controller + testControllerDisposable = setupTestController(initialLspClient) + context.subscriptions.push(testControllerDisposable, testController) + initialStartSucceeded = true + }) + if (!initialStartSucceeded) { + return // abort activation – nothing else to do + } // Register the rendered model provider const renderedModelProvider = new RenderedModelProvider() diff --git a/vscode/extension/src/utilities/common/configurationChange.test.ts b/vscode/extension/src/utilities/common/configurationChange.test.ts index 1a9c633ead..c381e5b2d0 100644 --- a/vscode/extension/src/utilities/common/configurationChange.test.ts +++ b/vscode/extension/src/utilities/common/configurationChange.test.ts @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 -import { describe, expect, it } from 'vitest' -import { requiresLspRestart } from './configurationChange' +import { describe, expect, it, vi } from 'vitest' +import { + requiresLspRestart, + restartLspOnTelemetryChange, +} from './configurationChange' /** * Build a stand-in for `vscode.ConfigurationChangeEvent` from the settings that @@ -58,3 +61,73 @@ describe('requiresLspRestart', () => { ).toBe(true) }) }) + +describe('restartLspOnTelemetryChange', () => { + const setup = () => { + let listener: ((enabled: boolean) => void) | undefined + const disposable = { dispose: vi.fn() } + const onDidChangeTelemetryEnabled = vi.fn( + (registeredListener: (enabled: boolean) => void) => { + listener = registeredListener + return disposable + }, + ) + const restartLsp = vi.fn(() => Promise.resolve()) + + const subscription = restartLspOnTelemetryChange( + onDidChangeTelemetryEnabled, + restartLsp, + ) + + return { disposable, listener: () => listener, restartLsp, subscription } + } + + it('defers and collapses telemetry changes during initial LSP startup', async () => { + const { listener, restartLsp, subscription } = setup() + + await subscription.runDuringInitialStart(() => { + listener()?.(false) + listener()?.(true) + expect(restartLsp).not.toHaveBeenCalled() + return Promise.resolve() + }) + + expect(restartLsp).toHaveBeenCalledOnce() + }) + + it('restarts for every telemetry preference change after initial startup', async () => { + const { listener, restartLsp, subscription } = setup() + await subscription.runDuringInitialStart(() => Promise.resolve()) + + expect(restartLsp).not.toHaveBeenCalled() + + listener()?.(false) + await vi.waitFor(() => expect(restartLsp).toHaveBeenCalledTimes(1)) + + listener()?.(true) + await vi.waitFor(() => expect(restartLsp).toHaveBeenCalledTimes(2)) + }) + + it('releases deferred restarts when initial startup fails', async () => { + const { listener, restartLsp, subscription } = setup() + + await expect( + subscription.runDuringInitialStart(() => { + listener()?.(false) + return Promise.reject(new Error('initial start failed')) + }), + ).rejects.toThrow('initial start failed') + + expect(restartLsp).toHaveBeenCalledOnce() + listener()?.(true) + await vi.waitFor(() => expect(restartLsp).toHaveBeenCalledTimes(2)) + }) + + it('disposes the underlying VS Code event listener', () => { + const { disposable, subscription } = setup() + + subscription.dispose() + + expect(disposable.dispose).toHaveBeenCalledOnce() + }) +}) diff --git a/vscode/extension/src/utilities/common/configurationChange.ts b/vscode/extension/src/utilities/common/configurationChange.ts index 7e3ab1f3e0..98a27bf4a7 100644 --- a/vscode/extension/src/utilities/common/configurationChange.ts +++ b/vscode/extension/src/utilities/common/configurationChange.ts @@ -21,6 +21,18 @@ export interface ConfigurationChange { affectsConfiguration(section: string): boolean } +interface Disposable { + dispose(): unknown +} + +export interface TelemetryRestartSubscription extends Disposable { + runDuringInitialStart(task: () => Promise): Promise +} + +type TelemetryChangeEvent = ( + listener: (enabled: boolean) => unknown, +) => TDisposable + /** * Whether a configuration change affects a setting the language server reads. * @@ -33,3 +45,34 @@ export function requiresLspRestart(event: ConfigurationChange): boolean { event.affectsConfiguration(section), ) } + +/** Restart the language server whenever VS Code's effective telemetry preference changes. */ +export function restartLspOnTelemetryChange( + event: TelemetryChangeEvent, + restartLsp: () => Promise, +): TelemetryRestartSubscription { + let initialStartComplete = false + let restartPending = false + const subscription = event(() => { + if (!initialStartComplete) { + restartPending = true + return + } + void restartLsp() + }) + + return { + dispose: () => subscription.dispose(), + runDuringInitialStart: async task => { + try { + return await task() + } finally { + initialStartComplete = true + if (restartPending) { + restartPending = false + await restartLsp() + } + } + }, + } +} diff --git a/vscode/extension/src/utilities/sqlmesh/sqlmesh.test.ts b/vscode/extension/src/utilities/sqlmesh/sqlmesh.test.ts new file mode 100644 index 0000000000..e9753dcb47 --- /dev/null +++ b/vscode/extension/src/utilities/sqlmesh/sqlmesh.test.ts @@ -0,0 +1,194 @@ +// SPDX-License-Identifier: Apache-2.0 + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest' + +// Mutable telemetry state — flip per test. +let _isTelemetryEnabled = true + +vi.mock('vscode', () => ({ + default: {}, + env: { + get isTelemetryEnabled(): boolean { + return _isTelemetryEnabled + }, + }, + ProgressLocation: { Notification: 15 }, + window: { + withProgress: vi.fn(), + showInformationMessage: vi.fn(), + }, +})) + +vi.mock('../common/python', () => ({ + getInterpreterDetails: vi.fn(), + getPythonEnvVariables: vi.fn(), +})) + +vi.mock('../common/log', () => ({ + traceInfo: vi.fn(), + traceLog: vi.fn(), + traceVerbose: vi.fn(), + traceError: vi.fn(), + traceWarn: vi.fn(), +})) + +vi.mock('../common/utilities', () => ({ + getProjectRoot: vi.fn(), +})) + +vi.mock('../python', () => ({ + isPythonModuleInstalled: vi.fn(), +})) + +vi.mock('../../auth/auth', () => ({ + isSignedIntoTobikoCloud: vi.fn(), +})) + +vi.mock('../exec', () => ({ + execAsync: vi.fn(), +})) + +vi.mock('../config', () => ({ + getSqlmeshLspEntryPoint: vi.fn().mockReturnValue(undefined), + resolveProjectPath: vi.fn(), +})) + +vi.mock('../isWindows', () => ({ + IS_WINDOWS: false, +})) + +import { getSqlmeshEnvironment, sqlmeshLspExec } from './sqlmesh' +import { getInterpreterDetails, getPythonEnvVariables } from '../common/python' +import { getProjectRoot } from '../common/utilities' +import { getSqlmeshLspEntryPoint, resolveProjectPath } from '../config' +import { ok } from '@bus/result' + +const ANALYTICS_KEY = 'SQLMESH__DISABLE_ANONYMIZED_ANALYTICS' +let originalAnalyticsValue: string | undefined +let originallyHadAnalyticsKey = false + +beforeEach(() => { + originallyHadAnalyticsKey = Object.prototype.hasOwnProperty.call( + process.env, + ANALYTICS_KEY, + ) + originalAnalyticsValue = process.env[ANALYTICS_KEY] + Reflect.deleteProperty(process.env, ANALYTICS_KEY) + _isTelemetryEnabled = true +}) + +afterEach(() => { + if (originallyHadAnalyticsKey) { + process.env[ANALYTICS_KEY] = originalAnalyticsValue + } else { + Reflect.deleteProperty(process.env, ANALYTICS_KEY) + } + vi.clearAllMocks() +}) + +describe('getSqlmeshEnvironment telemetry', () => { + const baseInterpreterDetails = { + path: ['/usr/bin/python3'], + binPath: undefined as string | undefined, + isVirtualEnvironment: false, + resource: undefined, + } + + beforeEach(() => { + vi.mocked(getInterpreterDetails).mockResolvedValue(baseInterpreterDetails) + vi.mocked(getPythonEnvVariables).mockResolvedValue(ok({})) + }) + + it('sets SQLMESH__DISABLE_ANONYMIZED_ANALYTICS=true when VS Code telemetry is disabled', async () => { + _isTelemetryEnabled = false + + const result = await getSqlmeshEnvironment() + + expect(result.ok).toBe(true) + if (result.ok) { + expect(result.value[ANALYTICS_KEY]).toBe('true') + } + expect(process.env[ANALYTICS_KEY]).toBeUndefined() + }) + + it('overrides a conflicting false value when VS Code telemetry is disabled', async () => { + _isTelemetryEnabled = false + // Simulate an inherited env that tries to re-enable analytics. + vi.mocked(getPythonEnvVariables).mockResolvedValue( + ok({ [ANALYTICS_KEY]: 'false' }), + ) + + const result = await getSqlmeshEnvironment() + + expect(result.ok).toBe(true) + if (result.ok) { + expect(result.value[ANALYTICS_KEY]).toBe('true') + } + }) + + it('does not add SQLMESH__DISABLE_ANONYMIZED_ANALYTICS when VS Code telemetry is enabled', async () => { + _isTelemetryEnabled = true + + const result = await getSqlmeshEnvironment() + + expect(result.ok).toBe(true) + if (result.ok) { + // The variable must be absent, not just falsy — any presence would + // disable analytics even when the user opted in. + expect(Object.prototype.hasOwnProperty.call(result.value, ANALYTICS_KEY)).toBe(false) + } + }) + + it('preserves a user-provided value unchanged when VS Code telemetry is enabled', async () => { + _isTelemetryEnabled = true + vi.mocked(getPythonEnvVariables).mockResolvedValue( + ok({ [ANALYTICS_KEY]: 'true' }), + ) + + const result = await getSqlmeshEnvironment() + + expect(result.ok).toBe(true) + if (result.ok) { + expect(result.value[ANALYTICS_KEY]).toBe('true') + } + }) +}) + +describe('sqlmeshLspExec telemetry', () => { + beforeEach(() => { + vi.mocked(getProjectRoot).mockResolvedValue({} as never) + vi.mocked(resolveProjectPath).mockReturnValue( + ok({ + workspaceFolder: '/workspace', + projectPaths: undefined, + }), + ) + vi.mocked(getSqlmeshLspEntryPoint).mockReturnValue({ + entrypoint: '/custom/sqlmesh_lsp', + args: ['--debug'], + }) + }) + + it('disables analytics for a configured LSP entry point when VS Code telemetry is disabled', async () => { + _isTelemetryEnabled = false + + const result = await sqlmeshLspExec() + + expect(result.ok).toBe(true) + if (result.ok) { + expect(result.value.env[ANALYTICS_KEY]).toBe('true') + } + expect(process.env[ANALYTICS_KEY]).toBeUndefined() + }) + + it('does not add an analytics override for a configured LSP entry point when VS Code telemetry is enabled', async () => { + _isTelemetryEnabled = true + + const result = await sqlmeshLspExec() + + expect(result.ok).toBe(true) + if (result.ok) { + expect(Object.prototype.hasOwnProperty.call(result.value.env, ANALYTICS_KEY)).toBe(false) + } + }) +}) diff --git a/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts b/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts index c9e181fc06..e52e4dc80b 100644 --- a/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts +++ b/vscode/extension/src/utilities/sqlmesh/sqlmesh.ts @@ -9,7 +9,7 @@ import { ErrorType } from '../errors' import { isSignedIntoTobikoCloud } from '../../auth/auth' import { execAsync } from '../exec' import z from 'zod' -import { ProgressLocation, window } from 'vscode' +import { ProgressLocation, window, env as vscodeEnv } from 'vscode' import { IS_WINDOWS } from '../isWindows' import { getSqlmeshLspEntryPoint, resolveProjectPath } from '../config' import { isSemVerGreaterThanOrEqual } from '../semver' @@ -21,6 +21,30 @@ export interface SqlmeshExecInfo { args: string[] } +/** + * Applies the VS Code telemetry preference to a copy of the given environment. + * + * When the user has disabled telemetry through VS Code, + * SQLMESH__DISABLE_ANONYMIZED_ANALYTICS is forced to "true" so that the + * language-server process cannot send analytics regardless of its own config. + * When telemetry is enabled the variable is left untouched, preserving any + * project-level SQLMesh configuration the user may have set. + * + * The function never mutates process.env or the caller's object. + */ +function applyTelemetryEnv(env: Record): Record +function applyTelemetryEnv( + env: Record, +): Record +function applyTelemetryEnv( + env: Record, +): Record { + if (vscodeEnv.isTelemetryEnabled) { + return { ...env } + } + return { ...env, SQLMESH__DISABLE_ANONYMIZED_ANALYTICS: 'true' } +} + /** * Gets the current SQLMesh environment variables that would be used for execution. * This is useful for debugging and understanding the environment configuration. @@ -53,7 +77,7 @@ export async function getSqlmeshEnvironment(): Promise { const entrypointFile = path.join(tempDir, entrypointFileName) const fileWhereStoredInputs = path.join(tempDir, 'inputs.txt') + const fileWhereStoredTelemetryPreference = path.join( + tempDir, + 'telemetry-preference.txt', + ) const sqlmeshLSPFile = path.join(tempDir, '.venv/bin/sqlmesh_lsp') // Create the entrypoint file @@ -50,6 +55,7 @@ const createEntrypointFile = ( entrypointFile, `#!/bin/bash echo "$@" > ${fileWhereStoredInputs} +echo "\${SQLMESH__DISABLE_ANONYMIZED_ANALYTICS:-}" > ${fileWhereStoredTelemetryPreference} # Strip bitToStripFromArgs from the beginning of the args if it matches if [[ "$1" == "${bitToStripFromArgs}" ]]; then shift @@ -62,6 +68,7 @@ ${sqlmeshLSPFile} "$@"`, return { entrypointFile, fileWhereStoredInputs, + fileWhereStoredTelemetryPreference, } } @@ -75,10 +82,8 @@ test.describe('Test LSP Entrypoint configuration', () => { await setupPythonEnvironment(tempDir) - const { fileWhereStoredInputs } = createEntrypointFile( - tempDir, - 'entrypoint.sh', - ) + const { fileWhereStoredInputs, fileWhereStoredTelemetryPreference } = + createEntrypointFile(tempDir, 'entrypoint.sh') const settings = { 'sqlmesh.lspEntrypoint': './entrypoint.sh', @@ -111,6 +116,12 @@ test.describe('Test LSP Entrypoint configuration', () => { expect(fs.existsSync(fileWhereStoredInputs)).toBe(true) expect(fs.readFileSync(fileWhereStoredInputs, 'utf8')).toBe(`--stdio `) + // The e2e code-server is launched with --disable-telemetry. Verify the + // packaged extension propagates that global preference to the real LSP + // child process, not only to a mocked environment builder. + expect(fs.readFileSync(fileWhereStoredTelemetryPreference, 'utf8')).toBe( + 'true\n', + ) }) test('specify one entrypoint absolute path', async ({ diff --git a/vscode/extension/tests/utils_code_server.ts b/vscode/extension/tests/utils_code_server.ts index 68bf2ed597..d35c2678bb 100644 --- a/vscode/extension/tests/utils_code_server.ts +++ b/vscode/extension/tests/utils_code_server.ts @@ -80,6 +80,13 @@ export async function startCodeServer({ const userDataDir = await fs.mkdtemp( path.join(os.tmpdir(), 'vscode-test-sushi-user-data-dir-'), ) + const userSettingsDir = path.join(userDataDir, 'User') + await fs.ensureDir(userSettingsDir) + await fs.writeJson(path.join(userSettingsDir, 'settings.json'), { + // `--disable-telemetry` disables code-server's own telemetry, but VS Code + // derives `env.isTelemetryEnabled` from this global user preference. + 'telemetry.telemetryLevel': 'off', + }) // Start code-server instance using the shared extensions directory const codeServerProcess = spawn( diff --git a/vscode/extension/vitest.config.ts b/vscode/extension/vitest.config.ts index 49fbea3be3..0399bf274d 100644 --- a/vscode/extension/vitest.config.ts +++ b/vscode/extension/vitest.config.ts @@ -1,6 +1,12 @@ import { defineConfig } from 'vitest/config' +import path from 'path' export default defineConfig({ + resolve: { + alias: { + '@bus': path.resolve(__dirname, '../bus/src'), + }, + }, test: { globals: true, include: ['src/**/*.test.ts'],