diff --git a/cspell.json b/cspell.json index 8aa54f8759..014f254e81 100644 --- a/cspell.json +++ b/cspell.json @@ -48,6 +48,7 @@ "esmock", "evalue", "findstr", + "fromlist", "getsitepackages", "IMAGENAME", "ipykernel", diff --git a/package-lock.json b/package-lock.json index f8e9ef939b..877b0a0722 100644 --- a/package-lock.json +++ b/package-lock.json @@ -41,6 +41,7 @@ "clsx": "^2.1.1", "cross-fetch": "^3.1.5", "d3-format": "^3.1.0", + "dotenv": "^17.2.3", "encoding": "^0.1.13", "express": "^5.2.1", "fast-deep-equal": "^2.0.1", @@ -15521,6 +15522,18 @@ "url": "https://github.com/fb55/domutils?sponsor=1" } }, + "node_modules/dotenv": { + "version": "17.4.2", + "resolved": "https://registry.npmjs.org/dotenv/-/dotenv-17.4.2.tgz", + "integrity": "sha512-nI4U3TottKAcAD9LLud4Cb7b2QztQMUEfHbvhTH09bqXTxnSie8WnjPALV/WMCrJZ6UV/qHJ6L03OqO3LcdYZw==", + "license": "BSD-2-Clause", + "engines": { + "node": ">=12" + }, + "funding": { + "url": "https://dotenvx.com" + } + }, "node_modules/dunder-proto": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/dunder-proto/-/dunder-proto-1.0.1.tgz", @@ -46513,6 +46526,11 @@ "domhandler": "^5.0.3" } }, + "dotenv": { + "version": "17.4.2", + "resolved": "https://registry.npmjs.org/dotenv/-/dotenv-17.4.2.tgz", + "integrity": "sha512-nI4U3TottKAcAD9LLud4Cb7b2QztQMUEfHbvhTH09bqXTxnSie8WnjPALV/WMCrJZ6UV/qHJ6L03OqO3LcdYZw==" + }, "dunder-proto": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/dunder-proto/-/dunder-proto-1.0.1.tgz", diff --git a/package.json b/package.json index 8dc22af21d..f37d31f179 100644 --- a/package.json +++ b/package.json @@ -1649,6 +1649,12 @@ "description": "When enabled, outputs are saved to separate snapshot files in a 'snapshots' folder instead of the main .deepnote file.", "scope": "resource" }, + "deepnote.integrations.envFile.enabled": { + "type": "boolean", + "default": true, + "description": "When enabled, integration credentials are also loaded from a '.deepnote.env.yaml' file (with 'env:' references resolved against '.env' and environment variables) next to the .deepnote file or at the workspace root.", + "scope": "resource" + }, "deepnote.experiments.enabled": { "type": "boolean", "default": true, @@ -2703,6 +2709,7 @@ "clsx": "^2.1.1", "cross-fetch": "^3.1.5", "d3-format": "^3.1.0", + "dotenv": "^17.2.3", "encoding": "^0.1.13", "express": "^5.2.1", "fast-deep-equal": "^2.0.1", diff --git a/src/kernels/deepnote/deepnoteLspClientManager.node.ts b/src/kernels/deepnote/deepnoteLspClientManager.node.ts index 907adf8f8f..8dee764518 100644 --- a/src/kernels/deepnote/deepnoteLspClientManager.node.ts +++ b/src/kernels/deepnote/deepnoteLspClientManager.node.ts @@ -1,7 +1,7 @@ import * as fs from 'fs'; import * as vscode from 'vscode'; import { CancellationError } from 'vscode'; -import { inject, injectable } from 'inversify'; +import { inject, injectable, optional } from 'inversify'; import type { LanguageClient as LanguageClientType, LanguageClientOptions, @@ -26,7 +26,8 @@ import { noop } from '../../platform/common/utils/misc'; import { IIntegrationStorage, IPlatformNotebookEditorProvider, - IPlatformDeepnoteNotebookManager + IPlatformDeepnoteNotebookManager, + ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; import { ConfigurableDatabaseIntegrationConfig } from '../../platform/notebooks/deepnote/integrationTypes'; import { SqlLspConnection, isSupportedBySqlLsp, convertToSqlLspConnection } from './sqlLspConnectionUtils'; @@ -64,7 +65,10 @@ export class DeepnoteLspClientManager @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, @inject(IPlatformNotebookEditorProvider) private readonly notebookEditorProvider: IPlatformNotebookEditorProvider, - @inject(IPlatformDeepnoteNotebookManager) private readonly notebookManager: IPlatformDeepnoteNotebookManager + @inject(IPlatformDeepnoteNotebookManager) private readonly notebookManager: IPlatformDeepnoteNotebookManager, + @inject(ISqlIntegrationEnvVarsProvider) + @optional() + private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider ) { this.disposables.push(this); } @@ -623,13 +627,19 @@ export class DeepnoteLspClientManager logger.trace(`SQL LSP: Found ${projectIntegrations.length} integrations in project ${projectId}`); - const projectIntegrationConfigs = ( - await Promise.all( - projectIntegrations.map((integration) => - this.integrationStorage.getIntegrationConfig(integration.id) - ) - ) - ).filter((config): config is ConfigurableDatabaseIntegrationConfig => config != null); + // Prefer the merged (SecretStorage + `.deepnote.env.yaml`) configs so file-configured databases also get + // LSP autocomplete/schema (F13); fall back to SecretStorage-only when the merged provider is unavailable (e.g. web). + const projectIntegrationConfigs = this.sqlIntegrationEnvVars + ? (await this.sqlIntegrationEnvVars.getMergedConfigs(notebookUri)).filter( + (config): config is ConfigurableDatabaseIntegrationConfig => config.type !== 'pandas-dataframe' + ) + : ( + await Promise.all( + projectIntegrations.map((integration) => + this.integrationStorage.getIntegrationConfig(integration.id) + ) + ) + ).filter((config): config is ConfigurableDatabaseIntegrationConfig => config != null); const connections = projectIntegrationConfigs .filter((config) => isSupportedBySqlLsp(config.type)) diff --git a/src/kernels/helpers.ts b/src/kernels/helpers.ts index 11c88baa21..b85555cdcd 100644 --- a/src/kernels/helpers.ts +++ b/src/kernels/helpers.ts @@ -747,6 +747,45 @@ export async function executeSilently( return outputs; } +/** + * Leak-safe silent execution, for injecting secrets into a live kernel. + * + * `silent: true` suppresses the `execute_input` broadcast, the output cache (`Out[]`/`_`) and the + * execution-count bump; `store_history: false` keeps the code out of `_ih`/`In` and `history.sqlite`. + * Unlike {@link executeSilently} the code body is never logged, and IOPub is deliberately not read so + * no traceback can echo the source line back to the caller. + * + * Success is *only* `status: 'ok'`. Every other outcome — `error`, `abort`, a missing status, or a + * rejected/disposed future (a kernel that died or restarted mid-request) — is reported as a single + * content-free error output, so nothing about the executed code leaks through the result either. + */ +export async function executeSilentlyLeakSafe( + kernelConnection: Kernel.IKernelConnection, + code: string +): Promise { + const errorMarker: nbformat.IError[] = [{ output_type: 'error', ename: 'error', evalue: '', traceback: [] }]; + + try { + const request = kernelConnection.requestExecute( + { + code: code.replace(/\r\n/g, '\n'), + silent: true, + stop_on_error: false, + allow_stdin: false, + store_history: false + }, + true + ); + + // The shell `execute_reply` still resolves under `silent: true`. + const reply = await request.done; + + return reply?.content.status === 'ok' ? [] : errorMarker; + } catch { + return errorMarker; + } +} + export function executeSilentlyAndEmitOutput( kernelConnection: Kernel.IKernelConnection, code: string, diff --git a/src/kernels/helpers.unit.test.ts b/src/kernels/helpers.unit.test.ts index 0d6d163881..3a1857e648 100644 --- a/src/kernels/helpers.unit.test.ts +++ b/src/kernels/helpers.unit.test.ts @@ -13,6 +13,7 @@ import { PythonKernelConnectionMetadata } from './types'; import { EnvironmentType, PythonEnvironment } from '../platform/pythonEnvironments/info'; +import { logger } from '../platform/logging'; import { PythonExtension, type Environment } from '@vscode/python-extension'; import { resolvableInstance } from '../test/datascience/helpers'; import { DisposableStore, dispose } from '../platform/common/utils/lifecycle'; @@ -958,3 +959,181 @@ suite('Kernel Connection Helpers', () => { }); }); }); + +suite('executeSilentlyLeakSafe', () => { + /** The single output every non-`ok` outcome must collapse to: no ename detail, no value, no traceback. */ + const CONTENT_FREE_ERROR = { output_type: 'error', ename: 'error', evalue: '', traceback: [] }; + + /** Stands in for a credential embedded in the executed snippet. */ + const SECRET_CODE = '__import__("deepnote_toolkit.env").set_env("SQL_DEMO", "p4ssw0rd-do-not-log")'; + + type ExecuteReply = { content?: { status?: string; [field: string]: unknown } }; + + let requests: { content: any; disposeOnDone: unknown }[]; + let logCalls: sinon.SinonStub[]; + + setup(() => { + requests = []; + logCalls = []; + }); + + teardown(() => { + sinon.restore(); + }); + + /** A kernel connection whose `execute_reply` resolves with `reply`, or rejects if `rejectWith` is given. */ + function createMockKernel(options: { reply?: ExecuteReply; rejectWith?: Error; throwSync?: Error }) { + return { + requestExecute: (content: any, disposeOnDone: unknown) => { + requests.push({ content, disposeOnDone }); + + if (options.throwSync) { + throw options.throwSync; + } + + return { + done: options.rejectWith ? Promise.reject(options.rejectWith) : Promise.resolve(options.reply) + }; + } + }; + } + + /** Stubs every logger method so any logged text can be inspected. */ + function stubLogger() { + logCalls = (['error', 'warn', 'info', 'debug', 'trace', 'ci'] as const).map((level) => + sinon.stub(logger, level) + ); + } + + /** Every argument passed to any logger method, flattened to strings. */ + function allLoggedText(): string { + return logCalls + .flatMap((stub) => stub.getCalls()) + .flatMap((call) => call.args) + .map((arg) => (typeof arg === 'string' ? arg : JSON.stringify(arg) ?? String(arg))) + .join('\n'); + } + + test('sends exactly the leak-safe request flags', async () => { + const mockKernel = createMockKernel({ reply: { content: { status: 'ok' } } }); + + const { executeSilentlyLeakSafe } = await import('./helpers'); + await executeSilentlyLeakSafe(mockKernel as any, SECRET_CODE); + + assert.equal(requests.length, 1); + assert.deepStrictEqual(requests[0].content, { + code: SECRET_CODE, + silent: true, + stop_on_error: false, + allow_stdin: false, + store_history: false + }); + assert.strictEqual(requests[0].disposeOnDone, true); + }); + + test('normalizes CRLF line endings in the code', async () => { + const mockKernel = createMockKernel({ reply: { content: { status: 'ok' } } }); + + const { executeSilentlyLeakSafe } = await import('./helpers'); + await executeSilentlyLeakSafe(mockKernel as any, 'a = 1\r\nb = 2\r\n'); + + assert.strictEqual(requests[0].content.code, 'a = 1\nb = 2\n'); + }); + + test('resolves to no outputs when the reply status is ok', async () => { + const mockKernel = createMockKernel({ reply: { content: { status: 'ok' } } }); + + const { executeSilentlyLeakSafe } = await import('./helpers'); + const result = await executeSilentlyLeakSafe(mockKernel as any, SECRET_CODE); + + assert.deepStrictEqual(result, []); + }); + + test('resolves to one content-free error output on status error', async () => { + const mockKernel = createMockKernel({ + reply: { content: { status: 'error', ename: 'ValueError', evalue: SECRET_CODE, traceback: [SECRET_CODE] } } + }); + + const { executeSilentlyLeakSafe } = await import('./helpers'); + const result = await executeSilentlyLeakSafe(mockKernel as any, SECRET_CODE); + + // Nothing from the reply is copied through — a traceback renders the failing source line. + assert.deepStrictEqual(result, [CONTENT_FREE_ERROR]); + }); + + test('resolves to one content-free error output on status abort', async () => { + // Reachable when a request is aborted around a restart or shutdown. Treating it as success + // would wrongly advance the caller's removal baseline. + const mockKernel = createMockKernel({ reply: { content: { status: 'abort' } } }); + + const { executeSilentlyLeakSafe } = await import('./helpers'); + const result = await executeSilentlyLeakSafe(mockKernel as any, SECRET_CODE); + + assert.deepStrictEqual(result, [CONTENT_FREE_ERROR]); + }); + + test('resolves to one content-free error output when the status is missing', async () => { + const mockKernel = createMockKernel({ reply: { content: {} } }); + + const { executeSilentlyLeakSafe } = await import('./helpers'); + const result = await executeSilentlyLeakSafe(mockKernel as any, SECRET_CODE); + + assert.deepStrictEqual(result, [CONTENT_FREE_ERROR]); + }); + + test('resolves to one content-free error output when there is no reply at all', async () => { + const mockKernel = createMockKernel({ reply: undefined }); + + const { executeSilentlyLeakSafe } = await import('./helpers'); + const result = await executeSilentlyLeakSafe(mockKernel as any, SECRET_CODE); + + assert.deepStrictEqual(result, [CONTENT_FREE_ERROR]); + }); + + test('resolves to the same marker instead of throwing when the request future rejects', async () => { + // A refresh racing a restart or shutdown rejects the disposed future. + const mockKernel = createMockKernel({ rejectWith: new Error('Kernel disposed') }); + + const { executeSilentlyLeakSafe } = await import('./helpers'); + const result = await executeSilentlyLeakSafe(mockKernel as any, SECRET_CODE); + + assert.deepStrictEqual(result, [CONTENT_FREE_ERROR]); + }); + + test('resolves to the same marker instead of throwing when requestExecute throws synchronously', async () => { + const mockKernel = createMockKernel({ throwSync: new Error('Kernel is dead') }); + + const { executeSilentlyLeakSafe } = await import('./helpers'); + const result = await executeSilentlyLeakSafe(mockKernel as any, SECRET_CODE); + + assert.deepStrictEqual(result, [CONTENT_FREE_ERROR]); + }); + + test('never logs the code body, on success or on any failure', async () => { + stubLogger(); + + const { executeSilentlyLeakSafe } = await import('./helpers'); + const outcomes = [ + createMockKernel({ reply: { content: { status: 'ok' } } }), + createMockKernel({ reply: { content: { status: 'error' } } }), + createMockKernel({ reply: { content: { status: 'abort' } } }), + createMockKernel({ reply: { content: {} } }), + createMockKernel({ rejectWith: new Error('Kernel disposed') }), + createMockKernel({ throwSync: new Error('Kernel is dead') }) + ]; + + for (const mockKernel of outcomes) { + await executeSilentlyLeakSafe(mockKernel as any, SECRET_CODE); + } + + // The regression this guards: executeSilently logs the first 100 characters of the code, and + // the full source on error when traceErrors is set. Neither may happen here. + const logged = allLoggedText(); + + assert.notInclude(logged, 'p4ssw0rd-do-not-log'); + assert.notInclude(logged, SECRET_CODE); + assert.notInclude(logged, SECRET_CODE.substring(0, 100)); + assert.notInclude(logged, 'set_env'); + assert.strictEqual(logged, '', 'the leak-safe primitive must log nothing at all'); + }); +}); diff --git a/src/kernels/kernelExecution.ts b/src/kernels/kernelExecution.ts index 99b36ce964..18dd522800 100644 --- a/src/kernels/kernelExecution.ts +++ b/src/kernels/kernelExecution.ts @@ -20,7 +20,7 @@ import { CellExecutionFactory } from './execution/cellExecution'; import { CellExecutionMessageHandlerService } from './execution/cellExecutionMessageHandlerService'; import { CellExecutionQueue } from './execution/cellExecutionQueue'; import { cellOutputToVSCCellOutput, traceCellMessage } from './execution/helpers'; -import { executeSilently } from './helpers'; +import { executeSilently, executeSilentlyLeakSafe } from './helpers'; import { sendKernelTelemetryEvent } from './telemetry/sendKernelTelemetryEvent'; import { IKernel, @@ -317,6 +317,12 @@ export class NotebookKernelExecution implements INotebookKernelExecution { session.kernel ? executeSilently(session.kernel, code) : Promise.reject(new SessionDisposedError()) ); } + executeHiddenSilent(code: string): Promise { + const sessionPromise = this.kernel.start(); + return sessionPromise.then((session) => + session.kernel ? executeSilentlyLeakSafe(session.kernel, code) : Promise.reject(new SessionDisposedError()) + ); + } private async onWillInterrupt() { const executionQueue = this.documentExecutions.get(this.notebook); if (!executionQueue && this.kernel.kernelConnectionMetadata.kind !== 'connectToLiveRemoteKernel') { diff --git a/src/kernels/raw/session/rawKernelConnection.node.ts b/src/kernels/raw/session/rawKernelConnection.node.ts index 5ffdb3f4a0..7346469a6b 100644 --- a/src/kernels/raw/session/rawKernelConnection.node.ts +++ b/src/kernels/raw/session/rawKernelConnection.node.ts @@ -42,9 +42,49 @@ import { KernelProcessExitedError } from '../../errors/kernelProcessExitedError' import { once } from '../../../platform/common/utils/functional'; import { disposeAsync } from '../../../platform/common/utils'; import { generateUuid } from '../../../platform/common/uuid'; +import { getParentHeaderMsgId } from '../../execution/cellExecutionMessageHandler'; let jupyterLabKernel: typeof import('@jupyterlab/services/lib/kernel/default'); +/** + * How many leak-safe request ids the wire dump remembers, so the set cannot grow for the life of the + * kernel. One id is added per integration-env refresh, which is far more history than correlating a + * reply ever needs. + * + * Ids expire by age rather than being dropped when a request completes, and that is deliberate: the + * shell `execute_reply` is not the last message for a request. IOPub travels on its own socket with no + * cross-channel ordering guarantee, and `silent: true` suppresses `execute_input` but *not* IOPub + * `error` — whose traceback renders the failing source line. Retiring an id on its reply would let + * exactly that message through. + */ +const MAX_TRACKED_LEAK_SAFE_MSG_IDS = 32; + +/** + * `silent` together with `store_history: false` is the signature of `executeSilentlyLeakSafe`, the only + * caller that combines them — every other execution path in the extension sends `silent: false`. + */ +function isLeakSafeExecuteRequest(msg: Readonly): boolean { + if (msg.header.msg_type !== 'execute_request') { + return false; + } + + const content = (msg as KernelMessage.IExecuteRequestMsg).content; + + return content.silent === true && content.store_history === false; +} + +function rememberLeakSafeMsgId(msgIds: Set, msgId: string): void { + msgIds.add(msgId); + + if (msgIds.size > MAX_TRACKED_LEAK_SAFE_MSG_IDS) { + // A Set iterates in insertion order, so this evicts the oldest request. + const oldest = msgIds.values().next().value; + if (oldest !== undefined) { + msgIds.delete(oldest); + } + } +} + /* RawKernel class represents the mapping from the JupyterLab services IKernel interface to a raw IPython kernel running on the local machine. RawKernel is in charge of taking @@ -701,7 +741,30 @@ function newRawKernel(kernelProcess: IKernelProcess, clientId: string, username: model }); if (workspace.getConfiguration('deepnote').get('enablePythonKernelLogging', false)) { + // Leak-safe executions (see `executeSilentlyLeakSafe`) embed credentials, so neither the request + // nor anything answering it may be dumped. Redacting the request's `code` alone is not enough: + // an error reply's `traceback` renders the failing source line via `linecache`, independent of + // `store_history`, so replies have to be dropped by parent message id too. + // + // Scope: this covers the raw-kernel path only. Deepnote kernels are remote connection metadata, + // so `KernelSessionFactory` routes them to `JupyterKernelSessionFactory` and `newRawKernel` + // never runs for them — but that path has no equivalent dump (`baseJupyterSessionConnection` + // only re-emits `anyMessage`), so there is nothing to suppress there. This still matters when a + // `.deepnote` file is opened against a plain local Python interpreter kernel. + const leakSafeMsgIds = new Set(); + realKernel.anyMessage.connect((_, msg) => { + if (isLeakSafeExecuteRequest(msg.msg)) { + rememberLeakSafeMsgId(leakSafeMsgIds, msg.msg.header.msg_id); + + return; + } + + const parentMsgId = getParentHeaderMsgId(msg.msg); + if (parentMsgId !== undefined && leakSafeMsgIds.has(parentMsgId)) { + return; + } + logger.trace(`[AnyMessage Event] [${msg.direction}] [${kernelProcess.pid}] ${JSON.stringify(msg.msg)}`); }); } diff --git a/src/kernels/types.ts b/src/kernels/types.ts index fbc44b2ebc..8312bf8a14 100644 --- a/src/kernels/types.ts +++ b/src/kernels/types.ts @@ -492,6 +492,13 @@ export interface INotebookKernelExecution { * Executes arbitrary code against the kernel without incrementing the execution count. */ executeHidden(code: string): Promise; + /** + * Executes arbitrary code against the kernel leaving no trace of it: no IPython history, no + * `execute_input` broadcast, no output cache, and nothing logged. Use this — not `executeHidden` — + * whenever the code embeds secrets. Resolves to `[]` on success, or a single content-free error + * output on any failure (see `executeSilentlyLeakSafe`). + */ + executeHiddenSilent(code: string): Promise; } /** * Kernels created by third party extensions. diff --git a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts index 0c3c048305..a55dfeb7c5 100644 --- a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts +++ b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts @@ -10,7 +10,7 @@ import { IFederatedAuthTokenStorage } from '../types'; /** * Node-only bridge that restarts kernels when a federated integration's token changes, clearing stale - * `os.environ` mutations and kernel globals. Separate from {@link IntegrationKernelRestartHandler} because + * `os.environ` mutations and kernel globals. Separate from {@link IntegrationEnvRefreshHandler} because * {@link IFederatedAuthTokenStorage} is node-only. */ @injectable() diff --git a/src/notebooks/deepnote/integrations/integrationDetector.ts b/src/notebooks/deepnote/integrations/integrationDetector.ts index 5715ead44c..9e872066df 100644 --- a/src/notebooks/deepnote/integrations/integrationDetector.ts +++ b/src/notebooks/deepnote/integrations/integrationDetector.ts @@ -1,4 +1,5 @@ -import { inject, injectable } from 'inversify'; +import { inject, injectable, optional } from 'inversify'; +import { workspace } from 'vscode'; import { logger } from '../../../platform/logging'; import { IDeepnoteNotebookManager } from '../../types'; @@ -8,7 +9,8 @@ import { IntegrationWithStatus } from '../../../platform/notebooks/deepnote/integrationTypes'; import { IIntegrationDetector, IIntegrationStorage } from './types'; -import { databaseIntegrationTypes } from '@deepnote/database-integrations'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; +import { DatabaseIntegrationConfig, databaseIntegrationTypes } from '@deepnote/database-integrations'; /** * Service for detecting integrations used in Deepnote notebooks @@ -17,7 +19,10 @@ import { databaseIntegrationTypes } from '@deepnote/database-integrations'; export class IntegrationDetector implements IIntegrationDetector { constructor( @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, - @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager + @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, + @inject(ISqlIntegrationEnvVarsProvider) + @optional() + private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider ) {} /** @@ -42,6 +47,19 @@ export class IntegrationDetector implements IIntegrationDetector { const projectIntegrations = project.project.integrations?.slice() ?? []; logger.debug(`IntegrationDetector: Found ${projectIntegrations.length} integrations in project.integrations`); + // Merged (SecretStorage + `.deepnote.env.yaml`) configs, so file-configured integrations are not shown as + // unconfigured (F13). Resolved from the open notebook document; when the merged provider is unavailable + // (e.g. web) this stays empty and detection falls back to SecretStorage only. + const mergedConfigsById = new Map(); + const notebookUri = workspace.notebookDocuments.find( + (nb) => nb.metadata?.deepnoteProjectId === projectId && nb.metadata?.deepnoteNotebookId === notebookId + )?.uri; + if (this.sqlIntegrationEnvVars && notebookUri) { + for (const config of await this.sqlIntegrationEnvVars.getMergedConfigs(notebookUri)) { + mergedConfigsById.set(config.id, config); + } + } + for (const projectIntegration of projectIntegrations) { const integrationId = projectIntegration.id; const integrationType = projectIntegration.type; @@ -53,11 +71,12 @@ export class IntegrationDetector implements IIntegrationDetector { continue; } - // Check if the integration is configured + // Configured if SecretStorage has it OR a `.deepnote.env.yaml` file config provides it. const config = await this.integrationStorage.getIntegrationConfig(integrationId); + const isConfigured = config != null || mergedConfigsById.has(integrationId); const status: IntegrationWithStatus = { config: config ?? null, - status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected, + status: isConfigured ? IntegrationStatus.Connected : IntegrationStatus.Disconnected, // Include integration metadata from project for prefilling when config is null integrationName: projectIntegration.name, integrationType: integrationType as ConfigurableDatabaseIntegrationType @@ -66,6 +85,23 @@ export class IntegrationDetector implements IIntegrationDetector { integrations.set(integrationId, status); } + // Append file-only integrations (present in `.deepnote.env.yaml` but not declared in project.integrations). + for (const [integrationId, fileConfig] of mergedConfigsById) { + if ( + integrations.has(integrationId) || + !(databaseIntegrationTypes as readonly string[]).includes(fileConfig.type) || + fileConfig.type === 'pandas-dataframe' + ) { + continue; + } + integrations.set(integrationId, { + config: null, + status: IntegrationStatus.Connected, + integrationName: fileConfig.name, + integrationType: fileConfig.type as ConfigurableDatabaseIntegrationType + }); + } + logger.debug(`IntegrationDetector: Found ${integrations.size} integrations`); return integrations; diff --git a/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.ts b/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.ts new file mode 100644 index 0000000000..277c606e96 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.ts @@ -0,0 +1,211 @@ +import { inject, injectable } from 'inversify'; +import { l10n, type NotebookDocument, window } from 'vscode'; + +import { IKernel, IKernelProvider } from '../../../kernels/types'; +import { IExtensionSyncActivationService } from '../../../platform/activation/types'; +import { IDisposableRegistry } from '../../../platform/common/types'; +import { noop } from '../../../platform/common/utils/misc'; +import { logger } from '../../../platform/logging'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; +import { buildIntegrationEnvRefreshSnippet, IntegrationEnvValidationError } from './integrationEnvSnippet'; +import { getStartupIntegrationEnvNames } from './startupIntegrationEnvTracker'; +import { IIntegrationEnvLiveRefresher } from './types'; + +/** How long the transient "environment updated" status-bar message stays visible. */ +const STATUS_BAR_MESSAGE_TIMEOUT_MS = 5000; + +/** + * Applies integration credential changes to already-running kernels without a restart, by executing + * the assignments directly in the kernel through the leak-safe execution primitive. + * + * Initial and post-restart delivery is not done here — `SqlIntegrationStartupCodeProvider` handles + * that as startup code. This class only needs to know *which* names that provider set, so it can + * remove the ones that later disappear from the configuration. + */ +@injectable() +export class IntegrationEnvLiveRefresher implements IIntegrationEnvLiveRefresher, IExtensionSyncActivationService { + /** Per kernel, the integration env-var names currently believed to be set in it. */ + private readonly lastSetNames = new WeakMap>(); + /** Tail of each kernel's work queue — see {@link enqueue}. */ + private readonly pendingByKernel = new WeakMap>(); + /** Monotonic per-kernel refresh counter used to drop refreshes that a newer one supersedes. */ + private readonly refreshGenerations = new WeakMap(); + + constructor( + @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider, + @inject(ISqlIntegrationEnvVarsProvider) private readonly envVarsProvider: ISqlIntegrationEnvVarsProvider, + @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry + ) {} + + public activate(): void { + // A restart fires onDidRestartKernel rather than onDidStartKernel, so both are subscribed. + // + // These events make the baseline available as early as possible, but they are not what + // guarantees it is there: `startedAtLeastOnce` — the sole gate on a refresh — is set at the top + // of `startJupyterSession`, before the session even exists, while `_onStarted` fires only at the + // very end of the start path. A refresh arriving in that window would otherwise read an empty + // baseline and emit no removals. `applyEnvToKernel` closes that by merging the startup names in + // at read time; these subscriptions are belt-and-braces on top. + // + // onDidPostInitializeKernel would be worse still — it early-returns for `disableUI` starts and + // for `ignoreTriggeringOnPostInitialized`. + this.kernelProvider.onDidStartKernel(this.seedRemovalBaseline, this, this.disposables); + this.kernelProvider.onDidRestartKernel(this.seedRemovalBaseline, this, this.disposables); + } + + public async refresh(notebooks: readonly NotebookDocument[]): Promise { + const results = await Promise.all(notebooks.map((notebook) => this.refreshNotebook(notebook))); + const refreshedCount = results.filter(Boolean).length; + + if (refreshedCount > 0) { + // Transient status-bar message rather than a persistent toast, so frequent env-file edits don't spam notifications (F2). + window.setStatusBarMessage( + l10n.t('Deepnote integration environment updated.'), + STATUS_BAR_MESSAGE_TIMEOUT_MS + ); + } + } + + /** + * Resolves this notebook's integration env and pushes the difference into the kernel. Always runs + * inside the kernel's queue, since it reads and then writes the removal baseline. + */ + private async applyEnvToKernel(kernel: IKernel, notebook: NotebookDocument): Promise { + const envVars = (await this.envVarsProvider.getEnvironmentVariables(notebook.uri)) ?? {}; + // The startup-code provider's names are merged in here, at read time, not just via the + // start/restart event: a refresh can beat that event (see `activate`), and reading an empty + // baseline would silently skip the removal of a credential the user just deleted. Repeating the + // merge on every read is idempotent, and the worst case — re-emitting `unset_env` for a name + // already removed — is a harmless `os.environ.pop(name, None)`. + const baseline = new Set([ + ...(this.lastSetNames.get(kernel) ?? []), + ...(getStartupIntegrationEnvNames(kernel) ?? []) + ]); + const hasEntries = Object.values(envVars).some((value) => typeof value === 'string'); + + // `getEnvironmentVariables` resolves to `{}` on its soft-failure paths (notebook not found, + // missing project/notebook metadata, project not cached), and a healthy open Deepnote notebook + // always yields at least the internal DuckDB variable. So an empty result against a non-empty + // baseline is a failed read, never a genuine "no integrations configured" — applying it would + // unset every tracked name and wipe the kernel's live credentials. + if (!hasEntries && baseline.size > 0) { + logger.warn( + `IntegrationEnvLiveRefresher: Resolved no integration environment for ${notebook.uri.toString()} while ${ + baseline.size + } variable(s) are set; keeping the current environment.` + ); + + return false; + } + + const { code, setNames } = buildIntegrationEnvRefreshSnippet(envVars, baseline); + + // Nothing to set and nothing to remove. Reporting success would put "environment updated" on + // the status bar for a refresh that delivered nothing — this is reached when the provider + // resolved no variables at all, which is a soft failure rather than a real empty configuration. + if (code.length === 0) { + return false; + } + + const outputs = await this.kernelProvider.getKernelExecution(kernel).executeHiddenSilent(code); + + // `outputs` is deliberately never logged: the code that produced it embeds credentials. + if (outputs.some((output) => output.output_type === 'error')) { + logger.warn( + `IntegrationEnvLiveRefresher: Failed to apply the integration environment for ${notebook.uri.toString()}` + ); + + // A snippet that raised part-way through still applied the lines before it — `stop_on_error` + // governs later queued requests, not statements within one request. Union rather than + // discard `setNames`, so a variable that did get set stays tracked and a later deletion can + // still unset it. This can only over-track, and an extra `unset_env` is harmless. + this.lastSetNames.set(kernel, new Set([...baseline, ...setNames])); + + return false; + } + + // Replace, rather than union, is correct here: every name dropped from the baseline was + // unset by this same snippet. + this.lastSetNames.set(kernel, setNames); + + return true; + } + + /** + * Serializes work per kernel: `work` starts only once everything already queued for that kernel has + * settled. Seeding and refreshing both read-modify-write the same baseline and so must not + * interleave, and two racing refreshes must not let the older provider read land last. + */ + private enqueue(kernel: IKernel, work: () => Promise): Promise { + // The stored tail never rejects, so `work` is always reached. + const previous = this.pendingByKernel.get(kernel) ?? Promise.resolve(); + const next = previous.then(work); + + this.pendingByKernel.set(kernel, next.then(noop, noop)); + + return next; + } + + /** Refreshes one kernel; never throws (per-notebook errors are logged), resolves true only on a clean run. */ + private async refreshNotebook(notebook: NotebookDocument): Promise { + try { + const kernel = this.kernelProvider.get(notebook); + if (!kernel || !kernel.startedAtLeastOnce) { + return false; + } + + const generation = (this.refreshGenerations.get(kernel) ?? 0) + 1; + this.refreshGenerations.set(kernel, generation); + + return await this.enqueue(kernel, async () => { + // Superseded before it got to start. Not a staleness guard — the provider read happens + // inside this closure, so this refresh would see the same fresh configuration as the + // one that superseded it. Skipping is pure de-duplication of identical work. + if (this.refreshGenerations.get(kernel) !== generation) { + return false; + } + + return this.applyEnvToKernel(kernel, notebook); + }); + } catch (err) { + if (err instanceof IntegrationEnvValidationError) { + // Nothing was emitted, so the kernel still holds the previous, working environment. + logger.warn( + `IntegrationEnvLiveRefresher: ${err.message} Refresh aborted for ${notebook.uri.toString()}` + ); + + return false; + } + + // The error type only, never the error itself: the logger renders an Error through + // `util.inspect`, which would emit its full stack and `[cause]`, and this path handles + // throws from the provider, whose messages are not guaranteed to be free of credentials. + logger.error( + `IntegrationEnvLiveRefresher: Failed to refresh integration env for ${notebook.uri.toString()} (${ + err instanceof Error ? err.name : 'unknown' + })` + ); + + return false; + } + } + + /** + * Adds the names the startup-code provider wrote for this kernel to its removal baseline. + * + * Unions, never replaces: over-tracking costs at most a harmless `os.environ.pop(name, None)`, + * whereas under-tracking would leave a credential the user deleted live in the kernel. + */ + private seedRemovalBaseline(kernel: IKernel): void { + const startupNames = getStartupIntegrationEnvNames(kernel); + if (!startupNames || startupNames.size === 0) { + return; + } + + // Queued, not applied inline, so it cannot land between an in-flight refresh reading the + // baseline and writing it back — which would silently discard the seeded names. + this.enqueue(kernel, async () => { + this.lastSetNames.set(kernel, new Set([...(this.lastSetNames.get(kernel) ?? []), ...startupNames])); + }).catch(noop); + } +} diff --git a/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.unit.test.ts b/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.unit.test.ts new file mode 100644 index 0000000000..5ac806db85 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.unit.test.ts @@ -0,0 +1,587 @@ +import { assert } from 'chai'; +import * as sinon from 'sinon'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; +import { Disposable, EventEmitter, NotebookDocument, Uri } from 'vscode'; + +import { IKernel, IKernelProvider, INotebookKernelExecution } from '../../../kernels/types'; +import { IDisposable } from '../../../platform/common/types'; +import { dispose } from '../../../platform/common/utils/lifecycle'; +import { logger } from '../../../platform/logging'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; +import { IntegrationEnvLiveRefresher } from './integrationEnvLiveRefresher.node'; +import { recordStartupIntegrationEnvNames } from './startupIntegrationEnvTracker'; + +const ENV_VAR_NAME = 'SQL_DEMO'; +const ENV_VAR_VALUE = 'postgres://demo'; + +const EXPECTED_NOTIFICATION = 'Deepnote integration environment updated.'; + +/** Every level the refresher could conceivably log through, so leak assertions cover all of them. */ +const LOG_LEVELS = ['error', 'warn', 'info', 'debug', 'trace', 'ci'] as const; + +suite('IntegrationEnvLiveRefresher', () => { + let refresher: IntegrationEnvLiveRefresher; + let kernelProvider: IKernelProvider; + let envVarsProvider: ISqlIntegrationEnvVarsProvider; + let executeHiddenSilentSpy: sinon.SinonStub; + let onDidStartKernel: EventEmitter; + let onDidRestartKernel: EventEmitter; + let disposables: IDisposable[]; + + setup(() => { + resetVSCodeMocks(); + disposables = [new Disposable(() => resetVSCodeMocks())]; + kernelProvider = mock(); + envVarsProvider = mock(); + + onDidStartKernel = new EventEmitter(); + onDidRestartKernel = new EventEmitter(); + disposables.push(onDidStartKernel, onDidRestartKernel); + when(kernelProvider.onDidStartKernel).thenReturn(onDidStartKernel.event); + when(kernelProvider.onDidRestartKernel).thenReturn(onDidRestartKernel.event); + + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ [ENV_VAR_NAME]: ENV_VAR_VALUE }); + + executeHiddenSilentSpy = sinon.stub().resolves([]); + when(kernelProvider.getKernelExecution(anything())).thenReturn({ + executeHiddenSilent: executeHiddenSilentSpy + } as unknown as INotebookKernelExecution); + + refresher = new IntegrationEnvLiveRefresher(instance(kernelProvider), instance(envVarsProvider), []); + }); + + teardown(() => { + sinon.restore(); + disposables = dispose(disposables); + }); + + /** A notebook whose kernel is started; `kernelProvider.get(notebook)` returns the returned kernel. */ + function createRunningKernel(uri: Uri): { notebook: NotebookDocument; kernel: IKernel } { + const notebookMock = mock(); + when(notebookMock.uri).thenReturn(uri); + const notebook = instance(notebookMock); + + const kernelMock = mock(); + when(kernelMock.startedAtLeastOnce).thenReturn(true); + const kernel = instance(kernelMock); + when(kernelProvider.get(notebook)).thenReturn(kernel); + + return { notebook, kernel }; + } + + function createRunningNotebook(uri: Uri): NotebookDocument { + return createRunningKernel(uri).notebook; + } + + /** The code passed to the nth (0-based) `executeHiddenSilent` call. */ + function executedCode(callIndex: number): string { + assert.isAbove( + executeHiddenSilentSpy.callCount, + callIndex, + `expected at least ${callIndex + 1} execution(s), saw ${executeHiddenSilentSpy.callCount}` + ); + + return executeHiddenSilentSpy.getCall(callIndex).args[0] as string; + } + + function stubAllLogLevels(): sinon.SinonStub[] { + return LOG_LEVELS.map((level) => sinon.stub(logger, level)); + } + + /** Everything passed to any stubbed logger method, flattened to one searchable string. */ + function loggedText(stubs: sinon.SinonStub[]): string { + return stubs + .flatMap((stub) => stub.getCalls()) + .flatMap((call) => call.args) + .map((arg) => (typeof arg === 'string' ? arg : JSON.stringify(arg) ?? String(arg))) + .join('\n'); + } + + test('applies the resolved env in a started kernel and shows one status-bar message', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + + await refresher.refresh([notebook]); + + assert.strictEqual(executeHiddenSilentSpy.callCount, 1, 'the env snippet should run once'); + const code = executedCode(0); + assert.include(code, `.set_env("${ENV_VAR_NAME}", "${ENV_VAR_VALUE}")`); + assert.notInclude(code, 'set_integration_env', 'the toolkit round-trip must no longer be used'); + + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).once(); + const [message] = capture(mockedVSCodeNamespaces.window.setStatusBarMessage).last(); + assert.strictEqual(message, EXPECTED_NOTIFICATION); + }); + + test('skips notebooks with no kernel and shows no status-bar message', async () => { + const notebookMock = mock(); + when(notebookMock.uri).thenReturn(Uri.file('/ws/a.deepnote')); + const notebook = instance(notebookMock); + when(kernelProvider.get(notebook)).thenReturn(undefined); + + await refresher.refresh([notebook]); + + assert.strictEqual(executeHiddenSilentSpy.callCount, 0, 'no kernel means nothing to refresh'); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).never(); + }); + + test('skips kernels that have not started and shows no status-bar message', async () => { + const notebookMock = mock(); + when(notebookMock.uri).thenReturn(Uri.file('/ws/a.deepnote')); + const notebook = instance(notebookMock); + + const kernelMock = mock(); + when(kernelMock.startedAtLeastOnce).thenReturn(false); + when(kernelProvider.get(notebook)).thenReturn(instance(kernelMock)); + + await refresher.refresh([notebook]); + + assert.strictEqual( + executeHiddenSilentSpy.callCount, + 0, + 'a kernel that has not started must not be executed against' + ); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).never(); + }); + + test('does not show a status-bar message when the env snippet produces an error output', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + executeHiddenSilentSpy.resolves([{ output_type: 'error', ename: 'error', evalue: '', traceback: [] }]); + + await refresher.refresh([notebook]); + + assert.strictEqual( + executeHiddenSilentSpy.callCount, + 1, + 'the snippet still runs, but its output signals failure' + ); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).never(); + }); + + test('shows exactly one status-bar message when multiple kernels are refreshed', async () => { + const notebookA = createRunningNotebook(Uri.file('/ws/a.deepnote')); + const notebookB = createRunningNotebook(Uri.file('/ws/b.deepnote')); + + await refresher.refresh([notebookA, notebookB]); + + assert.strictEqual(executeHiddenSilentSpy.callCount, 2, 'both started kernels are refreshed'); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).once(); + }); + + test('continues to the next notebook when one execution throws, and still notifies for the success', async () => { + const notebookA = createRunningNotebook(Uri.file('/ws/a.deepnote')); + const notebookB = createRunningNotebook(Uri.file('/ws/b.deepnote')); + executeHiddenSilentSpy.onFirstCall().rejects(new Error('kernel exploded')); + executeHiddenSilentSpy.onSecondCall().resolves([]); + + await refresher.refresh([notebookA, notebookB]); + + assert.strictEqual(executeHiddenSilentSpy.callCount, 2, 'a throw on the first must not stop the second'); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).once(); + }); + + test('refreshes kernels in parallel: both executions start before either resolves', async () => { + const notebookA = createRunningNotebook(Uri.file('/ws/a.deepnote')); + const notebookB = createRunningNotebook(Uri.file('/ws/b.deepnote')); + + // First execution resolves only once the second has been invoked; a sequential loop would deadlock (and time out). + let markSecondInvoked!: () => void; + const secondInvoked = new Promise((resolve) => (markSecondInvoked = resolve)); + executeHiddenSilentSpy.onFirstCall().callsFake(() => secondInvoked.then(() => [])); + executeHiddenSilentSpy.onSecondCall().callsFake(() => { + markSecondInvoked(); + + return Promise.resolve([]); + }); + + await refresher.refresh([notebookA, notebookB]); + + assert.strictEqual(executeHiddenSilentSpy.callCount, 2, 'both started kernels are refreshed'); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).once(); + }); + + suite('removal tracking', () => { + test('emits unset_env for a variable that disappeared since the previous refresh', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ SQL_A: '1', SQL_B: '2' }); + await refresher.refresh([notebook]); + + assert.notInclude(executedCode(0), '.unset_env(', 'nothing is tracked yet on the first refresh'); + + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ SQL_A: '1' }); + await refresher.refresh([notebook]); + + const code = executedCode(1); + assert.include(code, '.unset_env("SQL_B")', 'the removed variable must be unset in the kernel'); + assert.include(code, '.set_env("SQL_A", "1")'); + assert.notInclude(code, '.unset_env("SQL_A")'); + }); + + test('keeps tracking an already-set variable when a later push fails, so the removal is retried', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ SQL_A: '1', SQL_B: '2' }); + await refresher.refresh([notebook]); + + // SQL_B disappears, but the push fails — the kernel still holds both. + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ SQL_A: '1' }); + executeHiddenSilentSpy + .onSecondCall() + .resolves([{ output_type: 'error', ename: 'error', evalue: '', traceback: [] }]); + await refresher.refresh([notebook]); + + await refresher.refresh([notebook]); + + assert.include( + executedCode(2), + '.unset_env("SQL_B")', + 'a failed push must not make the refresher forget that SQL_B is still set' + ); + }); + + test('tracks the variables a partially-applied snippet may already have set', async () => { + // `stop_on_error: false` governs subsequent queued requests, not the statements within one + // request: an exception part-way through leaves every earlier line applied. So on failure + // the names the snippet tried to set have to stay tracked — otherwise a variable that did + // land in the kernel is invisible to the refresher and a later deletion never removes it. + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ SQL_A: '1' }); + await refresher.refresh([notebook]); + + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ + SQL_A: '1', + SQL_B: '2', + SQL_C: '3' + }); + executeHiddenSilentSpy + .onSecondCall() + .resolves([{ output_type: 'error', ename: 'error', evalue: '', traceback: [] }]); + await refresher.refresh([notebook]); + + // Both integrations are deleted again; whatever the failed snippet managed to apply must go. + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ SQL_A: '1' }); + await refresher.refresh([notebook]); + + const code = executedCode(2); + assert.include(code, '.unset_env("SQL_B")', 'SQL_B may have been applied before the error'); + assert.include(code, '.unset_env("SQL_C")', 'SQL_C may have been applied before the error'); + }); + + test('tracks each kernel separately', async () => { + const notebookA = createRunningNotebook(Uri.file('/ws/a.deepnote')); + + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ SQL_A: '1' }); + await refresher.refresh([notebookA]); + + const notebookB = createRunningNotebook(Uri.file('/ws/b.deepnote')); + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ SQL_B: '2' }); + await refresher.refresh([notebookB]); + + assert.notInclude( + executedCode(1), + '.unset_env(', + "a fresh kernel must not inherit another kernel's baseline" + ); + }); + }); + + suite('per-kernel serialization', () => { + test('does not interleave two overlapping refreshes of the same kernel', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + const events: string[] = []; + + let releaseFirst!: () => void; + const firstReleased = new Promise((resolve) => (releaseFirst = resolve)); + executeHiddenSilentSpy.onFirstCall().callsFake(() => { + events.push('start-1'); + + return firstReleased.then(() => { + events.push('end-1'); + + return []; + }); + }); + executeHiddenSilentSpy.onSecondCall().callsFake(() => { + events.push('start-2'); + + return Promise.resolve().then(() => { + events.push('end-2'); + + return []; + }); + }); + + const first = refresher.refresh([notebook]); + // Let the first refresh reach the kernel before the second is requested, so the second + // queues behind an in-flight execution rather than superseding it. + await new Promise((resolve) => setTimeout(resolve, 0)); + const second = refresher.refresh([notebook]); + releaseFirst(); + await Promise.all([first, second]); + + assert.deepStrictEqual( + events, + ['start-1', 'end-1', 'start-2', 'end-2'], + 'the second execution must not start until the first has finished' + ); + }); + + test('drops a queued refresh once a newer one supersedes it, so a stale read cannot land last', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + + // Both refreshes are requested before either reads configuration; only the newer one + // should reach the kernel, since it is the one that reads the fresher configuration. + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ [ENV_VAR_NAME]: 'newest' }); + + await Promise.all([refresher.refresh([notebook]), refresher.refresh([notebook])]); + + assert.strictEqual(executeHiddenSilentSpy.callCount, 1, 'the superseded refresh must not execute'); + assert.include(executedCode(0), `.set_env("${ENV_VAR_NAME}", "newest")`); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).once(); + }); + }); + + suite('empty-result guard', () => { + test('keeps the current environment when the provider resolves nothing but variables are set', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + const warnStub = sinon.stub(logger, 'warn'); + + await refresher.refresh([notebook]); + assert.strictEqual(executeHiddenSilentSpy.callCount, 1, 'the first refresh establishes the baseline'); + + // getEnvironmentVariables resolves to {} on its soft-failure paths, and a healthy Deepnote + // notebook always yields at least the internal DuckDB variable — so this is a failed read, + // and applying it would unset every tracked name. + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({}); + await refresher.refresh([notebook]); + + assert.strictEqual(executeHiddenSilentSpy.callCount, 1, 'an empty read must not execute anything'); + assert.strictEqual(warnStub.callCount, 1, 'the skipped refresh must be reported'); + assert.include(warnStub.firstCall.args[0], 'keeping the current environment'); + assert.notInclude(warnStub.firstCall.args[0], ENV_VAR_VALUE, 'the warning must carry no credential'); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).once(); + + // The baseline survived: a later good read still knows the variable is set in the kernel. + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ SQL_OTHER: 'x' }); + await refresher.refresh([notebook]); + + assert.include( + executedCode(1), + `.unset_env("${ENV_VAR_NAME}")`, + 'the baseline must have been preserved across the empty read' + ); + }); + + test('reports no update when the provider resolves nothing and nothing is tracked yet', async () => { + // The M4 guard above only covers a non-empty baseline. With an empty one the snippet comes + // out empty, and executing an empty snippet and reporting success would announce an + // environment update to the user that never happened — while the empty read is, as above, + // far more likely to be one of `getEnvironmentVariables`' soft failures than a real + // "no integrations configured" state. + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({}); + + await refresher.refresh([notebook]); + + assert.strictEqual(executeHiddenSilentSpy.callCount, 0, 'an empty snippet must not be sent to the kernel'); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).never(); + }); + + test('treats an undefined provider result as an empty one', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + + await refresher.refresh([notebook]); + + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve( + undefined as unknown as Record + ); + await refresher.refresh([notebook]); + + assert.strictEqual(executeHiddenSilentSpy.callCount, 1, 'an undefined read must not wipe the environment'); + }); + }); + + suite('leak safety', () => { + test('logs neither the credential nor the outputs when the snippet fails', async () => { + const secret = 'postgres://user:hunter2-token@host/db'; + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ [ENV_VAR_NAME]: secret }); + executeHiddenSilentSpy.resolves([ + { output_type: 'error', ename: 'UnicodeEncodeError', evalue: secret, traceback: [secret] } + ]); + const stubs = stubAllLogLevels(); + + await refresher.refresh([notebook]); + + const logged = loggedText(stubs); + assert.notInclude(logged, secret, 'the credential must never be logged'); + assert.notInclude(logged, 'hunter2-token'); + assert.notInclude(logged, 'set_env', 'the executed code must never be logged'); + assert.notInclude(logged, 'UnicodeEncodeError', 'the outputs object must never be logged'); + assert.notInclude(logged, 'traceback'); + assert.include(logged, 'Failed to apply the integration environment', 'the failure is still reported'); + }); + + test('logs no credential when an entry cannot be applied and the refresh is aborted', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + + // A lone surrogate cannot be assigned to os.environ, so the snippet builder fails closed. + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ + SQL_BROKEN: 'secret-\uD800-value' + }); + const stubs = stubAllLogLevels(); + + await refresher.refresh([notebook]); + + assert.strictEqual(executeHiddenSilentSpy.callCount, 0, 'nothing may be executed for an invalid entry'); + + const logged = loggedText(stubs); + assert.notInclude(logged, 'secret-', 'the rejected value must never be logged'); + assert.notInclude(logged, 'SQL_BROKEN', 'not even the variable name may be logged'); + assert.include(logged, 'Refresh aborted'); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).never(); + }); + + test('logs no credential when the provider itself rejects', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + + when(envVarsProvider.getEnvironmentVariables(anything())).thenReject(new Error('config unreadable')); + const stubs = stubAllLogLevels(); + + await refresher.refresh([notebook]); + + assert.strictEqual(executeHiddenSilentSpy.callCount, 0); + assert.notInclude(loggedText(stubs), ENV_VAR_VALUE); + }); + }); + + suite('removal baseline seeding', () => { + test('uses the startup names even when the start event has not been delivered yet', async () => { + // `kernel.ts` latches `_startedAtLeastOnce` — the only gate on a refresh — well before it + // fires `_onStarted`, which is what drives seeding. A refresh landing in that window would + // otherwise read an empty baseline and silently skip the removal of a just-deleted + // credential, so the startup names have to be consulted when the baseline is read, not + // only when the event arrives. + const { notebook, kernel } = createRunningKernel(Uri.file('/ws/a.deepnote')); + recordStartupIntegrationEnvNames(kernel, ['SQL_FROM_STARTUP']); + + refresher.activate(); + // Deliberately no onDidStartKernel.fire(kernel): this is the pre-event window. + + await refresher.refresh([notebook]); + + assert.include( + executedCode(0), + '.unset_env("SQL_FROM_STARTUP")', + 'a credential deleted between kernel start and the start event must still be removed' + ); + }); + + test('seeds the baseline when a kernel starts', async () => { + const { notebook, kernel } = createRunningKernel(Uri.file('/ws/a.deepnote')); + recordStartupIntegrationEnvNames(kernel, ['SQL_FROM_STARTUP']); + + refresher.activate(); + onDidStartKernel.fire(kernel); + + await refresher.refresh([notebook]); + + assert.include( + executedCode(0), + '.unset_env("SQL_FROM_STARTUP")', + 'the startup provider wrote this variable, so the refresher must be able to remove it' + ); + }); + + test('seeds the baseline when a kernel restarts', async () => { + // A restart fires onDidRestartKernel, not onDidStartKernel, and re-runs the startup code. + const { notebook, kernel } = createRunningKernel(Uri.file('/ws/a.deepnote')); + recordStartupIntegrationEnvNames(kernel, ['SQL_FROM_RESTART']); + + refresher.activate(); + onDidRestartKernel.fire(kernel); + + await refresher.refresh([notebook]); + + assert.include(executedCode(0), '.unset_env("SQL_FROM_RESTART")'); + }); + + test('unions the startup names into the baseline rather than replacing it', async () => { + const { notebook, kernel } = createRunningKernel(Uri.file('/ws/a.deepnote')); + + // A live refresh has already set SQL_DEMO in this kernel. + await refresher.refresh([notebook]); + + recordStartupIntegrationEnvNames(kernel, ['SQL_FROM_STARTUP']); + refresher.activate(); + onDidStartKernel.fire(kernel); + + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ SQL_NEW: 'v' }); + await refresher.refresh([notebook]); + + const code = executedCode(1); + assert.include(code, '.unset_env("SQL_FROM_STARTUP")', 'the seeded name must be tracked'); + assert.include(code, `.unset_env("${ENV_VAR_NAME}")`, 'seeding must not discard what was already tracked'); + }); + + test('keeps a name the startup provider wrote even after the configuration changed', async () => { + // The whole point of recording what the startup provider emitted: an edit landing between + // its read and any later re-read would drop SQL_EDITED, and it would then never be removed. + const { notebook, kernel } = createRunningKernel(Uri.file('/ws/a.deepnote')); + recordStartupIntegrationEnvNames(kernel, ['SQL_EDITED']); + + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ [ENV_VAR_NAME]: ENV_VAR_VALUE }); + + refresher.activate(); + onDidStartKernel.fire(kernel); + + await refresher.refresh([notebook]); + + assert.include( + executedCode(0), + '.unset_env("SQL_EDITED")', + 'a variable written at startup must stay removable even though the current config omits it' + ); + }); + + test('tracks nothing when the startup provider recorded no names for the kernel', async () => { + const { notebook, kernel } = createRunningKernel(Uri.file('/ws/a.deepnote')); + + refresher.activate(); + onDidStartKernel.fire(kernel); + + await refresher.refresh([notebook]); + + assert.notInclude(executedCode(0), '.unset_env('); + }); + + test('does not interleave seeding with an in-flight refresh', async () => { + const { notebook, kernel } = createRunningKernel(Uri.file('/ws/a.deepnote')); + recordStartupIntegrationEnvNames(kernel, ['SQL_FROM_STARTUP']); + refresher.activate(); + + let releaseFirst!: () => void; + const firstReleased = new Promise((resolve) => (releaseFirst = resolve)); + executeHiddenSilentSpy.onFirstCall().callsFake(() => firstReleased.then(() => [])); + + const inFlight = refresher.refresh([notebook]); + await new Promise((resolve) => setTimeout(resolve, 0)); + + // Seeds while the refresh is between reading and writing the baseline; queueing is what + // stops the refresh's write from discarding the seeded name. + onDidStartKernel.fire(kernel); + releaseFirst(); + await inFlight; + + when(envVarsProvider.getEnvironmentVariables(anything())).thenResolve({ SQL_NEW: 'v' }); + await refresher.refresh([notebook]); + + const code = executedCode(1); + assert.include(code, '.unset_env("SQL_FROM_STARTUP")', 'the seeded name must survive the refresh'); + assert.include(code, `.unset_env("${ENV_VAR_NAME}")`); + }); + }); +}); diff --git a/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.ts b/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.ts new file mode 100644 index 0000000000..c4806e87cf --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.ts @@ -0,0 +1,40 @@ +import { inject, injectable } from 'inversify'; +import { workspace } from 'vscode'; + +import { IExtensionSyncActivationService } from '../../../platform/activation/types'; +import { IDisposableRegistry } from '../../../platform/common/types'; +import { logger } from '../../../platform/logging'; +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; +import { IIntegrationEnvLiveRefresher, IIntegrationStorage } from './types'; + +/** Live-refreshes integration env in open Deepnote kernels when integration configs change (no restart). */ +@injectable() +export class IntegrationEnvRefreshHandler implements IExtensionSyncActivationService { + constructor( + @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, + @inject(IIntegrationEnvLiveRefresher) private readonly liveRefresher: IIntegrationEnvLiveRefresher, + @inject(IDisposableRegistry) disposables: IDisposableRegistry + ) { + logger.info('IntegrationEnvRefreshHandler: Initialized'); + + disposables.push( + this.integrationStorage.onDidChangeIntegrations(() => { + this.onIntegrationConfigurationChanged().catch((err) => + logger.error('IntegrationEnvRefreshHandler: Failed to handle integration change', err) + ); + }) + ); + } + + public activate(): void { + // Service is activated via constructor + } + + private async onIntegrationConfigurationChanged(): Promise { + const notebooks = workspace.notebookDocuments.filter( + (notebook) => notebook.notebookType === DEEPNOTE_NOTEBOOK_TYPE + ); + + await this.liveRefresher.refresh(notebooks); + } +} diff --git a/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.unit.test.ts b/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.unit.test.ts new file mode 100644 index 0000000000..d7dcb3ace7 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.unit.test.ts @@ -0,0 +1,125 @@ +import { assert } from 'chai'; +import * as sinon from 'sinon'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; +import { Disposable, EventEmitter, NotebookDocument, Uri } from 'vscode'; + +import { IDisposable } from '../../../platform/common/types'; +import { dispose } from '../../../platform/common/utils/lifecycle'; +import { logger } from '../../../platform/logging'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; +import { IIntegrationEnvLiveRefresher, IIntegrationStorage } from './types'; +import { IntegrationEnvRefreshHandler } from './integrationEnvRefreshHandler'; +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; + +suite('IntegrationEnvRefreshHandler', () => { + let handler: IntegrationEnvRefreshHandler; + let integrationStorage: IIntegrationStorage; + let liveRefresher: IIntegrationEnvLiveRefresher; + let disposables: IDisposable[]; + let onDidChangeIntegrations: EventEmitter; + + setup(() => { + resetVSCodeMocks(); + disposables = [new Disposable(() => resetVSCodeMocks())]; + integrationStorage = mock(); + liveRefresher = mock(); + onDidChangeIntegrations = new EventEmitter(); + disposables.push(onDidChangeIntegrations); + + when(integrationStorage.onDidChangeIntegrations).thenReturn(onDidChangeIntegrations.event); + when(liveRefresher.refresh(anything())).thenResolve(); + + handler = new IntegrationEnvRefreshHandler(instance(integrationStorage), instance(liveRefresher), disposables); + }); + + teardown(() => { + sinon.restore(); + disposables = dispose(disposables); + }); + + function createMockNotebook(notebookType: string, uri: Uri): NotebookDocument { + const notebook = mock(); + when(notebook.notebookType).thenReturn(notebookType); + when(notebook.uri).thenReturn(uri); + + return instance(notebook); + } + + test('refreshes integration env for open Deepnote notebooks when integrations change', async () => { + const notebook = createMockNotebook(DEEPNOTE_NOTEBOOK_TYPE, Uri.file('/test.deepnote')); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [notebook]); + }); + + test('passes only Deepnote notebooks to the refresher', async () => { + const deepnote = createMockNotebook(DEEPNOTE_NOTEBOOK_TYPE, Uri.file('/a.deepnote')); + const jupyter = createMockNotebook('jupyter-notebook', Uri.file('/b.ipynb')); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([deepnote, jupyter]); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [deepnote]); + }); + + test('refreshes multiple Deepnote notebooks in a single call', async () => { + const notebook1 = createMockNotebook(DEEPNOTE_NOTEBOOK_TYPE, Uri.file('/test1.deepnote')); + const notebook2 = createMockNotebook(DEEPNOTE_NOTEBOOK_TYPE, Uri.file('/test2.deepnote')); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [notebook1, notebook2]); + }); + + test('refreshes with an empty list when no Deepnote notebooks are open', async () => { + const jupyter = createMockNotebook('jupyter-notebook', Uri.file('/b.ipynb')); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([jupyter]); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], []); + }); + + test('catches and logs a rejected refresh instead of propagating it', async () => { + const notebook = createMockNotebook(DEEPNOTE_NOTEBOOK_TYPE, Uri.file('/test.deepnote')); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(liveRefresher.refresh(anything())).thenReject(new Error('refresh boom')); + const errorStub = sinon.stub(logger, 'error'); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + assert.strictEqual(errorStub.callCount, 1, 'the fire-and-forget rejection must be caught and logged'); + }); +}); diff --git a/src/notebooks/deepnote/integrations/integrationEnvSnippet.ts b/src/notebooks/deepnote/integrations/integrationEnvSnippet.ts new file mode 100644 index 0000000000..a3f471390a --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvSnippet.ts @@ -0,0 +1,98 @@ +import type { EnvironmentVariables } from '../../../platform/common/variables/types'; + +/** + * The toolkit's own env accessors, reached through `__import__` so the snippet binds no name in the + * kernel's user namespace. They must be used instead of a raw `os.environ` write: the toolkit's + * `get_env` — which SQL execution resolves credentials through — consults its in-memory `_STATE` map + * *before* `os.environ`, so a raw write is shadowed by any existing `_STATE` entry. `set_env` / + * `unset_env` update both stores, and both return `None`, so no line is ever a trailing expression + * whose value could be echoed. + */ +const SET_ENV_CALL = '__import__("deepnote_toolkit.env", fromlist=["set_env"]).set_env'; +const UNSET_ENV_CALL = '__import__("deepnote_toolkit.env", fromlist=["unset_env"]).unset_env'; + +/** + * Matches an unpaired surrogate. Under the `u` flag a well-formed pair is a single non-surrogate code + * point, so this only fires on lone surrogates — which `JSON.stringify` escapes to `\uXXXX` and Python + * decodes back into a surrogate that raises `UnicodeEncodeError` on assignment to `os.environ`. + */ +const LONE_SURROGATE_PATTERN = /\p{Surrogate}/u; + +/** + * Thrown by {@link buildIntegrationEnvRefreshSnippet} when an entry cannot be expressed as a kernel + * environment variable. Carries a count only — never a name or value — so it stays safe to log. + */ +export class IntegrationEnvValidationError extends Error { + constructor(public readonly invalidCount: number) { + super(`${invalidCount} integration environment variable(s) cannot be applied to a kernel environment.`); + this.name = 'IntegrationEnvValidationError'; + } +} + +/** + * Builds the Python snippet that brings a live kernel's integration environment in line with + * `envVars`, removing anything in `previousNames` that is no longer present. + * + * The result is meant for `executeHiddenSilent` only — it embeds credential values as literals. + * `setNames` is the set of names the snippet assigns, i.e. the caller's new removal baseline, and is + * only valid once the snippet has run successfully. + * + * @throws {IntegrationEnvValidationError} if any entry to be set cannot be represented as a kernel + * environment variable. Validation runs before a single line is emitted so a refresh is + * all-or-nothing: the alternative is a snippet that raises part-way through and leaves the kernel + * with a half-updated environment. + */ +export function buildIntegrationEnvRefreshSnippet( + envVars: EnvironmentVariables, + previousNames: Iterable +): { code: string; setNames: Set } { + // Filtered on `typeof value === 'string'`, so empty strings are kept. This matches the toolkit's + // own `if value is not None` guard; a truthiness check would silently skip `""` and leave the + // previous value of that variable live in the kernel. + const entries = Object.entries(envVars).filter((entry): entry is [string, string] => typeof entry[1] === 'string'); + + const invalidCount = entries.filter(([name, value]) => !isApplicableName(name) || !isApplicableValue(value)).length; + + if (invalidCount > 0) { + throw new IntegrationEnvValidationError(invalidCount); + } + + const setNames = new Set(entries.map(([name]) => name)); + // Unapplicable names could never have been set in the first place (the assignment would have + // raised), so there is nothing to remove for them — drop rather than fail the whole refresh. + const unsetNames = [...new Set(previousNames)].filter((name) => !setNames.has(name) && isApplicableName(name)); + + // Removals first, so a name that moved from removed to set ends up set. + const lines = [ + ...unsetNames.map((name) => `${UNSET_ENV_CALL}(${toPythonStringLiteral(name)})`), + ...entries.map( + ([name, value]) => `${SET_ENV_CALL}(${toPythonStringLiteral(name)}, ${toPythonStringLiteral(value)})` + ) + ]; + + return { code: lines.join('\n'), setNames }; +} + +/** + * Names are additionally rejected when empty or containing `=`, both of which CPython refuses (or + * mis-handles) when writing to the process environment. + */ +function isApplicableName(name: string): boolean { + return name.length > 0 && !name.includes('=') && isApplicableValue(name); +} + +/** A NUL byte raises `ValueError` and a lone surrogate raises `UnicodeEncodeError` on assignment. */ +function isApplicableValue(value: string): boolean { + return !value.includes('\0') && !LONE_SURROGATE_PATTERN.test(value); +} + +/** + * JSON string escaping is a subset of Python's: it escapes only `"`, `\`, `\b`, `\f`, `\n`, `\r`, + * `\t` and `\uXXXX` for the remaining control characters — all of which Python decodes identically — + * and leaves non-ASCII (including emoji) as literal UTF-8. So the literal round-trips to the same + * Python `str`. The name is encoded the same way, unlike `SqlIntegrationStartupCodeProvider`, which + * interpolates it raw and would emit broken code for a name containing a quote or a backslash. + */ +function toPythonStringLiteral(value: string): string { + return JSON.stringify(value); +} diff --git a/src/notebooks/deepnote/integrations/integrationEnvSnippet.unit.test.ts b/src/notebooks/deepnote/integrations/integrationEnvSnippet.unit.test.ts new file mode 100644 index 0000000000..ae467eea8e --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvSnippet.unit.test.ts @@ -0,0 +1,298 @@ +import { assert } from 'chai'; + +import { buildIntegrationEnvRefreshSnippet, IntegrationEnvValidationError } from './integrationEnvSnippet'; + +const SET_CALL = '__import__("deepnote_toolkit.env", fromlist=["set_env"]).set_env'; +const UNSET_CALL = '__import__("deepnote_toolkit.env", fromlist=["unset_env"]).unset_env'; + +/** + * `unset_env` ends with the substring `set_env`, so every check here must be anchored on the leading + * dot of the attribute access (`.set_env(` / `.unset_env(`) to tell the two call kinds apart. + */ +function isSetLine(line: string): boolean { + return line.includes('.set_env('); +} + +function isUnsetLine(line: string): boolean { + return line.includes('.unset_env('); +} + +function linesOf(code: string): string[] { + return code.length === 0 ? [] : code.split('\n'); +} + +/** + * Extracts the raw source of the value argument of a `set_env(, )` call, given the name + * it was built from. Anchored on the fully-known prefix rather than by splitting on `, `, which a + * value is free to contain. + */ +function valueLiteralOf(line: string, name: string): string { + const prefix = `${SET_CALL}(${JSON.stringify(name)}, `; + + assert.isTrue(line.startsWith(prefix), `expected ${line} to start with ${prefix}`); + assert.strictEqual(line[line.length - 1], ')', `expected ${line} to end the call`); + + return line.slice(prefix.length, -1); +} + +suite('buildIntegrationEnvRefreshSnippet', () => { + test('emits one set_env call per variable, through __import__ and with a JSON-encoded name and value', () => { + const { code } = buildIntegrationEnvRefreshSnippet({ SQL_DEMO: 'postgres://demo' }, []); + + assert.strictEqual(code, `${SET_CALL}("SQL_DEMO", "postgres://demo")`); + }); + + test('encodes the name as a literal too, so a quote or backslash in it cannot break the snippet', () => { + // SqlIntegrationStartupCodeProvider interpolates the key raw and would emit broken Python here. + const { code } = buildIntegrationEnvRefreshSnippet({ 'ODD"NAME\\': 'v' }, []); + + assert.strictEqual(code, `${SET_CALL}("ODD\\"NAME\\\\", "v")`); + assert.strictEqual(linesOf(code).length, 1); + }); + + suite('value encoding', () => { + /** + * The builder relies on JSON string escaping being a subset of Python's: every escape it can + * emit (`\\" \\\\ \\b \\f \\n \\r \\t` and `\\uXXXX`) decodes identically in Python, and + * non-ASCII is left as literal UTF-8. Running a real interpreter would make these tests + * depend on a Python install, so the round-trip is asserted on the JSON side — the literal + * must parse back to the exact original string — plus the structural property that no value + * can break out of its line. + */ + function assertValueRoundTrips(value: string, description: string) { + const { code } = buildIntegrationEnvRefreshSnippet({ NAME: value }, []); + const lines = linesOf(code); + + assert.strictEqual(lines.length, 1, `${description}: a value must never split the snippet across lines`); + assert.strictEqual( + JSON.parse(valueLiteralOf(lines[0], 'NAME')), + value, + `${description}: literal must decode to the original` + ); + } + + test('round-trips a double quote', () => { + assertValueRoundTrips('pa"ss', 'double quote'); + }); + + test('round-trips a backslash', () => { + assertValueRoundTrips('C:\\keys\\db', 'backslash'); + }); + + test('round-trips a backslash immediately before a quote', () => { + assertValueRoundTrips('trailing\\', 'trailing backslash'); + }); + + test('round-trips a newline without splitting the line', () => { + const key = '-----BEGIN-----\nabc\n-----END-----'; + const { code } = buildIntegrationEnvRefreshSnippet({ KEY: key }, []); + + assert.strictEqual(linesOf(code).length, 1, 'a newline in a value must be escaped, not emitted raw'); + assert.include(code, '\\n'); + assert.strictEqual(JSON.parse(valueLiteralOf(code, 'KEY')), key); + }); + + test('round-trips tabs and other control characters', () => { + assertValueRoundTrips('a\tb\rc', 'control characters'); + }); + + test('round-trips accented characters, CJK and emoji, leaving them as literal UTF-8', () => { + assertValueRoundTrips('naïve 世界 😀', 'non-ASCII'); + + const { code } = buildIntegrationEnvRefreshSnippet({ NAME: 'naïve 世界 😀' }, []); + + // Not \uXXXX-escaped: the snippet is sent to the kernel as UTF-8 already. + assert.include(code, 'naïve 世界 😀'); + }); + + test('round-trips a well-formed surrogate pair, which must not be mistaken for a lone surrogate', () => { + // The same emoji, written as the two code units it is stored as, to pin down that the + // lone-surrogate rejection below does not fire on a properly paired one. + assertValueRoundTrips('😀', 'surrogate pair'); + }); + }); + + suite('removals', () => { + test('emits unset_env for a name that is no longer present', () => { + const { code } = buildIntegrationEnvRefreshSnippet({ KEPT: '1' }, ['KEPT', 'REMOVED']); + + assert.include(code, `${UNSET_CALL}("REMOVED")`); + assert.notInclude(code, '.unset_env("KEPT")', 'a name that is still set must not be unset'); + }); + + test('emits every unset before any set', () => { + const { code } = buildIntegrationEnvRefreshSnippet({ A: '1', B: '2' }, ['GONE_1', 'GONE_2']); + const lines = linesOf(code); + + assert.strictEqual(lines.length, 4); + + const lastUnset = lines.map(isUnsetLine).lastIndexOf(true); + const firstSet = lines.findIndex(isSetLine); + + assert.isAbove(firstSet, lastUnset, 'a removal must never run after a set, or it would undo it'); + assert.deepStrictEqual(lines.map(isUnsetLine), [true, true, false, false]); + }); + + test('sets, rather than removes, a name that reappears in the new environment', () => { + const { code } = buildIntegrationEnvRefreshSnippet({ REVIVED: 'new' }, ['REVIVED']); + + assert.strictEqual(code, `${SET_CALL}("REVIVED", "new")`); + }); + + test('deduplicates repeated previous names', () => { + const { code } = buildIntegrationEnvRefreshSnippet({}, ['GONE', 'GONE']); + + assert.strictEqual(code, `${UNSET_CALL}("GONE")`); + }); + + test('accepts a Set as the previous names', () => { + const { code } = buildIntegrationEnvRefreshSnippet({}, new Set(['GONE'])); + + assert.strictEqual(code, `${UNSET_CALL}("GONE")`); + }); + + test('emits nothing when there is neither anything to set nor anything to remove', () => { + const { code, setNames } = buildIntegrationEnvRefreshSnippet({}, []); + + assert.strictEqual(code, ''); + assert.deepStrictEqual([...setNames], []); + }); + }); + + suite('entry filtering', () => { + test('keeps empty-string values', () => { + // Deliberately unlike SqlIntegrationStartupCodeProvider's `if (value)` guard: skipping '' + // would leave the variable's previous, stale value live in the kernel. + const { code, setNames } = buildIntegrationEnvRefreshSnippet({ EMPTY: '' }, []); + + assert.strictEqual(code, `${SET_CALL}("EMPTY", "")`); + assert.deepStrictEqual([...setNames], ['EMPTY']); + }); + + test('skips undefined values', () => { + const { code, setNames } = buildIntegrationEnvRefreshSnippet({ SET: 'v', MISSING: undefined }, []); + + assert.strictEqual(code, `${SET_CALL}("SET", "v")`); + assert.deepStrictEqual([...setNames], ['SET']); + }); + + test('skips non-string values', () => { + const envVars = { NUMERIC: 42, NULLED: null, SET: 'v' } as unknown as Record; + const { code, setNames } = buildIntegrationEnvRefreshSnippet(envVars, []); + + assert.strictEqual(code, `${SET_CALL}("SET", "v")`); + assert.deepStrictEqual([...setNames], ['SET']); + }); + + test('removes a previously set name whose value is now undefined', () => { + const { code, setNames } = buildIntegrationEnvRefreshSnippet({ DROPPED: undefined }, ['DROPPED']); + + assert.strictEqual(code, `${UNSET_CALL}("DROPPED")`); + assert.deepStrictEqual([...setNames], []); + }); + + test('setNames is exactly the set of included names', () => { + const { setNames } = buildIntegrationEnvRefreshSnippet({ A: 'a', EMPTY: '', SKIPPED: undefined, B: 'b' }, [ + 'OLD' + ]); + + assert.deepStrictEqual(setNames, new Set(['A', 'EMPTY', 'B'])); + }); + }); + + test('binds no name in the user namespace', () => { + const { code } = buildIntegrationEnvRefreshSnippet({ A: 'a', B: 'b' }, ['GONE']); + const lines = linesOf(code); + + for (const line of lines) { + // Every statement is a bare call expression reached through __import__, so nothing — not + // the module, not a value — is left bound in the kernel's user namespace. + assert.isTrue( + line.startsWith(SET_CALL + '(') || line.startsWith(UNSET_CALL + '('), + `unexpected statement shape: ${line}` + ); + } + + assert.notMatch(code, /^\s*import\s/m, 'a plain import would bind the module name'); + assert.notMatch(code, /^\s*from\s+\S+\s+import\s/m, 'a from-import would bind a name'); + }); + + suite('fail-closed validation', () => { + /** Asserts the builder throws and produces no code at all, rather than a half-applied snippet. */ + function assertRejects( + envVars: Record, + previousNames: Iterable, + expectedCount: number + ) { + let result: { code: string } | undefined; + let thrown: unknown; + + try { + result = buildIntegrationEnvRefreshSnippet(envVars, previousNames); + } catch (err) { + thrown = err; + } + + assert.isUndefined(result, 'no snippet — not even a partial one — may be produced'); + assert.instanceOf(thrown, IntegrationEnvValidationError); + assert.strictEqual((thrown as IntegrationEnvValidationError).invalidCount, expectedCount); + + return thrown as IntegrationEnvValidationError; + } + + test('rejects a lone high surrogate value', () => { + // Python decodes \uD800 back to a surrogate, which raises UnicodeEncodeError on assignment. + assertRejects({ BAD: 'prefix\uD800suffix' }, [], 1); + }); + + test('rejects a lone low surrogate value', () => { + assertRejects({ BAD: '\uDC00' }, [], 1); + }); + + test('rejects a NUL in a value', () => { + assertRejects({ BAD: 'a\0b' }, [], 1); + }); + + test('rejects a name containing =', () => { + assertRejects({ 'BAD=NAME': 'v' }, [], 1); + }); + + test('rejects a name containing NUL', () => { + assertRejects({ 'BAD\0NAME': 'v' }, [], 1); + }); + + test('rejects an empty name', () => { + assertRejects({ '': 'v' }, [], 1); + }); + + test('aborts the whole refresh, emitting nothing for the valid entries or the removals', () => { + // All-or-nothing: removals are emitted first, so a snippet that raised part-way through + // would leave the kernel with removals applied and the replacements missing. + assertRejects({ GOOD: 'v', BAD: '\uD800' }, ['GONE'], 1); + }); + + test('counts every invalid entry', () => { + assertRejects({ GOOD: 'v', BAD_1: '\uD800', BAD_2: 'a\0b', 'BAD=3': 'v' }, [], 3); + }); + + test('carries no name and no value in its message, so the refresher can log it safely', () => { + const name = 'SQL_SECRET_NAME'; + const value = 'super-secret-\uD800-value'; + const error = assertRejects({ [name]: value }, [], 1); + + assert.notInclude(error.message, name); + assert.notInclude(error.message, 'super-secret'); + assert.notInclude(error.message, value); + assert.include(error.message, '1'); + assert.strictEqual(error.name, 'IntegrationEnvValidationError'); + }); + + test('drops, rather than rejects, an unapplicable previous name', () => { + // Such a name could never have been set, so there is nothing to remove and no reason to + // fail an otherwise valid refresh. + const { code } = buildIntegrationEnvRefreshSnippet({ GOOD: 'v' }, ['BAD=NAME', 'GONE']); + + assert.strictEqual(code, `${UNSET_CALL}("GONE")\n${SET_CALL}("GOOD", "v")`); + }); + }); +}); diff --git a/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.ts b/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.ts deleted file mode 100644 index b082d0d192..0000000000 --- a/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.ts +++ /dev/null @@ -1,155 +0,0 @@ -import { inject, injectable } from 'inversify'; -import { l10n, NotebookDocument, workspace, window } from 'vscode'; - -import { IDisposableRegistry } from '../../../platform/common/types'; -import { IExtensionSyncActivationService } from '../../../platform/activation/types'; -import { logger } from '../../../platform/logging'; -import { IIntegrationStorage } from './types'; -import { IKernelProvider } from '../../../kernels/types'; -import { DATAFRAME_SQL_INTEGRATION_ID } from '../../../platform/notebooks/deepnote/integrationTypes'; - -/** - * Handles automatic kernel restart when integration configurations change. - * When a user saves/deletes an integration config, this service restarts all kernels - * that are using that integration so they pick up the new credentials. - */ -@injectable() -export class IntegrationKernelRestartHandler implements IExtensionSyncActivationService { - private isRestarting = false; - - constructor( - @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, - @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider, - @inject(IDisposableRegistry) disposables: IDisposableRegistry - ) { - logger.info('IntegrationKernelRestartHandler: Initialized'); - - // Listen for integration configuration changes - disposables.push( - this.integrationStorage.onDidChangeIntegrations(() => { - this.onIntegrationConfigurationChanged().catch((err) => - logger.error('IntegrationKernelRestartHandler: Failed to handle integration change', err) - ); - }) - ); - } - - public activate(): void { - // Service is activated via constructor - } - - /** - * Handle integration configuration changes by restarting affected kernels - */ - private async onIntegrationConfigurationChanged(): Promise { - // Prevent multiple simultaneous restart attempts - if (this.isRestarting) { - logger.debug('IntegrationKernelRestartHandler: Already restarting, skipping'); - return; - } - - try { - this.isRestarting = true; - - logger.info( - 'IntegrationKernelRestartHandler: Integration configuration changed, checking for affected kernels' - ); - - // Find all Deepnote notebooks with running kernels that use SQL integrations - const notebooksToRestart: NotebookDocument[] = []; - - for (const notebook of workspace.notebookDocuments) { - // Only process Deepnote notebooks - if (notebook.notebookType !== 'deepnote') { - continue; - } - - // Check if kernel is running - const kernel = this.kernelProvider.get(notebook); - if (!kernel || !kernel.startedAtLeastOnce) { - continue; - } - - // Check if notebook uses SQL integrations - const usesIntegrations = this.notebookUsesSqlIntegrations(notebook); - if (usesIntegrations) { - notebooksToRestart.push(notebook); - } - } - - if (notebooksToRestart.length === 0) { - logger.info( - 'IntegrationKernelRestartHandler: No running kernels use SQL integrations, no restart needed' - ); - return; - } - - logger.info( - `IntegrationKernelRestartHandler: Found ${notebooksToRestart.length} notebook(s) with kernels that need restart` - ); - - // Restart kernels for affected notebooks - const restartPromises = notebooksToRestart.map(async (notebook) => { - const kernel = this.kernelProvider.get(notebook); - if (kernel) { - try { - logger.info( - `IntegrationKernelRestartHandler: Restarting kernel for notebook: ${notebook.uri.toString()}` - ); - await kernel.restart(); - logger.info( - `IntegrationKernelRestartHandler: Successfully restarted kernel for: ${notebook.uri.toString()}` - ); - } catch (error) { - logger.error( - `IntegrationKernelRestartHandler: Failed to restart kernel for ${notebook.uri.toString()}`, - error - ); - // Don't throw - we want to continue restarting other kernels - } - } - }); - - await Promise.all(restartPromises); - - // Show a notification to the user - if (notebooksToRestart.length === 1) { - void window.showInformationMessage( - l10n.t('Integration configuration updated. Kernel restarted to apply changes.') - ); - } else { - void window.showInformationMessage( - l10n.t( - 'Integration configuration updated. {0} kernels restarted to apply changes.', - notebooksToRestart.length - ) - ); - } - } finally { - this.isRestarting = false; - } - } - - /** - * Check if a notebook uses SQL integrations by scanning cells for sql_integration_id metadata - */ - private notebookUsesSqlIntegrations(notebook: NotebookDocument): boolean { - for (const cell of notebook.getCells()) { - // Check for SQL cells - if (cell.document.languageId !== 'sql') { - continue; - } - - const metadata = cell.metadata; - if (metadata && typeof metadata === 'object') { - const integrationId = (metadata as Record).sql_integration_id; - if (typeof integrationId === 'string' && integrationId !== DATAFRAME_SQL_INTEGRATION_ID) { - // Found a SQL cell with an external integration - return true; - } - } - } - - return false; - } -} diff --git a/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.unit.test.ts b/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.unit.test.ts deleted file mode 100644 index 62888c3fea..0000000000 --- a/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.unit.test.ts +++ /dev/null @@ -1,243 +0,0 @@ -import { anything, instance, mock, verify, when } from 'ts-mockito'; -import { Disposable, EventEmitter, NotebookCell, NotebookDocument, TextDocument, Uri } from 'vscode'; - -import { IDisposable } from '../../../platform/common/types'; -import { dispose } from '../../../platform/common/utils/lifecycle'; -import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; -import { IIntegrationStorage } from './types'; -import { IKernel, IKernelProvider } from '../../../kernels/types'; -import { IntegrationKernelRestartHandler } from './integrationKernelRestartHandler'; -import { DATAFRAME_SQL_INTEGRATION_ID } from '../../../platform/notebooks/deepnote/integrationTypes'; - -suite('IntegrationKernelRestartHandler', () => { - let handler: IntegrationKernelRestartHandler; - let integrationStorage: IIntegrationStorage; - let kernelProvider: IKernelProvider; - let disposables: IDisposable[]; - let onDidChangeIntegrations: EventEmitter; - - setup(() => { - resetVSCodeMocks(); - disposables = [new Disposable(() => resetVSCodeMocks())]; - integrationStorage = mock(); - kernelProvider = mock(); - onDidChangeIntegrations = new EventEmitter(); - disposables.push(onDidChangeIntegrations); - - when(integrationStorage.onDidChangeIntegrations).thenReturn(onDidChangeIntegrations.event); - - handler = new IntegrationKernelRestartHandler( - instance(integrationStorage), - instance(kernelProvider), - disposables - ); - }); - - teardown(() => { - disposables = dispose(disposables); - }); - - function createMockNotebook( - notebookType: string, - uri: Uri, - cells: { languageId: string; metadata?: Record }[] - ): NotebookDocument { - const notebook = mock(); - const mockCells: NotebookCell[] = cells.map((cellConfig, index) => { - const cell = mock(); - const doc = mock(); - when(doc.languageId).thenReturn(cellConfig.languageId); - when(cell.document).thenReturn(instance(doc)); - when(cell.metadata).thenReturn(cellConfig.metadata || {}); - when(cell.index).thenReturn(index); - return instance(cell); - }); - - when(notebook.notebookType).thenReturn(notebookType); - when(notebook.uri).thenReturn(uri); - when(notebook.getCells()).thenReturn(mockCells); - - return instance(notebook); - } - - test('restarts kernel when integration changes and notebook uses SQL integrations', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).once(); - }); - - test('does not restart kernel for non-Deepnote notebooks', async () => { - const notebook = createMockNotebook('jupyter-notebook', Uri.file('/test.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).never(); - }); - - test('does not restart kernel that has not started', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(false); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).never(); - }); - - test('does not restart kernel when notebook has no SQL integrations', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'python', metadata: {} } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).never(); - }); - - test('does not restart kernel when notebook only uses internal DuckDB integration', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: DATAFRAME_SQL_INTEGRATION_ID } } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).never(); - }); - - test('restarts multiple kernels in parallel', async () => { - const notebook1 = createMockNotebook('deepnote', Uri.file('/test1.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const notebook2 = createMockNotebook('deepnote', Uri.file('/test2.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'bigquery-1' } } - ]); - const mockKernel1 = mock(); - const mockKernel2 = mock(); - when(mockKernel1.startedAtLeastOnce).thenReturn(true); - when(mockKernel1.restart()).thenResolve(); - when(mockKernel2.startedAtLeastOnce).thenReturn(true); - when(mockKernel2.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - when(kernelProvider.get(notebook1)).thenReturn(instance(mockKernel1)); - when(kernelProvider.get(notebook2)).thenReturn(instance(mockKernel2)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel1.restart()).once(); - verify(mockKernel2.restart()).once(); - }); - - test('continues restarting other kernels when one fails', async () => { - const notebook1 = createMockNotebook('deepnote', Uri.file('/test1.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const notebook2 = createMockNotebook('deepnote', Uri.file('/test2.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'bigquery-1' } } - ]); - const mockKernel1 = mock(); - const mockKernel2 = mock(); - when(mockKernel1.startedAtLeastOnce).thenReturn(true); - when(mockKernel1.restart()).thenReject(new Error('Restart failed')); - when(mockKernel2.startedAtLeastOnce).thenReturn(true); - when(mockKernel2.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - when(kernelProvider.get(notebook1)).thenReturn(instance(mockKernel1)); - when(kernelProvider.get(notebook2)).thenReturn(instance(mockKernel2)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel1.restart()).once(); - verify(mockKernel2.restart()).once(); - }); - - test('handles notebooks with mixed cell types', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'python', metadata: {} }, - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } }, - { languageId: 'markdown', metadata: {} } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).once(); - }); - - test('does not restart when no notebooks are open', async () => { - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(kernelProvider.get(anything())).never(); - }); -}); diff --git a/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.ts b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.ts new file mode 100644 index 0000000000..a0e237ca86 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.ts @@ -0,0 +1,193 @@ +import { inject, injectable } from 'inversify'; +import { NotebookDocument, RelativePattern, Uri, workspace } from 'vscode'; +import { DEFAULT_ENV_FILE, DEFAULT_INTEGRATIONS_FILE } from '@deepnote/database-integrations'; + +import { IExtensionSyncActivationService } from '../../../platform/activation/types'; +import { IFileSystem } from '../../../platform/common/platform/types'; +import { IDisposableRegistry } from '../../../platform/common/types'; +import { notebookPathToDeepnoteProjectFilePath } from '../../../platform/deepnote/deepnoteProjectUtils'; +import { logger } from '../../../platform/logging'; +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; +import { IIntegrationEnvLiveRefresher } from './types'; + +/** Trailing-edge debounce so a burst of edits (e.g. .env and .deepnote.env.yaml both saved) is handled once. */ +const debounceTimeInMilliseconds = 500; + +const watchedEnvFileNames = [DEFAULT_INTEGRATIONS_FILE, DEFAULT_ENV_FILE]; + +/** + * Watches `.deepnote.env.yaml` / `.env` and live-refreshes affected notebooks' kernels on change (no restart). + * Deleting `.deepnote.env.yaml` refreshes too, so the credentials it contributed are unset rather than left live. + */ +@injectable() +export class IntegrationsEnvFileWatcher implements IExtensionSyncActivationService { + private readonly changedDirs = new Set(); + private debounceTimer: ReturnType | undefined; + /** Subset of {@link changedDirs} where `.deepnote.env.yaml` itself was created, changed or deleted. */ + private readonly integrationsFileChangedDirs = new Set(); + private readonly watchedDirs = new Set(); + + constructor( + @inject(IIntegrationEnvLiveRefresher) private readonly liveRefresher: IIntegrationEnvLiveRefresher, + @inject(IFileSystem) private readonly fileSystem: IFileSystem, + @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry + ) {} + + public activate(): void { + for (const folder of workspace.workspaceFolders ?? []) { + this.watchDir(folder.uri); + } + + for (const notebook of workspace.notebookDocuments) { + this.watchNotebookDir(notebook); + } + + this.disposables.push( + workspace.onDidOpenNotebookDocument((notebook) => this.watchNotebookDir(notebook)), + workspace.onDidChangeWorkspaceFolders((event) => { + for (const folder of event.added) { + this.watchDir(folder.uri); + } + }), + { + dispose: () => { + if (this.debounceTimer) { + clearTimeout(this.debounceTimer); + this.debounceTimer = undefined; + } + } + } + ); + } + + /** Public so it can be unit-tested without real filesystem events. */ + public async handleChangedDirs( + changedDirs: Set, + integrationsFileChangedDirs: Set = new Set() + ): Promise { + const affected = await this.findAffectedNotebooks(changedDirs, integrationsFileChangedDirs); + if (affected.length === 0) { + return; + } + + await this.liveRefresher.refresh(affected); + } + + private async findAffectedNotebooks( + changedDirs: Set, + integrationsFileChangedDirs: Set + ): Promise { + const affected: NotebookDocument[] = []; + + for (const notebook of workspace.notebookDocuments) { + if (notebook.notebookType !== DEEPNOTE_NOTEBOOK_TYPE) { + continue; + } + + const deepnoteFileUri = notebookPathToDeepnoteProjectFilePath(notebook.uri); + + // Mirror IntegrationsFileConfigProvider's gate: a disabled feature must not trigger kernel refreshes. + const enabled = workspace + .getConfiguration('deepnote', deepnoteFileUri) + .get('integrations.envFile.enabled', true); + if (enabled === false) { + continue; + } + + const deepnoteDir = Uri.joinPath(deepnoteFileUri, '..'); + const workspaceRoot = workspace.getWorkspaceFolder(notebook.uri)?.uri; + + const changedInScope = + changedDirs.has(deepnoteDir.fsPath) || (workspaceRoot != null && changedDirs.has(workspaceRoot.fsPath)); + if (!changedInScope) { + continue; + } + + // An event on `.deepnote.env.yaml` itself is unambiguously ours and always refreshes, without + // consulting the filesystem. Requiring the file to exist would skip exactly the deletion case — + // the variables it contributed would stay live in the kernel, and deleting the file is the most + // direct way a user revokes them. + const integrationsFileChanged = + integrationsFileChangedDirs.has(deepnoteDir.fsPath) || + (workspaceRoot != null && integrationsFileChangedDirs.has(workspaceRoot.fsPath)); + + // A `.env` change, by contrast, only affects integration env when a `.deepnote.env.yaml` actually + // exists for this notebook; without one the refresh is a no-op and its status message misleading, + // so an unrelated `.env` (a very common non-Deepnote file) must not trigger hidden kernel + // executions (F2). + const candidateDirs = + workspaceRoot != null && workspaceRoot.fsPath !== deepnoteDir.fsPath + ? [deepnoteDir, workspaceRoot] + : [deepnoteDir]; + if (integrationsFileChanged || (await this.hasIntegrationsFile(candidateDirs))) { + affected.push(notebook); + } + } + + return affected; + } + + /** True when a `.deepnote.env.yaml` exists in any candidate dir (dir-then-root), mirroring the config provider's probe. */ + private async hasIntegrationsFile(dirs: Uri[]): Promise { + for (const dir of dirs) { + const candidate = Uri.joinPath(dir, DEFAULT_INTEGRATIONS_FILE); + if (await this.fileSystem.exists(candidate)) { + return true; + } + } + + return false; + } + + private onFileEvent(dir: Uri, fileName: string): void { + this.changedDirs.add(dir.fsPath); + + if (fileName === DEFAULT_INTEGRATIONS_FILE) { + this.integrationsFileChangedDirs.add(dir.fsPath); + } + + if (this.debounceTimer) { + clearTimeout(this.debounceTimer); + } + + this.debounceTimer = setTimeout(() => { + this.debounceTimer = undefined; + const dirs = new Set(this.changedDirs); + const integrationsFileDirs = new Set(this.integrationsFileChangedDirs); + this.changedDirs.clear(); + this.integrationsFileChangedDirs.clear(); + this.handleChangedDirs(dirs, integrationsFileDirs).catch((error) => + logger.error('IntegrationsEnvFileWatcher: Failed to handle env file change', error) + ); + }, debounceTimeInMilliseconds); + } + + private watchDir(dir: Uri): void { + const dirPath = dir.fsPath; + if (this.watchedDirs.has(dirPath)) { + return; + } + this.watchedDirs.add(dirPath); + + for (const fileName of watchedEnvFileNames) { + const pattern = new RelativePattern(dir, fileName); + const watcher = workspace.createFileSystemWatcher(pattern, false, false, false); + + this.disposables.push( + watcher, + watcher.onDidChange(() => this.onFileEvent(dir, fileName)), + watcher.onDidCreate(() => this.onFileEvent(dir, fileName)), + watcher.onDidDelete(() => this.onFileEvent(dir, fileName)) + ); + } + } + + private watchNotebookDir(notebook: NotebookDocument): void { + if (notebook.notebookType !== DEEPNOTE_NOTEBOOK_TYPE) { + return; + } + + const deepnoteDir = Uri.joinPath(notebookPathToDeepnoteProjectFilePath(notebook.uri), '..'); + this.watchDir(deepnoteDir); + } +} diff --git a/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.unit.test.ts b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.unit.test.ts new file mode 100644 index 0000000000..e148d3e9da --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.unit.test.ts @@ -0,0 +1,196 @@ +import { assert } from 'chai'; +import { Disposable, NotebookDocument, Uri } from 'vscode'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; + +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; +import { IDisposable } from '../../../platform/common/types'; +import { IFileSystem } from '../../../platform/common/platform/types'; +import { IIntegrationEnvLiveRefresher } from './types'; +import { IntegrationsEnvFileWatcher } from './integrationsEnvFileWatcher.node'; +import { dispose } from '../../../platform/common/utils/lifecycle'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; +import { notebookPathToDeepnoteProjectFilePath } from '../../../platform/deepnote/deepnoteProjectUtils'; + +suite('IntegrationsEnvFileWatcher', () => { + let watcher: IntegrationsEnvFileWatcher; + let liveRefresher: IIntegrationEnvLiveRefresher; + let fileSystem: IFileSystem; + let disposables: IDisposable[]; + + const workspaceRoot = Uri.file('/ws'); + + function createMockNotebook(uri: Uri, notebookType: string = DEEPNOTE_NOTEBOOK_TYPE): NotebookDocument { + const notebook = mock(); + when(notebook.notebookType).thenReturn(notebookType); + when(notebook.uri).thenReturn(uri); + + return instance(notebook); + } + + /** The dir fsPath the watcher derives from a notebook uri (the `.deepnote` file's dir). */ + function deepnoteDirOf(uri: Uri): string { + return Uri.joinPath(notebookPathToDeepnoteProjectFilePath(uri), '..').fsPath; + } + + setup(() => { + resetVSCodeMocks(); + disposables = [new Disposable(() => resetVSCodeMocks())]; + + liveRefresher = mock(); + when(liveRefresher.refresh(anything())).thenResolve(); + + // Default: a `.deepnote.env.yaml` exists, so a dir change refreshes; individual tests override this. + fileSystem = mock(); + when(fileSystem.exists(anything())).thenResolve(true); + + watcher = new IntegrationsEnvFileWatcher(instance(liveRefresher), instance(fileSystem), disposables); + }); + + teardown(() => { + disposables = dispose(disposables); + }); + + test('refreshes every notebook view whose .deepnote dir changed (no deduplication; the refresher gates each kernel)', async () => { + // Two open views of the SAME .deepnote file (differ only by notebook query) — both are refreshed. + const uriA = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const uriB = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-b' }); + const notebookA = createMockNotebook(uriA); + const notebookB = createMockNotebook(uriB); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookA, notebookB]); + + await watcher.handleChangedDirs(new Set([deepnoteDirOf(uriA)])); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [notebookA, notebookB]); + }); + + test('resolves affected notebooks via the workspace-folder root (dir-then-root fallback)', async () => { + // The .deepnote lives in a nested dir; only the workspace ROOT changed. + const uri = Uri.file('/ws/nested/deep/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(mockedVSCodeNamespaces.workspace.getWorkspaceFolder(anything())).thenReturn({ + uri: workspaceRoot + } as never); + + // changedDirs = workspace root, NOT the .deepnote dir (/ws/nested/deep). + await watcher.handleChangedDirs(new Set([workspaceRoot.fsPath])); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [notebook]); + }); + + test('does not refresh when the changed dir matches no open Deepnote notebook', async () => { + const uri = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + // An unrelated dir changed - the notebook's dir and workspace root are not in the set. + await watcher.handleChangedDirs(new Set([Uri.file('/some/other/dir').fsPath])); + + verify(liveRefresher.refresh(anything())).never(); + }); + + test('ignores non-Deepnote notebooks even when their dir changed', async () => { + const uri = Uri.file('/ws/proj/app.ipynb').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri, 'jupyter-notebook'); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + await watcher.handleChangedDirs(new Set([deepnoteDirOf(uri)])); + + verify(liveRefresher.refresh(anything())).never(); + }); + + test('does not refresh when the env-file feature is disabled for the notebook', async () => { + const uri = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote', anything())).thenReturn({ + get: () => false + } as never); + + await watcher.handleChangedDirs(new Set([deepnoteDirOf(uri)])); + + verify(liveRefresher.refresh(anything())).never(); + }); + + test('refreshes when the env-file feature is explicitly enabled for the notebook', async () => { + const uri = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote', anything())).thenReturn({ + get: () => true + } as never); + + await watcher.handleChangedDirs(new Set([deepnoteDirOf(uri)])); + + verify(liveRefresher.refresh(anything())).once(); + }); + + test('does not refresh when no .deepnote.env.yaml exists for the notebook (an unrelated .env change)', async () => { + const uri = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + // No integrations file present in any candidate dir — the change must be treated as unrelated. + when(fileSystem.exists(anything())).thenResolve(false); + + await watcher.handleChangedDirs(new Set([deepnoteDirOf(uri)])); + + verify(liveRefresher.refresh(anything())).never(); + }); + + suite('integrations-file deletion', () => { + test('refreshes after `.deepnote.env.yaml` is deleted, so the variables it contributed are unset', async () => { + const uri = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + // The file is gone, so the existence probe that gates a plain `.env` change would skip this + // notebook. Deleting the file is how a user revokes its credentials, so it must still refresh. + when(fileSystem.exists(anything())).thenResolve(false); + + await watcher.handleChangedDirs(new Set([deepnoteDirOf(uri)]), new Set([deepnoteDirOf(uri)])); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [notebook]); + }); + + test('refreshes after a `.deepnote.env.yaml` deletion in the workspace root', async () => { + const uri = Uri.file('/ws/nested/deep/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(mockedVSCodeNamespaces.workspace.getWorkspaceFolder(anything())).thenReturn({ + uri: workspaceRoot + } as never); + when(fileSystem.exists(anything())).thenResolve(false); + + await watcher.handleChangedDirs(new Set([workspaceRoot.fsPath]), new Set([workspaceRoot.fsPath])); + + verify(liveRefresher.refresh(anything())).once(); + }); + + test('a deleted `.env` alone still does not refresh when no `.deepnote.env.yaml` exists', async () => { + const uri = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(fileSystem.exists(anything())).thenResolve(false); + + // Empty second set: the event was for `.env`, not the integrations file, so the F2 guard applies. + await watcher.handleChangedDirs(new Set([deepnoteDirOf(uri)]), new Set()); + + verify(liveRefresher.refresh(anything())).never(); + }); + }); +}); diff --git a/src/notebooks/deepnote/integrations/sqlIntegrationStartupCodeProvider.ts b/src/notebooks/deepnote/integrations/sqlIntegrationStartupCodeProvider.ts index 693f1d8f11..2ebfbd4294 100644 --- a/src/notebooks/deepnote/integrations/sqlIntegrationStartupCodeProvider.ts +++ b/src/notebooks/deepnote/integrations/sqlIntegrationStartupCodeProvider.ts @@ -10,6 +10,7 @@ import { logger } from '../../../platform/logging'; import { isPythonKernelConnection } from '../../../kernels/helpers'; import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; import { workspace } from 'vscode'; +import { recordStartupIntegrationEnvNames } from './startupIntegrationEnvTracker'; /** * Provides startup code to inject SQL integration credentials into the kernel environment. @@ -89,14 +90,21 @@ export class SqlIntegrationStartupCodeProvider implements IStartupCodeProvider, code.push(` # [SQL Integration] Setting ${Object.keys(envVars).length} SQL integration env vars...`); // Set each environment variable directly in os.environ + const emittedNames: string[] = []; for (const [key, value] of Object.entries(envVars)) { if (value) { // Use JSON.stringify to properly escape the value const jsonEscaped = JSON.stringify(value); code.push(` os.environ['${key}'] = ${jsonEscaped}`); + emittedNames.push(key); } } + // Give IntegrationEnvLiveRefresher the removal baseline for this kernel. Recorded only on + // the path that emits code: the early returns above write nothing, and clearing a previous + // entry there would lose names that are still live in a restarted kernel's environment. + recordStartupIntegrationEnvNames(kernel, emittedNames); + code.push( ` # [SQL Integration] Successfully set ${Object.keys(envVars).length} SQL integration env vars` ); diff --git a/src/notebooks/deepnote/integrations/startupIntegrationEnvTracker.ts b/src/notebooks/deepnote/integrations/startupIntegrationEnvTracker.ts new file mode 100644 index 0000000000..2bf95f9013 --- /dev/null +++ b/src/notebooks/deepnote/integrations/startupIntegrationEnvTracker.ts @@ -0,0 +1,25 @@ +import type { IBaseKernel } from '../../../kernels/types'; + +/** + * The integration env-var names that `SqlIntegrationStartupCodeProvider` last emitted for a kernel. + * + * `IntegrationEnvLiveRefresher` needs a baseline of what is already set in a kernel to know what to + * remove on the first live refresh, and start/restart credentials are delivered by the startup-code + * provider rather than by the refresher. Re-reading the integration config to derive that baseline + * would be wrong: an edit landing between the provider's read and the re-read drops a name the + * provider actually wrote, and that variable would then never be removed. So the provider records + * exactly what it emitted and the refresher unions those names in. + * + * Keyed weakly by kernel, so entries disappear with the kernel; a restart overwrites the previous + * entry because the provider runs again on every kernel start and restart. + */ +const namesByKernel = new WeakMap>(); + +/** The names recorded for `kernel`, or `undefined` if the provider never emitted any for it. */ +export function getStartupIntegrationEnvNames(kernel: IBaseKernel): ReadonlySet | undefined { + return namesByKernel.get(kernel); +} + +export function recordStartupIntegrationEnvNames(kernel: IBaseKernel, names: Iterable): void { + namesByKernel.set(kernel, new Set(names)); +} diff --git a/src/notebooks/deepnote/integrations/types.ts b/src/notebooks/deepnote/integrations/types.ts index 3e151cb788..b7f02521c6 100644 --- a/src/notebooks/deepnote/integrations/types.ts +++ b/src/notebooks/deepnote/integrations/types.ts @@ -1,5 +1,5 @@ import type { DeepnoteBlock } from '@deepnote/blocks'; -import { Event, Uri } from 'vscode'; +import { Event, NotebookDocument, Uri } from 'vscode'; import { IntegrationWithStatus } from '../../../platform/notebooks/deepnote/integrationTypes'; @@ -46,6 +46,12 @@ export interface IIntegrationManager { activate(): void; } +export const IIntegrationEnvLiveRefresher = Symbol('IIntegrationEnvLiveRefresher'); +export interface IIntegrationEnvLiveRefresher { + /** Applies each notebook's current integration env to its running kernel (no restart); notifies once. */ + refresh(notebooks: readonly NotebookDocument[]): Promise; +} + /** Persisted federated-auth token entry; fingerprints `${clientId}|${clientSecret}|${project}` to detect stale tokens. Only the refresh token is persisted. */ export interface FederatedAuthTokenEntry { integrationId: string; diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts index 5bf312977c..4c94f450c8 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts @@ -16,7 +16,7 @@ import { window, workspace } from 'vscode'; -import { inject, injectable } from 'inversify'; +import { inject, injectable, optional } from 'inversify'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { IDisposableRegistry } from '../../platform/common/types'; @@ -27,6 +27,7 @@ import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes'; import { IDeepnoteNotebookManager } from '../types'; +import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; import { DatabaseIntegrationType, databaseIntegrationTypes } from '@deepnote/database-integrations'; /** @@ -69,7 +70,10 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid constructor( @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, - @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager + @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, + @inject(ISqlIntegrationEnvVarsProvider) + @optional() + private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider ) {} public activate(): void { @@ -229,12 +233,25 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid // Integration is configured, use the config name displayName = config.name; } else { - // Integration is not configured, try to get the name from the project's integration list - const notebookId = cell.notebook.metadata?.deepnoteNotebookId; - const project = notebookId ? this.notebookManager.getProjectForNotebook(projectId, notebookId) : undefined; - const projectIntegration = project?.project.integrations?.find((i) => i.id === integrationId); - const baseName = projectIntegration?.name || l10n.t('Unknown integration'); - displayName = l10n.t('{0} (configure)', baseName); + // Not in SecretStorage — a `.deepnote.env.yaml` file config still counts as configured, so check the + // merged configs before prompting the user to configure (F13). + const fileConfig = this.sqlIntegrationEnvVars + ? (await this.sqlIntegrationEnvVars.getMergedConfigs(cell.notebook.uri)).find( + (c) => c.id === integrationId + ) + : undefined; + if (fileConfig) { + displayName = fileConfig.name; + } else { + // Integration is not configured, try to get the name from the project's integration list + const notebookId = cell.notebook.metadata?.deepnoteNotebookId; + const project = notebookId + ? this.notebookManager.getProjectForNotebook(projectId, notebookId) + : undefined; + const projectIntegration = project?.project.integrations?.find((i) => i.id === integrationId); + const baseName = projectIntegration?.name || l10n.t('Unknown integration'); + displayName = l10n.t('{0} (configure)', baseName); + } } // Create a status bar item that opens the integration picker diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index eb87ce58a3..a99c39e985 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -44,6 +44,7 @@ import { DeepnoteActivationService } from './deepnote/deepnoteActivationService' import { DeepnoteNotebookManager } from './deepnote/deepnoteNotebookManager'; import { IDeepnoteNotebookManager } from './types'; import { IntegrationStorage } from '../platform/notebooks/deepnote/integrationStorage'; +import { IntegrationsFileConfigProvider } from '../platform/notebooks/deepnote/integrationsFileConfigProvider.node'; import { IntegrationDetector } from './deepnote/integrations/integrationDetector'; import { IntegrationManager } from './deepnote/integrations/integrationManager'; import { IntegrationWebviewProvider } from './deepnote/integrations/integrationWebview'; @@ -51,6 +52,7 @@ import { IFederatedAuthSqlBlockCodeGenerator, IFederatedAuthTokenStorage, IIntegrationDetector, + IIntegrationEnvLiveRefresher, IIntegrationManager, IIntegrationStorage, IIntegrationWebviewProvider @@ -61,6 +63,7 @@ import { FederatedAuthOrphanedTokenCleaner } from './deepnote/integrations/feder import { FederatedAuthSqlBlockCodeGenerator } from './deepnote/integrations/federatedAuth/federatedAuthSqlBlockCodeGenerator.node'; import { FederatedAuthTokenStorage } from './deepnote/integrations/federatedAuth/federatedAuthTokenStorage.node'; import { + IIntegrationsFileConfigProvider, IPlatformNotebookEditorProvider, IPlatformDeepnoteNotebookManager } from '../platform/notebooks/deepnote/types'; @@ -100,7 +103,9 @@ import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIn import { DeepnoteCellCopyHandler } from './deepnote/deepnoteCellCopyHandler'; import { DeepnoteEnvironmentTreeDataProvider } from '../kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.node'; import { OpenInDeepnoteHandler } from './deepnote/openInDeepnoteHandler.node'; -import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler'; +import { IntegrationEnvRefreshHandler } from './deepnote/integrations/integrationEnvRefreshHandler'; +import { IntegrationsEnvFileWatcher } from './deepnote/integrations/integrationsEnvFileWatcher.node'; +import { IntegrationEnvLiveRefresher } from './deepnote/integrations/integrationEnvLiveRefresher.node'; import { ISnapshotMetadataService, SnapshotService } from './deepnote/snapshots/snapshotService'; import { EnvironmentCapture, IEnvironmentCapture } from './deepnote/snapshots/environmentCapture.node'; import { DeepnoteFileChangeWatcher } from './deepnote/deepnoteFileChangeWatcher'; @@ -233,8 +238,23 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea ); serviceManager.addSingleton( IExtensionSyncActivationService, - IntegrationKernelRestartHandler + IntegrationEnvRefreshHandler ); + serviceManager.addSingleton( + IIntegrationsFileConfigProvider, + IntegrationsFileConfigProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + IntegrationsEnvFileWatcher + ); + serviceManager.addSingleton( + IIntegrationEnvLiveRefresher, + IntegrationEnvLiveRefresher + ); + // Activated in its own right: it subscribes to kernel start/restart to seed its removal baseline, + // which must not depend on some other service happening to inject it. + serviceManager.addBinding(IIntegrationEnvLiveRefresher, IExtensionSyncActivationService); // Deepnote kernel services serviceManager.addSingleton(DeepnoteAgentSkillsManager, DeepnoteAgentSkillsManager); diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index 28937d6b60..6d6b826180 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/src/notebooks/serviceRegistry.web.ts @@ -53,7 +53,6 @@ import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnote import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; -import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler'; import { FederatedAuthCommandHandlerWeb } from './deepnote/integrations/federatedAuth/federatedAuthCommandHandler.web'; import { DeepnoteFileChangeWatcher } from './deepnote/deepnoteFileChangeWatcher'; import { DeepnoteNotebookInfoStatusBar } from './deepnote/deepnoteNotebookInfoStatusBar'; @@ -139,10 +138,6 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, SqlCellStatusBarProvider ); - serviceManager.addSingleton( - IExtensionSyncActivationService, - IntegrationKernelRestartHandler - ); serviceManager.addSingleton( IExtensionSyncActivationService, FederatedAuthCommandHandlerWeb diff --git a/src/platform/notebooks/deepnote/integrationTypes.ts b/src/platform/notebooks/deepnote/integrationTypes.ts index 8eb24affc5..ad92029eac 100644 --- a/src/platform/notebooks/deepnote/integrationTypes.ts +++ b/src/platform/notebooks/deepnote/integrationTypes.ts @@ -74,7 +74,12 @@ export interface LegacyDuckDBIntegrationConfig extends BaseLegacyIntegrationConf type: LegacyIntegrationType.DuckDB; } -import { DatabaseIntegrationConfig, DatabaseIntegrationType } from '@deepnote/database-integrations'; +import { + DatabaseIntegrationConfig, + DatabaseIntegrationType, + FederatedAuthMethod, + isFederatedAuthMethod +} from '@deepnote/database-integrations'; // Import and re-export Snowflake auth constants from shared module import { type SnowflakeAuthMethod, @@ -173,3 +178,22 @@ export interface IntegrationWithStatus { /** Federated-auth token status; only meaningful for federated integrations (currently BigQuery + `google-oauth`). */ tokenStatus?: FederatedAuthTokenStatus; } + +/** + * Narrows integration metadata to the federated-auth variant. Shared by the file-config provider and the SQL + * env-vars provider (upstream `isFederatedAuthMetadata`'s generic doesn't unify with our + * `DatabaseIntegrationConfig['metadata']` union); delegates to the exported `isFederatedAuthMethod` at runtime. + */ +export function isFederatedAuthMetadata( + metadata: DatabaseIntegrationConfig['metadata'] +): metadata is Extract { + if (typeof metadata !== 'object' || metadata === null) { + return false; + } + if (!('authMethod' in metadata)) { + return false; + } + const authMethod = metadata.authMethod; + + return typeof authMethod === 'string' && isFederatedAuthMethod(authMethod); +} diff --git a/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.ts b/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.ts new file mode 100644 index 0000000000..b21319a339 --- /dev/null +++ b/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.ts @@ -0,0 +1,220 @@ +import dotenv from 'dotenv'; +import { inject, injectable } from 'inversify'; +import { Diagnostic, DiagnosticCollection, DiagnosticSeverity, languages, Range, Uri, workspace } from 'vscode'; + +import { + BUILTIN_INTEGRATIONS, + DatabaseIntegrationConfig, + DEFAULT_ENV_FILE, + DEFAULT_INTEGRATIONS_FILE, + parseIntegrations, + ValidationIssue +} from '@deepnote/database-integrations'; + +import { IFileSystem } from '../../common/platform/types'; +import { IDisposableRegistry } from '../../common/types'; +import { logger } from '../../logging'; +import { isFederatedAuthMetadata } from './integrationTypes'; +import { IIntegrationsFileConfigProvider } from './types'; + +/** + * Stateless loader that reads integration configs from a `.deepnote.env.yaml` file (CLI parity), + * resolving `env:` references against a sibling `.env` file and `process.env`. Replicates the Node + * filesystem/dotenv shell that `@deepnote/database-integrations` does not export, delegating parsing + * to the exported, environment-agnostic `parseIntegrations`. + * + * No caching, no watching: a fresh read happens on every call, since it is only invoked at + * kernel/server (re)start. + */ +@injectable() +export class IntegrationsFileConfigProvider implements IIntegrationsFileConfigProvider { + private readonly diagnostics: DiagnosticCollection | undefined; + + constructor( + @inject(IFileSystem) private readonly fileSystem: IFileSystem, + @inject(IDisposableRegistry) disposables: IDisposableRegistry + ) { + this.diagnostics = languages.createDiagnosticCollection('deepnote-integrations'); + if (this.diagnostics) { + disposables.push(this.diagnostics); + } + } + + public async getConfigsForFile( + deepnoteFileUri: Uri + ): Promise<{ configs: DatabaseIntegrationConfig[]; issues: ValidationIssue[] }> { + try { + const enabled = workspace + .getConfiguration('deepnote', deepnoteFileUri) + .get('integrations.envFile.enabled', true); + if (!enabled) { + return { configs: [], issues: [] }; + } + + const candidateDirs = this.getCandidateDirs(deepnoteFileUri); + + // Locate the integrations YAML (dir-then-root). A missing file is not an error. + const yamlUri = await this.findFirstExisting(candidateDirs, DEFAULT_INTEGRATIONS_FILE); + if (!yamlUri) { + return { configs: [], issues: [] }; + } + + const yaml = await this.fileSystem.readFile(yamlUri); + + // Locate the `.env` (dir-then-root) and resolve `env:` refs against it; real env wins over the file. + const envUri = await this.findFirstExisting(candidateDirs, DEFAULT_ENV_FILE); + const fileEnv = envUri ? dotenv.parse(await this.fileSystem.readFile(envUri)) : {}; + const env: Record = { ...fileEnv, ...this.getProcessEnvironment() }; + + const { integrations, issues } = parseIntegrations({ yaml, env }); + + const result = this.filterIntegrations(integrations, issues); + this.updateDiagnostics(yamlUri, result.issues); + + return result; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + const issue: ValidationIssue = { + path: '', + message: `Failed to read integrations file: ${message}`, + code: 'file_read_error' + }; + logger.error(`IntegrationsFileConfigProvider: ${issue.message}`); + + return { configs: [], issues: [issue] }; + } + } + + /** The process environment merged over the `.env` file; a seam tests override so they never touch the real `process.env`. */ + protected getProcessEnvironment(): Record { + return process.env; + } + + /** + * Filters parsed integrations into the configs we can inject, collecting an issue for each dropped + * entry: reserved ids, unsupported (dataframe) types, duplicate ids (first wins), and federated-auth + * configs (whose tokens are only available via SecretStorage, not the environment file). + */ + private filterIntegrations( + integrations: DatabaseIntegrationConfig[], + parseIssues: ValidationIssue[] + ): { configs: DatabaseIntegrationConfig[]; issues: ValidationIssue[] } { + const configs: DatabaseIntegrationConfig[] = []; + const issues: ValidationIssue[] = [...parseIssues]; + const seenIds = new Set(); + + integrations.forEach((integration, index) => { + const issuePath = `integrations[${index}]`; + + if (BUILTIN_INTEGRATIONS.has(integration.id)) { + issues.push({ + path: issuePath, + message: `Integration '${integration.id}' uses a reserved id and was ignored.`, + code: 'reserved_integration_id' + }); + + return; + } + + if (integration.type === 'pandas-dataframe') { + issues.push({ + path: issuePath, + message: `Integration '${integration.id}' has unsupported type '${integration.type}' and was ignored.`, + code: 'unsupported_integration_type' + }); + + return; + } + + if (seenIds.has(integration.id)) { + issues.push({ + path: issuePath, + message: `Integration '${integration.id}' has a duplicate id and was ignored.`, + code: 'duplicate_integration_id' + }); + + return; + } + + if (isFederatedAuthMetadata(integration.metadata)) { + issues.push({ + path: issuePath, + message: `Integration '${integration.id}' uses federated authentication, which is unsupported from the environment file, and was ignored.`, + code: 'unsupported_federated_integration' + }); + + return; + } + + seenIds.add(integration.id); + configs.push(integration); + }); + + issues.forEach((issue) => { + logger.warn(`IntegrationsFileConfigProvider: ${issue.code} at '${issue.path}': ${issue.message}`); + }); + + return { configs, issues }; + } + + private async findFirstExisting(dirs: Uri[], fileName: string): Promise { + for (const dir of dirs) { + const candidate = Uri.joinPath(dir, fileName); + if (await this.fileSystem.exists(candidate)) { + return candidate; + } + } + + return undefined; + } + + /** + * Candidate directories to look for the integration/env files in priority order: next to the + * `.deepnote` file first, then the workspace-folder root. Undefined entries are skipped and + * duplicates removed. + */ + private getCandidateDirs(deepnoteFileUri: Uri): Uri[] { + const dirs: Uri[] = [Uri.joinPath(deepnoteFileUri, '..')]; + const workspaceFolder = workspace.getWorkspaceFolder(deepnoteFileUri); + if (workspaceFolder) { + dirs.push(workspaceFolder.uri); + } + + const seen = new Set(); + + return dirs.filter((dir) => { + const key = dir.toString(); + if (seen.has(key)) { + return false; + } + seen.add(key); + + return true; + }); + } + + /** Surfaces validation issues in the Problems panel against the located `.deepnote.env.yaml` so a typo/missing key isn't silent (F6); a clean parse clears them. No-op when diagnostics are unavailable (e.g. web/tests). */ + private updateDiagnostics(yamlUri: Uri, issues: ValidationIssue[]): void { + if (!this.diagnostics) { + return; + } + + if (issues.length === 0) { + this.diagnostics.delete(yamlUri); + + return; + } + + const diagnostics = issues.map((issue) => { + const detail = issue.path + ? `${issue.code} at '${issue.path}': ${issue.message}` + : `${issue.code}: ${issue.message}`; + const diagnostic = new Diagnostic(new Range(0, 0, 0, 0), detail, DiagnosticSeverity.Warning); + diagnostic.source = 'Deepnote integrations'; + + return diagnostic; + }); + + this.diagnostics.set(yamlUri, diagnostics); + } +} diff --git a/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.unit.test.ts b/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.unit.test.ts new file mode 100644 index 0000000000..1d9a2af74d --- /dev/null +++ b/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.unit.test.ts @@ -0,0 +1,375 @@ +import assert from 'assert'; + +import { + DatabaseIntegrationConfig, + DEFAULT_ENV_FILE, + DEFAULT_INTEGRATIONS_FILE +} from '@deepnote/database-integrations'; +import dedent from 'dedent'; +import { anything, instance, mock, verify, when } from 'ts-mockito'; +import { Uri, WorkspaceConfiguration, WorkspaceFolder } from 'vscode'; + +import { IFileSystem } from '../../common/platform/types'; +import { IntegrationsFileConfigProvider } from './integrationsFileConfigProvider.node'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; + +/** A file present in the virtual filesystem: either readable `content`, or a `readError` that rejects. */ +interface VirtualFile { + path: string; + content?: string; + readError?: Error; +} + +/** Provider whose process environment is controllable, so tests never read or mutate the real `process.env`. */ +class TestableIntegrationsFileConfigProvider extends IntegrationsFileConfigProvider { + public processEnvironment: Record = {}; + + protected override getProcessEnvironment(): Record { + return this.processEnvironment; + } +} + +suite('IntegrationsFileConfigProvider', () => { + const deepnoteFileUri = Uri.file('/workspace/project/notebook.deepnote'); + const deepnoteDirUri = Uri.joinPath(deepnoteFileUri, '..'); + // Build the expected file paths exactly as the loader does (dir of the `.deepnote` file). + const yamlPath = Uri.joinPath(deepnoteDirUri, DEFAULT_INTEGRATIONS_FILE).fsPath; + const envPath = Uri.joinPath(deepnoteDirUri, DEFAULT_ENV_FILE).fsPath; + + let fileSystem: IFileSystem; + let provider: TestableIntegrationsFileConfigProvider; + let featureEnabled: boolean; + let workspaceFolder: WorkspaceFolder | undefined; + + setup(() => { + resetVSCodeMocks(); + + featureEnabled = true; + workspaceFolder = undefined; + + fileSystem = mock(); + provider = new TestableIntegrationsFileConfigProvider(instance(fileSystem), []); + + // The gate reads `deepnote.integrations.envFile.enabled`; return the current `featureEnabled` value. + when(mockedVSCodeNamespaces.workspace.getConfiguration(anything(), anything())).thenReturn({ + get: () => featureEnabled + } as unknown as WorkspaceConfiguration); + when(mockedVSCodeNamespaces.workspace.getWorkspaceFolder(anything())).thenCall(() => workspaceFolder); + }); + + /** Wires `IFileSystem.exists`/`readFile` to a small in-memory set of files keyed by fsPath. */ + function configureFileSystem(files: VirtualFile[]): void { + const byPath = new Map(files.map((file) => [file.path, file])); + + when(fileSystem.exists(anything())).thenCall((uri: Uri) => Promise.resolve(byPath.has(uri.fsPath))); + when(fileSystem.readFile(anything())).thenCall((uri: Uri) => { + const file = byPath.get(uri.fsPath); + if (!file) { + return Promise.reject(new Error(`ENOENT: ${uri.fsPath}`)); + } + if (file.readError) { + return Promise.reject(file.readError); + } + + return Promise.resolve(file.content ?? ''); + }); + } + + /** Reads a metadata field off a parsed config without narrowing the metadata union. */ + function metadataField(config: DatabaseIntegrationConfig, key: string): unknown { + return (config.metadata as unknown as Record)[key]; + } + + test('returns configs for a valid integrations file', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: my-postgres + name: My Postgres + type: pgsql + metadata: + host: localhost + port: "5432" + database: mydb + user: root + password: my-secret + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(configs[0].id, 'my-postgres'); + assert.strictEqual(configs[0].type, 'pgsql'); + assert.deepStrictEqual(issues, []); + }); + + test('resolves env: references from the .env file', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: dotenv-postgres + name: Dotenv Postgres + type: pgsql + metadata: + host: localhost + port: "5432" + database: mydb + user: root + password: "env:DEEPNOTE_TEST_DOTENV_PASSWORD" + ` + }, + { path: envPath, content: 'DEEPNOTE_TEST_DOTENV_PASSWORD=secret-from-dotenv\n' } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(metadataField(configs[0], 'password'), 'secret-from-dotenv'); + assert.deepStrictEqual(issues, []); + }); + + test('lets the process environment override values from the .env file', async () => { + provider.processEnvironment = { DEEPNOTE_TEST_OVERRIDE_PASSWORD: 'secret-from-process-env' }; + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: override-postgres + name: Override Postgres + type: pgsql + metadata: + host: localhost + port: "5432" + database: mydb + user: root + password: "env:DEEPNOTE_TEST_OVERRIDE_PASSWORD" + ` + }, + { path: envPath, content: 'DEEPNOTE_TEST_OVERRIDE_PASSWORD=stale-from-dotenv\n' } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(metadataField(configs[0], 'password'), 'secret-from-process-env'); + assert.deepStrictEqual(issues, []); + }); + + test('returns an empty result when the YAML file is missing', async () => { + configureFileSystem([]); + + const result = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(result, { configs: [], issues: [] }); + }); + + test('reports a yaml_parse_error for malformed YAML', async () => { + configureFileSystem([{ path: yamlPath, content: 'integrations:\n - id: "unclosed string' }]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'yaml_parse_error'); + }); + + test('drops an integration and reports env_var_not_defined for an unresolved env: reference', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: missing-env-postgres + name: Missing Env Postgres + type: pgsql + metadata: + host: localhost + port: "5432" + database: mydb + user: root + password: "env:DEEPNOTE_TEST_UNDEFINED_VAR_DEADBEEF" + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'env_var_not_defined'); + }); + + test('finds the YAML at the workspace-folder root when absent next to the .deepnote file', async () => { + const nestedDeepnoteUri = Uri.file('/workspace/project/sub/notebook.deepnote'); + const rootFolder: WorkspaceFolder = { uri: Uri.file('/workspace/project'), name: 'project', index: 0 }; + const rootYamlPath = Uri.joinPath(rootFolder.uri, DEFAULT_INTEGRATIONS_FILE).fsPath; + + workspaceFolder = rootFolder; + configureFileSystem([ + { + path: rootYamlPath, + content: dedent` + integrations: + - id: root-postgres + name: Root Postgres + type: pgsql + metadata: + host: localhost + port: "5432" + database: mydb + user: root + password: my-secret + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(nestedDeepnoteUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(configs[0].id, 'root-postgres'); + assert.deepStrictEqual(issues, []); + }); + + test('drops an integration whose id is reserved (reserved_integration_id)', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: deepnote-dataframe-sql + name: Reserved + type: pgsql + metadata: + host: localhost + port: "5432" + database: mydb + user: root + password: my-secret + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'reserved_integration_id'); + assert.strictEqual(issues[0].path, 'integrations[0]'); + }); + + test('drops an integration with an unsupported pandas-dataframe type (unsupported_integration_type)', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: my-dataframe + name: My Dataframe + type: pandas-dataframe + metadata: {} + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'unsupported_integration_type'); + assert.strictEqual(issues[0].path, 'integrations[0]'); + }); + + test('drops a duplicate id, keeping the first occurrence (duplicate_integration_id)', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: dup-postgres + name: First + type: pgsql + metadata: + host: first-host + port: "5432" + database: mydb + user: root + password: my-secret + - id: dup-postgres + name: Second + type: pgsql + metadata: + host: second-host + port: "5432" + database: mydb + user: root + password: my-secret + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(configs[0].id, 'dup-postgres'); + assert.strictEqual(metadataField(configs[0], 'host'), 'first-host'); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'duplicate_integration_id'); + assert.strictEqual(issues[0].path, 'integrations[1]'); + }); + + test('returns a file_read_error issue (and does not throw) when reading the file fails', async () => { + configureFileSystem([{ path: yamlPath, readError: new Error('disk failure') }]); + + // Must resolve, never reject: a read failure degrades to an issue, not a thrown error. + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'file_read_error'); + assert.strictEqual(issues[0].path, ''); + assert.ok(issues[0].message.includes('Failed to read integrations file')); + }); + + test('drops a federated (google-oauth) integration (unsupported_federated_integration)', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: bq-oauth + name: BigQuery OAuth + type: big-query + metadata: + authMethod: google-oauth + project: my-project + clientId: my-client-id + clientSecret: my-client-secret + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'unsupported_federated_integration'); + assert.strictEqual(issues[0].path, 'integrations[0]'); + }); + + test('returns an empty result without touching the filesystem when the feature is disabled', async () => { + featureEnabled = false; + + const result = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(result, { configs: [], issues: [] }); + verify(fileSystem.exists(anything())).never(); + verify(fileSystem.readFile(anything())).never(); + }); +}); diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts index 7a612e927e..8da5e213c7 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts @@ -1,37 +1,24 @@ -import { inject, injectable } from 'inversify'; -import { CancellationToken, Event, EventEmitter } from 'vscode'; +import { inject, injectable, optional } from 'inversify'; +import { CancellationToken, Event, EventEmitter, Uri } from 'vscode'; -import { - DatabaseIntegrationConfig, - FederatedAuthMethod, - getEnvironmentVariablesForIntegrations, - isFederatedAuthMethod -} from '@deepnote/database-integrations'; +import { DeepnoteFile } from '@deepnote/blocks'; +import { DatabaseIntegrationConfig, getEnvironmentVariablesForIntegrations } from '@deepnote/database-integrations'; import { IDisposableRegistry, Resource } from '../../common/types'; import { EnvironmentVariables } from '../../common/variables/types'; +import { notebookPathToDeepnoteProjectFilePath } from '../../deepnote/deepnoteProjectUtils'; import { logger } from '../../logging'; import { + IIntegrationsFileConfigProvider, IIntegrationStorage, ISqlIntegrationEnvVarsProvider, IPlatformNotebookEditorProvider, IPlatformDeepnoteNotebookManager } from './types'; -import { DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes'; - -/** Narrows metadata to the federated-auth variant; upstream `isFederatedAuthMetadata` can't be reused because its generic doesn't unify with our union. Delegates to upstream `isFederatedAuthMethod` at runtime. */ -function isFederatedAuthMetadata( - metadata: DatabaseIntegrationConfig['metadata'] -): metadata is Extract { - if (typeof metadata !== 'object' || metadata === null) { - return false; - } - if (!('authMethod' in metadata)) { - return false; - } - const authMethod = metadata.authMethod; - return typeof authMethod === 'string' && isFederatedAuthMethod(authMethod); -} +import { DATAFRAME_SQL_INTEGRATION_ID, isFederatedAuthMetadata } from './integrationTypes'; + +/** One entry of a Deepnote project's `integrations` list. */ +type ProjectIntegration = NonNullable[number]; /** * Provides environment variables for SQL integrations. @@ -49,7 +36,10 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati @inject(IPlatformNotebookEditorProvider) private readonly notebookEditorProvider: IPlatformNotebookEditorProvider, @inject(IPlatformDeepnoteNotebookManager) private readonly notebookManager: IPlatformDeepnoteNotebookManager, - @inject(IDisposableRegistry) disposables: IDisposableRegistry + @inject(IDisposableRegistry) disposables: IDisposableRegistry, + @inject(IIntegrationsFileConfigProvider) + @optional() + private readonly fileConfigProvider?: IIntegrationsFileConfigProvider ) { logger.info('SqlIntegrationEnvironmentVariablesProvider: Constructor called - provider is being instantiated'); // Dispose emitter when extension deactivates @@ -111,19 +101,8 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati `SqlIntegrationEnvironmentVariablesProvider: Found ${projectIntegrations.length} integrations in project` ); - const configResults = await Promise.allSettled( - projectIntegrations.map((integration) => this.integrationStorage.getIntegrationConfig(integration.id)) - ); - const allConfigs: Array = configResults.flatMap((result, index) => { - if (result.status === 'fulfilled') { - return result.value ? [result.value] : []; - } - logger.error( - `SqlIntegrationEnvironmentVariablesProvider: Failed to load integration config ${projectIntegrations[index].id}`, - result.reason - ); - return []; - }); + const fileConfigs = await this.loadFileConfigs(notebook.uri); + const allConfigs = await this.mergeIntegrationConfigs(projectIntegrations, fileConfigs); // Skip federated-auth integrations: tokens are fetched per-cell via per-cell codegen in `FederatedAuthSqlBlockCodeGenerator`, not baked into kernel env. const projectIntegrationConfigs: Array = []; @@ -159,4 +138,120 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati return envVars; } + + /** + * Project SecretStorage integrations merged with `.deepnote.env.yaml` file configs (file wins, additive + * file-only). The single source of truth so integration detection, the SQL status bar, and the SQL LSP agree + * with what kernel execution actually sees (F13). Excludes the internal DuckDB integration. + */ + public async getMergedConfigs(resource: Resource, token?: CancellationToken): Promise { + if (!resource || token?.isCancellationRequested) { + return []; + } + + const notebook = this.notebookEditorProvider.findAssociatedNotebookDocument(resource); + if (!notebook) { + return []; + } + + const projectId = notebook.metadata?.deepnoteProjectId as string | undefined; + const notebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; + if (!projectId || !notebookId) { + return []; + } + + const project = this.notebookManager.getProjectForNotebook(projectId, notebookId); + if (!project) { + return []; + } + + const projectIntegrations = project.project.integrations?.slice() ?? []; + const fileConfigs = await this.loadFileConfigs(notebook.uri); + + return this.mergeIntegrationConfigs(projectIntegrations, fileConfigs); + } + + /** Loads `.deepnote.env.yaml` configs (CLI parity) when a file source is present; failures — or no provider, e.g. web — degrade to []. */ + private async loadFileConfigs(notebookUri: Uri): Promise { + if (!this.fileConfigProvider) { + return []; + } + + try { + const result = await this.fileConfigProvider.getConfigsForFile( + notebookPathToDeepnoteProjectFilePath(notebookUri) + ); + result.issues.forEach((issue) => { + logger.warn( + `SqlIntegrationEnvironmentVariablesProvider: integrations file issue ${issue.code} at '${issue.path}': ${issue.message}` + ); + }); + + return result.configs; + } catch (error) { + logger.error( + 'SqlIntegrationEnvironmentVariablesProvider: file integrations source failed; falling back to SecretStorage', + error + ); + + return []; + } + } + + /** File config wins on id conflict; SecretStorage is the fallback for project ids the file lacks; file-only ids are appended additively (CLI parity). */ + private async mergeIntegrationConfigs( + projectIntegrations: ProjectIntegration[], + fileConfigs: DatabaseIntegrationConfig[] + ): Promise { + const fileConfigsById = new Map(fileConfigs.map((config) => [config.id, config])); + const consumedFileIds = new Set(); + + // Read from SecretStorage only the project integrations the file did not provide. + const secretStorageIds = projectIntegrations + .map((integration) => integration.id) + .filter((id) => !fileConfigsById.has(id)); + const secretStorageResults = await Promise.allSettled( + secretStorageIds.map((id) => this.integrationStorage.getIntegrationConfig(id)) + ); + const secretStorageConfigsById = new Map(); + secretStorageResults.forEach((result, index) => { + const id = secretStorageIds[index]; + if (result.status === 'fulfilled') { + if (result.value) { + secretStorageConfigsById.set(id, result.value); + } + + return; + } + logger.error( + `SqlIntegrationEnvironmentVariablesProvider: Failed to load integration config ${id}`, + result.reason + ); + }); + + // Resolve each project integration in declared order: file config wins, else the SecretStorage fallback. + const allConfigs: Array = []; + for (const integration of projectIntegrations) { + const fileConfig = fileConfigsById.get(integration.id); + if (fileConfig) { + consumedFileIds.add(integration.id); + allConfigs.push(fileConfig); + + continue; + } + const secretStorageConfig = secretStorageConfigsById.get(integration.id); + if (secretStorageConfig) { + allConfigs.push(secretStorageConfig); + } + } + + // Append file-only integrations (not declared in project.integrations) additively, deduped by the map. + for (const fileConfig of fileConfigsById.values()) { + if (!consumedFileIds.has(fileConfig.id)) { + allConfigs.push(fileConfig); + } + } + + return allConfigs; + } } diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts index 11b5e57eaa..ba07c58dd1 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts @@ -1,12 +1,17 @@ import assert from 'assert'; import type { DeepnoteFile } from '@deepnote/blocks'; -import { instance, mock, when } from 'ts-mockito'; +import { anything, instance, mock, verify, when } from 'ts-mockito'; import { CancellationTokenSource, EventEmitter, NotebookDocument, Uri } from 'vscode'; import { IDisposableRegistry } from '../../common/types'; import { SqlIntegrationEnvironmentVariablesProvider } from './sqlIntegrationEnvironmentVariablesProvider'; -import { IIntegrationStorage, IPlatformDeepnoteNotebookManager, IPlatformNotebookEditorProvider } from './types'; -import { DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes'; +import { + IIntegrationsFileConfigProvider, + IIntegrationStorage, + IPlatformDeepnoteNotebookManager, + IPlatformNotebookEditorProvider +} from './types'; +import { ConfigurableDatabaseIntegrationConfig, DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes'; import { DatabaseIntegrationConfig } from '@deepnote/database-integrations'; /** Create a minimal `DeepnoteFile` for tests. */ @@ -437,6 +442,216 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { }); }); + suite('File config source (.deepnote.env.yaml) merge', () => { + const notebookUri = Uri.file('/ws/project.deepnote'); + const duckDbEnvVar = `SQL_${DATAFRAME_SQL_INTEGRATION_ID.toUpperCase().replace(/-/g, '_')}`; + let fileConfigProvider: IIntegrationsFileConfigProvider; + let providerWithFile: SqlIntegrationEnvironmentVariablesProvider; + + /** A non-federated, non-reserved pgsql config whose host is embedded in the generated connection URL. */ + function pgConfig(id: string, host: string): ConfigurableDatabaseIntegrationConfig { + return { + id, + name: id, + type: 'pgsql', + metadata: { + host, + port: '5432', + database: 'db', + user: 'u', + password: 'p', + sslEnabled: false + } + }; + } + + function stubNotebookWithProject(project: DeepnoteFile): void { + const notebook = mock(); + when(notebook.uri).thenReturn(notebookUri); + when(notebook.metadata).thenReturn({ + deepnoteProjectId: 'project-123', + deepnoteNotebookId: 'notebook-123' + }); + when(notebookEditorProvider.findAssociatedNotebookDocument(notebookUri)).thenReturn(instance(notebook)); + when(notebookManager.getProjectForNotebook('project-123', 'notebook-123')).thenReturn(project); + } + + setup(() => { + fileConfigProvider = mock(); + providerWithFile = new SqlIntegrationEnvironmentVariablesProvider( + instance(integrationStorage), + instance(notebookEditorProvider), + instance(notebookManager), + disposables, + instance(fileConfigProvider) + ); + }); + + test('Merges file config with SecretStorage config: both are included', async () => { + stubNotebookWithProject( + createMockProject('project-123', [ + { id: 'file-db', name: 'file-db', type: 'pgsql' }, + { id: 'secret-db', name: 'secret-db', type: 'pgsql' } + ]) + ); + when(fileConfigProvider.getConfigsForFile(anything())).thenResolve({ + configs: [pgConfig('file-db', 'from-file.example.com')], + issues: [] + }); + when(integrationStorage.getIntegrationConfig('secret-db')).thenResolve( + pgConfig('secret-db', 'from-secret.example.com') + ); + + const result = await providerWithFile.getEnvironmentVariables(notebookUri); + + assert.ok(result['SQL_FILE_DB'], 'File-sourced integration env var should be present'); + assert.ok(result['SQL_SECRET_DB'], 'SecretStorage-sourced integration env var should be present'); + assert.ok( + JSON.parse(result['SQL_FILE_DB']!).url.includes('from-file.example.com'), + 'File config connection details should be used for the file-sourced integration' + ); + assert.ok( + JSON.parse(result['SQL_SECRET_DB']!).url.includes('from-secret.example.com'), + 'SecretStorage config connection details should be used for the SecretStorage-only integration' + ); + }); + + test('File wins on id conflict: file config used and SecretStorage is not queried for that id', async () => { + stubNotebookWithProject( + createMockProject('project-123', [ + { id: 'shared-db', name: 'shared-db', type: 'pgsql' }, + { id: 'secret-only', name: 'secret-only', type: 'pgsql' } + ]) + ); + when(fileConfigProvider.getConfigsForFile(anything())).thenResolve({ + configs: [pgConfig('shared-db', 'from-file.example.com')], + issues: [] + }); + // Stubbed with a different host to prove the file wins; the provider must never consult it for `shared-db`. + when(integrationStorage.getIntegrationConfig('shared-db')).thenResolve( + pgConfig('shared-db', 'from-secret.example.com') + ); + when(integrationStorage.getIntegrationConfig('secret-only')).thenResolve( + pgConfig('secret-only', 'secret-only.example.com') + ); + + const result = await providerWithFile.getEnvironmentVariables(notebookUri); + + const sharedUrl = JSON.parse(result['SQL_SHARED_DB']!).url as string; + assert.ok(sharedUrl.includes('from-file.example.com'), 'File config host should win the conflict'); + assert.ok(!sharedUrl.includes('from-secret.example.com'), 'SecretStorage host must not be used'); + assert.ok(result['SQL_SECRET_ONLY'], 'SecretStorage-only integration should still be resolved'); + + // The conflicting id must never hit SecretStorage; the SecretStorage-only id must be queried exactly once. + verify(integrationStorage.getIntegrationConfig('shared-db')).never(); + verify(integrationStorage.getIntegrationConfig('secret-only')).once(); + }); + + test('File-only integration (absent from project.integrations) is injected additively', async () => { + stubNotebookWithProject(createMockProject('project-123', [])); + when(fileConfigProvider.getConfigsForFile(anything())).thenResolve({ + configs: [pgConfig('file-only', 'file-only.example.com')], + issues: [] + }); + + const result = await providerWithFile.getEnvironmentVariables(notebookUri); + + assert.ok(result['SQL_FILE_ONLY'], 'File-only integration env var should be present'); + assert.ok( + JSON.parse(result['SQL_FILE_ONLY']!).url.includes('file-only.example.com'), + 'File-only integration should use its file config' + ); + }); + + test('getMergedConfigs returns the merged config list (file wins, SecretStorage fallback, file-only additive)', async () => { + stubNotebookWithProject( + createMockProject('project-123', [ + { id: 'shared-db', name: 'shared-db', type: 'pgsql' }, + { id: 'secret-only', name: 'secret-only', type: 'pgsql' } + ]) + ); + when(fileConfigProvider.getConfigsForFile(anything())).thenResolve({ + configs: [ + pgConfig('shared-db', 'from-file.example.com'), + pgConfig('file-only', 'file-only.example.com') + ], + issues: [] + }); + when(integrationStorage.getIntegrationConfig('secret-only')).thenResolve( + pgConfig('secret-only', 'secret-only.example.com') + ); + + const merged = await providerWithFile.getMergedConfigs(notebookUri); + const byId = new Map(merged.map((config) => [config.id, config])); + + assert.deepStrictEqual( + [...byId.keys()].sort(), + ['file-only', 'secret-only', 'shared-db'], + 'merged configs must include the file-won, SecretStorage-fallback, and file-only integrations' + ); + const sharedDb = byId.get('shared-db'); + assert.ok( + sharedDb && JSON.stringify(sharedDb.metadata).includes('from-file.example.com'), + 'file config must win the id conflict in the merged list' + ); + assert.ok( + !byId.has(DATAFRAME_SQL_INTEGRATION_ID), + 'the internal DuckDB integration is not part of the merged list' + ); + verify(integrationStorage.getIntegrationConfig('shared-db')).never(); + }); + + test('getMergedConfigs returns [] when the resource resolves to no project', async () => { + const merged = await providerWithFile.getMergedConfigs(undefined); + + assert.deepStrictEqual(merged, []); + }); + + test('File provider absent: behavior is SecretStorage-only (unchanged)', async () => { + const providerWithoutFile = new SqlIntegrationEnvironmentVariablesProvider( + instance(integrationStorage), + instance(notebookEditorProvider), + instance(notebookManager), + disposables, + undefined + ); + stubNotebookWithProject( + createMockProject('project-123', [{ id: 'secret-db', name: 'secret-db', type: 'pgsql' }]) + ); + when(integrationStorage.getIntegrationConfig('secret-db')).thenResolve( + pgConfig('secret-db', 'from-secret.example.com') + ); + + const result = await providerWithoutFile.getEnvironmentVariables(notebookUri); + + assert.ok(result['SQL_SECRET_DB'], 'SecretStorage integration should be resolved without a file provider'); + assert.ok( + JSON.parse(result['SQL_SECRET_DB']!).url.includes('from-secret.example.com'), + 'SecretStorage config should be used' + ); + assert.ok(result[duckDbEnvVar], 'DuckDB integration should always be included'); + verify(integrationStorage.getIntegrationConfig('secret-db')).once(); + }); + + test('File source throws: degrades to SecretStorage + DuckDB without rejecting', async () => { + stubNotebookWithProject( + createMockProject('project-123', [{ id: 'secret-db', name: 'secret-db', type: 'pgsql' }]) + ); + when(fileConfigProvider.getConfigsForFile(anything())).thenReject(new Error('boom')); + when(integrationStorage.getIntegrationConfig('secret-db')).thenResolve( + pgConfig('secret-db', 'from-secret.example.com') + ); + + const result = await providerWithFile.getEnvironmentVariables(notebookUri); + + assert.ok( + result['SQL_SECRET_DB'], + 'SecretStorage integration should still be resolved when the file source throws' + ); + assert.ok(result[duckDbEnvVar], 'DuckDB integration should still be included when the file source throws'); + }); + }); + suite('Federated-auth integrations are skipped', () => { test('Mixed project: federated integration is skipped, non-federated is included', async () => { const resource = Uri.file('/test/notebook.deepnote'); diff --git a/src/platform/notebooks/deepnote/types.ts b/src/platform/notebooks/deepnote/types.ts index 23be88f5da..1a5f618cb7 100644 --- a/src/platform/notebooks/deepnote/types.ts +++ b/src/platform/notebooks/deepnote/types.ts @@ -3,6 +3,7 @@ import { IDisposable, Resource } from '../../common/types'; import { EnvironmentVariables } from '../../common/variables/types'; import { ConfigurableDatabaseIntegrationConfig } from './integrationTypes'; import { DeepnoteFile } from '@deepnote/blocks'; +import { DatabaseIntegrationConfig, ValidationIssue } from '@deepnote/database-integrations'; /** * Settings for select input blocks @@ -96,6 +97,25 @@ export interface ISqlIntegrationEnvVarsProvider { * Get environment variables for SQL integrations used in the given notebook. */ getEnvironmentVariables(resource: Resource, token?: CancellationToken): Promise; + + /** + * Project SecretStorage integrations merged with `.deepnote.env.yaml` file configs (file wins, additive + * file-only), so integration detection, the SQL status bar, and the SQL LSP agree with kernel execution. + */ + getMergedConfigs(resource: Resource, token?: CancellationToken): Promise; +} + +export const IIntegrationsFileConfigProvider = Symbol('IIntegrationsFileConfigProvider'); +export interface IIntegrationsFileConfigProvider { + /** + * Loads integration configs from a `.deepnote.env.yaml` file located next to the given `.deepnote` + * project file (or at the workspace-folder root), resolving `env:` references against a sibling + * `.env` file and `process.env`. Returns the accepted configs plus any validation issues; a missing + * file yields an empty result and this never throws. + */ + getConfigsForFile( + deepnoteFileUri: Uri + ): Promise<{ configs: DatabaseIntegrationConfig[]; issues: ValidationIssue[] }>; } /** diff --git a/test/e2e/fixtures/integrations-env-file.deepnote b/test/e2e/fixtures/integrations-env-file.deepnote new file mode 100644 index 0000000000..5ca88bb04b --- /dev/null +++ b/test/e2e/fixtures/integrations-env-file.deepnote @@ -0,0 +1,26 @@ +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00.000Z' + modifiedAt: '2025-01-01T00:00:00.000Z' +project: + id: e2e-integrations-env-file-project + name: E2E Integrations Env File + integrations: + - id: e2e-pg-integration + name: Prod Postgres + type: pgsql + notebooks: + - id: e2e-integrations-env-file-notebook + name: Integrations Env File + blocks: + - id: e2e-integrations-block + blockGroup: e2e-group + type: code + content: |- + import os + print(os.environ.get('PROD_POSTGRES_HOST')) + sortingKey: a0 + metadata: {} + executionMode: block + isModule: false + settings: {} diff --git a/test/e2e/suite/integrationsEnvFileInjection.e2e.test.ts b/test/e2e/suite/integrationsEnvFileInjection.e2e.test.ts new file mode 100644 index 0000000000..f49c48c9e9 --- /dev/null +++ b/test/e2e/suite/integrationsEnvFileInjection.e2e.test.ts @@ -0,0 +1,142 @@ +/** + * ExTester E2E for `.deepnote.env.yaml` integration injection: writes a `.deepnote.env.yaml` + `.env`, opens the + * notebook, runs a cell printing `PROD_POSTGRES_HOST`, and asserts it resolves to the `.env` value (see consts below). + */ + +import { expect } from 'chai'; +import * as fs from 'fs'; +import * as path from 'path'; +import { EditorView, VSBrowser, WebView } from 'vscode-extension-tester'; + +import { + FIRST_RUN_OUTPUT_TIMEOUT, + SUITE_TIMEOUT, + WORKBENCH_TIMEOUT, + copyFixtureToTempDir, + createEnvironment, + openFolderViaDialog, + openWorkspaceFile, + runOnceAndAwaitOutput, + selectEnvironmentForNotebook +} from '../helpers'; + +const NOTEBOOK_FILE_NAME = 'integrations-env-file.deepnote'; +const EXPECTED_OUTPUT = 'injected-host.example.com'; + +const INTEGRATIONS_ENV_FILE_NAME = '.deepnote.env.yaml'; +const DOTENV_FILE_NAME = '.env'; + +// `.deepnote.env.yaml`: one `pgsql` integration whose `host` is an `env:` ref resolved against `.env`; the +// remaining metadata are literals so the config is complete. `host`/`port` are quoted to keep them strings +// (`env:...` contains a colon; `5432` would otherwise parse as a number). +const INTEGRATIONS_ENV_YAML = `integrations: + - id: e2e-pg-integration + name: Prod Postgres + type: pgsql + metadata: + host: "env:DEMO_DB_HOST" + port: "5432" + database: mydb + user: root + password: secret +`; + +// `.env`: dotenv source for the `env:DEMO_DB_HOST` ref above. Deliberately a DISTINCT key from the printed +// `PROD_POSTGRES_HOST`, so a passing assertion can only come from the `.yaml` `env:` resolution — not the +// extension's direct `.env` injection (which would set `DEMO_DB_HOST`, a key the cell never reads). +const DOTENV_CONTENT = 'DEMO_DB_HOST=injected-host.example.com\n'; + +describe('Deepnote E2E — inject integration env var from `.deepnote.env.yaml`', function () { + // Per-test timeout for the whole suite (overrides the mocharc default for these tests). + this.timeout(SUITE_TIMEOUT); + + // A stable name: createEnvironment is idempotent (it treats "already exists" as success), so a + // leftover environment from a previous or retried run is reused rather than colliding — which + // also lets a persistent test instance reuse the already-provisioned venv. + const environmentName = 'E2E Integrations Env'; + + // Captured in `before` and invoked in `after` to remove the throwaway temp dir. + let cleanupTempDir: (() => void) | undefined; + // The temp workspace dir, so the live-refresh assertion can rewrite `.env`. + let tempDir: string; + + before(async function () { + // Work on a throwaway copy so execution-dirtied notebook state never touches the source tree. + const copied = copyFixtureToTempDir(NOTEBOOK_FILE_NAME); + cleanupTempDir = copied.cleanup; + tempDir = copied.tempDir; + + // Write the env files next to the notebook BEFORE opening the workspace so the loader sees them + // when the kernel first starts. `.deepnote.env.yaml` carries the integration; `.env` resolves its + // `env:` ref. + fs.writeFileSync(path.join(tempDir, INTEGRATIONS_ENV_FILE_NAME), INTEGRATIONS_ENV_YAML); + fs.writeFileSync(path.join(tempDir, DOTENV_FILE_NAME), DOTENV_CONTENT); + + await VSBrowser.instance.waitForWorkbench(WORKBENCH_TIMEOUT); + + // Open the folder as the workspace FIRST (the serializer's snapshot read blocks headlessly without one), then re-wait for the workbench after the reload. + await openFolderViaDialog(tempDir); + await VSBrowser.instance.waitForWorkbench(WORKBENCH_TIMEOUT); + + // Now that the containing folder is the workspace, the notebook is reachable by name. + await openWorkspaceFile(NOTEBOOK_FILE_NAME); + + // The native notebook editor opens because the extension registers a serializer for the + // `deepnote` notebook type; a single-notebook file resolves to its default notebook. + await VSBrowser.instance.driver.wait( + async () => (await new EditorView().getOpenEditorTitles()).some((t) => t.includes(NOTEBOOK_FILE_NAME)), + WORKBENCH_TIMEOUT, + 'Deepnote notebook editor did not open' + ); + }); + + after(async function () { + // Defensive cleanup: never leave the driver stuck inside a webview frame, and close tabs. + await new WebView().switchBack().catch((error) => { + console.warn('[deepnote-e2e] switch back from webview during cleanup:', error); + }); + await new EditorView().closeAllEditors().catch((error) => { + console.warn('[deepnote-e2e] close all editors during cleanup:', error); + }); + + // Remove the throwaway temp dir last so a failure above can't leak it. + try { + cleanupTempDir?.(); + } catch (error) { + console.warn('[deepnote-e2e] remove temp workspace dir during cleanup:', error); + } + }); + + it('injects `PROD_POSTGRES_HOST` from `.env`, then live-refreshes it on a `.env` change without a restart', async function () { + await createEnvironment(environmentName); + await selectEnvironmentForNotebook(environmentName, NOTEBOOK_FILE_NAME); + + const first = await runOnceAndAwaitOutput(NOTEBOOK_FILE_NAME, EXPECTED_OUTPUT, FIRST_RUN_OUTPUT_TIMEOUT); + expect(first).to.contain(EXPECTED_OUTPUT); + + // Live refresh: rewrite `.env`; the watcher applies the new env directly in the SAME kernel (no restart) so a re-run reads the new value. + fs.writeFileSync(path.join(tempDir, DOTENV_FILE_NAME), 'DEMO_DB_HOST=refreshed-host.example.com\n'); + + // The refresh is asynchronous (watcher debounce + hidden kernel exec), so re-run a bounded number of times + // until the refreshed value renders — rather than a fixed sleep + single run, which is flaky and burns the + // full timeout on a slow box. Re-running is safe here: the kernel is already bound after the first run. + const REFRESH_MAX_ATTEMPTS = 6; + const REFRESH_ATTEMPT_TIMEOUT = 10_000; + let second = ''; + for (let attempt = 1; attempt <= REFRESH_MAX_ATTEMPTS; attempt++) { + try { + second = await runOnceAndAwaitOutput( + NOTEBOOK_FILE_NAME, + 'refreshed-host.example.com', + REFRESH_ATTEMPT_TIMEOUT + ); + break; + } catch (error) { + if (attempt === REFRESH_MAX_ATTEMPTS) { + throw error; + } + } + } + expect(second).to.contain('refreshed-host.example.com'); + }); +});