diff --git a/INTEGRATIONS_CREDENTIALS.md b/INTEGRATIONS_CREDENTIALS.md new file mode 100644 index 0000000000..108e33a5c0 --- /dev/null +++ b/INTEGRATIONS_CREDENTIALS.md @@ -0,0 +1,414 @@ +# Deepnote Integrations & Credentials System + +## Overview + +The integrations system enables Deepnote notebooks to connect to external data sources (PostgreSQL, BigQuery, etc.) by securely managing credentials and exposing them to SQL blocks. The system handles: + +1. **Credential Storage**: Secure storage using VSCode's SecretStorage API +2. **Integration Detection**: Automatic discovery of integrations used in notebooks +3. **UI Management**: Webview-based configuration interface +4. **Kernel Integration**: Injection of credentials into Jupyter kernel environment +5. **Toolkit Exposure**: Making credentials available to `deepnote-toolkit` for SQL execution + +## Architecture + +### Core Components + +#### 1. **Integration Storage** (`integrationStorage.ts`) + +Manages persistent storage of integration configurations using VSCode's encrypted SecretStorage API. + +**Key Features:** + +- Uses VSCode's `SecretStorage` API for secure credential storage +- Storage is scoped to the user's machine (shared across all Deepnote projects) +- In-memory caching for performance +- Event-driven updates via `onDidChangeIntegrations` event +- Index-based storage for efficient retrieval + +**Storage Format:** + +- Each integration config is stored as JSON under key: `deepnote-integrations.{integrationId}` +- An index is maintained at key: `deepnote-integrations.index` containing all integration IDs + +**Key Methods:** + +- `getAll()`: Retrieve all stored integration configurations +- `getIntegrationConfig(integrationId)`: Get a specific integration by ID +- `getProjectIntegrationConfig(projectId, integrationId)`: Get the effective project-scoped config +- `save(config)`: Save or update an integration configuration +- `delete(integrationId)`: Remove an integration configuration +- `exists(integrationId)`: Check if an integration is configured + +**Integration Config Types:** + +```typescript +// PostgreSQL +{ + id: string; + name: string; + type: 'postgres'; + host: string; + port: number; + database: string; + username: string; + password: string; + ssl?: boolean; +} + +// BigQuery +{ + id: string; + name: string; + type: 'bigquery'; + projectId: string; + credentials: string; // JSON string of service account credentials +} +``` + +#### 2. **Integration Detector** (`integrationDetector.ts`) + +Scans Deepnote projects to discover which integrations are used in SQL blocks. + +**Detection Process:** + +1. Retrieves the Deepnote project from `IDeepnoteNotebookManager` +2. Scans all notebooks in the project +3. Examines each code block for `metadata.sql_integration_id` +4. Checks if each integration is configured (has credentials) +5. Returns a map of integration IDs to their status + +**Integration Status:** + +- `Connected`: Integration has valid credentials stored +- `Disconnected`: Integration is used but not configured +- `Error`: Integration configuration is invalid + +**Special Cases:** + +- Excludes `deepnote-dataframe-sql` (internal DuckDB integration) +- Only processes code blocks with SQL integration metadata + +#### 3. **Integration Manager** (`integrationManager.ts`) + +Orchestrates the integration management UI and commands. + +**Responsibilities:** + +- Registers the `deepnote.manageIntegrations` command +- Updates VSCode context keys for UI visibility: + - `deepnote.hasIntegrations`: True if any integrations are detected + - `deepnote.hasUnconfiguredIntegrations`: True if any integrations lack credentials +- Handles notebook selection changes +- Opens the integration webview with detected integrations + +**Command Flow:** + +1. User triggers command (from command palette or SQL cell status bar) +2. Manager detects integrations in the active notebook +3. Manager opens webview with integration list +4. Optionally pre-selects a specific integration for configuration + +#### 4. **Integration Webview** (`integrationWebview.ts`) + +Provides the webview-based UI for managing integration credentials. + +**Features:** + +- Persistent webview panel (survives defocus) +- Real-time integration status updates +- Configuration forms for each integration type +- Delete/reset functionality + +**Message Protocol:** + +Extension → Webview: + +```typescript +// Update integration list +{ type: 'update', integrations: IntegrationWithStatus[] } + +// Show configuration form +{ type: 'showForm', integrationId: string, config: IntegrationConfig | null } + +// Status messages +{ type: 'success' | 'error', message: string } +``` + +Webview → Extension: + +```typescript +// Save configuration +{ type: 'save', integrationId: string, config: IntegrationConfig } + +// Delete configuration +{ type: 'delete', integrationId: string } + +// Request configuration form +{ type: 'configure', integrationId: string } +``` + +### UI Components (React) + +#### 5. **Integration Panel** (`IntegrationPanel.tsx`) + +Main React component that manages the webview UI state. + +**State Management:** + +- `integrations`: List of detected integrations with status +- `selectedIntegrationId`: Currently selected integration for configuration +- `selectedConfig`: Existing configuration being edited +- `message`: Success/error messages +- `confirmDelete`: Confirmation state for deletion + +**User Flows:** + +**Configure Integration:** + +1. User clicks "Configure" button +2. Panel shows configuration form overlay +3. User enters credentials +4. Panel sends save message to extension +5. Extension stores credentials and updates status +6. Panel shows success message and refreshes list + +**Delete Integration:** + +1. User clicks "Reset" button +2. Panel shows confirmation prompt (5 seconds) +3. User clicks again to confirm +4. Panel sends delete message to extension +5. Extension removes credentials +6. Panel updates status to "Disconnected" + +#### 6. **Configuration Forms** (`PostgresForm.tsx`, `BigQueryForm.tsx`) + +Type-specific forms for entering integration credentials. + +**PostgreSQL Form Fields:** + +- Name (display name) +- Host +- Port (default: 5432) +- Database +- Username +- Password +- SSL (checkbox) + +**BigQuery Form Fields:** + +- Name (display name) +- Project ID +- Service Account Credentials (JSON textarea) + +**Validation:** + +- All fields are required +- BigQuery credentials must be valid JSON +- Port must be a valid number + +### Kernel Integration + +#### 7. **SQL Integration Environment Variables Provider** (`sqlIntegrationEnvironmentVariablesProvider.ts`) + +Provides environment variables containing integration credentials for the Jupyter kernel. + +**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 + +**Environment Variable Format:** + +Variable name: `SQL_{INTEGRATION_ID}` (uppercased, special chars replaced with `_`) + +Example: Integration ID `my-postgres-db` → Environment variable `SQL_MY_POSTGRES_DB` + +**Credential JSON Format:** + +PostgreSQL: + +```json +{ + "url": "postgresql://username:password@host:port/database", + "params": { "sslmode": "require" }, + "param_style": "format" +} +``` + +BigQuery: + +```json +{ + "url": "bigquery://?user_supplied_client=true", + "params": { + "project_id": "my-project", + "credentials": { + /* service account JSON */ + } + }, + "param_style": "format" +} +``` + +**Integration Points:** + +- Registered as an environment variable provider in the kernel environment service +- Called when starting a Jupyter kernel for a Deepnote notebook +- Environment variables are passed to the kernel process at startup + +#### 8. **SQL Integration Startup Code Provider** (`sqlIntegrationStartupCodeProvider.ts`) + +Injects Python code into the kernel at startup to set environment variables. + +**Why This Is Needed:** +Jupyter doesn't automatically pass all environment variables from the server process to the kernel process. This provider ensures credentials are available in the kernel's `os.environ`. + +**Generated Code:** + +```python +try: + import os + # [SQL Integration] Setting N SQL integration env vars... + os.environ['SQL_MY_POSTGRES_DB'] = '{"url":"postgresql://...","params":{},"param_style":"format"}' + os.environ['SQL_MY_BIGQUERY'] = '{"url":"bigquery://...","params":{...},"param_style":"format"}' + # [SQL Integration] Successfully set N SQL integration env vars +except Exception as e: + import traceback + print(f"[SQL Integration] ERROR: Failed to set SQL integration env vars: {e}") + traceback.print_exc() +``` + +**Execution:** + +- Registered with `IStartupCodeProviders` for `JupyterNotebookView` +- Runs automatically when a Python kernel starts for a Deepnote notebook +- Priority: `StartupCodePriority.Base` (runs early) +- Only runs for Python kernels on Deepnote notebooks + +### Toolkit Integration + +#### 9. **How Credentials Are Exposed to deepnote-toolkit** + +The `deepnote-toolkit` Python package reads credentials from environment variables to execute SQL blocks. + +**Flow:** + +1. Extension detects SQL blocks in notebook +2. Extension retrieves credentials from secure storage +3. Extension converts credentials to JSON format +4. Extension injects credentials as environment variables (two methods): + - **Server Process**: Via `SqlIntegrationEnvironmentVariablesProvider` when starting Jupyter server + - **Kernel Process**: Via `SqlIntegrationStartupCodeProvider` when starting Python kernel +5. `deepnote-toolkit` reads environment variables when executing SQL blocks +6. Toolkit creates database connections using the credentials +7. Toolkit executes SQL queries and returns results + +**Environment Variable Lookup:** +When a SQL block with `sql_integration_id: "my-postgres-db"` is executed: + +1. Toolkit looks for environment variable `SQL_MY_POSTGRES_DB` +2. Toolkit parses the JSON value +3. Toolkit creates a SQLAlchemy connection using the `url` and `params` +4. Toolkit executes the SQL query +5. Toolkit returns results as a pandas DataFrame + +## Data Flow + +### Configuration Flow + +```text +User → IntegrationPanel (UI) + → vscodeApi.postMessage({ type: 'save', config }) + → IntegrationWebviewProvider.onMessage() + → IntegrationStorage.save(config) + → EncryptedStorage.store() [VSCode SecretStorage API] + → IntegrationStorage fires onDidChangeIntegrations event + → SqlIntegrationEnvironmentVariablesProvider fires onDidChangeEnvironmentVariables event +``` + +### Execution Flow + +```text +User executes SQL cell + → Kernel startup triggered + → SqlIntegrationEnvironmentVariablesProvider.getEnvironmentVariables() + → Scans notebook for SQL cells + → Retrieves credentials from IntegrationStorage + → Converts to JSON format + → Returns environment variables + → Environment variables passed to Jupyter server process + → SqlIntegrationStartupCodeProvider.getCode() + → Generates Python code to set os.environ + → Startup code executed in kernel + → deepnote-toolkit reads os.environ['SQL_*'] + → Toolkit executes SQL query + → Results returned to notebook +``` + +## Security Considerations + +1. **Encrypted Storage**: All credentials are stored using VSCode's SecretStorage API, which uses the OS keychain +2. **No Plaintext**: Credentials are never written to disk in plaintext +3. **Scoped Access**: Storage is scoped to the VSCode extension +4. **Environment Isolation**: Each notebook gets only the credentials it needs +5. **No Logging**: Credential values are never logged; only non-sensitive metadata (key names, counts) is logged + +## Adding New Integration Types + +To add a new integration type (e.g., MySQL, Snowflake): + +1. **Add type to `integrationTypes.ts`**: + + ```typescript + export enum IntegrationType { + Postgres = 'postgres', + BigQuery = 'bigquery', + MySQL = 'mysql' // New type + } + + export interface MySQLIntegrationConfig extends BaseIntegrationConfig { + type: IntegrationType.MySQL; + host: string; + port: number; + database: string; + username: string; + password: string; + } + + export type IntegrationConfig = PostgresIntegrationConfig | BigQueryIntegrationConfig | MySQLIntegrationConfig; + ``` + +2. **Add conversion logic in `sqlIntegrationEnvironmentVariablesProvider.ts`**: + + ```typescript + case IntegrationType.MySQL: { + const url = `mysql://${config.username}:${config.password}@${config.host}:${config.port}/${config.database}`; + return JSON.stringify({ url, params: {}, param_style: 'format' }); + } + ``` + +3. **Create UI form component** (`MySQLForm.tsx`) + +4. **Update `ConfigurationForm.tsx`** to render the new form + +5. **Update webview types** (`src/webviews/webview-side/integrations/types.ts`) + +6. **Add localization strings** for the new integration type + +## Testing + +Unit tests are located in: + +- `sqlIntegrationEnvironmentVariablesProvider.unit.test.ts` + +Tests cover: + +- Environment variable generation for each integration type +- Multiple integrations in a single notebook +- Missing credentials handling +- Integration ID to environment variable name conversion +- JSON format validation diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index ec6b242d3e..bee7c4f131 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -import { inject, injectable, named } from 'inversify'; +import { inject, injectable, named, optional } from 'inversify'; import { CancellationToken, Uri } from 'vscode'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { IDeepnoteServerStarter, IDeepnoteToolkitInstaller, DeepnoteServerInfo, DEEPNOTE_DEFAULT_PORT } from './types'; @@ -12,6 +12,7 @@ import { STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants'; import { sleep } from '../../platform/common/utils/async'; import { Cancellation, raceCancellationError } from '../../platform/common/cancellation'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; import getPort from 'get-port'; import * as fs from 'fs-extra'; import * as os from 'os'; @@ -47,7 +48,10 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension @inject(IDeepnoteToolkitInstaller) private readonly toolkitInstaller: IDeepnoteToolkitInstaller, @inject(IOutputChannel) @named(STANDARD_OUTPUT_CHANNEL) private readonly outputChannel: IOutputChannel, @inject(IHttpClient) private readonly httpClient: IHttpClient, - @inject(IAsyncDisposableRegistry) asyncRegistry: IAsyncDisposableRegistry + @inject(IAsyncDisposableRegistry) asyncRegistry: IAsyncDisposableRegistry, + @inject(ISqlIntegrationEnvVarsProvider) + @optional() + private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider ) { // Register for disposal when the extension deactivates asyncRegistry.push(this); @@ -149,10 +153,28 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension // Detached mode ensures no requests are made to the backend (directly, or via proxy) // as there is no backend running in the extension, therefore: - // 1. integration environment variables won't work / be injected + // 1. integration environment variables are injected here instead // 2. post start hooks won't work / are not executed env.DEEPNOTE_RUNTIME__RUNNING_IN_DETACHED_MODE = 'true'; + // Inject SQL integration environment variables + if (this.sqlIntegrationEnvVars) { + logger.debug(`DeepnoteServerStarter: Injecting SQL integration env vars for ${deepnoteFileUri.toString()}`); + try { + const sqlEnvVars = await this.sqlIntegrationEnvVars.getEnvironmentVariables(deepnoteFileUri, token); + if (sqlEnvVars && Object.keys(sqlEnvVars).length > 0) { + logger.debug(`DeepnoteServerStarter: Injecting ${Object.keys(sqlEnvVars).length} SQL env vars`); + Object.assign(env, sqlEnvVars); + } else { + logger.debug('DeepnoteServerStarter: No SQL integration env vars to inject'); + } + } catch (error) { + logger.error('DeepnoteServerStarter: Failed to get SQL integration env vars', error.message); + } + } else { + logger.debug('DeepnoteServerStarter: SqlIntegrationEnvironmentVariablesProvider not available'); + } + // Remove PYTHONHOME if it exists (can interfere with venv) delete env.PYTHONHOME; diff --git a/src/kernels/helpers.unit.test.ts b/src/kernels/helpers.unit.test.ts index e84f2fc5b1..0d6d163881 100644 --- a/src/kernels/helpers.unit.test.ts +++ b/src/kernels/helpers.unit.test.ts @@ -646,4 +646,315 @@ suite('Kernel Connection Helpers', () => { assert.strictEqual(name, '.env (Python 9.8.7)'); }); }); + + suite('executeSilently', () => { + interface MockKernelOptions { + status: 'ok' | 'error'; + messages?: Array<{ + msg_type: 'stream' | 'error' | 'display_data' | 'execute_result'; + content: any; + }>; + errorContent?: { + ename: string; + evalue: string; + traceback: string[]; + }; + } + + function createMockKernel(options: MockKernelOptions) { + return { + requestExecute: () => { + let resolvePromise: (value: any) => void; + + // Create a promise that will be resolved after IOPub messages are dispatched + const donePromise = new Promise((resolve) => { + resolvePromise = resolve; + }); + + return { + done: donePromise, + set onIOPub(cb: (msg: any) => void) { + // Invoke IOPub callback synchronously with all messages + if (options.messages && options.messages.length > 0) { + options.messages.forEach((msg) => { + cb({ + header: { msg_type: msg.msg_type }, + content: msg.content + }); + }); + } + // Resolve the done promise after messages are dispatched + resolvePromise({ + content: + options.status === 'ok' + ? { status: 'ok' as const } + : { + status: 'error' as const, + ...options.errorContent + } + }); + } + }; + } + }; + } + + test('Returns outputs from kernel execution', async () => { + const mockKernel = createMockKernel({ + status: 'ok', + messages: [ + { + msg_type: 'stream', + content: { + name: 'stdout', + text: 'hello\n' + } + } + ] + }); + + const code = 'print("hello")'; + const { executeSilently } = await import('./helpers'); + const result = await executeSilently(mockKernel as any, code); + + // executeSilently should return outputs array with collected stream output + assert.isArray(result); + assert.equal(result.length, 1); + assert.equal(result[0].output_type, 'stream'); + assert.equal((result[0] as any).name, 'stdout'); + assert.equal((result[0] as any).text, 'hello\n'); + }); + + test('Collects error outputs', async () => { + const mockKernel = createMockKernel({ + status: 'error', + errorContent: { + ename: 'NameError', + evalue: 'name not defined', + traceback: ['Traceback...'] + }, + messages: [ + { + msg_type: 'error', + content: { + ename: 'NameError', + evalue: 'name not defined', + traceback: ['Traceback...'] + } + } + ] + }); + + const code = 'undefined_variable'; + const { executeSilently } = await import('./helpers'); + const result = await executeSilently(mockKernel as any, code); + + assert.isArray(result); + assert.equal(result.length, 1); + assert.equal(result[0].output_type, 'error'); + assert.equal((result[0] as any).ename, 'NameError'); + assert.equal((result[0] as any).evalue, 'name not defined'); + assert.deepStrictEqual((result[0] as any).traceback, ['Traceback...']); + }); + + test('Collects display_data outputs', async () => { + const mockKernel = createMockKernel({ + status: 'ok', + messages: [ + { + msg_type: 'display_data', + content: { + data: { + 'text/plain': 'some data' + }, + metadata: {} + } + } + ] + }); + + const code = 'display("data")'; + const { executeSilently } = await import('./helpers'); + const result = await executeSilently(mockKernel as any, code); + + assert.isArray(result); + assert.equal(result.length, 1); + assert.equal(result[0].output_type, 'display_data'); + assert.deepStrictEqual((result[0] as any).data, { 'text/plain': 'some data' }); + assert.deepStrictEqual((result[0] as any).metadata, {}); + }); + + test('Handles multiple outputs', async () => { + const mockKernel = createMockKernel({ + status: 'ok', + messages: [ + { + msg_type: 'stream', + content: { + name: 'stdout', + text: 'output 1' + } + }, + { + msg_type: 'stream', + content: { + name: 'stdout', + text: 'output 2' + } + } + ] + }); + + const code = 'print("1"); print("2")'; + const { executeSilently } = await import('./helpers'); + const result = await executeSilently(mockKernel as any, code); + + assert.isArray(result); + // Consecutive stream messages with the same name are concatenated + assert.equal(result.length, 1); + assert.equal(result[0].output_type, 'stream'); + assert.equal((result[0] as any).name, 'stdout'); + assert.equal((result[0] as any).text, 'output 1output 2'); + }); + + test('Collects execute_result outputs', async () => { + const mockKernel = createMockKernel({ + status: 'ok', + messages: [ + { + msg_type: 'execute_result', + content: { + data: { + 'text/plain': '42' + }, + metadata: {}, + execution_count: 1 + } + } + ] + }); + + const code = '42'; + const { executeSilently } = await import('./helpers'); + const result = await executeSilently(mockKernel as any, code); + + assert.isArray(result); + assert.equal(result.length, 1); + assert.equal(result[0].output_type, 'execute_result'); + assert.deepStrictEqual((result[0] as any).data, { 'text/plain': '42' }); + assert.deepStrictEqual((result[0] as any).metadata, {}); + assert.equal((result[0] as any).execution_count, 1); + }); + + test('Stream messages with different names produce separate outputs', async () => { + const mockKernel = createMockKernel({ + status: 'ok', + messages: [ + { + msg_type: 'stream', + content: { + name: 'stdout', + text: 'standard output' + } + }, + { + msg_type: 'stream', + content: { + name: 'stderr', + text: 'error output' + } + }, + { + msg_type: 'stream', + content: { + name: 'stdout', + text: ' more stdout' + } + } + ] + }); + + const code = 'print("test")'; + const { executeSilently } = await import('./helpers'); + const result = await executeSilently(mockKernel as any, code); + + assert.isArray(result); + // Should have 3 outputs: stdout, stderr, stdout (not concatenated because stderr is in between) + assert.equal(result.length, 3); + assert.equal(result[0].output_type, 'stream'); + assert.equal((result[0] as any).name, 'stdout'); + assert.equal((result[0] as any).text, 'standard output'); + assert.equal(result[1].output_type, 'stream'); + assert.equal((result[1] as any).name, 'stderr'); + assert.equal((result[1] as any).text, 'error output'); + assert.equal(result[2].output_type, 'stream'); + assert.equal((result[2] as any).name, 'stdout'); + assert.equal((result[2] as any).text, ' more stdout'); + }); + + test('errorOptions with traceErrors logs errors', async () => { + const mockKernel = createMockKernel({ + status: 'error', + errorContent: { + ename: 'ValueError', + evalue: 'invalid value', + traceback: ['Traceback (most recent call last):', ' File "", line 1'] + }, + messages: [ + { + msg_type: 'error', + content: { + ename: 'ValueError', + evalue: 'invalid value', + traceback: ['Traceback (most recent call last):', ' File "", line 1'] + } + } + ] + }); + + const code = 'raise ValueError("invalid value")'; + const { executeSilently } = await import('./helpers'); + const result = await executeSilently(mockKernel as any, code, { + traceErrors: true, + traceErrorsMessage: 'Custom error message' + }); + + assert.isArray(result); + assert.equal(result.length, 1); + assert.equal(result[0].output_type, 'error'); + assert.equal((result[0] as any).ename, 'ValueError'); + }); + + test('errorOptions without traceErrors still collects errors', async () => { + const mockKernel = createMockKernel({ + status: 'error', + errorContent: { + ename: 'RuntimeError', + evalue: 'runtime issue', + traceback: ['Traceback...'] + }, + messages: [ + { + msg_type: 'error', + content: { + ename: 'RuntimeError', + evalue: 'runtime issue', + traceback: ['Traceback...'] + } + } + ] + }); + + const code = 'raise RuntimeError("runtime issue")'; + const { executeSilently } = await import('./helpers'); + const result = await executeSilently(mockKernel as any, code, { + traceErrors: false + }); + + assert.isArray(result); + assert.equal(result.length, 1); + assert.equal(result[0].output_type, 'error'); + assert.equal((result[0] as any).ename, 'RuntimeError'); + }); + }); }); diff --git a/src/kernels/kernel.ts b/src/kernels/kernel.ts index 0d1d1cb4db..6b0bed1a32 100644 --- a/src/kernels/kernel.ts +++ b/src/kernels/kernel.ts @@ -853,10 +853,22 @@ abstract class BaseKernel implements IBaseKernel { // Gather all of the startup code at one time and execute as one cell const startupCode = await this.gatherInternalStartupCode(); - await this.executeSilently(session, startupCode, { + logger.trace(`Executing startup code with ${startupCode.length} lines`); + + const outputs = await this.executeSilently(session, startupCode, { traceErrors: true, traceErrorsMessage: 'Error executing jupyter extension internal startup code' }); + logger.trace(`Startup code execution completed with ${outputs?.length || 0} outputs`); + if (outputs && outputs.length > 0) { + // Avoid logging content; output types only. + logger.trace( + `Startup code produced ${outputs.length} output(s): ${outputs + .map((o) => o.output_type) + .join(', ')}` + ); + } + // Run user specified startup commands await this.executeSilently(session, this.getUserStartupCommands(), { traceErrors: false }); } diff --git a/src/kernels/raw/launcher/kernelEnvVarsService.node.ts b/src/kernels/raw/launcher/kernelEnvVarsService.node.ts index f4629e77e5..51e88d9483 100644 --- a/src/kernels/raw/launcher/kernelEnvVarsService.node.ts +++ b/src/kernels/raw/launcher/kernelEnvVarsService.node.ts @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -import { inject, injectable } from 'inversify'; +import { inject, injectable, optional } from 'inversify'; import { logger } from '../../../platform/logging'; import { getDisplayPath } from '../../../platform/common/platform/fs-paths'; import { IConfigurationService, Resource, type ReadWrite } from '../../../platform/common/types'; @@ -18,6 +18,7 @@ import { IJupyterKernelSpec } from '../../types'; import { CancellationToken, Uri } from 'vscode'; import { PYTHON_LANGUAGE } from '../../../platform/common/constants'; import { trackKernelResourceInformation } from '../../telemetry/helper'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; /** * Class used to fetch environment variables for a kernel. @@ -30,8 +31,15 @@ export class KernelEnvironmentVariablesService { @inject(IEnvironmentVariablesService) private readonly envVarsService: IEnvironmentVariablesService, @inject(ICustomEnvironmentVariablesProvider) private readonly customEnvVars: ICustomEnvironmentVariablesProvider, - @inject(IConfigurationService) private readonly configService: IConfigurationService - ) {} + @inject(IConfigurationService) private readonly configService: IConfigurationService, + @inject(ISqlIntegrationEnvVarsProvider) + @optional() + private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider + ) { + logger.debug( + `KernelEnvironmentVariablesService: Constructor; SQL env provider present=${!!sqlIntegrationEnvVars}` + ); + } /** * Generates the environment variables for the kernel. * @@ -51,6 +59,11 @@ export class KernelEnvironmentVariablesService { kernelSpec: IJupyterKernelSpec, token?: CancellationToken ) { + logger.debug( + `KernelEnvVarsService.getEnvironmentVariables: Called for resource ${ + resource ? getDisplayPath(resource) : 'undefined' + }, sqlIntegrationEnvVars is ${this.sqlIntegrationEnvVars ? 'AVAILABLE' : 'UNDEFINED'}` + ); let kernelEnv = kernelSpec.env && Object.keys(kernelSpec.env).length > 0 ? (Object.assign({}, kernelSpec.env) as ReadWrite) @@ -68,7 +81,7 @@ export class KernelEnvironmentVariablesService { if (token?.isCancellationRequested) { return; } - let [customEnvVars, interpreterEnv] = await Promise.all([ + let [customEnvVars, interpreterEnv, sqlIntegrationEnvVars] = await Promise.all([ this.customEnvVars .getCustomEnvironmentVariables(resource, isPythonKernel ? 'RunPythonCode' : 'RunNonPythonCode', token) .catch(noop), @@ -82,6 +95,22 @@ export class KernelEnvironmentVariablesService { ); return undefined; }) + : undefined, + this.sqlIntegrationEnvVars + ? this.sqlIntegrationEnvVars + .getEnvironmentVariables(resource, token) + .then((vars) => { + if (vars && Object.keys(vars).length > 0) { + logger.debug( + `KernelEnvVarsService: Got ${Object.keys(vars).length} SQL integration env vars` + ); + } + return vars; + }) + .catch((ex) => { + logger.error('Failed to get SQL integration env variables for Kernel', ex); + return undefined; + }) : undefined ]); if (token?.isCancellationRequested) { @@ -111,7 +140,8 @@ export class KernelEnvironmentVariablesService { // Keep a list of the kernelSpec variables that need to be substituted. const kernelSpecVariablesRequiringSubstitution: Record = {}; for (const [key, value] of Object.entries(kernelEnv || {})) { - if (typeof value === 'string' && substituteEnvVars(key, value, process.env) !== value) { + // Detect placeholders regardless of current process.env; we'll resolve after merges. + if (typeof value === 'string' && /\${[A-Za-z]\w*(?:[^}\w].*)?}/.test(value)) { kernelSpecVariablesRequiringSubstitution[key] = value; delete kernelEnv[key]; } @@ -123,6 +153,16 @@ export class KernelEnvironmentVariablesService { Object.assign(mergedVars, interpreterEnv, kernelEnv); // kernels vars win over interpreter. + // Merge SQL integration environment variables + if (sqlIntegrationEnvVars) { + logger.debug( + `KernelEnvVarsService: Merging ${ + Object.keys(sqlIntegrationEnvVars).length + } SQL integration env vars into kernel env` + ); + this.envVarsService.mergeVariables(sqlIntegrationEnvVars, mergedVars); + } + // If user asks us to, set PYTHONNOUSERSITE // For more details see here https://github.com/microsoft/vscode-jupyter/issues/8553#issuecomment-997144591 // https://docs.python.org/3/library/site.html#site.ENABLE_USER_SITE @@ -138,11 +178,21 @@ export class KernelEnvironmentVariablesService { // We can support this, however since this has not been requested, lets not do it.' this.envVarsService.mergeVariables(kernelEnv, mergedVars); // kernels vars win over interpreter. this.envVarsService.mergeVariables(customEnvVars, mergedVars); // custom vars win over all. + + // Merge SQL integration environment variables + if (sqlIntegrationEnvVars) { + logger.debug( + `KernelEnvVarsService: Merging ${ + Object.keys(sqlIntegrationEnvVars).length + } SQL integration env vars into kernel env (non-python)` + ); + this.envVarsService.mergeVariables(sqlIntegrationEnvVars, mergedVars); + } } // env variables in kernelSpecs can contain variables that need to be substituted for (const [key, value] of Object.entries(kernelSpecVariablesRequiringSubstitution)) { - mergedVars[key] = substituteEnvVars(key, value, mergedVars); + mergedVars[key] = substituteEnvVars(value, mergedVars); } return mergedVars; @@ -151,7 +201,7 @@ export class KernelEnvironmentVariablesService { const SUBST_REGEX = /\${([a-zA-Z]\w*)?([^}\w].*)?}/g; -function substituteEnvVars(key: string, value: string, globalVars: EnvironmentVariables): string { +function substituteEnvVars(value: string, globalVars: EnvironmentVariables): string { if (!value.includes('$')) { return value; } @@ -171,7 +221,6 @@ function substituteEnvVars(key: string, value: string, globalVars: EnvironmentVa return globalVars[substName] || ''; }); if (!invalid && replacement !== value) { - logger.debug(`${key} value in kernelSpec updated from ${value} to ${replacement}`); value = replacement; } diff --git a/src/kernels/raw/launcher/kernelEnvVarsService.unit.test.ts b/src/kernels/raw/launcher/kernelEnvVarsService.unit.test.ts index 927290e91a..b4dd91bdc5 100644 --- a/src/kernels/raw/launcher/kernelEnvVarsService.unit.test.ts +++ b/src/kernels/raw/launcher/kernelEnvVarsService.unit.test.ts @@ -17,6 +17,7 @@ import { IJupyterKernelSpec } from '../../types'; import { Uri } from 'vscode'; import { IConfigurationService, IWatchableJupyterSettings, type ReadWrite } from '../../../platform/common/types'; import { JupyterSettings } from '../../../platform/common/configSettings'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; use(chaiAsPromised); @@ -29,6 +30,7 @@ suite('Kernel Environment Variables Service', () => { let interpreterService: IInterpreterService; let configService: IConfigurationService; let settings: IWatchableJupyterSettings; + let sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider; const pathFile = Uri.joinPath(Uri.file('foobar'), 'bar'); const interpreter: PythonEnvironment = { uri: pathFile, @@ -53,14 +55,8 @@ suite('Kernel Environment Variables Service', () => { variablesService = new EnvironmentVariablesService(instance(fs)); configService = mock(); settings = mock(JupyterSettings); + sqlIntegrationEnvVars = mock(); when(configService.getSettings(anything())).thenReturn(instance(settings)); - kernelVariablesService = new KernelEnvironmentVariablesService( - instance(interpreterService), - instance(envActivation), - variablesService, - instance(customVariablesService), - instance(configService) - ); if (process.platform === 'win32') { // Win32 will generate upper case all the time const entries = Object.entries(process.env); @@ -72,9 +68,34 @@ suite('Kernel Environment Variables Service', () => { processEnv = process.env; } processPath = Object.keys(processEnv).find((k) => k.toLowerCase() == 'path'); + kernelVariablesService = buildKernelEnvVarsService(); }); + teardown(() => Object.assign(process.env, originalEnvVars)); + /** + * Helper factory function to build KernelEnvironmentVariablesService with optional overrides. + * @param overrides Optional overrides for the service dependencies + * @returns A new instance of KernelEnvironmentVariablesService + */ + function buildKernelEnvVarsService(overrides?: { + sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider | undefined; + }): KernelEnvironmentVariablesService { + const sqlProvider = + overrides && 'sqlIntegrationEnvVars' in overrides + ? overrides.sqlIntegrationEnvVars + : instance(sqlIntegrationEnvVars); + + return new KernelEnvironmentVariablesService( + instance(interpreterService), + instance(envActivation), + variablesService, + instance(customVariablesService), + instance(configService), + sqlProvider + ); + } + test('Python Interpreter path trumps process', async () => { when(envActivation.getActivatedEnvironmentVariables(anything(), anything(), anything())).thenResolve({ PATH: 'foobar' @@ -84,6 +105,7 @@ suite('Kernel Environment Variables Service', () => { }); when(customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything())).thenResolve(); when(customVariablesService.getCustomEnvironmentVariables(anything(), anything())).thenResolve(); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); const vars = await kernelVariablesService.getEnvironmentVariables(undefined, interpreter, kernelSpec); @@ -100,6 +122,7 @@ suite('Kernel Environment Variables Service', () => { }); when(customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything())).thenResolve(); when(customVariablesService.getCustomEnvironmentVariables(anything(), anything())).thenResolve(); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); const vars = await kernelVariablesService.getEnvironmentVariables(undefined, interpreter, kernelSpec); @@ -119,6 +142,7 @@ suite('Kernel Environment Variables Service', () => { when(customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything())).thenResolve({ HELLO_VAR: 'new' }); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); const vars = await kernelVariablesService.getEnvironmentVariables(undefined, interpreter, kernelSpec); @@ -134,6 +158,7 @@ suite('Kernel Environment Variables Service', () => { when(customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything())).thenResolve({ HELLO_VAR: 'new' }); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); const vars = await kernelVariablesService.getEnvironmentVariables(undefined, undefined, kernelSpec); @@ -148,6 +173,7 @@ suite('Kernel Environment Variables Service', () => { test('Returns process.env vars if no interpreter and no kernelspec.env', async () => { delete kernelSpec.interpreterPath; when(customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything())).thenResolve(); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); const vars = await kernelVariablesService.getEnvironmentVariables(undefined, undefined, kernelSpec); @@ -161,6 +187,7 @@ suite('Kernel Environment Variables Service', () => { when(customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything())).thenResolve({ PATH: 'foobaz' }); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); const vars = await kernelVariablesService.getEnvironmentVariables(undefined, interpreter, kernelSpec); assert.isOk(processPath); @@ -178,6 +205,7 @@ suite('Kernel Environment Variables Service', () => { when(customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything())).thenResolve({ PATH: 'foobaz' }); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); // undefined for interpreter here, interpreterPath from the spec should be used const vars = await kernelVariablesService.getEnvironmentVariables(undefined, undefined, kernelSpec); @@ -196,6 +224,7 @@ suite('Kernel Environment Variables Service', () => { when(customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything())).thenResolve({ PATH: 'foobaz' }); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); kernelSpec.env = { ONE: '1', TWO: '2' @@ -216,6 +245,7 @@ suite('Kernel Environment Variables Service', () => { when(customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything())).thenResolve({ PATH: 'foobaz' }); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); kernelSpec.env = { ONE: '1', TWO: '2', @@ -241,6 +271,7 @@ suite('Kernel Environment Variables Service', () => { when(customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything())).thenResolve({ PATH: 'foobaz' }); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); when(settings.excludeUserSitePackages).thenReturn(shouldBeSet); // undefined for interpreter here, interpreterPath from the spec should be used @@ -262,4 +293,136 @@ suite('Kernel Environment Variables Service', () => { test('PYTHONNOUSERSITE should be set for Virtual Env', async () => { await testPYTHONNOUSERSITE(EnvironmentType.VirtualEnv, true); }); + + suite('SQL Integration Environment Variables', () => { + test('SQL integration env vars are merged for Python kernels', async () => { + const resource = Uri.file('test.ipynb'); + when(envActivation.getActivatedEnvironmentVariables(anything(), anything(), anything())).thenResolve({ + PATH: 'foobar' + }); + when( + customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything()) + ).thenResolve(); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({ + SQL_MY_DB: '{"url":"postgresql://user:pass@host:5432/db","params":{},"param_style":"format"}' + }); + + const vars = await kernelVariablesService.getEnvironmentVariables(resource, interpreter, kernelSpec); + + assert.strictEqual( + vars!['SQL_MY_DB'], + '{"url":"postgresql://user:pass@host:5432/db","params":{},"param_style":"format"}' + ); + }); + + test('SQL integration env vars are merged for non-Python kernels', async () => { + const resource = Uri.file('test.ipynb'); + delete kernelSpec.interpreterPath; + when( + customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything()) + ).thenResolve(); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({ + SQL_MY_DB: '{"url":"postgresql://user:pass@host:5432/db","params":{},"param_style":"format"}' + }); + + const vars = await kernelVariablesService.getEnvironmentVariables(resource, undefined, kernelSpec); + + assert.strictEqual( + vars!['SQL_MY_DB'], + '{"url":"postgresql://user:pass@host:5432/db","params":{},"param_style":"format"}' + ); + }); + + test('SQL integration env vars are not added when provider returns empty object', async () => { + const resource = Uri.file('test.ipynb'); + when(envActivation.getActivatedEnvironmentVariables(anything(), anything(), anything())).thenResolve({ + PATH: 'foobar' + }); + when( + customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything()) + ).thenResolve(); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); + + const vars = await kernelVariablesService.getEnvironmentVariables(resource, interpreter, kernelSpec); + + assert.isUndefined(vars!['SQL_MY_DB']); + }); + + test('SQL integration env vars are not added when provider returns undefined', async () => { + const resource = Uri.file('test.ipynb'); + when(envActivation.getActivatedEnvironmentVariables(anything(), anything(), anything())).thenResolve({ + PATH: 'foobar' + }); + when( + customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything()) + ).thenResolve(); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); + + const vars = await kernelVariablesService.getEnvironmentVariables(resource, interpreter, kernelSpec); + + assert.isUndefined(vars!['SQL_MY_DB']); + }); + + test('Multiple SQL integration env vars are merged correctly', async () => { + const resource = Uri.file('test.ipynb'); + when(envActivation.getActivatedEnvironmentVariables(anything(), anything(), anything())).thenResolve({ + PATH: 'foobar' + }); + when( + customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything()) + ).thenResolve(); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({ + SQL_MY_DB: '{"url":"postgresql://user:pass@host:5432/db","params":{},"param_style":"format"}', + SQL_ANOTHER_DB: '{"url":"postgresql://user2:pass2@host2:5432/db2","params":{},"param_style":"format"}' + }); + + const vars = await kernelVariablesService.getEnvironmentVariables(resource, interpreter, kernelSpec); + + assert.strictEqual( + vars!['SQL_MY_DB'], + '{"url":"postgresql://user:pass@host:5432/db","params":{},"param_style":"format"}' + ); + assert.strictEqual( + vars!['SQL_ANOTHER_DB'], + '{"url":"postgresql://user2:pass2@host2:5432/db2","params":{},"param_style":"format"}' + ); + }); + + test('SQL integration env vars work when provider is undefined (optional dependency)', async () => { + const resource = Uri.file('test.ipynb'); + when(envActivation.getActivatedEnvironmentVariables(anything(), anything(), anything())).thenResolve({ + PATH: 'foobar' + }); + when( + customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything()) + ).thenResolve(); + + // Create service without SQL integration provider + const serviceWithoutSql = buildKernelEnvVarsService({ sqlIntegrationEnvVars: undefined }); + + const vars = await serviceWithoutSql.getEnvironmentVariables(resource, interpreter, kernelSpec); + + assert.isOk(vars); + assert.isUndefined(vars!['SQL_MY_DB']); + }); + + test('SQL integration env vars handle errors gracefully', async () => { + const resource = Uri.file('test.ipynb'); + when(envActivation.getActivatedEnvironmentVariables(anything(), anything(), anything())).thenResolve({ + PATH: 'foobar' + }); + when( + customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything()) + ).thenResolve(); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenReject( + new Error('Failed to get SQL env vars') + ); + + const vars = await kernelVariablesService.getEnvironmentVariables(resource, interpreter, kernelSpec); + + // Should still return vars without SQL integration vars + assert.isOk(vars); + assert.isUndefined(vars!['SQL_MY_DB']); + }); + }); }); diff --git a/src/kernels/serviceRegistry.node.ts b/src/kernels/serviceRegistry.node.ts index 74a8c5594f..0d55ce9e91 100644 --- a/src/kernels/serviceRegistry.node.ts +++ b/src/kernels/serviceRegistry.node.ts @@ -48,6 +48,8 @@ import { LastCellExecutionTracker } from './execution/lastCellExecutionTracker'; import { ClearJupyterServersCommand } from './jupyter/clearJupyterServersCommand'; import { KernelChatStartupCodeProvider } from './chat/kernelStartupCodeProvider'; import { KernelWorkingDirectory } from './raw/session/kernelWorkingDirectory.node'; +import { SqlIntegrationEnvironmentVariablesProvider } from '../platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider'; +import { ISqlIntegrationEnvVarsProvider } from '../platform/notebooks/deepnote/types'; export function registerTypes(serviceManager: IServiceManager, isDevMode: boolean) { serviceManager.addSingleton(IExtensionSyncActivationService, Activation); @@ -62,6 +64,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea ); serviceManager.addSingleton(IRawKernelSessionFactory, RawKernelSessionFactory); serviceManager.addSingleton(IKernelLauncher, KernelLauncher); + serviceManager.addSingleton( + ISqlIntegrationEnvVarsProvider, + SqlIntegrationEnvironmentVariablesProvider + ); serviceManager.addSingleton( KernelEnvironmentVariablesService, KernelEnvironmentVariablesService diff --git a/src/notebooks/deepnote/integrations/integrationDetector.ts b/src/notebooks/deepnote/integrations/integrationDetector.ts index 974b4a24fd..dac3d044d4 100644 --- a/src/notebooks/deepnote/integrations/integrationDetector.ts +++ b/src/notebooks/deepnote/integrations/integrationDetector.ts @@ -2,7 +2,7 @@ import { inject, injectable } from 'inversify'; import { logger } from '../../../platform/logging'; import { IDeepnoteNotebookManager } from '../../types'; -import { IntegrationStatus, IntegrationWithStatus } from './integrationTypes'; +import { IntegrationStatus, IntegrationWithStatus } from '../../../platform/notebooks/deepnote/integrationTypes'; import { IIntegrationDetector, IIntegrationStorage } from './types'; import { BlockWithIntegration, scanBlocksForIntegrations } from './integrationUtils'; diff --git a/src/notebooks/deepnote/integrations/integrationManager.ts b/src/notebooks/deepnote/integrations/integrationManager.ts index 87ab98d905..c138cc42b6 100644 --- a/src/notebooks/deepnote/integrations/integrationManager.ts +++ b/src/notebooks/deepnote/integrations/integrationManager.ts @@ -5,7 +5,7 @@ import { IExtensionContext } from '../../../platform/common/types'; import { Commands } from '../../../platform/common/constants'; import { logger } from '../../../platform/logging'; import { IIntegrationDetector, IIntegrationManager, IIntegrationStorage, IIntegrationWebviewProvider } from './types'; -import { IntegrationStatus, IntegrationWithStatus } from './integrationTypes'; +import { IntegrationStatus, IntegrationWithStatus } from '../../../platform/notebooks/deepnote/integrationTypes'; import { BlockWithIntegration, scanBlocksForIntegrations } from './integrationUtils'; /** @@ -149,7 +149,7 @@ export class IntegrationManager implements IIntegrationManager { // ensure it's in the map even if not detected from the project if (selectedIntegrationId && !integrations.has(selectedIntegrationId)) { logger.debug(`IntegrationManager: Adding requested integration ${selectedIntegrationId} to the map`); - const config = await this.integrationStorage.get(selectedIntegrationId); + const config = await this.integrationStorage.getIntegrationConfig(selectedIntegrationId); integrations.set(selectedIntegrationId, { config: config || null, status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected diff --git a/src/notebooks/deepnote/integrations/integrationUtils.ts b/src/notebooks/deepnote/integrations/integrationUtils.ts index 01b62b081e..9e9015ca9e 100644 --- a/src/notebooks/deepnote/integrations/integrationUtils.ts +++ b/src/notebooks/deepnote/integrations/integrationUtils.ts @@ -1,6 +1,10 @@ import { logger } from '../../../platform/logging'; import { IIntegrationStorage } from './types'; -import { DATAFRAME_SQL_INTEGRATION_ID, IntegrationStatus, IntegrationWithStatus } from './integrationTypes'; +import { + DATAFRAME_SQL_INTEGRATION_ID, + IntegrationStatus, + IntegrationWithStatus +} from '../../../platform/notebooks/deepnote/integrationTypes'; /** * Represents a block with SQL integration metadata @@ -48,7 +52,7 @@ export async function scanBlocksForIntegrations( logger.debug(`${logContext}: Found integration: ${integrationId} in block ${block.id}`); // Check if the integration is configured - const config = await integrationStorage.get(integrationId); + const config = await integrationStorage.getIntegrationConfig(integrationId); const status: IntegrationWithStatus = { config: config || null, diff --git a/src/notebooks/deepnote/integrations/integrationWebview.ts b/src/notebooks/deepnote/integrations/integrationWebview.ts index 09c2134369..bbcbd27660 100644 --- a/src/notebooks/deepnote/integrations/integrationWebview.ts +++ b/src/notebooks/deepnote/integrations/integrationWebview.ts @@ -6,7 +6,11 @@ import * as localize from '../../../platform/common/utils/localize'; import { logger } from '../../../platform/logging'; import { LocalizedMessages, SharedMessages } from '../../../messageTypes'; import { IIntegrationStorage, IIntegrationWebviewProvider } from './types'; -import { IntegrationConfig, IntegrationStatus, IntegrationWithStatus } from './integrationTypes'; +import { + IntegrationConfig, + IntegrationStatus, + IntegrationWithStatus +} from '../../../platform/notebooks/deepnote/integrationTypes'; /** * Manages the webview panel for integration configuration diff --git a/src/notebooks/deepnote/integrations/sqlIntegrationStartupCodeProvider.ts b/src/notebooks/deepnote/integrations/sqlIntegrationStartupCodeProvider.ts new file mode 100644 index 0000000000..693f1d8f11 --- /dev/null +++ b/src/notebooks/deepnote/integrations/sqlIntegrationStartupCodeProvider.ts @@ -0,0 +1,116 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { inject, injectable } from 'inversify'; +import { IStartupCodeProvider, IStartupCodeProviders, StartupCodePriority, IKernel } from '../../../kernels/types'; +import { JupyterNotebookView } from '../../../platform/common/constants'; +import { IExtensionSyncActivationService } from '../../../platform/activation/types'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; +import { logger } from '../../../platform/logging'; +import { isPythonKernelConnection } from '../../../kernels/helpers'; +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; +import { workspace } from 'vscode'; + +/** + * Provides startup code to inject SQL integration credentials into the kernel environment. + * This is necessary because Jupyter doesn't automatically pass all environment variables + * from the server process to the kernel process. + */ +@injectable() +export class SqlIntegrationStartupCodeProvider implements IStartupCodeProvider, IExtensionSyncActivationService { + public priority = StartupCodePriority.Base; + + constructor( + @inject(IStartupCodeProviders) private readonly registry: IStartupCodeProviders, + @inject(ISqlIntegrationEnvVarsProvider) + private readonly envVarsProvider: ISqlIntegrationEnvVarsProvider + ) {} + + activate(): void { + logger.debug('SqlIntegrationStartupCodeProvider: Activating and registering with JupyterNotebookView'); + this.registry.register(this, JupyterNotebookView); + logger.debug('SqlIntegrationStartupCodeProvider: Successfully registered'); + } + + async getCode(kernel: IKernel): Promise { + logger.debug( + `SqlIntegrationStartupCodeProvider.getCode called for kernel ${ + kernel.id + }, resourceUri: ${kernel.resourceUri?.toString()}` + ); + + // Only run for Python kernels on Deepnote notebooks + if (!isPythonKernelConnection(kernel.kernelConnectionMetadata)) { + logger.debug(`SqlIntegrationStartupCodeProvider: Not a Python kernel, skipping`); + return []; + } + + // Check if this is a Deepnote notebook + if (!kernel.resourceUri) { + logger.debug(`SqlIntegrationStartupCodeProvider: No resourceUri, skipping`); + return []; + } + + const notebook = workspace.notebookDocuments.find((nb) => nb.uri.toString() === kernel.resourceUri?.toString()); + if (!notebook) { + logger.debug(`SqlIntegrationStartupCodeProvider: Notebook not found for ${kernel.resourceUri.toString()}`); + return []; + } + + logger.debug(`SqlIntegrationStartupCodeProvider: Found notebook with type: ${notebook.notebookType}`); + + if (notebook.notebookType !== DEEPNOTE_NOTEBOOK_TYPE) { + logger.debug(`SqlIntegrationStartupCodeProvider: Not a Deepnote notebook, skipping`); + return []; + } + + try { + // Get SQL integration environment variables for this notebook + const envVars = await this.envVarsProvider.getEnvironmentVariables(kernel.resourceUri); + + if (!envVars || Object.keys(envVars).length === 0) { + logger.trace( + `SqlIntegrationStartupCodeProvider: No SQL integration env vars for ${kernel.resourceUri.toString()}` + ); + return []; + } + + logger.debug( + `SqlIntegrationStartupCodeProvider: Injecting ${ + Object.keys(envVars).length + } SQL integration env vars into kernel` + ); + + // Generate Python code to set environment variables directly in os.environ + const code: string[] = []; + + code.push('try:'); + code.push(' import os'); + code.push(` # [SQL Integration] Setting ${Object.keys(envVars).length} SQL integration env vars...`); + + // Set each environment variable directly in os.environ + for (const [key, value] of Object.entries(envVars)) { + if (value) { + // Use JSON.stringify to properly escape the value + const jsonEscaped = JSON.stringify(value); + code.push(` os.environ['${key}'] = ${jsonEscaped}`); + } + } + + code.push( + ` # [SQL Integration] Successfully set ${Object.keys(envVars).length} SQL integration env vars` + ); + code.push('except Exception as e:'); + code.push(' import traceback'); + code.push(' print(f"[SQL Integration] ERROR: Failed to set SQL integration env vars: {e}")'); + code.push(' traceback.print_exc()'); + + logger.debug('SqlIntegrationStartupCodeProvider: Generated startup code'); + + return code; + } catch (error) { + logger.error('SqlIntegrationStartupCodeProvider: Failed to generate startup code', error); + return []; + } + } +} diff --git a/src/notebooks/deepnote/integrations/types.ts b/src/notebooks/deepnote/integrations/types.ts index f414e34583..254f3c8912 100644 --- a/src/notebooks/deepnote/integrations/types.ts +++ b/src/notebooks/deepnote/integrations/types.ts @@ -1,27 +1,7 @@ -import { Event } from 'vscode'; -import { IDisposable } from '../../../platform/common/types'; -import { IntegrationConfig, IntegrationWithStatus } from './integrationTypes'; +import { IntegrationWithStatus } from '../../../platform/notebooks/deepnote/integrationTypes'; -export const IIntegrationStorage = Symbol('IIntegrationStorage'); -export interface IIntegrationStorage extends IDisposable { - /** - * Event fired when integrations change - */ - readonly onDidChangeIntegrations: Event; - - getAll(): Promise; - get(integrationId: string): Promise; - - /** - * Get integration configuration for a specific project and integration - */ - getIntegrationConfig(projectId: string, integrationId: string): Promise; - - save(config: IntegrationConfig): Promise; - delete(integrationId: string): Promise; - exists(integrationId: string): Promise; - clear(): Promise; -} +// Re-export IIntegrationStorage from platform layer +export { IIntegrationStorage } from '../../../platform/notebooks/deepnote/types'; export const IIntegrationDetector = Symbol('IIntegrationDetector'); export interface IIntegrationDetector { diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts index 80740260a9..c6439f2d8e 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts @@ -15,7 +15,7 @@ import { IExtensionSyncActivationService } from '../../platform/activation/types import { IDisposableRegistry } from '../../platform/common/types'; import { Commands } from '../../platform/common/constants'; import { IIntegrationStorage } from './integrations/types'; -import { DATAFRAME_SQL_INTEGRATION_ID } from './integrations/integrationTypes'; +import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes'; /** * Provides status bar items for SQL cells showing the integration name @@ -96,7 +96,7 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid } // Get integration configuration to display the name - const config = await this.integrationStorage.getIntegrationConfig(projectId, integrationId); + const config = await this.integrationStorage.getProjectIntegrationConfig(projectId, integrationId); const displayName = config?.name || l10n.t('Unknown integration (configure)'); // Create a status bar item that opens the integration management UI diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts index 1949be6712..ac7cf16386 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts @@ -13,7 +13,7 @@ import { import { IDisposableRegistry } from '../../platform/common/types'; import { IIntegrationStorage } from './integrations/types'; import { SqlCellStatusBarProvider } from './sqlCellStatusBarProvider'; -import { DATAFRAME_SQL_INTEGRATION_ID, IntegrationType } from './integrations/integrationTypes'; +import { DATAFRAME_SQL_INTEGRATION_ID, IntegrationType } from '../../platform/notebooks/deepnote/integrationTypes'; suite('SqlCellStatusBarProvider', () => { let provider: SqlCellStatusBarProvider; @@ -68,7 +68,7 @@ suite('SqlCellStatusBarProvider', () => { } ); - when(integrationStorage.getIntegrationConfig(anything(), anything())).thenResolve({ + when(integrationStorage.getProjectIntegrationConfig(anything(), anything())).thenResolve({ id: integrationId, name: 'My Postgres DB', type: IntegrationType.Postgres, @@ -101,7 +101,7 @@ suite('SqlCellStatusBarProvider', () => { } ); - when(integrationStorage.getIntegrationConfig(anything(), anything())).thenResolve(undefined); + when(integrationStorage.getProjectIntegrationConfig(anything(), anything())).thenResolve(undefined); const result = await provider.provideCellStatusBarItems(cell, cancellationToken); diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index a42308f4f4..3bd4f9d236 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -43,7 +43,7 @@ import { INotebookEditorProvider, INotebookPythonEnvironmentService } from './ty import { DeepnoteActivationService } from './deepnote/deepnoteActivationService'; import { DeepnoteNotebookManager } from './deepnote/deepnoteNotebookManager'; import { IDeepnoteNotebookManager } from './types'; -import { IntegrationStorage } from './deepnote/integrations/integrationStorage'; +import { IntegrationStorage } from '../platform/notebooks/deepnote/integrationStorage'; import { IntegrationDetector } from './deepnote/integrations/integrationDetector'; import { IntegrationManager } from './deepnote/integrations/integrationManager'; import { IntegrationWebviewProvider } from './deepnote/integrations/integrationWebview'; @@ -66,6 +66,7 @@ import { DeepnoteKernelAutoSelector } from './deepnote/deepnoteKernelAutoSelecto import { DeepnoteServerProvider } from '../kernels/deepnote/deepnoteServerProvider.node'; import { DeepnoteInitNotebookRunner, IDeepnoteInitNotebookRunner } from './deepnote/deepnoteInitNotebookRunner.node'; import { DeepnoteRequirementsHelper, IDeepnoteRequirementsHelper } from './deepnote/deepnoteRequirementsHelper.node'; +import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIntegrationStartupCodeProvider'; export function registerTypes(serviceManager: IServiceManager, isDevMode: boolean) { registerControllerTypes(serviceManager, isDevMode); @@ -147,6 +148,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, SqlCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + SqlIntegrationStartupCodeProvider + ); // Deepnote kernel services serviceManager.addSingleton(IDeepnoteToolkitInstaller, DeepnoteToolkitInstaller); diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index 1fb50d004a..dc67a169af 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/src/notebooks/serviceRegistry.web.ts @@ -38,7 +38,7 @@ import { INotebookEditorProvider, INotebookPythonEnvironmentService } from './ty import { DeepnoteActivationService } from './deepnote/deepnoteActivationService'; import { DeepnoteNotebookManager } from './deepnote/deepnoteNotebookManager'; import { IDeepnoteNotebookManager } from './types'; -import { IntegrationStorage } from './deepnote/integrations/integrationStorage'; +import { IntegrationStorage } from '../platform/notebooks/deepnote/integrationStorage'; import { IntegrationDetector } from './deepnote/integrations/integrationDetector'; import { IntegrationManager } from './deepnote/integrations/integrationManager'; import { IntegrationWebviewProvider } from './deepnote/integrations/integrationWebview'; diff --git a/src/notebooks/deepnote/integrations/integrationStorage.ts b/src/platform/notebooks/deepnote/integrationStorage.ts similarity index 92% rename from src/notebooks/deepnote/integrations/integrationStorage.ts rename to src/platform/notebooks/deepnote/integrationStorage.ts index 17a8a2b15c..6b86d52f1b 100644 --- a/src/notebooks/deepnote/integrations/integrationStorage.ts +++ b/src/platform/notebooks/deepnote/integrationStorage.ts @@ -1,9 +1,9 @@ import { inject, injectable } from 'inversify'; import { EventEmitter } from 'vscode'; -import { IEncryptedStorage } from '../../../platform/common/application/types'; -import { IAsyncDisposableRegistry } from '../../../platform/common/types'; -import { logger } from '../../../platform/logging'; +import { IEncryptedStorage } from '../../common/application/types'; +import { IAsyncDisposableRegistry } from '../../common/types'; +import { logger } from '../../logging'; import { IntegrationConfig, IntegrationType } from './integrationTypes'; import { IIntegrationStorage } from './types'; @@ -43,7 +43,7 @@ export class IntegrationStorage implements IIntegrationStorage { /** * Get a specific integration configuration by ID */ - async get(integrationId: string): Promise { + async getIntegrationConfig(integrationId: string): Promise { await this.ensureCacheLoaded(); return this.cache.get(integrationId); } @@ -53,8 +53,11 @@ export class IntegrationStorage implements IIntegrationStorage { * Note: Currently integrations are stored globally, not per-project, * so this method ignores the projectId parameter */ - async getIntegrationConfig(_projectId: string, integrationId: string): Promise { - return this.get(integrationId); + async getProjectIntegrationConfig( + _projectId: string, + integrationId: string + ): Promise { + return this.getIntegrationConfig(integrationId); } /** diff --git a/src/notebooks/deepnote/integrations/integrationTypes.ts b/src/platform/notebooks/deepnote/integrationTypes.ts similarity index 100% rename from src/notebooks/deepnote/integrations/integrationTypes.ts rename to src/platform/notebooks/deepnote/integrationTypes.ts diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts new file mode 100644 index 0000000000..13da2ab4a3 --- /dev/null +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts @@ -0,0 +1,212 @@ +import { inject, injectable } from 'inversify'; +import { CancellationToken, Event, EventEmitter, NotebookDocument, workspace } from 'vscode'; + +import { IDisposableRegistry, Resource } from '../../common/types'; +import { EnvironmentVariables } from '../../common/variables/types'; +import { BaseError } from '../../errors/types'; +import { logger } from '../../logging'; +import { IIntegrationStorage, ISqlIntegrationEnvVarsProvider } from './types'; +import { DATAFRAME_SQL_INTEGRATION_ID, IntegrationConfig, IntegrationType } from './integrationTypes'; + +/** + * Error thrown when an unsupported integration type is encountered. + * + * Cause: + * An integration configuration has a type that is not supported by the SQL integration system. + * + * Handled by: + * Callers should handle this error and inform the user that the integration type is not supported. + */ +class UnsupportedIntegrationError extends BaseError { + constructor(public readonly integrationType: string) { + super('unknown', `Unsupported integration type: ${integrationType}`); + } +} + +/** + * 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: IntegrationConfig): string { + switch (config.type) { + 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' + }); + } + + default: + throw new UnsupportedIntegrationError((config as IntegrationConfig).type); + } +} + +/** + * Provides environment variables for SQL integrations. + * This service scans notebooks for SQL blocks and injects the necessary credentials + * as environment variables so they can be used during SQL block execution. + */ +@injectable() +export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrationEnvVarsProvider { + private readonly _onDidChangeEnvironmentVariables = new EventEmitter(); + + public readonly onDidChangeEnvironmentVariables: Event = this._onDidChangeEnvironmentVariables.event; + + constructor( + @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, + @inject(IDisposableRegistry) disposables: IDisposableRegistry + ) { + logger.info('SqlIntegrationEnvironmentVariablesProvider: Constructor called - provider is being instantiated'); + // Dispose emitter when extension deactivates + disposables.push(this._onDidChangeEnvironmentVariables); + // Listen for changes to integration storage and fire change event + disposables.push( + this.integrationStorage.onDidChangeIntegrations(() => { + // Fire change event for all notebooks + this._onDidChangeEnvironmentVariables.fire(undefined); + }) + ); + } + + /** + * Get environment variables for SQL integrations used in the given notebook. + */ + public async getEnvironmentVariables(resource: Resource, token?: CancellationToken): Promise { + const envVars: EnvironmentVariables = {}; + + if (!resource) { + return envVars; + } + + if (token?.isCancellationRequested) { + return envVars; + } + + 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; + } + + // 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()}` + ); + return envVars; + } + + logger.trace(`SqlIntegrationEnvironmentVariablesProvider: Found ${integrationIds.size} SQL integrations`); + + // Get credentials for each integration and add to environment variables + for (const integrationId of integrationIds) { + if (token?.isCancellationRequested) { + break; + } + + try { + const config = await this.integrationStorage.getIntegrationConfig(integrationId); + if (!config) { + logger.warn( + `SqlIntegrationEnvironmentVariablesProvider: No configuration found for integration ${integrationId}` + ); + continue; + } + + // Convert integration config to JSON and add as environment variable + const envVarName = convertToEnvironmentVariableName(getSqlEnvVarName(integrationId)); + const credentialsJson = convertIntegrationConfigToJson(config); + + envVars[envVarName] = credentialsJson; + logger.info( + `SqlIntegrationEnvironmentVariablesProvider: Added env var ${envVarName} for integration ${integrationId}` + ); + } catch (error) { + logger.error( + `SqlIntegrationEnvironmentVariablesProvider: Failed to get credentials for integration ${integrationId}`, + error + ); + } + } + + logger.trace(`SqlIntegrationEnvironmentVariablesProvider: Returning ${Object.keys(envVars).length} env vars`); + + 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') { + // Skip the internal DuckDB integration + if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) { + continue; + } + + 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 new file mode 100644 index 0000000000..f7348d124a --- /dev/null +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts @@ -0,0 +1,394 @@ +import { assert } from 'chai'; +import { instance, mock, when } from 'ts-mockito'; +import { CancellationTokenSource, EventEmitter, NotebookCell, NotebookCellKind, NotebookDocument, Uri } from 'vscode'; + +import { IDisposableRegistry } from '../../common/types'; +import { IntegrationStorage } from './integrationStorage'; +import { SqlIntegrationEnvironmentVariablesProvider } from './sqlIntegrationEnvironmentVariablesProvider'; +import { IntegrationType, PostgresIntegrationConfig, BigQueryIntegrationConfig } from './integrationTypes'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; + +suite('SqlIntegrationEnvironmentVariablesProvider', () => { + let provider: SqlIntegrationEnvironmentVariablesProvider; + let integrationStorage: IntegrationStorage; + let disposables: IDisposableRegistry; + + setup(() => { + resetVSCodeMocks(); + disposables = []; + integrationStorage = mock(IntegrationStorage); + when(integrationStorage.onDidChangeIntegrations).thenReturn(new EventEmitter().event); + + provider = new SqlIntegrationEnvironmentVariablesProvider(instance(integrationStorage), 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 notebook is not found', 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, {}); + }); + + 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]); + + const envVars = await provider.getEnvironmentVariables(uri); + assert.deepStrictEqual(envVars, {}); + }); + + test('Skips 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: 'deepnote-dataframe-sql' + }) + ]); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + const envVars = await provider.getEnvironmentVariables(uri); + assert.deepStrictEqual(envVars, {}); + }); + + test('Returns environment variable for PostgreSQL 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', + 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); + + 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: BigQueryIntegrationConfig = { + id: integrationId, + name: 'My BigQuery', + type: IntegrationType.BigQuery, + projectId: 'my-project', + 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); + + 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('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 + assert.property(envVars, 'SQL_MY_POSTGRES_DB'); + assert.strictEqual(Object.keys(envVars).length, 1); + }); + + test('Handles multiple SQL cells with different integrations', async () => { + const uri = Uri.file('/test/notebook.deepnote'); + const postgresId = 'my-postgres-db'; + const bigqueryId = 'my-bigquery'; + + const postgresConfig: PostgresIntegrationConfig = { + id: postgresId, + name: 'My Postgres DB', + type: IntegrationType.Postgres, + host: 'localhost', + port: 5432, + database: 'mydb', + username: 'user', + password: 'pass' + }; + + const bigqueryConfig: BigQueryIntegrationConfig = { + id: bigqueryId, + name: 'My BigQuery', + type: IntegrationType.BigQuery, + projectId: 'my-project', + 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); + + const envVars = await provider.getEnvironmentVariables(uri); + + // Should have two environment variables + assert.property(envVars, 'SQL_MY_POSTGRES_DB'); + assert.property(envVars, 'SQL_MY_BIGQUERY'); + assert.strictEqual(Object.keys(envVars).length, 2); + }); + + 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 empty object when integration config is missing + assert.deepStrictEqual(envVars, {}); + }); + + test('Properly encodes special characters in PostgreSQL credentials', async () => { + const uri = Uri.file('/test/notebook.deepnote'); + const integrationId = 'special-chars-db'; + const config: PostgresIntegrationConfig = { + id: integrationId, + name: 'Special Chars DB', + type: IntegrationType.Postgres, + host: 'db.example.com', + port: 5432, + database: 'my@db:name', + username: 'user@domain', + password: 'pa:ss@word!#$%', + 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); + + 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: PostgresIntegrationConfig = { + id: integrationId, + name: 'Production Database', + type: IntegrationType.Postgres, + host: 'prod.example.com', + port: 5432, + database: 'proddb', + username: 'admin', + password: 'secret', + 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); + + 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: PostgresIntegrationConfig = { + id: integrationId, + name: 'Test DB', + type: IntegrationType.Postgres, + host: 'localhost', + port: 5432, + database: 'testdb', + username: 'user', + password: 'pass', + 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); + + 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 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)) + ); + const cts = new CancellationTokenSource(); + cts.cancel(); + const envVars = await provider.getEnvironmentVariables(uri, cts.token); + assert.deepStrictEqual(envVars, {}); + }); +}); + +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; +} diff --git a/src/platform/notebooks/deepnote/types.ts b/src/platform/notebooks/deepnote/types.ts new file mode 100644 index 0000000000..6036737e2c --- /dev/null +++ b/src/platform/notebooks/deepnote/types.ts @@ -0,0 +1,55 @@ +import { CancellationToken, Event } from 'vscode'; +import { IDisposable, Resource } from '../../common/types'; +import { EnvironmentVariables } from '../../common/variables/types'; +import { IntegrationConfig } from './integrationTypes'; + +export const IIntegrationStorage = Symbol('IIntegrationStorage'); +export interface IIntegrationStorage extends IDisposable { + /** + * Event fired when integrations change + */ + readonly onDidChangeIntegrations: Event; + + getAll(): Promise; + + /** + * Retrieves the global (non-project-scoped) integration configuration by integration ID. + * + * This method returns integration configurations that are stored globally and shared + * across all projects. These configurations are stored in VSCode's SecretStorage and + * are scoped to the user's machine. + * + * This differs from `getProjectIntegrationConfig()` which returns project-scoped + * configurations that are specific to a particular Deepnote project and stored + * within the project's YAML file. + * + * @param integrationId - The unique identifier of the integration to retrieve + * @returns A Promise that resolves to: + * - 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; + + /** + * Get integration configuration for a specific project and integration + */ + getProjectIntegrationConfig(projectId: string, integrationId: string): Promise; + + save(config: IntegrationConfig): Promise; + delete(integrationId: string): Promise; + exists(integrationId: string): Promise; + clear(): Promise; +} + +export const ISqlIntegrationEnvVarsProvider = Symbol('ISqlIntegrationEnvVarsProvider'); +export interface ISqlIntegrationEnvVarsProvider { + /** + * Event fired when environment variables change + */ + readonly onDidChangeEnvironmentVariables: Event; + + /** + * Get environment variables for SQL integrations used in the given notebook. + */ + getEnvironmentVariables(resource: Resource, token?: CancellationToken): Promise; +}