diff --git a/cspell.json b/cspell.json index 8aa54f875..abf2ce302 100644 --- a/cspell.json +++ b/cspell.json @@ -98,6 +98,8 @@ "Unconfigured", "unuse", "unittests", + "userpod", + "Userpod", "vegalite", "venv", "Venv", diff --git a/package-lock.json b/package-lock.json index f8e9ef939..877b0a072 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 8dc22af21..f37d31f17 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 907adf8f8..8dee76451 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/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 112f7f0e0..207c9d800 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -19,9 +19,10 @@ import { IProcessServiceFactory } from '../../platform/common/process/types.node import { IAsyncDisposableRegistry, IDisposable, IOutputChannel } from '../../platform/common/types'; import { sleep } from '../../platform/common/utils/async'; import { generateUuid } from '../../platform/common/uuid'; +import { resolveProjectIdForFile } from '../../platform/deepnote/deepnoteProjectIdResolver'; import { DeepnoteServerStartupError } from '../../platform/errors/deepnoteKernelErrors'; import { logger } from '../../platform/logging'; -import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; +import { ISqlIntegrationEnvVarsProvider, IUserpodApiEndpoints } from '../../platform/notebooks/deepnote/types'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import * as path from '../../platform/vscode-path/path'; import { DeepnoteServerInfo, IDeepnoteServerStarter, IDeepnoteToolkitInstaller } from './types'; @@ -77,7 +78,10 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension @inject(IAsyncDisposableRegistry) asyncRegistry: IAsyncDisposableRegistry, @inject(ISqlIntegrationEnvVarsProvider) @optional() - private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider + private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider, + @inject(IUserpodApiEndpoints) + @optional() + private readonly userpodApiEndpoints?: IUserpodApiEndpoints ) { asyncRegistry.push(this); } @@ -262,6 +266,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension this.outputChannel.appendLine(l10n.t('Starting Deepnote server...')); const extraEnv = await this.gatherSqlIntegrationEnvVars(deepnoteFileUri, environmentId, token); + await this.applyIntegrationEndpointEnv(extraEnv, deepnoteFileUri); // Initialize output tracking for error reporting this.serverOutputByFile.set(fileKey, { stdout: '', stderr: '' }); @@ -397,6 +402,41 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension return extraEnv; } + // Skipped unless the endpoint is up and the file has a project id — else the toolkit raises on an unreachable URL. + private async applyIntegrationEndpointEnv(extraEnv: Record, deepnoteFileUri: Uri): Promise { + const endpoint = this.userpodApiEndpoints; + + if (!endpoint) { + return; + } + + // Wait for the initial bind so a kernel starting before the loopback endpoint is listening still gets the env (F3). + await endpoint.ready; + + const baseUrl = endpoint.baseUrl; + if (!baseUrl) { + logger.warn( + 'DeepnoteServerStarter: integration endpoint is not listening; skipping live integration env injection.' + ); + + return; + } + + const projectId = await resolveProjectIdForFile(deepnoteFileUri); + + if (!projectId) { + return; + } + + extraEnv['DEEPNOTE_RUNTIME__ENV_INTEGRATION_ENABLED'] = 'true'; + extraEnv['DEEPNOTE_RUNTIME__RUNNING_IN_DETACHED_MODE'] = 'true'; + extraEnv['DEEPNOTE_RUNTIME__WEBAPP_URL'] = baseUrl; + // 2.1.1 dereferences project_secret without a null-check in detached mode; also the endpoint's per-project bearer token. + extraEnv['DEEPNOTE_RUNTIME__PROJECT_SECRET'] = endpoint.getAuthToken(projectId); + // Legacy key (not __PROJECT_ID): also satisfies set_notebook_path's has_env check, avoiding a session-name parse. + extraEnv['DEEPNOTE_PROJECT_ID'] = projectId; + } + /** * Stream stdout/stderr from the server process to the VSCode output channel. */ diff --git a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts index d0eb9070c..fa9ee5413 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts @@ -1,21 +1,24 @@ import { assert } from 'chai'; import * as fakeTimers from '@sinonjs/fake-timers'; import * as sinon from 'sinon'; -import { anything, instance, mock, when } from 'ts-mockito'; +import { anything, instance, mock, verify, when } from 'ts-mockito'; import { CancellationError, Uri } from 'vscode'; +import { serializeDeepnoteFile, type DeepnoteFile } from '@deepnote/blocks'; + import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; import { DeepnoteServerStarter } from './deepnoteServerStarter.node'; import { IProcessServiceFactory } from '../../platform/common/process/types.node'; import { IAsyncDisposableRegistry, IOutputChannel } from '../../platform/common/types'; import { IDeepnoteToolkitInstaller } from './types'; -import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; +import { ISqlIntegrationEnvVarsProvider, IUserpodApiEndpoints } from '../../platform/notebooks/deepnote/types'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { __getStartServerCalls, __getStopServerCalls, __resetRuntimeCoreMock } from '../../test/mocks/deepnoteRuntimeCore'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; /** * Unit tests for DeepnoteServerStarter. @@ -264,4 +267,153 @@ suite('DeepnoteServerStarter', () => { }); }); }); + + /** + * Direct tests for the private `applyIntegrationEndpointEnv`, exercised via a cast so we mutate a + * plain env object rather than intercepting the third-party `startServer`. The four integration env + * vars are only injected when the loopback endpoint is listening AND the file resolves to a project + * id; every other path degrades gracefully to whatever SQL_* vars were already gathered. + */ + suite('applyIntegrationEndpointEnv (integration endpoint env injection)', () => { + const projectFileUri = Uri.file('/workspace/project/notebook-a.deepnote'); + const baseUrl = 'http://127.0.0.1:5555'; + // A pre-seeded SQL_* var stands in for the env already gathered before the endpoint step runs. + const sqlEnvKey = 'SQL_DEEPNOTE_INTEGRATION_ABC'; + const sqlEnvValue = 'postgres://localhost:5432/db'; + + type WithApplyIntegrationEndpointEnv = { + applyIntegrationEndpointEnv(extraEnv: Record, deepnoteFileUri: Uri): Promise; + }; + + // Starters built with an endpoint arg are tracked so each is disposed after its test. + let extraStarters: DeepnoteServerStarter[]; + + setup(() => { + resetVSCodeMocks(); + extraStarters = []; + }); + + teardown(async () => { + await Promise.all(extraStarters.map((starter) => starter.dispose())); + }); + + /** + * Builds a starter wired with the integration endpoint (optional ctor arg at positional index 6), + * whose `baseUrl` getter yields `endpointBaseUrl`. Tracked for disposal. + */ + function createStarterWithEndpoint(endpointBaseUrl: string | undefined): DeepnoteServerStarter { + const endpoint: IUserpodApiEndpoints = { + baseUrl: endpointBaseUrl, + ready: Promise.resolve(), + getAuthToken: () => 'endpoint-token' + }; + const starter = new DeepnoteServerStarter( + instance(mockProcessServiceFactory), + instance(mockToolkitInstaller), + instance(mockAgentSkillsManager), + instance(mockOutputChannel), + instance(mockAsyncRegistry), + instance(mockSqlIntegrationEnvVars), + endpoint + ); + extraStarters.push(starter); + + return starter; + } + + function applyIntegrationEndpointEnv( + starter: DeepnoteServerStarter, + extraEnv: Record, + deepnoteFileUri: Uri + ): Promise { + return (starter as unknown as WithApplyIntegrationEndpointEnv).applyIntegrationEndpointEnv( + extraEnv, + deepnoteFileUri + ); + } + + /** + * Stubs `workspace.fs.readFile` — the read behind `resolveProjectIdForFile` — to yield the + * serialized `.deepnote` bytes. Returns the mock so tests can assert whether the read happened. + */ + function stubReadFile(fileContents: string): typeof import('vscode').workspace.fs { + const mockFs = mock(); + + when(mockFs.readFile(anything())).thenReturn(Promise.resolve(new TextEncoder().encode(fileContents))); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + return mockFs; + } + + function serializeProjectFile(projectId: string): string { + const file: DeepnoteFile = { + metadata: { + createdAt: '2023-01-01T00:00:00Z', + modifiedAt: '2023-01-02T00:00:00Z' + }, + project: { + id: projectId, + name: 'Project', + notebooks: [{ id: 'notebook-1', name: 'Notebook One', blocks: [] }], + settings: {} + }, + version: '1.0.0' + }; + + return serializeDeepnoteFile(file); + } + + test('injects all five integration env vars (preserving pre-seeded SQL_* keys) when the endpoint is listening and the file has a project id', async () => { + const mockFs = stubReadFile(serializeProjectFile('the-project-id')); + const starter = createStarterWithEndpoint(baseUrl); + + const extraEnv: Record = { [sqlEnvKey]: sqlEnvValue }; + await applyIntegrationEndpointEnv(starter, extraEnv, projectFileUri); + + assert.deepStrictEqual(extraEnv, { + [sqlEnvKey]: sqlEnvValue, + DEEPNOTE_RUNTIME__ENV_INTEGRATION_ENABLED: 'true', + DEEPNOTE_RUNTIME__RUNNING_IN_DETACHED_MODE: 'true', + DEEPNOTE_RUNTIME__WEBAPP_URL: baseUrl, + DEEPNOTE_RUNTIME__PROJECT_SECRET: 'endpoint-token', + DEEPNOTE_PROJECT_ID: 'the-project-id' + }); + // The enabled path must resolve the project id from the file. + verify(mockFs.readFile(anything())).once(); + }); + + test('injects nothing and does NOT read the file when the endpoint has no baseUrl', async () => { + const mockFs = stubReadFile(serializeProjectFile('the-project-id')); + const starter = createStarterWithEndpoint(undefined); + + const extraEnv: Record = { [sqlEnvKey]: sqlEnvValue }; + await applyIntegrationEndpointEnv(starter, extraEnv, projectFileUri); + + assert.deepStrictEqual(extraEnv, { [sqlEnvKey]: sqlEnvValue }); + // A missing baseUrl short-circuits BEFORE resolving the project id — no file read. + verify(mockFs.readFile(anything())).never(); + }); + + test('injects nothing and does NOT read the file when no endpoint was injected at all', async () => { + // The outer-suite `serverStarter` is constructed WITHOUT the optional endpoint arg. + const mockFs = stubReadFile(serializeProjectFile('the-project-id')); + + const extraEnv: Record = { [sqlEnvKey]: sqlEnvValue }; + await applyIntegrationEndpointEnv(serverStarter, extraEnv, projectFileUri); + + assert.deepStrictEqual(extraEnv, { [sqlEnvKey]: sqlEnvValue }); + verify(mockFs.readFile(anything())).never(); + }); + + test('injects nothing when the endpoint is listening but the file resolves to no project id', async () => { + // A schema-valid `.deepnote` whose project.id is empty — `resolveProjectIdForFile` yields a falsy id. + stubReadFile(serializeProjectFile('')); + const starter = createStarterWithEndpoint(baseUrl); + + const extraEnv: Record = { [sqlEnvKey]: sqlEnvValue }; + await applyIntegrationEndpointEnv(starter, extraEnv, projectFileUri); + + assert.deepStrictEqual(extraEnv, { [sqlEnvKey]: sqlEnvValue }); + }); + }); }); diff --git a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts index 0c3c04830..a55dfeb7c 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 5715ead44..9e872066d 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 000000000..cf6e946bc --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.ts @@ -0,0 +1,64 @@ +import { inject, injectable } from 'inversify'; +import { l10n, type NotebookDocument, window } from 'vscode'; + +import { IKernelProvider } from '../../../kernels/types'; +import { logger } from '../../../platform/logging'; +import { IIntegrationEnvLiveRefresher } from './types'; + +/** `_dntk` isn't guaranteed here, so import the package to re-fetch + apply integration env in the kernel. */ +const REFRESH_INTEGRATION_ENV_SNIPPET = `import deepnote_toolkit +deepnote_toolkit.set_integration_env()`; + +/** How long the transient "environment updated" status-bar message stays visible. */ +const STATUS_BAR_MESSAGE_TIMEOUT_MS = 5000; + +@injectable() +export class IntegrationEnvLiveRefresher implements IIntegrationEnvLiveRefresher { + constructor(@inject(IKernelProvider) private readonly kernelProvider: IKernelProvider) {} + + 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 + ); + } + } + + /** 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 outputs = await this.kernelProvider + .getKernelExecution(kernel) + .executeHidden(REFRESH_INTEGRATION_ENV_SNIPPET); + + const errors = outputs.filter((output) => output.output_type === 'error'); + if (errors.length > 0) { + logger.warn( + `IntegrationEnvLiveRefresher: Refresh snippet produced errors for ${notebook.uri.toString()}`, + errors + ); + + return false; + } + + return true; + } catch (err) { + logger.error( + `IntegrationEnvLiveRefresher: Failed to refresh integration env for ${notebook.uri.toString()}`, + err + ); + + return false; + } + } +} 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 000000000..ad1c418e0 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.unit.test.ts @@ -0,0 +1,149 @@ +import { assert } from 'chai'; +import * as sinon from 'sinon'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; +import { Disposable, 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 { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; +import { IntegrationEnvLiveRefresher } from './integrationEnvLiveRefresher.node'; + +// Must match REFRESH_INTEGRATION_ENV_SNIPPET in integrationEnvLiveRefresher.node.ts exactly (real newline). +const EXPECTED_SNIPPET = `import deepnote_toolkit +deepnote_toolkit.set_integration_env()`; + +const EXPECTED_NOTIFICATION = 'Deepnote integration environment updated.'; + +suite('IntegrationEnvLiveRefresher', () => { + let refresher: IntegrationEnvLiveRefresher; + let kernelProvider: IKernelProvider; + let executeHiddenSpy: sinon.SinonStub; + let disposables: IDisposable[]; + + setup(() => { + resetVSCodeMocks(); + disposables = [new Disposable(() => resetVSCodeMocks())]; + kernelProvider = mock(); + + executeHiddenSpy = sinon.stub().resolves([]); + when(kernelProvider.getKernelExecution(anything())).thenReturn({ + executeHidden: executeHiddenSpy + } as unknown as INotebookKernelExecution); + + refresher = new IntegrationEnvLiveRefresher(instance(kernelProvider)); + }); + + teardown(() => { + disposables = dispose(disposables); + }); + + /** A notebook whose kernel is started; `kernelProvider.get(notebook)` returns it. */ + function createRunningNotebook(uri: Uri): NotebookDocument { + const notebookMock = mock(); + when(notebookMock.uri).thenReturn(uri); + const notebook = instance(notebookMock); + + const kernelMock = mock(); + when(kernelMock.startedAtLeastOnce).thenReturn(true); + when(kernelProvider.get(notebook)).thenReturn(instance(kernelMock)); + + return notebook; + } + + test('runs the exact refresh snippet 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(executeHiddenSpy.callCount, 1, 'the refresh snippet should run once'); + assert.strictEqual( + executeHiddenSpy.firstCall.args[0], + EXPECTED_SNIPPET, + 'executeHidden must receive the toolkit set_integration_env() snippet verbatim' + ); + + 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(executeHiddenSpy.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(executeHiddenSpy.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 refresh snippet produces an error output', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + executeHiddenSpy.resolves([{ output_type: 'error', ename: 'RuntimeError', evalue: 'boom', traceback: [] }]); + + await refresher.refresh([notebook]); + + assert.strictEqual(executeHiddenSpy.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(executeHiddenSpy.callCount, 2, 'both started kernels are refreshed'); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).once(); + }); + + test('continues to the next notebook when one executeHidden throws, and still notifies for the success', async () => { + const notebookA = createRunningNotebook(Uri.file('/ws/a.deepnote')); + const notebookB = createRunningNotebook(Uri.file('/ws/b.deepnote')); + executeHiddenSpy.onFirstCall().rejects(new Error('kernel exploded')); + executeHiddenSpy.onSecondCall().resolves([]); + + await refresher.refresh([notebookA, notebookB]); + + assert.strictEqual(executeHiddenSpy.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)); + executeHiddenSpy.onFirstCall().callsFake(() => secondInvoked.then(() => [])); + executeHiddenSpy.onSecondCall().callsFake(() => { + markSecondInvoked(); + + return Promise.resolve([]); + }); + + await refresher.refresh([notebookA, notebookB]); + + assert.strictEqual(executeHiddenSpy.callCount, 2, 'both started kernels are refreshed'); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).once(); + }); +}); diff --git a/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.ts b/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.ts new file mode 100644 index 000000000..c4806e87c --- /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 000000000..d7dcb3ace --- /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/integrationKernelRestartHandler.ts b/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.ts deleted file mode 100644 index b082d0d19..000000000 --- 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 62888c3fe..000000000 --- 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 000000000..cdae27fda --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.ts @@ -0,0 +1,167 @@ +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). */ +@injectable() +export class IntegrationsEnvFileWatcher implements IExtensionSyncActivationService { + private readonly changedDirs = new Set(); + private debounceTimer: ReturnType | undefined; + 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): Promise { + const affected = await this.findAffectedNotebooks(changedDirs); + if (affected.length === 0) { + return; + } + + await this.liveRefresher.refresh(affected); + } + + private async findAffectedNotebooks(changedDirs: 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; + } + + // A `.env` change 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 (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): void { + this.changedDirs.add(dir.fsPath); + + if (this.debounceTimer) { + clearTimeout(this.debounceTimer); + } + + this.debounceTimer = setTimeout(() => { + this.debounceTimer = undefined; + const dirs = new Set(this.changedDirs); + this.changedDirs.clear(); + this.handleChangedDirs(dirs).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)), + watcher.onDidCreate(() => this.onFileEvent(dir)), + watcher.onDidDelete(() => this.onFileEvent(dir)) + ); + } + } + + 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 000000000..7457a2f03 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.unit.test.ts @@ -0,0 +1,150 @@ +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(); + }); +}); diff --git a/src/notebooks/deepnote/integrations/types.ts b/src/notebooks/deepnote/integrations/types.ts index 3e151cb78..2196ef61b 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 { + /** Re-runs the toolkit's `set_integration_env()` in each notebook's 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/integrations/userpodApiEndpoints.node.ts b/src/notebooks/deepnote/integrations/userpodApiEndpoints.node.ts new file mode 100644 index 000000000..d2af72e50 --- /dev/null +++ b/src/notebooks/deepnote/integrations/userpodApiEndpoints.node.ts @@ -0,0 +1,215 @@ +import { timingSafeEqual } from 'crypto'; +import express, { type Express, type Request, type Response } from 'express'; +import * as http from 'http'; +import { inject, injectable } from 'inversify'; +import { commands, l10n, window, workspace } from 'vscode'; + +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; +import { IExtensionSyncActivationService } from '../../../platform/activation/types'; +import { IDisposableRegistry } from '../../../platform/common/types'; +import { generateUuid } from '../../../platform/common/uuid'; +import { logger } from '../../../platform/logging'; +import { ISqlIntegrationEnvVarsProvider, IUserpodApiEndpoints } from '../../../platform/notebooks/deepnote/types'; + +/** Loopback host for the toolkit's `userpod-api` calls; currently serves integration env vars for `set_integration_env()` (as `[{name,value}]`). */ +@injectable() +export class UserpodApiEndpoints implements IUserpodApiEndpoints, IExtensionSyncActivationService { + private readonly authTokensByProject = new Map(); + private isListening = false; + private server: http.Server | undefined; + private serverBaseUrl: string | undefined; + private startAttempt: Promise = Promise.resolve(); + + constructor( + @inject(ISqlIntegrationEnvVarsProvider) + private readonly sqlIntegrationEnvVarsProvider: ISqlIntegrationEnvVarsProvider, + @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry + ) {} + + public get baseUrl(): string | undefined { + return this.serverBaseUrl; + } + + /** Settles (never rejects) once the initial bind attempt completes, so callers can await readiness before reading `baseUrl`. */ + public get ready(): Promise { + return this.startAttempt; + } + + /** Per-project bearer token, generated on first use, so a kernel can only read its own project's credentials. */ + public getAuthToken(projectId: string): string { + let token = this.authTokensByProject.get(projectId); + if (token === undefined) { + token = generateUuid(); + this.authTokensByProject.set(projectId, token); + } + + return token; + } + + public activate(): void { + this.disposables.push({ dispose: () => this.stop() }); + this.startAttempt = this.start().catch((err) => + logger.error('UserpodApiEndpoints: Failed to start integration env vars endpoint', err) + ); + } + + private isAuthorized(authorizationHeader: string | undefined, projectId: string): boolean { + if (authorizationHeader === undefined) { + return false; + } + + const expectedToken = this.authTokensByProject.get(projectId); + if (expectedToken === undefined) { + return false; + } + + // Constant-time compare: the response carries credentials, so don't leak the token via early-exit timing. + const expected = Buffer.from(`Bearer ${expectedToken}`); + const actual = Buffer.from(authorizationHeader); + + return actual.length === expected.length && timingSafeEqual(actual, expected); + } + + private onServerError(err: Error): void { + logger.error('UserpodApiEndpoints: HTTP server error', err); + + // Startup failures are surfaced by start()'s rejection; only recover from errors once actually listening. + if (!this.isListening) { + return; + } + this.isListening = false; + this.serverBaseUrl = undefined; + + const restart = l10n.t('Restart'); + const reloadWindow = l10n.t('Reload Window'); + + void window + .showErrorMessage( + l10n.t( + 'The Deepnote integrations service stopped unexpectedly. Integration environment variables will not update until it restarts.' + ), + restart, + reloadWindow + ) + .then((choice) => { + if (choice === restart) { + this.restart(); + } else if (choice === reloadWindow) { + void commands.executeCommand('workbench.action.reloadWindow'); + } + }); + } + + private restart(): void { + this.stop(); + this.startAttempt = this.start().catch((err) => + logger.error('UserpodApiEndpoints: Failed to restart integration env vars endpoint', err) + ); + } + + private async start(): Promise { + const app: Express = express(); + + app.get('/userpod-api/:projectId/integrations/environment-variables', async (req: Request, res: Response) => { + try { + // The `:projectId` route segment is always a single value at runtime; narrow the over-broad express type. + const projectId = Array.isArray(req.params.projectId) ? req.params.projectId[0] : req.params.projectId; + + if (!this.isAuthorized(req.headers.authorization, projectId)) { + res.status(401).json([]); + + return; + } + + // Filter (not find): sibling `.deepnote` files can share a project id, so serve the project's env + // vars rather than an arbitrary first match, merged deterministically by notebook uri (F1). + const notebooks = workspace.notebookDocuments.filter( + (nb) => nb.notebookType === DEEPNOTE_NOTEBOOK_TYPE && nb.metadata?.deepnoteProjectId === projectId + ); + + if (notebooks.length === 0) { + res.json([]); + + return; + } + + if (notebooks.length > 1) { + logger.warn( + `UserpodApiEndpoints: ${notebooks.length} open notebooks share project '${projectId}'; merging their integration env vars.` + ); + } + + const ordered = [...notebooks].sort((a, b) => a.uri.toString().localeCompare(b.uri.toString())); + const resolved = await Promise.all( + ordered.map((notebook) => this.sqlIntegrationEnvVarsProvider.getEnvironmentVariables(notebook.uri)) + ); + + const merged = new Map(); + for (const envVars of resolved) { + for (const [name, value] of Object.entries(envVars ?? {})) { + if (!merged.has(name)) { + merged.set(name, value); + } + } + } + + const payload = [...merged].map(([name, value]) => ({ name, value })); + + res.json(payload); + } catch (err) { + logger.error('UserpodApiEndpoints: Failed to resolve integration environment variables', err); + res.status(500).send('Failed to resolve integration environment variables'); + } + }); + + const server = http.createServer(app); + this.server = server; + + // Persistent handler: an 'error' with no listener is re-thrown as an uncaught exception that crashes the host. + server.on('error', (err) => this.onServerError(err)); + + server.listen(0, '127.0.0.1'); + + const port = await new Promise((resolve, reject) => { + let onStartupError: (err: Error) => void = () => undefined; + let onListening: () => void = () => undefined; + onStartupError = (err: Error) => { + server.removeListener('listening', onListening); + reject(err); + }; + onListening = () => { + server.removeListener('error', onStartupError); + const address = server.address(); + if (!address || typeof address === 'string') { + reject(new Error('Integration env vars endpoint did not bind a port.')); + + return; + } + this.isListening = true; + resolve(address.port); + }; + server.once('error', onStartupError); + server.once('listening', onListening); + }); + + this.serverBaseUrl = `http://127.0.0.1:${port}`; + logger.info(`UserpodApiEndpoints: Listening on ${this.serverBaseUrl}`); + } + + private stop(): void { + const server = this.server; + this.server = undefined; + this.serverBaseUrl = undefined; + this.isListening = false; + if (!server) { + return; + } + + try { + server.closeAllConnections(); + server.close(); + } catch (err) { + logger.warn('UserpodApiEndpoints: Error while stopping server', err); + } + } +} diff --git a/src/notebooks/deepnote/integrations/userpodApiEndpoints.node.unit.test.ts b/src/notebooks/deepnote/integrations/userpodApiEndpoints.node.unit.test.ts new file mode 100644 index 000000000..79a7ee3f4 --- /dev/null +++ b/src/notebooks/deepnote/integrations/userpodApiEndpoints.node.unit.test.ts @@ -0,0 +1,240 @@ +import * as http from 'http'; + +import { assert } from 'chai'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; +import { Disposable, NotebookDocument, Uri } from 'vscode'; + +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; +import { IDisposable } from '../../../platform/common/types'; +import { dispose } from '../../../platform/common/utils/lifecycle'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; +import { UserpodApiEndpoints } from './userpodApiEndpoints.node'; + +suite('UserpodApiEndpoints', () => { + let endpoint: UserpodApiEndpoints; + let provider: ISqlIntegrationEnvVarsProvider; + let disposables: IDisposable[]; + + setup(() => { + resetVSCodeMocks(); + disposables = [new Disposable(() => resetVSCodeMocks())]; + provider = mock(); + + endpoint = new UserpodApiEndpoints(instance(provider), disposables); + }); + + teardown(() => { + // Disposing tears down the HTTP server (a dispose handler is pushed into `disposables` at start). + disposables = dispose(disposables); + }); + + function createMockNotebook( + projectId: string | undefined, + uri: Uri, + notebookType: string = DEEPNOTE_NOTEBOOK_TYPE + ): NotebookDocument { + const notebook = mock(); + when(notebook.notebookType).thenReturn(notebookType); + when(notebook.metadata).thenReturn(projectId === undefined ? {} : { deepnoteProjectId: projectId }); + when(notebook.uri).thenReturn(uri); + + return instance(notebook); + } + + /** `activate()` is fire-and-forget; poll `baseUrl` until the server has bound a port. */ + async function waitForBaseUrl(timeoutMs = 3000): Promise { + const start = Date.now(); + while (endpoint.baseUrl === undefined) { + if (Date.now() - start > timeoutMs) { + throw new Error('UserpodApiEndpoints did not start listening in time'); + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + + return endpoint.baseUrl; + } + + function envVarsUrl(baseUrl: string, projectId: string): string { + return `${baseUrl}/userpod-api/${projectId}/integrations/environment-variables`; + } + + /** GET the endpoint carrying the per-project bearer token it requires. */ + function authedFetch(url: string, projectId: string): Promise { + return fetch(url, { headers: { Authorization: `Bearer ${endpoint.getAuthToken(projectId)}` } }); + } + + test('returns the provider env map as [{name,value}] for a matching open deepnote notebook', async () => { + const notebook = createMockNotebook('project-1', Uri.file('/ws/app.deepnote')); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(provider.getEnvironmentVariables(anything())).thenResolve({ FOO: 'bar', BAZ: 'qux' }); + + endpoint.activate(); + const baseUrl = await waitForBaseUrl(); + + const response = await authedFetch(envVarsUrl(baseUrl, 'project-1'), 'project-1'); + + assert.strictEqual(response.status, 200); + const body = await response.json(); + assert.deepStrictEqual(body, [ + { name: 'FOO', value: 'bar' }, + { name: 'BAZ', value: 'qux' } + ]); + }); + + test('returns an empty array when the matching notebook has no integration env vars', async () => { + const notebook = createMockNotebook('project-1', Uri.file('/ws/app.deepnote')); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(provider.getEnvironmentVariables(anything())).thenResolve({}); + + endpoint.activate(); + const baseUrl = await waitForBaseUrl(); + + const response = await authedFetch(envVarsUrl(baseUrl, 'project-1'), 'project-1'); + + assert.strictEqual(response.status, 200); + assert.deepStrictEqual(await response.json(), []); + }); + + test('routes by projectId: queries only the notebook whose metadata matches the URL param', async () => { + const uriTwo = Uri.file('/ws/two.deepnote'); + const notebookOne = createMockNotebook('project-one', Uri.file('/ws/one.deepnote')); + const notebookTwo = createMockNotebook('project-two', uriTwo); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookOne, notebookTwo]); + when(provider.getEnvironmentVariables(anything())).thenResolve({ FROM: 'two' }); + + endpoint.activate(); + const baseUrl = await waitForBaseUrl(); + + const response = await authedFetch(envVarsUrl(baseUrl, 'project-two'), 'project-two'); + const body = await response.json(); + + assert.deepStrictEqual(body, [{ name: 'FROM', value: 'two' }]); + // The endpoint must resolve env vars for the matching notebook's uri, not the other project's. + verify(provider.getEnvironmentVariables(anything())).once(); + const [uriArg] = capture(provider.getEnvironmentVariables).last(); + assert.strictEqual((uriArg as Uri).toString(), uriTwo.toString()); + }); + + test('returns an empty array and never queries the provider when no notebook matches the projectId', async () => { + const notebook = createMockNotebook('some-other-project', Uri.file('/ws/app.deepnote')); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + endpoint.activate(); + const baseUrl = await waitForBaseUrl(); + + const response = await authedFetch(envVarsUrl(baseUrl, 'project-1'), 'project-1'); + + assert.strictEqual(response.status, 200); + assert.deepStrictEqual(await response.json(), []); + verify(provider.getEnvironmentVariables(anything())).never(); + }); + + test('ignores non-Deepnote notebooks even when their projectId matches', async () => { + const notebook = createMockNotebook('project-1', Uri.file('/ws/app.ipynb'), 'jupyter-notebook'); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + endpoint.activate(); + const baseUrl = await waitForBaseUrl(); + + const response = await authedFetch(envVarsUrl(baseUrl, 'project-1'), 'project-1'); + + assert.deepStrictEqual(await response.json(), []); + verify(provider.getEnvironmentVariables(anything())).never(); + }); + + test('responds 500 when the provider rejects', async () => { + const notebook = createMockNotebook('project-1', Uri.file('/ws/app.deepnote')); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(provider.getEnvironmentVariables(anything())).thenReject(new Error('resolution failed')); + + endpoint.activate(); + const baseUrl = await waitForBaseUrl(); + + const response = await authedFetch(envVarsUrl(baseUrl, 'project-1'), 'project-1'); + + assert.strictEqual(response.status, 500); + }); + + test('responds 401 and never queries the provider when the bearer token is missing or wrong', async () => { + const notebook = createMockNotebook('project-1', Uri.file('/ws/app.deepnote')); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + endpoint.activate(); + const baseUrl = await waitForBaseUrl(); + + const noHeader = await fetch(envVarsUrl(baseUrl, 'project-1')); + assert.strictEqual(noHeader.status, 401); + + const wrongToken = await fetch(envVarsUrl(baseUrl, 'project-1'), { + headers: { Authorization: 'Bearer wrong-token' } + }); + assert.strictEqual(wrongToken.status, 401); + + verify(provider.getEnvironmentVariables(anything())).never(); + }); + + test('responds 401 for a wrong token of the SAME length (exercises the constant-time compare path)', async () => { + const notebook = createMockNotebook('project-1', Uri.file('/ws/app.deepnote')); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + endpoint.activate(); + const baseUrl = await waitForBaseUrl(); + + // Issue the project's token, then present a DIFFERENT value of the same byte length — timingSafeEqual must reject it. + const realToken = endpoint.getAuthToken('project-1'); + const sameLengthWrong = `Bearer ${'x'.repeat(realToken.length)}`; + const response = await fetch(envVarsUrl(baseUrl, 'project-1'), { + headers: { Authorization: sameLengthWrong } + }); + + assert.strictEqual(response.status, 401); + verify(provider.getEnvironmentVariables(anything())).never(); + }); + + test('cross-project: a token issued for one project cannot read another project', async () => { + const notebookA = createMockNotebook('project-a', Uri.file('/ws/a.deepnote')); + const notebookB = createMockNotebook('project-b', Uri.file('/ws/b.deepnote')); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookA, notebookB]); + + endpoint.activate(); + const baseUrl = await waitForBaseUrl(); + + // Issue tokens for both projects, then use project A's token against project B's URL. + const tokenA = endpoint.getAuthToken('project-a'); + endpoint.getAuthToken('project-b'); + + const response = await fetch(envVarsUrl(baseUrl, 'project-b'), { + headers: { Authorization: `Bearer ${tokenA}` } + }); + + assert.strictEqual(response.status, 401, "project A's token must not authorize a read of project B"); + verify(provider.getEnvironmentVariables(anything())).never(); + }); + + test('ready resolves once the server is listening', async () => { + endpoint.activate(); + + await endpoint.ready; + + assert.isString(endpoint.baseUrl, 'baseUrl must be set once ready resolves'); + }); + + test('logs and prompts the user to recover when the server errors after startup, without crashing', async () => { + const notebook = createMockNotebook('project-1', Uri.file('/ws/app.deepnote')); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(mockedVSCodeNamespaces.window.showErrorMessage(anything(), anything(), anything())).thenResolve( + undefined as never + ); + + endpoint.activate(); + await waitForBaseUrl(); + + // Reach the running server to simulate a post-listen failure (e.g. an accept error under fd exhaustion). + const server = (endpoint as unknown as { server: http.Server }).server; + server.emit('error', new Error('accept failed')); + + assert.strictEqual(endpoint.baseUrl, undefined, 'a crashed endpoint must stop advertising its base URL'); + verify(mockedVSCodeNamespaces.window.showErrorMessage(anything(), anything(), anything())).once(); + }); +}); diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts index 5bf312977..4c94f450c 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 eb87ce58a..bf07345c7 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,8 +63,10 @@ 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 + IPlatformDeepnoteNotebookManager, + IUserpodApiEndpoints } from '../platform/notebooks/deepnote/types'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; import { DirtyInputBlockStatusBarProvider } from './deepnote/dirtyInputBlockStatusBarProvider'; @@ -100,7 +104,10 @@ 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 { UserpodApiEndpoints } from './deepnote/integrations/userpodApiEndpoints.node'; import { ISnapshotMetadataService, SnapshotService } from './deepnote/snapshots/snapshotService'; import { EnvironmentCapture, IEnvironmentCapture } from './deepnote/snapshots/environmentCapture.node'; import { DeepnoteFileChangeWatcher } from './deepnote/deepnoteFileChangeWatcher'; @@ -233,8 +240,22 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea ); serviceManager.addSingleton( IExtensionSyncActivationService, - IntegrationKernelRestartHandler + IntegrationEnvRefreshHandler ); + serviceManager.addSingleton( + IIntegrationsFileConfigProvider, + IntegrationsFileConfigProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + IntegrationsEnvFileWatcher + ); + serviceManager.addSingleton( + IIntegrationEnvLiveRefresher, + IntegrationEnvLiveRefresher + ); + serviceManager.addSingleton(IUserpodApiEndpoints, UserpodApiEndpoints); + serviceManager.addBinding(IUserpodApiEndpoints, IExtensionSyncActivationService); // Deepnote kernel services serviceManager.addSingleton(DeepnoteAgentSkillsManager, DeepnoteAgentSkillsManager); diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index 28937d6b6..6d6b82618 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 8eb24affc..ad92029ea 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 000000000..b21319a33 --- /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 000000000..1d9a2af74 --- /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 7a612e927..8da5e213c 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 11b5e57ea..ba07c58dd 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 23be88f5d..f0cf89e6f 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,37 @@ 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 IUserpodApiEndpoints = Symbol('IUserpodApiEndpoints'); +export interface IUserpodApiEndpoints { + /** Loopback base URL the toolkit fetches integration env vars from; `undefined` until the server is listening. */ + readonly baseUrl: string | undefined; + + /** Settles (never rejects) once the initial bind attempt completes, so callers can await readiness before reading `baseUrl`. */ + readonly ready: Promise; + + /** Per-project bearer token the endpoint requires (it serves credentials); injected into the toolkit as `DEEPNOTE_RUNTIME__PROJECT_SECRET`. */ + getAuthToken(projectId: string): string; +} + +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 000000000..5ca88bb04 --- /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 000000000..d79feb59f --- /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 runs set_integration_env() 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'); + }); +});