diff --git a/cspell.json b/cspell.json index 8aa54f8759..abf2ce3021 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 f8e9ef939b..877b0a0722 100644 --- a/package-lock.json +++ b/package-lock.json @@ -41,6 +41,7 @@ "clsx": "^2.1.1", "cross-fetch": "^3.1.5", "d3-format": "^3.1.0", + "dotenv": "^17.2.3", "encoding": "^0.1.13", "express": "^5.2.1", "fast-deep-equal": "^2.0.1", @@ -15521,6 +15522,18 @@ "url": "https://github.com/fb55/domutils?sponsor=1" } }, + "node_modules/dotenv": { + "version": "17.4.2", + "resolved": "https://registry.npmjs.org/dotenv/-/dotenv-17.4.2.tgz", + "integrity": "sha512-nI4U3TottKAcAD9LLud4Cb7b2QztQMUEfHbvhTH09bqXTxnSie8WnjPALV/WMCrJZ6UV/qHJ6L03OqO3LcdYZw==", + "license": "BSD-2-Clause", + "engines": { + "node": ">=12" + }, + "funding": { + "url": "https://dotenvx.com" + } + }, "node_modules/dunder-proto": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/dunder-proto/-/dunder-proto-1.0.1.tgz", @@ -46513,6 +46526,11 @@ "domhandler": "^5.0.3" } }, + "dotenv": { + "version": "17.4.2", + "resolved": "https://registry.npmjs.org/dotenv/-/dotenv-17.4.2.tgz", + "integrity": "sha512-nI4U3TottKAcAD9LLud4Cb7b2QztQMUEfHbvhTH09bqXTxnSie8WnjPALV/WMCrJZ6UV/qHJ6L03OqO3LcdYZw==" + }, "dunder-proto": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/dunder-proto/-/dunder-proto-1.0.1.tgz", diff --git a/package.json b/package.json index 8dc22af21d..f37d31f179 100644 --- a/package.json +++ b/package.json @@ -1649,6 +1649,12 @@ "description": "When enabled, outputs are saved to separate snapshot files in a 'snapshots' folder instead of the main .deepnote file.", "scope": "resource" }, + "deepnote.integrations.envFile.enabled": { + "type": "boolean", + "default": true, + "description": "When enabled, integration credentials are also loaded from a '.deepnote.env.yaml' file (with 'env:' references resolved against '.env' and environment variables) next to the .deepnote file or at the workspace root.", + "scope": "resource" + }, "deepnote.experiments.enabled": { "type": "boolean", "default": true, @@ -2703,6 +2709,7 @@ "clsx": "^2.1.1", "cross-fetch": "^3.1.5", "d3-format": "^3.1.0", + "dotenv": "^17.2.3", "encoding": "^0.1.13", "express": "^5.2.1", "fast-deep-equal": "^2.0.1", diff --git a/specs/INTEGRATIONS_CREDENTIALS.md b/specs/INTEGRATIONS_CREDENTIALS.md index 5b414342c6..6e6e193ca2 100644 --- a/specs/INTEGRATIONS_CREDENTIALS.md +++ b/specs/INTEGRATIONS_CREDENTIALS.md @@ -397,28 +397,23 @@ Handles migration of legacy integration configurations to the new `@deepnote/dat #### 2. **Integration Detector** (`integrationDetector.ts`) -Scans Deepnote projects to discover which integrations are used in SQL blocks. +Lists the integrations a Deepnote project declares, paired with the credentials stored for each. **Detection Process:** 1. Retrieves the Deepnote project from `IDeepnoteNotebookManager` -2. Scans all notebooks in the project -3. Examines each code block for `metadata.sql_integration_id` -4. Maps Deepnote integration types to `DatabaseIntegrationType` using the project's integration list -5. Checks if each integration is configured (has credentials) -6. Returns a map of integration IDs to their status - -**Integration Status:** - -- `Connected`: Integration has valid credentials stored -- `Disconnected`: Integration is used but not configured -- `Error`: Integration configuration is invalid +2. Iterates the project's `integrations` roster (ids, names and types only — never credentials) +3. Skips entries whose type is not a configurable `DatabaseIntegrationType` +4. Reads each integration's config from `SecretStorage`, or `null` when nothing is stored +5. Returns a map of integration IDs to their config and declared name/type **Special Cases:** - Excludes `deepnote-dataframe-sql` (internal DuckDB integration) -- Only processes code blocks with SQL integration metadata -- Uses project integration metadata to determine integration types +- Integrations configured only in `.deepnote.env.yaml` have a `null` config here, since those configs are never + written through `SecretStorage`. They still resolve normally for kernel execution and SQL autocomplete, which + read the merged configs directly — but the panel labels them "Not Configured" (see below), because it has no + visibility into the merged configs. #### 3. **Integration Manager** (`integrationManager.ts`) @@ -427,10 +422,6 @@ Orchestrates the integration management UI and commands. **Responsibilities:** - Registers the `deepnote.manageIntegrations` command -- Updates VSCode context keys for UI visibility: - - `deepnote.hasIntegrations`: True if any integrations are detected - - `deepnote.hasUnconfiguredIntegrations`: True if any integrations lack credentials -- Handles notebook selection changes - Opens the integration webview with detected integrations **Command Flow:** @@ -447,8 +438,16 @@ Provides the webview-based UI for managing integration credentials. **Features:** - Persistent webview panel (survives defocus) -- Real-time integration status updates +- Real-time updates as credentials are saved or cleared - Configuration forms for each integration type + +**Connection Status:** + +`IntegrationItem.tsx` derives the status pill from `config` alone — "Connected" when SecretStorage holds +credentials, "Not Configured" otherwise. The panel is a SecretStorage editor and is not given the merged +configs, so an integration configured only in `.deepnote.env.yaml` reads "Not Configured" here even though its +SQL cells run fine. The separate federated-auth pill (BigQuery + `google-oauth`) is independent and tracks +whether a refresh token is stored. - Delete/reset functionality **Message Protocol:** @@ -457,7 +456,7 @@ Extension → Webview: ```typescript // Update integration list -{ type: 'update', integrations: IntegrationWithStatus[] } +{ type: 'update', integrations: DetectedIntegration[] } // Show configuration form { type: 'showForm', integrationId: string, config: IntegrationConfig | null } @@ -487,7 +486,7 @@ Main React component that manages the webview UI state. **State Management:** -- `integrations`: List of detected integrations with status +- `integrations`: List of detected integrations with their stored configs - `selectedIntegrationId`: Currently selected integration for configuration - `selectedConfig`: Existing configuration being edited - `message`: Success/error messages @@ -501,7 +500,7 @@ Main React component that manages the webview UI state. 2. Panel shows configuration form overlay 3. User enters credentials 4. Panel sends save message to extension -5. Extension stores credentials and updates status +5. Extension stores credentials 6. Panel shows success message and refreshes list **Delete Integration:** @@ -511,7 +510,7 @@ Main React component that manages the webview UI state. 3. User clicks again to confirm 4. Panel sends delete message to extension 5. Extension removes credentials -6. Panel updates status to "Disconnected" +6. Panel clears the stored config, so the item offers "Configure" again #### 6. **Configuration Forms** diff --git a/src/kernels/deepnote/deepnoteIntegrationEndpointEnv.ts b/src/kernels/deepnote/deepnoteIntegrationEndpointEnv.ts new file mode 100644 index 0000000000..3da80bdbda --- /dev/null +++ b/src/kernels/deepnote/deepnoteIntegrationEndpointEnv.ts @@ -0,0 +1,47 @@ +import { Uri } from 'vscode'; + +import { resolveProjectIdForFile } from '../../platform/deepnote/deepnoteProjectIdResolver'; +import { logger } from '../../platform/logging'; +import { IUserpodApiEndpoints } from '../../platform/notebooks/deepnote/types'; + +/** + * Injects live-integration env vars into `extraEnv` when the loopback endpoint is listening and + * `deepnoteFileUri` resolves to a project id. Skipped otherwise — the toolkit raises on an unreachable URL. + * + * Mutates `extraEnv` in place. + */ +export async function applyIntegrationEndpointEnv({ + deepnoteFileUri, + endpoint, + extraEnv +}: { + deepnoteFileUri: Uri; + endpoint: IUserpodApiEndpoints; + extraEnv: Record; +}): Promise { + // Wait for the initial bind so a kernel starting before the loopback endpoint is listening still gets the env. + await endpoint.ready; + + const baseUrl = endpoint.baseUrl; + if (!baseUrl) { + logger.warn( + 'applyIntegrationEndpointEnv: 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; +} diff --git a/src/kernels/deepnote/deepnoteIntegrationEndpointEnv.unit.test.ts b/src/kernels/deepnote/deepnoteIntegrationEndpointEnv.unit.test.ts new file mode 100644 index 0000000000..c7044a70d1 --- /dev/null +++ b/src/kernels/deepnote/deepnoteIntegrationEndpointEnv.unit.test.ts @@ -0,0 +1,115 @@ +import { assert } from 'chai'; +import { anything, instance, mock, verify, when } from 'ts-mockito'; +import { Uri } from 'vscode'; + +import { serializeDeepnoteFile, type DeepnoteFile } from '@deepnote/blocks'; + +import { IUserpodApiEndpoints } from '../../platform/notebooks/deepnote/types'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; +import { applyIntegrationEndpointEnv } from './deepnoteIntegrationEndpointEnv'; + +/** + * 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 leaves `extraEnv` unchanged. + */ +suite('applyIntegrationEndpointEnv', () => { + 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'; + + setup(() => { + resetVSCodeMocks(); + }); + + function createEndpoint(endpointBaseUrl: string | undefined): IUserpodApiEndpoints { + return { + baseUrl: endpointBaseUrl, + ready: Promise.resolve(), + getAuthToken: () => 'endpoint-token' + }; + } + + /** + * 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 extraEnv: Record = { [sqlEnvKey]: sqlEnvValue }; + + await applyIntegrationEndpointEnv({ + deepnoteFileUri: projectFileUri, + endpoint: createEndpoint(baseUrl), + extraEnv + }); + + 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 extraEnv: Record = { [sqlEnvKey]: sqlEnvValue }; + + await applyIntegrationEndpointEnv({ + deepnoteFileUri: projectFileUri, + endpoint: createEndpoint(undefined), + extraEnv + }); + + 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 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 extraEnv: Record = { [sqlEnvKey]: sqlEnvValue }; + + await applyIntegrationEndpointEnv({ + deepnoteFileUri: projectFileUri, + endpoint: createEndpoint(baseUrl), + extraEnv + }); + + assert.deepStrictEqual(extraEnv, { [sqlEnvKey]: sqlEnvValue }); + }); +}); diff --git a/src/kernels/deepnote/deepnoteLspClientManager.node.ts b/src/kernels/deepnote/deepnoteLspClientManager.node.ts index 907adf8f8f..8d14766729 100644 --- a/src/kernels/deepnote/deepnoteLspClientManager.node.ts +++ b/src/kernels/deepnote/deepnoteLspClientManager.node.ts @@ -24,11 +24,10 @@ import { logger } from '../../platform/logging'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; 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'; interface LspClientInfo { @@ -61,10 +60,11 @@ export class DeepnoteLspClientManager constructor( @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, - @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) + private readonly sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider ) { this.disposables.push(this); } @@ -623,18 +623,16 @@ 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); + // Merged (SecretStorage + `.deepnote.env.yaml`) configs, so file-configured databases also get + // LSP autocomplete/schema. + const projectIntegrationConfigs = (await this.sqlIntegrationEnvVars.getMergedConfigs(notebookUri)).filter( + (config) => config.type !== 'pandas-dataframe' + ); const connections = projectIntegrationConfigs .filter((config) => isSupportedBySqlLsp(config.type)) .map((config) => convertToSqlLspConnection(config)) - .filter((conn): conn is SqlLspConnection => conn !== null); + .filter((conn) => conn !== null); logger.trace( `SQL LSP: Found ${connections.length} SQL LSP-compatible integrations for project ${projectId}` diff --git a/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts b/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts index fc2838bf66..8f9c8ca0f2 100644 --- a/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts +++ b/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts @@ -1,4 +1,5 @@ import { assert } from 'chai'; +import { anything, instance, mock, when } from 'ts-mockito'; import { Uri } from 'vscode'; import { DeepnoteLspClientManager } from './deepnoteLspClientManager.node'; @@ -6,7 +7,7 @@ import { createMockChildProcess } from './deepnoteTestHelpers.node'; import { IDisposableRegistry } from '../../platform/common/types'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import * as path from '../../platform/vscode-path/path'; -import { noop } from '../../platform/common/utils/misc'; +import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; suite('DeepnoteLspClientManager Integration Tests', () => { let lspClientManager: DeepnoteLspClientManager; @@ -29,16 +30,6 @@ suite('DeepnoteLspClientManager Integration Tests', () => { dispose: () => Promise.resolve() } as any; - // Mock integration storage - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const mockIntegrationStorage = { - getAll: async () => [], - getIntegrationConfig: async () => undefined, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - onDidChangeIntegrations: { dispose: noop } as any, - dispose: noop - } as any; - // Mock notebook editor provider // eslint-disable-next-line @typescript-eslint/no-explicit-any const mockNotebookEditorProvider = { @@ -52,11 +43,15 @@ suite('DeepnoteLspClientManager Integration Tests', () => { } as any; setup(() => { + // No integrations configured, so the LSP resolves an empty connection list. + const sqlIntegrationEnvVars = mock(); + when(sqlIntegrationEnvVars.getMergedConfigs(anything())).thenResolve([]); + lspClientManager = new DeepnoteLspClientManager( mockDisposableRegistry, - mockIntegrationStorage, mockNotebookEditorProvider, - mockNotebookManager + mockNotebookManager, + instance(sqlIntegrationEnvVars) ); lspClientManager.activate(); }); diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 112f7f0e09..26e782498c 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -6,7 +6,7 @@ */ import * as fs from 'fs-extra'; -import { inject, injectable, named, optional } from 'inversify'; +import { inject, injectable, named } from 'inversify'; import * as os from 'os'; import { CancellationToken, l10n, Uri } from 'vscode'; @@ -21,11 +21,12 @@ import { sleep } from '../../platform/common/utils/async'; import { generateUuid } from '../../platform/common/uuid'; 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'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; +import { applyIntegrationEndpointEnv } from './deepnoteIntegrationEndpointEnv'; +import { DeepnoteServerInfo, IDeepnoteServerStarter, IDeepnoteToolkitInstaller } from './types'; const MAX_OUTPUT_TRACKING_LENGTH = 5000; const SERVER_STARTUP_TIMEOUT_MS = 120_000; @@ -76,8 +77,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension @inject(IOutputChannel) @named(STANDARD_OUTPUT_CHANNEL) private readonly outputChannel: IOutputChannel, @inject(IAsyncDisposableRegistry) asyncRegistry: IAsyncDisposableRegistry, @inject(ISqlIntegrationEnvVarsProvider) - @optional() - private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider + private readonly sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider, + @inject(IUserpodApiEndpoints) + private readonly userpodApiEndpoints: IUserpodApiEndpoints ) { asyncRegistry.push(this); } @@ -262,6 +264,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: '' }); @@ -371,12 +374,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension token?: CancellationToken ): Promise> { const extraEnv: Record = {}; - - if (!this.sqlIntegrationEnvVars) { - logger.debug('DeepnoteServerStarter: SqlIntegrationEnvironmentVariablesProvider not available'); - return extraEnv; - } - const fileKey = deepnoteFileUri.fsPath; logger.debug( @@ -397,6 +394,10 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension return extraEnv; } + private async applyIntegrationEndpointEnv(extraEnv: Record, deepnoteFileUri: Uri): Promise { + await applyIntegrationEndpointEnv({ deepnoteFileUri, endpoint: this.userpodApiEndpoints, extraEnv }); + } + /** * 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 d0eb9070cc..dbcee7178c 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts @@ -4,18 +4,18 @@ import * as sinon from 'sinon'; import { anything, instance, mock, when } from 'ts-mockito'; import { CancellationError, Uri } from 'vscode'; -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 { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; +import { DeepnoteServerStarter } from './deepnoteServerStarter.node'; +import { IDeepnoteToolkitInstaller } from './types'; /** * Unit tests for DeepnoteServerStarter. @@ -42,6 +42,7 @@ suite('DeepnoteServerStarter', () => { let mockOutputChannel: IOutputChannel; let mockAsyncRegistry: IAsyncDisposableRegistry; let mockSqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider; + let mockUserpodApiEndpoints: IUserpodApiEndpoints; setup(() => { __resetRuntimeCoreMock(); @@ -52,10 +53,14 @@ suite('DeepnoteServerStarter', () => { mockOutputChannel = mock(); mockAsyncRegistry = mock(); mockSqlIntegrationEnvVars = mock(); + mockUserpodApiEndpoints = mock(); when(mockAsyncRegistry.push(anything())).thenReturn(); when(mockOutputChannel.appendLine(anything())).thenReturn(); + when(mockUserpodApiEndpoints.ready).thenReturn(Promise.resolve()); + when(mockUserpodApiEndpoints.baseUrl).thenReturn(undefined); + // The toolkit install step runs before runtime-core's startServer; stub it so the // start path reaches startServer. (ts-mockito methods that are not stubbed return null.) when(mockToolkitInstaller.ensureVenvAndToolkit(anything(), anything(), anything(), anything())).thenResolve({ @@ -71,7 +76,8 @@ suite('DeepnoteServerStarter', () => { instance(mockAgentSkillsManager), instance(mockOutputChannel), instance(mockAsyncRegistry), - instance(mockSqlIntegrationEnvVars) + instance(mockSqlIntegrationEnvVars), + instance(mockUserpodApiEndpoints) ); }); @@ -81,26 +87,15 @@ suite('DeepnoteServerStarter', () => { }); suite('SQL integration env vars', () => { - test('starts the server without SQL env vars when no provider is available', async () => { - // The provider is @optional() — construct a starter without it. - const starterWithoutSql = new DeepnoteServerStarter( - instance(mockProcessServiceFactory), - instance(mockToolkitInstaller), - instance(mockAgentSkillsManager), - instance(mockOutputChannel), - instance(mockAsyncRegistry) - ); + test('starts the server without SQL env vars when the provider yields none (e.g. web)', async () => { + when(mockSqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); - try { - await starterWithoutSql.startServer(interpreter, venvPath, true, [], 'env1', uriA); + await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); - assert.deepStrictEqual( - __getStartServerCalls().map((c) => c.env), - [{}] - ); - } finally { - await starterWithoutSql.dispose(); - } + assert.deepStrictEqual( + __getStartServerCalls().map((c) => c.env), + [{}] + ); }); test('starts the server without SQL env vars when the provider rejects with a cancellation error', async () => { diff --git a/src/kernels/jupyter/session/jupyterKernelService.unit.test.ts b/src/kernels/jupyter/session/jupyterKernelService.unit.test.ts index af86e50ffa..fe58ee4754 100644 --- a/src/kernels/jupyter/session/jupyterKernelService.unit.test.ts +++ b/src/kernels/jupyter/session/jupyterKernelService.unit.test.ts @@ -30,6 +30,7 @@ import { ICustomEnvironmentVariablesProvider } from '../../../platform/common/va import { EnvironmentVariablesService } from '../../../platform/common/variables/environment.node'; import { isWeb } from '../../../platform/common/utils/misc'; import { isPythonKernelConnection } from '../../helpers'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; // eslint-disable-next-line suite('JupyterKernelService', () => { @@ -442,12 +443,16 @@ suite('JupyterKernelService', () => { const configService = mock(ConfigurationService); settings = mock(JupyterSettings); when(configService.getSettings(anything())).thenReturn(instance(settings)); + // These tests do not exercise SQL integrations; the provider just has to yield nothing. + const sqlIntegrationEnvVars = mock(); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); const kernelEnvService = new KernelEnvironmentVariablesService( instance(interpreterService), instance(appEnv), variablesService, instance(customEnvVars), - instance(configService) + instance(configService), + instance(sqlIntegrationEnvVars) ); testWorkspaceFolder = Uri.file(path.join(EXTENSION_ROOT_DIR, 'src', 'test', 'datascience')); const jupyterPaths = mock(); diff --git a/src/kernels/raw/launcher/kernelEnvVarsService.node.ts b/src/kernels/raw/launcher/kernelEnvVarsService.node.ts index 11b7eb4293..c4f6ce626f 100644 --- a/src/kernels/raw/launcher/kernelEnvVarsService.node.ts +++ b/src/kernels/raw/launcher/kernelEnvVarsService.node.ts @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -import { inject, injectable, optional } from 'inversify'; +import { inject, injectable } from 'inversify'; import { logger } from '../../../platform/logging'; import { getDisplayPath } from '../../../platform/common/platform/fs-paths.node'; import { IConfigurationService, Resource, type ReadWrite } from '../../../platform/common/types'; @@ -33,13 +33,8 @@ export class KernelEnvironmentVariablesService { private readonly customEnvVars: ICustomEnvironmentVariablesProvider, @inject(IConfigurationService) private readonly configService: IConfigurationService, @inject(ISqlIntegrationEnvVarsProvider) - @optional() - private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider - ) { - logger.debug( - `KernelEnvironmentVariablesService: Constructor; SQL env provider present=${!!sqlIntegrationEnvVars}` - ); - } + private readonly sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider + ) {} /** * Generates the environment variables for the kernel. * @@ -62,7 +57,7 @@ export class KernelEnvironmentVariablesService { logger.debug( `KernelEnvVarsService.getEnvironmentVariables: Called for resource ${ resource ? getDisplayPath(resource) : 'undefined' - }, sqlIntegrationEnvVars is ${this.sqlIntegrationEnvVars ? 'AVAILABLE' : 'UNDEFINED'}` + }` ); let kernelEnv = kernelSpec.env && Object.keys(kernelSpec.env).length > 0 @@ -97,21 +92,17 @@ export class KernelEnvironmentVariablesService { }) : undefined, this.sqlIntegrationEnvVars - ? this.sqlIntegrationEnvVars - .getEnvironmentVariables(resource, token) - .then((vars) => { - if (vars && Object.keys(vars).length > 0) { - logger.debug( - `KernelEnvVarsService: Got ${Object.keys(vars).length} SQL integration env vars` - ); - } - return vars; - }) - .catch((ex) => { - logger.error('Failed to get SQL integration env variables for Kernel', ex); - return undefined; - }) - : undefined + .getEnvironmentVariables(resource, token) + .then((vars) => { + if (vars && Object.keys(vars).length > 0) { + logger.debug(`KernelEnvVarsService: Got ${Object.keys(vars).length} SQL integration env vars`); + } + return vars; + }) + .catch((ex) => { + logger.error('Failed to get SQL integration env variables for Kernel', ex); + return undefined; + }) ]); if (token?.isCancellationRequested) { return; diff --git a/src/kernels/raw/launcher/kernelEnvVarsService.unit.test.ts b/src/kernels/raw/launcher/kernelEnvVarsService.unit.test.ts index b4dd91bdc5..a4934213ec 100644 --- a/src/kernels/raw/launcher/kernelEnvVarsService.unit.test.ts +++ b/src/kernels/raw/launcher/kernelEnvVarsService.unit.test.ts @@ -73,26 +73,15 @@ suite('Kernel Environment Variables Service', () => { teardown(() => Object.assign(process.env, originalEnvVars)); - /** - * Helper factory function to build KernelEnvironmentVariablesService with optional overrides. - * @param overrides Optional overrides for the service dependencies - * @returns A new instance of KernelEnvironmentVariablesService - */ - function buildKernelEnvVarsService(overrides?: { - sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider | undefined; - }): KernelEnvironmentVariablesService { - const sqlProvider = - overrides && 'sqlIntegrationEnvVars' in overrides - ? overrides.sqlIntegrationEnvVars - : instance(sqlIntegrationEnvVars); - + /** Builds the service under test from the suite's mocks. */ + function buildKernelEnvVarsService(): KernelEnvironmentVariablesService { return new KernelEnvironmentVariablesService( instance(interpreterService), instance(envActivation), variablesService, instance(customVariablesService), instance(configService), - sqlProvider + instance(sqlIntegrationEnvVars) ); } @@ -388,7 +377,7 @@ suite('Kernel Environment Variables Service', () => { ); }); - test('SQL integration env vars work when provider is undefined (optional dependency)', async () => { + test('SQL integration env vars work when the provider yields none (e.g. web)', async () => { const resource = Uri.file('test.ipynb'); when(envActivation.getActivatedEnvironmentVariables(anything(), anything(), anything())).thenResolve({ PATH: 'foobar' @@ -397,10 +386,9 @@ suite('Kernel Environment Variables Service', () => { customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything()) ).thenResolve(); - // Create service without SQL integration provider - const serviceWithoutSql = buildKernelEnvVarsService({ sqlIntegrationEnvVars: undefined }); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); - const vars = await serviceWithoutSql.getEnvironmentVariables(resource, interpreter, kernelSpec); + const vars = await kernelVariablesService.getEnvironmentVariables(resource, interpreter, kernelSpec); assert.isOk(vars); assert.isUndefined(vars!['SQL_MY_DB']); diff --git a/src/kernels/serviceRegistry.web.ts b/src/kernels/serviceRegistry.web.ts index 7779e8c213..5f276b6fb7 100644 --- a/src/kernels/serviceRegistry.web.ts +++ b/src/kernels/serviceRegistry.web.ts @@ -35,6 +35,8 @@ import { KernelStartupCodeProviders } from './kernelStartupCodeProviders.web'; import { LastCellExecutionTracker } from './execution/lastCellExecutionTracker'; import { ClearJupyterServersCommand } from './jupyter/clearJupyterServersCommand'; import { KernelChatStartupCodeProvider } from './chat/kernelStartupCodeProvider'; +import { SqlIntegrationEnvironmentVariablesProviderWeb } from '../platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.web'; +import { ISqlIntegrationEnvVarsProvider } from '../platform/notebooks/deepnote/types'; @injectable() class RawNotebookSupportedService implements IRawNotebookSupportedService { @@ -83,6 +85,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea ); serviceManager.addSingleton(IKernelFinder, KernelFinder); serviceManager.addSingleton(IKernelDependencyService, KernelDependencyService); + serviceManager.addSingleton( + ISqlIntegrationEnvVarsProvider, + SqlIntegrationEnvironmentVariablesProviderWeb + ); serviceManager.addSingleton( IExtensionSyncActivationService, diff --git a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts index 0c3c048305..a55dfeb7c5 100644 --- a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts +++ b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts @@ -10,7 +10,7 @@ import { IFederatedAuthTokenStorage } from '../types'; /** * Node-only bridge that restarts kernels when a federated integration's token changes, clearing stale - * `os.environ` mutations and kernel globals. Separate from {@link IntegrationKernelRestartHandler} because + * `os.environ` mutations and kernel globals. Separate from {@link IntegrationEnvRefreshHandler} because * {@link IFederatedAuthTokenStorage} is node-only. */ @injectable() diff --git a/src/notebooks/deepnote/integrations/integrationDetector.ts b/src/notebooks/deepnote/integrations/integrationDetector.ts index 5715ead44c..3f3a2bcfd0 100644 --- a/src/notebooks/deepnote/integrations/integrationDetector.ts +++ b/src/notebooks/deepnote/integrations/integrationDetector.ts @@ -3,12 +3,10 @@ import { inject, injectable } from 'inversify'; import { logger } from '../../../platform/logging'; import { IDeepnoteNotebookManager } from '../../types'; import { - ConfigurableDatabaseIntegrationType, - IntegrationStatus, - IntegrationWithStatus + DetectedIntegration, + isConfigurableDatabaseIntegrationType } from '../../../platform/notebooks/deepnote/integrationTypes'; -import { IIntegrationDetector, IIntegrationStorage } from './types'; -import { databaseIntegrationTypes } from '@deepnote/database-integrations'; +import { IIntegrationDetector, IIntegrationStorage, IntegrationDetectionInput } from './types'; /** * Service for detecting integrations used in Deepnote notebooks @@ -21,68 +19,48 @@ export class IntegrationDetector implements IIntegrationDetector { ) {} /** - * Detect all integrations used in the given project. - * Uses the project's integrations field as the source of truth. + * Detect all integrations for the notebook's project. Two inputs, two roles: + * - `project.integrations` is the roster (ids, names and types only — never credentials), so it decides + * which integrations the panel lists at all. + * - SecretStorage supplies the editable config for each one; integrations configured only in + * `.deepnote.env.yaml` stay `null` here, since those configs are never persisted through it. */ - async detectIntegrations(projectId: string, notebookId: string): Promise> { - // Get the project + async detectIntegrations(input: IntegrationDetectionInput): Promise> { + const { projectId, notebookId } = input; + const project = this.notebookManager.getProjectForNotebook(projectId, notebookId); if (!project) { logger.warn( `IntegrationDetector: No project found for ID: ${projectId}. The project may not have been loaded yet.` ); + return new Map(); } - logger.debug(`IntegrationDetector: Scanning project ${projectId} for integrations`); + const declarations = project.project.integrations ?? []; - const integrations = new Map(); + logger.debug(`IntegrationDetector: Project ${projectId} declares ${declarations.length} integrations`); - // Use the project's integrations field as the source of truth - const projectIntegrations = project.project.integrations?.slice() ?? []; - logger.debug(`IntegrationDetector: Found ${projectIntegrations.length} integrations in project.integrations`); + const integrations = new Map(); - for (const projectIntegration of projectIntegrations) { - const integrationId = projectIntegration.id; - const integrationType = projectIntegration.type; - if ( - !(databaseIntegrationTypes as readonly string[]).includes(integrationType) || - integrationType === 'pandas-dataframe' - ) { + for (const declaration of declarations) { + const integrationType = declaration.type; + if (!isConfigurableDatabaseIntegrationType(integrationType)) { logger.debug(`IntegrationDetector: Skipping unsupported integration type: ${integrationType}`); continue; } - // Check if the integration is configured - const config = await this.integrationStorage.getIntegrationConfig(integrationId); - const status: IntegrationWithStatus = { - config: config ?? null, - status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected, - // Include integration metadata from project for prefilling when config is null - integrationName: projectIntegration.name, - integrationType: integrationType as ConfigurableDatabaseIntegrationType - }; + const storedConfig = await this.integrationStorage.getIntegrationConfig(declaration.id); - integrations.set(integrationId, status); + integrations.set(declaration.id, { + config: storedConfig ?? null, + integrationName: declaration.name, + integrationType + }); } logger.debug(`IntegrationDetector: Found ${integrations.size} integrations`); return integrations; } - - /** - * Check if a project has any unconfigured integrations - */ - async hasUnconfiguredIntegrations(projectId: string, notebookId: string): Promise { - const integrations = await this.detectIntegrations(projectId, notebookId); - - for (const integration of integrations.values()) { - if (integration.status === IntegrationStatus.Disconnected) { - return true; - } - } - - return false; - } } diff --git a/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.ts b/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.ts new file mode 100644 index 0000000000..1abe38e314 --- /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. + 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 0000000000..9e6cce71f2 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.unit.test.ts @@ -0,0 +1,151 @@ +import { assert } from 'chai'; +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 kernelExecution: INotebookKernelExecution; + let disposables: IDisposable[]; + + setup(() => { + resetVSCodeMocks(); + disposables = [new Disposable(() => resetVSCodeMocks())]; + kernelProvider = mock(); + + // Every notebook resolves to the same execution mock, so call counts below are totals across notebooks. + kernelExecution = mock(); + when(kernelExecution.executeHidden(anything())).thenResolve([]); + when(kernelProvider.getKernelExecution(anything())).thenReturn(instance(kernelExecution)); + + 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]); + + verify(kernelExecution.executeHidden(anything())).once(); + const [code] = capture(kernelExecution.executeHidden).first(); + assert.strictEqual( + code, + 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]); + + verify(kernelExecution.executeHidden(anything())).never(); + 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]); + + verify(kernelExecution.executeHidden(anything())).never(); + 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')); + when(kernelExecution.executeHidden(anything())).thenResolve([ + { output_type: 'error', ename: 'RuntimeError', evalue: 'boom', traceback: [] } + ]); + + await refresher.refresh([notebook]); + + verify(kernelExecution.executeHidden(anything())).once(); // the snippet still runs; 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]); + + verify(kernelExecution.executeHidden(anything())).twice(); // 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')); + when(kernelExecution.executeHidden(anything())).thenReject(new Error('kernel exploded')).thenResolve([]); + + await refresher.refresh([notebookA, notebookB]); + + verify(kernelExecution.executeHidden(anything())).twice(); // 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)); + when(kernelExecution.executeHidden(anything())) + .thenCall(() => secondInvoked.then(() => [])) + .thenCall(() => { + markSecondInvoked(); + + return Promise.resolve([]); + }); + + await refresher.refresh([notebookA, notebookB]); + + verify(kernelExecution.executeHidden(anything())).twice(); // 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 0000000000..c4806e87cf --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.ts @@ -0,0 +1,40 @@ +import { inject, injectable } from 'inversify'; +import { workspace } from 'vscode'; + +import { IExtensionSyncActivationService } from '../../../platform/activation/types'; +import { IDisposableRegistry } from '../../../platform/common/types'; +import { logger } from '../../../platform/logging'; +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; +import { IIntegrationEnvLiveRefresher, IIntegrationStorage } from './types'; + +/** Live-refreshes integration env in open Deepnote kernels when integration configs change (no restart). */ +@injectable() +export class IntegrationEnvRefreshHandler implements IExtensionSyncActivationService { + constructor( + @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, + @inject(IIntegrationEnvLiveRefresher) private readonly liveRefresher: IIntegrationEnvLiveRefresher, + @inject(IDisposableRegistry) disposables: IDisposableRegistry + ) { + logger.info('IntegrationEnvRefreshHandler: Initialized'); + + disposables.push( + this.integrationStorage.onDidChangeIntegrations(() => { + this.onIntegrationConfigurationChanged().catch((err) => + logger.error('IntegrationEnvRefreshHandler: Failed to handle integration change', err) + ); + }) + ); + } + + public activate(): void { + // Service is activated via constructor + } + + private async onIntegrationConfigurationChanged(): Promise { + const notebooks = workspace.notebookDocuments.filter( + (notebook) => notebook.notebookType === DEEPNOTE_NOTEBOOK_TYPE + ); + + await this.liveRefresher.refresh(notebooks); + } +} diff --git a/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.unit.test.ts b/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.unit.test.ts new file mode 100644 index 0000000000..d7dcb3ace7 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.unit.test.ts @@ -0,0 +1,125 @@ +import { assert } from 'chai'; +import * as sinon from 'sinon'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; +import { Disposable, EventEmitter, NotebookDocument, Uri } from 'vscode'; + +import { IDisposable } from '../../../platform/common/types'; +import { dispose } from '../../../platform/common/utils/lifecycle'; +import { logger } from '../../../platform/logging'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; +import { IIntegrationEnvLiveRefresher, IIntegrationStorage } from './types'; +import { IntegrationEnvRefreshHandler } from './integrationEnvRefreshHandler'; +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; + +suite('IntegrationEnvRefreshHandler', () => { + let handler: IntegrationEnvRefreshHandler; + let integrationStorage: IIntegrationStorage; + let liveRefresher: IIntegrationEnvLiveRefresher; + let disposables: IDisposable[]; + let onDidChangeIntegrations: EventEmitter; + + setup(() => { + resetVSCodeMocks(); + disposables = [new Disposable(() => resetVSCodeMocks())]; + integrationStorage = mock(); + liveRefresher = mock(); + onDidChangeIntegrations = new EventEmitter(); + disposables.push(onDidChangeIntegrations); + + when(integrationStorage.onDidChangeIntegrations).thenReturn(onDidChangeIntegrations.event); + when(liveRefresher.refresh(anything())).thenResolve(); + + handler = new IntegrationEnvRefreshHandler(instance(integrationStorage), instance(liveRefresher), disposables); + }); + + teardown(() => { + sinon.restore(); + disposables = dispose(disposables); + }); + + function createMockNotebook(notebookType: string, uri: Uri): NotebookDocument { + const notebook = mock(); + when(notebook.notebookType).thenReturn(notebookType); + when(notebook.uri).thenReturn(uri); + + return instance(notebook); + } + + test('refreshes integration env for open Deepnote notebooks when integrations change', async () => { + const notebook = createMockNotebook(DEEPNOTE_NOTEBOOK_TYPE, Uri.file('/test.deepnote')); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [notebook]); + }); + + test('passes only Deepnote notebooks to the refresher', async () => { + const deepnote = createMockNotebook(DEEPNOTE_NOTEBOOK_TYPE, Uri.file('/a.deepnote')); + const jupyter = createMockNotebook('jupyter-notebook', Uri.file('/b.ipynb')); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([deepnote, jupyter]); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [deepnote]); + }); + + test('refreshes multiple Deepnote notebooks in a single call', async () => { + const notebook1 = createMockNotebook(DEEPNOTE_NOTEBOOK_TYPE, Uri.file('/test1.deepnote')); + const notebook2 = createMockNotebook(DEEPNOTE_NOTEBOOK_TYPE, Uri.file('/test2.deepnote')); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [notebook1, notebook2]); + }); + + test('refreshes with an empty list when no Deepnote notebooks are open', async () => { + const jupyter = createMockNotebook('jupyter-notebook', Uri.file('/b.ipynb')); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([jupyter]); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], []); + }); + + test('catches and logs a rejected refresh instead of propagating it', async () => { + const notebook = createMockNotebook(DEEPNOTE_NOTEBOOK_TYPE, Uri.file('/test.deepnote')); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(liveRefresher.refresh(anything())).thenReject(new Error('refresh boom')); + const errorStub = sinon.stub(logger, 'error'); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + assert.strictEqual(errorStub.callCount, 1, 'the fire-and-forget rejection must be caught and logged'); + }); +}); diff --git a/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.ts b/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.ts deleted file mode 100644 index b082d0d192..0000000000 --- a/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.ts +++ /dev/null @@ -1,155 +0,0 @@ -import { inject, injectable } from 'inversify'; -import { l10n, NotebookDocument, workspace, window } from 'vscode'; - -import { IDisposableRegistry } from '../../../platform/common/types'; -import { IExtensionSyncActivationService } from '../../../platform/activation/types'; -import { logger } from '../../../platform/logging'; -import { IIntegrationStorage } from './types'; -import { IKernelProvider } from '../../../kernels/types'; -import { DATAFRAME_SQL_INTEGRATION_ID } from '../../../platform/notebooks/deepnote/integrationTypes'; - -/** - * Handles automatic kernel restart when integration configurations change. - * When a user saves/deletes an integration config, this service restarts all kernels - * that are using that integration so they pick up the new credentials. - */ -@injectable() -export class IntegrationKernelRestartHandler implements IExtensionSyncActivationService { - private isRestarting = false; - - constructor( - @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, - @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider, - @inject(IDisposableRegistry) disposables: IDisposableRegistry - ) { - logger.info('IntegrationKernelRestartHandler: Initialized'); - - // Listen for integration configuration changes - disposables.push( - this.integrationStorage.onDidChangeIntegrations(() => { - this.onIntegrationConfigurationChanged().catch((err) => - logger.error('IntegrationKernelRestartHandler: Failed to handle integration change', err) - ); - }) - ); - } - - public activate(): void { - // Service is activated via constructor - } - - /** - * Handle integration configuration changes by restarting affected kernels - */ - private async onIntegrationConfigurationChanged(): Promise { - // Prevent multiple simultaneous restart attempts - if (this.isRestarting) { - logger.debug('IntegrationKernelRestartHandler: Already restarting, skipping'); - return; - } - - try { - this.isRestarting = true; - - logger.info( - 'IntegrationKernelRestartHandler: Integration configuration changed, checking for affected kernels' - ); - - // Find all Deepnote notebooks with running kernels that use SQL integrations - const notebooksToRestart: NotebookDocument[] = []; - - for (const notebook of workspace.notebookDocuments) { - // Only process Deepnote notebooks - if (notebook.notebookType !== 'deepnote') { - continue; - } - - // Check if kernel is running - const kernel = this.kernelProvider.get(notebook); - if (!kernel || !kernel.startedAtLeastOnce) { - continue; - } - - // Check if notebook uses SQL integrations - const usesIntegrations = this.notebookUsesSqlIntegrations(notebook); - if (usesIntegrations) { - notebooksToRestart.push(notebook); - } - } - - if (notebooksToRestart.length === 0) { - logger.info( - 'IntegrationKernelRestartHandler: No running kernels use SQL integrations, no restart needed' - ); - return; - } - - logger.info( - `IntegrationKernelRestartHandler: Found ${notebooksToRestart.length} notebook(s) with kernels that need restart` - ); - - // Restart kernels for affected notebooks - const restartPromises = notebooksToRestart.map(async (notebook) => { - const kernel = this.kernelProvider.get(notebook); - if (kernel) { - try { - logger.info( - `IntegrationKernelRestartHandler: Restarting kernel for notebook: ${notebook.uri.toString()}` - ); - await kernel.restart(); - logger.info( - `IntegrationKernelRestartHandler: Successfully restarted kernel for: ${notebook.uri.toString()}` - ); - } catch (error) { - logger.error( - `IntegrationKernelRestartHandler: Failed to restart kernel for ${notebook.uri.toString()}`, - error - ); - // Don't throw - we want to continue restarting other kernels - } - } - }); - - await Promise.all(restartPromises); - - // Show a notification to the user - if (notebooksToRestart.length === 1) { - void window.showInformationMessage( - l10n.t('Integration configuration updated. Kernel restarted to apply changes.') - ); - } else { - void window.showInformationMessage( - l10n.t( - 'Integration configuration updated. {0} kernels restarted to apply changes.', - notebooksToRestart.length - ) - ); - } - } finally { - this.isRestarting = false; - } - } - - /** - * Check if a notebook uses SQL integrations by scanning cells for sql_integration_id metadata - */ - private notebookUsesSqlIntegrations(notebook: NotebookDocument): boolean { - for (const cell of notebook.getCells()) { - // Check for SQL cells - if (cell.document.languageId !== 'sql') { - continue; - } - - const metadata = cell.metadata; - if (metadata && typeof metadata === 'object') { - const integrationId = (metadata as Record).sql_integration_id; - if (typeof integrationId === 'string' && integrationId !== DATAFRAME_SQL_INTEGRATION_ID) { - // Found a SQL cell with an external integration - return true; - } - } - } - - return false; - } -} diff --git a/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.unit.test.ts b/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.unit.test.ts deleted file mode 100644 index 62888c3fea..0000000000 --- a/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.unit.test.ts +++ /dev/null @@ -1,243 +0,0 @@ -import { anything, instance, mock, verify, when } from 'ts-mockito'; -import { Disposable, EventEmitter, NotebookCell, NotebookDocument, TextDocument, Uri } from 'vscode'; - -import { IDisposable } from '../../../platform/common/types'; -import { dispose } from '../../../platform/common/utils/lifecycle'; -import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; -import { IIntegrationStorage } from './types'; -import { IKernel, IKernelProvider } from '../../../kernels/types'; -import { IntegrationKernelRestartHandler } from './integrationKernelRestartHandler'; -import { DATAFRAME_SQL_INTEGRATION_ID } from '../../../platform/notebooks/deepnote/integrationTypes'; - -suite('IntegrationKernelRestartHandler', () => { - let handler: IntegrationKernelRestartHandler; - let integrationStorage: IIntegrationStorage; - let kernelProvider: IKernelProvider; - let disposables: IDisposable[]; - let onDidChangeIntegrations: EventEmitter; - - setup(() => { - resetVSCodeMocks(); - disposables = [new Disposable(() => resetVSCodeMocks())]; - integrationStorage = mock(); - kernelProvider = mock(); - onDidChangeIntegrations = new EventEmitter(); - disposables.push(onDidChangeIntegrations); - - when(integrationStorage.onDidChangeIntegrations).thenReturn(onDidChangeIntegrations.event); - - handler = new IntegrationKernelRestartHandler( - instance(integrationStorage), - instance(kernelProvider), - disposables - ); - }); - - teardown(() => { - disposables = dispose(disposables); - }); - - function createMockNotebook( - notebookType: string, - uri: Uri, - cells: { languageId: string; metadata?: Record }[] - ): NotebookDocument { - const notebook = mock(); - const mockCells: NotebookCell[] = cells.map((cellConfig, index) => { - const cell = mock(); - const doc = mock(); - when(doc.languageId).thenReturn(cellConfig.languageId); - when(cell.document).thenReturn(instance(doc)); - when(cell.metadata).thenReturn(cellConfig.metadata || {}); - when(cell.index).thenReturn(index); - return instance(cell); - }); - - when(notebook.notebookType).thenReturn(notebookType); - when(notebook.uri).thenReturn(uri); - when(notebook.getCells()).thenReturn(mockCells); - - return instance(notebook); - } - - test('restarts kernel when integration changes and notebook uses SQL integrations', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).once(); - }); - - test('does not restart kernel for non-Deepnote notebooks', async () => { - const notebook = createMockNotebook('jupyter-notebook', Uri.file('/test.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).never(); - }); - - test('does not restart kernel that has not started', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(false); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).never(); - }); - - test('does not restart kernel when notebook has no SQL integrations', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'python', metadata: {} } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).never(); - }); - - test('does not restart kernel when notebook only uses internal DuckDB integration', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: DATAFRAME_SQL_INTEGRATION_ID } } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).never(); - }); - - test('restarts multiple kernels in parallel', async () => { - const notebook1 = createMockNotebook('deepnote', Uri.file('/test1.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const notebook2 = createMockNotebook('deepnote', Uri.file('/test2.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'bigquery-1' } } - ]); - const mockKernel1 = mock(); - const mockKernel2 = mock(); - when(mockKernel1.startedAtLeastOnce).thenReturn(true); - when(mockKernel1.restart()).thenResolve(); - when(mockKernel2.startedAtLeastOnce).thenReturn(true); - when(mockKernel2.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - when(kernelProvider.get(notebook1)).thenReturn(instance(mockKernel1)); - when(kernelProvider.get(notebook2)).thenReturn(instance(mockKernel2)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel1.restart()).once(); - verify(mockKernel2.restart()).once(); - }); - - test('continues restarting other kernels when one fails', async () => { - const notebook1 = createMockNotebook('deepnote', Uri.file('/test1.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const notebook2 = createMockNotebook('deepnote', Uri.file('/test2.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'bigquery-1' } } - ]); - const mockKernel1 = mock(); - const mockKernel2 = mock(); - when(mockKernel1.startedAtLeastOnce).thenReturn(true); - when(mockKernel1.restart()).thenReject(new Error('Restart failed')); - when(mockKernel2.startedAtLeastOnce).thenReturn(true); - when(mockKernel2.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - when(kernelProvider.get(notebook1)).thenReturn(instance(mockKernel1)); - when(kernelProvider.get(notebook2)).thenReturn(instance(mockKernel2)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel1.restart()).once(); - verify(mockKernel2.restart()).once(); - }); - - test('handles notebooks with mixed cell types', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'python', metadata: {} }, - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } }, - { languageId: 'markdown', metadata: {} } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).once(); - }); - - test('does not restart when no notebooks are open', async () => { - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(kernelProvider.get(anything())).never(); - }); -}); diff --git a/src/notebooks/deepnote/integrations/integrationManager.ts b/src/notebooks/deepnote/integrations/integrationManager.ts index a98f3763bb..cf8f163e4a 100644 --- a/src/notebooks/deepnote/integrations/integrationManager.ts +++ b/src/notebooks/deepnote/integrations/integrationManager.ts @@ -1,11 +1,10 @@ import { inject, injectable } from 'inversify'; -import { commands, l10n, window, workspace } from 'vscode'; +import { commands, l10n, window } from 'vscode'; import { IExtensionContext } from '../../../platform/common/types'; import { Commands } from '../../../platform/common/constants'; import { logger } from '../../../platform/logging'; import { IIntegrationDetector, IIntegrationManager, IIntegrationStorage, IIntegrationWebviewProvider } from './types'; -import { IntegrationStatus } from '../../../platform/notebooks/deepnote/integrationTypes'; import { IDeepnoteNotebookManager } from '../../types'; import { DatabaseIntegrationType, databaseIntegrationTypes } from '@deepnote/database-integrations'; @@ -14,10 +13,6 @@ import { DatabaseIntegrationType, databaseIntegrationTypes } from '@deepnote/dat */ @injectable() export class IntegrationManager implements IIntegrationManager { - private hasIntegrationsContext = 'deepnote.hasIntegrations'; - - private hasUnconfiguredIntegrationsContext = 'deepnote.hasUnconfiguredIntegrations'; - constructor( @inject(IExtensionContext) private readonly extensionContext: IExtensionContext, @inject(IIntegrationDetector) private readonly integrationDetector: IIntegrationDetector, @@ -50,68 +45,6 @@ export class IntegrationManager implements IIntegrationManager { return this.showIntegrationsUI(integrationId); }) ); - - // Listen for active notebook changes to update context - this.extensionContext.subscriptions.push( - window.onDidChangeActiveNotebookEditor(() => - this.updateContext().catch((err) => - logger.error('IntegrationManager: Failed to update context on notebook editor change', err) - ) - ) - ); - - // Listen for notebook document changes - this.extensionContext.subscriptions.push( - workspace.onDidOpenNotebookDocument(() => - this.updateContext().catch((err) => - logger.error('IntegrationManager: Failed to update context on notebook open', err) - ) - ) - ); - - this.extensionContext.subscriptions.push( - workspace.onDidCloseNotebookDocument(() => - this.updateContext().catch((err) => - logger.error('IntegrationManager: Failed to update context on notebook close', err) - ) - ) - ); - - // Initial context update - this.updateContext().catch((err) => - logger.error('IntegrationManager: Failed to update context on activation', err) - ); - } - - /** - * Update the context keys based on the active notebook - */ - private async updateContext(): Promise { - const activeNotebook = window.activeNotebookEditor?.notebook; - - if (!activeNotebook || activeNotebook.notebookType !== 'deepnote') { - await commands.executeCommand('setContext', this.hasIntegrationsContext, false); - await commands.executeCommand('setContext', this.hasUnconfiguredIntegrationsContext, false); - return; - } - - const projectId = activeNotebook.metadata?.deepnoteProjectId; - const notebookId = activeNotebook.metadata?.deepnoteNotebookId; - if (!projectId || !notebookId) { - await commands.executeCommand('setContext', this.hasIntegrationsContext, false); - await commands.executeCommand('setContext', this.hasUnconfiguredIntegrationsContext, false); - return; - } - - // Detect integrations in the project - const integrations = await this.integrationDetector.detectIntegrations(projectId, notebookId); - const hasIntegrations = integrations.size > 0; - const hasUnconfigured = Array.from(integrations.values()).some( - (integration) => integration.status === IntegrationStatus.Disconnected - ); - - await commands.executeCommand('setContext', this.hasIntegrationsContext, hasIntegrations); - await commands.executeCommand('setContext', this.hasUnconfiguredIntegrationsContext, hasUnconfigured); } /** @@ -137,7 +70,10 @@ export class IntegrationManager implements IIntegrationManager { logger.trace(`IntegrationManager: Notebook metadata:`, activeNotebook.metadata); // First try to detect integrations from the stored project - let integrations = await this.integrationDetector.detectIntegrations(projectId, notebookId); + let integrations = await this.integrationDetector.detectIntegrations({ + projectId, + notebookId + }); logger.debug(`IntegrationManager: Found ${integrations.size} integrations`); // If a specific integration was requested (e.g., from status bar click), @@ -167,7 +103,6 @@ export class IntegrationManager implements IIntegrationManager { } else { integrations.set(selectedIntegrationId, { config: config || null, - status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected, integrationName, integrationType }); diff --git a/src/notebooks/deepnote/integrations/integrationWebview.ts b/src/notebooks/deepnote/integrations/integrationWebview.ts index 0ba25f8c5f..1a091e667f 100644 --- a/src/notebooks/deepnote/integrations/integrationWebview.ts +++ b/src/notebooks/deepnote/integrations/integrationWebview.ts @@ -14,8 +14,7 @@ import { IFederatedAuthTokenStorage, IIntegrationStorage, IIntegrationWebviewPro import { ConfigurableDatabaseIntegrationConfig, FederatedAuthTokenStatus, - IntegrationStatus, - IntegrationWithStatus + DetectedIntegration } from '../../../platform/notebooks/deepnote/integrationTypes'; /** @@ -29,7 +28,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { private readonly disposables: Disposable[] = []; - private integrations: Map = new Map(); + private integrations: Map = new Map(); private projectId: string | undefined; @@ -69,7 +68,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { */ public async show( projectId: string, - integrations: Map, + integrations: Map, activeFileUri: Uri, selectedIntegrationId?: string, projectName?: string @@ -467,7 +466,6 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { id, integrationName: integration.integrationName, integrationType: integration.integrationType, - status: integration.status, tokenStatus: await this.deriveTokenStatus(id, integration.config) })) ); @@ -641,7 +639,6 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { if (integration) { // Existing integration - update it integration.config = config; - integration.status = IntegrationStatus.Connected; integration.integrationName = config.name; integration.integrationType = config.type; this.integrations.set(integrationId, integration); @@ -649,7 +646,6 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { // New integration - add it to the map this.integrations.set(integrationId, { config, - status: IntegrationStatus.Connected, integrationName: config.name, integrationType: config.type }); @@ -691,7 +687,6 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { const integration = this.integrations.get(integrationId); if (integration) { integration.config = null; - integration.status = IntegrationStatus.Disconnected; this.integrations.set(integrationId, integration); } diff --git a/src/notebooks/deepnote/integrations/integrationWebview.unit.test.ts b/src/notebooks/deepnote/integrations/integrationWebview.unit.test.ts index 51d2b67bbb..27e1e78b6e 100644 --- a/src/notebooks/deepnote/integrations/integrationWebview.unit.test.ts +++ b/src/notebooks/deepnote/integrations/integrationWebview.unit.test.ts @@ -11,8 +11,7 @@ import { FederatedAuthTokenEntry, IFederatedAuthTokenStorage, IIntegrationStorag import { computeMetadataFingerprint } from './federatedAuth/federatedAuthTokenStorage.node'; import { ConfigurableDatabaseIntegrationConfig, - IntegrationStatus, - IntegrationWithStatus + DetectedIntegration } from '../../../platform/notebooks/deepnote/integrationTypes'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; import { @@ -164,11 +163,11 @@ suite('IntegrationWebviewProvider', () => { function singleIntegrationMap( id: string, config: ConfigurableDatabaseIntegrationConfig - ): Map { - return new Map([[id, { config, status: IntegrationStatus.Connected }]]); + ): Map { + return new Map([[id, { config }]]); } - async function show(provider: IntegrationWebviewProvider, integrations: Map) { + async function show(provider: IntegrationWebviewProvider, integrations: Map) { await provider.show(PROJECT_ID, integrations, Uri.file('/ws/active.deepnote')); } diff --git a/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.ts b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.ts new file mode 100644 index 0000000000..1fef96a0d2 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.ts @@ -0,0 +1,166 @@ +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. */ +export 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; + } + } + } + ); + } + + 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. + const candidateDirs = + workspaceRoot != null && workspaceRoot.fsPath !== deepnoteDir.fsPath + ? [deepnoteDir, workspaceRoot] + : [deepnoteDir]; + if (await this.hasIntegrationsFile(candidateDirs)) { + affected.push(notebook); + } + } + + return affected; + } + + private async handleChangedDirs(changedDirs: Set): Promise { + const affected = await this.findAffectedNotebooks(changedDirs); + if (affected.length === 0) { + return; + } + + await this.liveRefresher.refresh(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 0000000000..7540fdd2dc --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.unit.test.ts @@ -0,0 +1,254 @@ +import { assert } from 'chai'; +import { + Disposable, + EventEmitter, + FileSystemWatcher, + NotebookDocument, + RelativePattern, + Uri, + WorkspaceFolder, + WorkspaceFoldersChangeEvent +} from 'vscode'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; +import * as sinon from 'sinon'; +import { DEFAULT_ENV_FILE, DEFAULT_INTEGRATIONS_FILE } from '@deepnote/database-integrations'; + +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 { debounceTimeInMilliseconds, 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[]; + let clock: sinon.SinonFakeTimers; + /** One emitter per file system watcher the code under test asked VS Code to create, keyed by dir + file name. */ + let fileEvents: Map>; + let onDidOpenNotebookDocument: EventEmitter; + + const workspaceFolder: WorkspaceFolder = { index: 0, name: 'ws', uri: Uri.file('/ws') }; + + /** Hands back a fake watcher whose events the test can fire, standing in for real filesystem events. */ + function createFileSystemWatcher(pattern: RelativePattern): FileSystemWatcher { + const emitter = new EventEmitter(); + fileEvents.set(watcherKey(pattern.baseUri, pattern.pattern), emitter); + disposables.push(emitter); + + const fsWatcher = mock(); + when(fsWatcher.onDidChange).thenReturn(emitter.event); + when(fsWatcher.onDidCreate).thenReturn(emitter.event); + when(fsWatcher.onDidDelete).thenReturn(emitter.event); + when(fsWatcher.dispose()).thenReturn(); + + return instance(fsWatcher); + } + + 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 the watcher derives from a notebook uri (the `.deepnote` file's dir). */ + function deepnoteDirOf(uri: Uri): Uri { + return Uri.joinPath(notebookPathToDeepnoteProjectFilePath(uri), '..'); + } + + /** Fires a watcher event and lets the debounce elapse, as a real edit of the file would. */ + async function fireEnvFileChange(dir: Uri, fileName: string = DEFAULT_INTEGRATIONS_FILE): Promise { + fireEnvFileEvent(dir, fileName); + + await clock.tickAsync(debounceTimeInMilliseconds); + } + + /** Fires a watcher event without advancing the clock, so callers can queue a burst within one debounce window. */ + function fireEnvFileEvent(dir: Uri, fileName: string): void { + const emitter = fileEvents.get(watcherKey(dir, fileName)); + if (!emitter) { + assert.fail(`No file system watcher was created for ${fileName} in ${dir.fsPath}`); + } + + emitter.fire(Uri.joinPath(dir, fileName)); + } + + function watcherKey(dir: Uri, fileName: string): string { + return `${dir.fsPath}::${fileName}`; + } + + setup(() => { + resetVSCodeMocks(); + disposables = [new Disposable(() => resetVSCodeMocks())]; + clock = sinon.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }); + + fileEvents = new Map>(); + when( + mockedVSCodeNamespaces.workspace.createFileSystemWatcher(anything(), anything(), anything(), anything()) + ).thenCall(createFileSystemWatcher); + + onDidOpenNotebookDocument = new EventEmitter(); + disposables.push(onDidOpenNotebookDocument); + when(mockedVSCodeNamespaces.workspace.onDidOpenNotebookDocument).thenReturn(onDidOpenNotebookDocument.event); + when(mockedVSCodeNamespaces.workspace.onDidChangeWorkspaceFolders).thenReturn( + new EventEmitter().event + ); + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder]); + + 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(() => { + clock.restore(); + 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]); + watcher.activate(); + + await fireEnvFileChange(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(workspaceFolder); + watcher.activate(); + + // The changed dir is the workspace root, NOT the .deepnote dir (/ws/nested/deep). + await fireEnvFileChange(workspaceFolder.uri); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [notebook]); + }); + + test('coalesces a burst of env file events into a single refresh', async () => { + const uri = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + watcher.activate(); + + // Both watched files saved at once — the debounce must collapse them into one refresh. + fireEnvFileEvent(deepnoteDirOf(uri), DEFAULT_INTEGRATIONS_FILE); + fireEnvFileEvent(deepnoteDirOf(uri), DEFAULT_ENV_FILE); + await clock.tickAsync(debounceTimeInMilliseconds); + + verify(liveRefresher.refresh(anything())).once(); + }); + + test('watches the dir of a notebook opened after activation', async () => { + // Outside the workspace folder, so the dir is watched only because the notebook was opened. + const uri = Uri.file('/elsewhere/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri); + + watcher.activate(); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + onDidOpenNotebookDocument.fire(notebook); + + await fireEnvFileChange(deepnoteDirOf(uri)); + + verify(liveRefresher.refresh(anything())).once(); + }); + + test('does not refresh when the changed dir matches no open Deepnote notebook', async () => { + const otherFolder: WorkspaceFolder = { index: 1, name: 'other', uri: Uri.file('/other') }; + const uri = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri); + + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder, otherFolder]); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + watcher.activate(); + + // An unrelated workspace folder changed — the notebook's dir and its workspace root are untouched. + await fireEnvFileChange(otherFolder.uri); + + verify(liveRefresher.refresh(anything())).never(); + }); + + test('ignores non-Deepnote notebooks even when their dir changed', async () => { + const uri = Uri.file('/ws/app.ipynb').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook(uri, 'jupyter-notebook'); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + watcher.activate(); + + // The dir is watched as a workspace folder, so the event lands — the notebook type must rule it out. + await fireEnvFileChange(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); + watcher.activate(); + + await fireEnvFileChange(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); + watcher.activate(); + + await fireEnvFileChange(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); + watcher.activate(); + + await fireEnvFileChange(deepnoteDirOf(uri), DEFAULT_ENV_FILE); + + verify(liveRefresher.refresh(anything())).never(); + }); +}); diff --git a/src/notebooks/deepnote/integrations/types.ts b/src/notebooks/deepnote/integrations/types.ts index 3e151cb788..437d95882d 100644 --- a/src/notebooks/deepnote/integrations/types.ts +++ b/src/notebooks/deepnote/integrations/types.ts @@ -1,22 +1,24 @@ import type { DeepnoteBlock } from '@deepnote/blocks'; -import { Event, Uri } from 'vscode'; +import { Event, NotebookDocument, Uri } from 'vscode'; -import { IntegrationWithStatus } from '../../../platform/notebooks/deepnote/integrationTypes'; +import { DetectedIntegration } from '../../../platform/notebooks/deepnote/integrationTypes'; // Re-export IIntegrationStorage from platform layer export { IIntegrationStorage } from '../../../platform/notebooks/deepnote/types'; export const IIntegrationDetector = Symbol('IIntegrationDetector'); + +/** Inputs for {@link IIntegrationDetector.detectIntegrations}: the open notebook's already-validated Deepnote IDs. */ +export interface IntegrationDetectionInput { + projectId: string; + notebookId: string; +} + export interface IIntegrationDetector { /** * Detect all integrations used in the given project */ - detectIntegrations(projectId: string, notebookId: string): Promise>; - - /** - * Check if a project has any unconfigured integrations - */ - hasUnconfiguredIntegrations(projectId: string, notebookId: string): Promise; + detectIntegrations(input: IntegrationDetectionInput): Promise>; } export const IIntegrationWebviewProvider = Symbol('IIntegrationWebviewProvider'); @@ -24,14 +26,14 @@ export interface IIntegrationWebviewProvider { /** * Show the integration management webview * @param projectId The Deepnote project ID - * @param integrations Map of integration IDs to their status + * @param integrations Map of integration IDs to their detected config and metadata * @param activeFileUri The `.deepnote` file being edited — always persisted to disk on save * @param selectedIntegrationId Optional integration ID to select/configure immediately * @param projectName Optional project display name (sourced from the active notebook's metadata) */ show( projectId: string, - integrations: Map, + integrations: Map, activeFileUri: Uri, selectedIntegrationId?: string, projectName?: string @@ -46,6 +48,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 0000000000..0c8e4bb740 --- /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. + 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 0000000000..79a7ee3f49 --- /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 5bf312977c..4bb7ee213e 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts @@ -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,9 @@ 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) + private readonly sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider ) {} public activate(): void { @@ -229,12 +232,23 @@ 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. + const fileConfig = (await this.sqlIntegrationEnvVars.getMergedConfigs(cell.notebook.uri)).find( + (c) => c.id === integrationId + ); + 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/deepnote/sqlCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts index f729453582..be773611e6 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts @@ -11,6 +11,18 @@ import { createEventHandler } from '../../test/common'; import { Commands } from '../../platform/common/constants'; import { IDeepnoteNotebookManager } from '../types'; import { createMockCell } from './deepnoteTestHelpers'; +import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; + +/** + * A merged-config source that yields nothing, so integrations resolve from SecretStorage alone. + * These suites never assert against it, hence an instance rather than a mock handle. + */ +function emptySqlIntegrationEnvVars(): ISqlIntegrationEnvVarsProvider { + const provider = mock(); + when(provider.getMergedConfigs(anything())).thenResolve([]); + + return instance(provider); +} suite('SqlCellStatusBarProvider', () => { let provider: SqlCellStatusBarProvider; @@ -23,7 +35,12 @@ suite('SqlCellStatusBarProvider', () => { disposables = []; integrationStorage = mock(); notebookManager = mock(); - provider = new SqlCellStatusBarProvider(disposables, instance(integrationStorage), instance(notebookManager)); + provider = new SqlCellStatusBarProvider( + disposables, + instance(integrationStorage), + instance(notebookManager), + emptySqlIntegrationEnvVars() + ); const tokenSource = new CancellationTokenSource(); cancellationToken = tokenSource.token; @@ -304,7 +321,8 @@ suite('SqlCellStatusBarProvider', () => { activateProvider = new SqlCellStatusBarProvider( activateDisposables, instance(activateIntegrationStorage), - instance(activateNotebookManager) + instance(activateNotebookManager), + emptySqlIntegrationEnvVars() ); }); @@ -540,7 +558,8 @@ suite('SqlCellStatusBarProvider', () => { eventProvider = new SqlCellStatusBarProvider( eventDisposables, instance(eventIntegrationStorage), - instance(eventNotebookManager) + instance(eventNotebookManager), + emptySqlIntegrationEnvVars() ); }); @@ -668,7 +687,8 @@ suite('SqlCellStatusBarProvider', () => { commandProvider = new SqlCellStatusBarProvider( commandDisposables, instance(commandIntegrationStorage), - instance(commandNotebookManager) + instance(commandNotebookManager), + emptySqlIntegrationEnvVars() ); // Capture the command handler @@ -809,7 +829,8 @@ suite('SqlCellStatusBarProvider', () => { commandProvider = new SqlCellStatusBarProvider( commandDisposables, instance(commandIntegrationStorage), - instance(commandNotebookManager) + instance(commandNotebookManager), + emptySqlIntegrationEnvVars() ); // Capture the command handler diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index eb87ce58a3..bf07345c76 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 28937d6b60..6d6b826180 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/src/notebooks/serviceRegistry.web.ts @@ -53,7 +53,6 @@ import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnote import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; -import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler'; import { FederatedAuthCommandHandlerWeb } from './deepnote/integrations/federatedAuth/federatedAuthCommandHandler.web'; import { DeepnoteFileChangeWatcher } from './deepnote/deepnoteFileChangeWatcher'; import { DeepnoteNotebookInfoStatusBar } from './deepnote/deepnoteNotebookInfoStatusBar'; @@ -139,10 +138,6 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, SqlCellStatusBarProvider ); - serviceManager.addSingleton( - IExtensionSyncActivationService, - IntegrationKernelRestartHandler - ); serviceManager.addSingleton( IExtensionSyncActivationService, FederatedAuthCommandHandlerWeb diff --git a/src/platform/notebooks/deepnote/integrationTypes.ts b/src/platform/notebooks/deepnote/integrationTypes.ts index 8eb24affc5..06ce5121d1 100644 --- a/src/platform/notebooks/deepnote/integrationTypes.ts +++ b/src/platform/notebooks/deepnote/integrationTypes.ts @@ -74,7 +74,13 @@ export interface LegacyDuckDBIntegrationConfig extends BaseLegacyIntegrationConf type: LegacyIntegrationType.DuckDB; } -import { DatabaseIntegrationConfig, DatabaseIntegrationType } from '@deepnote/database-integrations'; +import { + DatabaseIntegrationConfig, + DatabaseIntegrationType, + databaseIntegrationTypes, + FederatedAuthMethod, + isFederatedAuthMethod +} from '@deepnote/database-integrations'; // Import and re-export Snowflake auth constants from shared module import { type SnowflakeAuthMethod, @@ -143,25 +149,25 @@ export type ConfigurableDatabaseIntegrationConfig = Extract< export type ConfigurableDatabaseIntegrationType = Exclude; -/** - * Integration connection status - */ -export enum IntegrationStatus { - Connected = 'connected', - Disconnected = 'disconnected', - Error = 'error' +/** Narrows a raw type string to one the webview can configure; excludes the internal DuckDB integration. */ +export function isConfigurableDatabaseIntegrationType( + type: string | undefined +): type is ConfigurableDatabaseIntegrationType { + return ( + type !== undefined && + type !== 'pandas-dataframe' && + (databaseIntegrationTypes as readonly string[]).includes(type) + ); } /** Federated-auth token status: `'authenticated'`, `'disconnected'` (federated but no token), or `'unsupported'` (non-federated or web/remote). */ export type FederatedAuthTokenStatus = 'authenticated' | 'disconnected' | 'unsupported'; /** - * Integration with its current status + * An integration declared by a project, paired with the credentials stored for it (if any) */ -export interface IntegrationWithStatus { +export interface DetectedIntegration { config: ConfigurableDatabaseIntegrationConfig | null; - status: IntegrationStatus; - error?: string; /** * Name from the project's integrations list (used for prefilling when config is null) */ @@ -173,3 +179,22 @@ export interface IntegrationWithStatus { /** Federated-auth token status; only meaningful for federated integrations (currently BigQuery + `google-oauth`). */ tokenStatus?: FederatedAuthTokenStatus; } + +/** + * Narrows integration metadata to the federated-auth variant. Shared by the file-config provider and the SQL + * env-vars provider (upstream `isFederatedAuthMetadata`'s generic doesn't unify with our + * `DatabaseIntegrationConfig['metadata']` union); delegates to the exported `isFederatedAuthMethod` at runtime. + */ +export function isFederatedAuthMetadata( + metadata: DatabaseIntegrationConfig['metadata'] +): metadata is Extract { + if (typeof metadata !== 'object' || metadata === null) { + return false; + } + if (!('authMethod' in metadata)) { + return false; + } + const authMethod = metadata.authMethod; + + return typeof authMethod === 'string' && isFederatedAuthMethod(authMethod); +} diff --git a/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.ts b/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.ts new file mode 100644 index 0000000000..a8fd53b3b4 --- /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; a clean parse clears them. No-op when diagnostics are unavailable (e.g. web/tests). */ + private updateDiagnostics(yamlUri: Uri, issues: ValidationIssue[]): void { + if (!this.diagnostics) { + return; + } + + if (issues.length === 0) { + this.diagnostics.delete(yamlUri); + + return; + } + + const diagnostics = issues.map((issue) => { + const detail = issue.path + ? `${issue.code} at '${issue.path}': ${issue.message}` + : `${issue.code}: ${issue.message}`; + const diagnostic = new Diagnostic(new Range(0, 0, 0, 0), detail, DiagnosticSeverity.Warning); + diagnostic.source = 'Deepnote integrations'; + + return diagnostic; + }); + + this.diagnostics.set(yamlUri, diagnostics); + } +} diff --git a/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.unit.test.ts b/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.unit.test.ts new file mode 100644 index 0000000000..1d9a2af74d --- /dev/null +++ b/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.unit.test.ts @@ -0,0 +1,375 @@ +import assert from 'assert'; + +import { + DatabaseIntegrationConfig, + DEFAULT_ENV_FILE, + DEFAULT_INTEGRATIONS_FILE +} from '@deepnote/database-integrations'; +import dedent from 'dedent'; +import { anything, instance, mock, verify, when } from 'ts-mockito'; +import { Uri, WorkspaceConfiguration, WorkspaceFolder } from 'vscode'; + +import { IFileSystem } from '../../common/platform/types'; +import { IntegrationsFileConfigProvider } from './integrationsFileConfigProvider.node'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; + +/** A file present in the virtual filesystem: either readable `content`, or a `readError` that rejects. */ +interface VirtualFile { + path: string; + content?: string; + readError?: Error; +} + +/** Provider whose process environment is controllable, so tests never read or mutate the real `process.env`. */ +class TestableIntegrationsFileConfigProvider extends IntegrationsFileConfigProvider { + public processEnvironment: Record = {}; + + protected override getProcessEnvironment(): Record { + return this.processEnvironment; + } +} + +suite('IntegrationsFileConfigProvider', () => { + const deepnoteFileUri = Uri.file('/workspace/project/notebook.deepnote'); + const deepnoteDirUri = Uri.joinPath(deepnoteFileUri, '..'); + // Build the expected file paths exactly as the loader does (dir of the `.deepnote` file). + const yamlPath = Uri.joinPath(deepnoteDirUri, DEFAULT_INTEGRATIONS_FILE).fsPath; + const envPath = Uri.joinPath(deepnoteDirUri, DEFAULT_ENV_FILE).fsPath; + + let fileSystem: IFileSystem; + let provider: TestableIntegrationsFileConfigProvider; + let featureEnabled: boolean; + let workspaceFolder: WorkspaceFolder | undefined; + + setup(() => { + resetVSCodeMocks(); + + featureEnabled = true; + workspaceFolder = undefined; + + fileSystem = mock(); + provider = new TestableIntegrationsFileConfigProvider(instance(fileSystem), []); + + // The gate reads `deepnote.integrations.envFile.enabled`; return the current `featureEnabled` value. + when(mockedVSCodeNamespaces.workspace.getConfiguration(anything(), anything())).thenReturn({ + get: () => featureEnabled + } as unknown as WorkspaceConfiguration); + when(mockedVSCodeNamespaces.workspace.getWorkspaceFolder(anything())).thenCall(() => workspaceFolder); + }); + + /** Wires `IFileSystem.exists`/`readFile` to a small in-memory set of files keyed by fsPath. */ + function configureFileSystem(files: VirtualFile[]): void { + const byPath = new Map(files.map((file) => [file.path, file])); + + when(fileSystem.exists(anything())).thenCall((uri: Uri) => Promise.resolve(byPath.has(uri.fsPath))); + when(fileSystem.readFile(anything())).thenCall((uri: Uri) => { + const file = byPath.get(uri.fsPath); + if (!file) { + return Promise.reject(new Error(`ENOENT: ${uri.fsPath}`)); + } + if (file.readError) { + return Promise.reject(file.readError); + } + + return Promise.resolve(file.content ?? ''); + }); + } + + /** Reads a metadata field off a parsed config without narrowing the metadata union. */ + function metadataField(config: DatabaseIntegrationConfig, key: string): unknown { + return (config.metadata as unknown as Record)[key]; + } + + test('returns configs for a valid integrations file', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: my-postgres + name: My Postgres + type: pgsql + metadata: + host: localhost + port: "5432" + database: mydb + user: root + password: my-secret + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(configs[0].id, 'my-postgres'); + assert.strictEqual(configs[0].type, 'pgsql'); + assert.deepStrictEqual(issues, []); + }); + + test('resolves env: references from the .env file', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: dotenv-postgres + name: Dotenv Postgres + type: pgsql + metadata: + host: localhost + port: "5432" + database: mydb + user: root + password: "env:DEEPNOTE_TEST_DOTENV_PASSWORD" + ` + }, + { path: envPath, content: 'DEEPNOTE_TEST_DOTENV_PASSWORD=secret-from-dotenv\n' } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(metadataField(configs[0], 'password'), 'secret-from-dotenv'); + assert.deepStrictEqual(issues, []); + }); + + test('lets the process environment override values from the .env file', async () => { + provider.processEnvironment = { DEEPNOTE_TEST_OVERRIDE_PASSWORD: 'secret-from-process-env' }; + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: override-postgres + name: Override Postgres + type: pgsql + metadata: + host: localhost + port: "5432" + database: mydb + user: root + password: "env:DEEPNOTE_TEST_OVERRIDE_PASSWORD" + ` + }, + { path: envPath, content: 'DEEPNOTE_TEST_OVERRIDE_PASSWORD=stale-from-dotenv\n' } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(metadataField(configs[0], 'password'), 'secret-from-process-env'); + assert.deepStrictEqual(issues, []); + }); + + test('returns an empty result when the YAML file is missing', async () => { + configureFileSystem([]); + + const result = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(result, { configs: [], issues: [] }); + }); + + test('reports a yaml_parse_error for malformed YAML', async () => { + configureFileSystem([{ path: yamlPath, content: 'integrations:\n - id: "unclosed string' }]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'yaml_parse_error'); + }); + + test('drops an integration and reports env_var_not_defined for an unresolved env: reference', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: missing-env-postgres + name: Missing Env Postgres + type: pgsql + metadata: + host: localhost + port: "5432" + database: mydb + user: root + password: "env:DEEPNOTE_TEST_UNDEFINED_VAR_DEADBEEF" + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'env_var_not_defined'); + }); + + test('finds the YAML at the workspace-folder root when absent next to the .deepnote file', async () => { + const nestedDeepnoteUri = Uri.file('/workspace/project/sub/notebook.deepnote'); + const rootFolder: WorkspaceFolder = { uri: Uri.file('/workspace/project'), name: 'project', index: 0 }; + const rootYamlPath = Uri.joinPath(rootFolder.uri, DEFAULT_INTEGRATIONS_FILE).fsPath; + + workspaceFolder = rootFolder; + configureFileSystem([ + { + path: rootYamlPath, + content: dedent` + integrations: + - id: root-postgres + name: Root Postgres + type: pgsql + metadata: + host: localhost + port: "5432" + database: mydb + user: root + password: my-secret + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(nestedDeepnoteUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(configs[0].id, 'root-postgres'); + assert.deepStrictEqual(issues, []); + }); + + test('drops an integration whose id is reserved (reserved_integration_id)', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: deepnote-dataframe-sql + name: Reserved + type: pgsql + metadata: + host: localhost + port: "5432" + database: mydb + user: root + password: my-secret + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'reserved_integration_id'); + assert.strictEqual(issues[0].path, 'integrations[0]'); + }); + + test('drops an integration with an unsupported pandas-dataframe type (unsupported_integration_type)', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: my-dataframe + name: My Dataframe + type: pandas-dataframe + metadata: {} + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'unsupported_integration_type'); + assert.strictEqual(issues[0].path, 'integrations[0]'); + }); + + test('drops a duplicate id, keeping the first occurrence (duplicate_integration_id)', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: dup-postgres + name: First + type: pgsql + metadata: + host: first-host + port: "5432" + database: mydb + user: root + password: my-secret + - id: dup-postgres + name: Second + type: pgsql + metadata: + host: second-host + port: "5432" + database: mydb + user: root + password: my-secret + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(configs[0].id, 'dup-postgres'); + assert.strictEqual(metadataField(configs[0], 'host'), 'first-host'); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'duplicate_integration_id'); + assert.strictEqual(issues[0].path, 'integrations[1]'); + }); + + test('returns a file_read_error issue (and does not throw) when reading the file fails', async () => { + configureFileSystem([{ path: yamlPath, readError: new Error('disk failure') }]); + + // Must resolve, never reject: a read failure degrades to an issue, not a thrown error. + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'file_read_error'); + assert.strictEqual(issues[0].path, ''); + assert.ok(issues[0].message.includes('Failed to read integrations file')); + }); + + test('drops a federated (google-oauth) integration (unsupported_federated_integration)', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: bq-oauth + name: BigQuery OAuth + type: big-query + metadata: + authMethod: google-oauth + project: my-project + clientId: my-client-id + clientSecret: my-client-secret + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'unsupported_federated_integration'); + assert.strictEqual(issues[0].path, 'integrations[0]'); + }); + + test('returns an empty result without touching the filesystem when the feature is disabled', async () => { + featureEnabled = false; + + const result = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(result, { configs: [], issues: [] }); + verify(fileSystem.exists(anything())).never(); + verify(fileSystem.readFile(anything())).never(); + }); +}); diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts index 7a612e927e..22f630b273 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 { 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,9 @@ 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) + private readonly fileConfigProvider: IIntegrationsFileConfigProvider ) { logger.info('SqlIntegrationEnvironmentVariablesProvider: Constructor called - provider is being instantiated'); // Dispose emitter when extension deactivates @@ -111,19 +100,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 +137,116 @@ 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. 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); failures degrade to [] so SecretStorage still applies. */ + private async loadFileConfigs(notebookUri: Uri): Promise { + try { + const result = await this.fileConfigProvider.getConfigsForFile( + notebookPathToDeepnoteProjectFilePath(notebookUri) + ); + result.issues.forEach((issue) => { + logger.warn( + `SqlIntegrationEnvironmentVariablesProvider: integrations file issue ${issue.code} at '${issue.path}': ${issue.message}` + ); + }); + + return result.configs; + } catch (error) { + logger.error( + 'SqlIntegrationEnvironmentVariablesProvider: file integrations source failed; falling back to SecretStorage', + error + ); + + return []; + } + } + + /** File config wins on id conflict; SecretStorage is the fallback for project ids the file lacks; file-only ids are appended additively (CLI parity). */ + private async mergeIntegrationConfigs( + projectIntegrations: ProjectIntegration[], + fileConfigs: DatabaseIntegrationConfig[] + ): Promise { + const fileConfigsById = new Map(fileConfigs.map((config) => [config.id, config])); + const consumedFileIds = new Set(); + + // Read from SecretStorage only the project integrations the file did not provide. + const secretStorageIds = projectIntegrations + .map((integration) => integration.id) + .filter((id) => !fileConfigsById.has(id)); + const secretStorageResults = await Promise.allSettled( + secretStorageIds.map((id) => this.integrationStorage.getIntegrationConfig(id)) + ); + const secretStorageConfigsById = new Map(); + secretStorageResults.forEach((result, index) => { + const id = secretStorageIds[index]; + if (result.status === 'fulfilled') { + if (result.value) { + secretStorageConfigsById.set(id, result.value); + } + + return; + } + logger.error( + `SqlIntegrationEnvironmentVariablesProvider: Failed to load integration config ${id}`, + result.reason + ); + }); + + // Resolve each project integration in declared order: file config wins, else the SecretStorage fallback. + const allConfigs: Array = []; + for (const integration of projectIntegrations) { + const fileConfig = fileConfigsById.get(integration.id); + if (fileConfig) { + consumedFileIds.add(integration.id); + allConfigs.push(fileConfig); + + continue; + } + const secretStorageConfig = secretStorageConfigsById.get(integration.id); + if (secretStorageConfig) { + allConfigs.push(secretStorageConfig); + } + } + + // Append file-only integrations (not declared in project.integrations) additively, deduped by the map. + for (const fileConfig of fileConfigsById.values()) { + if (!consumedFileIds.has(fileConfig.id)) { + allConfigs.push(fileConfig); + } + } + + return allConfigs; + } } diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts index 11b5e57eaa..a99ea636c9 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. */ @@ -29,6 +34,11 @@ function createMockProject( }; } +/** A file-config source that yields nothing — what callers see when no `.deepnote.env.yaml` exists. */ +function emptyFileConfigProvider(): IIntegrationsFileConfigProvider { + return { getConfigsForFile: async () => ({ configs: [], issues: [] }) }; +} + suite('SqlIntegrationEnvironmentVariablesProvider', () => { let provider: SqlIntegrationEnvironmentVariablesProvider; let integrationStorage: IIntegrationStorage; @@ -50,7 +60,8 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { instance(integrationStorage), instance(notebookEditorProvider), instance(notebookManager), - disposables + disposables, + emptyFileConfigProvider() ); }); @@ -437,6 +448,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 source yields nothing: behavior is SecretStorage-only (unchanged)', async () => { + const providerWithoutFile = new SqlIntegrationEnvironmentVariablesProvider( + instance(integrationStorage), + instance(notebookEditorProvider), + instance(notebookManager), + disposables, + emptyFileConfigProvider() + ); + 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/sqlIntegrationEnvironmentVariablesProvider.web.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.web.ts new file mode 100644 index 0000000000..d629cef359 --- /dev/null +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.web.ts @@ -0,0 +1,32 @@ +import { injectable } from 'inversify'; +import { EventEmitter } from 'vscode'; + +import { DatabaseIntegrationConfig } from '@deepnote/database-integrations'; + +import { Resource } from '../../common/types'; +import { DisposableBase } from '../../common/utils/lifecycle'; +import { EnvironmentVariables } from '../../common/variables/types'; +import { ISqlIntegrationEnvVarsProvider } from './types'; + +/** + * Web stub for `ISqlIntegrationEnvVarsProvider`. Integration credentials are only resolved for locally + * launched kernels, which web does not have, so this yields no env vars and no configs. Binding it keeps + * the dependency non-optional for every consumer instead of making them branch on a missing provider. + */ +@injectable() +export class SqlIntegrationEnvironmentVariablesProviderWeb + extends DisposableBase + implements ISqlIntegrationEnvVarsProvider +{ + private readonly _onDidChangeEnvironmentVariables = this._register(new EventEmitter()); + + public readonly onDidChangeEnvironmentVariables = this._onDidChangeEnvironmentVariables.event; + + public async getEnvironmentVariables(): Promise { + return {}; + } + + public async getMergedConfigs(): Promise { + return []; + } +} diff --git a/src/platform/notebooks/deepnote/types.ts b/src/platform/notebooks/deepnote/types.ts index 23be88f5da..f0cf89e6f7 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/src/webviews/webview-side/integrations/IntegrationItem.tsx b/src/webviews/webview-side/integrations/IntegrationItem.tsx index 4c3a206727..1fe1726a47 100644 --- a/src/webviews/webview-side/integrations/IntegrationItem.tsx +++ b/src/webviews/webview-side/integrations/IntegrationItem.tsx @@ -2,11 +2,11 @@ import * as React from 'react'; import { BigQueryAuthMethods } from '@deepnote/database-integrations'; import { getLocString } from '../react-common/locReactSide'; -import { ConfigurableDatabaseIntegrationType, IntegrationWithStatus } from './types'; +import { ConfigurableDatabaseIntegrationType, DetectedIntegration } from './types'; import { integrationTypeIcons } from './integrationUtils'; export interface IIntegrationItemProps { - integration: IntegrationWithStatus; + integration: DetectedIntegration; onConfigure: (integrationId: string) => void; onReset: (integrationId: string) => void; onDelete: (integrationId: string) => void; @@ -63,11 +63,11 @@ export const IntegrationItem: React.FC = ({ onDelete, onAuthenticate }) => { - const statusClass = integration.status === 'connected' ? 'status-connected' : 'status-disconnected'; - const statusText = - integration.status === 'connected' - ? getLocString('integrationsConnected', 'Connected') - : getLocString('integrationsNotConfigured', 'Not Configured'); + // Credentials the panel can edit live in SecretStorage, which is exactly what `config` holds. + const statusClass = integration.config ? 'status-connected' : 'status-disconnected'; + const statusText = integration.config + ? getLocString('integrationsConnected', 'Connected') + : getLocString('integrationsNotConfigured', 'Not Configured'); const configureText = integration.config ? getLocString('integrationsReconfigure', 'Reconfigure') : getLocString('integrationsConfigure', 'Configure'); diff --git a/src/webviews/webview-side/integrations/IntegrationList.tsx b/src/webviews/webview-side/integrations/IntegrationList.tsx index d87a7a14a6..be9aef954b 100644 --- a/src/webviews/webview-side/integrations/IntegrationList.tsx +++ b/src/webviews/webview-side/integrations/IntegrationList.tsx @@ -1,10 +1,10 @@ import * as React from 'react'; import { getLocString } from '../react-common/locReactSide'; import { IntegrationItem } from './IntegrationItem'; -import { IntegrationWithStatus } from './types'; +import { DetectedIntegration } from './types'; export interface IIntegrationListProps { - integrations: IntegrationWithStatus[]; + integrations: DetectedIntegration[]; onConfigure: (integrationId: string) => void; onReset: (integrationId: string) => void; onDelete: (integrationId: string) => void; diff --git a/src/webviews/webview-side/integrations/IntegrationPanel.tsx b/src/webviews/webview-side/integrations/IntegrationPanel.tsx index 6cf4349e0b..45ce11c2a8 100644 --- a/src/webviews/webview-side/integrations/IntegrationPanel.tsx +++ b/src/webviews/webview-side/integrations/IntegrationPanel.tsx @@ -9,7 +9,7 @@ import { ConfigurationForm } from './ConfigurationForm'; import { ConfigurableDatabaseIntegrationConfig, ConfigurableDatabaseIntegrationType, - IntegrationWithStatus, + DetectedIntegration, WebviewMessage, WebviewOutboundMessage } from './types'; @@ -20,7 +20,7 @@ export interface IIntegrationPanelProps { } export const IntegrationPanel: React.FC = ({ baseTheme, vscodeApi }) => { - const [integrations, setIntegrations] = React.useState([]); + const [integrations, setIntegrations] = React.useState([]); const [projectName, setProjectName] = React.useState(undefined); const [selectedIntegrationId, setSelectedIntegrationId] = React.useState(null); const [selectedConfig, setSelectedConfig] = React.useState(null); diff --git a/src/webviews/webview-side/integrations/types.ts b/src/webviews/webview-side/integrations/types.ts index 2b9b259cd4..247cd3a2bb 100644 --- a/src/webviews/webview-side/integrations/types.ts +++ b/src/webviews/webview-side/integrations/types.ts @@ -4,15 +4,12 @@ export type ConfigurableDatabaseIntegrationType = Exclude; -export type IntegrationStatus = 'connected' | 'disconnected' | 'error'; - /** Federated-auth token status; mirrors `FederatedAuthTokenStatus` in platform/integrationTypes.ts (duplicated because the webview bundles separately). */ export type FederatedAuthTokenStatus = 'authenticated' | 'disconnected' | 'unsupported'; -export interface IntegrationWithStatus { +export interface DetectedIntegration { id: string; config: ConfigurableDatabaseIntegrationConfig | null; - status: IntegrationStatus; integrationName?: string; integrationType?: ConfigurableDatabaseIntegrationType; tokenStatus?: FederatedAuthTokenStatus; @@ -26,7 +23,7 @@ export interface IVsCodeMessage { export interface UpdateMessage { type: 'update'; - integrations: IntegrationWithStatus[]; + integrations: DetectedIntegration[]; projectName?: string; } diff --git a/test/e2e/fixtures/integrations-env-file.deepnote b/test/e2e/fixtures/integrations-env-file.deepnote new file mode 100644 index 0000000000..5ca88bb04b --- /dev/null +++ b/test/e2e/fixtures/integrations-env-file.deepnote @@ -0,0 +1,26 @@ +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00.000Z' + modifiedAt: '2025-01-01T00:00:00.000Z' +project: + id: e2e-integrations-env-file-project + name: E2E Integrations Env File + integrations: + - id: e2e-pg-integration + name: Prod Postgres + type: pgsql + notebooks: + - id: e2e-integrations-env-file-notebook + name: Integrations Env File + blocks: + - id: e2e-integrations-block + blockGroup: e2e-group + type: code + content: |- + import os + print(os.environ.get('PROD_POSTGRES_HOST')) + sortingKey: a0 + metadata: {} + executionMode: block + isModule: false + settings: {} diff --git a/test/e2e/suite/integrationsEnvFileInjection.e2e.test.ts b/test/e2e/suite/integrationsEnvFileInjection.e2e.test.ts new file mode 100644 index 0000000000..d79feb59f4 --- /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'); + }); +});