From 9bc4e4485b9a0ed9dae0c4f66c0b5b24335c05c1 Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 13:32:51 +0100 Subject: [PATCH 01/23] refactor: rename IntegrationConfig to LegacyIntegrationConfig --- .../deepnote/integrations/integrationWebview.ts | 6 +++--- .../notebooks/deepnote/integrationStorage.ts | 16 ++++++++-------- .../notebooks/deepnote/integrationTypes.ts | 7 +++++-- ...sqlIntegrationEnvironmentVariablesProvider.ts | 6 +++--- src/platform/notebooks/deepnote/types.ts | 10 +++++----- 5 files changed, 24 insertions(+), 21 deletions(-) diff --git a/src/notebooks/deepnote/integrations/integrationWebview.ts b/src/notebooks/deepnote/integrations/integrationWebview.ts index c94e87ade3..3feb8e54e2 100644 --- a/src/notebooks/deepnote/integrations/integrationWebview.ts +++ b/src/notebooks/deepnote/integrations/integrationWebview.ts @@ -9,7 +9,7 @@ import { IDeepnoteNotebookManager, ProjectIntegration } from '../../types'; import { IIntegrationStorage, IIntegrationWebviewProvider } from './types'; import { INTEGRATION_TYPE_TO_DEEPNOTE, - IntegrationConfig, + LegacyIntegrationConfig, IntegrationStatus, IntegrationWithStatus, RawIntegrationType @@ -221,7 +221,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { private async handleMessage(message: { type: string; integrationId?: string; - config?: IntegrationConfig; + config?: LegacyIntegrationConfig; }): Promise { switch (message.type) { case 'configure': @@ -263,7 +263,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { /** * Save the configuration for an integration */ - private async saveConfiguration(integrationId: string, config: IntegrationConfig): Promise { + private async saveConfiguration(integrationId: string, config: LegacyIntegrationConfig): Promise { try { await this.integrationStorage.save(config); diff --git a/src/platform/notebooks/deepnote/integrationStorage.ts b/src/platform/notebooks/deepnote/integrationStorage.ts index 6b86d52f1b..ef10acd7eb 100644 --- a/src/platform/notebooks/deepnote/integrationStorage.ts +++ b/src/platform/notebooks/deepnote/integrationStorage.ts @@ -4,7 +4,7 @@ import { EventEmitter } from 'vscode'; import { IEncryptedStorage } from '../../common/application/types'; import { IAsyncDisposableRegistry } from '../../common/types'; import { logger } from '../../logging'; -import { IntegrationConfig, IntegrationType } from './integrationTypes'; +import { LegacyIntegrationConfig, IntegrationType } from './integrationTypes'; import { IIntegrationStorage } from './types'; const INTEGRATION_SERVICE_NAME = 'deepnote-integrations'; @@ -16,7 +16,7 @@ const INTEGRATION_SERVICE_NAME = 'deepnote-integrations'; */ @injectable() export class IntegrationStorage implements IIntegrationStorage { - private readonly cache: Map = new Map(); + private readonly cache: Map = new Map(); private cacheLoaded = false; @@ -35,7 +35,7 @@ export class IntegrationStorage implements IIntegrationStorage { /** * Get all stored integration configurations */ - async getAll(): Promise { + async getAll(): Promise { await this.ensureCacheLoaded(); return Array.from(this.cache.values()); } @@ -43,7 +43,7 @@ export class IntegrationStorage implements IIntegrationStorage { /** * Get a specific integration configuration by ID */ - async getIntegrationConfig(integrationId: string): Promise { + async getIntegrationConfig(integrationId: string): Promise { await this.ensureCacheLoaded(); return this.cache.get(integrationId); } @@ -56,14 +56,14 @@ export class IntegrationStorage implements IIntegrationStorage { async getProjectIntegrationConfig( _projectId: string, integrationId: string - ): Promise { + ): Promise { return this.getIntegrationConfig(integrationId); } /** * Get all integrations of a specific type */ - async getByType(type: IntegrationType): Promise { + async getByType(type: IntegrationType): Promise { await this.ensureCacheLoaded(); return Array.from(this.cache.values()).filter((config) => config.type === type); } @@ -71,7 +71,7 @@ export class IntegrationStorage implements IIntegrationStorage { /** * Save or update an integration configuration */ - async save(config: IntegrationConfig): Promise { + async save(config: LegacyIntegrationConfig): Promise { await this.ensureCacheLoaded(); // Store the configuration as JSON in encrypted storage @@ -160,7 +160,7 @@ export class IntegrationStorage implements IIntegrationStorage { const configJson = await this.encryptedStorage.retrieve(INTEGRATION_SERVICE_NAME, id); if (configJson) { try { - const config: IntegrationConfig = JSON.parse(configJson); + const config: LegacyIntegrationConfig = JSON.parse(configJson); this.cache.set(id, config); } catch (error) { logger.error(`Failed to parse integration config for ${id}:`, error); diff --git a/src/platform/notebooks/deepnote/integrationTypes.ts b/src/platform/notebooks/deepnote/integrationTypes.ts index 8c4442e8bc..9e1298bafe 100644 --- a/src/platform/notebooks/deepnote/integrationTypes.ts +++ b/src/platform/notebooks/deepnote/integrationTypes.ts @@ -119,7 +119,10 @@ export type SnowflakeIntegrationConfig = BaseSnowflakeConfig & /** * Union type of all integration configurations */ -export type IntegrationConfig = PostgresIntegrationConfig | BigQueryIntegrationConfig | SnowflakeIntegrationConfig; +export type LegacyIntegrationConfig = + | PostgresIntegrationConfig + | BigQueryIntegrationConfig + | SnowflakeIntegrationConfig; /** * Integration connection status @@ -134,7 +137,7 @@ export enum IntegrationStatus { * Integration with its current status */ export interface IntegrationWithStatus { - config: IntegrationConfig | null; + config: LegacyIntegrationConfig | null; status: IntegrationStatus; error?: string; /** diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts index cfd9df1bdd..bac055f149 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts @@ -8,7 +8,7 @@ import { logger } from '../../logging'; import { IIntegrationStorage, ISqlIntegrationEnvVarsProvider } from './types'; import { DATAFRAME_SQL_INTEGRATION_ID, - IntegrationConfig, + LegacyIntegrationConfig, IntegrationType, SnowflakeAuthMethods } from './integrationTypes'; @@ -34,7 +34,7 @@ function getSqlEnvVarName(integrationId: string): string { * "param_style": "qmark" | "format" | etc. * } */ -function convertIntegrationConfigToJson(config: IntegrationConfig): string { +function convertIntegrationConfigToJson(config: LegacyIntegrationConfig): string { switch (config.type) { case IntegrationType.Postgres: { // Build PostgreSQL connection URL @@ -141,7 +141,7 @@ function convertIntegrationConfigToJson(config: IntegrationConfig): string { default: throw new UnsupportedIntegrationError( - l10n.t('Unsupported integration type: {0}', (config as IntegrationConfig).type) + l10n.t('Unsupported integration type: {0}', (config as LegacyIntegrationConfig).type) ); } } diff --git a/src/platform/notebooks/deepnote/types.ts b/src/platform/notebooks/deepnote/types.ts index ff2a1cd0fe..98de30ccde 100644 --- a/src/platform/notebooks/deepnote/types.ts +++ b/src/platform/notebooks/deepnote/types.ts @@ -1,7 +1,7 @@ import { CancellationToken, Event } from 'vscode'; import { IDisposable, Resource } from '../../common/types'; import { EnvironmentVariables } from '../../common/variables/types'; -import { IntegrationConfig } from './integrationTypes'; +import { LegacyIntegrationConfig } from './integrationTypes'; /** * Settings for select input blocks @@ -30,7 +30,7 @@ export interface IIntegrationStorage extends IDisposable { */ readonly onDidChangeIntegrations: Event; - getAll(): Promise; + getAll(): Promise; /** * Retrieves the global (non-project-scoped) integration configuration by integration ID. @@ -48,14 +48,14 @@ export interface IIntegrationStorage extends IDisposable { * - The `IntegrationConfig` object if a global configuration exists for the given ID * - `undefined` if no global configuration exists for the given integration ID */ - getIntegrationConfig(integrationId: string): Promise; + getIntegrationConfig(integrationId: string): Promise; /** * Get integration configuration for a specific project and integration */ - getProjectIntegrationConfig(projectId: string, integrationId: string): Promise; + getProjectIntegrationConfig(projectId: string, integrationId: string): Promise; - save(config: IntegrationConfig): Promise; + save(config: LegacyIntegrationConfig): Promise; delete(integrationId: string): Promise; exists(integrationId: string): Promise; clear(): Promise; From 37815078b90ad96be90fe13711b1ea9a6002a696 Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 14:31:24 +0100 Subject: [PATCH 02/23] replace notebook scanning for integrations with full project integration list --- INTEGRATIONS_CREDENTIALS.md | 9 +- ...IntegrationEnvironmentVariablesProvider.ts | 79 ++--- ...nEnvironmentVariablesProvider.unit.test.ts | 270 ++---------------- 3 files changed, 43 insertions(+), 315 deletions(-) diff --git a/INTEGRATIONS_CREDENTIALS.md b/INTEGRATIONS_CREDENTIALS.md index 108e33a5c0..e5f54df12c 100644 --- a/INTEGRATIONS_CREDENTIALS.md +++ b/INTEGRATIONS_CREDENTIALS.md @@ -216,10 +216,11 @@ Provides environment variables containing integration credentials for the Jupyte **Process:** -1. Scans the notebook for SQL cells with `sql_integration_id` metadata -2. Retrieves credentials for each detected integration -3. Converts credentials to the format expected by `deepnote-toolkit` -4. Returns environment variables to be injected into the kernel process +1. Retrieves all configured integrations from `IIntegrationStorage` +2. Converts credentials to the format expected by `deepnote-toolkit` +3. Returns environment variables to be injected into the kernel process + +**Note:** This provider makes credentials for ALL configured integrations available as environment variables, not just those used in the current notebook. This ensures that integrations are available project-wide, matching Deepnote's behavior where integrations are project-scoped. **Environment Variable Format:** diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts index bac055f149..c27f3e2de1 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts @@ -1,5 +1,5 @@ import { inject, injectable } from 'inversify'; -import { CancellationToken, Event, EventEmitter, l10n, NotebookDocument, workspace } from 'vscode'; +import { CancellationToken, Event, EventEmitter, l10n } from 'vscode'; import { IDisposableRegistry, Resource } from '../../common/types'; import { EnvironmentVariables } from '../../common/variables/types'; @@ -148,7 +148,7 @@ function convertIntegrationConfigToJson(config: LegacyIntegrationConfig): string /** * Provides environment variables for SQL integrations. - * This service scans notebooks for SQL blocks and injects the necessary credentials + * This service provides credentials for all configured integrations in the project * as environment variables so they can be used during SQL block execution. */ @injectable() @@ -174,7 +174,8 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati } /** - * Get environment variables for SQL integrations used in the given notebook. + * Get environment variables for SQL integrations. + * Provides credentials for all configured integrations in the project. */ public async getEnvironmentVariables(resource: Resource, token?: CancellationToken): Promise { const envVars: EnvironmentVariables = {}; @@ -188,16 +189,6 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati } logger.trace(`SqlIntegrationEnvironmentVariablesProvider: Getting env vars for resource`); - logger.trace( - `SqlIntegrationEnvironmentVariablesProvider: Available notebooks count: ${workspace.notebookDocuments.length}` - ); - - // Find the notebook document for this resource - const notebook = workspace.notebookDocuments.find((nb) => nb.uri.toString() === resource.toString()); - if (!notebook) { - logger.warn(`SqlIntegrationEnvironmentVariablesProvider: No notebook found for ${resource.toString()}`); - return envVars; - } // Always add the internal DuckDB integration const dataframeSqlIntegrationEnvVarName = convertToEnvironmentVariableName( @@ -212,49 +203,40 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati envVars[dataframeSqlIntegrationEnvVarName] = dataframeSqlIntegrationCredentialsJson; logger.debug(`SqlIntegrationEnvironmentVariablesProvider: Added env var for dataframe SQL integration`); - // Scan all cells for SQL integration IDs - const integrationIds = this.scanNotebookForIntegrations(notebook); - if (integrationIds.size === 0) { - logger.info( - `SqlIntegrationEnvironmentVariablesProvider: No SQL integrations found in ${resource.toString()}` - ); + // Get all configured integrations from storage + const allIntegrations = await this.integrationStorage.getAll(); + if (allIntegrations.length === 0) { + logger.info(`SqlIntegrationEnvironmentVariablesProvider: No configured integrations found`); return envVars; } - logger.trace(`SqlIntegrationEnvironmentVariablesProvider: Found ${integrationIds.size} SQL integrations`); + logger.trace( + `SqlIntegrationEnvironmentVariablesProvider: Found ${allIntegrations.length} configured integrations` + ); // Get credentials for each integration and add to environment variables - for (const integrationId of integrationIds) { + for (const config of allIntegrations) { if (token?.isCancellationRequested) { break; } try { - // Handle internal DuckDB integration specially - if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) { - // Internal DuckDB integration is handled above - continue; - } - - const config = await this.integrationStorage.getIntegrationConfig(integrationId); - if (!config) { - logger.warn( - `SqlIntegrationEnvironmentVariablesProvider: No configuration found for integration ${integrationId}` - ); + // Skip internal DuckDB integration (already handled above) + if (config.id === DATAFRAME_SQL_INTEGRATION_ID) { continue; } // Convert integration config to JSON and add as environment variable - const envVarName = convertToEnvironmentVariableName(getSqlEnvVarName(integrationId)); + const envVarName = convertToEnvironmentVariableName(getSqlEnvVarName(config.id)); const credentialsJson = convertIntegrationConfigToJson(config); envVars[envVarName] = credentialsJson; logger.debug( - `SqlIntegrationEnvironmentVariablesProvider: Added env var ${envVarName} for integration ${integrationId}` + `SqlIntegrationEnvironmentVariablesProvider: Added env var ${envVarName} for integration ${config.id}` ); } catch (error) { logger.error( - `SqlIntegrationEnvironmentVariablesProvider: Failed to get credentials for integration ${integrationId}`, + `SqlIntegrationEnvironmentVariablesProvider: Failed to get credentials for integration ${config.id}`, error ); } @@ -264,31 +246,4 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati return envVars; } - - /** - * Scan a notebook for SQL integration IDs. - */ - private scanNotebookForIntegrations(notebook: NotebookDocument): Set { - const integrationIds = new Set(); - - for (const cell of notebook.getCells()) { - // Only check 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') { - integrationIds.add(integrationId); - logger.trace( - `SqlIntegrationEnvironmentVariablesProvider: Found integration ${integrationId} in cell ${cell.index}` - ); - } - } - } - - return integrationIds; - } } diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts index 92026b7638..bc50a7626e 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts @@ -1,6 +1,6 @@ import { assert } from 'chai'; import { instance, mock, when } from 'ts-mockito'; -import { CancellationTokenSource, EventEmitter, NotebookCell, NotebookCellKind, NotebookDocument, Uri } from 'vscode'; +import { CancellationTokenSource, EventEmitter, Uri } from 'vscode'; import { IDisposableRegistry } from '../../common/types'; import { IntegrationStorage } from './integrationStorage'; @@ -9,11 +9,9 @@ import { IntegrationType, PostgresIntegrationConfig, BigQueryIntegrationConfig, - DATAFRAME_SQL_INTEGRATION_ID, SnowflakeIntegrationConfig, SnowflakeAuthMethods } from './integrationTypes'; -import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; const EXPECTED_DATAFRAME_ONLY_ENV_VARS = { SQL_DEEPNOTE_DATAFRAME_SQL: '{"url":"deepnote+duckdb:///:memory:","params":{},"param_style":"qmark"}' @@ -25,7 +23,6 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { let disposables: IDisposableRegistry; setup(() => { - resetVSCodeMocks(); disposables = []; integrationStorage = mock(IntegrationStorage); when(integrationStorage.onDidChangeIntegrations).thenReturn(new EventEmitter().event); @@ -42,34 +39,9 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { assert.deepStrictEqual(envVars, {}); }); - test('Returns empty object when notebook is not found', async () => { + test('Returns only dataframe integration when no integrations are configured', async () => { const uri = Uri.file('/test/notebook.deepnote'); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); - - const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, {}); - }); - - test('Returns empty object when notebook has no SQL cells', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'python', 'print("hello")'), - createMockCell(1, NotebookCellKind.Markup, 'markdown', '# Title') - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - - const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); - }); - - test('Returns empty object when SQL cells have no integration ID', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM users', {}) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(integrationStorage.getAll()).thenResolve([]); const envVars = await provider.getEnvironmentVariables(uri); assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); @@ -77,13 +49,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Returns environment variable for internal DuckDB integration', async () => { const uri = Uri.file('/test/notebook.deepnote'); - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM df', { - sql_integration_id: DATAFRAME_SQL_INTEGRATION_ID - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(integrationStorage.getAll()).thenResolve([]); const envVars = await provider.getEnvironmentVariables(uri); @@ -110,14 +76,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { ssl: true }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM users', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); @@ -141,14 +100,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { credentials: serviceAccountJson }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM dataset.table', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); @@ -163,40 +115,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { assert.strictEqual(credentialsJson.param_style, 'format'); }); - test('Handles multiple SQL cells with same integration', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'my-postgres-db'; - const config: PostgresIntegrationConfig = { - id: integrationId, - name: 'My Postgres DB', - type: IntegrationType.Postgres, - host: 'localhost', - port: 5432, - database: 'mydb', - username: 'user', - password: 'pass' - }; - - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM users', { - sql_integration_id: integrationId - }), - createMockCell(1, NotebookCellKind.Code, 'sql', 'SELECT * FROM orders', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - - // Should only have one environment variable apart from the internal DuckDB integration - assert.property(envVars, 'SQL_MY_POSTGRES_DB'); - assert.strictEqual(Object.keys(envVars).length, 2); - }); - - test('Handles multiple SQL cells with different integrations', async () => { + test('Returns environment variables for all configured integrations', async () => { const uri = Uri.file('/test/notebook.deepnote'); const postgresId = 'my-postgres-db'; const bigqueryId = 'my-bigquery'; @@ -220,18 +139,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { credentials: JSON.stringify({ type: 'service_account' }) }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM users', { - sql_integration_id: postgresId - }), - createMockCell(1, NotebookCellKind.Code, 'sql', 'SELECT * FROM dataset.table', { - sql_integration_id: bigqueryId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(postgresId)).thenResolve(postgresConfig); - when(integrationStorage.getIntegrationConfig(bigqueryId)).thenResolve(bigqueryConfig); + when(integrationStorage.getAll()).thenResolve([postgresConfig, bigqueryConfig]); const envVars = await provider.getEnvironmentVariables(uri); @@ -241,25 +149,6 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { assert.strictEqual(Object.keys(envVars).length, 3); }); - test('Handles missing integration configuration gracefully', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'missing-integration'; - - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM users', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(undefined); - - const envVars = await provider.getEnvironmentVariables(uri); - - // Should return only dataframe integration when integration config is missing - assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); - }); - test('Properly encodes special characters in PostgreSQL credentials', async () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'special-chars-db'; @@ -275,14 +164,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { ssl: false }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM users', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); @@ -314,14 +196,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { ssl: true }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM products', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); @@ -349,14 +224,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { ssl: false }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); @@ -369,15 +237,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Honors CancellationToken (returns empty when cancelled early)', async () => { const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'cancel-me'; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', { sql_integration_id: integrationId }) - ]); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - // Return a slow promise to ensure cancellation path is hit - when(integrationStorage.getIntegrationConfig(integrationId)).thenCall( - () => new Promise((resolve) => setTimeout(() => resolve(undefined), 50)) - ); + when(integrationStorage.getAll()).thenCall(() => new Promise((resolve) => setTimeout(() => resolve([]), 50))); const cts = new CancellationTokenSource(); cts.cancel(); const envVars = await provider.getEnvironmentVariables(uri, cts.token); @@ -401,14 +261,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'secret123' }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM customers', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); @@ -437,14 +290,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'pass' }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); @@ -476,14 +322,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { privateKeyPassphrase: 'passphrase123' }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT * FROM events', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); @@ -517,14 +356,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { privateKey: privateKey }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); @@ -555,14 +387,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'p@ss:word!#$%' }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); @@ -588,14 +413,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'pass' }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); @@ -617,14 +435,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { authMethod: SnowflakeAuthMethods.OKTA }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); // Should return only dataframe integration when unsupported auth method is encountered const envVars = await provider.getEnvironmentVariables(uri); @@ -642,14 +453,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { authMethod: SnowflakeAuthMethods.AZURE_AD }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); @@ -666,42 +470,10 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { authMethod: SnowflakeAuthMethods.KEY_PAIR }; - const notebook = createMockNotebook(uri, [ - createMockCell(0, NotebookCellKind.Code, 'sql', 'SELECT 1', { - sql_integration_id: integrationId - }) - ]); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); + when(integrationStorage.getAll()).thenResolve([config]); const envVars = await provider.getEnvironmentVariables(uri); assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); }); }); }); - -function createMockNotebook(uri: Uri, cells: NotebookCell[]): NotebookDocument { - return { - uri, - getCells: () => cells - } as NotebookDocument; -} - -function createMockCell( - index: number, - kind: NotebookCellKind, - languageId: string, - value: string, - metadata?: Record -): NotebookCell { - return { - index, - kind, - document: { - languageId, - getText: () => value - }, - metadata: metadata || {} - } as NotebookCell; -} From 702cb0a4d91a6de84e4cbbe430f469087385162c Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 14:32:22 +0100 Subject: [PATCH 03/23] move duckdb integration inclusion from env var provider to integration storage --- .../integrations/integrationWebview.ts | 7 +++ .../notebooks/deepnote/integrationStorage.ts | 21 +++++++-- .../notebooks/deepnote/integrationTypes.ts | 16 +++++-- ...IntegrationEnvironmentVariablesProvider.ts | 41 +++++----------- ...nEnvironmentVariablesProvider.unit.test.ts | 47 ++++++++++++------- 5 files changed, 79 insertions(+), 53 deletions(-) diff --git a/src/notebooks/deepnote/integrations/integrationWebview.ts b/src/notebooks/deepnote/integrations/integrationWebview.ts index 3feb8e54e2..3b740a4cc4 100644 --- a/src/notebooks/deepnote/integrations/integrationWebview.ts +++ b/src/notebooks/deepnote/integrations/integrationWebview.ts @@ -9,6 +9,7 @@ import { IDeepnoteNotebookManager, ProjectIntegration } from '../../types'; import { IIntegrationStorage, IIntegrationWebviewProvider } from './types'; import { INTEGRATION_TYPE_TO_DEEPNOTE, + IntegrationType, LegacyIntegrationConfig, IntegrationStatus, IntegrationWithStatus, @@ -349,6 +350,12 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { return null; } + // Skip DuckDB integration (internal, not a real Deepnote integration) + if (type === IntegrationType.DuckDB) { + logger.trace(`IntegrationWebviewProvider: Skipping internal DuckDB integration ${id}`); + return null; + } + // Map to Deepnote integration type const deepnoteType: RawIntegrationType | undefined = INTEGRATION_TYPE_TO_DEEPNOTE[type]; if (!deepnoteType) { diff --git a/src/platform/notebooks/deepnote/integrationStorage.ts b/src/platform/notebooks/deepnote/integrationStorage.ts index ef10acd7eb..68b6130f52 100644 --- a/src/platform/notebooks/deepnote/integrationStorage.ts +++ b/src/platform/notebooks/deepnote/integrationStorage.ts @@ -4,7 +4,12 @@ import { EventEmitter } from 'vscode'; import { IEncryptedStorage } from '../../common/application/types'; import { IAsyncDisposableRegistry } from '../../common/types'; import { logger } from '../../logging'; -import { LegacyIntegrationConfig, IntegrationType } from './integrationTypes'; +import { + LegacyIntegrationConfig, + IntegrationType, + DuckDBIntegrationConfig, + DATAFRAME_SQL_INTEGRATION_ID +} from './integrationTypes'; import { IIntegrationStorage } from './types'; const INTEGRATION_SERVICE_NAME = 'deepnote-integrations'; @@ -33,11 +38,21 @@ export class IntegrationStorage implements IIntegrationStorage { } /** - * Get all stored integration configurations + * Get all stored integration configurations. + * Always includes the internal DuckDB integration at the end of the list. */ async getAll(): Promise { await this.ensureCacheLoaded(); - return Array.from(this.cache.values()); + const configs = Array.from(this.cache.values()); + + // Always add the internal DuckDB integration at the end + const duckdbConfig: DuckDBIntegrationConfig = { + id: DATAFRAME_SQL_INTEGRATION_ID, + name: 'Dataframe SQL (DuckDB)', + type: IntegrationType.DuckDB + }; + + return [...configs, duckdbConfig]; } /** diff --git a/src/platform/notebooks/deepnote/integrationTypes.ts b/src/platform/notebooks/deepnote/integrationTypes.ts index 9e1298bafe..14aa2afbff 100644 --- a/src/platform/notebooks/deepnote/integrationTypes.ts +++ b/src/platform/notebooks/deepnote/integrationTypes.ts @@ -10,17 +10,19 @@ export const DATAFRAME_SQL_INTEGRATION_ID = 'deepnote-dataframe-sql'; export enum IntegrationType { Postgres = 'postgres', BigQuery = 'bigquery', - Snowflake = 'snowflake' + Snowflake = 'snowflake', + DuckDB = 'duckdb' } /** * Map our IntegrationType enum to Deepnote integration type strings + * Note: DuckDB is not included as it's an internal integration that doesn't exist in Deepnote */ export const INTEGRATION_TYPE_TO_DEEPNOTE = { [IntegrationType.Postgres]: 'pgsql', [IntegrationType.BigQuery]: 'big-query', [IntegrationType.Snowflake]: 'snowflake' -} as const satisfies { [type in IntegrationType]: string }; +} as const satisfies { [type in Exclude]: string }; export type RawIntegrationType = (typeof INTEGRATION_TYPE_TO_DEEPNOTE)[keyof typeof INTEGRATION_TYPE_TO_DEEPNOTE]; @@ -64,6 +66,13 @@ export interface BigQueryIntegrationConfig extends BaseIntegrationConfig { credentials: string; // JSON string of service account credentials } +/** + * DuckDB integration configuration (internal, always available) + */ +export interface DuckDBIntegrationConfig extends BaseIntegrationConfig { + type: IntegrationType.DuckDB; +} + // Import and re-export Snowflake auth constants from shared module import { type SnowflakeAuthMethod, @@ -122,7 +131,8 @@ export type SnowflakeIntegrationConfig = BaseSnowflakeConfig & export type LegacyIntegrationConfig = | PostgresIntegrationConfig | BigQueryIntegrationConfig - | SnowflakeIntegrationConfig; + | SnowflakeIntegrationConfig + | DuckDBIntegrationConfig; /** * Integration connection status diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts index c27f3e2de1..781adfebf1 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts @@ -6,12 +6,7 @@ import { EnvironmentVariables } from '../../common/variables/types'; import { UnsupportedIntegrationError } from '../../errors/unsupportedIntegrationError'; import { logger } from '../../logging'; import { IIntegrationStorage, ISqlIntegrationEnvVarsProvider } from './types'; -import { - DATAFRAME_SQL_INTEGRATION_ID, - LegacyIntegrationConfig, - IntegrationType, - SnowflakeAuthMethods -} from './integrationTypes'; +import { LegacyIntegrationConfig, IntegrationType, SnowflakeAuthMethods } from './integrationTypes'; /** * Converts an integration ID to the environment variable name format expected by SQL blocks. @@ -36,6 +31,15 @@ function getSqlEnvVarName(integrationId: string): string { */ function convertIntegrationConfigToJson(config: LegacyIntegrationConfig): string { switch (config.type) { + case IntegrationType.DuckDB: { + // Internal DuckDB integration for querying dataframes + return JSON.stringify({ + url: 'deepnote+duckdb:///:memory:', + params: {}, + param_style: 'qmark' + }); + } + case IntegrationType.Postgres: { // Build PostgreSQL connection URL // Format: postgresql://username:password@host:port/database @@ -176,6 +180,7 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati /** * Get environment variables for SQL integrations. * Provides credentials for all configured integrations in the project. + * The internal DuckDB integration is always included via IntegrationStorage.getAll(). */ public async getEnvironmentVariables(resource: Resource, token?: CancellationToken): Promise { const envVars: EnvironmentVariables = {}; @@ -190,25 +195,8 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati logger.trace(`SqlIntegrationEnvironmentVariablesProvider: Getting env vars for resource`); - // Always add the internal DuckDB integration - const dataframeSqlIntegrationEnvVarName = convertToEnvironmentVariableName( - getSqlEnvVarName(DATAFRAME_SQL_INTEGRATION_ID) - ); - const dataframeSqlIntegrationCredentialsJson = JSON.stringify({ - url: 'deepnote+duckdb:///:memory:', - params: {}, - param_style: 'qmark' - }); - - envVars[dataframeSqlIntegrationEnvVarName] = dataframeSqlIntegrationCredentialsJson; - logger.debug(`SqlIntegrationEnvironmentVariablesProvider: Added env var for dataframe SQL integration`); - - // Get all configured integrations from storage + // Get all configured integrations from storage (includes DuckDB integration) const allIntegrations = await this.integrationStorage.getAll(); - if (allIntegrations.length === 0) { - logger.info(`SqlIntegrationEnvironmentVariablesProvider: No configured integrations found`); - return envVars; - } logger.trace( `SqlIntegrationEnvironmentVariablesProvider: Found ${allIntegrations.length} configured integrations` @@ -221,11 +209,6 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati } try { - // Skip internal DuckDB integration (already handled above) - if (config.id === DATAFRAME_SQL_INTEGRATION_ID) { - continue; - } - // Convert integration config to JSON and add as environment variable const envVarName = convertToEnvironmentVariableName(getSqlEnvVarName(config.id)); const credentialsJson = convertIntegrationConfigToJson(config); diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts index bc50a7626e..f011744a9a 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts @@ -10,13 +10,22 @@ import { PostgresIntegrationConfig, BigQueryIntegrationConfig, SnowflakeIntegrationConfig, - SnowflakeAuthMethods + SnowflakeAuthMethods, + DuckDBIntegrationConfig, + DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes'; const EXPECTED_DATAFRAME_ONLY_ENV_VARS = { SQL_DEEPNOTE_DATAFRAME_SQL: '{"url":"deepnote+duckdb:///:memory:","params":{},"param_style":"qmark"}' }; +// Helper to create the DuckDB integration config that's always included +const DUCKDB_INTEGRATION: DuckDBIntegrationConfig = { + id: DATAFRAME_SQL_INTEGRATION_ID, + name: 'Dataframe SQL (DuckDB)', + type: IntegrationType.DuckDB +}; + suite('SqlIntegrationEnvironmentVariablesProvider', () => { let provider: SqlIntegrationEnvironmentVariablesProvider; let integrationStorage: IntegrationStorage; @@ -41,7 +50,8 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Returns only dataframe integration when no integrations are configured', async () => { const uri = Uri.file('/test/notebook.deepnote'); - when(integrationStorage.getAll()).thenResolve([]); + // IntegrationStorage.getAll() always includes the DuckDB integration + when(integrationStorage.getAll()).thenResolve([DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); @@ -49,7 +59,8 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Returns environment variable for internal DuckDB integration', async () => { const uri = Uri.file('/test/notebook.deepnote'); - when(integrationStorage.getAll()).thenResolve([]); + // IntegrationStorage.getAll() always includes the DuckDB integration + when(integrationStorage.getAll()).thenResolve([DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -76,7 +87,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { ssl: true }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -100,7 +111,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { credentials: serviceAccountJson }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -139,7 +150,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { credentials: JSON.stringify({ type: 'service_account' }) }; - when(integrationStorage.getAll()).thenResolve([postgresConfig, bigqueryConfig]); + when(integrationStorage.getAll()).thenResolve([postgresConfig, bigqueryConfig, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -164,7 +175,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { ssl: false }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -196,7 +207,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { ssl: true }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -224,7 +235,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { ssl: false }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -261,7 +272,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'secret123' }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -290,7 +301,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'pass' }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -322,7 +333,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { privateKeyPassphrase: 'passphrase123' }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -356,7 +367,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { privateKey: privateKey }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -387,7 +398,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'p@ss:word!#$%' }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -413,7 +424,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'pass' }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); @@ -435,7 +446,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { authMethod: SnowflakeAuthMethods.OKTA }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); // Should return only dataframe integration when unsupported auth method is encountered const envVars = await provider.getEnvironmentVariables(uri); @@ -453,7 +464,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { authMethod: SnowflakeAuthMethods.AZURE_AD }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); @@ -470,7 +481,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { authMethod: SnowflakeAuthMethods.KEY_PAIR }; - when(integrationStorage.getAll()).thenResolve([config]); + when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); const envVars = await provider.getEnvironmentVariables(uri); assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); From dae0131774c73b0e384595ac26ea804c6a709cc6 Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 14:51:11 +0100 Subject: [PATCH 04/23] fix: scope integrations for env vars by the open project --- .../notebooks/deepnote/integrationStorage.ts | 19 +- ...IntegrationEnvironmentVariablesProvider.ts | 86 +++++++- ...nEnvironmentVariablesProvider.unit.test.ts | 200 +++++++++++++++--- 3 files changed, 248 insertions(+), 57 deletions(-) diff --git a/src/platform/notebooks/deepnote/integrationStorage.ts b/src/platform/notebooks/deepnote/integrationStorage.ts index 68b6130f52..aa8250ba8b 100644 --- a/src/platform/notebooks/deepnote/integrationStorage.ts +++ b/src/platform/notebooks/deepnote/integrationStorage.ts @@ -4,12 +4,7 @@ import { EventEmitter } from 'vscode'; import { IEncryptedStorage } from '../../common/application/types'; import { IAsyncDisposableRegistry } from '../../common/types'; import { logger } from '../../logging'; -import { - LegacyIntegrationConfig, - IntegrationType, - DuckDBIntegrationConfig, - DATAFRAME_SQL_INTEGRATION_ID -} from './integrationTypes'; +import { LegacyIntegrationConfig, IntegrationType } from './integrationTypes'; import { IIntegrationStorage } from './types'; const INTEGRATION_SERVICE_NAME = 'deepnote-integrations'; @@ -39,20 +34,10 @@ export class IntegrationStorage implements IIntegrationStorage { /** * Get all stored integration configurations. - * Always includes the internal DuckDB integration at the end of the list. */ async getAll(): Promise { await this.ensureCacheLoaded(); - const configs = Array.from(this.cache.values()); - - // Always add the internal DuckDB integration at the end - const duckdbConfig: DuckDBIntegrationConfig = { - id: DATAFRAME_SQL_INTEGRATION_ID, - name: 'Dataframe SQL (DuckDB)', - type: IntegrationType.DuckDB - }; - - return [...configs, duckdbConfig]; + return Array.from(this.cache.values()); } /** diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts index 781adfebf1..1d00e8873b 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts @@ -6,7 +6,14 @@ import { EnvironmentVariables } from '../../common/variables/types'; import { UnsupportedIntegrationError } from '../../errors/unsupportedIntegrationError'; import { logger } from '../../logging'; import { IIntegrationStorage, ISqlIntegrationEnvVarsProvider } from './types'; -import { LegacyIntegrationConfig, IntegrationType, SnowflakeAuthMethods } from './integrationTypes'; +import { + LegacyIntegrationConfig, + IntegrationType, + SnowflakeAuthMethods, + DuckDBIntegrationConfig, + DATAFRAME_SQL_INTEGRATION_ID +} from './integrationTypes'; +import { INotebookEditorProvider, IDeepnoteNotebookManager } from '../../../notebooks/types'; /** * Converts an integration ID to the environment variable name format expected by SQL blocks. @@ -163,6 +170,8 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati constructor( @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, + @inject(INotebookEditorProvider) private readonly notebookEditorProvider: INotebookEditorProvider, + @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, @inject(IDisposableRegistry) disposables: IDisposableRegistry ) { logger.info('SqlIntegrationEnvironmentVariablesProvider: Constructor called - provider is being instantiated'); @@ -179,8 +188,8 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati /** * Get environment variables for SQL integrations. - * Provides credentials for all configured integrations in the project. - * The internal DuckDB integration is always included via IntegrationStorage.getAll(). + * Provides credentials for all integrations in the Deepnote project. + * The internal DuckDB integration is always included. */ public async getEnvironmentVariables(resource: Resource, token?: CancellationToken): Promise { const envVars: EnvironmentVariables = {}; @@ -195,20 +204,79 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati logger.trace(`SqlIntegrationEnvironmentVariablesProvider: Getting env vars for resource`); - // Get all configured integrations from storage (includes DuckDB integration) - const allIntegrations = await this.integrationStorage.getAll(); + // Get the notebook document from the resource + const notebook = this.notebookEditorProvider.findAssociatedNotebookDocument(resource); + if (!notebook) { + logger.trace(`SqlIntegrationEnvironmentVariablesProvider: No notebook found for resource`); + return envVars; + } + + // Get the project ID from the notebook metadata + const projectId = notebook.metadata?.deepnoteProjectId as string | undefined; + if (!projectId) { + logger.trace(`SqlIntegrationEnvironmentVariablesProvider: No project ID found in notebook metadata`); + return envVars; + } + logger.trace(`SqlIntegrationEnvironmentVariablesProvider: Project ID: ${projectId}`); + + // Get the project from the notebook manager + const project = this.notebookManager.getOriginalProject(projectId); + if (!project) { + logger.trace(`SqlIntegrationEnvironmentVariablesProvider: No project found for ID: ${projectId}`); + return envVars; + } + + // Get the list of integrations from the project + const projectIntegrations = project.project.integrations || []; logger.trace( - `SqlIntegrationEnvironmentVariablesProvider: Found ${allIntegrations.length} configured integrations` + `SqlIntegrationEnvironmentVariablesProvider: Found ${projectIntegrations.length} integrations in project` ); - // Get credentials for each integration and add to environment variables - for (const config of allIntegrations) { + // Always add the internal DuckDB integration + const duckdbConfig: DuckDBIntegrationConfig = { + id: DATAFRAME_SQL_INTEGRATION_ID, + name: 'Dataframe SQL (DuckDB)', + type: IntegrationType.DuckDB + }; + + try { + const envVarName = convertToEnvironmentVariableName(getSqlEnvVarName(duckdbConfig.id)); + const credentialsJson = convertIntegrationConfigToJson(duckdbConfig); + envVars[envVarName] = credentialsJson; + logger.debug( + `SqlIntegrationEnvironmentVariablesProvider: Added env var ${envVarName} for DuckDB integration` + ); + } catch (error) { + logger.error( + `SqlIntegrationEnvironmentVariablesProvider: Failed to get credentials for DuckDB integration`, + error + ); + } + + // Get credentials for each project integration and add to environment variables + for (const projectIntegration of projectIntegrations) { if (token?.isCancellationRequested) { break; } + const integrationId = projectIntegration.id; + + // Skip the internal DuckDB integration (already added above) + if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) { + continue; + } + try { + // Get the integration configuration from storage + const config = await this.integrationStorage.getIntegrationConfig(integrationId); + if (!config) { + logger.debug( + `SqlIntegrationEnvironmentVariablesProvider: No configuration found for integration ${integrationId}, skipping` + ); + continue; + } + // Convert integration config to JSON and add as environment variable const envVarName = convertToEnvironmentVariableName(getSqlEnvVarName(config.id)); const credentialsJson = convertIntegrationConfigToJson(config); @@ -219,7 +287,7 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati ); } catch (error) { logger.error( - `SqlIntegrationEnvironmentVariablesProvider: Failed to get credentials for integration ${config.id}`, + `SqlIntegrationEnvironmentVariablesProvider: Failed to get credentials for integration ${integrationId}`, error ); } diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts index f011744a9a..5ccae17475 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts @@ -1,6 +1,6 @@ import { assert } from 'chai'; import { instance, mock, when } from 'ts-mockito'; -import { CancellationTokenSource, EventEmitter, Uri } from 'vscode'; +import { CancellationTokenSource, EventEmitter, NotebookDocument, Uri } from 'vscode'; import { IDisposableRegistry } from '../../common/types'; import { IntegrationStorage } from './integrationStorage'; @@ -10,33 +10,56 @@ import { PostgresIntegrationConfig, BigQueryIntegrationConfig, SnowflakeIntegrationConfig, - SnowflakeAuthMethods, - DuckDBIntegrationConfig, - DATAFRAME_SQL_INTEGRATION_ID + SnowflakeAuthMethods } from './integrationTypes'; +import { INotebookEditorProvider, IDeepnoteNotebookManager } from '../../../notebooks/types'; const EXPECTED_DATAFRAME_ONLY_ENV_VARS = { SQL_DEEPNOTE_DATAFRAME_SQL: '{"url":"deepnote+duckdb:///:memory:","params":{},"param_style":"qmark"}' }; -// Helper to create the DuckDB integration config that's always included -const DUCKDB_INTEGRATION: DuckDBIntegrationConfig = { - id: DATAFRAME_SQL_INTEGRATION_ID, - name: 'Dataframe SQL (DuckDB)', - type: IntegrationType.DuckDB -}; +// Helper to create a mock project with integrations +function createMockProject(integrations: Array<{ id: string; name: string; type: string }>) { + return { + project: { + id: 'test-project-id', + name: 'Test Project', + integrations + } + }; +} + +// Helper to create a mock notebook document +function createMockNotebook(projectId: string): NotebookDocument { + return { + metadata: { + deepnoteProjectId: projectId, + deepnoteNotebookId: 'test-notebook-id' + } + } as unknown as NotebookDocument; +} suite('SqlIntegrationEnvironmentVariablesProvider', () => { let provider: SqlIntegrationEnvironmentVariablesProvider; let integrationStorage: IntegrationStorage; + let notebookEditorProvider: INotebookEditorProvider; + let notebookManager: IDeepnoteNotebookManager; let disposables: IDisposableRegistry; setup(() => { disposables = []; integrationStorage = mock(IntegrationStorage); + notebookEditorProvider = mock(); + notebookManager = mock(); + when(integrationStorage.onDidChangeIntegrations).thenReturn(new EventEmitter().event); - provider = new SqlIntegrationEnvironmentVariablesProvider(instance(integrationStorage), disposables); + provider = new SqlIntegrationEnvironmentVariablesProvider( + instance(integrationStorage), + instance(notebookEditorProvider), + instance(notebookManager), + disposables + ); }); teardown(() => { @@ -48,10 +71,40 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { assert.deepStrictEqual(envVars, {}); }); + test('Returns empty object when no notebook is found', async () => { + const uri = Uri.file('/test/notebook.deepnote'); + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(undefined); + + const envVars = await provider.getEnvironmentVariables(uri); + assert.deepStrictEqual(envVars, {}); + }); + + test('Returns empty object when notebook has no project ID', async () => { + const uri = Uri.file('/test/notebook.deepnote'); + const notebook = { metadata: {} } as NotebookDocument; + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + + const envVars = await provider.getEnvironmentVariables(uri); + assert.deepStrictEqual(envVars, {}); + }); + + test('Returns empty object when project is not found', async () => { + const uri = Uri.file('/test/notebook.deepnote'); + const notebook = createMockNotebook('test-project-id'); + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(undefined); + + const envVars = await provider.getEnvironmentVariables(uri); + assert.deepStrictEqual(envVars, {}); + }); + test('Returns only dataframe integration when no integrations are configured', async () => { const uri = Uri.file('/test/notebook.deepnote'); - // IntegrationStorage.getAll() always includes the DuckDB integration - when(integrationStorage.getAll()).thenResolve([DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); const envVars = await provider.getEnvironmentVariables(uri); assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); @@ -59,8 +112,11 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Returns environment variable for internal DuckDB integration', async () => { const uri = Uri.file('/test/notebook.deepnote'); - // IntegrationStorage.getAll() always includes the DuckDB integration - when(integrationStorage.getAll()).thenResolve([DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); const envVars = await provider.getEnvironmentVariables(uri); @@ -87,7 +143,12 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { ssl: true }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([{ id: integrationId, name: 'My Postgres DB', type: 'pgsql' }]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); @@ -111,7 +172,12 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { credentials: serviceAccountJson }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([{ id: integrationId, name: 'My BigQuery', type: 'big-query' }]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); @@ -150,7 +216,16 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { credentials: JSON.stringify({ type: 'service_account' }) }; - when(integrationStorage.getAll()).thenResolve([postgresConfig, bigqueryConfig, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([ + { id: postgresId, name: 'My Postgres DB', type: 'pgsql' }, + { id: bigqueryId, name: 'My BigQuery', type: 'big-query' } + ]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(postgresId)).thenResolve(postgresConfig); + when(integrationStorage.getIntegrationConfig(bigqueryId)).thenResolve(bigqueryConfig); const envVars = await provider.getEnvironmentVariables(uri); @@ -175,7 +250,12 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { ssl: false }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([{ id: integrationId, name: 'Special Chars DB', type: 'pgsql' }]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); @@ -207,7 +287,12 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { ssl: true }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([{ id: integrationId, name: 'Production Database', type: 'pgsql' }]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); @@ -235,7 +320,12 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { ssl: false }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([{ id: integrationId, name: 'Test DB', type: 'pgsql' }]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); @@ -248,7 +338,6 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Honors CancellationToken (returns empty when cancelled early)', async () => { const uri = Uri.file('/test/notebook.deepnote'); - when(integrationStorage.getAll()).thenCall(() => new Promise((resolve) => setTimeout(() => resolve([]), 50))); const cts = new CancellationTokenSource(); cts.cancel(); const envVars = await provider.getEnvironmentVariables(uri, cts.token); @@ -272,7 +361,12 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'secret123' }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([{ id: integrationId, name: 'My Snowflake', type: 'snowflake' }]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); @@ -301,7 +395,12 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'pass' }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([{ id: integrationId, name: 'Legacy Snowflake', type: 'snowflake' }]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); @@ -333,7 +432,12 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { privateKeyPassphrase: 'passphrase123' }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([{ id: integrationId, name: 'Snowflake KeyPair', type: 'snowflake' }]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); @@ -367,7 +471,14 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { privateKey: privateKey }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([ + { id: integrationId, name: 'Snowflake KeyPair No Pass', type: 'snowflake' } + ]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); @@ -398,7 +509,12 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'p@ss:word!#$%' }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([{ id: integrationId, name: 'Snowflake Special', type: 'snowflake' }]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); @@ -424,7 +540,12 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'pass' }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([{ id: integrationId, name: 'Snowflake Minimal', type: 'snowflake' }]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); @@ -446,7 +567,12 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { authMethod: SnowflakeAuthMethods.OKTA }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([{ id: integrationId, name: 'Snowflake OKTA', type: 'snowflake' }]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); // Should return only dataframe integration when unsupported auth method is encountered const envVars = await provider.getEnvironmentVariables(uri); @@ -464,7 +590,12 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { authMethod: SnowflakeAuthMethods.AZURE_AD }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([{ id: integrationId, name: 'Snowflake Azure', type: 'snowflake' }]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); @@ -481,7 +612,14 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { authMethod: SnowflakeAuthMethods.KEY_PAIR }; - when(integrationStorage.getAll()).thenResolve([config, DUCKDB_INTEGRATION]); + const notebook = createMockNotebook('test-project-id'); + const project = createMockProject([ + { id: integrationId, name: 'Snowflake KeyPair User', type: 'snowflake' } + ]); + + when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); + when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); + when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); const envVars = await provider.getEnvironmentVariables(uri); assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); From 558a6f03f37c5418d859288a5684073f6add799e Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 14:53:50 +0100 Subject: [PATCH 05/23] consolidate duckdb and other integration processing --- ...IntegrationEnvironmentVariablesProvider.ts | 27 +++++-------------- 1 file changed, 6 insertions(+), 21 deletions(-) diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts index 1d00e8873b..6eb4cff12d 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts @@ -228,7 +228,7 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati } // Get the list of integrations from the project - const projectIntegrations = project.project.integrations || []; + const projectIntegrations = project.project.integrations?.slice() ?? []; logger.trace( `SqlIntegrationEnvironmentVariablesProvider: Found ${projectIntegrations.length} integrations in project` ); @@ -239,20 +239,7 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati name: 'Dataframe SQL (DuckDB)', type: IntegrationType.DuckDB }; - - try { - const envVarName = convertToEnvironmentVariableName(getSqlEnvVarName(duckdbConfig.id)); - const credentialsJson = convertIntegrationConfigToJson(duckdbConfig); - envVars[envVarName] = credentialsJson; - logger.debug( - `SqlIntegrationEnvironmentVariablesProvider: Added env var ${envVarName} for DuckDB integration` - ); - } catch (error) { - logger.error( - `SqlIntegrationEnvironmentVariablesProvider: Failed to get credentials for DuckDB integration`, - error - ); - } + projectIntegrations.push(duckdbConfig); // Get credentials for each project integration and add to environment variables for (const projectIntegration of projectIntegrations) { @@ -262,14 +249,12 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati const integrationId = projectIntegration.id; - // Skip the internal DuckDB integration (already added above) - if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) { - continue; - } - try { // Get the integration configuration from storage - const config = await this.integrationStorage.getIntegrationConfig(integrationId); + const config = + integrationId === DATAFRAME_SQL_INTEGRATION_ID + ? duckdbConfig + : await this.integrationStorage.getIntegrationConfig(integrationId); if (!config) { logger.debug( `SqlIntegrationEnvironmentVariablesProvider: No configuration found for integration ${integrationId}, skipping` From 594624de51a88de75b898357c6b59f964f78088d Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 15:11:50 +0100 Subject: [PATCH 06/23] add `@deepnote/database-integrations` package --- package-lock.json | 47 +++++++++++++++++++++++++++++++++++++++++++++++ package.json | 1 + 2 files changed, 48 insertions(+) diff --git a/package-lock.json b/package-lock.json index 612c821116..eae592ca46 100644 --- a/package-lock.json +++ b/package-lock.json @@ -13,6 +13,7 @@ "@c4312/evt": "^0.1.1", "@deepnote/blocks": "^1.3.5", "@deepnote/convert": "^1.2.0", + "@deepnote/database-integrations": "^1.1.0", "@enonic/fnv-plus": "^1.3.0", "@jupyter-widgets/base": "^6.0.8", "@jupyter-widgets/controls": "^5.0.9", @@ -1417,6 +1418,25 @@ "url": "https://github.com/chalk/chalk?sponsor=1" } }, + "node_modules/@deepnote/database-integrations": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/@deepnote/database-integrations/-/database-integrations-1.1.0.tgz", + "integrity": "sha512-LD21aXg9YgMroWPhSxRZ7bZ6qQOfypsEv7v+I1jcdDlexLdQVZtG0VoZ6TUD4jRWU9Lf7QpjzgFlQW9OdbR8BQ==", + "license": "Apache-2.0", + "dependencies": { + "build-url-ts": "^6.1.7", + "zod": "3.25.76" + } + }, + "node_modules/@deepnote/database-integrations/node_modules/zod": { + "version": "3.25.76", + "resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz", + "integrity": "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==", + "license": "MIT", + "funding": { + "url": "https://github.com/sponsors/colinhacks" + } + }, "node_modules/@enonic/fnv-plus": { "version": "1.3.0", "resolved": "https://registry.npmjs.org/@enonic/fnv-plus/-/fnv-plus-1.3.0.tgz", @@ -6108,6 +6128,12 @@ "node": ">=6.14.2" } }, + "node_modules/build-url-ts": { + "version": "6.2.0", + "resolved": "https://registry.npmjs.org/build-url-ts/-/build-url-ts-6.2.0.tgz", + "integrity": "sha512-bwF2Cytry4/QWshF9ukPtinIGgBUuqbjVVVv7mpk6gmUCi0tkTorE98/fmu76t+VF10kKDIL9S+bQt1/7z8KUA==", + "license": "MIT" + }, "node_modules/builtin-modules": { "version": "3.3.0", "resolved": "https://registry.npmjs.org/builtin-modules/-/builtin-modules-3.3.0.tgz", @@ -21610,6 +21636,22 @@ } } }, + "@deepnote/database-integrations": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/@deepnote/database-integrations/-/database-integrations-1.1.0.tgz", + "integrity": "sha512-LD21aXg9YgMroWPhSxRZ7bZ6qQOfypsEv7v+I1jcdDlexLdQVZtG0VoZ6TUD4jRWU9Lf7QpjzgFlQW9OdbR8BQ==", + "requires": { + "build-url-ts": "^6.1.7", + "zod": "3.25.76" + }, + "dependencies": { + "zod": { + "version": "3.25.76", + "resolved": "https://registry.npmjs.org/zod/-/zod-3.25.76.tgz", + "integrity": "sha512-gzUt/qt81nXsFGKIFcC3YnfEAx5NkunCfnDlvuBSSFS02bcXu4Lmea0AFIUwbLWxWPx3d9p8S5QoaujKcNQxcQ==" + } + } + }, "@enonic/fnv-plus": { "version": "1.3.0", "resolved": "https://registry.npmjs.org/@enonic/fnv-plus/-/fnv-plus-1.3.0.tgz", @@ -25024,6 +25066,11 @@ "node-gyp-build": "^4.3.0" } }, + "build-url-ts": { + "version": "6.2.0", + "resolved": "https://registry.npmjs.org/build-url-ts/-/build-url-ts-6.2.0.tgz", + "integrity": "sha512-bwF2Cytry4/QWshF9ukPtinIGgBUuqbjVVVv7mpk6gmUCi0tkTorE98/fmu76t+VF10kKDIL9S+bQt1/7z8KUA==" + }, "builtin-modules": { "version": "3.3.0", "resolved": "https://registry.npmjs.org/builtin-modules/-/builtin-modules-3.3.0.tgz", diff --git a/package.json b/package.json index 626012ddab..43571694a7 100644 --- a/package.json +++ b/package.json @@ -2455,6 +2455,7 @@ "@c4312/evt": "^0.1.1", "@deepnote/blocks": "^1.3.5", "@deepnote/convert": "^1.2.0", + "@deepnote/database-integrations": "^1.1.0", "@enonic/fnv-plus": "^1.3.0", "@jupyter-widgets/base": "^6.0.8", "@jupyter-widgets/controls": "^5.0.9", From ebbbf01d564f7564a972f0acd4a386cd63740e3b Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 15:17:22 +0100 Subject: [PATCH 07/23] fix import restriction violations in integration env var logic --- src/notebooks/serviceRegistry.node.ts | 8 +++++++ ...IntegrationEnvironmentVariablesProvider.ts | 13 +++++++---- ...nEnvironmentVariablesProvider.unit.test.ts | 10 ++++---- src/platform/notebooks/deepnote/types.ts | 23 ++++++++++++++++++- 4 files changed, 44 insertions(+), 10 deletions(-) diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index bf9ba6186b..c316bc97ae 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -53,6 +53,10 @@ import { IIntegrationStorage, IIntegrationWebviewProvider } from './deepnote/integrations/types'; +import { + IPlatformNotebookEditorProvider, + IPlatformDeepnoteNotebookManager +} from '../platform/notebooks/deepnote/types'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; import { IDeepnoteToolkitInstaller, @@ -78,6 +82,8 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea serviceManager.addSingleton(INotebookCommandHandler, NotebookCommandListener); serviceManager.addBinding(INotebookCommandHandler, IExtensionSyncActivationService); serviceManager.addSingleton(INotebookEditorProvider, NotebookEditorProvider); + // Bind the platform-layer interface to the same implementation + serviceManager.addBinding(INotebookEditorProvider, IPlatformNotebookEditorProvider); serviceManager.addSingleton( IExtensionSyncActivationService, RemoteKernelControllerWatcher @@ -148,6 +154,8 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea DeepnoteNotebookCommandListener ); serviceManager.addSingleton(IDeepnoteNotebookManager, DeepnoteNotebookManager); + // Bind the platform-layer interface to the same implementation + serviceManager.addBinding(IDeepnoteNotebookManager, IPlatformDeepnoteNotebookManager); serviceManager.addSingleton(IIntegrationStorage, IntegrationStorage); serviceManager.addSingleton(IIntegrationDetector, IntegrationDetector); serviceManager.addSingleton(IIntegrationWebviewProvider, IntegrationWebviewProvider); diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts index 6eb4cff12d..6cbd5b69b5 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts @@ -5,7 +5,12 @@ import { IDisposableRegistry, Resource } from '../../common/types'; import { EnvironmentVariables } from '../../common/variables/types'; import { UnsupportedIntegrationError } from '../../errors/unsupportedIntegrationError'; import { logger } from '../../logging'; -import { IIntegrationStorage, ISqlIntegrationEnvVarsProvider } from './types'; +import { + IIntegrationStorage, + ISqlIntegrationEnvVarsProvider, + IPlatformNotebookEditorProvider, + IPlatformDeepnoteNotebookManager +} from './types'; import { LegacyIntegrationConfig, IntegrationType, @@ -13,7 +18,6 @@ import { DuckDBIntegrationConfig, DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes'; -import { INotebookEditorProvider, IDeepnoteNotebookManager } from '../../../notebooks/types'; /** * Converts an integration ID to the environment variable name format expected by SQL blocks. @@ -170,8 +174,9 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati constructor( @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, - @inject(INotebookEditorProvider) private readonly notebookEditorProvider: INotebookEditorProvider, - @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, + @inject(IPlatformNotebookEditorProvider) + private readonly notebookEditorProvider: IPlatformNotebookEditorProvider, + @inject(IPlatformDeepnoteNotebookManager) private readonly notebookManager: IPlatformDeepnoteNotebookManager, @inject(IDisposableRegistry) disposables: IDisposableRegistry ) { logger.info('SqlIntegrationEnvironmentVariablesProvider: Constructor called - provider is being instantiated'); diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts index 5ccae17475..8b7a8cc698 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts @@ -12,7 +12,7 @@ import { SnowflakeIntegrationConfig, SnowflakeAuthMethods } from './integrationTypes'; -import { INotebookEditorProvider, IDeepnoteNotebookManager } from '../../../notebooks/types'; +import { IPlatformNotebookEditorProvider, IPlatformDeepnoteNotebookManager } from './types'; const EXPECTED_DATAFRAME_ONLY_ENV_VARS = { SQL_DEEPNOTE_DATAFRAME_SQL: '{"url":"deepnote+duckdb:///:memory:","params":{},"param_style":"qmark"}' @@ -42,15 +42,15 @@ function createMockNotebook(projectId: string): NotebookDocument { suite('SqlIntegrationEnvironmentVariablesProvider', () => { let provider: SqlIntegrationEnvironmentVariablesProvider; let integrationStorage: IntegrationStorage; - let notebookEditorProvider: INotebookEditorProvider; - let notebookManager: IDeepnoteNotebookManager; + let notebookEditorProvider: IPlatformNotebookEditorProvider; + let notebookManager: IPlatformDeepnoteNotebookManager; let disposables: IDisposableRegistry; setup(() => { disposables = []; integrationStorage = mock(IntegrationStorage); - notebookEditorProvider = mock(); - notebookManager = mock(); + notebookEditorProvider = mock(); + notebookManager = mock(); when(integrationStorage.onDidChangeIntegrations).thenReturn(new EventEmitter().event); diff --git a/src/platform/notebooks/deepnote/types.ts b/src/platform/notebooks/deepnote/types.ts index 98de30ccde..9147496835 100644 --- a/src/platform/notebooks/deepnote/types.ts +++ b/src/platform/notebooks/deepnote/types.ts @@ -1,7 +1,8 @@ -import { CancellationToken, Event } from 'vscode'; +import { CancellationToken, Event, NotebookDocument, Uri } from 'vscode'; import { IDisposable, Resource } from '../../common/types'; import { EnvironmentVariables } from '../../common/variables/types'; import { LegacyIntegrationConfig } from './integrationTypes'; +import { DeepnoteProject } from '../../deepnote/deepnoteTypes'; /** * Settings for select input blocks @@ -73,3 +74,23 @@ export interface ISqlIntegrationEnvVarsProvider { */ getEnvironmentVariables(resource: Resource, token?: CancellationToken): Promise; } + +/** + * Platform-layer interface for accessing notebook documents. + * This is a subset of the full INotebookEditorProvider interface from the notebooks layer. + * The implementation in the notebooks layer should be bound to this symbol as well. + */ +export const IPlatformNotebookEditorProvider = Symbol('IPlatformNotebookEditorProvider'); +export interface IPlatformNotebookEditorProvider { + findAssociatedNotebookDocument(uri: Uri): NotebookDocument | undefined; +} + +/** + * Platform-layer interface for accessing Deepnote project data. + * This is a subset of the full IDeepnoteNotebookManager interface from the notebooks layer. + * The implementation in the notebooks layer should be bound to this symbol as well. + */ +export const IPlatformDeepnoteNotebookManager = Symbol('IPlatformDeepnoteNotebookManager'); +export interface IPlatformDeepnoteNotebookManager { + getOriginalProject(projectId: string): DeepnoteProject | undefined; +} From ff380a16a587603d824064ce45c9a03285421809 Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 15:39:40 +0100 Subject: [PATCH 08/23] rename IntegrationType to LegacyIntegrationType --- .../integrations/integrationDetector.ts | 4 +-- .../integrations/integrationManager.ts | 11 +++--- .../integrations/integrationWebview.ts | 8 ++--- .../deepnote/sqlCellStatusBarProvider.ts | 14 ++++---- .../sqlCellStatusBarProvider.unit.test.ts | 9 +++-- .../notebooks/deepnote/integrationStorage.ts | 4 +-- .../notebooks/deepnote/integrationTypes.ts | 35 ++++++++++--------- ...IntegrationEnvironmentVariablesProvider.ts | 2 +- ...nEnvironmentVariablesProvider.unit.test.ts | 34 +++++++++--------- 9 files changed, 63 insertions(+), 58 deletions(-) diff --git a/src/notebooks/deepnote/integrations/integrationDetector.ts b/src/notebooks/deepnote/integrations/integrationDetector.ts index 55e627b734..558439db45 100644 --- a/src/notebooks/deepnote/integrations/integrationDetector.ts +++ b/src/notebooks/deepnote/integrations/integrationDetector.ts @@ -4,7 +4,7 @@ import { logger } from '../../../platform/logging'; import { IDeepnoteNotebookManager } from '../../types'; import { DATAFRAME_SQL_INTEGRATION_ID, - DEEPNOTE_TO_INTEGRATION_TYPE, + DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE, IntegrationStatus, IntegrationWithStatus, RawIntegrationType @@ -54,7 +54,7 @@ export class IntegrationDetector implements IIntegrationDetector { logger.debug(`IntegrationDetector: Found integration: ${integrationId} (${projectIntegration.type})`); // Map the Deepnote integration type to our IntegrationType - const integrationType = DEEPNOTE_TO_INTEGRATION_TYPE[projectIntegration.type as RawIntegrationType]; + const integrationType = DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE[projectIntegration.type as RawIntegrationType]; // Skip unknown integration types if (!integrationType) { diff --git a/src/notebooks/deepnote/integrations/integrationManager.ts b/src/notebooks/deepnote/integrations/integrationManager.ts index e9dd729c47..00ac2dc4f1 100644 --- a/src/notebooks/deepnote/integrations/integrationManager.ts +++ b/src/notebooks/deepnote/integrations/integrationManager.ts @@ -6,9 +6,9 @@ import { Commands } from '../../../platform/common/constants'; import { logger } from '../../../platform/logging'; import { IIntegrationDetector, IIntegrationManager, IIntegrationStorage, IIntegrationWebviewProvider } from './types'; import { - DEEPNOTE_TO_INTEGRATION_TYPE, + DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE, IntegrationStatus, - IntegrationType, + LegacyIntegrationType, IntegrationWithStatus, RawIntegrationType } from '../../../platform/notebooks/deepnote/integrationTypes'; @@ -164,15 +164,16 @@ export class IntegrationManager implements IIntegrationManager { const projectIntegration = project?.project.integrations?.find((i) => i.id === selectedIntegrationId); let integrationName: string | undefined; - let integrationType: IntegrationType | undefined; + let integrationType: LegacyIntegrationType | undefined; if (projectIntegration) { integrationName = projectIntegration.name; // Validate that projectIntegration.type exists in the mapping before lookup - if (projectIntegration.type in DEEPNOTE_TO_INTEGRATION_TYPE) { + if (projectIntegration.type in DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE) { // Map the Deepnote integration type to our IntegrationType - integrationType = DEEPNOTE_TO_INTEGRATION_TYPE[projectIntegration.type as RawIntegrationType]; + integrationType = + DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE[projectIntegration.type as RawIntegrationType]; } else { logger.warn( `IntegrationManager: Unknown integration type '${projectIntegration.type}' for integration ID '${selectedIntegrationId}' in project '${projectId}'. Integration type will be undefined.` diff --git a/src/notebooks/deepnote/integrations/integrationWebview.ts b/src/notebooks/deepnote/integrations/integrationWebview.ts index 3b740a4cc4..018012675c 100644 --- a/src/notebooks/deepnote/integrations/integrationWebview.ts +++ b/src/notebooks/deepnote/integrations/integrationWebview.ts @@ -8,8 +8,8 @@ import { LocalizedMessages, SharedMessages } from '../../../messageTypes'; import { IDeepnoteNotebookManager, ProjectIntegration } from '../../types'; import { IIntegrationStorage, IIntegrationWebviewProvider } from './types'; import { - INTEGRATION_TYPE_TO_DEEPNOTE, - IntegrationType, + LEGACY_INTEGRATION_TYPE_TO_DEEPNOTE, + LegacyIntegrationType, LegacyIntegrationConfig, IntegrationStatus, IntegrationWithStatus, @@ -351,13 +351,13 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { } // Skip DuckDB integration (internal, not a real Deepnote integration) - if (type === IntegrationType.DuckDB) { + if (type === LegacyIntegrationType.DuckDB) { logger.trace(`IntegrationWebviewProvider: Skipping internal DuckDB integration ${id}`); return null; } // Map to Deepnote integration type - const deepnoteType: RawIntegrationType | undefined = INTEGRATION_TYPE_TO_DEEPNOTE[type]; + const deepnoteType: RawIntegrationType | undefined = LEGACY_INTEGRATION_TYPE_TO_DEEPNOTE[type]; if (!deepnoteType) { logger.warn(`IntegrationWebviewProvider: Cannot map type ${type} for integration ${id}, skipping`); return null; diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts index b0ba8897ce..adf79bb358 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts @@ -24,8 +24,8 @@ import { IIntegrationStorage } from './integrations/types'; import { Commands } from '../../platform/common/constants'; import { DATAFRAME_SQL_INTEGRATION_ID, - DEEPNOTE_TO_INTEGRATION_TYPE, - IntegrationType, + DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE, + LegacyIntegrationType, RawIntegrationType } from '../../platform/notebooks/deepnote/integrationTypes'; import { IDeepnoteNotebookManager } from '../types'; @@ -347,7 +347,7 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid continue; } - const integrationType = DEEPNOTE_TO_INTEGRATION_TYPE[projectIntegration.type as RawIntegrationType]; + const integrationType = DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE[projectIntegration.type as RawIntegrationType]; const typeLabel = integrationType ? this.getIntegrationTypeLabel(integrationType) : projectIntegration.type; const item: LocalQuickPickItem = { @@ -431,13 +431,13 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid this._onDidChangeCellStatusBarItems.fire(); } - private getIntegrationTypeLabel(type: IntegrationType): string { + private getIntegrationTypeLabel(type: LegacyIntegrationType): string { switch (type) { - case IntegrationType.Postgres: + case LegacyIntegrationType.Postgres: return l10n.t('PostgreSQL'); - case IntegrationType.BigQuery: + case LegacyIntegrationType.BigQuery: return l10n.t('BigQuery'); - case IntegrationType.Snowflake: + case LegacyIntegrationType.Snowflake: return l10n.t('Snowflake'); default: return String(type); diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts index 51289fed34..1487339620 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts @@ -14,7 +14,10 @@ import { import { IDisposableRegistry } from '../../platform/common/types'; import { IIntegrationStorage } from './integrations/types'; import { SqlCellStatusBarProvider } from './sqlCellStatusBarProvider'; -import { DATAFRAME_SQL_INTEGRATION_ID, IntegrationType } from '../../platform/notebooks/deepnote/integrationTypes'; +import { + DATAFRAME_SQL_INTEGRATION_ID, + LegacyIntegrationType +} from '../../platform/notebooks/deepnote/integrationTypes'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; import { createEventHandler } from '../../test/common'; import { Commands } from '../../platform/common/constants'; @@ -134,7 +137,7 @@ suite('SqlCellStatusBarProvider', () => { when(integrationStorage.getProjectIntegrationConfig(anything(), anything())).thenResolve({ id: integrationId, name: 'My Postgres DB', - type: IntegrationType.Postgres, + type: LegacyIntegrationType.Postgres, host: 'localhost', port: 5432, database: 'test', @@ -282,7 +285,7 @@ suite('SqlCellStatusBarProvider', () => { when(integrationStorage.getProjectIntegrationConfig(anything(), anything())).thenResolve({ id: integrationId, name: 'My Postgres DB', - type: IntegrationType.Postgres, + type: LegacyIntegrationType.Postgres, host: 'localhost', port: 5432, database: 'test', diff --git a/src/platform/notebooks/deepnote/integrationStorage.ts b/src/platform/notebooks/deepnote/integrationStorage.ts index aa8250ba8b..afd53c81c9 100644 --- a/src/platform/notebooks/deepnote/integrationStorage.ts +++ b/src/platform/notebooks/deepnote/integrationStorage.ts @@ -4,7 +4,7 @@ import { EventEmitter } from 'vscode'; import { IEncryptedStorage } from '../../common/application/types'; import { IAsyncDisposableRegistry } from '../../common/types'; import { logger } from '../../logging'; -import { LegacyIntegrationConfig, IntegrationType } from './integrationTypes'; +import { LegacyIntegrationConfig, LegacyIntegrationType } from './integrationTypes'; import { IIntegrationStorage } from './types'; const INTEGRATION_SERVICE_NAME = 'deepnote-integrations'; @@ -63,7 +63,7 @@ export class IntegrationStorage implements IIntegrationStorage { /** * Get all integrations of a specific type */ - async getByType(type: IntegrationType): Promise { + async getByType(type: LegacyIntegrationType): Promise { await this.ensureCacheLoaded(); return Array.from(this.cache.values()).filter((config) => config.type === type); } diff --git a/src/platform/notebooks/deepnote/integrationTypes.ts b/src/platform/notebooks/deepnote/integrationTypes.ts index 14aa2afbff..4d2ac4b853 100644 --- a/src/platform/notebooks/deepnote/integrationTypes.ts +++ b/src/platform/notebooks/deepnote/integrationTypes.ts @@ -7,7 +7,7 @@ export const DATAFRAME_SQL_INTEGRATION_ID = 'deepnote-dataframe-sql'; /** * Supported integration types */ -export enum IntegrationType { +export enum LegacyIntegrationType { Postgres = 'postgres', BigQuery = 'bigquery', Snowflake = 'snowflake', @@ -18,21 +18,22 @@ export enum IntegrationType { * Map our IntegrationType enum to Deepnote integration type strings * Note: DuckDB is not included as it's an internal integration that doesn't exist in Deepnote */ -export const INTEGRATION_TYPE_TO_DEEPNOTE = { - [IntegrationType.Postgres]: 'pgsql', - [IntegrationType.BigQuery]: 'big-query', - [IntegrationType.Snowflake]: 'snowflake' -} as const satisfies { [type in Exclude]: string }; +export const LEGACY_INTEGRATION_TYPE_TO_DEEPNOTE = { + [LegacyIntegrationType.Postgres]: 'pgsql', + [LegacyIntegrationType.BigQuery]: 'big-query', + [LegacyIntegrationType.Snowflake]: 'snowflake' +} as const satisfies { [type in Exclude]: string }; -export type RawIntegrationType = (typeof INTEGRATION_TYPE_TO_DEEPNOTE)[keyof typeof INTEGRATION_TYPE_TO_DEEPNOTE]; +export type RawIntegrationType = + (typeof LEGACY_INTEGRATION_TYPE_TO_DEEPNOTE)[keyof typeof LEGACY_INTEGRATION_TYPE_TO_DEEPNOTE]; /** * Map Deepnote integration type strings to our IntegrationType enum */ -export const DEEPNOTE_TO_INTEGRATION_TYPE: Record = { - pgsql: IntegrationType.Postgres, - 'big-query': IntegrationType.BigQuery, - snowflake: IntegrationType.Snowflake +export const DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE: Record = { + pgsql: LegacyIntegrationType.Postgres, + 'big-query': LegacyIntegrationType.BigQuery, + snowflake: LegacyIntegrationType.Snowflake }; /** @@ -41,14 +42,14 @@ export const DEEPNOTE_TO_INTEGRATION_TYPE: Record { const config: PostgresIntegrationConfig = { id: integrationId, name: 'My Postgres DB', - type: IntegrationType.Postgres, + type: LegacyIntegrationType.Postgres, host: 'localhost', port: 5432, database: 'mydb', @@ -167,7 +167,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: BigQueryIntegrationConfig = { id: integrationId, name: 'My BigQuery', - type: IntegrationType.BigQuery, + type: LegacyIntegrationType.BigQuery, projectId: 'my-project', credentials: serviceAccountJson }; @@ -200,7 +200,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const postgresConfig: PostgresIntegrationConfig = { id: postgresId, name: 'My Postgres DB', - type: IntegrationType.Postgres, + type: LegacyIntegrationType.Postgres, host: 'localhost', port: 5432, database: 'mydb', @@ -211,7 +211,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const bigqueryConfig: BigQueryIntegrationConfig = { id: bigqueryId, name: 'My BigQuery', - type: IntegrationType.BigQuery, + type: LegacyIntegrationType.BigQuery, projectId: 'my-project', credentials: JSON.stringify({ type: 'service_account' }) }; @@ -241,7 +241,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: PostgresIntegrationConfig = { id: integrationId, name: 'Special Chars DB', - type: IntegrationType.Postgres, + type: LegacyIntegrationType.Postgres, host: 'db.example.com', port: 5432, database: 'my@db:name', @@ -278,7 +278,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: PostgresIntegrationConfig = { id: integrationId, name: 'Production Database', - type: IntegrationType.Postgres, + type: LegacyIntegrationType.Postgres, host: 'prod.example.com', port: 5432, database: 'proddb', @@ -311,7 +311,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: PostgresIntegrationConfig = { id: integrationId, name: 'Test DB', - type: IntegrationType.Postgres, + type: LegacyIntegrationType.Postgres, host: 'localhost', port: 5432, database: 'testdb', @@ -351,7 +351,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: SnowflakeIntegrationConfig = { id: integrationId, name: 'My Snowflake', - type: IntegrationType.Snowflake, + type: LegacyIntegrationType.Snowflake, account: 'myorg-myaccount', warehouse: 'COMPUTE_WH', database: 'MYDB', @@ -386,7 +386,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: SnowflakeIntegrationConfig = { id: integrationId, name: 'Legacy Snowflake', - type: IntegrationType.Snowflake, + type: LegacyIntegrationType.Snowflake, account: 'legacy-account', warehouse: 'WH', database: 'DB', @@ -421,7 +421,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: SnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake KeyPair', - type: IntegrationType.Snowflake, + type: LegacyIntegrationType.Snowflake, account: 'keypair-account', warehouse: 'ETL_WH', database: 'PROD_DB', @@ -462,7 +462,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: SnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake KeyPair No Pass', - type: IntegrationType.Snowflake, + type: LegacyIntegrationType.Snowflake, account: 'account123', warehouse: 'WH', database: 'DB', @@ -499,7 +499,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: SnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake Special', - type: IntegrationType.Snowflake, + type: LegacyIntegrationType.Snowflake, account: 'my-org.account', warehouse: 'WH@2024', database: 'DB:TEST', @@ -533,7 +533,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: SnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake Minimal', - type: IntegrationType.Snowflake, + type: LegacyIntegrationType.Snowflake, account: 'minimal-account', authMethod: SnowflakeAuthMethods.PASSWORD, username: 'user', @@ -562,7 +562,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: SnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake OKTA', - type: IntegrationType.Snowflake, + type: LegacyIntegrationType.Snowflake, account: 'okta-account', authMethod: SnowflakeAuthMethods.OKTA }; @@ -585,7 +585,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: SnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake Azure', - type: IntegrationType.Snowflake, + type: LegacyIntegrationType.Snowflake, account: 'azure-account', authMethod: SnowflakeAuthMethods.AZURE_AD }; @@ -607,7 +607,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const config: SnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake KeyPair User', - type: IntegrationType.Snowflake, + type: LegacyIntegrationType.Snowflake, account: 'keypair-user-account', authMethod: SnowflakeAuthMethods.KEY_PAIR }; From f6141563028c18071b54883b4c9e6d0b2ca8b611 Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 16:06:44 +0100 Subject: [PATCH 09/23] add legacy integration config upgrading logic --- .../deepnote/legacyIntegrationConfigUtils.ts | 89 +++++++++++++++++++ 1 file changed, 89 insertions(+) create mode 100644 src/platform/notebooks/deepnote/legacyIntegrationConfigUtils.ts diff --git a/src/platform/notebooks/deepnote/legacyIntegrationConfigUtils.ts b/src/platform/notebooks/deepnote/legacyIntegrationConfigUtils.ts new file mode 100644 index 0000000000..cd85a00d32 --- /dev/null +++ b/src/platform/notebooks/deepnote/legacyIntegrationConfigUtils.ts @@ -0,0 +1,89 @@ +import { LegacyIntegrationConfig, LegacyIntegrationType } from './integrationTypes'; +import { + BigQueryAuthMethods, + DatabaseIntegrationConfig, + databaseMetadataSchemasByType, + SnowflakeAuthMethods +} from '@deepnote/database-integrations'; +import { SnowflakeAuthMethods as LegacySnowflakeAuthMethods } from './snowflakeAuthConstants'; + +export async function upgradeLegacyIntegrationConfig( + config: LegacyIntegrationConfig +): Promise { + switch (config.type) { + case LegacyIntegrationType.Postgres: { + const metadata = databaseMetadataSchemasByType.pgsql.safeParse({ + host: config.host, + port: config.port ? String(config.port) : undefined, + database: config.database, + user: config.username, + password: config.password, + sslEnabled: config.ssl + }).data; + + return metadata + ? { + id: config.id, + name: config.name, + type: 'pgsql', + metadata + } + : null; + } + case LegacyIntegrationType.BigQuery: { + const metadata = databaseMetadataSchemasByType['big-query'].safeParse({ + authMethod: BigQueryAuthMethods.ServiceAccount, + service_account: config.credentials + }).data; + + return metadata + ? { + id: config.id, + name: config.name, + type: 'big-query', + metadata + } + : null; + } + case LegacyIntegrationType.Snowflake: { + const metadata = (() => { + switch (config.authMethod) { + case LegacySnowflakeAuthMethods.PASSWORD: + return databaseMetadataSchemasByType.snowflake.safeParse({ + authMethod: SnowflakeAuthMethods.Password, + accountName: config.account, + warehouse: config.warehouse, + database: config.database, + role: config.role, + username: config.username, + password: config.password + }).data; + case LegacySnowflakeAuthMethods.SERVICE_ACCOUNT_KEY_PAIR: + return databaseMetadataSchemasByType.snowflake.safeParse({ + authMethod: SnowflakeAuthMethods.ServiceAccountKeyPair, + accountName: config.account, + warehouse: config.warehouse, + database: config.database, + role: config.role, + username: config.username, + privateKey: config.privateKey, + privateKeyPassphrase: config.privateKeyPassphrase + }).data; + default: + return null; + } + })(); + + return metadata + ? { + id: config.id, + name: config.name, + type: 'snowflake', + metadata + } + : null; + } + default: + return null; + } +} From c98fe32d5760eb6fe40206ae34f5ba0e5f6a5073 Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 16:44:03 +0100 Subject: [PATCH 10/23] prefix integration config types with "legacy" to differentiate from new --- .../integrations/integrationDetector.ts | 5 ++- .../integrations/integrationManager.ts | 4 +- .../integrations/integrationWebview.ts | 4 +- .../deepnote/sqlCellStatusBarProvider.ts | 5 ++- .../notebooks/deepnote/integrationTypes.ts | 24 ++++++------ ...IntegrationEnvironmentVariablesProvider.ts | 2 +- ...nEnvironmentVariablesProvider.unit.test.ts | 38 +++++++++---------- 7 files changed, 42 insertions(+), 40 deletions(-) diff --git a/src/notebooks/deepnote/integrations/integrationDetector.ts b/src/notebooks/deepnote/integrations/integrationDetector.ts index 558439db45..2c2308ae0c 100644 --- a/src/notebooks/deepnote/integrations/integrationDetector.ts +++ b/src/notebooks/deepnote/integrations/integrationDetector.ts @@ -7,7 +7,7 @@ import { DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE, IntegrationStatus, IntegrationWithStatus, - RawIntegrationType + RawLegacyIntegrationType } from '../../../platform/notebooks/deepnote/integrationTypes'; import { IIntegrationDetector, IIntegrationStorage } from './types'; @@ -54,7 +54,8 @@ export class IntegrationDetector implements IIntegrationDetector { logger.debug(`IntegrationDetector: Found integration: ${integrationId} (${projectIntegration.type})`); // Map the Deepnote integration type to our IntegrationType - const integrationType = DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE[projectIntegration.type as RawIntegrationType]; + const integrationType = + DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE[projectIntegration.type as RawLegacyIntegrationType]; // Skip unknown integration types if (!integrationType) { diff --git a/src/notebooks/deepnote/integrations/integrationManager.ts b/src/notebooks/deepnote/integrations/integrationManager.ts index 00ac2dc4f1..7391367e97 100644 --- a/src/notebooks/deepnote/integrations/integrationManager.ts +++ b/src/notebooks/deepnote/integrations/integrationManager.ts @@ -10,7 +10,7 @@ import { IntegrationStatus, LegacyIntegrationType, IntegrationWithStatus, - RawIntegrationType + RawLegacyIntegrationType } from '../../../platform/notebooks/deepnote/integrationTypes'; import { BlockWithIntegration, scanBlocksForIntegrations } from './integrationUtils'; import { IDeepnoteNotebookManager } from '../../types'; @@ -173,7 +173,7 @@ export class IntegrationManager implements IIntegrationManager { if (projectIntegration.type in DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE) { // Map the Deepnote integration type to our IntegrationType integrationType = - DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE[projectIntegration.type as RawIntegrationType]; + DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE[projectIntegration.type as RawLegacyIntegrationType]; } else { logger.warn( `IntegrationManager: Unknown integration type '${projectIntegration.type}' for integration ID '${selectedIntegrationId}' in project '${projectId}'. Integration type will be undefined.` diff --git a/src/notebooks/deepnote/integrations/integrationWebview.ts b/src/notebooks/deepnote/integrations/integrationWebview.ts index 018012675c..d44ac9c0c8 100644 --- a/src/notebooks/deepnote/integrations/integrationWebview.ts +++ b/src/notebooks/deepnote/integrations/integrationWebview.ts @@ -13,7 +13,7 @@ import { LegacyIntegrationConfig, IntegrationStatus, IntegrationWithStatus, - RawIntegrationType + RawLegacyIntegrationType } from '../../../platform/notebooks/deepnote/integrationTypes'; /** @@ -357,7 +357,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { } // Map to Deepnote integration type - const deepnoteType: RawIntegrationType | undefined = LEGACY_INTEGRATION_TYPE_TO_DEEPNOTE[type]; + const deepnoteType: RawLegacyIntegrationType | undefined = LEGACY_INTEGRATION_TYPE_TO_DEEPNOTE[type]; if (!deepnoteType) { logger.warn(`IntegrationWebviewProvider: Cannot map type ${type} for integration ${id}, skipping`); return null; diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts index adf79bb358..262b3a4418 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts @@ -26,7 +26,7 @@ import { DATAFRAME_SQL_INTEGRATION_ID, DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE, LegacyIntegrationType, - RawIntegrationType + RawLegacyIntegrationType } from '../../platform/notebooks/deepnote/integrationTypes'; import { IDeepnoteNotebookManager } from '../types'; @@ -347,7 +347,8 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid continue; } - const integrationType = DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE[projectIntegration.type as RawIntegrationType]; + const integrationType = + DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE[projectIntegration.type as RawLegacyIntegrationType]; const typeLabel = integrationType ? this.getIntegrationTypeLabel(integrationType) : projectIntegration.type; const item: LocalQuickPickItem = { diff --git a/src/platform/notebooks/deepnote/integrationTypes.ts b/src/platform/notebooks/deepnote/integrationTypes.ts index 4d2ac4b853..4aaa09e3b9 100644 --- a/src/platform/notebooks/deepnote/integrationTypes.ts +++ b/src/platform/notebooks/deepnote/integrationTypes.ts @@ -24,13 +24,13 @@ export const LEGACY_INTEGRATION_TYPE_TO_DEEPNOTE = { [LegacyIntegrationType.Snowflake]: 'snowflake' } as const satisfies { [type in Exclude]: string }; -export type RawIntegrationType = +export type RawLegacyIntegrationType = (typeof LEGACY_INTEGRATION_TYPE_TO_DEEPNOTE)[keyof typeof LEGACY_INTEGRATION_TYPE_TO_DEEPNOTE]; /** * Map Deepnote integration type strings to our IntegrationType enum */ -export const DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE: Record = { +export const DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE: Record = { pgsql: LegacyIntegrationType.Postgres, 'big-query': LegacyIntegrationType.BigQuery, snowflake: LegacyIntegrationType.Snowflake @@ -39,7 +39,7 @@ export const DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE: Record { test('Returns environment variable for PostgreSQL integration', async () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'my-postgres-db'; - const config: PostgresIntegrationConfig = { + const config: LegacyPostgresIntegrationConfig = { id: integrationId, name: 'My Postgres DB', type: LegacyIntegrationType.Postgres, @@ -164,7 +164,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'my-bigquery'; const serviceAccountJson = JSON.stringify({ type: 'service_account', project_id: 'my-project' }); - const config: BigQueryIntegrationConfig = { + const config: LegacyBigQueryIntegrationConfig = { id: integrationId, name: 'My BigQuery', type: LegacyIntegrationType.BigQuery, @@ -197,7 +197,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const postgresId = 'my-postgres-db'; const bigqueryId = 'my-bigquery'; - const postgresConfig: PostgresIntegrationConfig = { + const postgresConfig: LegacyPostgresIntegrationConfig = { id: postgresId, name: 'My Postgres DB', type: LegacyIntegrationType.Postgres, @@ -208,7 +208,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { password: 'pass' }; - const bigqueryConfig: BigQueryIntegrationConfig = { + const bigqueryConfig: LegacyBigQueryIntegrationConfig = { id: bigqueryId, name: 'My BigQuery', type: LegacyIntegrationType.BigQuery, @@ -238,7 +238,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Properly encodes special characters in PostgreSQL credentials', async () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'special-chars-db'; - const config: PostgresIntegrationConfig = { + const config: LegacyPostgresIntegrationConfig = { id: integrationId, name: 'Special Chars DB', type: LegacyIntegrationType.Postgres, @@ -275,7 +275,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Normalizes integration ID with spaces and mixed case for env var name', async () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'My Production DB'; - const config: PostgresIntegrationConfig = { + const config: LegacyPostgresIntegrationConfig = { id: integrationId, name: 'Production Database', type: LegacyIntegrationType.Postgres, @@ -308,7 +308,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Normalizes integration ID with special characters for env var name', async () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'my-db@2024!'; - const config: PostgresIntegrationConfig = { + const config: LegacyPostgresIntegrationConfig = { id: integrationId, name: 'Test DB', type: LegacyIntegrationType.Postgres, @@ -348,7 +348,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Returns environment variable for Snowflake with PASSWORD auth', async () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'my-snowflake'; - const config: SnowflakeIntegrationConfig = { + const config: LegacySnowflakeIntegrationConfig = { id: integrationId, name: 'My Snowflake', type: LegacyIntegrationType.Snowflake, @@ -383,7 +383,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Returns environment variable for Snowflake with legacy null auth (username+password)', async () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'legacy-snowflake'; - const config: SnowflakeIntegrationConfig = { + const config: LegacySnowflakeIntegrationConfig = { id: integrationId, name: 'Legacy Snowflake', type: LegacyIntegrationType.Snowflake, @@ -418,7 +418,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const integrationId = 'snowflake-keypair'; const privateKey = '-----BEGIN ' + 'PRIVATE KEY-----\nfakekey-MIIEvQIBADANBg...\n-----END ' + 'PRIVATE KEY-----'; - const config: SnowflakeIntegrationConfig = { + const config: LegacySnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake KeyPair', type: LegacyIntegrationType.Snowflake, @@ -459,7 +459,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { const integrationId = 'snowflake-keypair-no-pass'; const privateKey = '-----BEGIN ' + 'PRIVATE KEY-----\nfakekey-MIIEvQIBADANBg...\n-----END ' + 'PRIVATE KEY-----'; - const config: SnowflakeIntegrationConfig = { + const config: LegacySnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake KeyPair No Pass', type: LegacyIntegrationType.Snowflake, @@ -496,7 +496,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Properly encodes special characters in Snowflake credentials', async () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'snowflake-special'; - const config: SnowflakeIntegrationConfig = { + const config: LegacySnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake Special', type: LegacyIntegrationType.Snowflake, @@ -530,7 +530,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Handles Snowflake with minimal optional fields', async () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'snowflake-minimal'; - const config: SnowflakeIntegrationConfig = { + const config: LegacySnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake Minimal', type: LegacyIntegrationType.Snowflake, @@ -559,7 +559,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Skips unsupported Snowflake auth method (OKTA)', async () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'snowflake-okta'; - const config: SnowflakeIntegrationConfig = { + const config: LegacySnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake OKTA', type: LegacyIntegrationType.Snowflake, @@ -582,7 +582,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Skips unsupported Snowflake auth method (AZURE_AD)', async () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'snowflake-azure'; - const config: SnowflakeIntegrationConfig = { + const config: LegacySnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake Azure', type: LegacyIntegrationType.Snowflake, @@ -604,7 +604,7 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { test('Skips unsupported Snowflake auth method (KEY_PAIR)', async () => { const uri = Uri.file('/test/notebook.deepnote'); const integrationId = 'snowflake-keypair-user'; - const config: SnowflakeIntegrationConfig = { + const config: LegacySnowflakeIntegrationConfig = { id: integrationId, name: 'Snowflake KeyPair User', type: LegacyIntegrationType.Snowflake, From 70580c1c5c208409c6b363c5857ee3f45343757c Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 16:44:46 +0100 Subject: [PATCH 11/23] migrate integration storage to format of deepnote/database-integrations --- .../notebooks/deepnote/integrationStorage.ts | 98 +++++++++++++++---- src/platform/notebooks/deepnote/types.ts | 13 ++- 2 files changed, 87 insertions(+), 24 deletions(-) diff --git a/src/platform/notebooks/deepnote/integrationStorage.ts b/src/platform/notebooks/deepnote/integrationStorage.ts index afd53c81c9..41b2158e06 100644 --- a/src/platform/notebooks/deepnote/integrationStorage.ts +++ b/src/platform/notebooks/deepnote/integrationStorage.ts @@ -4,11 +4,28 @@ import { EventEmitter } from 'vscode'; import { IEncryptedStorage } from '../../common/application/types'; import { IAsyncDisposableRegistry } from '../../common/types'; import { logger } from '../../logging'; -import { LegacyIntegrationConfig, LegacyIntegrationType } from './integrationTypes'; import { IIntegrationStorage } from './types'; +import { upgradeLegacyIntegrationConfig } from './legacyIntegrationConfigUtils'; +import { + DatabaseIntegrationConfig, + DatabaseIntegrationType, + databaseIntegrationTypes, + databaseMetadataSchemasByType +} from '@deepnote/database-integrations'; const INTEGRATION_SERVICE_NAME = 'deepnote-integrations'; +// NOTE: We need a way to upgrade existing configurations to the new format of deepnote/database-integrations. +type VersionedDatabaseIntegrationConfig = DatabaseIntegrationConfig & { version: 1 }; + +function storeEncryptedIntegrationConfig( + encryptedStorage: IEncryptedStorage, + integrationId: string, + config: VersionedDatabaseIntegrationConfig +): Promise { + return encryptedStorage.store(INTEGRATION_SERVICE_NAME, integrationId, JSON.stringify(config)); +} + /** * Storage service for integration configurations. * Uses VSCode's SecretStorage API to securely store credentials. @@ -16,7 +33,7 @@ const INTEGRATION_SERVICE_NAME = 'deepnote-integrations'; */ @injectable() export class IntegrationStorage implements IIntegrationStorage { - private readonly cache: Map = new Map(); + private readonly cache: Map = new Map(); private cacheLoaded = false; @@ -35,7 +52,7 @@ export class IntegrationStorage implements IIntegrationStorage { /** * Get all stored integration configurations. */ - async getAll(): Promise { + async getAll(): Promise { await this.ensureCacheLoaded(); return Array.from(this.cache.values()); } @@ -43,7 +60,7 @@ export class IntegrationStorage implements IIntegrationStorage { /** * Get a specific integration configuration by ID */ - async getIntegrationConfig(integrationId: string): Promise { + async getIntegrationConfig(integrationId: string): Promise { await this.ensureCacheLoaded(); return this.cache.get(integrationId); } @@ -56,28 +73,18 @@ export class IntegrationStorage implements IIntegrationStorage { async getProjectIntegrationConfig( _projectId: string, integrationId: string - ): Promise { + ): Promise { return this.getIntegrationConfig(integrationId); } - /** - * Get all integrations of a specific type - */ - async getByType(type: LegacyIntegrationType): Promise { - await this.ensureCacheLoaded(); - return Array.from(this.cache.values()).filter((config) => config.type === type); - } - /** * Save or update an integration configuration */ - async save(config: LegacyIntegrationConfig): Promise { + async save(config: DatabaseIntegrationConfig): Promise { await this.ensureCacheLoaded(); // Store the configuration as JSON in encrypted storage - const configJson = JSON.stringify(config); - await this.encryptedStorage.store(INTEGRATION_SERVICE_NAME, config.id, configJson); - + await storeEncryptedIntegrationConfig(this.encryptedStorage, config.id, { ...config, version: 1 }); // Update cache this.cache.set(config.id, config); @@ -154,19 +161,72 @@ export class IntegrationStorage implements IIntegrationStorage { try { const integrationIds: string[] = JSON.parse(indexJson); + const idsToDelete: string[] = []; // Load each integration configuration for (const id of integrationIds) { const configJson = await this.encryptedStorage.retrieve(INTEGRATION_SERVICE_NAME, id); if (configJson) { try { - const config: LegacyIntegrationConfig = JSON.parse(configJson); - this.cache.set(id, config); + const parsedData = JSON.parse(configJson); + + // Check if this is a legacy config (missing 'version' field) + if (!('version' in parsedData)) { + logger.info(`Upgrading legacy integration config for ${id}`); + + // Attempt to upgrade the legacy config + const upgradedConfig = await upgradeLegacyIntegrationConfig(parsedData); + + if (upgradedConfig) { + // Successfully upgraded - save the new config + logger.info(`Successfully upgraded integration config for ${id}`); + await storeEncryptedIntegrationConfig(this.encryptedStorage, id, { + ...upgradedConfig, + version: 1 + }); + this.cache.set(id, upgradedConfig); + } else { + // Upgrade failed - mark for deletion + logger.warn(`Failed to upgrade integration ${id}, marking for deletion`); + idsToDelete.push(id); + } + } else { + // Already versioned config - validate against current schema + const { version: _version, ...rawConfig } = parsedData; + const config = databaseIntegrationTypes.includes(rawConfig.type) + ? (rawConfig as DatabaseIntegrationConfig) + : null; + const validMetadata = config + ? databaseMetadataSchemasByType[config.type].safeParse(config.metadata).data + : null; + if (config && validMetadata) { + this.cache.set( + id, + // NOTE: We must cast here because there is no union-wide schema parser at the moment. + { ...config, metadata: validMetadata } as DatabaseIntegrationConfig + ); + } else { + logger.warn(`Invalid integration config for ${id}, marking for deletion`); + idsToDelete.push(id); + } + } } catch (error) { logger.error(`Failed to parse integration config for ${id}:`, error); + // Mark corrupted configs for deletion + idsToDelete.push(id); } } } + + // Delete any configs that failed to upgrade or were corrupted + if (idsToDelete.length > 0) { + logger.info(`Deleting ${idsToDelete.length} invalid integration config(s)`); + for (const id of idsToDelete) { + await this.encryptedStorage.store(INTEGRATION_SERVICE_NAME, id, undefined); + } + // Update the index to remove deleted IDs + await this.updateIndex(); + } } catch (error) { logger.error('Failed to parse integration index:', error); } diff --git a/src/platform/notebooks/deepnote/types.ts b/src/platform/notebooks/deepnote/types.ts index 9147496835..2632483c99 100644 --- a/src/platform/notebooks/deepnote/types.ts +++ b/src/platform/notebooks/deepnote/types.ts @@ -1,8 +1,8 @@ import { CancellationToken, Event, NotebookDocument, Uri } from 'vscode'; import { IDisposable, Resource } from '../../common/types'; import { EnvironmentVariables } from '../../common/variables/types'; -import { LegacyIntegrationConfig } from './integrationTypes'; import { DeepnoteProject } from '../../deepnote/deepnoteTypes'; +import { DatabaseIntegrationConfig } from '@deepnote/database-integrations'; /** * Settings for select input blocks @@ -31,7 +31,7 @@ export interface IIntegrationStorage extends IDisposable { */ readonly onDidChangeIntegrations: Event; - getAll(): Promise; + getAll(): Promise; /** * Retrieves the global (non-project-scoped) integration configuration by integration ID. @@ -49,14 +49,17 @@ export interface IIntegrationStorage extends IDisposable { * - The `IntegrationConfig` object if a global configuration exists for the given ID * - `undefined` if no global configuration exists for the given integration ID */ - getIntegrationConfig(integrationId: string): Promise; + getIntegrationConfig(integrationId: string): Promise; /** * Get integration configuration for a specific project and integration */ - getProjectIntegrationConfig(projectId: string, integrationId: string): Promise; + getProjectIntegrationConfig( + projectId: string, + integrationId: string + ): Promise; - save(config: LegacyIntegrationConfig): Promise; + save(config: DatabaseIntegrationConfig): Promise; delete(integrationId: string): Promise; exists(integrationId: string): Promise; clear(): Promise; From 392f52cf09df4c70d26c88cacb07caf445238056 Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 16:45:18 +0100 Subject: [PATCH 12/23] replace integration env var logic with deepnote/database-integrations --- ...IntegrationEnvironmentVariablesProvider.ts | 213 ++---------------- 1 file changed, 25 insertions(+), 188 deletions(-) diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts index d89ffb7d53..649a2f78a9 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts @@ -18,148 +18,7 @@ import { LegacyDuckDBIntegrationConfig, DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes'; - -/** - * Converts an integration ID to the environment variable name format expected by SQL blocks. - * Example: 'my-postgres-db' -> 'SQL_MY_POSTGRES_DB' - */ -function convertToEnvironmentVariableName(str: string): string { - return (/^\d/.test(str) ? `_${str}` : str).toUpperCase().replace(/[^\w]/g, '_'); -} - -function getSqlEnvVarName(integrationId: string): string { - return `SQL_${integrationId}`; -} - -/** - * Converts integration configuration to the JSON format expected by the SQL execution code. - * The format must match what deepnote_toolkit expects: - * { - * "url": "sqlalchemy_connection_url", - * "params": {}, - * "param_style": "qmark" | "format" | etc. - * } - */ -function convertIntegrationConfigToJson(config: LegacyIntegrationConfig): string { - switch (config.type) { - case IntegrationType.DuckDB: { - // Internal DuckDB integration for querying dataframes - return JSON.stringify({ - url: 'deepnote+duckdb:///:memory:', - params: {}, - param_style: 'qmark' - }); - } - - case IntegrationType.Postgres: { - // Build PostgreSQL connection URL - // Format: postgresql://username:password@host:port/database - const encodedUsername = encodeURIComponent(config.username); - const encodedPassword = encodeURIComponent(config.password); - const encodedDatabase = encodeURIComponent(config.database); - const url = `postgresql://${encodedUsername}:${encodedPassword}@${config.host}:${config.port}/${encodedDatabase}`; - - return JSON.stringify({ - url: url, - params: config.ssl ? { sslmode: 'require' } : {}, - param_style: 'format' - }); - } - - case IntegrationType.BigQuery: { - // BigQuery uses a special URL format - return JSON.stringify({ - url: 'bigquery://?user_supplied_client=true', - params: { - project_id: config.projectId, - credentials: JSON.parse(config.credentials) - }, - param_style: 'format' - }); - } - - case IntegrationType.Snowflake: { - // Build Snowflake connection URL - // Format depends on auth method: - // Username+password: snowflake://{username}:{password}@{account}/{database}?warehouse={warehouse}&role={role}&application=YourApp - // Service account key-pair: snowflake://{username}@{account}/{database}?warehouse={warehouse}&role={role}&authenticator=snowflake_jwt&application=YourApp - const encodedAccount = encodeURIComponent(config.account); - - let url: string; - const params: Record = {}; - - if (config.authMethod === null || config.authMethod === SnowflakeAuthMethods.PASSWORD) { - // Username+password authentication - const encodedUsername = encodeURIComponent(config.username); - const encodedPassword = encodeURIComponent(config.password); - const database = config.database ? `/${encodeURIComponent(config.database)}` : ''; - url = `snowflake://${encodedUsername}:${encodedPassword}@${encodedAccount}${database}`; - - const queryParams = new URLSearchParams(); - if (config.warehouse) { - queryParams.set('warehouse', config.warehouse); - } - if (config.role) { - queryParams.set('role', config.role); - } - queryParams.set('application', 'Deepnote'); - - const queryString = queryParams.toString(); - if (queryString) { - url += `?${queryString}`; - } - } else { - // Service account key-pair authentication (the only other supported method) - // TypeScript needs help narrowing the type here - if (config.authMethod !== SnowflakeAuthMethods.SERVICE_ACCOUNT_KEY_PAIR) { - // This should never happen due to the type guard above, but TypeScript needs this - throw new UnsupportedIntegrationError( - l10n.t( - "Snowflake integration with auth method '{0}' is not supported in VSCode", - config.authMethod - ) - ); - } - - const encodedUsername = encodeURIComponent(config.username); - const database = config.database ? `/${encodeURIComponent(config.database)}` : ''; - url = `snowflake://${encodedUsername}@${encodedAccount}${database}`; - - const queryParams = new URLSearchParams(); - if (config.warehouse) { - queryParams.set('warehouse', config.warehouse); - } - if (config.role) { - queryParams.set('role', config.role); - } - queryParams.set('authenticator', 'snowflake_jwt'); - queryParams.set('application', 'Deepnote'); - - const queryString = queryParams.toString(); - if (queryString) { - url += `?${queryString}`; - } - - // For key-pair auth, pass the private key and passphrase as params - params.snowflake_private_key = btoa(config.privateKey); - if (config.privateKeyPassphrase) { - params.snowflake_private_key_passphrase = config.privateKeyPassphrase; - } - } - - return JSON.stringify({ - url: url, - params: params, - param_style: 'pyformat' - }); - } - - default: - throw new UnsupportedIntegrationError( - l10n.t('Unsupported integration type: {0}', (config as LegacyIntegrationConfig).type) - ); - } -} +import { getEnvironmentVariablesForIntegrations } from '@deepnote/database-integrations'; /** * Provides environment variables for SQL integrations. @@ -197,14 +56,12 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati * The internal DuckDB integration is always included. */ public async getEnvironmentVariables(resource: Resource, token?: CancellationToken): Promise { - const envVars: EnvironmentVariables = {}; - if (!resource) { - return envVars; + return {}; } if (token?.isCancellationRequested) { - return envVars; + return {}; } logger.trace(`SqlIntegrationEnvironmentVariablesProvider: Getting env vars for resource`); @@ -213,14 +70,14 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati const notebook = this.notebookEditorProvider.findAssociatedNotebookDocument(resource); if (!notebook) { logger.trace(`SqlIntegrationEnvironmentVariablesProvider: No notebook found for resource`); - return envVars; + return {}; } // Get the project ID from the notebook metadata const projectId = notebook.metadata?.deepnoteProjectId as string | undefined; if (!projectId) { logger.trace(`SqlIntegrationEnvironmentVariablesProvider: No project ID found in notebook metadata`); - return envVars; + return {}; } logger.trace(`SqlIntegrationEnvironmentVariablesProvider: Project ID: ${projectId}`); @@ -229,7 +86,7 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati const project = this.notebookManager.getOriginalProject(projectId); if (!project) { logger.trace(`SqlIntegrationEnvironmentVariablesProvider: No project found for ID: ${projectId}`); - return envVars; + return {}; } // Get the list of integrations from the project @@ -238,51 +95,31 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati `SqlIntegrationEnvironmentVariablesProvider: Found ${projectIntegrations.length} integrations in project` ); + const projectIntegrationConfigs = ( + await Promise.all( + projectIntegrations.map((integration) => { + return this.integrationStorage.getIntegrationConfig(integration.id); + }) + ) + ).filter((config) => config != null); + // Always add the internal DuckDB integration - const duckdbConfig: DuckDBIntegrationConfig = { + projectIntegrationConfigs.push({ id: DATAFRAME_SQL_INTEGRATION_ID, name: 'Dataframe SQL (DuckDB)', - type: IntegrationType.DuckDB - }; - projectIntegrations.push(duckdbConfig); + type: 'pandas-dataframe', + metadata: {} + }); - // Get credentials for each project integration and add to environment variables - for (const projectIntegration of projectIntegrations) { - if (token?.isCancellationRequested) { - break; - } + const { envVars: envVarList, errors } = getEnvironmentVariablesForIntegrations(projectIntegrationConfigs, { + projectRootDirectory: '' + }); - const integrationId = projectIntegration.id; - - try { - // Get the integration configuration from storage - const config = - integrationId === DATAFRAME_SQL_INTEGRATION_ID - ? duckdbConfig - : await this.integrationStorage.getIntegrationConfig(integrationId); - if (!config) { - logger.debug( - `SqlIntegrationEnvironmentVariablesProvider: No configuration found for integration ${integrationId}, skipping` - ); - continue; - } - - // Convert integration config to JSON and add as environment variable - const envVarName = convertToEnvironmentVariableName(getSqlEnvVarName(config.id)); - const credentialsJson = convertIntegrationConfigToJson(config); - - envVars[envVarName] = credentialsJson; - logger.debug( - `SqlIntegrationEnvironmentVariablesProvider: Added env var ${envVarName} for integration ${config.id}` - ); - } catch (error) { - logger.error( - `SqlIntegrationEnvironmentVariablesProvider: Failed to get credentials for integration ${integrationId}`, - error - ); - } - } + errors.forEach((error) => { + logger.error(`SqlIntegrationEnvironmentVariablesProvider: ${error.message}`); + }); + const envVars: EnvironmentVariables = Object.fromEntries(envVarList.map(({ name, value }) => [name, value])); logger.trace(`SqlIntegrationEnvironmentVariablesProvider: Returning ${Object.keys(envVars).length} env vars`); return envVars; From d1e8f493df809409580c32259cebd1ba2707197f Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 17:10:49 +0100 Subject: [PATCH 13/23] prevent saving and loading of duckdb integrations via integration storage --- .../notebooks/deepnote/integrationStorage.ts | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/src/platform/notebooks/deepnote/integrationStorage.ts b/src/platform/notebooks/deepnote/integrationStorage.ts index 41b2158e06..2d57fbf132 100644 --- a/src/platform/notebooks/deepnote/integrationStorage.ts +++ b/src/platform/notebooks/deepnote/integrationStorage.ts @@ -8,10 +8,10 @@ import { IIntegrationStorage } from './types'; import { upgradeLegacyIntegrationConfig } from './legacyIntegrationConfigUtils'; import { DatabaseIntegrationConfig, - DatabaseIntegrationType, databaseIntegrationTypes, databaseMetadataSchemasByType } from '@deepnote/database-integrations'; +import { DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes'; const INTEGRATION_SERVICE_NAME = 'deepnote-integrations'; @@ -81,6 +81,11 @@ export class IntegrationStorage implements IIntegrationStorage { * Save or update an integration configuration */ async save(config: DatabaseIntegrationConfig): Promise { + if (config.type === 'pandas-dataframe' || config.id === DATAFRAME_SQL_INTEGRATION_ID) { + logger.warn(`IntegrationStorage: Skipping save for internal DuckDB integration ${config.id}`); + return; + } + await this.ensureCacheLoaded(); // Store the configuration as JSON in encrypted storage @@ -165,6 +170,10 @@ export class IntegrationStorage implements IIntegrationStorage { // Load each integration configuration for (const id of integrationIds) { + if (id === DATAFRAME_SQL_INTEGRATION_ID) { + continue; + } + const configJson = await this.encryptedStorage.retrieve(INTEGRATION_SERVICE_NAME, id); if (configJson) { try { @@ -178,6 +187,11 @@ export class IntegrationStorage implements IIntegrationStorage { const upgradedConfig = await upgradeLegacyIntegrationConfig(parsedData); if (upgradedConfig) { + if (upgradedConfig.type === 'pandas-dataframe') { + logger.warn(`IntegrationStorage: Skipping internal DuckDB integration ${id}`); + continue; + } + // Successfully upgraded - save the new config logger.info(`Successfully upgraded integration config for ${id}`); await storeEncryptedIntegrationConfig(this.encryptedStorage, id, { @@ -193,9 +207,11 @@ export class IntegrationStorage implements IIntegrationStorage { } else { // Already versioned config - validate against current schema const { version: _version, ...rawConfig } = parsedData; - const config = databaseIntegrationTypes.includes(rawConfig.type) - ? (rawConfig as DatabaseIntegrationConfig) - : null; + const config = + databaseIntegrationTypes.includes(rawConfig.type) && + rawConfig.type !== 'pandas-dataframe' + ? (rawConfig as DatabaseIntegrationConfig) + : null; const validMetadata = config ? databaseMetadataSchemasByType[config.type].safeParse(config.metadata).data : null; From 4a7c399b18f7317855bc4a5098d4c88c6b6a3af4 Mon Sep 17 00:00:00 2001 From: jankuca Date: Sat, 1 Nov 2025 18:42:41 +0100 Subject: [PATCH 14/23] migrate all components to the new database integration config format --- .../integrations/integrationDetector.ts | 36 +- .../integrations/integrationManager.ts | 71 +- .../deepnote/integrations/integrationUtils.ts | 68 -- .../integrations/integrationWebview.ts | 25 +- .../deepnote/sqlCellStatusBarProvider.ts | 24 +- .../sqlCellStatusBarProvider.unit.test.ts | 35 +- .../notebooks/deepnote/integrationTypes.ts | 5 +- ...IntegrationEnvironmentVariablesProvider.ts | 11 +- ...nEnvironmentVariablesProvider.unit.test.ts | 628 ------------------ .../integrations/BigQueryForm.tsx | 121 ++-- .../integrations/ConfigurationForm.tsx | 98 ++- .../integrations/IntegrationItem.tsx | 9 +- .../integrations/IntegrationPanel.tsx | 19 +- .../integrations/PostgresForm.tsx | 180 +++-- .../integrations/SnowflakeForm.tsx | 618 +++++++++-------- .../webview-side/integrations/types.ts | 89 +-- 16 files changed, 644 insertions(+), 1393 deletions(-) delete mode 100644 src/notebooks/deepnote/integrations/integrationUtils.ts delete mode 100644 src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts diff --git a/src/notebooks/deepnote/integrations/integrationDetector.ts b/src/notebooks/deepnote/integrations/integrationDetector.ts index 2c2308ae0c..a17fc2ba30 100644 --- a/src/notebooks/deepnote/integrations/integrationDetector.ts +++ b/src/notebooks/deepnote/integrations/integrationDetector.ts @@ -2,14 +2,9 @@ import { inject, injectable } from 'inversify'; import { logger } from '../../../platform/logging'; import { IDeepnoteNotebookManager } from '../../types'; -import { - DATAFRAME_SQL_INTEGRATION_ID, - DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE, - IntegrationStatus, - IntegrationWithStatus, - RawLegacyIntegrationType -} from '../../../platform/notebooks/deepnote/integrationTypes'; +import { IntegrationStatus, IntegrationWithStatus } from '../../../platform/notebooks/deepnote/integrationTypes'; import { IIntegrationDetector, IIntegrationStorage } from './types'; +import { DatabaseIntegrationType, databaseIntegrationTypes } from '@deepnote/database-integrations'; /** * Service for detecting integrations used in Deepnote notebooks @@ -40,40 +35,25 @@ export class IntegrationDetector implements IIntegrationDetector { const integrations = new Map(); // Use the project's integrations field as the source of truth - const projectIntegrations = project.project.integrations || []; + const projectIntegrations = project.project.integrations?.slice() ?? []; logger.debug(`IntegrationDetector: Found ${projectIntegrations.length} integrations in project.integrations`); for (const projectIntegration of projectIntegrations) { const integrationId = projectIntegration.id; - - // Skip the internal DuckDB integration - if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) { - continue; - } - - logger.debug(`IntegrationDetector: Found integration: ${integrationId} (${projectIntegration.type})`); - - // Map the Deepnote integration type to our IntegrationType - const integrationType = - DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE[projectIntegration.type as RawLegacyIntegrationType]; - - // Skip unknown integration types - if (!integrationType) { - logger.warn( - `IntegrationDetector: Unknown integration type '${projectIntegration.type}' for integration ID '${integrationId}'. Skipping.` - ); + const integrationType = projectIntegration.type; + if (!(databaseIntegrationTypes as readonly string[]).includes(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, + 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 + integrationType: integrationType as DatabaseIntegrationType }; integrations.set(integrationId, status); diff --git a/src/notebooks/deepnote/integrations/integrationManager.ts b/src/notebooks/deepnote/integrations/integrationManager.ts index 7391367e97..24456c59b7 100644 --- a/src/notebooks/deepnote/integrations/integrationManager.ts +++ b/src/notebooks/deepnote/integrations/integrationManager.ts @@ -1,19 +1,13 @@ import { inject, injectable } from 'inversify'; -import { commands, l10n, NotebookDocument, window, workspace } from 'vscode'; +import { commands, l10n, window, workspace } 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 { - DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE, - IntegrationStatus, - LegacyIntegrationType, - IntegrationWithStatus, - RawLegacyIntegrationType -} from '../../../platform/notebooks/deepnote/integrationTypes'; -import { BlockWithIntegration, scanBlocksForIntegrations } from './integrationUtils'; +import { IntegrationStatus } from '../../../platform/notebooks/deepnote/integrationTypes'; import { IDeepnoteNotebookManager } from '../../types'; +import { DatabaseIntegrationType, databaseIntegrationTypes } from '@deepnote/database-integrations'; /** * Manages integration UI and commands for Deepnote notebooks @@ -143,14 +137,6 @@ export class IntegrationManager implements IIntegrationManager { // First try to detect integrations from the stored project let integrations = await this.integrationDetector.detectIntegrations(projectId); - - // If no integrations found in stored project, scan cells directly - // This handles the case where the notebook was already open when the extension loaded - if (integrations.size === 0) { - logger.debug(`IntegrationManager: No integrations found in stored project, scanning cells directly`); - integrations = await this.detectIntegrationsFromCells(activeNotebook); - } - logger.debug(`IntegrationManager: Found ${integrations.size} integrations`); // If a specific integration was requested (e.g., from status bar click), @@ -164,21 +150,15 @@ export class IntegrationManager implements IIntegrationManager { const projectIntegration = project?.project.integrations?.find((i) => i.id === selectedIntegrationId); let integrationName: string | undefined; - let integrationType: LegacyIntegrationType | undefined; + let integrationType: DatabaseIntegrationType | undefined; - if (projectIntegration) { + // Validate that projectIntegration.type against supported types + if ( + projectIntegration && + (databaseIntegrationTypes as readonly string[]).includes(projectIntegration.type) + ) { integrationName = projectIntegration.name; - - // Validate that projectIntegration.type exists in the mapping before lookup - if (projectIntegration.type in DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE) { - // Map the Deepnote integration type to our IntegrationType - integrationType = - DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE[projectIntegration.type as RawLegacyIntegrationType]; - } else { - logger.warn( - `IntegrationManager: Unknown integration type '${projectIntegration.type}' for integration ID '${selectedIntegrationId}' in project '${projectId}'. Integration type will be undefined.` - ); - } + integrationType = projectIntegration.type as DatabaseIntegrationType; } integrations.set(selectedIntegrationId, { @@ -197,35 +177,4 @@ export class IntegrationManager implements IIntegrationManager { // Show the webview with optional selected integration await this.webviewProvider.show(projectId, integrations, selectedIntegrationId); } - - /** - * Detect integrations by scanning cells directly (fallback method) - * This is used when the project isn't stored in the notebook manager - */ - private async detectIntegrationsFromCells(notebook: NotebookDocument): Promise> { - // Collect all cells with SQL integration metadata - const blocksWithIntegrations: BlockWithIntegration[] = []; - - for (const cell of notebook.getCells()) { - const metadata = cell.metadata; - logger.trace(`IntegrationManager: Cell ${cell.index} metadata:`, metadata); - - // Check cell metadata for sql_integration_id - if (metadata && typeof metadata === 'object') { - const integrationId = (metadata as Record).sql_integration_id; - if (typeof integrationId === 'string') { - logger.debug(`IntegrationManager: Found integration ${integrationId} in cell ${cell.index}`); - blocksWithIntegrations.push({ - id: `cell-${cell.index}`, - sql_integration_id: integrationId - }); - } - } - } - - logger.debug(`IntegrationManager: Found ${blocksWithIntegrations.length} cells with integrations`); - - // Use the shared utility to scan blocks and build the status map - return scanBlocksForIntegrations(blocksWithIntegrations, this.integrationStorage, 'IntegrationManager'); - } } diff --git a/src/notebooks/deepnote/integrations/integrationUtils.ts b/src/notebooks/deepnote/integrations/integrationUtils.ts deleted file mode 100644 index 9e9015ca9e..0000000000 --- a/src/notebooks/deepnote/integrations/integrationUtils.ts +++ /dev/null @@ -1,68 +0,0 @@ -import { logger } from '../../../platform/logging'; -import { IIntegrationStorage } from './types'; -import { - DATAFRAME_SQL_INTEGRATION_ID, - IntegrationStatus, - IntegrationWithStatus -} from '../../../platform/notebooks/deepnote/integrationTypes'; - -/** - * Represents a block with SQL integration metadata - */ -export interface BlockWithIntegration { - id: string; - sql_integration_id: string; -} - -/** - * Scans blocks for SQL integrations and builds a status map. - * This is the core logic shared between IntegrationDetector and IntegrationManager. - * - * @param blocks - Iterator of blocks to scan (can be from Deepnote project or VSCode notebook cells) - * @param integrationStorage - Storage service to check configuration status - * @param logContext - Context string for logging (e.g., "IntegrationDetector", "IntegrationManager") - * @returns Map of integration IDs to their status - */ -export async function scanBlocksForIntegrations( - blocks: Iterable, - integrationStorage: IIntegrationStorage, - logContext: string -): Promise> { - const integrations = new Map(); - - for (const block of blocks) { - const integrationId = block.sql_integration_id; - - // Skip blocks without integration IDs - if (!integrationId) { - continue; - } - - // Skip excluded integrations (e.g., internal DuckDB integration) - if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) { - logger.trace(`${logContext}: Skipping excluded integration: ${integrationId} in block ${block.id}`); - continue; - } - - // Skip if we've already detected this integration - if (integrations.has(integrationId)) { - continue; - } - - logger.debug(`${logContext}: Found integration: ${integrationId} in block ${block.id}`); - - // Check if the integration is configured - const config = await integrationStorage.getIntegrationConfig(integrationId); - - const status: IntegrationWithStatus = { - config: config || null, - status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected - }; - - integrations.set(integrationId, status); - } - - logger.debug(`${logContext}: Found ${integrations.size} integrations`); - - return integrations; -} diff --git a/src/notebooks/deepnote/integrations/integrationWebview.ts b/src/notebooks/deepnote/integrations/integrationWebview.ts index d44ac9c0c8..b652ff1690 100644 --- a/src/notebooks/deepnote/integrations/integrationWebview.ts +++ b/src/notebooks/deepnote/integrations/integrationWebview.ts @@ -7,14 +7,8 @@ import { logger } from '../../../platform/logging'; import { LocalizedMessages, SharedMessages } from '../../../messageTypes'; import { IDeepnoteNotebookManager, ProjectIntegration } from '../../types'; import { IIntegrationStorage, IIntegrationWebviewProvider } from './types'; -import { - LEGACY_INTEGRATION_TYPE_TO_DEEPNOTE, - LegacyIntegrationType, - LegacyIntegrationConfig, - IntegrationStatus, - IntegrationWithStatus, - RawLegacyIntegrationType -} from '../../../platform/notebooks/deepnote/integrationTypes'; +import { IntegrationStatus, IntegrationWithStatus } from '../../../platform/notebooks/deepnote/integrationTypes'; +import { DatabaseIntegrationConfig } from '@deepnote/database-integrations'; /** * Manages the webview panel for integration configuration @@ -222,7 +216,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { private async handleMessage(message: { type: string; integrationId?: string; - config?: LegacyIntegrationConfig; + config?: DatabaseIntegrationConfig; }): Promise { switch (message.type) { case 'configure': @@ -264,7 +258,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { /** * Save the configuration for an integration */ - private async saveConfiguration(integrationId: string, config: LegacyIntegrationConfig): Promise { + private async saveConfiguration(integrationId: string, config: DatabaseIntegrationConfig): Promise { try { await this.integrationStorage.save(config); @@ -351,22 +345,15 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { } // Skip DuckDB integration (internal, not a real Deepnote integration) - if (type === LegacyIntegrationType.DuckDB) { + if (type === 'pandas-dataframe') { logger.trace(`IntegrationWebviewProvider: Skipping internal DuckDB integration ${id}`); return null; } - // Map to Deepnote integration type - const deepnoteType: RawLegacyIntegrationType | undefined = LEGACY_INTEGRATION_TYPE_TO_DEEPNOTE[type]; - if (!deepnoteType) { - logger.warn(`IntegrationWebviewProvider: Cannot map type ${type} for integration ${id}, skipping`); - return null; - } - return { id, name: integration.config?.name || integration.integrationName || id, - type: deepnoteType + type }; }) .filter((integration): integration is ProjectIntegration => integration !== null); diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts index 262b3a4418..20ca8f45ec 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts @@ -22,13 +22,9 @@ import { IExtensionSyncActivationService } from '../../platform/activation/types import { IDisposableRegistry } from '../../platform/common/types'; import { IIntegrationStorage } from './integrations/types'; import { Commands } from '../../platform/common/constants'; -import { - DATAFRAME_SQL_INTEGRATION_ID, - DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE, - LegacyIntegrationType, - RawLegacyIntegrationType -} from '../../platform/notebooks/deepnote/integrationTypes'; +import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes'; import { IDeepnoteNotebookManager } from '../types'; +import { DatabaseIntegrationType, databaseIntegrationTypes } from '@deepnote/database-integrations'; /** * QuickPick item with an integration ID @@ -347,9 +343,11 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid continue; } - const integrationType = - DEEPNOTE_TO_LEGACY_INTEGRATION_TYPE[projectIntegration.type as RawLegacyIntegrationType]; - const typeLabel = integrationType ? this.getIntegrationTypeLabel(integrationType) : projectIntegration.type; + const integrationType = projectIntegration.type; + const typeLabel = + integrationType && (databaseIntegrationTypes as readonly string[]).includes(integrationType) + ? this.getIntegrationTypeLabel(integrationType as DatabaseIntegrationType) + : projectIntegration.type; const item: LocalQuickPickItem = { label: projectIntegration.name || projectIntegration.id, @@ -432,13 +430,13 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid this._onDidChangeCellStatusBarItems.fire(); } - private getIntegrationTypeLabel(type: LegacyIntegrationType): string { + private getIntegrationTypeLabel(type: DatabaseIntegrationType): string { switch (type) { - case LegacyIntegrationType.Postgres: + case 'pgsql': return l10n.t('PostgreSQL'); - case LegacyIntegrationType.BigQuery: + case 'big-query': return l10n.t('BigQuery'); - case LegacyIntegrationType.Snowflake: + case 'snowflake': return l10n.t('Snowflake'); default: return String(type); diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts index 1487339620..40a72119b0 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts @@ -14,10 +14,7 @@ import { import { IDisposableRegistry } from '../../platform/common/types'; import { IIntegrationStorage } from './integrations/types'; import { SqlCellStatusBarProvider } from './sqlCellStatusBarProvider'; -import { - DATAFRAME_SQL_INTEGRATION_ID, - LegacyIntegrationType -} from '../../platform/notebooks/deepnote/integrationTypes'; +import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; import { createEventHandler } from '../../test/common'; import { Commands } from '../../platform/common/constants'; @@ -137,12 +134,15 @@ suite('SqlCellStatusBarProvider', () => { when(integrationStorage.getProjectIntegrationConfig(anything(), anything())).thenResolve({ id: integrationId, name: 'My Postgres DB', - type: LegacyIntegrationType.Postgres, - host: 'localhost', - port: 5432, - database: 'test', - username: 'user', - password: 'pass' + type: 'pgsql', + metadata: { + host: 'localhost', + port: '5432', + database: 'test', + user: 'user', + password: 'pass', + sslEnabled: false + } }); const result = await provider.provideCellStatusBarItems(cell, cancellationToken); @@ -285,12 +285,15 @@ suite('SqlCellStatusBarProvider', () => { when(integrationStorage.getProjectIntegrationConfig(anything(), anything())).thenResolve({ id: integrationId, name: 'My Postgres DB', - type: LegacyIntegrationType.Postgres, - host: 'localhost', - port: 5432, - database: 'test', - username: 'user', - password: 'pass' + type: 'pgsql', + metadata: { + host: 'localhost', + port: '5432', + database: 'test', + user: 'user', + password: 'pass', + sslEnabled: false + } }); const result = await provider.provideCellStatusBarItems(cell, cancellationToken); diff --git a/src/platform/notebooks/deepnote/integrationTypes.ts b/src/platform/notebooks/deepnote/integrationTypes.ts index 4aaa09e3b9..19da67bcdf 100644 --- a/src/platform/notebooks/deepnote/integrationTypes.ts +++ b/src/platform/notebooks/deepnote/integrationTypes.ts @@ -74,6 +74,7 @@ export interface LegacyDuckDBIntegrationConfig extends BaseLegacyIntegrationConf type: LegacyIntegrationType.DuckDB; } +import { DatabaseIntegrationConfig, DatabaseIntegrationType } from '@deepnote/database-integrations'; // Import and re-export Snowflake auth constants from shared module import { type SnowflakeAuthMethod, @@ -148,7 +149,7 @@ export enum IntegrationStatus { * Integration with its current status */ export interface IntegrationWithStatus { - config: LegacyIntegrationConfig | null; + config: DatabaseIntegrationConfig | null; status: IntegrationStatus; error?: string; /** @@ -158,5 +159,5 @@ export interface IntegrationWithStatus { /** * Type from the project's integrations list (used for prefilling when config is null) */ - integrationType?: LegacyIntegrationType; + integrationType?: DatabaseIntegrationType; } diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts index 649a2f78a9..54f2a69295 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts @@ -1,9 +1,8 @@ import { inject, injectable } from 'inversify'; -import { CancellationToken, Event, EventEmitter, l10n } from 'vscode'; +import { CancellationToken, Event, EventEmitter } from 'vscode'; import { IDisposableRegistry, Resource } from '../../common/types'; import { EnvironmentVariables } from '../../common/variables/types'; -import { UnsupportedIntegrationError } from '../../errors/unsupportedIntegrationError'; import { logger } from '../../logging'; import { IIntegrationStorage, @@ -11,13 +10,7 @@ import { IPlatformNotebookEditorProvider, IPlatformDeepnoteNotebookManager } from './types'; -import { - LegacyIntegrationConfig, - LegacyIntegrationType, - SnowflakeAuthMethods, - LegacyDuckDBIntegrationConfig, - DATAFRAME_SQL_INTEGRATION_ID -} from './integrationTypes'; +import { DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes'; import { getEnvironmentVariablesForIntegrations } from '@deepnote/database-integrations'; /** diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts deleted file mode 100644 index 6dc5aa39cc..0000000000 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts +++ /dev/null @@ -1,628 +0,0 @@ -import { assert } from 'chai'; -import { instance, mock, when } from 'ts-mockito'; -import { CancellationTokenSource, EventEmitter, NotebookDocument, Uri } from 'vscode'; - -import { IDisposableRegistry } from '../../common/types'; -import { IntegrationStorage } from './integrationStorage'; -import { SqlIntegrationEnvironmentVariablesProvider } from './sqlIntegrationEnvironmentVariablesProvider'; -import { - LegacyIntegrationType, - LegacyPostgresIntegrationConfig, - LegacyBigQueryIntegrationConfig, - LegacySnowflakeIntegrationConfig, - SnowflakeAuthMethods -} from './integrationTypes'; -import { IPlatformNotebookEditorProvider, IPlatformDeepnoteNotebookManager } from './types'; - -const EXPECTED_DATAFRAME_ONLY_ENV_VARS = { - SQL_DEEPNOTE_DATAFRAME_SQL: '{"url":"deepnote+duckdb:///:memory:","params":{},"param_style":"qmark"}' -}; - -// Helper to create a mock project with integrations -function createMockProject(integrations: Array<{ id: string; name: string; type: string }>) { - return { - project: { - id: 'test-project-id', - name: 'Test Project', - integrations - } - }; -} - -// Helper to create a mock notebook document -function createMockNotebook(projectId: string): NotebookDocument { - return { - metadata: { - deepnoteProjectId: projectId, - deepnoteNotebookId: 'test-notebook-id' - } - } as unknown as NotebookDocument; -} - -suite('SqlIntegrationEnvironmentVariablesProvider', () => { - let provider: SqlIntegrationEnvironmentVariablesProvider; - let integrationStorage: IntegrationStorage; - let notebookEditorProvider: IPlatformNotebookEditorProvider; - let notebookManager: IPlatformDeepnoteNotebookManager; - let disposables: IDisposableRegistry; - - setup(() => { - disposables = []; - integrationStorage = mock(IntegrationStorage); - notebookEditorProvider = mock(); - notebookManager = mock(); - - when(integrationStorage.onDidChangeIntegrations).thenReturn(new EventEmitter().event); - - provider = new SqlIntegrationEnvironmentVariablesProvider( - instance(integrationStorage), - instance(notebookEditorProvider), - instance(notebookManager), - disposables - ); - }); - - teardown(() => { - disposables.forEach((d) => d.dispose()); - }); - - test('Returns empty object when resource is undefined', async () => { - const envVars = await provider.getEnvironmentVariables(undefined); - assert.deepStrictEqual(envVars, {}); - }); - - test('Returns empty object when no notebook is found', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(undefined); - - const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, {}); - }); - - test('Returns empty object when notebook has no project ID', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const notebook = { metadata: {} } as NotebookDocument; - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - - const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, {}); - }); - - test('Returns empty object when project is not found', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const notebook = createMockNotebook('test-project-id'); - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(undefined); - - const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, {}); - }); - - test('Returns only dataframe integration when no integrations are configured', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - - const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); - }); - - test('Returns environment variable for internal DuckDB integration', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - - const envVars = await provider.getEnvironmentVariables(uri); - - // Check that the environment variable is set for dataframe SQL - assert.property(envVars, 'SQL_DEEPNOTE_DATAFRAME_SQL'); - const credentialsJson = JSON.parse(envVars['SQL_DEEPNOTE_DATAFRAME_SQL']!); - assert.strictEqual(credentialsJson.url, 'deepnote+duckdb:///:memory:'); - assert.deepStrictEqual(credentialsJson.params, {}); - assert.strictEqual(credentialsJson.param_style, 'qmark'); - }); - - test('Returns environment variable for PostgreSQL integration', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'my-postgres-db'; - const config: LegacyPostgresIntegrationConfig = { - id: integrationId, - name: 'My Postgres DB', - type: LegacyIntegrationType.Postgres, - host: 'localhost', - port: 5432, - database: 'mydb', - username: 'user', - password: 'pass', - ssl: true - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([{ id: integrationId, name: 'My Postgres DB', type: 'pgsql' }]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - - // Check that the environment variable is set - assert.property(envVars, 'SQL_MY_POSTGRES_DB'); - const credentialsJson = JSON.parse(envVars['SQL_MY_POSTGRES_DB']!); - assert.strictEqual(credentialsJson.url, 'postgresql://user:pass@localhost:5432/mydb'); - assert.deepStrictEqual(credentialsJson.params, { sslmode: 'require' }); - assert.strictEqual(credentialsJson.param_style, 'format'); - }); - - test('Returns environment variable for BigQuery integration', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'my-bigquery'; - const serviceAccountJson = JSON.stringify({ type: 'service_account', project_id: 'my-project' }); - const config: LegacyBigQueryIntegrationConfig = { - id: integrationId, - name: 'My BigQuery', - type: LegacyIntegrationType.BigQuery, - projectId: 'my-project', - credentials: serviceAccountJson - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([{ id: integrationId, name: 'My BigQuery', type: 'big-query' }]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - - // Check that the environment variable is set - assert.property(envVars, 'SQL_MY_BIGQUERY'); - const credentialsJson = JSON.parse(envVars['SQL_MY_BIGQUERY']!); - assert.strictEqual(credentialsJson.url, 'bigquery://?user_supplied_client=true'); - assert.deepStrictEqual(credentialsJson.params, { - project_id: 'my-project', - credentials: { type: 'service_account', project_id: 'my-project' } - }); - assert.strictEqual(credentialsJson.param_style, 'format'); - }); - - test('Returns environment variables for all configured integrations', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const postgresId = 'my-postgres-db'; - const bigqueryId = 'my-bigquery'; - - const postgresConfig: LegacyPostgresIntegrationConfig = { - id: postgresId, - name: 'My Postgres DB', - type: LegacyIntegrationType.Postgres, - host: 'localhost', - port: 5432, - database: 'mydb', - username: 'user', - password: 'pass' - }; - - const bigqueryConfig: LegacyBigQueryIntegrationConfig = { - id: bigqueryId, - name: 'My BigQuery', - type: LegacyIntegrationType.BigQuery, - projectId: 'my-project', - credentials: JSON.stringify({ type: 'service_account' }) - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([ - { id: postgresId, name: 'My Postgres DB', type: 'pgsql' }, - { id: bigqueryId, name: 'My BigQuery', type: 'big-query' } - ]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(postgresId)).thenResolve(postgresConfig); - when(integrationStorage.getIntegrationConfig(bigqueryId)).thenResolve(bigqueryConfig); - - const envVars = await provider.getEnvironmentVariables(uri); - - // Should have two environment variables apart from the internal DuckDB integration - assert.property(envVars, 'SQL_MY_POSTGRES_DB'); - assert.property(envVars, 'SQL_MY_BIGQUERY'); - assert.strictEqual(Object.keys(envVars).length, 3); - }); - - test('Properly encodes special characters in PostgreSQL credentials', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'special-chars-db'; - const config: LegacyPostgresIntegrationConfig = { - id: integrationId, - name: 'Special Chars DB', - type: LegacyIntegrationType.Postgres, - host: 'db.example.com', - port: 5432, - database: 'my@db:name', - username: 'user@domain', - password: 'pa:ss@word!#$%', - ssl: false - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([{ id: integrationId, name: 'Special Chars DB', type: 'pgsql' }]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - - // Check that the environment variable is set - assert.property(envVars, 'SQL_SPECIAL_CHARS_DB'); - const credentialsJson = JSON.parse(envVars['SQL_SPECIAL_CHARS_DB']!); - - // Verify that special characters are properly URL-encoded - assert.strictEqual( - credentialsJson.url, - 'postgresql://user%40domain:pa%3Ass%40word!%23%24%25@db.example.com:5432/my%40db%3Aname' - ); - assert.deepStrictEqual(credentialsJson.params, {}); - assert.strictEqual(credentialsJson.param_style, 'format'); - }); - - test('Normalizes integration ID with spaces and mixed case for env var name', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'My Production DB'; - const config: LegacyPostgresIntegrationConfig = { - id: integrationId, - name: 'Production Database', - type: LegacyIntegrationType.Postgres, - host: 'prod.example.com', - port: 5432, - database: 'proddb', - username: 'admin', - password: 'secret', - ssl: true - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([{ id: integrationId, name: 'Production Database', type: 'pgsql' }]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - - // Check that the environment variable name is properly normalized - // Spaces should be converted to underscores and uppercased - assert.property(envVars, 'SQL_MY_PRODUCTION_DB'); - const credentialsJson = JSON.parse(envVars['SQL_MY_PRODUCTION_DB']!); - assert.strictEqual(credentialsJson.url, 'postgresql://admin:secret@prod.example.com:5432/proddb'); - assert.deepStrictEqual(credentialsJson.params, { sslmode: 'require' }); - assert.strictEqual(credentialsJson.param_style, 'format'); - }); - - test('Normalizes integration ID with special characters for env var name', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'my-db@2024!'; - const config: LegacyPostgresIntegrationConfig = { - id: integrationId, - name: 'Test DB', - type: LegacyIntegrationType.Postgres, - host: 'localhost', - port: 5432, - database: 'testdb', - username: 'user', - password: 'pass', - ssl: false - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([{ id: integrationId, name: 'Test DB', type: 'pgsql' }]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - - // Check that special characters in integration ID are normalized for env var name - // Non-alphanumeric characters should be converted to underscores - assert.property(envVars, 'SQL_MY_DB_2024_'); - const credentialsJson = JSON.parse(envVars['SQL_MY_DB_2024_']!); - assert.strictEqual(credentialsJson.url, 'postgresql://user:pass@localhost:5432/testdb'); - }); - - test('Honors CancellationToken (returns empty when cancelled early)', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const cts = new CancellationTokenSource(); - cts.cancel(); - const envVars = await provider.getEnvironmentVariables(uri, cts.token); - assert.deepStrictEqual(envVars, {}); - }); - - suite('Snowflake Integration', () => { - test('Returns environment variable for Snowflake with PASSWORD auth', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'my-snowflake'; - const config: LegacySnowflakeIntegrationConfig = { - id: integrationId, - name: 'My Snowflake', - type: LegacyIntegrationType.Snowflake, - account: 'myorg-myaccount', - warehouse: 'COMPUTE_WH', - database: 'MYDB', - role: 'ANALYST', - authMethod: SnowflakeAuthMethods.PASSWORD, - username: 'john.doe', - password: 'secret123' - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([{ id: integrationId, name: 'My Snowflake', type: 'snowflake' }]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - - assert.property(envVars, 'SQL_MY_SNOWFLAKE'); - const credentialsJson = JSON.parse(envVars['SQL_MY_SNOWFLAKE']!); - assert.strictEqual( - credentialsJson.url, - 'snowflake://john.doe:secret123@myorg-myaccount/MYDB?warehouse=COMPUTE_WH&role=ANALYST&application=Deepnote' - ); - assert.deepStrictEqual(credentialsJson.params, {}); - assert.strictEqual(credentialsJson.param_style, 'pyformat'); - }); - - test('Returns environment variable for Snowflake with legacy null auth (username+password)', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'legacy-snowflake'; - const config: LegacySnowflakeIntegrationConfig = { - id: integrationId, - name: 'Legacy Snowflake', - type: LegacyIntegrationType.Snowflake, - account: 'legacy-account', - warehouse: 'WH', - database: 'DB', - authMethod: null, - username: 'user', - password: 'pass' - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([{ id: integrationId, name: 'Legacy Snowflake', type: 'snowflake' }]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - - assert.property(envVars, 'SQL_LEGACY_SNOWFLAKE'); - const credentialsJson = JSON.parse(envVars['SQL_LEGACY_SNOWFLAKE']!); - assert.strictEqual( - credentialsJson.url, - 'snowflake://user:pass@legacy-account/DB?warehouse=WH&application=Deepnote' - ); - assert.deepStrictEqual(credentialsJson.params, {}); - }); - - test('Returns environment variable for Snowflake with SERVICE_ACCOUNT_KEY_PAIR auth', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'snowflake-keypair'; - const privateKey = - '-----BEGIN ' + 'PRIVATE KEY-----\nfakekey-MIIEvQIBADANBg...\n-----END ' + 'PRIVATE KEY-----'; - const config: LegacySnowflakeIntegrationConfig = { - id: integrationId, - name: 'Snowflake KeyPair', - type: LegacyIntegrationType.Snowflake, - account: 'keypair-account', - warehouse: 'ETL_WH', - database: 'PROD_DB', - role: 'ETL_ROLE', - authMethod: SnowflakeAuthMethods.SERVICE_ACCOUNT_KEY_PAIR, - username: 'service_account', - privateKey: privateKey, - privateKeyPassphrase: 'passphrase123' - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([{ id: integrationId, name: 'Snowflake KeyPair', type: 'snowflake' }]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - - assert.property(envVars, 'SQL_SNOWFLAKE_KEYPAIR'); - const credentialsJson = JSON.parse(envVars['SQL_SNOWFLAKE_KEYPAIR']!); - assert.strictEqual( - credentialsJson.url, - 'snowflake://service_account@keypair-account/PROD_DB?warehouse=ETL_WH&role=ETL_ROLE&authenticator=snowflake_jwt&application=Deepnote' - ); - assert.deepStrictEqual(credentialsJson.params, { - snowflake_private_key: Buffer.from(privateKey).toString('base64'), - snowflake_private_key_passphrase: 'passphrase123' - }); - assert.strictEqual(credentialsJson.param_style, 'pyformat'); - }); - - test('Returns environment variable for Snowflake with SERVICE_ACCOUNT_KEY_PAIR auth without passphrase', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'snowflake-keypair-no-pass'; - const privateKey = - '-----BEGIN ' + 'PRIVATE KEY-----\nfakekey-MIIEvQIBADANBg...\n-----END ' + 'PRIVATE KEY-----'; - const config: LegacySnowflakeIntegrationConfig = { - id: integrationId, - name: 'Snowflake KeyPair No Pass', - type: LegacyIntegrationType.Snowflake, - account: 'account123', - warehouse: 'WH', - database: 'DB', - authMethod: SnowflakeAuthMethods.SERVICE_ACCOUNT_KEY_PAIR, - username: 'svc_user', - privateKey: privateKey - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([ - { id: integrationId, name: 'Snowflake KeyPair No Pass', type: 'snowflake' } - ]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - - assert.property(envVars, 'SQL_SNOWFLAKE_KEYPAIR_NO_PASS'); - const credentialsJson = JSON.parse(envVars['SQL_SNOWFLAKE_KEYPAIR_NO_PASS']!); - assert.strictEqual( - credentialsJson.url, - 'snowflake://svc_user@account123/DB?warehouse=WH&authenticator=snowflake_jwt&application=Deepnote' - ); - assert.deepStrictEqual(credentialsJson.params, { - snowflake_private_key: Buffer.from(privateKey).toString('base64') - }); - }); - - test('Properly encodes special characters in Snowflake credentials', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'snowflake-special'; - const config: LegacySnowflakeIntegrationConfig = { - id: integrationId, - name: 'Snowflake Special', - type: LegacyIntegrationType.Snowflake, - account: 'my-org.account', - warehouse: 'WH@2024', - database: 'DB:TEST', - role: 'ROLE#1', - authMethod: SnowflakeAuthMethods.PASSWORD, - username: 'user@domain.com', - password: 'p@ss:word!#$%' - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([{ id: integrationId, name: 'Snowflake Special', type: 'snowflake' }]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - - assert.property(envVars, 'SQL_SNOWFLAKE_SPECIAL'); - const credentialsJson = JSON.parse(envVars['SQL_SNOWFLAKE_SPECIAL']!); - // Verify URL encoding of special characters - assert.strictEqual( - credentialsJson.url, - 'snowflake://user%40domain.com:p%40ss%3Aword!%23%24%25@my-org.account/DB%3ATEST?warehouse=WH%402024&role=ROLE%231&application=Deepnote' - ); - }); - - test('Handles Snowflake with minimal optional fields', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'snowflake-minimal'; - const config: LegacySnowflakeIntegrationConfig = { - id: integrationId, - name: 'Snowflake Minimal', - type: LegacyIntegrationType.Snowflake, - account: 'minimal-account', - authMethod: SnowflakeAuthMethods.PASSWORD, - username: 'user', - password: 'pass' - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([{ id: integrationId, name: 'Snowflake Minimal', type: 'snowflake' }]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - - assert.property(envVars, 'SQL_SNOWFLAKE_MINIMAL'); - const credentialsJson = JSON.parse(envVars['SQL_SNOWFLAKE_MINIMAL']!); - // Should not include warehouse, database, or role in URL when not provided - assert.strictEqual(credentialsJson.url, 'snowflake://user:pass@minimal-account?application=Deepnote'); - assert.strictEqual(credentialsJson.param_style, 'pyformat'); - }); - - test('Skips unsupported Snowflake auth method (OKTA)', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'snowflake-okta'; - const config: LegacySnowflakeIntegrationConfig = { - id: integrationId, - name: 'Snowflake OKTA', - type: LegacyIntegrationType.Snowflake, - account: 'okta-account', - authMethod: SnowflakeAuthMethods.OKTA - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([{ id: integrationId, name: 'Snowflake OKTA', type: 'snowflake' }]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - // Should return only dataframe integration when unsupported auth method is encountered - const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); - }); - - test('Skips unsupported Snowflake auth method (AZURE_AD)', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'snowflake-azure'; - const config: LegacySnowflakeIntegrationConfig = { - id: integrationId, - name: 'Snowflake Azure', - type: LegacyIntegrationType.Snowflake, - account: 'azure-account', - authMethod: SnowflakeAuthMethods.AZURE_AD - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([{ id: integrationId, name: 'Snowflake Azure', type: 'snowflake' }]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); - }); - - test('Skips unsupported Snowflake auth method (KEY_PAIR)', async () => { - const uri = Uri.file('/test/notebook.deepnote'); - const integrationId = 'snowflake-keypair-user'; - const config: LegacySnowflakeIntegrationConfig = { - id: integrationId, - name: 'Snowflake KeyPair User', - type: LegacyIntegrationType.Snowflake, - account: 'keypair-user-account', - authMethod: SnowflakeAuthMethods.KEY_PAIR - }; - - const notebook = createMockNotebook('test-project-id'); - const project = createMockProject([ - { id: integrationId, name: 'Snowflake KeyPair User', type: 'snowflake' } - ]); - - when(notebookEditorProvider.findAssociatedNotebookDocument(uri)).thenReturn(notebook); - when(notebookManager.getOriginalProject('test-project-id')).thenReturn(project as any); - when(integrationStorage.getIntegrationConfig(integrationId)).thenResolve(config); - - const envVars = await provider.getEnvironmentVariables(uri); - assert.deepStrictEqual(envVars, EXPECTED_DATAFRAME_ONLY_ENV_VARS); - }); - }); -}); diff --git a/src/webviews/webview-side/integrations/BigQueryForm.tsx b/src/webviews/webview-side/integrations/BigQueryForm.tsx index 60710ca912..54cf987898 100644 --- a/src/webviews/webview-side/integrations/BigQueryForm.tsx +++ b/src/webviews/webview-side/integrations/BigQueryForm.tsx @@ -1,41 +1,67 @@ import * as React from 'react'; import { format, getLocString } from '../react-common/locReactSide'; -import { BigQueryIntegrationConfig } from './types'; +import { BigQueryAuthMethods, DatabaseIntegrationConfig } from '@deepnote/database-integrations'; + +type BigQueryConfig = Extract; + +function createEmptyBigQueryConfig(params: { id: string; name?: string }): BigQueryConfig { + const unnamedIntegration = getLocString('integrationsUnnamedIntegration', 'Unnamed Integration ({0})'); + + return { + id: params.id, + name: (params.name || format(unnamedIntegration, params.id)).trim(), + type: 'big-query', + metadata: { + authMethod: BigQueryAuthMethods.ServiceAccount, + service_account: '' + } + }; +} export interface IBigQueryFormProps { integrationId: string; - existingConfig: BigQueryIntegrationConfig | null; - integrationName?: string; - onSave: (config: BigQueryIntegrationConfig) => void; + existingConfig: BigQueryConfig | null; + defaultName?: string; + onSave: (config: BigQueryConfig) => void; onCancel: () => void; } export const BigQueryForm: React.FC = ({ integrationId, existingConfig, - integrationName, + defaultName, onSave, onCancel }) => { - const [name, setName] = React.useState(existingConfig?.name || integrationName || ''); - const [projectId, setProjectId] = React.useState(existingConfig?.projectId || ''); - const [credentials, setCredentials] = React.useState(existingConfig?.credentials || ''); + const [pendingConfig, setPendingConfig] = React.useState( + existingConfig && existingConfig.metadata.authMethod === BigQueryAuthMethods.ServiceAccount + ? structuredClone(existingConfig) + : createEmptyBigQueryConfig({ id: integrationId, name: defaultName }) + ); + const [credentialsError, setCredentialsError] = React.useState(null); - // Update form fields when existingConfig or integrationName changes React.useEffect(() => { - if (existingConfig) { - setName(existingConfig.name || ''); - setProjectId(existingConfig.projectId || ''); - setCredentials(existingConfig.credentials || ''); - setCredentialsError(null); - } else { - setName(integrationName || ''); - setProjectId(''); - setCredentials(''); - setCredentialsError(null); - } - }, [existingConfig, integrationName]); + setPendingConfig( + existingConfig && existingConfig.metadata.authMethod === BigQueryAuthMethods.ServiceAccount + ? structuredClone(existingConfig) + : createEmptyBigQueryConfig({ id: integrationId, name: defaultName }) + ); + setCredentialsError(null); + }, [existingConfig, integrationId, defaultName]); + + // Extract service account value with proper type narrowing + const serviceAccountValue = + pendingConfig.metadata.authMethod === BigQueryAuthMethods.ServiceAccount + ? pendingConfig.metadata.service_account + : ''; + + const handleNameChange = (e: React.ChangeEvent) => { + setPendingConfig((prev) => ({ + ...prev, + name: e.target.value + })); + }; const validateCredentials = (value: string): boolean => { if (!value.trim()) { @@ -57,31 +83,32 @@ export const BigQueryForm: React.FC = ({ const handleCredentialsChange = (e: React.ChangeEvent) => { const value = e.target.value; - setCredentials(value); + + setPendingConfig((prev) => { + if (prev.metadata.authMethod === BigQueryAuthMethods.ServiceAccount) { + return { + ...prev, + metadata: { + ...prev.metadata, + service_account: value + } + }; + } + return prev; + }); + validateCredentials(value); }; const handleSubmit = (e: React.FormEvent) => { e.preventDefault(); - const trimmedCredentials = credentials.trim(); - // Validate credentials before submitting - if (!validateCredentials(trimmedCredentials)) { + if (!validateCredentials(serviceAccountValue)) { return; } - const unnamedIntegration = getLocString('integrationsUnnamedIntegration', 'Unnamed Integration ({0})'); - - const config: BigQueryIntegrationConfig = { - id: integrationId, - name: (name || format(unnamedIntegration, integrationId)).trim(), - type: 'bigquery', - projectId: projectId.trim(), - credentials: trimmedCredentials - }; - - onSave(config); + onSave(pendingConfig); }; return ( @@ -91,29 +118,13 @@ export const BigQueryForm: React.FC = ({ setName(e.target.value)} + value={pendingConfig.name} + onChange={handleNameChange} placeholder={getLocString('integrationsBigQueryNamePlaceholder', 'My BigQuery Project')} autoComplete="off" /> -
- - setProjectId(e.target.value)} - placeholder={getLocString('integrationsBigQueryProjectIdPlaceholder', 'my-project-id')} - autoComplete="off" - required - /> -
-