diff --git a/cspell.json b/cspell.json index 8aa54f8759..abf2ce3021 100644 --- a/cspell.json +++ b/cspell.json @@ -98,6 +98,8 @@ "Unconfigured", "unuse", "unittests", + "userpod", + "Userpod", "vegalite", "venv", "Venv", diff --git a/package-lock.json b/package-lock.json index a3dcd43d19..add0b330cf 100644 --- a/package-lock.json +++ b/package-lock.json @@ -41,6 +41,7 @@ "clsx": "^2.1.1", "cross-fetch": "^3.1.5", "d3-format": "^3.1.0", + "dotenv": "^17.2.3", "encoding": "^0.1.13", "express": "^5.2.1", "fast-deep-equal": "^2.0.1", @@ -15519,6 +15520,18 @@ "url": "https://github.com/fb55/domutils?sponsor=1" } }, + "node_modules/dotenv": { + "version": "17.4.2", + "resolved": "https://registry.npmjs.org/dotenv/-/dotenv-17.4.2.tgz", + "integrity": "sha512-nI4U3TottKAcAD9LLud4Cb7b2QztQMUEfHbvhTH09bqXTxnSie8WnjPALV/WMCrJZ6UV/qHJ6L03OqO3LcdYZw==", + "license": "BSD-2-Clause", + "engines": { + "node": ">=12" + }, + "funding": { + "url": "https://dotenvx.com" + } + }, "node_modules/dunder-proto": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/dunder-proto/-/dunder-proto-1.0.1.tgz", @@ -46524,6 +46537,11 @@ "domhandler": "^5.0.3" } }, + "dotenv": { + "version": "17.4.2", + "resolved": "https://registry.npmjs.org/dotenv/-/dotenv-17.4.2.tgz", + "integrity": "sha512-nI4U3TottKAcAD9LLud4Cb7b2QztQMUEfHbvhTH09bqXTxnSie8WnjPALV/WMCrJZ6UV/qHJ6L03OqO3LcdYZw==" + }, "dunder-proto": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/dunder-proto/-/dunder-proto-1.0.1.tgz", diff --git a/package.json b/package.json index d2adcbfaff..05407a933f 100644 --- a/package.json +++ b/package.json @@ -1649,6 +1649,12 @@ "description": "When enabled, outputs are saved to separate snapshot files in a 'snapshots' folder instead of the main .deepnote file.", "scope": "resource" }, + "deepnote.integrations.envFile.enabled": { + "type": "boolean", + "default": true, + "description": "When enabled, integration credentials are also loaded from a '.deepnote.env.yaml' file (with 'env:' references resolved against '.env' and environment variables) next to the .deepnote file or at the workspace root.", + "scope": "resource" + }, "deepnote.experiments.enabled": { "type": "boolean", "default": true, @@ -2703,6 +2709,7 @@ "clsx": "^2.1.1", "cross-fetch": "^3.1.5", "d3-format": "^3.1.0", + "dotenv": "^17.2.3", "encoding": "^0.1.13", "express": "^5.2.1", "fast-deep-equal": "^2.0.1", diff --git a/specs/INTEGRATIONS_CREDENTIALS.md b/specs/INTEGRATIONS_CREDENTIALS.md index 5b414342c6..8edc5e2233 100644 --- a/specs/INTEGRATIONS_CREDENTIALS.md +++ b/specs/INTEGRATIONS_CREDENTIALS.md @@ -397,28 +397,25 @@ Handles migration of legacy integration configurations to the new `@deepnote/dat #### 2. **Integration Detector** (`integrationDetector.ts`) -Scans Deepnote projects to discover which integrations are used in SQL blocks. +Lists the integrations a Deepnote project declares, paired with the credentials stored for each. **Detection Process:** 1. Retrieves the Deepnote project from `IDeepnoteNotebookManager` -2. Scans all notebooks in the project -3. Examines each code block for `metadata.sql_integration_id` -4. Maps Deepnote integration types to `DatabaseIntegrationType` using the project's integration list -5. Checks if each integration is configured (has credentials) -6. Returns a map of integration IDs to their status - -**Integration Status:** - -- `Connected`: Integration has valid credentials stored -- `Disconnected`: Integration is used but not configured -- `Error`: Integration configuration is invalid +2. Iterates the project's `integrations` roster (ids, names and types only — never credentials) +3. Skips entries whose type is not a configurable `DatabaseIntegrationType` +4. Reads each integration's config from `SecretStorage`, or `null` when nothing is stored +5. Appends `.deepnote.env.yaml` integrations the roster omits, so the panel lists what actually applies at + execution time; a failed lookup leaves the roster-only result rather than blocking the panel +6. Returns a map of integration IDs to their config and declared name/type **Special Cases:** - Excludes `deepnote-dataframe-sql` (internal DuckDB integration) -- Only processes code blocks with SQL integration metadata -- Uses project integration metadata to determine integration types +- Integrations configured in `.deepnote.env.yaml` have a `null` config here, since those configs are never + written through `SecretStorage`. They still resolve normally for kernel execution and SQL autocomplete, which + read the merged configs directly, and the panel marks them read-only rather than unconfigured (see below): + `IntegrationWebviewProvider` pairs the detection result with the file-configured ids. #### 3. **Integration Manager** (`integrationManager.ts`) @@ -427,10 +424,6 @@ Orchestrates the integration management UI and commands. **Responsibilities:** - Registers the `deepnote.manageIntegrations` command -- Updates VSCode context keys for UI visibility: - - `deepnote.hasIntegrations`: True if any integrations are detected - - `deepnote.hasUnconfiguredIntegrations`: True if any integrations lack credentials -- Handles notebook selection changes - Opens the integration webview with detected integrations **Command Flow:** @@ -447,17 +440,35 @@ Provides the webview-based UI for managing integration credentials. **Features:** - Persistent webview panel (survives defocus) -- Real-time integration status updates +- Real-time updates as credentials are saved or cleared - Configuration forms for each integration type - Delete/reset functionality +**Connection Status:** + +`IntegrationItem.tsx` derives the status pill from two inputs: `config` (the SecretStorage credentials the panel +can edit) and `isFileConfigured`, which `IntegrationWebviewProvider` computes from +`ISqlIntegrationEnvVarsProvider.getFileConfiguredIntegrationIds()`. That call returns ids only, so neither the +file's config nor its credentials ever reach the webview. + +- **"Configured in file"** — `.deepnote.env.yaml` configures this id. Rendered with the connected styling, and + the Configure/Reset/Delete actions are hidden: the panel writes `SecretStorage` only and the file wins the + merge, so those edits would be silent no-ops at runtime. This holds even when SecretStorage also happens to + hold a config for the same id. +- **"Connected"** — SecretStorage holds credentials and no file config overrides them. +- **"Not Configured"** — neither layer configures the integration. + +The separate federated-auth pill (BigQuery + `google-oauth`) is independent and tracks whether a refresh token +is stored. It is driven by `tokenStatus` alone, which the extension derives from its own candidate set, so the +Authenticate action stays available for file-configured integrations — exactly the rows that carry no `config`. + **Message Protocol:** Extension → Webview: ```typescript // Update integration list -{ type: 'update', integrations: IntegrationWithStatus[] } +{ type: 'update', integrations: DetectedIntegration[] } // Show configuration form { type: 'showForm', integrationId: string, config: IntegrationConfig | null } @@ -487,7 +498,7 @@ Main React component that manages the webview UI state. **State Management:** -- `integrations`: List of detected integrations with status +- `integrations`: List of detected integrations with their stored configs - `selectedIntegrationId`: Currently selected integration for configuration - `selectedConfig`: Existing configuration being edited - `message`: Success/error messages @@ -501,7 +512,7 @@ Main React component that manages the webview UI state. 2. Panel shows configuration form overlay 3. User enters credentials 4. Panel sends save message to extension -5. Extension stores credentials and updates status +5. Extension stores credentials 6. Panel shows success message and refreshes list **Delete Integration:** @@ -511,7 +522,7 @@ Main React component that manages the webview UI state. 3. User clicks again to confirm 4. Panel sends delete message to extension 5. Extension removes credentials -6. Panel updates status to "Disconnected" +6. Panel clears the stored config, so the item offers "Configure" again #### 6. **Configuration Forms** @@ -678,55 +689,23 @@ DuckDB (internal): **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 +- Backs the loopback `userpod-api` endpoint (`userpodApiEndpoints.node.ts`), which serves the resolved credentials to `deepnote-toolkit` as `[{name, value}]` +- Queried per request rather than at kernel start, so a kernel always reads the credentials that are current at the moment it asks - Fires `onDidChangeEnvironmentVariables` event when integration storage changes -#### 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** +#### 8. **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 +1. Extension starts the Jupyter server with the integration endpoint env vars applied by `applyIntegrationEndpointEnv` (`DEEPNOTE_RUNTIME__ENV_INTEGRATION_ENABLED`, `DEEPNOTE_RUNTIME__RUNNING_IN_DETACHED_MODE`, `DEEPNOTE_RUNTIME__WEBAPP_URL`, `DEEPNOTE_RUNTIME__PROJECT_SECRET`, `DEEPNOTE_PROJECT_ID`) — no credentials are injected by the extension +2. At kernel init, `deepnote-toolkit`'s `set_integration_env()` calls the extension's loopback `userpod-api` endpoint at `DEEPNOTE_RUNTIME__WEBAPP_URL`, authenticating with the per-project bearer token in `DEEPNOTE_RUNTIME__PROJECT_SECRET` +3. The endpoint resolves the project's integrations via `SqlIntegrationEnvironmentVariablesProvider` — reading secure storage and converting the configs to `SQL_*` JSON values — and returns them as `[{name, value}]` +4. Toolkit sets the returned variables into the kernel's `os.environ` +5. `deepnote-toolkit` reads those environment variables when executing SQL blocks 6. Toolkit creates database connections using the credentials 7. Toolkit executes SQL queries and returns results @@ -760,21 +739,31 @@ User → IntegrationPanel (UI) ### Execution Flow ```text +Deepnote server starts + → applyIntegrationEndpointEnv() + → Awaits UserpodApiEndpoints readiness, reads its loopback baseUrl + → Resolves the project id from the .deepnote file + → Sets DEEPNOTE_RUNTIME__* + DEEPNOTE_PROJECT_ID on the server process + → Jupyter server process starts with those env vars (no credentials) + User executes SQL cell → Kernel startup triggered - → SqlIntegrationEnvironmentVariablesProvider.getEnvironmentVariables() - → Identifies Deepnote project from notebook resource - → Retrieves project integrations from notebook manager - → Fetches configured credentials from IntegrationStorage - → Adds internal DuckDB integration - → Calls getEnvironmentVariablesForIntegrations() from @deepnote/database-integrations - → Converts configs to environment variable format - → Generates SQL_* environment variables - → 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 set_integration_env() + → GET {DEEPNOTE_RUNTIME__WEBAPP_URL}/userpod-api/{projectId}/integrations/environment-variables + Authorization: Bearer {DEEPNOTE_RUNTIME__PROJECT_SECRET} + → UserpodApiEndpoints handles the request + → Verifies the per-project bearer token + → Finds the open deepnote notebook(s) for that project id + → SqlIntegrationEnvironmentVariablesProvider.getEnvironmentVariables() + → Retrieves project integrations from notebook manager + → Fetches configured credentials from IntegrationStorage + → Adds internal DuckDB integration + → Calls getEnvironmentVariablesForIntegrations() from @deepnote/database-integrations + → Converts configs to environment variable format + → Generates SQL_* environment variables + → Returns environment variables + → Responds with [{name, value}] + → Toolkit sets them into the kernel's os.environ → deepnote-toolkit reads os.environ['SQL_*'] → Toolkit executes SQL query → Results returned to notebook diff --git a/specs/INTEGRATION_ENV_LIVE_REFRESH.md b/specs/INTEGRATION_ENV_LIVE_REFRESH.md new file mode 100644 index 0000000000..f5a1fb238b --- /dev/null +++ b/specs/INTEGRATION_ENV_LIVE_REFRESH.md @@ -0,0 +1,154 @@ +# Live Integration Environment Refresh + +## Overview + +When integration credentials change — either in the extension's SecretStorage or in a `.deepnote.env.yaml` / +`.env` file — running Jupyter kernels pick up the new values **without a restart**. The kernel is asked to +re-fetch its integration environment in place, so variables, imports and loaded data survive the change. + +This replaces the earlier design, which restarted every kernel that used SQL integrations and re-injected +credentials through startup code. + +## Why no restart + +Credentials are no longer baked into the kernel at startup. Instead the extension runs a small loopback HTTP +server (`UserpodApiEndpoints`) and `deepnote-toolkit` fetches integration env vars from it on demand. Because +the source of truth lives on the extension side, refreshing is just a matter of telling the toolkit to fetch +again — restarting the kernel would throw away user state for no benefit. + +## Components + +### 1. `IntegrationEnvRefreshHandler` (`integrationEnvRefreshHandler.ts`) + +Reacts to **SecretStorage** changes. + +- Subscribes to `IIntegrationStorage.onDidChangeIntegrations` (fired by `IntegrationStorage.save()` / `delete()`) +- Collects every open notebook of type `deepnote` +- Hands them to `IIntegrationEnvLiveRefresher` + +### 2. `IntegrationsEnvFileWatcher` (`integrationsEnvFileWatcher.node.ts`) + +Reacts to **file** changes. + +- Watches `.deepnote.env.yaml` and `.env` in each workspace folder and in every open notebook's directory +- Debounces bursts (500 ms trailing edge) so saving both files is handled once +- For each open Deepnote notebook, keeps it only when all of the following hold: + 1. `deepnote.integrations.envFile.enabled` is on for that file — the same gate the config provider uses, via + the shared `isIntegrationsEnvFileEnabled()` helper + 2. the changed directory is the notebook's own directory or its workspace root + 3. a `.deepnote.env.yaml` actually exists for it — otherwise an unrelated `.env` (a very common + non-Deepnote file) would trigger hidden kernel executions and a misleading status message +- Hands the survivors to `IIntegrationEnvLiveRefresher` + +### 3. `IntegrationEnvLiveRefresher` (`integrationEnvLiveRefresher.node.ts`) + +Performs the refresh, for both triggers. + +- Skips notebooks with no kernel, or whose kernel has never started +- Runs a hidden execution in each remaining kernel: + + ```python + import deepnote_toolkit + deepnote_toolkit.set_integration_env() + ``` + +- Treats `error` outputs as a failed refresh and logs them +- On at least one success, shows a transient status-bar message ("Deepnote integration environment updated.") + for 5 seconds rather than a toast, so frequent env-file edits don't spam notifications + +### 4. `UserpodApiEndpoints` (`userpodApiEndpoints.node.ts`) + +Serves the credentials the toolkit fetches. + +- Binds an ephemeral port on `127.0.0.1` +- `GET /userpod-api/:projectId/integrations/environment-variables` returns `[{ name, value }]` +- Requires a per-project bearer token, compared in constant time +- Resolves values through `ISqlIntegrationEnvVarsProvider.getEnvironmentVariables()`, which merges SecretStorage + with `.deepnote.env.yaml` (the file wins) + +The endpoint URL and token reach the kernel via `applyIntegrationEndpointEnv()` +(`src/kernels/deepnote/deepnoteIntegrationEndpointEnv.ts`), which sets `DEEPNOTE_RUNTIME__WEBAPP_URL`, +`DEEPNOTE_RUNTIME__PROJECT_SECRET`, `DEEPNOTE_RUNTIME__ENV_INTEGRATION_ENABLED`, +`DEEPNOTE_RUNTIME__RUNNING_IN_DETACHED_MODE` and `DEEPNOTE_PROJECT_ID` at spawn time. + +## Flow + +``` +┌──────────────────────────────────┐ ┌──────────────────────────────────┐ +│ User saves config in the panel │ │ `.deepnote.env.yaml` / `.env` │ +│ │ │ created, changed or deleted │ +└────────────────┬─────────────────┘ └────────────────┬─────────────────┘ + │ │ + ▼ ▼ +┌──────────────────────────────────┐ ┌──────────────────────────────────┐ +│ IntegrationStorage │ │ IntegrationsEnvFileWatcher │ +│ - Writes to SecretStorage │ │ - Debounces 500 ms │ +│ - Fires onDidChangeIntegrations │ │ - Filters to affected notebooks │ +└────────────────┬─────────────────┘ └────────────────┬─────────────────┘ + │ │ + ▼ │ +┌──────────────────────────────────┐ │ +│ IntegrationEnvRefreshHandler │ │ +│ - Collects open .deepnote docs │ │ +└────────────────┬─────────────────┘ │ + │ │ + └──────────────┬───────────────────────┘ + ▼ + ┌──────────────────────────────────┐ + │ IntegrationEnvLiveRefresher │ + │ - Hidden exec per live kernel: │ + │ set_integration_env() │ + └────────────────┬─────────────────┘ + │ HTTP (loopback, bearer token) + ▼ + ┌──────────────────────────────────┐ + │ UserpodApiEndpoints │ + │ - Merges SecretStorage + file │ + │ - Returns [{name, value}] │ + └────────────────┬─────────────────┘ + ▼ + ┌──────────────────────────────────┐ + │ Kernel env updated in place │ + │ - SQL cells use new credentials │ + │ - No restart, no state lost │ + └──────────────────────────────────┘ +``` + +## Service Registration + +All four services are **node-only** — they are registered in `src/notebooks/serviceRegistry.node.ts` and have no +web counterpart, since the loopback server and the file reads both require Node APIs. The web registry +(`serviceRegistry.web.ts`) registers only the storage, detector, webview and manager. + +`IntegrationEnvRefreshHandler`, `IntegrationsEnvFileWatcher` and `UserpodApiEndpoints` are registered as +`IExtensionSyncActivationService`, so they activate with the extension. + +## Error Handling + +- `IntegrationEnvLiveRefresher.refreshNotebook()` never throws; a per-notebook failure is logged and the other + notebooks still refresh +- Both triggers wrap their async entry point and log rather than surfacing an unhandled rejection +- A refresh that produces `error` outputs is counted as failed, so the success message only appears when at + least one kernel actually applied the new environment +- `UserpodApiEndpoints` keeps a persistent `error` listener (an `error` with no listener would crash the + extension host) and, once it has been listening, offers a window reload if the server dies — restarting in + place would bind a new port that already-running kernels don't know about + +## Testing + +- `integrationEnvLiveRefresher.node.unit.test.ts` — kernel selection, hidden execution, error outputs +- `integrationEnvRefreshHandler.unit.test.ts` — SecretStorage event → refresh +- `integrationsEnvFileWatcher.node.unit.test.ts` — debounce, scoping, the `enabled` gate and the + "no `.deepnote.env.yaml`, no refresh" rule +- `userpodApiEndpoints.node.unit.test.ts` — auth, project scoping, payload shape + +## Related Files + +- `src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.ts` — SecretStorage trigger +- `src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.ts` — file trigger +- `src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.ts` — performs the refresh +- `src/notebooks/deepnote/integrations/userpodApiEndpoints.node.ts` — serves credentials over loopback +- `src/kernels/deepnote/deepnoteIntegrationEndpointEnv.ts` — injects the endpoint URL and token into kernel env +- `src/platform/notebooks/deepnote/integrationStorage.ts` — storage layer that fires the change event +- `src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts` — merges SecretStorage and file +- `src/platform/notebooks/deepnote/integrationsEnvFileSettings.ts` — the shared `envFile.enabled` gate diff --git a/specs/KERNEL_RESTART_ON_INTEGRATION_CHANGE.md b/specs/KERNEL_RESTART_ON_INTEGRATION_CHANGE.md deleted file mode 100644 index d0a45be342..0000000000 --- a/specs/KERNEL_RESTART_ON_INTEGRATION_CHANGE.md +++ /dev/null @@ -1,163 +0,0 @@ -# Automatic Kernel Restart on Integration Configuration Changes - -## Overview - -This feature automatically restarts Jupyter kernels when integration configurations (e.g., PostgreSQL, BigQuery, Snowflake credentials) are changed. This ensures that running kernels immediately pick up new credentials without requiring manual intervention. - -## Implementation - -### New Service: `IntegrationKernelRestartHandler` - -**Location**: `src/notebooks/deepnote/integrations/integrationKernelRestartHandler.ts` - -**Purpose**: Listens for integration configuration changes and automatically restarts affected kernels. - -**Key Features**: -- Listens to `onDidChangeIntegrations` event from `IntegrationStorage` -- Scans all open Deepnote notebooks for SQL cells that use integrations -- Identifies running kernels that need to be restarted -- Restarts affected kernels in parallel -- Shows user-friendly notifications about the restart - -### How It Works - -1. **Configuration Change Detection** - - When a user saves or deletes an integration configuration in the webview - - `IntegrationStorage.save()` or `IntegrationStorage.delete()` is called - - This fires the `onDidChangeIntegrations` event - -2. **Event Handling** - - `IntegrationKernelRestartHandler` receives the event - - It scans all open notebook documents with type `'deepnote'` - - For each notebook, it checks if there's a running kernel - - It examines cells for `sql_integration_id` metadata to determine if the notebook uses SQL integrations - -3. **Kernel Restart** - - Kernels for notebooks that use SQL integrations are restarted using `kernel.restart()` - - Restarts happen in parallel for better performance - - User receives a notification: "Integration configuration updated. N kernel(s) restarted to apply changes." - -4. **Credential Injection** - - When the kernel restarts, `SqlIntegrationStartupCodeProvider` automatically injects the new credentials - - Environment variables are updated with the new integration configurations - - SQL cells can immediately use the updated credentials - -### Architecture Changes - -``` -┌─────────────────────────────────────────────────────────────────┐ -│ User saves integration config in webview │ -└──────────────────────┬──────────────────────────────────────────┘ - │ - ▼ -┌─────────────────────────────────────────────────────────────────┐ -│ IntegrationStorage.save(config) │ -│ - Stores to encrypted storage (VSCode SecretStorage API) │ -│ - Updates in-memory cache │ -│ - Fires onDidChangeIntegrations event ────────────────────┐ │ -└──────────────────────────────────────────────────────────┼─────┘ - │ - ┌────────────────────────────────────────────────┤ - │ │ - ▼ ▼ -┌──────────────────────────────────┐ ┌──────────────────────────────────┐ -│ SqlIntegrationEnvironmentVars │ │ IntegrationKernelRestartHandler │ -│ VariablesProvider │ │ (NEW) │ -│ - Fires onDidChangeEnvVars │ │ - Scans open notebooks │ -│ │ │ - Finds kernels using SQL │ -│ │ │ - Restarts affected kernels │ -└──────────────────────────────────┘ └─────────────┬────────────────────┘ - │ - ▼ - ┌──────────────────────────────┐ - │ Kernel restarts │ - │ - SqlIntegrationStartup │ - │ CodeProvider injects │ - │ new credentials │ - │ - SQL cells work with new │ - │ credentials │ - └──────────────────────────────┘ -``` - -### Service Registration - -The service is registered in both node and web environments: -- `src/notebooks/serviceRegistry.node.ts` -- `src/notebooks/serviceRegistry.web.ts` - -Registered as an `IExtensionSyncActivationService`, which means it's automatically activated when the extension loads. - -## Benefits - -1. **Seamless Experience**: Users don't need to manually restart kernels after changing credentials -2. **Immediate Effect**: New credentials are available immediately after configuration -3. **Smart Detection**: Only restarts kernels that actually use SQL integrations -4. **User Feedback**: Clear notifications inform users about the restart -5. **Error Resilient**: If one kernel fails to restart, others continue - -## Technical Details - -### Dependencies -- `IIntegrationStorage`: To listen for configuration changes -- `IKernelProvider`: To access and restart kernels -- `IDisposableRegistry`: To manage event subscriptions - -### Key Methods - -**`onIntegrationConfigurationChanged()`** -- Main handler for integration changes -- Prevents concurrent restart attempts using `isRestarting` flag -- Scans workspace notebooks for Deepnote notebooks with running kernels -- Filters to only notebooks that use SQL integrations - -**`notebookUsesSqlIntegrations(notebook)`** -- Scans notebook cells for SQL language -- Checks cell metadata for `sql_integration_id` -- Excludes internal DuckDB integration (`deepnote-dataframe-sql`) -- Returns true if notebook uses external SQL integrations - -### Error Handling -- Individual kernel restart failures don't stop other restarts -- Errors are logged but don't throw exceptions -- User is still notified of successful restarts - -## Testing - -The service can be tested similar to `SqlCellStatusBarProvider`: -1. Mock `IIntegrationStorage` with an `EventEmitter` for `onDidChangeIntegrations` -2. Mock `IKernelProvider` to return test kernels -3. Fire the event and verify `kernel.restart()` is called for appropriate notebooks - -Example test structure: -```typescript -test('restarts kernels when integration changes', async () => { - const onDidChangeIntegrations = new EventEmitter(); - when(integrationStorage.onDidChangeIntegrations).thenReturn(onDidChangeIntegrations.event); - - handler.activate(); - - // Fire integration change event - onDidChangeIntegrations.fire(); - - // Verify kernel.restart() was called - verify(mockKernel.restart()).once(); -}); -``` - -## Future Enhancements - -Potential improvements: -1. **Selective Restart**: Only restart kernels that use the specific integration that changed -2. **Confirmation Dialog**: Ask user before restarting (optional setting) -3. **Restart Queue**: Batch multiple integration changes to avoid multiple restarts -4. **Active Execution Check**: Warn if cells are currently executing -5. **Kernel State Preservation**: Try to preserve kernel variables across restart (advanced) - -## Related Files - -- `src/notebooks/deepnote/integrations/integrationWebview.ts` - Webview that triggers config saves -- `src/platform/notebooks/deepnote/integrationStorage.ts` - Storage layer that fires events -- `src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts` - Environment variable provider -- `src/notebooks/deepnote/integrations/sqlIntegrationStartupCodeProvider.ts` - Injects credentials on kernel start -- `src/kernels/kernel.ts` - Kernel restart implementation - diff --git a/src/kernels/deepnote/deepnoteIntegrationEndpointEnv.ts b/src/kernels/deepnote/deepnoteIntegrationEndpointEnv.ts new file mode 100644 index 0000000000..3da80bdbda --- /dev/null +++ b/src/kernels/deepnote/deepnoteIntegrationEndpointEnv.ts @@ -0,0 +1,47 @@ +import { Uri } from 'vscode'; + +import { resolveProjectIdForFile } from '../../platform/deepnote/deepnoteProjectIdResolver'; +import { logger } from '../../platform/logging'; +import { IUserpodApiEndpoints } from '../../platform/notebooks/deepnote/types'; + +/** + * Injects live-integration env vars into `extraEnv` when the loopback endpoint is listening and + * `deepnoteFileUri` resolves to a project id. Skipped otherwise — the toolkit raises on an unreachable URL. + * + * Mutates `extraEnv` in place. + */ +export async function applyIntegrationEndpointEnv({ + deepnoteFileUri, + endpoint, + extraEnv +}: { + deepnoteFileUri: Uri; + endpoint: IUserpodApiEndpoints; + extraEnv: Record; +}): Promise { + // Wait for the initial bind so a kernel starting before the loopback endpoint is listening still gets the env. + await endpoint.ready; + + const baseUrl = endpoint.baseUrl; + if (!baseUrl) { + logger.warn( + 'applyIntegrationEndpointEnv: integration endpoint is not listening; skipping live integration env injection.' + ); + + return; + } + + const projectId = await resolveProjectIdForFile(deepnoteFileUri); + + if (!projectId) { + return; + } + + extraEnv['DEEPNOTE_RUNTIME__ENV_INTEGRATION_ENABLED'] = 'true'; + extraEnv['DEEPNOTE_RUNTIME__RUNNING_IN_DETACHED_MODE'] = 'true'; + extraEnv['DEEPNOTE_RUNTIME__WEBAPP_URL'] = baseUrl; + // 2.1.1 dereferences project_secret without a null-check in detached mode; also the endpoint's per-project bearer token. + extraEnv['DEEPNOTE_RUNTIME__PROJECT_SECRET'] = endpoint.getAuthToken(projectId); + // Legacy key (not __PROJECT_ID): also satisfies set_notebook_path's has_env check, avoiding a session-name parse. + extraEnv['DEEPNOTE_PROJECT_ID'] = projectId; +} diff --git a/src/kernels/deepnote/deepnoteIntegrationEndpointEnv.unit.test.ts b/src/kernels/deepnote/deepnoteIntegrationEndpointEnv.unit.test.ts new file mode 100644 index 0000000000..6f4d741738 --- /dev/null +++ b/src/kernels/deepnote/deepnoteIntegrationEndpointEnv.unit.test.ts @@ -0,0 +1,84 @@ +import { assert } from 'chai'; +import { anything, verify } from 'ts-mockito'; +import { Uri } from 'vscode'; + +import { serializeProjectFile } from '../../notebooks/deepnote/deepnoteTestHelpers'; +import { IUserpodApiEndpoints } from '../../platform/notebooks/deepnote/types'; +import { stubReadFile } from '../../test/mocks/vscodeFs'; +import { resetVSCodeMocks } from '../../test/vscode-mock'; +import { applyIntegrationEndpointEnv } from './deepnoteIntegrationEndpointEnv'; + +/** + * The five integration env vars are only injected when the loopback endpoint is listening AND the + * file resolves to a project id; every other path leaves `extraEnv` unchanged. + */ +suite('applyIntegrationEndpointEnv', () => { + const projectFileUri = Uri.file('/workspace/project/notebook-a.deepnote'); + const baseUrl = 'http://127.0.0.1:5555'; + // A pre-seeded key stands in for env a caller may already hold; the helper must add to it, not replace it. + const sqlEnvKey = 'SQL_DEEPNOTE_INTEGRATION_ABC'; + const sqlEnvValue = 'postgres://localhost:5432/db'; + + setup(() => { + resetVSCodeMocks(); + }); + + function createEndpoint(endpointBaseUrl: string | undefined): IUserpodApiEndpoints { + return { + baseUrl: endpointBaseUrl, + ready: Promise.resolve(), + getAuthToken: () => 'endpoint-token' + }; + } + + test('injects all five integration env vars (preserving pre-existing keys) when the endpoint is listening and the file has a project id', async () => { + const mockFs = stubReadFile(serializeProjectFile('the-project-id')); + const extraEnv: Record = { [sqlEnvKey]: sqlEnvValue }; + + await applyIntegrationEndpointEnv({ + deepnoteFileUri: projectFileUri, + endpoint: createEndpoint(baseUrl), + extraEnv + }); + + assert.deepStrictEqual(extraEnv, { + [sqlEnvKey]: sqlEnvValue, + DEEPNOTE_RUNTIME__ENV_INTEGRATION_ENABLED: 'true', + DEEPNOTE_RUNTIME__RUNNING_IN_DETACHED_MODE: 'true', + DEEPNOTE_RUNTIME__WEBAPP_URL: baseUrl, + DEEPNOTE_RUNTIME__PROJECT_SECRET: 'endpoint-token', + DEEPNOTE_PROJECT_ID: 'the-project-id' + }); + // The enabled path must resolve the project id from the file. + verify(mockFs.readFile(anything())).once(); + }); + + test('injects nothing and does NOT read the file when the endpoint has no baseUrl', async () => { + const mockFs = stubReadFile(serializeProjectFile('the-project-id')); + const extraEnv: Record = { [sqlEnvKey]: sqlEnvValue }; + + await applyIntegrationEndpointEnv({ + deepnoteFileUri: projectFileUri, + endpoint: createEndpoint(undefined), + extraEnv + }); + + assert.deepStrictEqual(extraEnv, { [sqlEnvKey]: sqlEnvValue }); + // A missing baseUrl short-circuits BEFORE resolving the project id — no file read. + verify(mockFs.readFile(anything())).never(); + }); + + test('injects nothing when the endpoint is listening but the file resolves to no project id', async () => { + // A schema-valid `.deepnote` whose project.id is empty — `resolveProjectIdForFile` yields a falsy id. + stubReadFile(serializeProjectFile('')); + const extraEnv: Record = { [sqlEnvKey]: sqlEnvValue }; + + await applyIntegrationEndpointEnv({ + deepnoteFileUri: projectFileUri, + endpoint: createEndpoint(baseUrl), + extraEnv + }); + + assert.deepStrictEqual(extraEnv, { [sqlEnvKey]: sqlEnvValue }); + }); +}); diff --git a/src/kernels/deepnote/deepnoteLspClientManager.node.ts b/src/kernels/deepnote/deepnoteLspClientManager.node.ts index 907adf8f8f..fd464cd1f9 100644 --- a/src/kernels/deepnote/deepnoteLspClientManager.node.ts +++ b/src/kernels/deepnote/deepnoteLspClientManager.node.ts @@ -24,11 +24,10 @@ import { logger } from '../../platform/logging'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; import { noop } from '../../platform/common/utils/misc'; import { - IIntegrationStorage, IPlatformNotebookEditorProvider, - IPlatformDeepnoteNotebookManager + IPlatformDeepnoteNotebookManager, + ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; -import { ConfigurableDatabaseIntegrationConfig } from '../../platform/notebooks/deepnote/integrationTypes'; import { SqlLspConnection, isSupportedBySqlLsp, convertToSqlLspConnection } from './sqlLspConnectionUtils'; interface LspClientInfo { @@ -61,10 +60,11 @@ export class DeepnoteLspClientManager constructor( @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, - @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, @inject(IPlatformNotebookEditorProvider) private readonly notebookEditorProvider: IPlatformNotebookEditorProvider, - @inject(IPlatformDeepnoteNotebookManager) private readonly notebookManager: IPlatformDeepnoteNotebookManager + @inject(IPlatformDeepnoteNotebookManager) private readonly notebookManager: IPlatformDeepnoteNotebookManager, + @inject(ISqlIntegrationEnvVarsProvider) + private readonly sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider ) { this.disposables.push(this); } @@ -623,18 +623,16 @@ export class DeepnoteLspClientManager logger.trace(`SQL LSP: Found ${projectIntegrations.length} integrations in project ${projectId}`); + // Merged (SecretStorage + `.deepnote.env.yaml`) configs, so file-configured databases also get + // LSP autocomplete/schema. const projectIntegrationConfigs = ( - await Promise.all( - projectIntegrations.map((integration) => - this.integrationStorage.getIntegrationConfig(integration.id) - ) - ) - ).filter((config): config is ConfigurableDatabaseIntegrationConfig => config != null); + await this.sqlIntegrationEnvVars.getMergedIntegrationConfigs(notebookUri) + ).filter((config) => config.type !== 'pandas-dataframe'); const connections = projectIntegrationConfigs .filter((config) => isSupportedBySqlLsp(config.type)) .map((config) => convertToSqlLspConnection(config)) - .filter((conn): conn is SqlLspConnection => conn !== null); + .filter((conn) => conn !== null); logger.trace( `SQL LSP: Found ${connections.length} SQL LSP-compatible integrations for project ${projectId}` diff --git a/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts b/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts index fc2838bf66..bebbeff2b0 100644 --- a/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts +++ b/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts @@ -1,4 +1,5 @@ import { assert } from 'chai'; +import { anything, instance, mock, when } from 'ts-mockito'; import { Uri } from 'vscode'; import { DeepnoteLspClientManager } from './deepnoteLspClientManager.node'; @@ -6,7 +7,7 @@ import { createMockChildProcess } from './deepnoteTestHelpers.node'; import { IDisposableRegistry } from '../../platform/common/types'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import * as path from '../../platform/vscode-path/path'; -import { noop } from '../../platform/common/utils/misc'; +import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; suite('DeepnoteLspClientManager Integration Tests', () => { let lspClientManager: DeepnoteLspClientManager; @@ -29,16 +30,6 @@ suite('DeepnoteLspClientManager Integration Tests', () => { dispose: () => Promise.resolve() } as any; - // Mock integration storage - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const mockIntegrationStorage = { - getAll: async () => [], - getIntegrationConfig: async () => undefined, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - onDidChangeIntegrations: { dispose: noop } as any, - dispose: noop - } as any; - // Mock notebook editor provider // eslint-disable-next-line @typescript-eslint/no-explicit-any const mockNotebookEditorProvider = { @@ -52,11 +43,15 @@ suite('DeepnoteLspClientManager Integration Tests', () => { } as any; setup(() => { + // No integrations configured, so the LSP resolves an empty connection list. + const sqlIntegrationEnvVars = mock(); + when(sqlIntegrationEnvVars.getMergedIntegrationConfigs(anything())).thenResolve([]); + lspClientManager = new DeepnoteLspClientManager( mockDisposableRegistry, - mockIntegrationStorage, mockNotebookEditorProvider, - mockNotebookManager + mockNotebookManager, + instance(sqlIntegrationEnvVars) ); lspClientManager.activate(); }); diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 112f7f0e09..fb469b694f 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -6,7 +6,7 @@ */ import * as fs from 'fs-extra'; -import { inject, injectable, named, optional } from 'inversify'; +import { inject, injectable, named } from 'inversify'; import * as os from 'os'; import { CancellationToken, l10n, Uri } from 'vscode'; @@ -21,11 +21,12 @@ import { sleep } from '../../platform/common/utils/async'; import { generateUuid } from '../../platform/common/uuid'; import { DeepnoteServerStartupError } from '../../platform/errors/deepnoteKernelErrors'; import { logger } from '../../platform/logging'; -import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; +import { IUserpodApiEndpoints } from '../../platform/notebooks/deepnote/types'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import * as path from '../../platform/vscode-path/path'; -import { DeepnoteServerInfo, IDeepnoteServerStarter, IDeepnoteToolkitInstaller } from './types'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; +import { applyIntegrationEndpointEnv } from './deepnoteIntegrationEndpointEnv'; +import { DeepnoteServerInfo, IDeepnoteServerStarter, IDeepnoteToolkitInstaller } from './types'; const MAX_OUTPUT_TRACKING_LENGTH = 5000; const SERVER_STARTUP_TIMEOUT_MS = 120_000; @@ -57,8 +58,8 @@ interface ProjectContext { * * Uses @deepnote/runtime-core's `startServer`/`stopServer` for the core server * lifecycle (process spawn, port discovery, health checks, shutdown), and layers - * extension-specific concerns on top: lock files, orphan cleanup, SQL integration - * env vars, output channel logging, and multi-server concurrency control. + * extension-specific concerns on top: lock files, orphan cleanup, integration + * endpoint env vars, output channel logging, and multi-server concurrency control. */ @injectable() export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtensionSyncActivationService { @@ -75,9 +76,8 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension @inject(DeepnoteAgentSkillsManager) private readonly agentSkillsManager: DeepnoteAgentSkillsManager, @inject(IOutputChannel) @named(STANDARD_OUTPUT_CHANNEL) private readonly outputChannel: IOutputChannel, @inject(IAsyncDisposableRegistry) asyncRegistry: IAsyncDisposableRegistry, - @inject(ISqlIntegrationEnvVarsProvider) - @optional() - private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider + @inject(IUserpodApiEndpoints) + private readonly userpodApiEndpoints: IUserpodApiEndpoints ) { asyncRegistry.push(this); } @@ -224,7 +224,8 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension * * Extension-specific layers: * - Toolkit/venv installation (before start) - * - SQL integration env var injection (via ServerOptions.env) + * - Integration endpoint env var injection (via ServerOptions.env) — these point the toolkit at the + * extension's loopback `userpod-api` endpoint, which is how it fetches SQL credentials at kernel init * - Lock file creation (after start, using returned PID) * - Output channel logging (via process stdout/stderr streams) */ @@ -261,7 +262,13 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension logger.info(`Starting deepnote-toolkit server for ${fileKey} (environmentId ${environmentId})`); this.outputChannel.appendLine(l10n.t('Starting Deepnote server...')); - const extraEnv = await this.gatherSqlIntegrationEnvVars(deepnoteFileUri, environmentId, token); + const extraEnv: Record = {}; + + await applyIntegrationEndpointEnv({ + deepnoteFileUri, + endpoint: this.userpodApiEndpoints, + extraEnv + }); // Initialize output tracking for error reporting this.serverOutputByFile.set(fileKey, { stdout: '', stderr: '' }); @@ -362,41 +369,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } - /** - * Gather SQL integration environment variables for the deepnote-toolkit server. - */ - private async gatherSqlIntegrationEnvVars( - deepnoteFileUri: Uri, - environmentId: string, - token?: CancellationToken - ): Promise> { - const extraEnv: Record = {}; - - if (!this.sqlIntegrationEnvVars) { - logger.debug('DeepnoteServerStarter: SqlIntegrationEnvironmentVariablesProvider not available'); - return extraEnv; - } - - const fileKey = deepnoteFileUri.fsPath; - - logger.debug( - `DeepnoteServerStarter: Injecting SQL integration env vars for ${fileKey} with environmentId ${environmentId}` - ); - 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(extraEnv, 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); - } - - return extraEnv; - } - /** * Stream stdout/stderr from the server process to the VSCode output channel. */ diff --git a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts index d0eb9070cc..d5547f1306 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts @@ -2,20 +2,23 @@ import { assert } from 'chai'; import * as fakeTimers from '@sinonjs/fake-timers'; import * as sinon from 'sinon'; import { anything, instance, mock, when } from 'ts-mockito'; -import { CancellationError, Uri } from 'vscode'; +import { Uri } from 'vscode'; -import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; -import { DeepnoteServerStarter } from './deepnoteServerStarter.node'; +import { serializeProjectFile } from '../../notebooks/deepnote/deepnoteTestHelpers'; import { IProcessServiceFactory } from '../../platform/common/process/types.node'; import { IAsyncDisposableRegistry, IOutputChannel } from '../../platform/common/types'; -import { IDeepnoteToolkitInstaller } from './types'; -import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; +import { IUserpodApiEndpoints } from '../../platform/notebooks/deepnote/types'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { __getStartServerCalls, __getStopServerCalls, __resetRuntimeCoreMock } from '../../test/mocks/deepnoteRuntimeCore'; +import { stubReadFile } from '../../test/mocks/vscodeFs'; +import { resetVSCodeMocks } from '../../test/vscode-mock'; +import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; +import { DeepnoteServerStarter } from './deepnoteServerStarter.node'; +import { IDeepnoteToolkitInstaller } from './types'; /** * Unit tests for DeepnoteServerStarter. @@ -41,21 +44,25 @@ suite('DeepnoteServerStarter', () => { let mockAgentSkillsManager: DeepnoteAgentSkillsManager; let mockOutputChannel: IOutputChannel; let mockAsyncRegistry: IAsyncDisposableRegistry; - let mockSqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider; + let mockUserpodApiEndpoints: IUserpodApiEndpoints; setup(() => { __resetRuntimeCoreMock(); + resetVSCodeMocks(); mockProcessServiceFactory = mock(); mockToolkitInstaller = mock(); mockAgentSkillsManager = mock(); mockOutputChannel = mock(); mockAsyncRegistry = mock(); - mockSqlIntegrationEnvVars = mock(); + mockUserpodApiEndpoints = mock(); when(mockAsyncRegistry.push(anything())).thenReturn(); when(mockOutputChannel.appendLine(anything())).thenReturn(); + when(mockUserpodApiEndpoints.ready).thenReturn(Promise.resolve()); + when(mockUserpodApiEndpoints.baseUrl).thenReturn(undefined); + // The toolkit install step runs before runtime-core's startServer; stub it so the // start path reaches startServer. (ts-mockito methods that are not stubbed return null.) when(mockToolkitInstaller.ensureVenvAndToolkit(anything(), anything(), anything(), anything())).thenResolve({ @@ -71,7 +78,7 @@ suite('DeepnoteServerStarter', () => { instance(mockAgentSkillsManager), instance(mockOutputChannel), instance(mockAsyncRegistry), - instance(mockSqlIntegrationEnvVars) + instance(mockUserpodApiEndpoints) ); }); @@ -80,65 +87,42 @@ suite('DeepnoteServerStarter', () => { await serverStarter.dispose(); }); - suite('SQL integration env vars', () => { - test('starts the server without SQL env vars when no provider is available', async () => { - // The provider is @optional() — construct a starter without it. - const starterWithoutSql = new DeepnoteServerStarter( - instance(mockProcessServiceFactory), - instance(mockToolkitInstaller), - instance(mockAgentSkillsManager), - instance(mockOutputChannel), - instance(mockAsyncRegistry) + suite('integration endpoint env vars', () => { + // The env shape and the empty-env paths belong to applyIntegrationEndpointEnv and are covered by + // deepnoteIntegrationEndpointEnv.unit.test.ts; this test covers the per-notebook wiring only. + test('does NOT leak one project id or bearer token into a sibling notebook server', async () => { + // Sibling files in the same directory resolving to DIFFERENT projects, each with its own token. + stubReadFile((uri) => + uri.toString() === uriA.toString() + ? serializeProjectFile('project-a') + : serializeProjectFile('project-b') ); - - try { - await starterWithoutSql.startServer(interpreter, venvPath, true, [], 'env1', uriA); - - assert.deepStrictEqual( - __getStartServerCalls().map((c) => c.env), - [{}] - ); - } finally { - await starterWithoutSql.dispose(); - } - }); - - test('starts the server without SQL env vars when the provider rejects with a cancellation error', async () => { - when(mockSqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenReject( - new CancellationError() - ); - - await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); - - assert.deepStrictEqual( - __getStartServerCalls().map((c) => c.env), - [{}] - ); - }); - - test('forwards the SQL integration provider env vars into the started server', async () => { - when(mockSqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({ FOO: 'bar' }); - - await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); - - assert.deepStrictEqual( - __getStartServerCalls().map((c) => c.env), - [{ FOO: 'bar' }] - ); - }); - - test('does NOT leak one notebook SQL env vars into a sibling whose provider yields none', async () => { - // The provider is keyed by notebook URI: it yields env vars for A but nothing for its sibling B. - when(mockSqlIntegrationEnvVars.getEnvironmentVariables(uriA, anything())).thenResolve({ FOO: 'bar' }); - when(mockSqlIntegrationEnvVars.getEnvironmentVariables(uriB, anything())).thenResolve({}); + when(mockUserpodApiEndpoints.baseUrl).thenReturn('http://127.0.0.1:5555'); + when(mockUserpodApiEndpoints.getAuthToken('project-a')).thenReturn('token-a'); + when(mockUserpodApiEndpoints.getAuthToken('project-b')).thenReturn('token-b'); await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriA); await serverStarter.startServer(interpreter, venvPath, true, [], 'env1', uriB); assert.deepStrictEqual( __getStartServerCalls().map((c) => c.env), - [{ FOO: 'bar' }, {}], - "notebook B's server must NOT inherit A's SQL env vars" + [ + { + DEEPNOTE_RUNTIME__ENV_INTEGRATION_ENABLED: 'true', + DEEPNOTE_RUNTIME__RUNNING_IN_DETACHED_MODE: 'true', + DEEPNOTE_RUNTIME__WEBAPP_URL: 'http://127.0.0.1:5555', + DEEPNOTE_RUNTIME__PROJECT_SECRET: 'token-a', + DEEPNOTE_PROJECT_ID: 'project-a' + }, + { + DEEPNOTE_RUNTIME__ENV_INTEGRATION_ENABLED: 'true', + DEEPNOTE_RUNTIME__RUNNING_IN_DETACHED_MODE: 'true', + DEEPNOTE_RUNTIME__WEBAPP_URL: 'http://127.0.0.1:5555', + DEEPNOTE_RUNTIME__PROJECT_SECRET: 'token-b', + DEEPNOTE_PROJECT_ID: 'project-b' + } + ], + "notebook B's server must carry ONLY project B's id and bearer token" ); }); }); @@ -234,11 +218,11 @@ suite('DeepnoteServerStarter', () => { }); test('waits for an in-flight start operation before completing', async () => { - // Park the start mid-flight: its SQL env-var gathering resolves only via releaseStart. + // Park the start mid-flight: the integration endpoint's readiness settles only via releaseStart. let releaseStart!: () => void; - when(mockSqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenReturn( - new Promise((resolve) => { - releaseStart = () => resolve({}); + when(mockUserpodApiEndpoints.ready).thenReturn( + new Promise((resolve) => { + releaseStart = resolve; }) ); diff --git a/src/kernels/deepnote/sqlLspConnectionUtils.ts b/src/kernels/deepnote/sqlLspConnectionUtils.ts index cece95d409..aaa65529e9 100644 --- a/src/kernels/deepnote/sqlLspConnectionUtils.ts +++ b/src/kernels/deepnote/sqlLspConnectionUtils.ts @@ -1,4 +1,7 @@ -import { ConfigurableDatabaseIntegrationConfig } from '../../platform/notebooks/deepnote/integrationTypes'; +import { + ConfigurableDatabaseIntegrationConfig, + isFederatedAuthMetadata +} from '../../platform/notebooks/deepnote/integrationTypes'; /** * SQL LSP connection configuration format expected by sql-language-server @@ -36,39 +39,31 @@ export function isSupportedBySqlLsp(type: string): boolean { */ export function convertToSqlLspConnection(config: ConfigurableDatabaseIntegrationConfig): SqlLspConnection | null { try { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const metadata = config.metadata as Record; - const type = config.type as string; + if (isFederatedAuthMetadata(config.metadata)) { + return null; + } - switch (type) { + switch (config.type) { case 'pgsql': return { name: config.name || 'postgres', adapter: 'postgres', - host: metadata.host || 'localhost', - port: Number(metadata.port) || 5432, - user: metadata.username || metadata.user, - password: metadata.password, - database: metadata.database || metadata.dbname + host: config.metadata.host || 'localhost', + port: Number(config.metadata.port) || 5432, + user: config.metadata.user, + password: config.metadata.password, + database: config.metadata.database }; case 'mysql': return { name: config.name || 'mysql', adapter: 'mysql', - host: metadata.host || 'localhost', - port: Number(metadata.port) || 3306, - user: metadata.username || metadata.user, - password: metadata.password, - database: metadata.database || metadata.dbname - }; - - case 'big-query': - return { - name: config.name || 'bigquery', - adapter: 'bigquery', - projectId: metadata.projectId || metadata.project_id, - keyFilename: metadata.keyFilename || metadata.key_filename + host: config.metadata.host || 'localhost', + port: Number(config.metadata.port) || 3306, + user: config.metadata.user, + password: config.metadata.password, + database: config.metadata.database }; default: diff --git a/src/kernels/deepnote/sqlLspConnectionUtils.unit.test.ts b/src/kernels/deepnote/sqlLspConnectionUtils.unit.test.ts index 4ac5f1cc14..e19c6fe59a 100644 --- a/src/kernels/deepnote/sqlLspConnectionUtils.unit.test.ts +++ b/src/kernels/deepnote/sqlLspConnectionUtils.unit.test.ts @@ -42,7 +42,7 @@ suite('SQL LSP Connection Utils Unit Tests', () => { const config = createTestConfig('pgsql', 'My Postgres', { host: 'db.example.com', port: '5433', - username: 'admin', + user: 'admin', password: 'secret123', database: 'mydb' }); @@ -76,21 +76,6 @@ suite('SQL LSP Connection Utils Unit Tests', () => { }); }); - test('should handle alternative field names (user, dbname)', () => { - const config = createTestConfig('pgsql', 'Alt Fields', { - host: 'localhost', - port: 5432, - user: 'altuser', - password: 'pass', - dbname: 'altdb' - }); - - const result = convertToSqlLspConnection(config); - - assert.strictEqual(result?.user, 'altuser'); - assert.strictEqual(result?.database, 'altdb'); - }); - test('should convert string port to number', () => { const config = createTestConfig('pgsql', 'String Port', { port: '5432' @@ -108,7 +93,7 @@ suite('SQL LSP Connection Utils Unit Tests', () => { const config = createTestConfig('mysql', 'My MySQL', { host: 'mysql.example.com', port: '3307', - username: 'root', + user: 'root', password: 'rootpass', database: 'testdb' }); @@ -135,8 +120,8 @@ suite('SQL LSP Connection Utils Unit Tests', () => { }); }); - suite('BigQuery conversion', () => { - test('should convert big-query config with all fields', () => { + suite('BigQuery', () => { + test('should return null for big-query (not converted to SQL LSP connection)', () => { const config = createTestConfig('big-query', 'My BigQuery', { projectId: 'my-gcp-project', keyFilename: '/path/to/key.json' @@ -144,24 +129,32 @@ suite('SQL LSP Connection Utils Unit Tests', () => { const result = convertToSqlLspConnection(config); - assert.deepStrictEqual(result, { - name: 'My BigQuery', - adapter: 'bigquery', - projectId: 'my-gcp-project', - keyFilename: '/path/to/key.json' + assert.isNull(result); + }); + }); + + suite('Federated auth', () => { + test('should return null for federated-auth metadata', () => { + const config = createTestConfig('big-query', 'Federated BigQuery', { + authMethod: 'google-oauth' }); + + const result = convertToSqlLspConnection(config); + + assert.isNull(result); }); - test('should handle alternative field names (project_id, key_filename)', () => { - const config = createTestConfig('big-query', 'Alt BigQuery', { - project_id: 'alt-project', - key_filename: '/alt/path.json' + test('should return null for federated-auth pgsql metadata', () => { + const config = createTestConfig('pgsql', 'Federated Postgres', { + authMethod: 'google-oauth', + host: 'db.example.com', + user: 'admin', + database: 'mydb' }); const result = convertToSqlLspConnection(config); - assert.strictEqual(result?.projectId, 'alt-project'); - assert.strictEqual(result?.keyFilename, '/alt/path.json'); + assert.isNull(result); }); }); @@ -191,12 +184,9 @@ suite('SQL LSP Connection Utils Unit Tests', () => { // Missing metadata entirely } as ConfigurableDatabaseIntegrationConfig; - // Should not throw const result = convertToSqlLspConnection(config); - // Behavior depends on implementation - either null or partial result - // The important thing is it doesn't crash - assert.ok(result === null || typeof result === 'object'); + assert.isNull(result); }); }); }); diff --git a/src/kernels/execution/cellExecution.federatedAuth.unit.test.ts b/src/kernels/execution/cellExecution.federatedAuth.unit.test.ts index 54323afcb6..24a15dd870 100644 --- a/src/kernels/execution/cellExecution.federatedAuth.unit.test.ts +++ b/src/kernels/execution/cellExecution.federatedAuth.unit.test.ts @@ -168,15 +168,22 @@ suite('CellExecution federated-auth branch', () => { assertMainExecuteShape(calls[0]); }); - test('when generate() returns undefined: single requestExecute with silent=false, store_history=true', async () => { - const generator: IFederatedAuthSqlBlockCodeGenerator = { - generate: sinon.stub().resolves(undefined) - }; + test("generate() is passed the converted block and the executing cell's own notebook URI", async () => { + // Catches: dropping the URI (or sourcing it from the active editor) — the integration config is resolved + // per notebook from `.deepnote.env.yaml` merged over SecretStorage, so the wrong notebook resolves wrong. + const generateStub = sinon.stub().resolves(undefined); + const generator: IFederatedAuthSqlBlockCodeGenerator = { generate: generateStub }; const execution = createExecution(generator); + await execution.start(instance(session)); await execution.result.catch(() => undefined); - sinon.assert.calledOnce(generator.generate as sinon.SinonStub); + sinon.assert.calledOnce(generateStub); + const [block, notebookUri] = generateStub.firstCall.args; + assert.isObject(block, 'first argument must be the converted Deepnote block'); + assert.strictEqual(notebookUri, cell.notebook.uri); + + // `undefined` from the generator falls back to `createPythonCode` — still exactly one execute. const calls = requestExecuteSpy.getCalls(); assert.strictEqual(calls.length, 1, `expected exactly 1 requestExecute call, got ${calls.length}`); assertMainExecuteShape(calls[0]); diff --git a/src/kernels/execution/cellExecution.ts b/src/kernels/execution/cellExecution.ts index ac417abd04..d35a8fcebf 100644 --- a/src/kernels/execution/cellExecution.ts +++ b/src/kernels/execution/cellExecution.ts @@ -483,7 +483,10 @@ export class CellExecution implements ICellExecution, IDisposable { // Federated-auth (BigQuery + google-oauth): generator returns a single Python string with the connection JSON embedded as a literal (containing the fresh access token). `undefined` means non-federated or web — fall back to `createPythonCode`. let federatedCode: string | undefined; try { - federatedCode = await this.federatedAuthSqlBlockCodeGenerator?.generate(deepnoteBlock); + federatedCode = await this.federatedAuthSqlBlockCodeGenerator?.generate( + deepnoteBlock, + this.cell.notebook.uri + ); } catch (ex) { await this.handleFederatedGenerateError(ex); return; diff --git a/src/kernels/jupyter/session/jupyterKernelService.unit.test.ts b/src/kernels/jupyter/session/jupyterKernelService.unit.test.ts index af86e50ffa..fe58ee4754 100644 --- a/src/kernels/jupyter/session/jupyterKernelService.unit.test.ts +++ b/src/kernels/jupyter/session/jupyterKernelService.unit.test.ts @@ -30,6 +30,7 @@ import { ICustomEnvironmentVariablesProvider } from '../../../platform/common/va import { EnvironmentVariablesService } from '../../../platform/common/variables/environment.node'; import { isWeb } from '../../../platform/common/utils/misc'; import { isPythonKernelConnection } from '../../helpers'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; // eslint-disable-next-line suite('JupyterKernelService', () => { @@ -442,12 +443,16 @@ suite('JupyterKernelService', () => { const configService = mock(ConfigurationService); settings = mock(JupyterSettings); when(configService.getSettings(anything())).thenReturn(instance(settings)); + // These tests do not exercise SQL integrations; the provider just has to yield nothing. + const sqlIntegrationEnvVars = mock(); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); const kernelEnvService = new KernelEnvironmentVariablesService( instance(interpreterService), instance(appEnv), variablesService, instance(customEnvVars), - instance(configService) + instance(configService), + instance(sqlIntegrationEnvVars) ); testWorkspaceFolder = Uri.file(path.join(EXTENSION_ROOT_DIR, 'src', 'test', 'datascience')); const jupyterPaths = mock(); diff --git a/src/kernels/raw/launcher/kernelEnvVarsService.node.ts b/src/kernels/raw/launcher/kernelEnvVarsService.node.ts index 11b7eb4293..c4f6ce626f 100644 --- a/src/kernels/raw/launcher/kernelEnvVarsService.node.ts +++ b/src/kernels/raw/launcher/kernelEnvVarsService.node.ts @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -import { inject, injectable, optional } from 'inversify'; +import { inject, injectable } from 'inversify'; import { logger } from '../../../platform/logging'; import { getDisplayPath } from '../../../platform/common/platform/fs-paths.node'; import { IConfigurationService, Resource, type ReadWrite } from '../../../platform/common/types'; @@ -33,13 +33,8 @@ export class KernelEnvironmentVariablesService { private readonly customEnvVars: ICustomEnvironmentVariablesProvider, @inject(IConfigurationService) private readonly configService: IConfigurationService, @inject(ISqlIntegrationEnvVarsProvider) - @optional() - private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider - ) { - logger.debug( - `KernelEnvironmentVariablesService: Constructor; SQL env provider present=${!!sqlIntegrationEnvVars}` - ); - } + private readonly sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider + ) {} /** * Generates the environment variables for the kernel. * @@ -62,7 +57,7 @@ export class KernelEnvironmentVariablesService { logger.debug( `KernelEnvVarsService.getEnvironmentVariables: Called for resource ${ resource ? getDisplayPath(resource) : 'undefined' - }, sqlIntegrationEnvVars is ${this.sqlIntegrationEnvVars ? 'AVAILABLE' : 'UNDEFINED'}` + }` ); let kernelEnv = kernelSpec.env && Object.keys(kernelSpec.env).length > 0 @@ -97,21 +92,17 @@ export class KernelEnvironmentVariablesService { }) : undefined, this.sqlIntegrationEnvVars - ? this.sqlIntegrationEnvVars - .getEnvironmentVariables(resource, token) - .then((vars) => { - if (vars && Object.keys(vars).length > 0) { - logger.debug( - `KernelEnvVarsService: Got ${Object.keys(vars).length} SQL integration env vars` - ); - } - return vars; - }) - .catch((ex) => { - logger.error('Failed to get SQL integration env variables for Kernel', ex); - return undefined; - }) - : undefined + .getEnvironmentVariables(resource, token) + .then((vars) => { + if (vars && Object.keys(vars).length > 0) { + logger.debug(`KernelEnvVarsService: Got ${Object.keys(vars).length} SQL integration env vars`); + } + return vars; + }) + .catch((ex) => { + logger.error('Failed to get SQL integration env variables for Kernel', ex); + return undefined; + }) ]); if (token?.isCancellationRequested) { return; diff --git a/src/kernels/raw/launcher/kernelEnvVarsService.unit.test.ts b/src/kernels/raw/launcher/kernelEnvVarsService.unit.test.ts index b4dd91bdc5..a4934213ec 100644 --- a/src/kernels/raw/launcher/kernelEnvVarsService.unit.test.ts +++ b/src/kernels/raw/launcher/kernelEnvVarsService.unit.test.ts @@ -73,26 +73,15 @@ suite('Kernel Environment Variables Service', () => { teardown(() => Object.assign(process.env, originalEnvVars)); - /** - * Helper factory function to build KernelEnvironmentVariablesService with optional overrides. - * @param overrides Optional overrides for the service dependencies - * @returns A new instance of KernelEnvironmentVariablesService - */ - function buildKernelEnvVarsService(overrides?: { - sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider | undefined; - }): KernelEnvironmentVariablesService { - const sqlProvider = - overrides && 'sqlIntegrationEnvVars' in overrides - ? overrides.sqlIntegrationEnvVars - : instance(sqlIntegrationEnvVars); - + /** Builds the service under test from the suite's mocks. */ + function buildKernelEnvVarsService(): KernelEnvironmentVariablesService { return new KernelEnvironmentVariablesService( instance(interpreterService), instance(envActivation), variablesService, instance(customVariablesService), instance(configService), - sqlProvider + instance(sqlIntegrationEnvVars) ); } @@ -388,7 +377,7 @@ suite('Kernel Environment Variables Service', () => { ); }); - test('SQL integration env vars work when provider is undefined (optional dependency)', async () => { + test('SQL integration env vars work when the provider yields none (e.g. web)', async () => { const resource = Uri.file('test.ipynb'); when(envActivation.getActivatedEnvironmentVariables(anything(), anything(), anything())).thenResolve({ PATH: 'foobar' @@ -397,10 +386,9 @@ suite('Kernel Environment Variables Service', () => { customVariablesService.getCustomEnvironmentVariables(anything(), anything(), anything()) ).thenResolve(); - // Create service without SQL integration provider - const serviceWithoutSql = buildKernelEnvVarsService({ sqlIntegrationEnvVars: undefined }); + when(sqlIntegrationEnvVars.getEnvironmentVariables(anything(), anything())).thenResolve({}); - const vars = await serviceWithoutSql.getEnvironmentVariables(resource, interpreter, kernelSpec); + const vars = await kernelVariablesService.getEnvironmentVariables(resource, interpreter, kernelSpec); assert.isOk(vars); assert.isUndefined(vars!['SQL_MY_DB']); diff --git a/src/kernels/serviceRegistry.web.ts b/src/kernels/serviceRegistry.web.ts index 7779e8c213..5f276b6fb7 100644 --- a/src/kernels/serviceRegistry.web.ts +++ b/src/kernels/serviceRegistry.web.ts @@ -35,6 +35,8 @@ import { KernelStartupCodeProviders } from './kernelStartupCodeProviders.web'; import { LastCellExecutionTracker } from './execution/lastCellExecutionTracker'; import { ClearJupyterServersCommand } from './jupyter/clearJupyterServersCommand'; import { KernelChatStartupCodeProvider } from './chat/kernelStartupCodeProvider'; +import { SqlIntegrationEnvironmentVariablesProviderWeb } from '../platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.web'; +import { ISqlIntegrationEnvVarsProvider } from '../platform/notebooks/deepnote/types'; @injectable() class RawNotebookSupportedService implements IRawNotebookSupportedService { @@ -83,6 +85,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea ); serviceManager.addSingleton(IKernelFinder, KernelFinder); serviceManager.addSingleton(IKernelDependencyService, KernelDependencyService); + serviceManager.addSingleton( + ISqlIntegrationEnvVarsProvider, + SqlIntegrationEnvironmentVariablesProviderWeb + ); serviceManager.addSingleton( IExtensionSyncActivationService, diff --git a/src/messageTypes.ts b/src/messageTypes.ts index b9940320ce..e552da5378 100644 --- a/src/messageTypes.ts +++ b/src/messageTypes.ts @@ -169,6 +169,7 @@ export type LocalizedMessages = { integrationsTitle: string; integrationsNoIntegrationsFound: string; integrationsConnected: string; + integrationsConfiguredInFile: string; integrationsNotConfigured: string; integrationsConfigure: string; integrationsReconfigure: string; diff --git a/src/notebooks/deepnote/deepnoteTestHelpers.ts b/src/notebooks/deepnote/deepnoteTestHelpers.ts index c28ecc04bf..2f80248795 100644 --- a/src/notebooks/deepnote/deepnoteTestHelpers.ts +++ b/src/notebooks/deepnote/deepnoteTestHelpers.ts @@ -1,4 +1,4 @@ -import { DeepnoteFile } from '@deepnote/blocks'; +import { DeepnoteFile, serializeDeepnoteFile } from '@deepnote/blocks'; import { NotebookCell, NotebookCellKind, @@ -178,3 +178,8 @@ export function createDeepnoteFile(overrides: Partial = {}): Deepn export function createWorkspaceFolder(uri: Uri, index = 0): WorkspaceFolder { return { uri, name: uri.path.split('/').pop() ?? '', index }; } + +/** The serialized YAML of a minimal `.deepnote` file carrying `projectId`, for stubbing a file read. */ +export function serializeProjectFile(projectId: string): string { + return serializeDeepnoteFile(createDeepnoteFile({ project: createDeepnoteProject({ id: projectId }) })); +} diff --git a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthCommandHandler.node.ts b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthCommandHandler.node.ts index cc696a395c..fd4c27ecf1 100644 --- a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthCommandHandler.node.ts +++ b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthCommandHandler.node.ts @@ -8,7 +8,7 @@ import { Commands } from '../../../../platform/common/constants'; import { IExtensionContext } from '../../../../platform/common/types'; import { Integrations } from '../../../../platform/common/utils/localize'; import { logger } from '../../../../platform/logging'; -import { IIntegrationStorage } from '../../../../platform/notebooks/deepnote/types'; +import { ISqlIntegrationEnvVarsProvider } from '../../../../platform/notebooks/deepnote/types'; import { IFederatedAuthTokenStorage, type FederatedAuthTokenEntry } from '../types'; import { generateOAuthStateNonce, generatePkcePair } from './googleOAuthProvider.node'; import { computeMetadataFingerprint } from './federatedAuthTokenStorage.node'; @@ -30,21 +30,28 @@ export type RunOAuthFlowFn = (params: RunOAuthFlowParams) => Promise<{ refreshTo export class FederatedAuthCommandHandlerNode implements IExtensionSyncActivationService { constructor( @inject(IExtensionContext) private readonly extensionContext: IExtensionContext, - @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, + @inject(ISqlIntegrationEnvVarsProvider) + private readonly sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider, @inject(IFederatedAuthTokenStorage) private readonly tokenStorage: IFederatedAuthTokenStorage, private readonly runOAuthFlowFn: RunOAuthFlowFn = runOAuthFlow ) {} public activate(): void { this.extensionContext.subscriptions.push( - commands.registerCommand(Commands.AuthenticateIntegration, (integrationId: string) => - this.authenticate(integrationId) + commands.registerCommand(Commands.AuthenticateIntegration, (integrationId: string, resource: Uri) => + this.authenticate(integrationId, resource) ) ); } - /** Core flow. Public so tests can drive the handler without `commands.executeCommand`. */ - public async authenticate(integrationId: string): Promise { + /** + * Core flow. Public so tests can drive the handler without `commands.executeCommand`. + * + * `resource` is the notebook the request came from. Required, because the config is resolved per notebook + * from `.deepnote.env.yaml` merged over SecretStorage — falling back to whichever notebook happens to be + * focused would let the caller's eligibility check and this lookup mean different notebooks. + */ + public async authenticate(integrationId: string, resource: Uri): Promise { if (typeof integrationId !== 'string' || integrationId.length === 0) { logger.warn( `FederatedAuthCommandHandlerNode: invoked without a valid integrationId (received: ${String( @@ -54,7 +61,16 @@ export class FederatedAuthCommandHandlerNode implements IExtensionSyncActivation return; } - const integration = await this.integrationStorage.getIntegrationConfig(integrationId); + if (!resource) { + logger.warn( + `FederatedAuthCommandHandlerNode: invoked without a resource for integration "${integrationId}"` + ); + return; + } + + const integration = (await this.sqlIntegrationEnvVars.getMergedIntegrationConfigs(resource)).find( + (config) => config.id === integrationId + ); if (!integration) { logger.warn(`FederatedAuthCommandHandlerNode: integration "${integrationId}" not found.`); void window.showErrorMessage(Integrations.federatedAuthIntegrationNotFound(integrationId)); @@ -72,7 +88,7 @@ export class FederatedAuthCommandHandlerNode implements IExtensionSyncActivation const { clientId, clientSecret, project } = integration.metadata; const state = generateOAuthStateNonce(); const { challenge: codeChallenge, verifier: codeVerifier } = generatePkcePair(); - const deepnoteDomain = getDeepnoteDomain(getConfigurationResource()); + const deepnoteDomain = getDeepnoteDomain(resource); const proxyCallbackUrl = `https://${deepnoteDomain}/auth/bigquery/google-oauth-callback`; try { @@ -141,21 +157,14 @@ export class FederatedAuthCommandHandlerNode implements IExtensionSyncActivation } } -/** Reads the deepnote-host override (`deepnote.domain` setting); default `deepnote.com`. Mirrors `importClient.node.ts`. */ -function getDeepnoteDomain(resource?: Uri): string { +/** + * Reads the deepnote-host override (`deepnote.domain` setting); default `deepnote.com`. Mirrors + * `importClient.node.ts`. Scoped to the requesting notebook so workspace/folder overrides apply. + */ +function getDeepnoteDomain(resource: Uri): string { return workspace.getConfiguration('deepnote', resource).get('domain') ?? 'deepnote.com'; } -/** Prefer the active Deepnote notebook URI so workspace/folder `deepnote.domain` overrides apply. */ -function getConfigurationResource(): Uri | undefined { - const notebook = window.activeNotebookEditor?.notebook; - if (notebook?.notebookType === 'deepnote') { - return notebook.uri; - } - - return undefined; -} - /** Builds the proxy-start URL the user's browser will open. Public for unit tests. */ export function buildExtensionStartUrl(params: { clientId: string; diff --git a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthCommandHandler.node.unit.test.ts b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthCommandHandler.node.unit.test.ts index b9026bb7df..2b3dbaa574 100644 --- a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthCommandHandler.node.unit.test.ts +++ b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthCommandHandler.node.unit.test.ts @@ -2,13 +2,16 @@ import { assert } from 'chai'; import sinon from 'sinon'; import { CancellationError, Uri } from 'vscode'; import { anyString, anything, capture, deepEqual, instance, mock, verify, when } from 'ts-mockito'; +import type { DatabaseIntegrationConfig } from '@deepnote/database-integrations'; +import { Commands } from '../../../../platform/common/constants'; import { IExtensionContext, IDisposable } from '../../../../platform/common/types'; import { FederatedAuthCommandHandlerNode, buildExtensionStartUrl } from './federatedAuthCommandHandler.node'; import { IFederatedAuthTokenStorage } from '../types'; -import { IIntegrationStorage } from '../../../../platform/notebooks/deepnote/types'; +import { ISqlIntegrationEnvVarsProvider } from '../../../../platform/notebooks/deepnote/types'; import { computeMetadataFingerprint } from './federatedAuthTokenStorage.node'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../../test/vscode-mock'; +import { uriEquals } from '../../../../test/datascience/helpers'; import { FED_AUTH_FIXTURE, buildGoogleOauthIntegration, @@ -18,8 +21,14 @@ import { import type { RunOAuthFlowParams } from './oauthLoopbackFlow.node'; suite('FederatedAuthCommandHandlerNode', () => { + const NOTEBOOK_URI = Uri.file('/workspace/project.deepnote'); + const OTHER_NOTEBOOK_URI = Uri.file('/workspace/other/project.deepnote'); + const EXTERNAL_CALLBACK_URL = 'http://127.0.0.1:54321/auth/callback'; + let extensionContext: IExtensionContext; - let integrationStorage: IIntegrationStorage; + /** Merged (`.deepnote.env.yaml` over SecretStorage) configs keyed by the notebook they resolve for. */ + let mergedIntegrationConfigs: Map; + let sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider; let tokenStorage: IFederatedAuthTokenStorage; let subscriptions: IDisposable[]; let runOAuthFlowStub: sinon.SinonStub<[RunOAuthFlowParams], Promise<{ refreshToken: string }>>; @@ -28,11 +37,16 @@ suite('FederatedAuthCommandHandlerNode', () => { setup(() => { resetVSCodeMocks(); subscriptions = []; + mergedIntegrationConfigs = new Map(); extensionContext = mock(); - integrationStorage = mock(); + sqlIntegrationEnvVars = mock(); tokenStorage = mock(); when(extensionContext.subscriptions).thenReturn(subscriptions); + // A single matcher that dispatches on the URI: a per-URI `when` would be shadowed by matcher ordering. + when(sqlIntegrationEnvVars.getMergedIntegrationConfigs(anything())).thenCall( + async (resource: Uri) => mergedIntegrationConfigs.get(resource.toString()) ?? [] + ); runOAuthFlowStub = sinon.stub<[RunOAuthFlowParams], Promise<{ refreshToken: string }>>(); runOAuthFlowStub.resolves({ refreshToken: FED_AUTH_FIXTURE.REFRESH_TOKEN }); @@ -41,12 +55,30 @@ suite('FederatedAuthCommandHandlerNode', () => { handler = new FederatedAuthCommandHandlerNode( instance(extensionContext), - instance(integrationStorage), + instance(sqlIntegrationEnvVars), instance(tokenStorage), runOAuthFlowStub ); }); + teardown(() => { + sinon.restore(); + }); + + /** Publishes `configs` as what `.deepnote.env.yaml` + SecretStorage merge to for `uri`. */ + function setMergedIntegrationConfigs(uri: Uri, ...configs: DatabaseIntegrationConfig[]) { + mergedIntegrationConfigs.set(uri.toString(), configs); + } + + /** Drives the stubbed flow through `onListening`, the callback that opens the browser at the start URL. */ + function driveOnListening() { + runOAuthFlowStub.callsFake(async (params: RunOAuthFlowParams) => { + await params.onListening(EXTERNAL_CALLBACK_URL); + + return { refreshToken: FED_AUTH_FIXTURE.REFRESH_TOKEN }; + }); + } + ( [ ['unknown integration id', () => undefined, 'unknown-id'], @@ -57,9 +89,9 @@ suite('FederatedAuthCommandHandlerNode', () => { test(`skips OAuth flow for ${label}`, async () => { const config = build(); const id = lookupId ?? FED_AUTH_FIXTURE.INTEGRATION_ID; - when(integrationStorage.getIntegrationConfig(id)).thenResolve(config); + setMergedIntegrationConfigs(NOTEBOOK_URI, ...(config ? [config] : [])); - await handler.authenticate(id); + await handler.authenticate(id, NOTEBOOK_URI); assert.strictEqual(runOAuthFlowStub.callCount, 0); verify(tokenStorage.save(anything())).never(); @@ -67,12 +99,12 @@ suite('FederatedAuthCommandHandlerNode', () => { }); test('happy path: saves the captured refresh token with a fresh fingerprint', async () => { - when(integrationStorage.getIntegrationConfig(FED_AUTH_FIXTURE.INTEGRATION_ID)).thenResolve( - buildGoogleOauthIntegration() - ); + // The config lives only in that notebook's `.deepnote.env.yaml`; SecretStorage has nothing for the id. + setMergedIntegrationConfigs(NOTEBOOK_URI, buildGoogleOauthIntegration()); - await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID); + await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID, NOTEBOOK_URI); + verify(sqlIntegrationEnvVars.getMergedIntegrationConfigs(NOTEBOOK_URI)).once(); assert.strictEqual(runOAuthFlowStub.callCount, 1); verify( tokenStorage.save( @@ -89,12 +121,32 @@ suite('FederatedAuthCommandHandlerNode', () => { ).once(); }); + test('skips OAuth flow when the supplied notebook resolves no config for the id, even if another notebook does', async () => { + // Catches: resolving against an ambient/active notebook instead of the one the request came from. + setMergedIntegrationConfigs(NOTEBOOK_URI, buildGoogleOauthIntegration()); + + await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID, OTHER_NOTEBOOK_URI); + + verify(sqlIntegrationEnvVars.getMergedIntegrationConfigs(OTHER_NOTEBOOK_URI)).once(); + assert.strictEqual(runOAuthFlowStub.callCount, 0); + verify(tokenStorage.save(anything())).never(); + }); + + test('returns without a lookup when invoked without a resource', async () => { + // `executeCommand` callers are untyped, so the guard has to hold at runtime. + setMergedIntegrationConfigs(NOTEBOOK_URI, buildGoogleOauthIntegration()); + + await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID, undefined as unknown as Uri); + + verify(sqlIntegrationEnvVars.getMergedIntegrationConfigs(anything())).never(); + assert.strictEqual(runOAuthFlowStub.callCount, 0); + verify(tokenStorage.save(anything())).never(); + }); + test('runOAuthFlow is called with clientId, clientSecret, state, codeVerifier, and the deepnote-callback redirectUri', async () => { - when(integrationStorage.getIntegrationConfig(FED_AUTH_FIXTURE.INTEGRATION_ID)).thenResolve( - buildGoogleOauthIntegration() - ); + setMergedIntegrationConfigs(NOTEBOOK_URI, buildGoogleOauthIntegration()); - await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID); + await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID, NOTEBOOK_URI); assert.strictEqual(runOAuthFlowStub.callCount, 1); const callArg = runOAuthFlowStub.firstCall.args[0]; @@ -110,17 +162,10 @@ suite('FederatedAuthCommandHandlerNode', () => { }); test('onListening opens the deepnote.com start URL with the externalized callback as finalRedirect', async () => { - when(integrationStorage.getIntegrationConfig(FED_AUTH_FIXTURE.INTEGRATION_ID)).thenResolve( - buildGoogleOauthIntegration() - ); - - runOAuthFlowStub.callsFake(async (params: RunOAuthFlowParams) => { - await params.onListening('http://127.0.0.1:54321/auth/callback'); - - return { refreshToken: FED_AUTH_FIXTURE.REFRESH_TOKEN }; - }); + setMergedIntegrationConfigs(NOTEBOOK_URI, buildGoogleOauthIntegration()); + driveOnListening(); - await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID); + await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID, NOTEBOOK_URI); const [openedUri] = capture(mockedVSCodeNamespaces.env.openExternal).last(); // Inspect the Uri directly — going through `Uri.parse(...).toString()` would mangle percent-encoded characters in the query (mock decodes during parse). @@ -131,7 +176,7 @@ suite('FederatedAuthCommandHandlerNode', () => { const params = new URLSearchParams(uri.query); assert.strictEqual(params.get('client_id'), FED_AUTH_FIXTURE.CLIENT_ID); - assert.strictEqual(params.get('final_redirect'), 'http://127.0.0.1:54321/auth/callback'); + assert.strictEqual(params.get('final_redirect'), EXTERNAL_CALLBACK_URL); assert.isString(params.get('state')); assert.isString(params.get('code_challenge')); @@ -140,38 +185,66 @@ suite('FederatedAuthCommandHandlerNode', () => { assert.strictEqual(params.get('state'), callArg.state); }); - test('silently returns when the user cancels the flow', async () => { - when(integrationStorage.getIntegrationConfig(FED_AUTH_FIXTURE.INTEGRATION_ID)).thenResolve( - buildGoogleOauthIntegration() + test('scopes the deepnote.domain setting to the notebook the command was invoked for', async () => { + // Catches a revert to an ambient `window.activeNotebookEditor` lookup: a folder-scoped override only + // resolves if the setting is read against the supplied resource. + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote', uriEquals(NOTEBOOK_URI))).thenReturn({ + get: () => 'staging.deepnote.com' + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } as any); + setMergedIntegrationConfigs(NOTEBOOK_URI, buildGoogleOauthIntegration()); + driveOnListening(); + + await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID, NOTEBOOK_URI); + + const [openedUri] = capture(mockedVSCodeNamespaces.env.openExternal).last(); + assert.strictEqual((openedUri as Uri).authority, 'staging.deepnote.com'); + assert.strictEqual( + runOAuthFlowStub.firstCall.args[0].redirectUri, + 'https://staging.deepnote.com/auth/bigquery/google-oauth-callback' ); + }); + + test('silently returns when the user cancels the flow', async () => { + setMergedIntegrationConfigs(NOTEBOOK_URI, buildGoogleOauthIntegration()); runOAuthFlowStub.rejects(new CancellationError()); - await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID); + await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID, NOTEBOOK_URI); assert.strictEqual(runOAuthFlowStub.callCount, 1); verify(tokenStorage.save(anything())).never(); }); test('surfaces a generic OAuth error via the failure toast and does not save a token', async () => { - when(integrationStorage.getIntegrationConfig(FED_AUTH_FIXTURE.INTEGRATION_ID)).thenResolve( - buildGoogleOauthIntegration() - ); + setMergedIntegrationConfigs(NOTEBOOK_URI, buildGoogleOauthIntegration()); runOAuthFlowStub.rejects(new Error('boom')); - await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID); + await handler.authenticate(FED_AUTH_FIXTURE.INTEGRATION_ID, NOTEBOOK_URI); assert.strictEqual(runOAuthFlowStub.callCount, 1); verify(tokenStorage.save(anything())).never(); }); - test('activate registers the AuthenticateIntegration command and pushes a disposable', () => { + test('the registered command forwards both the integration id and the resource', async () => { + // Catches: dropping the second argument would silently fall back to no resource and fail the guard. when(mockedVSCodeNamespaces.commands.registerCommand(anyString(), anything())).thenReturn({ dispose: () => undefined } as IDisposable); + const authenticateStub = sinon.stub(handler, 'authenticate').resolves(); handler.activate(); assert.strictEqual(subscriptions.length, 1, 'one disposable subscription should be registered'); + + const [commandId, callback] = capture(mockedVSCodeNamespaces.commands.registerCommand).last(); + assert.strictEqual(commandId, Commands.AuthenticateIntegration); + + await (callback as (integrationId: string, resource: Uri) => Promise)( + FED_AUTH_FIXTURE.INTEGRATION_ID, + NOTEBOOK_URI + ); + + sinon.assert.calledOnceWithExactly(authenticateStub, FED_AUTH_FIXTURE.INTEGRATION_ID, NOTEBOOK_URI); }); }); @@ -193,15 +266,4 @@ suite('buildExtensionStartUrl', () => { assert.strictEqual(parsed.searchParams.get('code_challenge'), 'pkce-challenge'); assert.strictEqual(parsed.searchParams.get('final_redirect'), 'http://127.0.0.1:54321/auth/callback'); }); - - test('honors the deepnoteDomain override (for dev/staging hosts)', () => { - const url = buildExtensionStartUrl({ - clientId: 'c', - codeChallenge: 'c', - deepnoteDomain: 'dev.deepnote.org', - finalRedirect: 'http://127.0.0.1:1/cb', - state: 's' - }); - assert.strictEqual(new URL(url).host, 'dev.deepnote.org'); - }); }); diff --git a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts index 0c3c048305..a55dfeb7c5 100644 --- a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts +++ b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node.ts @@ -10,7 +10,7 @@ import { IFederatedAuthTokenStorage } from '../types'; /** * Node-only bridge that restarts kernels when a federated integration's token changes, clearing stale - * `os.environ` mutations and kernel globals. Separate from {@link IntegrationKernelRestartHandler} because + * `os.environ` mutations and kernel globals. Separate from {@link IntegrationEnvRefreshHandler} because * {@link IFederatedAuthTokenStorage} is node-only. */ @injectable() diff --git a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthOrphanedTokenCleaner.node.ts b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthOrphanedTokenCleaner.node.ts deleted file mode 100644 index 936e1b033c..0000000000 --- a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthOrphanedTokenCleaner.node.ts +++ /dev/null @@ -1,74 +0,0 @@ -import { inject, injectable } from 'inversify'; - -import { IExtensionSyncActivationService } from '../../../../platform/activation/types'; -import { IDisposableRegistry } from '../../../../platform/common/types'; -import { IIntegrationStorage } from '../../../../platform/notebooks/deepnote/types'; -import { IFederatedAuthTokenStorage } from '../types'; -import { logger } from '../../../../platform/logging'; - -/** - * Node-only listener that prunes federated-auth tokens when an integration is deleted: subscribes to - * `onDidChangeIntegrations` and diffs current IDs against {@link IFederatedAuthTokenStorage.listIntegrationIds}. - */ -@injectable() -export class FederatedAuthOrphanedTokenCleaner implements IExtensionSyncActivationService { - constructor( - @inject(IFederatedAuthTokenStorage) private readonly tokenStorage: IFederatedAuthTokenStorage, - @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, - @inject(IDisposableRegistry) disposables: IDisposableRegistry - ) { - logger.info('FederatedAuthOrphanedTokenCleaner: Initialized'); - - disposables.push( - this.integrationStorage.onDidChangeIntegrations(() => { - this.cleanupOrphanedTokens().catch((err) => - logger.error('FederatedAuthOrphanedTokenCleaner: Failed to clean up orphaned tokens', err) - ); - }) - ); - } - - public activate(): void { - this.cleanupOrphanedTokens().catch((err) => - logger.error('FederatedAuthOrphanedTokenCleaner: Initial orphaned token cleanup failed', err) - ); - } - - private async cleanupOrphanedTokens(): Promise { - const [tokenIds, integrations] = await Promise.all([ - this.tokenStorage.listIntegrationIds(), - this.integrationStorage.getAll() - ]); - - if (tokenIds.length === 0) { - logger.debug('FederatedAuthOrphanedTokenCleaner: No federated tokens stored, nothing to clean up.'); - return; - } - - const currentIntegrationIds = new Set(integrations.map((integration) => integration.id)); - const orphanedIds = tokenIds.filter((id) => !currentIntegrationIds.has(id)); - - if (orphanedIds.length === 0) { - logger.debug('FederatedAuthOrphanedTokenCleaner: No orphaned tokens to clean up.'); - return; - } - - logger.info( - `FederatedAuthOrphanedTokenCleaner: Cleaning up ${orphanedIds.length} orphaned token(s): ${orphanedIds.join( - ', ' - )}` - ); - - for (const id of orphanedIds) { - try { - await this.tokenStorage.delete(id); - logger.debug(`FederatedAuthOrphanedTokenCleaner: Deleted orphaned token for integration ${id}`); - } catch (error) { - logger.error( - `FederatedAuthOrphanedTokenCleaner: Failed to delete orphaned token for integration ${id}`, - error - ); - } - } - } -} diff --git a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthOrphanedTokenCleaner.node.unit.test.ts b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthOrphanedTokenCleaner.node.unit.test.ts deleted file mode 100644 index d8f85c8a05..0000000000 --- a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthOrphanedTokenCleaner.node.unit.test.ts +++ /dev/null @@ -1,140 +0,0 @@ -import sinon from 'sinon'; -import { Disposable, EventEmitter } from 'vscode'; - -import { FederatedAuthOrphanedTokenCleaner } from './federatedAuthOrphanedTokenCleaner.node'; -import { FederatedAuthTokenEntry, IFederatedAuthTokenStorage } from '../types'; -import { IDisposable } from '../../../../platform/common/types'; -import { IIntegrationStorage } from '../../../../platform/notebooks/deepnote/types'; -import { dispose } from '../../../../platform/common/utils/lifecycle'; -import { buildGoogleOauthIntegration, buildTokenEntry, settleAsyncHandlers } from './federatedAuthTestHelpers'; -import type { ConfigurableDatabaseIntegrationConfig } from '../../../../platform/notebooks/deepnote/integrationTypes'; - -suite('FederatedAuthOrphanedTokenCleaner', () => { - let disposables: IDisposable[]; - let integrations: Map; - let tokens: Map; - let onDidChangeIntegrations: EventEmitter; - let onDidChangeTokens: EventEmitter; - let deleteSpy: sinon.SinonSpy<[string], Promise>; - let integrationStorage: IIntegrationStorage; - let tokenStorage: IFederatedAuthTokenStorage; - - function buildTokenStorage(throwOnDelete?: Set): IFederatedAuthTokenStorage { - deleteSpy = sinon.spy(async (id: string) => { - if (throwOnDelete?.has(id)) { - throw new Error(`forced throw on delete: ${id}`); - } - tokens.delete(id); - }); - return { - onDidChangeTokens: onDidChangeTokens.event, - computeMetadataFingerprint: () => 'fp', - delete: deleteSpy, - get: async (id) => tokens.get(id), - has: async (id) => tokens.has(id), - listIntegrationIds: async () => Array.from(tokens.keys()), - save: async (entry) => { - tokens.set(entry.integrationId, entry); - } - }; - } - - function fireChangeAndWait(): Promise { - onDidChangeIntegrations.fire(); - return settleAsyncHandlers(); - } - - setup(() => { - disposables = []; - integrations = new Map(); - tokens = new Map(); - onDidChangeIntegrations = new EventEmitter(); - onDidChangeTokens = new EventEmitter(); - integrationStorage = { - onDidChangeIntegrations: onDidChangeIntegrations.event, - dispose: () => onDidChangeIntegrations.dispose(), - async clear() { - integrations.clear(); - }, - async delete(id) { - integrations.delete(id); - }, - async exists(id) { - return integrations.has(id); - }, - async getAll() { - return Array.from(integrations.values()); - }, - async getIntegrationConfig(id) { - return integrations.get(id); - }, - async getProjectIntegrationConfig() { - return undefined; - }, - async save(config) { - integrations.set(config.id, config); - } - }; - tokenStorage = buildTokenStorage(); - disposables.push(new Disposable(() => onDidChangeIntegrations.dispose())); - disposables.push(new Disposable(() => onDidChangeTokens.dispose())); - }); - - teardown(() => { - disposables = dispose(disposables); - }); - - test('does not call delete when every stored token has a matching integration', async () => { - integrations.set('bq-1', buildGoogleOauthIntegration({ id: 'bq-1' })); - integrations.set('bq-2', buildGoogleOauthIntegration({ id: 'bq-2' })); - tokens.set('bq-1', buildTokenEntry({ integrationId: 'bq-1' })); - tokens.set('bq-2', buildTokenEntry({ integrationId: 'bq-2' })); - - new FederatedAuthOrphanedTokenCleaner(tokenStorage, integrationStorage, disposables); - - await fireChangeAndWait(); - - sinon.assert.notCalled(deleteSpy); - }); - - test('deletes tokens for integrations that no longer exist', async () => { - integrations.set('bq-1', buildGoogleOauthIntegration({ id: 'bq-1' })); - tokens.set('bq-1', buildTokenEntry({ integrationId: 'bq-1' })); - tokens.set('orphan-a', buildTokenEntry({ integrationId: 'orphan-a' })); - tokens.set('orphan-b', buildTokenEntry({ integrationId: 'orphan-b' })); - - new FederatedAuthOrphanedTokenCleaner(tokenStorage, integrationStorage, disposables); - - await fireChangeAndWait(); - - sinon.assert.calledWith(deleteSpy, 'orphan-a'); - sinon.assert.calledWith(deleteSpy, 'orphan-b'); - sinon.assert.neverCalledWith(deleteSpy, 'bq-1'); - }); - - test('no-op when there are no stored tokens at all', async () => { - integrations.set('bq-1', buildGoogleOauthIntegration({ id: 'bq-1' })); - - new FederatedAuthOrphanedTokenCleaner(tokenStorage, integrationStorage, disposables); - - await fireChangeAndWait(); - - sinon.assert.notCalled(deleteSpy); - }); - - test('continues deleting other orphans when one delete fails', async () => { - tokenStorage = buildTokenStorage(new Set(['orphan-a'])); - integrations.set('bq-1', buildGoogleOauthIntegration({ id: 'bq-1' })); - tokens.set('bq-1', buildTokenEntry({ integrationId: 'bq-1' })); - tokens.set('orphan-a', buildTokenEntry({ integrationId: 'orphan-a' })); - tokens.set('orphan-b', buildTokenEntry({ integrationId: 'orphan-b' })); - - new FederatedAuthOrphanedTokenCleaner(tokenStorage, integrationStorage, disposables); - - await fireChangeAndWait(); - - // Both attempts must have happened, even after orphan-a's failure. - sinon.assert.calledWith(deleteSpy, 'orphan-a'); - sinon.assert.calledWith(deleteSpy, 'orphan-b'); - }); -}); diff --git a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthSqlBlockCodeGenerator.node.ts b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthSqlBlockCodeGenerator.node.ts index c952ecbefc..8013c10ee9 100644 --- a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthSqlBlockCodeGenerator.node.ts +++ b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthSqlBlockCodeGenerator.node.ts @@ -1,8 +1,9 @@ import type { DeepnoteBlock } from '@deepnote/blocks'; import { BigQueryAuthMethods } from '@deepnote/database-integrations'; import { inject, injectable } from 'inversify'; +import { Uri } from 'vscode'; -import { IIntegrationStorage } from '../../../../platform/notebooks/deepnote/types'; +import { ISqlIntegrationEnvVarsProvider } from '../../../../platform/notebooks/deepnote/types'; import { fetchFreshAccessToken, InvalidClientError, @@ -28,7 +29,8 @@ import { @injectable() export class FederatedAuthSqlBlockCodeGenerator implements IFederatedAuthSqlBlockCodeGenerator { constructor( - @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, + @inject(ISqlIntegrationEnvVarsProvider) + private readonly sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider, @inject(IFederatedAuthTokenStorage) private readonly tokenStorage: IFederatedAuthTokenStorage ) {} @@ -40,7 +42,7 @@ export class FederatedAuthSqlBlockCodeGenerator implements IFederatedAuthSqlBloc return fetchFreshAccessToken(entry, oauthConfig); } - public async generate(block: DeepnoteBlock): Promise { + public async generate(block: DeepnoteBlock, notebookUri: Uri): Promise { if (block.type !== 'sql') { return undefined; } @@ -52,7 +54,11 @@ export class FederatedAuthSqlBlockCodeGenerator implements IFederatedAuthSqlBloc return undefined; } - const integration = await this.integrationStorage.getIntegrationConfig(integrationId); + // Merged configs, not SecretStorage: a federated integration can be declared purely in `.deepnote.env.yaml`, + // and when both sources have it the file wins — same resolution the kernel and the SQL LSP use. + const integration = (await this.sqlIntegrationEnvVars.getMergedIntegrationConfigs(notebookUri)).find( + (config) => config.id === integrationId + ); if (!integration || integration.type !== 'big-query') { return undefined; } diff --git a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthSqlBlockCodeGenerator.node.unit.test.ts b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthSqlBlockCodeGenerator.node.unit.test.ts index 46f74db2c9..2214714ebe 100644 --- a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthSqlBlockCodeGenerator.node.unit.test.ts +++ b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthSqlBlockCodeGenerator.node.unit.test.ts @@ -1,6 +1,7 @@ import { assert } from 'chai'; import sinon from 'sinon'; -import { EventEmitter } from 'vscode'; +import { EventEmitter, Uri } from 'vscode'; +import type { DatabaseIntegrationConfig } from '@deepnote/database-integrations'; import { FederatedAuthTokenEntry, @@ -9,7 +10,9 @@ import { OAuthClientMisconfiguredError } from '../types'; import { FederatedAuthSqlBlockCodeGenerator } from './federatedAuthSqlBlockCodeGenerator.node'; -import { IIntegrationStorage } from '../../../../platform/notebooks/deepnote/types'; +import { ISqlIntegrationEnvVarsProvider } from '../../../../platform/notebooks/deepnote/types'; +import { Resource } from '../../../../platform/common/types'; +import { EnvironmentVariables } from '../../../../platform/common/variables/types'; import { InvalidClientError, InvalidGrantError, computeMetadataFingerprint } from './federatedAuthTokenStorage.node'; import { FED_AUTH_FIXTURE, @@ -21,13 +24,18 @@ import { buildTokenEntry, parsePythonSingleQuoted } from './federatedAuthTestHelpers'; -import type { ConfigurableDatabaseIntegrationConfig } from '../../../../platform/notebooks/deepnote/integrationTypes'; type FetcherFn = ( entry: FederatedAuthTokenEntry, oauthConfig: { tokenUrl: string; clientId: string; clientSecret: string } ) => Promise<{ accessToken: string; newRefreshToken?: string }>; +/** The OAuth-client metadata the fingerprint is computed over; derived from the library so it cannot drift. */ +type GoogleOauthMetadata = Extract< + Extract['metadata'], + { authMethod: 'google-oauth' } +>; + suite('FederatedAuthSqlBlockCodeGenerator', () => { const { INTEGRATION_ID, PROJECT, CLIENT_ID, CLIENT_SECRET, REFRESH_TOKEN, ACCESS_TOKEN } = FED_AUTH_FIXTURE; const VALID_FINGERPRINT = computeMetadataFingerprint({ @@ -35,55 +43,53 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { clientSecret: CLIENT_SECRET, project: PROJECT }); + const NOTEBOOK_URI = Uri.file('/workspace/project.deepnote'); + const OTHER_NOTEBOOK_URI = Uri.file('/workspace/other/project.deepnote'); - let integrations: Map; + /** Merged (`.deepnote.env.yaml` over SecretStorage) configs keyed by the notebook they resolve for. */ + let mergedConfigs: Map; let tokens: Map; let onDidChangeTokens: EventEmitter; - let onDidChangeIntegrations: EventEmitter; + let onDidChangeEnvironmentVariables: EventEmitter; let saveSpy: sinon.SinonSpy<[FederatedAuthTokenEntry, { silent?: boolean }?], Promise>; let deleteSpy: sinon.SinonSpy<[string], Promise>; - let integrationStorage: IIntegrationStorage; + let getMergedIntegrationConfigsSpy: sinon.SinonSpy<[Resource], Promise>; + let sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider; let tokenStorage: IFederatedAuthTokenStorage; let fetcher: sinon.SinonStub, ReturnType>; let generator: FederatedAuthSqlBlockCodeGenerator; setup(() => { - integrations = new Map(); + mergedConfigs = new Map(); tokens = new Map(); onDidChangeTokens = new EventEmitter(); - onDidChangeIntegrations = new EventEmitter(); + onDidChangeEnvironmentVariables = new EventEmitter(); saveSpy = sinon.spy(async (entry: FederatedAuthTokenEntry, _options?: { silent?: boolean }) => { tokens.set(entry.integrationId, entry); }); deleteSpy = sinon.spy(async (id: string) => { tokens.delete(id); }); + getMergedIntegrationConfigsSpy = sinon.spy(async (resource: Resource) => + resource ? mergedConfigs.get(resource.toString()) ?? [] : [] + ); - integrationStorage = { - onDidChangeIntegrations: onDidChangeIntegrations.event, - dispose: () => onDidChangeIntegrations.dispose(), - async clear() { - integrations.clear(); - }, - async delete(id) { - integrations.delete(id); + // Declared as a plain object (not a typed literal) so the extra members the provider grows for other + // consumers stay assignable here; the generator only ever calls `getMergedIntegrationConfigs`. + const envVarsProvider = { + onDidChangeEnvironmentVariables: onDidChangeEnvironmentVariables.event, + async getEnvironmentVariables(): Promise { + return {}; }, - async exists(id) { - return integrations.has(id); + async getFederatedAuthCandidates(): Promise> { + return new Set(); }, - async getAll() { - return Array.from(integrations.values()); + async getFileConfiguredIntegrationIds(): Promise> { + return new Set(); }, - async getIntegrationConfig(id) { - return integrations.get(id); - }, - async getProjectIntegrationConfig() { - return undefined; - }, - async save(config) { - integrations.set(config.id, config); - } + getMergedIntegrationConfigs: getMergedIntegrationConfigsSpy }; + sqlIntegrationEnvVars = envVarsProvider; tokenStorage = { onDidChangeTokens: onDidChangeTokens.event, @@ -91,11 +97,10 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { delete: deleteSpy, get: async (id) => tokens.get(id), has: async (id) => tokens.has(id), - listIntegrationIds: async () => Array.from(tokens.keys()), save: saveSpy }; - generator = new FederatedAuthSqlBlockCodeGenerator(integrationStorage, tokenStorage); + generator = new FederatedAuthSqlBlockCodeGenerator(sqlIntegrationEnvVars, tokenStorage); fetcher = sinon.stub(generator, 'fetchFreshAccessToken'); fetcher.resolves({ accessToken: ACCESS_TOKEN }); @@ -104,11 +109,16 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { teardown(() => { sinon.restore(); onDidChangeTokens.dispose(); - onDidChangeIntegrations.dispose(); + onDidChangeEnvironmentVariables.dispose(); }); + /** Publishes `configs` as what `.deepnote.env.yaml` + SecretStorage merge to for `uri`. */ + function setMergedConfigs(uri: Uri, ...configs: DatabaseIntegrationConfig[]) { + mergedConfigs.set(uri.toString(), configs); + } + function setupValidFederatedIntegration() { - integrations.set(INTEGRATION_ID, buildGoogleOauthIntegration()); + setMergedConfigs(NOTEBOOK_URI, buildGoogleOauthIntegration()); tokens.set( INTEGRATION_ID, buildTokenEntry({ refreshToken: REFRESH_TOKEN, metadataFingerprint: VALID_FINGERPRINT }) @@ -152,19 +162,30 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { test(`returns undefined for ${label}`, async () => { const integration = buildIntegration(); if (integration) { - integrations.set(integration.id, integration); + setMergedConfigs(NOTEBOOK_URI, integration); } - const result = await generator.generate(buildBlock()); + const result = await generator.generate(buildBlock(), NOTEBOOK_URI); assert.strictEqual(result, undefined); sinon.assert.notCalled(fetcher); }); }); + test('returns undefined when the supplied notebook resolves no config for the id, even if another notebook does', async () => { + // Catches: resolving against an ambient/active notebook instead of the one the cell belongs to. + setupValidFederatedIntegration(); + + const result = await generator.generate(buildSqlBlock(), OTHER_NOTEBOOK_URI); + + assert.strictEqual(result, undefined); + sinon.assert.calledOnceWithExactly(getMergedIntegrationConfigsSpy, OTHER_NOTEBOOK_URI); + sinon.assert.notCalled(fetcher); + }); + test('throws NotAuthenticatedError when federated integration has no stored token', async () => { - integrations.set(INTEGRATION_ID, buildGoogleOauthIntegration()); + setMergedConfigs(NOTEBOOK_URI, buildGoogleOauthIntegration()); try { - await generator.generate(buildSqlBlock()); + await generator.generate(buildSqlBlock(), NOTEBOOK_URI); assert.fail('Expected NotAuthenticatedError'); } catch (err) { assert(err instanceof NotAuthenticatedError); @@ -173,15 +194,25 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { sinon.assert.notCalled(fetcher); }); - test('throws NotAuthenticatedError and deletes the token when the metadata fingerprint is stale', async () => { - setupValidFederatedIntegration(); + test('a .deepnote.env.yaml edit to clientId invalidates the stored token: deletes it and throws NotAuthenticatedError', async () => { + // The token was bound to the OAuth client the file declared when it was saved. Editing the YAML rebinds + // the integration to a different client, so the fingerprint check must drop the now-unusable token. + // One perturbed field is enough: the generator hashes clientId/clientSecret/project through a single + // `computeMetadataFingerprint` call, and per-field sensitivity is covered in the token-storage suite. + const metadata: GoogleOauthMetadata = { + authMethod: 'google-oauth', + clientId: 'edited-in-yaml', + clientSecret: CLIENT_SECRET, + project: PROJECT + }; + setMergedConfigs(NOTEBOOK_URI, buildGoogleOauthIntegration({ metadata })); tokens.set( INTEGRATION_ID, - buildTokenEntry({ refreshToken: REFRESH_TOKEN, metadataFingerprint: 'stale-fingerprint' }) + buildTokenEntry({ refreshToken: REFRESH_TOKEN, metadataFingerprint: VALID_FINGERPRINT }) ); try { - await generator.generate(buildSqlBlock()); + await generator.generate(buildSqlBlock(), NOTEBOOK_URI); assert.fail('Expected NotAuthenticatedError'); } catch (err) { assert.instanceOf(err, NotAuthenticatedError); @@ -194,7 +225,7 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { // Mirrors deepnote-internal: one `_dntk.execute_sql_with_connection_json(...)` call with the connection JSON as a literal containing the fresh access token. Token is expected in the single execute payload, same as cloud. setupValidFederatedIntegration(); - const result = await generator.generate(buildSqlBlock()); + const result = await generator.generate(buildSqlBlock(), NOTEBOOK_URI); if (typeof result !== 'string') { throw new Error(`expected a string result, got ${typeof result}`); } @@ -213,18 +244,13 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { test('connection JSON literal round-trips through Python+json.loads when integration name contains backslash, newline, and single quote', async () => { // Catches: regressing `escapePythonString` (e.g. swapping in a single-char `\\'` escape) would leave `\\`/`\n` undecoded and break `json.loads` at the kernel. The literal lives inside the execute call now (no separate prelude assignment). const hostileProject = "gcp-with-\\-and-\n-and-'-project"; - integrations.set( - INTEGRATION_ID, - buildGoogleOauthIntegration({ - metadata: { - authMethod: 'google-oauth', - project: hostileProject, - clientId: CLIENT_ID, - clientSecret: CLIENT_SECRET - } - // eslint-disable-next-line @typescript-eslint/no-explicit-any - } as any) - ); + const metadata: GoogleOauthMetadata = { + authMethod: 'google-oauth', + clientId: CLIENT_ID, + clientSecret: CLIENT_SECRET, + project: hostileProject + }; + setMergedConfigs(NOTEBOOK_URI, buildGoogleOauthIntegration({ metadata })); const hostileFingerprint = computeMetadataFingerprint({ clientId: CLIENT_ID, clientSecret: CLIENT_SECRET, @@ -235,7 +261,7 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { buildTokenEntry({ refreshToken: REFRESH_TOKEN, metadataFingerprint: hostileFingerprint }) ); - const result = await generator.generate(buildSqlBlock()); + const result = await generator.generate(buildSqlBlock(), NOTEBOOK_URI); if (typeof result !== 'string') { throw new Error(`expected a string result, got ${typeof result}`); } @@ -255,8 +281,8 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { fetcher.onFirstCall().resolves({ accessToken: 'token-1' }); fetcher.onSecondCall().resolves({ accessToken: 'token-2' }); - const first = await generator.generate(buildSqlBlock()); - const second = await generator.generate(buildSqlBlock()); + const first = await generator.generate(buildSqlBlock(), NOTEBOOK_URI); + const second = await generator.generate(buildSqlBlock(), NOTEBOOK_URI); sinon.assert.calledTwice(fetcher); assert.notStrictEqual(first, second); @@ -269,7 +295,7 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { fetcher.rejects(new InvalidGrantError()); try { - await generator.generate(buildSqlBlock()); + await generator.generate(buildSqlBlock(), NOTEBOOK_URI); assert.fail('Expected NotAuthenticatedError'); } catch (err) { assert.instanceOf(err, NotAuthenticatedError); @@ -283,7 +309,7 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { fetcher.rejects(new InvalidClientError()); try { - await generator.generate(buildSqlBlock()); + await generator.generate(buildSqlBlock(), NOTEBOOK_URI); assert.fail('Expected OAuthClientMisconfiguredError'); } catch (err) { assert(err instanceof OAuthClientMisconfiguredError); @@ -299,7 +325,7 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { setupValidFederatedIntegration(); fetcher.resolves({ accessToken: ACCESS_TOKEN, newRefreshToken: 'new-refresh-token' }); - await generator.generate(buildSqlBlock()); + await generator.generate(buildSqlBlock(), NOTEBOOK_URI); sinon.assert.calledOnce(saveSpy); const [savedEntry, options] = saveSpy.firstCall.args; @@ -315,14 +341,14 @@ suite('FederatedAuthSqlBlockCodeGenerator', () => { setupValidFederatedIntegration(); fetcher.resolves({ accessToken: ACCESS_TOKEN, newRefreshToken: REFRESH_TOKEN }); - await generator.generate(buildSqlBlock()); + await generator.generate(buildSqlBlock(), NOTEBOOK_URI); sinon.assert.notCalled(saveSpy); }); test('honors deepnote_variable_name by emitting an assignment in the generated code', async () => { setupValidFederatedIntegration(); - const result = await generator.generate(buildSqlBlock({ deepnote_variable_name: 'my_df' })); + const result = await generator.generate(buildSqlBlock({ deepnote_variable_name: 'my_df' }), NOTEBOOK_URI); if (typeof result !== 'string') { throw new Error(`expected a string result, got ${typeof result}`); } diff --git a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthTokenStorage.node.ts b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthTokenStorage.node.ts index 5c644f02aa..f8d6eec289 100644 --- a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthTokenStorage.node.ts +++ b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthTokenStorage.node.ts @@ -171,11 +171,6 @@ export class FederatedAuthTokenStorage implements IFederatedAuthTokenStorage { return this.cache.has(integrationId); } - public async listIntegrationIds(): Promise { - await this.ensureCacheLoaded(); - return Array.from(this.cache.keys()); - } - public async save(entry: FederatedAuthTokenEntry, options?: { silent?: boolean }): Promise { await this.ensureCacheLoaded(); diff --git a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthTokenStorage.node.unit.test.ts b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthTokenStorage.node.unit.test.ts index 8046842c1c..fbeea21da4 100644 --- a/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthTokenStorage.node.unit.test.ts +++ b/src/notebooks/deepnote/integrations/federatedAuth/federatedAuthTokenStorage.node.unit.test.ts @@ -107,24 +107,6 @@ suite('federatedAuthTokenStorage', () => { assert.strictEqual(await storage.has(entry.integrationId), true); }); - test('listIntegrationIds returns all stored integration ids', async () => { - await storage.save(sampleEntry('integration-a')); - await storage.save(sampleEntry('integration-b')); - await storage.save(sampleEntry('integration-c')); - - const ids = await storage.listIntegrationIds(); - assert.deepStrictEqual(ids.sort(), ['integration-a', 'integration-b', 'integration-c']); - }); - - test('listIntegrationIds reflects deletions', async () => { - await storage.save(sampleEntry('integration-a')); - await storage.save(sampleEntry('integration-b')); - await storage.delete('integration-a'); - - const ids = await storage.listIntegrationIds(); - assert.deepStrictEqual(ids, ['integration-b']); - }); - test('delete removes the entry', async () => { const entry = sampleEntry(); await storage.save(entry); diff --git a/src/notebooks/deepnote/integrations/integrationDetector.ts b/src/notebooks/deepnote/integrations/integrationDetector.ts index 5715ead44c..702fa53d02 100644 --- a/src/notebooks/deepnote/integrations/integrationDetector.ts +++ b/src/notebooks/deepnote/integrations/integrationDetector.ts @@ -1,14 +1,16 @@ import { inject, injectable } from 'inversify'; +import { Uri } from 'vscode'; + +import { DatabaseIntegrationConfig } from '@deepnote/database-integrations'; import { logger } from '../../../platform/logging'; import { IDeepnoteNotebookManager } from '../../types'; import { - ConfigurableDatabaseIntegrationType, - IntegrationStatus, - IntegrationWithStatus + DetectedIntegration, + isConfigurableDatabaseIntegrationType } from '../../../platform/notebooks/deepnote/integrationTypes'; -import { IIntegrationDetector, IIntegrationStorage } from './types'; -import { databaseIntegrationTypes } from '@deepnote/database-integrations'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; +import { IIntegrationDetector, IIntegrationStorage, IntegrationDetectionInput } from './types'; /** * Service for detecting integrations used in Deepnote notebooks @@ -17,72 +19,93 @@ import { databaseIntegrationTypes } from '@deepnote/database-integrations'; export class IntegrationDetector implements IIntegrationDetector { constructor( @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, - @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager + @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, + @inject(ISqlIntegrationEnvVarsProvider) + private readonly sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider ) {} /** - * Detect all integrations used in the given project. - * Uses the project's integrations field as the source of truth. + * Detect all integrations for the notebook's project. Three inputs, three roles: + * - `project.integrations` is the roster (ids, names and types only — never credentials), so it decides + * the order and the names the panel shows. + * - SecretStorage supplies the editable config for each one; integrations configured only in + * `.deepnote.env.yaml` stay `null` here, since those configs are never persisted through it. + * - `.deepnote.env.yaml` entries missing from the roster are appended, matching what actually applies at + * execution time. Without this a file-only integration works but is invisible, and a federated one is + * unusable outright — its Authenticate action exists only as a row in this panel. */ - async detectIntegrations(projectId: string, notebookId: string): Promise> { - // Get the project + async detectIntegrations(input: IntegrationDetectionInput): Promise> { + const { projectId, notebookId } = input; + const project = this.notebookManager.getProjectForNotebook(projectId, notebookId); if (!project) { logger.warn( `IntegrationDetector: No project found for ID: ${projectId}. The project may not have been loaded yet.` ); + return new Map(); } - logger.debug(`IntegrationDetector: Scanning project ${projectId} for integrations`); + const projectIntegrations = project.project.integrations ?? []; - const integrations = new Map(); + logger.debug(`IntegrationDetector: Project ${projectId} declares ${projectIntegrations.length} integrations`); - // Use the project's integrations field as the source of truth - const projectIntegrations = project.project.integrations?.slice() ?? []; - logger.debug(`IntegrationDetector: Found ${projectIntegrations.length} integrations in project.integrations`); + const integrations = new Map(); for (const projectIntegration of projectIntegrations) { - const integrationId = projectIntegration.id; const integrationType = projectIntegration.type; - if ( - !(databaseIntegrationTypes as readonly string[]).includes(integrationType) || - integrationType === 'pandas-dataframe' - ) { + if (!isConfigurableDatabaseIntegrationType(integrationType)) { logger.debug(`IntegrationDetector: Skipping unsupported integration type: ${integrationType}`); continue; } - // Check if the integration is configured - const config = await this.integrationStorage.getIntegrationConfig(integrationId); - const status: IntegrationWithStatus = { - config: config ?? null, - status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected, - // Include integration metadata from project for prefilling when config is null - integrationName: projectIntegration.name, - integrationType: integrationType as ConfigurableDatabaseIntegrationType - }; + const storedConfig = await this.integrationStorage.getIntegrationConfig(projectIntegration.id); - integrations.set(integrationId, status); + integrations.set(projectIntegration.id, { + config: storedConfig ?? null, + integrationName: projectIntegration.name, + integrationType + }); } + await this.appendFileOnlyIntegrations(input.notebookUri, integrations); + logger.debug(`IntegrationDetector: Found ${integrations.size} integrations`); return integrations; } /** - * Check if a project has any unconfigured integrations + * Adds `.deepnote.env.yaml` integrations the roster omits. `config` stays `null` because the panel edits + * SecretStorage only and the file layer cannot be written back; the name and type are carried so the row + * renders. A failed lookup leaves the roster-only result rather than blocking the panel. */ - async hasUnconfiguredIntegrations(projectId: string, notebookId: string): Promise { - const integrations = await this.detectIntegrations(projectId, notebookId); + private async appendFileOnlyIntegrations( + notebookUri: Uri, + integrations: Map + ): Promise { + let mergedConfigs: DatabaseIntegrationConfig[]; - for (const integration of integrations.values()) { - if (integration.status === IntegrationStatus.Disconnected) { - return true; - } + try { + mergedConfigs = await this.sqlIntegrationEnvVars.getMergedIntegrationConfigs(notebookUri); + } catch (error) { + logger.error('IntegrationDetector: failed to read file integrations; listing the roster only', error); + + return; } - return false; + for (const config of mergedConfigs) { + // Anything merged but absent here came from the file alone — the merge resolves roster ids first. + if (integrations.has(config.id) || !isConfigurableDatabaseIntegrationType(config.type)) { + continue; + } + + logger.debug(`IntegrationDetector: Adding file-only integration ${config.id}`); + integrations.set(config.id, { + config: null, + integrationName: config.name, + integrationType: config.type + }); + } } } diff --git a/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.ts b/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.ts new file mode 100644 index 0000000000..1abe38e314 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.ts @@ -0,0 +1,64 @@ +import { inject, injectable } from 'inversify'; +import { l10n, type NotebookDocument, window } from 'vscode'; + +import { IKernelProvider } from '../../../kernels/types'; +import { logger } from '../../../platform/logging'; +import { IIntegrationEnvLiveRefresher } from './types'; + +/** `_dntk` isn't guaranteed here, so import the package to re-fetch + apply integration env in the kernel. */ +const REFRESH_INTEGRATION_ENV_SNIPPET = `import deepnote_toolkit +deepnote_toolkit.set_integration_env()`; + +/** How long the transient "environment updated" status-bar message stays visible. */ +const STATUS_BAR_MESSAGE_TIMEOUT_MS = 5000; + +@injectable() +export class IntegrationEnvLiveRefresher implements IIntegrationEnvLiveRefresher { + constructor(@inject(IKernelProvider) private readonly kernelProvider: IKernelProvider) {} + + public async refresh(notebooks: readonly NotebookDocument[]): Promise { + const results = await Promise.all(notebooks.map((notebook) => this.refreshNotebook(notebook))); + const refreshedCount = results.filter(Boolean).length; + + if (refreshedCount > 0) { + // Transient status-bar message rather than a persistent toast, so frequent env-file edits don't spam notifications. + window.setStatusBarMessage( + l10n.t('Deepnote integration environment updated.'), + STATUS_BAR_MESSAGE_TIMEOUT_MS + ); + } + } + + /** Refreshes one kernel; never throws (per-notebook errors are logged), resolves true only on a clean run. */ + private async refreshNotebook(notebook: NotebookDocument): Promise { + try { + const kernel = this.kernelProvider.get(notebook); + if (!kernel || !kernel.startedAtLeastOnce) { + return false; + } + + const outputs = await this.kernelProvider + .getKernelExecution(kernel) + .executeHidden(REFRESH_INTEGRATION_ENV_SNIPPET); + + const errors = outputs.filter((output) => output.output_type === 'error'); + if (errors.length > 0) { + logger.warn( + `IntegrationEnvLiveRefresher: Refresh snippet produced errors for ${notebook.uri.toString()}`, + errors + ); + + return false; + } + + return true; + } catch (err) { + logger.error( + `IntegrationEnvLiveRefresher: Failed to refresh integration env for ${notebook.uri.toString()}`, + err + ); + + return false; + } + } +} diff --git a/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.unit.test.ts b/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.unit.test.ts new file mode 100644 index 0000000000..1b39e45990 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvLiveRefresher.node.unit.test.ts @@ -0,0 +1,141 @@ +import { assert } from 'chai'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; +import { Disposable, NotebookDocument, Uri } from 'vscode'; + +import { IKernel, IKernelProvider, INotebookKernelExecution } from '../../../kernels/types'; +import { IDisposable } from '../../../platform/common/types'; +import { dispose } from '../../../platform/common/utils/lifecycle'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; +import { IntegrationEnvLiveRefresher } from './integrationEnvLiveRefresher.node'; + +// Must match REFRESH_INTEGRATION_ENV_SNIPPET in integrationEnvLiveRefresher.node.ts exactly (real newline). +const EXPECTED_SNIPPET = `import deepnote_toolkit +deepnote_toolkit.set_integration_env()`; + +const EXPECTED_NOTIFICATION = 'Deepnote integration environment updated.'; + +suite('IntegrationEnvLiveRefresher', () => { + let refresher: IntegrationEnvLiveRefresher; + let kernelProvider: IKernelProvider; + let kernelExecution: INotebookKernelExecution; + let disposables: IDisposable[]; + + setup(() => { + resetVSCodeMocks(); + disposables = [new Disposable(() => resetVSCodeMocks())]; + kernelProvider = mock(); + + // Every notebook resolves to the same execution mock, so call counts below are totals across notebooks. + kernelExecution = mock(); + when(kernelExecution.executeHidden(anything())).thenResolve([]); + when(kernelProvider.getKernelExecution(anything())).thenReturn(instance(kernelExecution)); + + refresher = new IntegrationEnvLiveRefresher(instance(kernelProvider)); + }); + + teardown(() => { + disposables = dispose(disposables); + }); + + /** A notebook whose kernel is started; `kernelProvider.get(notebook)` returns it. */ + function createRunningNotebook(uri: Uri): NotebookDocument { + const notebookMock = mock(); + when(notebookMock.uri).thenReturn(uri); + const notebook = instance(notebookMock); + + const kernelMock = mock(); + when(kernelMock.startedAtLeastOnce).thenReturn(true); + when(kernelProvider.get(notebook)).thenReturn(instance(kernelMock)); + + return notebook; + } + + test('runs the exact refresh snippet in a started kernel and shows one status-bar message', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + + await refresher.refresh([notebook]); + + verify(kernelExecution.executeHidden(anything())).once(); + const [code] = capture(kernelExecution.executeHidden).first(); + assert.strictEqual( + code, + EXPECTED_SNIPPET, + 'executeHidden must receive the toolkit set_integration_env() snippet verbatim' + ); + + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).once(); + const [message] = capture(mockedVSCodeNamespaces.window.setStatusBarMessage).last(); + assert.strictEqual(message, EXPECTED_NOTIFICATION); + }); + + test('skips notebooks with no kernel and shows no status-bar message', async () => { + const notebookMock = mock(); + when(notebookMock.uri).thenReturn(Uri.file('/ws/a.deepnote')); + const notebook = instance(notebookMock); + when(kernelProvider.get(notebook)).thenReturn(undefined); + + await refresher.refresh([notebook]); + + verify(kernelExecution.executeHidden(anything())).never(); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).never(); + }); + + test('skips kernels that have not started and shows no status-bar message', async () => { + const notebookMock = mock(); + when(notebookMock.uri).thenReturn(Uri.file('/ws/a.deepnote')); + const notebook = instance(notebookMock); + + const kernelMock = mock(); + when(kernelMock.startedAtLeastOnce).thenReturn(false); + when(kernelProvider.get(notebook)).thenReturn(instance(kernelMock)); + + await refresher.refresh([notebook]); + + verify(kernelExecution.executeHidden(anything())).never(); + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).never(); + }); + + test('does not show a status-bar message when the refresh snippet produces an error output', async () => { + const notebook = createRunningNotebook(Uri.file('/ws/a.deepnote')); + when(kernelExecution.executeHidden(anything())).thenResolve([ + { output_type: 'error', ename: 'RuntimeError', evalue: 'boom', traceback: [] } + ]); + + await refresher.refresh([notebook]); + + verify(kernelExecution.executeHidden(anything())).once(); // the snippet still runs; its output signals failure + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).never(); + }); + + test('continues to the next notebook when one executeHidden throws, and still notifies for the success', async () => { + const notebookA = createRunningNotebook(Uri.file('/ws/a.deepnote')); + const notebookB = createRunningNotebook(Uri.file('/ws/b.deepnote')); + when(kernelExecution.executeHidden(anything())).thenReject(new Error('kernel exploded')).thenResolve([]); + + await refresher.refresh([notebookA, notebookB]); + + verify(kernelExecution.executeHidden(anything())).twice(); // a throw on the first must not stop the second + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).once(); + }); + + test('refreshes kernels in parallel: both executions start before either resolves', async () => { + const notebookA = createRunningNotebook(Uri.file('/ws/a.deepnote')); + const notebookB = createRunningNotebook(Uri.file('/ws/b.deepnote')); + + // First execution resolves only once the second has been invoked; a sequential loop would deadlock (and time out). + let markSecondInvoked!: () => void; + const secondInvoked = new Promise((resolve) => (markSecondInvoked = resolve)); + when(kernelExecution.executeHidden(anything())) + .thenCall(() => secondInvoked.then(() => [])) + .thenCall(() => { + markSecondInvoked(); + + return Promise.resolve([]); + }); + + await refresher.refresh([notebookA, notebookB]); + + verify(kernelExecution.executeHidden(anything())).twice(); // both started kernels are refreshed + verify(mockedVSCodeNamespaces.window.setStatusBarMessage(anything(), anything())).once(); + }); +}); diff --git a/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.ts b/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.ts new file mode 100644 index 0000000000..c4806e87cf --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.ts @@ -0,0 +1,40 @@ +import { inject, injectable } from 'inversify'; +import { workspace } from 'vscode'; + +import { IExtensionSyncActivationService } from '../../../platform/activation/types'; +import { IDisposableRegistry } from '../../../platform/common/types'; +import { logger } from '../../../platform/logging'; +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; +import { IIntegrationEnvLiveRefresher, IIntegrationStorage } from './types'; + +/** Live-refreshes integration env in open Deepnote kernels when integration configs change (no restart). */ +@injectable() +export class IntegrationEnvRefreshHandler implements IExtensionSyncActivationService { + constructor( + @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, + @inject(IIntegrationEnvLiveRefresher) private readonly liveRefresher: IIntegrationEnvLiveRefresher, + @inject(IDisposableRegistry) disposables: IDisposableRegistry + ) { + logger.info('IntegrationEnvRefreshHandler: Initialized'); + + disposables.push( + this.integrationStorage.onDidChangeIntegrations(() => { + this.onIntegrationConfigurationChanged().catch((err) => + logger.error('IntegrationEnvRefreshHandler: Failed to handle integration change', err) + ); + }) + ); + } + + public activate(): void { + // Service is activated via constructor + } + + private async onIntegrationConfigurationChanged(): Promise { + const notebooks = workspace.notebookDocuments.filter( + (notebook) => notebook.notebookType === DEEPNOTE_NOTEBOOK_TYPE + ); + + await this.liveRefresher.refresh(notebooks); + } +} diff --git a/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.unit.test.ts b/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.unit.test.ts new file mode 100644 index 0000000000..26bb41ba6e --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationEnvRefreshHandler.unit.test.ts @@ -0,0 +1,87 @@ +import { assert } from 'chai'; +import * as sinon from 'sinon'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; +import { Disposable, EventEmitter, Uri } from 'vscode'; + +import { IDisposable } from '../../../platform/common/types'; +import { dispose } from '../../../platform/common/utils/lifecycle'; +import { logger } from '../../../platform/logging'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; +import { IIntegrationEnvLiveRefresher, IIntegrationStorage } from './types'; +import { IntegrationEnvRefreshHandler } from './integrationEnvRefreshHandler'; +import { createMockNotebook } from '../deepnoteTestHelpers'; + +suite('IntegrationEnvRefreshHandler', () => { + let handler: IntegrationEnvRefreshHandler; + let integrationStorage: IIntegrationStorage; + let liveRefresher: IIntegrationEnvLiveRefresher; + let disposables: IDisposable[]; + let onDidChangeIntegrations: EventEmitter; + + setup(() => { + resetVSCodeMocks(); + disposables = [new Disposable(() => resetVSCodeMocks())]; + integrationStorage = mock(); + liveRefresher = mock(); + onDidChangeIntegrations = new EventEmitter(); + disposables.push(onDidChangeIntegrations); + + when(integrationStorage.onDidChangeIntegrations).thenReturn(onDidChangeIntegrations.event); + when(liveRefresher.refresh(anything())).thenResolve(); + + handler = new IntegrationEnvRefreshHandler(instance(integrationStorage), instance(liveRefresher), disposables); + }); + + teardown(() => { + sinon.restore(); + disposables = dispose(disposables); + }); + + test('passes only Deepnote notebooks to the refresher when integrations change', async () => { + const deepnote = createMockNotebook({ uri: Uri.file('/a.deepnote') }); + const jupyter = createMockNotebook({ notebookType: 'jupyter-notebook', uri: Uri.file('/b.ipynb') }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([deepnote, jupyter]); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [deepnote]); + }); + + test('refreshes multiple Deepnote notebooks in a single call', async () => { + const notebook1 = createMockNotebook({ uri: Uri.file('/test1.deepnote') }); + const notebook2 = createMockNotebook({ uri: Uri.file('/test2.deepnote') }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [notebook1, notebook2]); + }); + + test('catches and logs a rejected refresh instead of propagating it', async () => { + const notebook = createMockNotebook({ uri: Uri.file('/test.deepnote') }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(liveRefresher.refresh(anything())).thenReject(new Error('refresh boom')); + const errorStub = sinon.stub(logger, 'error'); + + handler.activate(); + onDidChangeIntegrations.fire(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + verify(liveRefresher.refresh(anything())).once(); + assert.strictEqual(errorStub.callCount, 1, 'the fire-and-forget rejection must be caught and logged'); + }); +}); diff --git a/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.ts b/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.ts deleted file mode 100644 index b082d0d192..0000000000 --- a/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.ts +++ /dev/null @@ -1,155 +0,0 @@ -import { inject, injectable } from 'inversify'; -import { l10n, NotebookDocument, workspace, window } from 'vscode'; - -import { IDisposableRegistry } from '../../../platform/common/types'; -import { IExtensionSyncActivationService } from '../../../platform/activation/types'; -import { logger } from '../../../platform/logging'; -import { IIntegrationStorage } from './types'; -import { IKernelProvider } from '../../../kernels/types'; -import { DATAFRAME_SQL_INTEGRATION_ID } from '../../../platform/notebooks/deepnote/integrationTypes'; - -/** - * Handles automatic kernel restart when integration configurations change. - * When a user saves/deletes an integration config, this service restarts all kernels - * that are using that integration so they pick up the new credentials. - */ -@injectable() -export class IntegrationKernelRestartHandler implements IExtensionSyncActivationService { - private isRestarting = false; - - constructor( - @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, - @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider, - @inject(IDisposableRegistry) disposables: IDisposableRegistry - ) { - logger.info('IntegrationKernelRestartHandler: Initialized'); - - // Listen for integration configuration changes - disposables.push( - this.integrationStorage.onDidChangeIntegrations(() => { - this.onIntegrationConfigurationChanged().catch((err) => - logger.error('IntegrationKernelRestartHandler: Failed to handle integration change', err) - ); - }) - ); - } - - public activate(): void { - // Service is activated via constructor - } - - /** - * Handle integration configuration changes by restarting affected kernels - */ - private async onIntegrationConfigurationChanged(): Promise { - // Prevent multiple simultaneous restart attempts - if (this.isRestarting) { - logger.debug('IntegrationKernelRestartHandler: Already restarting, skipping'); - return; - } - - try { - this.isRestarting = true; - - logger.info( - 'IntegrationKernelRestartHandler: Integration configuration changed, checking for affected kernels' - ); - - // Find all Deepnote notebooks with running kernels that use SQL integrations - const notebooksToRestart: NotebookDocument[] = []; - - for (const notebook of workspace.notebookDocuments) { - // Only process Deepnote notebooks - if (notebook.notebookType !== 'deepnote') { - continue; - } - - // Check if kernel is running - const kernel = this.kernelProvider.get(notebook); - if (!kernel || !kernel.startedAtLeastOnce) { - continue; - } - - // Check if notebook uses SQL integrations - const usesIntegrations = this.notebookUsesSqlIntegrations(notebook); - if (usesIntegrations) { - notebooksToRestart.push(notebook); - } - } - - if (notebooksToRestart.length === 0) { - logger.info( - 'IntegrationKernelRestartHandler: No running kernels use SQL integrations, no restart needed' - ); - return; - } - - logger.info( - `IntegrationKernelRestartHandler: Found ${notebooksToRestart.length} notebook(s) with kernels that need restart` - ); - - // Restart kernels for affected notebooks - const restartPromises = notebooksToRestart.map(async (notebook) => { - const kernel = this.kernelProvider.get(notebook); - if (kernel) { - try { - logger.info( - `IntegrationKernelRestartHandler: Restarting kernel for notebook: ${notebook.uri.toString()}` - ); - await kernel.restart(); - logger.info( - `IntegrationKernelRestartHandler: Successfully restarted kernel for: ${notebook.uri.toString()}` - ); - } catch (error) { - logger.error( - `IntegrationKernelRestartHandler: Failed to restart kernel for ${notebook.uri.toString()}`, - error - ); - // Don't throw - we want to continue restarting other kernels - } - } - }); - - await Promise.all(restartPromises); - - // Show a notification to the user - if (notebooksToRestart.length === 1) { - void window.showInformationMessage( - l10n.t('Integration configuration updated. Kernel restarted to apply changes.') - ); - } else { - void window.showInformationMessage( - l10n.t( - 'Integration configuration updated. {0} kernels restarted to apply changes.', - notebooksToRestart.length - ) - ); - } - } finally { - this.isRestarting = false; - } - } - - /** - * Check if a notebook uses SQL integrations by scanning cells for sql_integration_id metadata - */ - private notebookUsesSqlIntegrations(notebook: NotebookDocument): boolean { - for (const cell of notebook.getCells()) { - // Check for SQL cells - if (cell.document.languageId !== 'sql') { - continue; - } - - const metadata = cell.metadata; - if (metadata && typeof metadata === 'object') { - const integrationId = (metadata as Record).sql_integration_id; - if (typeof integrationId === 'string' && integrationId !== DATAFRAME_SQL_INTEGRATION_ID) { - // Found a SQL cell with an external integration - return true; - } - } - } - - return false; - } -} diff --git a/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.unit.test.ts b/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.unit.test.ts deleted file mode 100644 index 62888c3fea..0000000000 --- a/src/notebooks/deepnote/integrations/integrationKernelRestartHandler.unit.test.ts +++ /dev/null @@ -1,243 +0,0 @@ -import { anything, instance, mock, verify, when } from 'ts-mockito'; -import { Disposable, EventEmitter, NotebookCell, NotebookDocument, TextDocument, Uri } from 'vscode'; - -import { IDisposable } from '../../../platform/common/types'; -import { dispose } from '../../../platform/common/utils/lifecycle'; -import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; -import { IIntegrationStorage } from './types'; -import { IKernel, IKernelProvider } from '../../../kernels/types'; -import { IntegrationKernelRestartHandler } from './integrationKernelRestartHandler'; -import { DATAFRAME_SQL_INTEGRATION_ID } from '../../../platform/notebooks/deepnote/integrationTypes'; - -suite('IntegrationKernelRestartHandler', () => { - let handler: IntegrationKernelRestartHandler; - let integrationStorage: IIntegrationStorage; - let kernelProvider: IKernelProvider; - let disposables: IDisposable[]; - let onDidChangeIntegrations: EventEmitter; - - setup(() => { - resetVSCodeMocks(); - disposables = [new Disposable(() => resetVSCodeMocks())]; - integrationStorage = mock(); - kernelProvider = mock(); - onDidChangeIntegrations = new EventEmitter(); - disposables.push(onDidChangeIntegrations); - - when(integrationStorage.onDidChangeIntegrations).thenReturn(onDidChangeIntegrations.event); - - handler = new IntegrationKernelRestartHandler( - instance(integrationStorage), - instance(kernelProvider), - disposables - ); - }); - - teardown(() => { - disposables = dispose(disposables); - }); - - function createMockNotebook( - notebookType: string, - uri: Uri, - cells: { languageId: string; metadata?: Record }[] - ): NotebookDocument { - const notebook = mock(); - const mockCells: NotebookCell[] = cells.map((cellConfig, index) => { - const cell = mock(); - const doc = mock(); - when(doc.languageId).thenReturn(cellConfig.languageId); - when(cell.document).thenReturn(instance(doc)); - when(cell.metadata).thenReturn(cellConfig.metadata || {}); - when(cell.index).thenReturn(index); - return instance(cell); - }); - - when(notebook.notebookType).thenReturn(notebookType); - when(notebook.uri).thenReturn(uri); - when(notebook.getCells()).thenReturn(mockCells); - - return instance(notebook); - } - - test('restarts kernel when integration changes and notebook uses SQL integrations', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).once(); - }); - - test('does not restart kernel for non-Deepnote notebooks', async () => { - const notebook = createMockNotebook('jupyter-notebook', Uri.file('/test.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).never(); - }); - - test('does not restart kernel that has not started', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(false); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).never(); - }); - - test('does not restart kernel when notebook has no SQL integrations', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'python', metadata: {} } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).never(); - }); - - test('does not restart kernel when notebook only uses internal DuckDB integration', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: DATAFRAME_SQL_INTEGRATION_ID } } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).never(); - }); - - test('restarts multiple kernels in parallel', async () => { - const notebook1 = createMockNotebook('deepnote', Uri.file('/test1.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const notebook2 = createMockNotebook('deepnote', Uri.file('/test2.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'bigquery-1' } } - ]); - const mockKernel1 = mock(); - const mockKernel2 = mock(); - when(mockKernel1.startedAtLeastOnce).thenReturn(true); - when(mockKernel1.restart()).thenResolve(); - when(mockKernel2.startedAtLeastOnce).thenReturn(true); - when(mockKernel2.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - when(kernelProvider.get(notebook1)).thenReturn(instance(mockKernel1)); - when(kernelProvider.get(notebook2)).thenReturn(instance(mockKernel2)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel1.restart()).once(); - verify(mockKernel2.restart()).once(); - }); - - test('continues restarting other kernels when one fails', async () => { - const notebook1 = createMockNotebook('deepnote', Uri.file('/test1.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } } - ]); - const notebook2 = createMockNotebook('deepnote', Uri.file('/test2.ipynb'), [ - { languageId: 'sql', metadata: { sql_integration_id: 'bigquery-1' } } - ]); - const mockKernel1 = mock(); - const mockKernel2 = mock(); - when(mockKernel1.startedAtLeastOnce).thenReturn(true); - when(mockKernel1.restart()).thenReject(new Error('Restart failed')); - when(mockKernel2.startedAtLeastOnce).thenReturn(true); - when(mockKernel2.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - when(kernelProvider.get(notebook1)).thenReturn(instance(mockKernel1)); - when(kernelProvider.get(notebook2)).thenReturn(instance(mockKernel2)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel1.restart()).once(); - verify(mockKernel2.restart()).once(); - }); - - test('handles notebooks with mixed cell types', async () => { - const notebook = createMockNotebook('deepnote', Uri.file('/test.ipynb'), [ - { languageId: 'python', metadata: {} }, - { languageId: 'sql', metadata: { sql_integration_id: 'postgres-1' } }, - { languageId: 'markdown', metadata: {} } - ]); - const mockKernel = mock(); - when(mockKernel.startedAtLeastOnce).thenReturn(true); - when(mockKernel.restart()).thenResolve(); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - when(kernelProvider.get(notebook)).thenReturn(instance(mockKernel)); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(mockKernel.restart()).once(); - }); - - test('does not restart when no notebooks are open', async () => { - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); - - handler.activate(); - onDidChangeIntegrations.fire(); - - await new Promise((resolve) => setTimeout(resolve, 10)); - - verify(kernelProvider.get(anything())).never(); - }); -}); diff --git a/src/notebooks/deepnote/integrations/integrationManager.ts b/src/notebooks/deepnote/integrations/integrationManager.ts index a98f3763bb..72b58dfada 100644 --- a/src/notebooks/deepnote/integrations/integrationManager.ts +++ b/src/notebooks/deepnote/integrations/integrationManager.ts @@ -1,11 +1,10 @@ import { inject, injectable } from 'inversify'; -import { commands, l10n, window, workspace } from 'vscode'; +import { commands, l10n, NotebookDocument, 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 { IntegrationStatus } from '../../../platform/notebooks/deepnote/integrationTypes'; import { IDeepnoteNotebookManager } from '../../types'; import { DatabaseIntegrationType, databaseIntegrationTypes } from '@deepnote/database-integrations'; @@ -14,10 +13,6 @@ import { DatabaseIntegrationType, databaseIntegrationTypes } from '@deepnote/dat */ @injectable() export class IntegrationManager implements IIntegrationManager { - private hasIntegrationsContext = 'deepnote.hasIntegrations'; - - private hasUnconfiguredIntegrationsContext = 'deepnote.hasUnconfiguredIntegrations'; - constructor( @inject(IExtensionContext) private readonly extensionContext: IExtensionContext, @inject(IIntegrationDetector) private readonly integrationDetector: IIntegrationDetector, @@ -38,90 +33,71 @@ export class IntegrationManager implements IIntegrationManager { // Find the integration ID from the arguments // It could be the first arg (if called directly) or in the args array (if called from UI) let integrationId: string | undefined; + let notebookUri: string | undefined; for (const arg of args) { if (typeof arg === 'string') { - integrationId = arg; - break; + integrationId ??= arg; + continue; } + notebookUri ??= this.extractNotebookUri(arg); } - logger.debug(`IntegrationManager: Extracted integrationId: ${integrationId}`); - return this.showIntegrationsUI(integrationId); + logger.debug(`IntegrationManager: Extracted integrationId: ${integrationId}, notebook: ${notebookUri}`); + + return this.showIntegrationsUI(integrationId, notebookUri); }) ); + } - // Listen for active notebook changes to update context - this.extensionContext.subscriptions.push( - window.onDidChangeActiveNotebookEditor(() => - this.updateContext().catch((err) => - logger.error('IntegrationManager: Failed to update context on notebook editor change', err) - ) - ) - ); + /** The notebook URI a menu contribution passed; `notebook/toolbar` sends `{ notebookEditor: { notebookUri } }`. */ + private extractNotebookUri(arg: unknown): string | undefined { + if (!arg || typeof arg !== 'object') { + return undefined; + } - // Listen for notebook document changes - this.extensionContext.subscriptions.push( - workspace.onDidOpenNotebookDocument(() => - this.updateContext().catch((err) => - logger.error('IntegrationManager: Failed to update context on notebook open', err) - ) - ) - ); + const candidate = arg as { notebookUri?: unknown; notebookEditor?: { notebookUri?: unknown } }; + const uri = candidate.notebookEditor?.notebookUri ?? candidate.notebookUri; - this.extensionContext.subscriptions.push( - workspace.onDidCloseNotebookDocument(() => - this.updateContext().catch((err) => - logger.error('IntegrationManager: Failed to update context on notebook close', err) - ) - ) - ); - - // Initial context update - this.updateContext().catch((err) => - logger.error('IntegrationManager: Failed to update context on activation', err) - ); + return uri ? String(uri) : undefined; } /** - * Update the context keys based on the active notebook + * The Deepnote notebook to act on: `window.activeNotebookEditor` is unset until an editor is focused, so a + * restored but not yet focused notebook resolves via the menu's URI or the one visible editor instead. */ - private async updateContext(): Promise { - const activeNotebook = window.activeNotebookEditor?.notebook; - - if (!activeNotebook || activeNotebook.notebookType !== 'deepnote') { - await commands.executeCommand('setContext', this.hasIntegrationsContext, false); - await commands.executeCommand('setContext', this.hasUnconfiguredIntegrationsContext, false); - return; + private resolveDeepnoteNotebook(notebookUri: string | undefined): NotebookDocument | undefined { + if (notebookUri) { + const fromUri = workspace.notebookDocuments.find( + (notebook) => notebook.notebookType === 'deepnote' && notebook.uri.toString() === notebookUri + ); + if (fromUri) { + return fromUri; + } } - const projectId = activeNotebook.metadata?.deepnoteProjectId; - const notebookId = activeNotebook.metadata?.deepnoteNotebookId; - if (!projectId || !notebookId) { - await commands.executeCommand('setContext', this.hasIntegrationsContext, false); - await commands.executeCommand('setContext', this.hasUnconfiguredIntegrationsContext, false); - return; + const active = window.activeNotebookEditor?.notebook; + if (active?.notebookType === 'deepnote') { + return active; } - // Detect integrations in the project - const integrations = await this.integrationDetector.detectIntegrations(projectId, notebookId); - const hasIntegrations = integrations.size > 0; - const hasUnconfigured = Array.from(integrations.values()).some( - (integration) => integration.status === IntegrationStatus.Disconnected - ); + // Only when unambiguous: several visible Deepnote editors give no basis for a guess. + const visible = window.visibleNotebookEditors + .map((editor) => editor.notebook) + .filter((notebook) => notebook.notebookType === 'deepnote'); - await commands.executeCommand('setContext', this.hasIntegrationsContext, hasIntegrations); - await commands.executeCommand('setContext', this.hasUnconfiguredIntegrationsContext, hasUnconfigured); + return visible.length === 1 ? visible[0] : undefined; } /** * Show the integrations management UI * @param selectedIntegrationId Optional integration ID to select/configure immediately + * @param notebookUri Optional notebook URI passed by the invoking menu contribution */ - private async showIntegrationsUI(selectedIntegrationId?: string): Promise { - const activeNotebook = window.activeNotebookEditor?.notebook; + private async showIntegrationsUI(selectedIntegrationId?: string, notebookUri?: string): Promise { + const activeNotebook = this.resolveDeepnoteNotebook(notebookUri); - if (!activeNotebook || activeNotebook.notebookType !== 'deepnote') { + if (!activeNotebook) { void window.showErrorMessage(l10n.t('No active Deepnote notebook')); return; } @@ -137,7 +113,11 @@ export class IntegrationManager implements IIntegrationManager { logger.trace(`IntegrationManager: Notebook metadata:`, activeNotebook.metadata); // First try to detect integrations from the stored project - let integrations = await this.integrationDetector.detectIntegrations(projectId, notebookId); + let integrations = await this.integrationDetector.detectIntegrations({ + notebookId, + notebookUri: activeNotebook.uri, + projectId + }); logger.debug(`IntegrationManager: Found ${integrations.size} integrations`); // If a specific integration was requested (e.g., from status bar click), @@ -167,7 +147,6 @@ export class IntegrationManager implements IIntegrationManager { } else { integrations.set(selectedIntegrationId, { config: config || null, - status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected, integrationName, integrationType }); diff --git a/src/notebooks/deepnote/integrations/integrationWebview.ts b/src/notebooks/deepnote/integrations/integrationWebview.ts index 0ba25f8c5f..99eaef3308 100644 --- a/src/notebooks/deepnote/integrations/integrationWebview.ts +++ b/src/notebooks/deepnote/integrations/integrationWebview.ts @@ -8,14 +8,14 @@ import { IDisposableRegistry, IExtensionContext } from '../../../platform/common import * as localize from '../../../platform/common/utils/localize'; import { logger } from '../../../platform/logging'; import { LocalizedMessages, SharedMessages } from '../../../messageTypes'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; import { IDeepnoteNotebookManager, ProjectIntegration } from '../../types'; import { persistProjectIntegrations } from './projectIntegrationsWriter'; import { IFederatedAuthTokenStorage, IIntegrationStorage, IIntegrationWebviewProvider } from './types'; import { ConfigurableDatabaseIntegrationConfig, FederatedAuthTokenStatus, - IntegrationStatus, - IntegrationWithStatus + DetectedIntegration } from '../../../platform/notebooks/deepnote/integrationTypes'; /** @@ -29,7 +29,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { private readonly disposables: Disposable[] = []; - private integrations: Map = new Map(); + private integrations: Map = new Map(); private projectId: string | undefined; @@ -43,6 +43,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, @inject(IDisposableRegistry) private readonly disposableRegistry: IDisposableRegistry, + @inject(ISqlIntegrationEnvVarsProvider) private readonly sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider, @inject(IFederatedAuthTokenStorage) @optional() private readonly tokenStorage?: IFederatedAuthTokenStorage @@ -69,7 +70,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { */ public async show( projectId: string, - integrations: Map, + integrations: Map, activeFileUri: Uri, selectedIntegrationId?: string, projectName?: string @@ -152,6 +153,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { integrationsTitle: localize.Integrations.title, integrationsNoIntegrationsFound: localize.Integrations.noIntegrationsFound, integrationsConnected: localize.Integrations.connected, + integrationsConfiguredInFile: localize.Integrations.configuredInFile, integrationsNotConfigured: localize.Integrations.notConfigured, integrationsConfigure: localize.Integrations.configure, integrationsReconfigure: localize.Integrations.reconfigure, @@ -458,21 +460,30 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { return; } + // Bumped before any await so the newest call always owns the highest generation; every earlier call + // then loses the comparison below no matter which of them resumes last. this.updateGeneration += 1; const generation = this.updateGeneration; + const [candidates, fileConfiguredIds] = await Promise.all([ + this.resolveFederatedAuthCandidates(), + this.resolveFileConfiguredIds() + ]); + const integrationsData = await Promise.all( Array.from(this.integrations.entries()).map(async ([id, integration]) => ({ + // SecretStorage only: the webview learns *that* an integration is file-configured, never the + // file's config or credentials. config: integration.config, id, integrationName: integration.integrationName, integrationType: integration.integrationType, - status: integration.status, - tokenStatus: await this.deriveTokenStatus(id, integration.config) + isFileConfigured: fileConfiguredIds.has(id), + tokenStatus: candidates.has(id) ? await this.deriveTokenStatus(id) : 'unsupported' })) ); - // Bail if the panel was disposed during the `tokenStorage.has()` await. + // Bail if the panel was disposed during the candidate/file-configured lookups or the `tokenStorage.has()` await. if (!this.currentPanel) { logger.debug('IntegrationWebviewProvider: Panel disposed during update, skipping postMessage'); return; @@ -514,12 +525,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { logger.info( `IntegrationWebviewProvider: deleting stale federated token for ${integrationId} (auth method changed).` ); - await this.tokenStorage.delete(integrationId).catch((err) => { - logger.warn( - `IntegrationWebviewProvider: failed to delete stale federated token for ${integrationId}`, - err - ); - }); + await this.tokenStorage.delete(integrationId); return; } @@ -530,34 +536,93 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { logger.info( `IntegrationWebviewProvider: deleting stale federated token for ${integrationId} (fingerprint changed).` ); - await this.tokenStorage.delete(integrationId).catch((err) => { - logger.warn( - `IntegrationWebviewProvider: failed to delete stale federated token for ${integrationId}`, - err - ); - }); + await this.tokenStorage.delete(integrationId); } } - /** Federated-auth token status: `'unsupported'` for non-BigQuery, non-google-oauth, or web; else `'authenticated'`/`'disconnected'`. */ - private async deriveTokenStatus( - integrationId: string, - config: ConfigurableDatabaseIntegrationConfig | null - ): Promise { - if (!this.tokenStorage) { - return 'unsupported'; + /** + * Ids eligible for federated auth in the active notebook — derived state only, so no `.deepnote.env.yaml` + * credentials enter the panel. A failed lookup degrades to "none eligible" rather than blocking the render. + */ + private async resolveFederatedAuthCandidates(): Promise> { + if (!this.activeFileUri) { + return new Set(); + } + + try { + return await this.sqlIntegrationEnvVars.getFederatedAuthCandidates(this.activeFileUri); + } catch (err) { + logger.warn('IntegrationWebviewProvider: failed to resolve federated auth candidates.', err); + + return new Set(); + } + } + + /** + * Ids `.deepnote.env.yaml` configures for the active notebook — ids only, so the panel can mark those rows + * read-only without ever holding file config. A failed lookup degrades to "none" rather than blocking the render. + */ + private async resolveFileConfiguredIds(): Promise> { + if (!this.activeFileUri) { + return new Set(); + } + + try { + return await this.sqlIntegrationEnvVars.getFileConfiguredIntegrationIds(this.activeFileUri); + } catch (err) { + logger.warn('IntegrationWebviewProvider: failed to resolve file-configured integration ids.', err); + + return new Set(); + } + } + + /** + * Refuses an edit to an integration `.deepnote.env.yaml` owns, and says so; returns `true` when it did. + * The panel writes SecretStorage, which the file-wins merge overrides, so going through with the edit would + * report success and change nothing at runtime. Each mutating action calls this itself rather than trusting + * the webview to have hidden its button — `showConfigurationForm` is reachable from the SQL status bar too. + * + * Re-resolves the file per action instead of reading what the last render saw, so it cannot be defeated by a + * stale snapshot, a lost `updateWebview()` generation race, or a switch to another notebook. That costs one + * extra `.deepnote.env.yaml` read per user-initiated edit (two on the `show(selectedIntegrationId)` path, which + * refreshes first) — rare and off the render path, so it is worth paying for a guard that cannot go stale. + * A read failure fails open: a hiccup must not block a real edit. + */ + private async refuseEditIfFileConfigured(integrationId: string): Promise { + const fileConfiguredIds = await this.resolveFileConfiguredIds(); + if (!fileConfiguredIds.has(integrationId)) { + return false; } - if (!config || config.type !== 'big-query' || config.metadata.authMethod !== BigQueryAuthMethods.GoogleOauth) { + + const name = this.integrations.get(integrationId)?.integrationName || integrationId; + + logger.debug( + `IntegrationWebviewProvider: Refused edit of ${integrationId}; it is configured in .deepnote.env.yaml` + ); + void window.showInformationMessage( + l10n.t( + "'{0}' is configured in .deepnote.env.yaml, which takes precedence over anything saved here. Edit that file to change it.", + name + ) + ); + + return true; + } + + /** Whether a federated-auth-eligible integration currently holds a token. Eligibility is the caller's call. */ + private async deriveTokenStatus(integrationId: string): Promise { + if (!this.tokenStorage) { return 'unsupported'; } + try { - const hasToken = await this.tokenStorage.has(integrationId); - return hasToken ? 'authenticated' : 'disconnected'; + return (await this.tokenStorage.has(integrationId)) ? 'authenticated' : 'disconnected'; } catch (err) { logger.warn( `IntegrationWebviewProvider: failed to check token for ${integrationId}; reporting disconnected.`, err ); + return 'disconnected'; } } @@ -592,7 +657,13 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { case 'authenticate': if (message.integrationId) { try { - await commands.executeCommand(Commands.AuthenticateIntegration, message.integrationId); + // Same URI the candidate set was derived from, so the button's eligibility and the + // command's config lookup cannot disagree about which notebook they mean. + await commands.executeCommand( + Commands.AuthenticateIntegration, + message.integrationId, + this.activeFileUri + ); } catch (error) { // Command handler shows its own toasts; log here to avoid an unhandled-rejection. logger.error( @@ -614,6 +685,10 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { return; } + if (await this.refuseEditIfFileConfigured(integrationId)) { + return; + } + await this.currentPanel?.webview.postMessage({ config: integration.config, integrationId, @@ -630,6 +705,10 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { integrationId: string, config: ConfigurableDatabaseIntegrationConfig ): Promise { + if (await this.refuseEditIfFileConfigured(integrationId)) { + return; + } + try { // Invalidate stale federated tokens before saving (fingerprint change or auth-method switch). await this.invalidateStaleFederatedToken(integrationId, config); @@ -641,7 +720,6 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { if (integration) { // Existing integration - update it integration.config = config; - integration.status = IntegrationStatus.Connected; integration.integrationName = config.name; integration.integrationType = config.type; this.integrations.set(integrationId, integration); @@ -649,7 +727,6 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { // New integration - add it to the map this.integrations.set(integrationId, { config, - status: IntegrationStatus.Connected, integrationName: config.name, integrationType: config.type }); @@ -681,17 +758,20 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { * Reset the configuration for an integration (clears credentials but keeps the integration entry) */ private async resetConfiguration(integrationId: string): Promise { + if (await this.refuseEditIfFileConfigured(integrationId)) { + return; + } + try { + // Token first: a failure here has to abort before the config is committed, otherwise the token is + // stranded with no integration left in the panel to retry from. + await this.tokenStorage?.delete(integrationId); await this.integrationStorage.delete(integrationId); - await this.tokenStorage?.delete(integrationId).catch((error) => { - logger.warn(`IntegrationWebviewProvider: failed to delete federated token for ${integrationId}`, error); - }); // Update local state const integration = this.integrations.get(integrationId); if (integration) { integration.config = null; - integration.status = IntegrationStatus.Disconnected; this.integrations.set(integrationId, integration); } @@ -721,11 +801,15 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider { * Delete the integration completely (removes credentials and integration entry) */ private async deleteConfiguration(integrationId: string): Promise { + if (await this.refuseEditIfFileConfigured(integrationId)) { + return; + } + try { + // Token first: a failure here has to abort before the config is committed, otherwise the token is + // stranded with no integration left in the panel to retry from. + await this.tokenStorage?.delete(integrationId); await this.integrationStorage.delete(integrationId); - await this.tokenStorage?.delete(integrationId).catch((error) => { - logger.warn(`IntegrationWebviewProvider: failed to delete federated token for ${integrationId}`, error); - }); // Remove from local state this.integrations.delete(integrationId); diff --git a/src/notebooks/deepnote/integrations/integrationWebview.unit.test.ts b/src/notebooks/deepnote/integrations/integrationWebview.unit.test.ts index 51d2b67bbb..6b432f86c8 100644 --- a/src/notebooks/deepnote/integrations/integrationWebview.unit.test.ts +++ b/src/notebooks/deepnote/integrations/integrationWebview.unit.test.ts @@ -3,16 +3,16 @@ import sinon from 'sinon'; import { EventEmitter, Uri } from 'vscode'; import { anyString, anything, instance, mock, reset, verify, when } from 'ts-mockito'; -import { IExtensionContext, IDisposable } from '../../../platform/common/types'; +import { IExtensionContext, IDisposable, Resource } from '../../../platform/common/types'; import { Commands } from '../../../platform/common/constants'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; import { IDeepnoteNotebookManager } from '../../types'; import { IntegrationWebviewProvider } from './integrationWebview'; import { FederatedAuthTokenEntry, IFederatedAuthTokenStorage, IIntegrationStorage } from './types'; import { computeMetadataFingerprint } from './federatedAuth/federatedAuthTokenStorage.node'; import { ConfigurableDatabaseIntegrationConfig, - IntegrationStatus, - IntegrationWithStatus + DetectedIntegration } from '../../../platform/notebooks/deepnote/integrationTypes'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; import { @@ -23,7 +23,7 @@ import { interface CapturedMessage { type: string; - integrations?: Array<{ id: string; tokenStatus?: string }>; + integrations?: Array<{ id: string; config?: unknown; isFileConfigured?: boolean; tokenStatus?: string }>; [key: string]: unknown; } @@ -46,7 +46,8 @@ function createFakeWebviewPanel(): FakeWebviewPanel { const webview = { html: '', cspSource: 'mock-csp', - asWebviewUri: (uri: unknown) => uri, + options: {}, + asWebviewUri: (uri: Uri) => uri, postMessage: (message: CapturedMessage) => postMessageImpl(message), onDidReceiveMessage: ( cb: (message: unknown) => Promise | void, @@ -59,8 +60,17 @@ function createFakeWebviewPanel(): FakeWebviewPanel { return disposable; } }; - const panel = { + const panel: import('vscode').WebviewPanel = { webview, + viewType: '', + title: '', + options: {}, + viewColumn: 1, + active: true, + visible: true, + onDidChangeViewState: function () { + return this; + }, reveal: () => undefined, dispose: () => undefined, onDidDispose: (cb: () => void, _thisArg?: unknown, disposables?: IDisposable[]): IDisposable => { @@ -71,7 +81,7 @@ function createFakeWebviewPanel(): FakeWebviewPanel { } }; return { - panel: panel as unknown as import('vscode').WebviewPanel, + panel, posted, onDidReceiveMessage: async (message: unknown) => { if (messageHandler) { @@ -90,11 +100,17 @@ function createFakeWebviewPanel(): FakeWebviewPanel { } suite('IntegrationWebviewProvider', () => { + const ACTIVE_FILE_URI = Uri.file('/ws/active.deepnote'); const PROJECT_ID = 'project-id-1'; let extensionContext: IExtensionContext; let integrationStorage: IIntegrationStorage; let notebookManager: IDeepnoteNotebookManager; + let federatedAuthCandidates: Set; + let candidatesSpy: sinon.SinonSpy<[Resource], Promise>>; + let fileConfiguredIds: Set; + let onDidChangeEnvironmentVariables: EventEmitter; + let sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider; let tokens: Map; let onDidChangeTokens: EventEmitter; let tokenSaveSpy: sinon.SinonSpy<[FederatedAuthTokenEntry, { silent?: boolean }?], Promise>; @@ -112,6 +128,20 @@ suite('IntegrationWebviewProvider', () => { when(extensionContext.subscriptions).thenReturn(extensionSubscriptions); when(extensionContext.extensionUri).thenReturn(Uri.file('/ext')); + // Federated-auth eligibility is derived state: the provider hands back ids only, never config. + federatedAuthCandidates = new Set(); + candidatesSpy = sinon.spy(async (_resource: Resource): Promise> => federatedAuthCandidates); + // Same shape for the `.deepnote.env.yaml` ids: read-only rows are derived from ids alone. + fileConfiguredIds = new Set(); + onDidChangeEnvironmentVariables = new EventEmitter(); + sqlIntegrationEnvVars = { + onDidChangeEnvironmentVariables: onDidChangeEnvironmentVariables.event, + getEnvironmentVariables: async () => ({}), + getFederatedAuthCandidates: candidatesSpy, + getFileConfiguredIntegrationIds: async () => fileConfiguredIds, + getMergedIntegrationConfigs: async () => [] + }; + tokens = new Map(); onDidChangeTokens = new EventEmitter(); tokenSaveSpy = sinon.spy(async (entry: FederatedAuthTokenEntry, options?: { silent?: boolean }) => { @@ -131,7 +161,6 @@ suite('IntegrationWebviewProvider', () => { delete: tokenDeleteSpy, get: async (id) => tokens.get(id), has: async (id) => tokens.has(id), - listIntegrationIds: async () => Array.from(tokens.keys()), save: tokenSaveSpy }; @@ -145,10 +174,12 @@ suite('IntegrationWebviewProvider', () => { reset(mockedVSCodeNamespaces.window); reset(mockedVSCodeNamespaces.commands); onDidChangeTokens.dispose(); + onDidChangeEnvironmentVariables.dispose(); }); function buildProvider( opts: { + sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider; tokenStorage?: IFederatedAuthTokenStorage; } = {} ): IntegrationWebviewProvider { @@ -157,6 +188,7 @@ suite('IntegrationWebviewProvider', () => { instance(integrationStorage), instance(notebookManager), extensionSubscriptions, + opts.sqlIntegrationEnvVars ?? sqlIntegrationEnvVars, opts.tokenStorage ); } @@ -164,12 +196,12 @@ suite('IntegrationWebviewProvider', () => { function singleIntegrationMap( id: string, config: ConfigurableDatabaseIntegrationConfig - ): Map { - return new Map([[id, { config, status: IntegrationStatus.Connected }]]); + ): Map { + return new Map([[id, { config }]]); } - async function show(provider: IntegrationWebviewProvider, integrations: Map) { - await provider.show(PROJECT_ID, integrations, Uri.file('/ws/active.deepnote')); + async function show(provider: IntegrationWebviewProvider, integrations: Map) { + await provider.show(PROJECT_ID, integrations, ACTIVE_FILE_URI); } function lastUpdate(): CapturedMessage { @@ -184,80 +216,185 @@ suite('IntegrationWebviewProvider', () => { }); } - suite('updateWebview tokenStatus matrix', () => { - ( - [ - { - name: 'no tokenStorage → unsupported', - tokenStorage: false as const, - config: () => buildGoogleOauthIntegration({ id: 'bq-1' }), - storeToken: false, - expected: 'unsupported' - }, - { - name: 'service-account BigQuery → unsupported', - tokenStorage: true as const, - config: () => buildServiceAccountIntegration({ id: 'bq-sa' }), - storeToken: false, - expected: 'unsupported' - }, - { - name: 'Postgres → unsupported', - tokenStorage: true as const, - config: () => buildPostgresIntegration({ id: 'pg-1' }), - storeToken: false, - expected: 'unsupported' - }, - { - name: 'BigQuery + google-oauth + stored token → authenticated', - tokenStorage: true as const, - config: () => buildGoogleOauthIntegration({ id: 'bq-2' }), - storeToken: true, - expected: 'authenticated' - }, - { - name: 'BigQuery + google-oauth + no stored token → disconnected', - tokenStorage: true as const, - config: () => buildGoogleOauthIntegration({ id: 'bq-3' }), - storeToken: false, - expected: 'disconnected' - } - ] as const - ).forEach((row) => { - test(row.name, async () => { - const config = row.config(); - if (row.storeToken) { - preStoreToken(config.id); - } - const provider = buildProvider({ - tokenStorage: row.tokenStorage ? tokenStorage : undefined - }); - await show(provider, singleIntegrationMap(config.id, config)); + suite('updateWebview tokenStatus', () => { + // Eligibility now comes entirely from the candidate set; `config` no longer gates the status. + test('candidate but no tokenStorage → unsupported', async () => { + const config = buildGoogleOauthIntegration({ id: 'bq-1' }); + federatedAuthCandidates.add(config.id); + + const provider = buildProvider(); + await show(provider, singleIntegrationMap(config.id, config)); + + const item = (lastUpdate().integrations || []).find((i) => i.id === config.id); + assert.strictEqual(item?.tokenStatus, 'unsupported'); + }); + + test('candidate + stored token → authenticated', async () => { + const config = buildGoogleOauthIntegration({ id: 'bq-2' }); + federatedAuthCandidates.add(config.id); + preStoreToken(config.id); + + const provider = buildProvider({ tokenStorage }); + await show(provider, singleIntegrationMap(config.id, config)); - const item = (lastUpdate().integrations || []).find((i) => i.id === config.id); - assert.strictEqual(item?.tokenStatus, row.expected); + const item = (lastUpdate().integrations || []).find((i) => i.id === config.id); + assert.strictEqual(item?.tokenStatus, 'authenticated'); + }); + + test('a candidate with no SecretStorage config gets a status while `config` stays null', async () => { + // A `.deepnote.env.yaml`-declared integration: authenticatable, but the panel holds no credentials + // for it and must not receive any from the file layer. + const integrationId = 'bq-file-only'; + federatedAuthCandidates.add(integrationId); + + const provider = buildProvider({ tokenStorage }); + await show( + provider, + new Map([ + [integrationId, { config: null, integrationName: 'File BigQuery', integrationType: 'big-query' }] + ]) + ); + + const item = (lastUpdate().integrations || []).find((i) => i.id === integrationId); + assert.strictEqual(item?.tokenStatus, 'disconnected'); + assert.isNull(item?.config, 'no `.deepnote.env.yaml` config may reach the webview payload'); + sinon.assert.calledWith(candidatesSpy, ACTIVE_FILE_URI); + }); + + test('a file-configured candidate keeps a real tokenStatus: read-only must not disable Authenticate', async () => { + // BigQuery + `google-oauth` declared in `.deepnote.env.yaml`: the config is read-only, but the OAuth + // token lives in SecretStorage, so authenticating is the one action the panel can still perform. + // Deriving federated-auth visibility from `isFileConfigured` would break exactly this row. + const integrationId = 'bq-file-configured-candidate'; + federatedAuthCandidates.add(integrationId); + fileConfiguredIds.add(integrationId); + + const provider = buildProvider({ tokenStorage }); + await show( + provider, + new Map([ + [integrationId, { config: null, integrationName: 'File BigQuery', integrationType: 'big-query' }] + ]) + ); + + const item = (lastUpdate().integrations || []).find((i) => i.id === integrationId); + assert.isTrue(item?.isFileConfigured, 'the row is file-configured, hence read-only'); + assert.strictEqual( + item?.tokenStatus, + 'disconnected', + 'a live token status must survive alongside `isFileConfigured`, or the Authenticate button disappears' + ); + }); + + test('a non-candidate reports unsupported even when a token exists', async () => { + const config = buildGoogleOauthIntegration({ id: 'bq-not-a-candidate' }); + preStoreToken(config.id); + + const provider = buildProvider({ tokenStorage }); + await show(provider, singleIntegrationMap(config.id, config)); + + const item = (lastUpdate().integrations || []).find((i) => i.id === config.id); + assert.strictEqual(item?.tokenStatus, 'unsupported'); + }); + + test('a rejected candidate lookup still renders the panel', async () => { + const config = buildGoogleOauthIntegration({ id: 'bq-lookup-fails' }); + preStoreToken(config.id); + + const provider = buildProvider({ + sqlIntegrationEnvVars: { + ...sqlIntegrationEnvVars, + getFederatedAuthCandidates: async () => { + throw new Error('merge failed'); + } + }, + tokenStorage }); + await show(provider, singleIntegrationMap(config.id, config)); + + const item = (lastUpdate().integrations || []).find((i) => i.id === config.id); + assert.strictEqual(item?.tokenStatus, 'unsupported', 'a failed lookup degrades to "no candidates"'); }); }); - test('handleMessage: "authenticate" → commands.executeCommand(AuthenticateIntegration, integrationId)', async () => { - const executeCommandStub = sinon.stub().resolves(undefined); - when(mockedVSCodeNamespaces.commands.executeCommand(anyString(), anything())).thenCall((command, arg) => - executeCommandStub(command, arg) + test('updateWebview flags `.deepnote.env.yaml`-configured ids as read-only, others not', async () => { + const fileConfig = buildPostgresIntegration({ id: 'pg-from-file' }); + const secretConfig = buildPostgresIntegration({ id: 'pg-from-secret-storage' }); + fileConfiguredIds.add(fileConfig.id); + + const provider = buildProvider({ tokenStorage }); + await show( + provider, + new Map([ + [fileConfig.id, { config: null, integrationName: fileConfig.name, integrationType: 'pgsql' }], + [secretConfig.id, { config: secretConfig }] + ]) ); - when(mockedVSCodeNamespaces.commands.executeCommand(anyString())).thenCall((command) => - executeCommandStub(command) + + const items = lastUpdate().integrations || []; + assert.isTrue( + items.find((i) => i.id === fileConfig.id)?.isFileConfigured, + 'a file-configured id must be marked read-only' + ); + assert.isFalse( + items.find((i) => i.id === secretConfig.id)?.isFileConfigured, + 'a SecretStorage-only id stays editable' + ); + }); + + test('a save for a file-configured id never reaches SecretStorage, whatever the webview sent', async () => { + // Read-only is enforced here, not just rendered: the SQL status bar's "Configure current integration" + // reaches `showConfigurationForm` directly, so a save can arrive for a row that has no Configure button. + // Letting it through would report success and change nothing, since the file wins the merge. + const integrationSaveSpy = sinon.spy(); + when(integrationStorage.save(anything())).thenCall(integrationSaveSpy); + + const config = buildPostgresIntegration({ id: 'pg-managed-by-file' }); + fileConfiguredIds.add(config.id); + + const provider = buildProvider({ tokenStorage }); + await show(provider, singleIntegrationMap(config.id, config)); + + await fakePanel.onDidReceiveMessage({ type: 'save', integrationId: config.id, config }); + + sinon.assert.notCalled(integrationSaveSpy); + assert.isFalse( + fakePanel.posted.some((message) => message.type === 'success'), + 'a refused edit must not be reported as saved' + ); + }); + + test('show() with a file-configured selectedIntegrationId opens no configuration form', async () => { + // The SQL status bar's "Configure current integration" routes through `Commands.ManageIntegrations` and + // lands on `showConfigurationForm` directly, so the panel's hidden Configure button never gets a say. + const config = buildPostgresIntegration({ id: 'pg-file-form' }); + fileConfiguredIds.add(config.id); + + const provider = buildProvider({ tokenStorage }); + await provider.show(PROJECT_ID, singleIntegrationMap(config.id, config), ACTIVE_FILE_URI, config.id); + + assert.isFalse( + fakePanel.posted.some((message) => message.type === 'showForm'), + 'an editable form must not open for an integration `.deepnote.env.yaml` owns' + ); + }); + + test('handleMessage: "authenticate" → executeCommand(AuthenticateIntegration, integrationId, activeFileUri)', async () => { + const executeCommandStub = sinon.stub().resolves(undefined); + when(mockedVSCodeNamespaces.commands.executeCommand(anyString(), anything(), anything())).thenCall( + (command, integrationId, resource) => executeCommandStub(command, integrationId, resource) ); const provider = buildProvider({ tokenStorage }); const integrationId = 'bq-auth'; + federatedAuthCandidates.add(integrationId); await show(provider, singleIntegrationMap(integrationId, buildGoogleOauthIntegration({ id: integrationId }))); await fakePanel.onDidReceiveMessage({ type: 'authenticate', integrationId }); assert.isTrue( - executeCommandStub.calledWith(Commands.AuthenticateIntegration, integrationId), - 'expected executeCommand to be called with AuthenticateIntegration and the integration id' + executeCommandStub.calledWith(Commands.AuthenticateIntegration, integrationId, ACTIVE_FILE_URI), + 'expected executeCommand to receive the id and the URI the candidate set was derived from' ); }); @@ -278,6 +415,38 @@ suite('IntegrationWebviewProvider', () => { sinon.assert.calledWith(tokenDeleteSpy, integrationId); verify(integrationStorage.delete(integrationId)).once(); }); + + test(`${messageType}Configuration: a failed token delete aborts before the config is removed`, async () => { + when(integrationStorage.delete(anyString())).thenResolve(); + + const provider = buildProvider({ + tokenStorage: { + ...tokenStorage, + delete: async () => { + throw new Error('keychain unavailable'); + } + } + }); + const integrationId = `bq-${messageType}-fails`; + preStoreToken(integrationId); + + await show( + provider, + singleIntegrationMap(integrationId, buildGoogleOauthIntegration({ id: integrationId })) + ); + await fakePanel.onDidReceiveMessage({ type: messageType, integrationId }); + + // Nothing is committed, so the integration stays in the panel for the user to retry from. + verify(integrationStorage.delete(integrationId)).never(); + assert.isTrue( + fakePanel.posted.some((message) => message.type === 'error'), + 'the failure must reach the panel' + ); + assert.isFalse( + fakePanel.posted.some((message) => message.type === 'success'), + 'a partial failure must not be reported as success' + ); + }); }); test('saveConfiguration: deletes the token BEFORE save when fingerprint changes', async () => { @@ -299,7 +468,7 @@ suite('IntegrationWebviewProvider', () => { clientId: 'new-client', clientSecret: 'new-secret' } - } as ConfigurableDatabaseIntegrationConfig); + }); await fakePanel.onDidReceiveMessage({ type: 'save', integrationId, config: newConfig }); @@ -308,6 +477,43 @@ suite('IntegrationWebviewProvider', () => { assert.isTrue(tokenDeleteSpy.calledBefore(integrationSaveSpy), 'token.delete must occur BEFORE storage.save'); }); + test('saveConfiguration: a failed token invalidation aborts the save', async () => { + const integrationId = 'bq-save-fails'; + const integrationSaveSpy = sinon.spy(); + when(integrationStorage.save(anything())).thenCall(integrationSaveSpy); + + const provider = buildProvider({ + tokenStorage: { + ...tokenStorage, + delete: async () => { + throw new Error('keychain unavailable'); + } + } + }); + preStoreToken(integrationId, 'old-fingerprint'); + await show(provider, singleIntegrationMap(integrationId, buildGoogleOauthIntegration({ id: integrationId }))); + + const newConfig = buildGoogleOauthIntegration({ + id: integrationId, + name: 'New name', + metadata: { + authMethod: 'google-oauth', + project: 'new-proj', + clientId: 'new-client', + clientSecret: 'new-secret' + } + }); + + await fakePanel.onDidReceiveMessage({ type: 'save', integrationId, config: newConfig }); + + // Saving anyway would pair the new client's config with a token issued against the old one. + sinon.assert.notCalled(integrationSaveSpy); + assert.isTrue( + fakePanel.posted.some((message) => message.type === 'error'), + 'the failure must reach the panel' + ); + }); + test('saveConfiguration: deletes the token when authMethod switches away from google-oauth', async () => { const integrationId = 'bq-switch'; when(integrationStorage.save(anything())).thenResolve(); @@ -400,9 +606,6 @@ suite('IntegrationWebviewProvider', () => { async delete() { /* no-op */ }, - async listIntegrationIds() { - return []; - }, computeMetadataFingerprint() { return 'fp'; } @@ -410,6 +613,8 @@ suite('IntegrationWebviewProvider', () => { const provider = buildProvider({ tokenStorage: slowTokenStorage }); const integrationId = 'bq-disposed-during-update'; + // Only candidates reach `deriveTokenStatus`, so the update parks on `has()` only if this id is one. + federatedAuthCandidates.add(integrationId); const integrations = singleIntegrationMap(integrationId, buildGoogleOauthIntegration({ id: integrationId })); const allPostedMessages: CapturedMessage[] = []; diff --git a/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.ts b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.ts new file mode 100644 index 0000000000..27ca87aea2 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.ts @@ -0,0 +1,164 @@ +import { inject, injectable } from 'inversify'; +import { NotebookDocument, RelativePattern, Uri, workspace } from 'vscode'; +import { DEFAULT_ENV_FILE, DEFAULT_INTEGRATIONS_FILE } from '@deepnote/database-integrations'; + +import { IExtensionSyncActivationService } from '../../../platform/activation/types'; +import { IFileSystem } from '../../../platform/common/platform/types'; +import { IDisposableRegistry } from '../../../platform/common/types'; +import { notebookPathToDeepnoteProjectFilePath } from '../../../platform/deepnote/deepnoteProjectUtils'; +import { isIntegrationsEnvFileEnabled } from '../../../platform/notebooks/deepnote/integrationsEnvFileSettings'; +import { logger } from '../../../platform/logging'; +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; +import { IIntegrationEnvLiveRefresher } from './types'; + +/** Trailing-edge debounce so a burst of edits (e.g. .env and .deepnote.env.yaml both saved) is handled once. */ +export const debounceTimeInMilliseconds = 500; + +const watchedEnvFileNames = [DEFAULT_INTEGRATIONS_FILE, DEFAULT_ENV_FILE]; + +/** Watches `.deepnote.env.yaml` / `.env` and live-refreshes affected notebooks' kernels on change (no restart). */ +@injectable() +export class IntegrationsEnvFileWatcher implements IExtensionSyncActivationService { + private readonly changedDirs = new Set(); + private debounceTimer: ReturnType | undefined; + private readonly watchedDirs = new Set(); + + constructor( + @inject(IIntegrationEnvLiveRefresher) private readonly liveRefresher: IIntegrationEnvLiveRefresher, + @inject(IFileSystem) private readonly fileSystem: IFileSystem, + @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry + ) {} + + public activate(): void { + for (const folder of workspace.workspaceFolders ?? []) { + this.watchDir(folder.uri); + } + + for (const notebook of workspace.notebookDocuments) { + this.watchNotebookDir(notebook); + } + + this.disposables.push( + workspace.onDidOpenNotebookDocument((notebook) => this.watchNotebookDir(notebook)), + workspace.onDidChangeWorkspaceFolders((event) => { + for (const folder of event.added) { + this.watchDir(folder.uri); + } + }), + { + dispose: () => { + if (this.debounceTimer) { + clearTimeout(this.debounceTimer); + this.debounceTimer = undefined; + } + } + } + ); + } + + private async findAffectedNotebooks(changedDirs: Set): Promise { + const affected: NotebookDocument[] = []; + + for (const notebook of workspace.notebookDocuments) { + if (notebook.notebookType !== DEEPNOTE_NOTEBOOK_TYPE) { + continue; + } + + const deepnoteFileUri = notebookPathToDeepnoteProjectFilePath(notebook.uri); + + // Same gate as IntegrationsFileConfigProvider: a disabled feature must not trigger kernel refreshes. + if (!isIntegrationsEnvFileEnabled(deepnoteFileUri)) { + continue; + } + + const deepnoteDir = Uri.joinPath(deepnoteFileUri, '..'); + const workspaceRoot = workspace.getWorkspaceFolder(notebook.uri)?.uri; + + const changedInScope = + changedDirs.has(deepnoteDir.fsPath) || (workspaceRoot != null && changedDirs.has(workspaceRoot.fsPath)); + if (!changedInScope) { + continue; + } + + // A `.env` change only affects integration env when a `.deepnote.env.yaml` actually exists for this + // notebook; without one the refresh is a no-op and its status message misleading, so an unrelated + // `.env` (a very common non-Deepnote file) must not trigger hidden kernel executions. + const candidateDirs = + workspaceRoot != null && workspaceRoot.fsPath !== deepnoteDir.fsPath + ? [deepnoteDir, workspaceRoot] + : [deepnoteDir]; + if (await this.hasIntegrationsFile(candidateDirs)) { + affected.push(notebook); + } + } + + return affected; + } + + private async handleChangedDirs(changedDirs: Set): Promise { + const affected = await this.findAffectedNotebooks(changedDirs); + if (affected.length === 0) { + return; + } + + await this.liveRefresher.refresh(affected); + } + + /** True when a `.deepnote.env.yaml` exists in any candidate dir (dir-then-root), mirroring the config provider's probe. */ + private async hasIntegrationsFile(dirs: Uri[]): Promise { + for (const dir of dirs) { + const candidate = Uri.joinPath(dir, DEFAULT_INTEGRATIONS_FILE); + if (await this.fileSystem.exists(candidate)) { + return true; + } + } + + return false; + } + + private onFileEvent(dir: Uri): void { + this.changedDirs.add(dir.fsPath); + + if (this.debounceTimer) { + clearTimeout(this.debounceTimer); + } + + this.debounceTimer = setTimeout(() => { + this.debounceTimer = undefined; + const dirs = new Set(this.changedDirs); + this.changedDirs.clear(); + this.handleChangedDirs(dirs).catch((error) => + logger.error('IntegrationsEnvFileWatcher: Failed to handle env file change', error) + ); + }, debounceTimeInMilliseconds); + } + + private watchDir(dir: Uri): void { + const dirPath = dir.fsPath; + if (this.watchedDirs.has(dirPath)) { + return; + } + this.watchedDirs.add(dirPath); + + for (const fileName of watchedEnvFileNames) { + const pattern = new RelativePattern(dir, fileName); + const watcher = workspace.createFileSystemWatcher(pattern, false, false, false); + + this.disposables.push( + watcher, + watcher.onDidChange(() => this.onFileEvent(dir)), + watcher.onDidCreate(() => this.onFileEvent(dir)), + watcher.onDidDelete(() => this.onFileEvent(dir)) + ); + } + } + + private watchNotebookDir(notebook: NotebookDocument): void { + if (notebook.notebookType !== DEEPNOTE_NOTEBOOK_TYPE) { + return; + } + + const deepnoteDir = Uri.joinPath(notebookPathToDeepnoteProjectFilePath(notebook.uri), '..'); + this.watchDir(deepnoteDir); + } +} diff --git a/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.unit.test.ts b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.unit.test.ts new file mode 100644 index 0000000000..222379dc90 --- /dev/null +++ b/src/notebooks/deepnote/integrations/integrationsEnvFileWatcher.node.unit.test.ts @@ -0,0 +1,237 @@ +import { assert } from 'chai'; +import { + Disposable, + EventEmitter, + FileSystemWatcher, + NotebookDocument, + RelativePattern, + Uri, + WorkspaceFolder, + WorkspaceFoldersChangeEvent +} from 'vscode'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; +import * as sinon from 'sinon'; +import { DEFAULT_ENV_FILE, DEFAULT_INTEGRATIONS_FILE } from '@deepnote/database-integrations'; + +import { IDisposable } from '../../../platform/common/types'; +import { IFileSystem } from '../../../platform/common/platform/types'; +import { IIntegrationEnvLiveRefresher } from './types'; +import { createMockNotebook } from '../deepnoteTestHelpers'; +import { debounceTimeInMilliseconds, IntegrationsEnvFileWatcher } from './integrationsEnvFileWatcher.node'; +import { dispose } from '../../../platform/common/utils/lifecycle'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; +import { notebookPathToDeepnoteProjectFilePath } from '../../../platform/deepnote/deepnoteProjectUtils'; + +suite('IntegrationsEnvFileWatcher', () => { + let watcher: IntegrationsEnvFileWatcher; + let liveRefresher: IIntegrationEnvLiveRefresher; + let fileSystem: IFileSystem; + let disposables: IDisposable[]; + let clock: sinon.SinonFakeTimers; + /** One emitter per file system watcher the code under test asked VS Code to create, keyed by dir + file name. */ + let fileEvents: Map>; + let onDidOpenNotebookDocument: EventEmitter; + + const workspaceFolder: WorkspaceFolder = { index: 0, name: 'ws', uri: Uri.file('/ws') }; + + /** Hands back a fake watcher whose events the test can fire, standing in for real filesystem events. */ + function createFileSystemWatcher(pattern: RelativePattern): FileSystemWatcher { + const emitter = new EventEmitter(); + fileEvents.set(watcherKey(pattern.baseUri, pattern.pattern), emitter); + disposables.push(emitter); + + const fsWatcher = mock(); + when(fsWatcher.onDidChange).thenReturn(emitter.event); + when(fsWatcher.onDidCreate).thenReturn(emitter.event); + when(fsWatcher.onDidDelete).thenReturn(emitter.event); + when(fsWatcher.dispose()).thenReturn(); + + return instance(fsWatcher); + } + + /** The dir the watcher derives from a notebook uri (the `.deepnote` file's dir). */ + function deepnoteDirOf(uri: Uri): Uri { + return Uri.joinPath(notebookPathToDeepnoteProjectFilePath(uri), '..'); + } + + /** Fires a watcher event and lets the debounce elapse, as a real edit of the file would. */ + async function fireEnvFileChange(dir: Uri, fileName: string = DEFAULT_INTEGRATIONS_FILE): Promise { + fireEnvFileEvent(dir, fileName); + + await clock.tickAsync(debounceTimeInMilliseconds); + } + + /** Fires a watcher event without advancing the clock, so callers can queue a burst within one debounce window. */ + function fireEnvFileEvent(dir: Uri, fileName: string): void { + const emitter = fileEvents.get(watcherKey(dir, fileName)); + if (!emitter) { + assert.fail(`No file system watcher was created for ${fileName} in ${dir.fsPath}`); + } + + emitter.fire(Uri.joinPath(dir, fileName)); + } + + function watcherKey(dir: Uri, fileName: string): string { + return `${dir.fsPath}::${fileName}`; + } + + setup(() => { + resetVSCodeMocks(); + disposables = [new Disposable(() => resetVSCodeMocks())]; + clock = sinon.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }); + + fileEvents = new Map>(); + when( + mockedVSCodeNamespaces.workspace.createFileSystemWatcher(anything(), anything(), anything(), anything()) + ).thenCall(createFileSystemWatcher); + + onDidOpenNotebookDocument = new EventEmitter(); + disposables.push(onDidOpenNotebookDocument); + when(mockedVSCodeNamespaces.workspace.onDidOpenNotebookDocument).thenReturn(onDidOpenNotebookDocument.event); + when(mockedVSCodeNamespaces.workspace.onDidChangeWorkspaceFolders).thenReturn( + new EventEmitter().event + ); + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder]); + + // The shared mock's `get()` ignores the default argument; the real API returns it when the setting is + // unset, which is what the `envFile.enabled` gate relies on. Individual tests override this. + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote', anything())).thenReturn({ + get: (_section: string, defaultValue?: unknown) => defaultValue + } as never); + + liveRefresher = mock(); + when(liveRefresher.refresh(anything())).thenResolve(); + + // Default: a `.deepnote.env.yaml` exists, so a dir change refreshes; individual tests override this. + fileSystem = mock(); + when(fileSystem.exists(anything())).thenResolve(true); + + watcher = new IntegrationsEnvFileWatcher(instance(liveRefresher), instance(fileSystem), disposables); + }); + + teardown(() => { + clock.restore(); + disposables = dispose(disposables); + }); + + test('refreshes every notebook view whose .deepnote dir changed (no deduplication; the refresher gates each kernel)', async () => { + // Two open views of the SAME .deepnote file (differ only by notebook query) — both are refreshed. + const uriA = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const uriB = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-b' }); + const notebookA = createMockNotebook({ uri: uriA }); + const notebookB = createMockNotebook({ uri: uriB }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookA, notebookB]); + watcher.activate(); + + await fireEnvFileChange(deepnoteDirOf(uriA)); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [notebookA, notebookB]); + }); + + test('resolves affected notebooks via the workspace-folder root (dir-then-root fallback)', async () => { + // The .deepnote lives in a nested dir; only the workspace ROOT changed. + const uri = Uri.file('/ws/nested/deep/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook({ uri }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(mockedVSCodeNamespaces.workspace.getWorkspaceFolder(anything())).thenReturn(workspaceFolder); + watcher.activate(); + + // The changed dir is the workspace root, NOT the .deepnote dir (/ws/nested/deep). + await fireEnvFileChange(workspaceFolder.uri); + + verify(liveRefresher.refresh(anything())).once(); + const [refreshed] = capture(liveRefresher.refresh).last(); + assert.deepStrictEqual([...refreshed], [notebook]); + }); + + test('coalesces a burst of env file events into a single refresh', async () => { + const uri = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook({ uri }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + watcher.activate(); + + // Both watched files saved at once — the debounce must collapse them into one refresh. + fireEnvFileEvent(deepnoteDirOf(uri), DEFAULT_INTEGRATIONS_FILE); + fireEnvFileEvent(deepnoteDirOf(uri), DEFAULT_ENV_FILE); + await clock.tickAsync(debounceTimeInMilliseconds); + + verify(liveRefresher.refresh(anything())).once(); + }); + + test('watches the dir of a notebook opened after activation', async () => { + // Outside the workspace folder, so the dir is watched only because the notebook was opened. + const uri = Uri.file('/elsewhere/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook({ uri }); + + watcher.activate(); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + onDidOpenNotebookDocument.fire(notebook); + + await fireEnvFileChange(deepnoteDirOf(uri)); + + verify(liveRefresher.refresh(anything())).once(); + }); + + test('does not refresh when the changed dir matches no open Deepnote notebook', async () => { + const otherFolder: WorkspaceFolder = { index: 1, name: 'other', uri: Uri.file('/other') }; + const uri = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook({ uri }); + + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder, otherFolder]); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + watcher.activate(); + + // An unrelated workspace folder changed — the notebook's dir and its workspace root are untouched. + await fireEnvFileChange(otherFolder.uri); + + verify(liveRefresher.refresh(anything())).never(); + }); + + test('ignores non-Deepnote notebooks even when their dir changed', async () => { + const uri = Uri.file('/ws/app.ipynb').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook({ uri, notebookType: 'jupyter-notebook' }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + watcher.activate(); + + // The dir is watched as a workspace folder, so the event lands — the notebook type must rule it out. + await fireEnvFileChange(deepnoteDirOf(uri)); + + verify(liveRefresher.refresh(anything())).never(); + }); + + test('does not refresh when the env-file feature is disabled for the notebook', async () => { + const uri = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook({ uri }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote', anything())).thenReturn({ + get: () => false + } as never); + watcher.activate(); + + await fireEnvFileChange(deepnoteDirOf(uri)); + + verify(liveRefresher.refresh(anything())).never(); + }); + + test('does not refresh when no .deepnote.env.yaml exists for the notebook (an unrelated .env change)', async () => { + const uri = Uri.file('/ws/proj/app.deepnote').with({ query: 'notebook=nb-a' }); + const notebook = createMockNotebook({ uri }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + // No integrations file present in any candidate dir — the change must be treated as unrelated. + when(fileSystem.exists(anything())).thenResolve(false); + watcher.activate(); + + await fireEnvFileChange(deepnoteDirOf(uri), DEFAULT_ENV_FILE); + + verify(liveRefresher.refresh(anything())).never(); + }); +}); diff --git a/src/notebooks/deepnote/integrations/sqlIntegrationStartupCodeProvider.ts b/src/notebooks/deepnote/integrations/sqlIntegrationStartupCodeProvider.ts deleted file mode 100644 index 693f1d8f11..0000000000 --- a/src/notebooks/deepnote/integrations/sqlIntegrationStartupCodeProvider.ts +++ /dev/null @@ -1,116 +0,0 @@ -// 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 3e151cb788..89ecf366aa 100644 --- a/src/notebooks/deepnote/integrations/types.ts +++ b/src/notebooks/deepnote/integrations/types.ts @@ -1,22 +1,28 @@ import type { DeepnoteBlock } from '@deepnote/blocks'; -import { Event, Uri } from 'vscode'; +import { Event, NotebookDocument, Uri } from 'vscode'; -import { IntegrationWithStatus } from '../../../platform/notebooks/deepnote/integrationTypes'; +import { DetectedIntegration } from '../../../platform/notebooks/deepnote/integrationTypes'; // Re-export IIntegrationStorage from platform layer export { IIntegrationStorage } from '../../../platform/notebooks/deepnote/types'; export const IIntegrationDetector = Symbol('IIntegrationDetector'); + +/** + * Inputs for {@link IIntegrationDetector.detectIntegrations}: the open notebook's already-validated Deepnote IDs, + * plus its URI — `.deepnote.env.yaml` is located relative to the notebook, so the ids alone cannot find it. + */ +export interface IntegrationDetectionInput { + notebookId: string; + notebookUri: Uri; + projectId: string; +} + export interface IIntegrationDetector { /** * Detect all integrations used in the given project */ - detectIntegrations(projectId: string, notebookId: string): Promise>; - - /** - * Check if a project has any unconfigured integrations - */ - hasUnconfiguredIntegrations(projectId: string, notebookId: string): Promise; + detectIntegrations(input: IntegrationDetectionInput): Promise>; } export const IIntegrationWebviewProvider = Symbol('IIntegrationWebviewProvider'); @@ -24,14 +30,14 @@ export interface IIntegrationWebviewProvider { /** * Show the integration management webview * @param projectId The Deepnote project ID - * @param integrations Map of integration IDs to their status + * @param integrations Map of integration IDs to their detected config and metadata * @param activeFileUri The `.deepnote` file being edited — always persisted to disk on save * @param selectedIntegrationId Optional integration ID to select/configure immediately * @param projectName Optional project display name (sourced from the active notebook's metadata) */ show( projectId: string, - integrations: Map, + integrations: Map, activeFileUri: Uri, selectedIntegrationId?: string, projectName?: string @@ -46,6 +52,12 @@ export interface IIntegrationManager { activate(): void; } +export const IIntegrationEnvLiveRefresher = Symbol('IIntegrationEnvLiveRefresher'); +export interface IIntegrationEnvLiveRefresher { + /** Re-runs the toolkit's `set_integration_env()` in each notebook's running kernel (no restart); notifies once. */ + refresh(notebooks: readonly NotebookDocument[]): Promise; +} + /** Persisted federated-auth token entry; fingerprints `${clientId}|${clientSecret}|${project}` to detect stale tokens. Only the refresh token is persisted. */ export interface FederatedAuthTokenEntry { integrationId: string; @@ -71,15 +83,21 @@ export interface IFederatedAuthTokenStorage { delete(integrationId: string): Promise; get(integrationId: string): Promise; has(integrationId: string): Promise; - /** All integration IDs with a stored token entry; used for orphaned-token cleanup. */ - listIntegrationIds(): Promise; /** Persists a token entry. Pass `silent: true` for refresh-token rotation to skip `onDidChangeTokens` (avoids interrupting in-flight SQL cells). */ save(entry: FederatedAuthTokenEntry, options?: { silent?: boolean }): Promise; } export const IFederatedAuthSqlBlockCodeGenerator = Symbol('IFederatedAuthSqlBlockCodeGenerator'); export interface IFederatedAuthSqlBlockCodeGenerator { - generate(block: DeepnoteBlock): Promise; + /** + * Python for a federated SQL block, or `undefined` when the block is not one (non-SQL, no integration id, or + * an integration that is not BigQuery + `google-oauth`). + * + * `notebookUri` is required, not optional: the integration config is resolved per notebook as + * `.deepnote.env.yaml` merged over SecretStorage, so the same integration id can carry different OAuth-client + * metadata in two notebooks and there is no ambient answer to fall back on. + */ + generate(block: DeepnoteBlock, notebookUri: Uri): Promise; } /** Thrown when a federated integration has no usable refresh token (not authenticated yet, fingerprint mismatch, or `invalid_grant`). */ diff --git a/src/notebooks/deepnote/integrations/userpodApiEndpoints.node.ts b/src/notebooks/deepnote/integrations/userpodApiEndpoints.node.ts new file mode 100644 index 0000000000..273a4a0006 --- /dev/null +++ b/src/notebooks/deepnote/integrations/userpodApiEndpoints.node.ts @@ -0,0 +1,222 @@ +import { timingSafeEqual } from 'crypto'; +import express, { type Express, type Request, type Response } from 'express'; +import * as http from 'http'; +import { inject, injectable } from 'inversify'; +import { commands, l10n, window, workspace } from 'vscode'; + +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; +import { IExtensionSyncActivationService } from '../../../platform/activation/types'; +import { IDisposableRegistry } from '../../../platform/common/types'; +import { generateUuid } from '../../../platform/common/uuid'; +import { logger } from '../../../platform/logging'; +import { ISqlIntegrationEnvVarsProvider, IUserpodApiEndpoints } from '../../../platform/notebooks/deepnote/types'; + +/** Loopback host for the toolkit's `userpod-api` calls; currently serves integration env vars for `set_integration_env()` (as `[{name,value}]`). */ +@injectable() +export class UserpodApiEndpoints implements IUserpodApiEndpoints, IExtensionSyncActivationService { + private readonly authTokensByProject = new Map(); + private isListening = false; + private server: http.Server | undefined; + private serverBaseUrl: string | undefined; + private startAttempt: Promise = Promise.resolve(); + + constructor( + @inject(ISqlIntegrationEnvVarsProvider) + private readonly sqlIntegrationEnvVarsProvider: ISqlIntegrationEnvVarsProvider, + @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry + ) {} + + public get baseUrl(): string | undefined { + return this.serverBaseUrl; + } + + /** Settles (never rejects) once the initial bind attempt completes, so callers can await readiness before reading `baseUrl`. */ + public get ready(): Promise { + return this.startAttempt; + } + + /** Per-project bearer token, generated on first use, so a kernel can only read its own project's credentials. */ + public getAuthToken(projectId: string): string { + let token = this.authTokensByProject.get(projectId); + if (token === undefined) { + token = generateUuid(); + this.authTokensByProject.set(projectId, token); + } + + return token; + } + + public activate(): void { + this.disposables.push({ dispose: () => this.stop() }); + this.startAttempt = this.startAndNotifyOnFailure(); + } + + private isAuthorized(authorizationHeader: string | undefined, projectId: string): boolean { + if (authorizationHeader === undefined) { + return false; + } + + const expectedToken = this.authTokensByProject.get(projectId); + if (expectedToken === undefined) { + return false; + } + + // Constant-time compare: the response carries credentials, so don't leak the token via early-exit timing. + const expected = Buffer.from(`Bearer ${expectedToken}`); + const actual = Buffer.from(authorizationHeader); + + return actual.length === expected.length && timingSafeEqual(actual, expected); + } + + private onServerError(err: Error): void { + logger.error('UserpodApiEndpoints: HTTP server error', err); + + // Startup failures are reported by startAndNotifyOnFailure(); only recover from errors once actually listening. + if (!this.isListening) { + return; + } + this.isListening = false; + this.serverBaseUrl = undefined; + + this.promptToRecover( + l10n.t( + 'The Deepnote integrations service stopped unexpectedly. SQL integrations will no longer receive their credentials. Reload the window to restore it.' + ) + ); + } + + /** + * Reload is the only offered remedy: restarting in place binds a NEW ephemeral port, and the endpoint's + * URL is baked into a toolkit server's env at spawn time, so already-running servers would keep pointing + * at the dead port and stay silently broken. + */ + private promptToRecover(message: string): void { + const reloadWindow = l10n.t('Reload Window'); + + void window.showErrorMessage(message, reloadWindow).then((choice) => { + if (choice === reloadWindow) { + void commands.executeCommand('workbench.action.reloadWindow'); + } + }); + } + + private async start(): Promise { + const app: Express = express(); + + app.get('/userpod-api/:projectId/integrations/environment-variables', async (req: Request, res: Response) => { + try { + // The `:projectId` route segment is always a single value at runtime; narrow the over-broad express type. + const projectId = Array.isArray(req.params.projectId) ? req.params.projectId[0] : req.params.projectId; + + if (!this.isAuthorized(req.headers.authorization, projectId)) { + res.status(401).json([]); + + return; + } + + // Filter (not find): sibling `.deepnote` files can share a project id, so serve the project's env + // vars rather than an arbitrary first match, merged deterministically by notebook uri. + const notebooks = workspace.notebookDocuments.filter( + (nb) => nb.notebookType === DEEPNOTE_NOTEBOOK_TYPE && nb.metadata?.deepnoteProjectId === projectId + ); + + if (notebooks.length === 0) { + res.json([]); + + return; + } + + if (notebooks.length > 1) { + logger.warn( + `UserpodApiEndpoints: ${notebooks.length} open notebooks share project '${projectId}'; merging their integration env vars.` + ); + } + + const ordered = [...notebooks].sort((a, b) => a.uri.toString().localeCompare(b.uri.toString())); + const resolved = await Promise.all( + ordered.map((notebook) => this.sqlIntegrationEnvVarsProvider.getEnvironmentVariables(notebook.uri)) + ); + + const merged = new Map(); + for (const envVars of resolved) { + for (const [name, value] of Object.entries(envVars ?? {})) { + if (!merged.has(name)) { + merged.set(name, value); + } + } + } + + const payload = [...merged].map(([name, value]) => ({ name, value })); + + res.json(payload); + } catch (err) { + logger.error('UserpodApiEndpoints: Failed to resolve integration environment variables', err); + res.status(500).send('Failed to resolve integration environment variables'); + } + }); + + const server = http.createServer(app); + this.server = server; + + // Persistent handler: an 'error' with no listener is re-thrown as an uncaught exception that crashes the host. + server.on('error', (err) => this.onServerError(err)); + + server.listen(0, '127.0.0.1'); + + const port = await new Promise((resolve, reject) => { + let onStartupError: (err: Error) => void = () => undefined; + let onListening: () => void = () => undefined; + onStartupError = (err: Error) => { + server.removeListener('listening', onListening); + reject(err); + }; + onListening = () => { + server.removeListener('error', onStartupError); + const address = server.address(); + if (!address || typeof address === 'string') { + reject(new Error('Integration env vars endpoint did not bind a port.')); + + return; + } + this.isListening = true; + resolve(address.port); + }; + server.once('error', onStartupError); + server.once('listening', onListening); + }); + + this.serverBaseUrl = `http://127.0.0.1:${port}`; + logger.info(`UserpodApiEndpoints: Listening on ${this.serverBaseUrl}`); + } + + /** Never rejects — `ready` is documented as settling regardless of the outcome. */ + private async startAndNotifyOnFailure(): Promise { + try { + await this.start(); + } catch (err) { + logger.error('UserpodApiEndpoints: Failed to start integration env vars endpoint', err); + this.promptToRecover( + l10n.t( + 'The Deepnote integrations service could not start. SQL integrations will not receive their credentials. Reload the window to try again.' + ) + ); + } + } + + private stop(): void { + const server = this.server; + this.server = undefined; + this.serverBaseUrl = undefined; + this.isListening = false; + if (!server) { + return; + } + + try { + server.closeAllConnections(); + server.close(); + } catch (err) { + logger.warn('UserpodApiEndpoints: Error while stopping server', err); + } + } +} diff --git a/src/notebooks/deepnote/integrations/userpodApiEndpoints.node.unit.test.ts b/src/notebooks/deepnote/integrations/userpodApiEndpoints.node.unit.test.ts new file mode 100644 index 0000000000..ed94716c56 --- /dev/null +++ b/src/notebooks/deepnote/integrations/userpodApiEndpoints.node.unit.test.ts @@ -0,0 +1,213 @@ +import * as http from 'http'; + +import { assert } from 'chai'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; +import { Disposable, NotebookDocument, Uri } from 'vscode'; + +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../../kernels/deepnote/types'; +import { IDisposable } from '../../../platform/common/types'; +import { dispose } from '../../../platform/common/utils/lifecycle'; +import { ISqlIntegrationEnvVarsProvider } from '../../../platform/notebooks/deepnote/types'; +import { createMockNotebook } from '../deepnoteTestHelpers'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; +import { UserpodApiEndpoints } from './userpodApiEndpoints.node'; + +suite('UserpodApiEndpoints', () => { + let endpoint: UserpodApiEndpoints; + let provider: ISqlIntegrationEnvVarsProvider; + let disposables: IDisposable[]; + + setup(() => { + resetVSCodeMocks(); + disposables = [new Disposable(() => resetVSCodeMocks())]; + provider = mock(); + + endpoint = new UserpodApiEndpoints(instance(provider), disposables); + }); + + teardown(() => { + // Disposing tears down the HTTP server (a dispose handler is pushed into `disposables` at start). + disposables = dispose(disposables); + }); + + /** A notebook carrying `projectId` in its metadata, which is how the endpoint's route lookup finds it. */ + function createNotebook(projectId: string, uri: Uri, notebookType: string = DEEPNOTE_NOTEBOOK_TYPE) { + return createMockNotebook({ notebookType, uri, metadata: { deepnoteProjectId: projectId } }); + } + + /** `activate()` is fire-and-forget; poll `baseUrl` until the server has bound a port. */ + async function waitForBaseUrl(timeoutMs = 3000): Promise { + const start = Date.now(); + while (endpoint.baseUrl === undefined) { + if (Date.now() - start > timeoutMs) { + throw new Error('UserpodApiEndpoints did not start listening in time'); + } + await new Promise((resolve) => setTimeout(resolve, 10)); + } + + return endpoint.baseUrl; + } + + /** Registers `notebooks` as the open documents, starts the endpoint, and returns its base URL. */ + function startWith(...notebooks: NotebookDocument[]): Promise { + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn(notebooks); + endpoint.activate(); + + return waitForBaseUrl(); + } + + function envVarsUrl(baseUrl: string, projectId: string): string { + return `${baseUrl}/userpod-api/${projectId}/integrations/environment-variables`; + } + + /** GET the endpoint carrying the per-project bearer token it requires. */ + function authedFetch(url: string, projectId: string): Promise { + return fetch(url, { headers: { Authorization: `Bearer ${endpoint.getAuthToken(projectId)}` } }); + } + + test('returns the provider env map as [{name,value}] for a matching open deepnote notebook', async () => { + when(provider.getEnvironmentVariables(anything())).thenResolve({ FOO: 'bar', BAZ: 'qux' }); + + const baseUrl = await startWith(createNotebook('project-1', Uri.file('/ws/app.deepnote'))); + + const response = await authedFetch(envVarsUrl(baseUrl, 'project-1'), 'project-1'); + + assert.strictEqual(response.status, 200); + const body = await response.json(); + assert.deepStrictEqual(body, [ + { name: 'FOO', value: 'bar' }, + { name: 'BAZ', value: 'qux' } + ]); + }); + + test('routes by projectId: queries only the notebook whose metadata matches the URL param', async () => { + const uriTwo = Uri.file('/ws/two.deepnote'); + when(provider.getEnvironmentVariables(anything())).thenResolve({ FROM: 'two' }); + + const baseUrl = await startWith( + createNotebook('project-one', Uri.file('/ws/one.deepnote')), + createNotebook('project-two', uriTwo) + ); + + const response = await authedFetch(envVarsUrl(baseUrl, 'project-two'), 'project-two'); + const body = await response.json(); + + assert.deepStrictEqual(body, [{ name: 'FROM', value: 'two' }]); + // The endpoint must resolve env vars for the matching notebook's uri, not the other project's. + verify(provider.getEnvironmentVariables(anything())).once(); + const [uriArg] = capture(provider.getEnvironmentVariables).last(); + assert.strictEqual((uriArg as Uri).toString(), uriTwo.toString()); + }); + + test('returns an empty array and never queries the provider when no notebook matches the projectId', async () => { + const baseUrl = await startWith(createNotebook('some-other-project', Uri.file('/ws/app.deepnote'))); + + const response = await authedFetch(envVarsUrl(baseUrl, 'project-1'), 'project-1'); + + assert.strictEqual(response.status, 200); + assert.deepStrictEqual(await response.json(), []); + verify(provider.getEnvironmentVariables(anything())).never(); + }); + + test('ignores non-Deepnote notebooks even when their projectId matches', async () => { + const baseUrl = await startWith(createNotebook('project-1', Uri.file('/ws/app.ipynb'), 'jupyter-notebook')); + + const response = await authedFetch(envVarsUrl(baseUrl, 'project-1'), 'project-1'); + + assert.deepStrictEqual(await response.json(), []); + verify(provider.getEnvironmentVariables(anything())).never(); + }); + + test('responds 500 when the provider rejects', async () => { + when(provider.getEnvironmentVariables(anything())).thenReject(new Error('resolution failed')); + + const baseUrl = await startWith(createNotebook('project-1', Uri.file('/ws/app.deepnote'))); + + const response = await authedFetch(envVarsUrl(baseUrl, 'project-1'), 'project-1'); + + assert.strictEqual(response.status, 500); + }); + + test('responds 401 and never queries the provider when the bearer token is missing or wrong', async () => { + const baseUrl = await startWith(createNotebook('project-1', Uri.file('/ws/app.deepnote'))); + + const noHeader = await fetch(envVarsUrl(baseUrl, 'project-1')); + assert.strictEqual(noHeader.status, 401); + + const wrongToken = await fetch(envVarsUrl(baseUrl, 'project-1'), { + headers: { Authorization: 'Bearer wrong-token' } + }); + assert.strictEqual(wrongToken.status, 401); + + verify(provider.getEnvironmentVariables(anything())).never(); + }); + + test('responds 401 for a wrong token of the SAME length (exercises the constant-time compare path)', async () => { + const baseUrl = await startWith(createNotebook('project-1', Uri.file('/ws/app.deepnote'))); + + // Issue the project's token, then present a DIFFERENT value of the same byte length — timingSafeEqual must reject it. + const realToken = endpoint.getAuthToken('project-1'); + const sameLengthWrong = `Bearer ${'x'.repeat(realToken.length)}`; + const response = await fetch(envVarsUrl(baseUrl, 'project-1'), { + headers: { Authorization: sameLengthWrong } + }); + + assert.strictEqual(response.status, 401); + verify(provider.getEnvironmentVariables(anything())).never(); + }); + + test('cross-project: a token issued for one project cannot read another project', async () => { + const baseUrl = await startWith( + createNotebook('project-a', Uri.file('/ws/a.deepnote')), + createNotebook('project-b', Uri.file('/ws/b.deepnote')) + ); + + // Issue tokens for both projects, then use project A's token against project B's URL. + const tokenA = endpoint.getAuthToken('project-a'); + endpoint.getAuthToken('project-b'); + + const response = await fetch(envVarsUrl(baseUrl, 'project-b'), { + headers: { Authorization: `Bearer ${tokenA}` } + }); + + assert.strictEqual(response.status, 401, "project A's token must not authorize a read of project B"); + verify(provider.getEnvironmentVariables(anything())).never(); + }); + + test('ready resolves once the server is listening', async () => { + endpoint.activate(); + + await endpoint.ready; + + assert.isString(endpoint.baseUrl, 'baseUrl must be set once ready resolves'); + }); + + test('prompts the user to recover when the initial bind fails, and never advertises a base URL', async () => { + when(mockedVSCodeNamespaces.window.showErrorMessage(anything(), anything())).thenResolve(undefined as never); + + endpoint.activate(); + + // `start()` assigns the server synchronously before its first await, so the bind failure can be + // simulated here — while `isListening` is still false, i.e. the initial-bind path. + const server = (endpoint as unknown as { server: http.Server }).server; + server.emit('error', new Error('EADDRINUSE')); + + await endpoint.ready; + + assert.strictEqual(endpoint.baseUrl, undefined, 'a failed bind must not advertise a base URL'); + verify(mockedVSCodeNamespaces.window.showErrorMessage(anything(), anything())).once(); + }); + + test('logs and prompts the user to recover when the server errors after startup, without crashing', async () => { + when(mockedVSCodeNamespaces.window.showErrorMessage(anything(), anything())).thenResolve(undefined as never); + + await startWith(createNotebook('project-1', Uri.file('/ws/app.deepnote'))); + + // Reach the running server to simulate a post-listen failure (e.g. an accept error under fd exhaustion). + const server = (endpoint as unknown as { server: http.Server }).server; + server.emit('error', new Error('accept failed')); + + assert.strictEqual(endpoint.baseUrl, undefined, 'a crashed endpoint must stop advertising its base URL'); + verify(mockedVSCodeNamespaces.window.showErrorMessage(anything(), anything())).once(); + }); +}); diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts index 5bf312977c..ec787a32db 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts @@ -24,11 +24,19 @@ import { IIntegrationStorage } from './integrations/types'; import { Commands } from '../../platform/common/constants'; import { ConfigurableDatabaseIntegrationType, - DATAFRAME_SQL_INTEGRATION_ID + DATAFRAME_SQL_INTEGRATION_ID, + isConfigurableDatabaseIntegrationType } from '../../platform/notebooks/deepnote/integrationTypes'; -import { IDeepnoteNotebookManager } from '../types'; +import { persistProjectIntegrations } from './integrations/projectIntegrationsWriter'; +import { IDeepnoteNotebookManager, ProjectIntegration } from '../types'; +import { logger } from '../../platform/logging'; +import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; +import type { DeepnoteFile } from '@deepnote/blocks'; import { DatabaseIntegrationType, databaseIntegrationTypes } from '@deepnote/database-integrations'; +/** One entry of a project's `integrations` list as it appears in the file, where `type` is not yet validated. */ +type RawProjectIntegration = NonNullable[number]; + /** * QuickPick item with an integration ID */ @@ -69,7 +77,9 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid constructor( @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, @inject(IIntegrationStorage) private readonly integrationStorage: IIntegrationStorage, - @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager + @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, + @inject(ISqlIntegrationEnvVarsProvider) + private readonly sqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider ) {} public activate(): void { @@ -229,12 +239,23 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid // Integration is configured, use the config name displayName = config.name; } else { - // Integration is not configured, try to get the name from the project's integration list - const notebookId = cell.notebook.metadata?.deepnoteNotebookId; - const project = notebookId ? this.notebookManager.getProjectForNotebook(projectId, notebookId) : undefined; - const projectIntegration = project?.project.integrations?.find((i) => i.id === integrationId); - const baseName = projectIntegration?.name || l10n.t('Unknown integration'); - displayName = l10n.t('{0} (configure)', baseName); + // Not in SecretStorage — a `.deepnote.env.yaml` file config still counts as configured, so check the + // merged configs before prompting the user to configure. + const fileConfig = (await this.sqlIntegrationEnvVars.getMergedIntegrationConfigs(cell.notebook.uri)).find( + (c) => c.id === integrationId + ); + if (fileConfig) { + displayName = fileConfig.name; + } else { + // Integration is not configured, try to get the name from the project's integration list + const notebookId = cell.notebook.metadata?.deepnoteNotebookId; + const project = notebookId + ? this.notebookManager.getProjectForNotebook(projectId, notebookId) + : undefined; + const projectIntegration = project?.project.integrations?.find((i) => i.id === integrationId); + const baseName = projectIntegration?.name || l10n.t('Unknown integration'); + displayName = l10n.t('{0} (configure)', baseName); + } } // Create a status bar item that opens the integration picker @@ -322,6 +343,65 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid this._onDidChangeCellStatusBarItems.fire(); } + /** + * Appends a picked integration to the project roster so the `.deepnote` file records what it uses. + * Additive only: existing entries pass through verbatim, never filtered, because a project's integrations + * are shared with sibling notebooks whose blocks this cannot see — dropping one would be data loss. + */ + private async addToProjectIntegrations( + cell: NotebookCell, + projectId: string, + roster: RawProjectIntegration[], + selected: RawProjectIntegration + ): Promise { + // No usable `type` means no valid roster entry; leave it out rather than guessing one. + if (!isConfigurableDatabaseIntegrationType(selected.type)) { + return; + } + + try { + await persistProjectIntegrations({ + notebookManager: this.notebookManager, + projectId, + // Cast rather than narrow: validating the existing entries would silently drop any type this + // build does not know about, which is pruning by another name. + integrations: [...roster, selected] as ProjectIntegration[], + activeFileUri: cell.notebook.uri + }); + } catch (error) { + // The cell metadata edit already succeeded, so the selection stands either way. + logger.error(`SqlCellStatusBarProvider: failed to add ${selected.id} to the project integrations`, error); + } + } + + /** + * The roster plus any `.deepnote.env.yaml` integrations it omits, so a file-only one can be picked here + * instead of only by hand-editing `sql_integration_id`. A failed lookup falls back to the roster alone. + */ + private async getSelectableIntegrations( + cell: NotebookCell, + projectIntegrations: RawProjectIntegration[] + ): Promise { + let mergedConfigs; + + try { + mergedConfigs = await this.sqlIntegrationEnvVars.getMergedIntegrationConfigs(cell.notebook.uri); + } catch (error) { + logger.error('SqlCellStatusBarProvider: failed to read file integrations; offering the roster only', error); + + return projectIntegrations; + } + + const rosterIds = new Set(projectIntegrations.map((integration) => integration.id)); + + // Anything merged but absent from the roster came from the file alone — the merge resolves roster ids first. + const fileOnly = mergedConfigs + .filter((config) => !rosterIds.has(config.id)) + .map((config) => ({ id: config.id, name: config.name, type: config.type })); + + return fileOnly.length > 0 ? [...projectIntegrations, ...fileOnly] : projectIntegrations; + } + private async switchIntegration(cell: NotebookCell): Promise { const currentIntegrationId = this.getIntegrationId(cell); @@ -343,7 +423,8 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid // Build quick pick items from project integrations const items: (QuickPickItem | LocalQuickPickItem)[] = []; - const projectIntegrations = project.project.integrations || []; + const roster = project.project.integrations || []; + const projectIntegrations = await this.getSelectableIntegrations(cell, roster); // Check if current integration is unknown (not in the project's list) const isCurrentIntegrationUnknown = @@ -457,6 +538,12 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid return; } + // Picking a file-declared integration is the one place the roster drifts, so reconcile it here. + const selectedIntegration = projectIntegrations.find((integration) => integration.id === selectedId); + if (selectedIntegration && !roster.some((integration) => integration.id === selectedId)) { + await this.addToProjectIntegrations(cell, projectId, roster, selectedIntegration); + } + // Trigger status bar update this._onDidChangeCellStatusBarItems.fire(); } diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts index f729453582..f608ea7850 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts @@ -1,5 +1,5 @@ import { assert } from 'chai'; -import { anything, instance, mock, verify, when } from 'ts-mockito'; +import { anything, capture, instance, mock, verify, when } from 'ts-mockito'; import { CancellationToken, CancellationTokenSource, EventEmitter, NotebookCell } from 'vscode'; import { IDisposableRegistry } from '../../platform/common/types'; @@ -11,6 +11,18 @@ import { createEventHandler } from '../../test/common'; import { Commands } from '../../platform/common/constants'; import { IDeepnoteNotebookManager } from '../types'; import { createMockCell } from './deepnoteTestHelpers'; +import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; + +/** + * A merged-config source that yields nothing, so integrations resolve from SecretStorage alone. + * These suites never assert against it, hence an instance rather than a mock handle. + */ +function emptySqlIntegrationEnvVars(): ISqlIntegrationEnvVarsProvider { + const provider = mock(); + when(provider.getMergedIntegrationConfigs(anything())).thenResolve([]); + + return instance(provider); +} suite('SqlCellStatusBarProvider', () => { let provider: SqlCellStatusBarProvider; @@ -23,7 +35,12 @@ suite('SqlCellStatusBarProvider', () => { disposables = []; integrationStorage = mock(); notebookManager = mock(); - provider = new SqlCellStatusBarProvider(disposables, instance(integrationStorage), instance(notebookManager)); + provider = new SqlCellStatusBarProvider( + disposables, + instance(integrationStorage), + instance(notebookManager), + emptySqlIntegrationEnvVars() + ); const tokenSource = new CancellationTokenSource(); cancellationToken = tokenSource.token; @@ -304,7 +321,8 @@ suite('SqlCellStatusBarProvider', () => { activateProvider = new SqlCellStatusBarProvider( activateDisposables, instance(activateIntegrationStorage), - instance(activateNotebookManager) + instance(activateNotebookManager), + emptySqlIntegrationEnvVars() ); }); @@ -448,6 +466,50 @@ suite('SqlCellStatusBarProvider', () => { verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).once(); }); + test('switchSqlIntegration offers `.deepnote.env.yaml` integrations the project roster omits', async () => { + let commandHandler: ((cell?: NotebookCell) => Promise) | undefined; + when(mockedVSCodeNamespaces.commands.registerCommand('deepnote.switchSqlIntegration', anything())).thenCall( + (_name, handler) => { + commandHandler = handler; + return { dispose: () => undefined }; + } + ); + + // A merged source reporting one integration the roster below never declares. + const envVars = mock(); + when(envVars.getMergedIntegrationConfigs(anything())).thenResolve([ + { id: 'file-only-bq', name: 'BigQuery from file', type: 'big-query', metadata: {} } as any + ]); + const fileAwareProvider = new SqlCellStatusBarProvider( + [], + instance(activateIntegrationStorage), + instance(activateNotebookManager), + instance(envVars) + ); + + const notebookMetadata = { deepnoteProjectId: 'project-1', deepnoteNotebookId: 'notebook-1' }; + const cell = createMockCell({ languageId: 'sql', notebookMetadata }); + when(activateNotebookManager.getProjectForNotebook('project-1', 'notebook-1')).thenReturn({ + project: { integrations: [] } + } as any); + + let offeredIds: string[] = []; + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenCall((items) => { + offeredIds = (items as { id?: string }[]).map((item) => item.id ?? ''); + + return Promise.resolve(undefined); + }); + + fileAwareProvider.activate(); + await commandHandler!(cell); + + assert.include( + offeredIds, + 'file-only-bq', + 'a file-only integration must be selectable, not reachable only by hand-editing sql_integration_id' + ); + }); + test('switchSqlIntegration command handler shows error when no cell and no active editor', async () => { let commandHandler: ((cell?: NotebookCell) => Promise) | undefined; when(mockedVSCodeNamespaces.commands.registerCommand('deepnote.switchSqlIntegration', anything())).thenCall( @@ -540,7 +602,8 @@ suite('SqlCellStatusBarProvider', () => { eventProvider = new SqlCellStatusBarProvider( eventDisposables, instance(eventIntegrationStorage), - instance(eventNotebookManager) + instance(eventNotebookManager), + emptySqlIntegrationEnvVars() ); }); @@ -668,7 +731,8 @@ suite('SqlCellStatusBarProvider', () => { commandProvider = new SqlCellStatusBarProvider( commandDisposables, instance(commandIntegrationStorage), - instance(commandNotebookManager) + instance(commandNotebookManager), + emptySqlIntegrationEnvVars() ); // Capture the command handler @@ -809,7 +873,8 @@ suite('SqlCellStatusBarProvider', () => { commandProvider = new SqlCellStatusBarProvider( commandDisposables, instance(commandIntegrationStorage), - instance(commandNotebookManager) + instance(commandNotebookManager), + emptySqlIntegrationEnvVars() ); // Capture the command handler @@ -864,6 +929,38 @@ suite('SqlCellStatusBarProvider', () => { verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).once(); }); + test('picking a `.deepnote.env.yaml` integration adds it to the project integrations list', async () => { + const notebookMetadata = { deepnoteProjectId: 'project-1', deepnoteNotebookId: 'notebook-1' }; + const cell = createMockCell({ languageId: 'sql', metadata: {}, notebookMetadata }); + const fileOnlyId = 'file-only-bq'; + + // Declared in the file only — the roster below never mentions it. + const envVars = mock(); + when(envVars.getMergedIntegrationConfigs(anything())).thenResolve([ + { id: fileOnlyId, name: 'BigQuery from file', type: 'big-query', metadata: {} } as any + ]); + new SqlCellStatusBarProvider( + [], + instance(commandIntegrationStorage), + instance(commandNotebookManager), + instance(envVars) + ).activate(); + + when(commandNotebookManager.getProjectForNotebook('project-1', 'notebook-1')).thenReturn({ + project: { integrations: [] } + } as any); + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( + Promise.resolve({ id: fileOnlyId, label: 'BigQuery from file' } as any) + ); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(true)); + + await switchIntegrationHandler(cell); + + const [projectId, integrations] = capture(commandNotebookManager.updateProjectIntegrations).last(); + assert.strictEqual(projectId, 'project-1'); + assert.deepStrictEqual(integrations, [{ id: fileOnlyId, name: 'BigQuery from file', type: 'big-query' }]); + }); + test('does not update if user cancels quick pick', async () => { const notebookMetadata = { deepnoteProjectId: 'project-1', deepnoteNotebookId: 'notebook-1' }; const cell = createMockCell({ diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index eb87ce58a3..38f7a8cdaf 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -44,6 +44,7 @@ import { DeepnoteActivationService } from './deepnote/deepnoteActivationService' import { DeepnoteNotebookManager } from './deepnote/deepnoteNotebookManager'; import { IDeepnoteNotebookManager } from './types'; import { IntegrationStorage } from '../platform/notebooks/deepnote/integrationStorage'; +import { IntegrationsFileConfigProvider } from '../platform/notebooks/deepnote/integrationsFileConfigProvider.node'; import { IntegrationDetector } from './deepnote/integrations/integrationDetector'; import { IntegrationManager } from './deepnote/integrations/integrationManager'; import { IntegrationWebviewProvider } from './deepnote/integrations/integrationWebview'; @@ -51,18 +52,20 @@ import { IFederatedAuthSqlBlockCodeGenerator, IFederatedAuthTokenStorage, IIntegrationDetector, + IIntegrationEnvLiveRefresher, IIntegrationManager, IIntegrationStorage, IIntegrationWebviewProvider } from './deepnote/integrations/types'; import { FederatedAuthCommandHandlerNode } from './deepnote/integrations/federatedAuth/federatedAuthCommandHandler.node'; import { FederatedAuthKernelRestartBridge } from './deepnote/integrations/federatedAuth/federatedAuthKernelRestartBridge.node'; -import { FederatedAuthOrphanedTokenCleaner } from './deepnote/integrations/federatedAuth/federatedAuthOrphanedTokenCleaner.node'; import { FederatedAuthSqlBlockCodeGenerator } from './deepnote/integrations/federatedAuth/federatedAuthSqlBlockCodeGenerator.node'; import { FederatedAuthTokenStorage } from './deepnote/integrations/federatedAuth/federatedAuthTokenStorage.node'; import { + IIntegrationsFileConfigProvider, IPlatformNotebookEditorProvider, - IPlatformDeepnoteNotebookManager + IPlatformDeepnoteNotebookManager, + IUserpodApiEndpoints } from '../platform/notebooks/deepnote/types'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; import { DirtyInputBlockStatusBarProvider } from './deepnote/dirtyInputBlockStatusBarProvider'; @@ -96,11 +99,13 @@ import { DeepnoteNotebookCommandListener } from './deepnote/deepnoteNotebookComm import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; -import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIntegrationStartupCodeProvider'; import { DeepnoteCellCopyHandler } from './deepnote/deepnoteCellCopyHandler'; import { DeepnoteEnvironmentTreeDataProvider } from '../kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.node'; import { OpenInDeepnoteHandler } from './deepnote/openInDeepnoteHandler.node'; -import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler'; +import { IntegrationEnvRefreshHandler } from './deepnote/integrations/integrationEnvRefreshHandler'; +import { IntegrationsEnvFileWatcher } from './deepnote/integrations/integrationsEnvFileWatcher.node'; +import { IntegrationEnvLiveRefresher } from './deepnote/integrations/integrationEnvLiveRefresher.node'; +import { UserpodApiEndpoints } from './deepnote/integrations/userpodApiEndpoints.node'; import { ISnapshotMetadataService, SnapshotService } from './deepnote/snapshots/snapshotService'; import { EnvironmentCapture, IEnvironmentCapture } from './deepnote/snapshots/environmentCapture.node'; import { DeepnoteFileChangeWatcher } from './deepnote/deepnoteFileChangeWatcher'; @@ -203,10 +208,6 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, FederatedAuthKernelRestartBridge ); - serviceManager.addSingleton( - IExtensionSyncActivationService, - FederatedAuthOrphanedTokenCleaner - ); serviceManager.addSingleton( IExtensionSyncActivationService, SqlCellStatusBarProvider @@ -221,20 +222,30 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea ); serviceManager.addSingleton( IExtensionSyncActivationService, - SqlIntegrationStartupCodeProvider + DeepnoteCellCopyHandler ); serviceManager.addSingleton( IExtensionSyncActivationService, - DeepnoteCellCopyHandler + OpenInDeepnoteHandler ); serviceManager.addSingleton( IExtensionSyncActivationService, - OpenInDeepnoteHandler + IntegrationEnvRefreshHandler + ); + serviceManager.addSingleton( + IIntegrationsFileConfigProvider, + IntegrationsFileConfigProvider ); serviceManager.addSingleton( IExtensionSyncActivationService, - IntegrationKernelRestartHandler + IntegrationsEnvFileWatcher + ); + serviceManager.addSingleton( + IIntegrationEnvLiveRefresher, + IntegrationEnvLiveRefresher ); + serviceManager.addSingleton(IUserpodApiEndpoints, UserpodApiEndpoints); + serviceManager.addBinding(IUserpodApiEndpoints, IExtensionSyncActivationService); // Deepnote kernel services serviceManager.addSingleton(DeepnoteAgentSkillsManager, DeepnoteAgentSkillsManager); diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index 28937d6b60..6d6b826180 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/src/notebooks/serviceRegistry.web.ts @@ -53,7 +53,6 @@ import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnote import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; -import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler'; import { FederatedAuthCommandHandlerWeb } from './deepnote/integrations/federatedAuth/federatedAuthCommandHandler.web'; import { DeepnoteFileChangeWatcher } from './deepnote/deepnoteFileChangeWatcher'; import { DeepnoteNotebookInfoStatusBar } from './deepnote/deepnoteNotebookInfoStatusBar'; @@ -139,10 +138,6 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, SqlCellStatusBarProvider ); - serviceManager.addSingleton( - IExtensionSyncActivationService, - IntegrationKernelRestartHandler - ); serviceManager.addSingleton( IExtensionSyncActivationService, FederatedAuthCommandHandlerWeb diff --git a/src/platform/common/utils/localize.ts b/src/platform/common/utils/localize.ts index ce2d7d7304..4f3c8f12e4 100644 --- a/src/platform/common/utils/localize.ts +++ b/src/platform/common/utils/localize.ts @@ -818,6 +818,7 @@ export namespace Integrations { export const title = l10n.t('Deepnote Integrations'); export const noIntegrationsFound = l10n.t('No integrations found in this project.'); export const connected = l10n.t('Connected'); + export const configuredInFile = l10n.t('Configured in file'); export const notConfigured = l10n.t('Not Configured'); export const configure = l10n.t('Configure'); export const reconfigure = l10n.t('Reconfigure'); diff --git a/src/platform/notebooks/deepnote/integrationTypes.ts b/src/platform/notebooks/deepnote/integrationTypes.ts index 8eb24affc5..c75018eb02 100644 --- a/src/platform/notebooks/deepnote/integrationTypes.ts +++ b/src/platform/notebooks/deepnote/integrationTypes.ts @@ -74,7 +74,14 @@ export interface LegacyDuckDBIntegrationConfig extends BaseLegacyIntegrationConf type: LegacyIntegrationType.DuckDB; } -import { DatabaseIntegrationConfig, DatabaseIntegrationType } from '@deepnote/database-integrations'; +import { + BigQueryAuthMethods, + DatabaseIntegrationConfig, + DatabaseIntegrationType, + databaseIntegrationTypes, + FederatedAuthMethod, + isFederatedAuthMethod +} from '@deepnote/database-integrations'; // Import and re-export Snowflake auth constants from shared module import { type SnowflakeAuthMethod, @@ -143,25 +150,25 @@ export type ConfigurableDatabaseIntegrationConfig = Extract< export type ConfigurableDatabaseIntegrationType = Exclude; -/** - * Integration connection status - */ -export enum IntegrationStatus { - Connected = 'connected', - Disconnected = 'disconnected', - Error = 'error' +/** Narrows a raw type string to one the webview can configure; excludes the internal DuckDB integration. */ +export function isConfigurableDatabaseIntegrationType( + type: string | undefined +): type is ConfigurableDatabaseIntegrationType { + return ( + type !== undefined && + type !== 'pandas-dataframe' && + (databaseIntegrationTypes as readonly string[]).includes(type) + ); } /** Federated-auth token status: `'authenticated'`, `'disconnected'` (federated but no token), or `'unsupported'` (non-federated or web/remote). */ export type FederatedAuthTokenStatus = 'authenticated' | 'disconnected' | 'unsupported'; /** - * Integration with its current status + * An integration declared by a project, paired with the credentials stored for it (if any) */ -export interface IntegrationWithStatus { +export interface DetectedIntegration { config: ConfigurableDatabaseIntegrationConfig | null; - status: IntegrationStatus; - error?: string; /** * Name from the project's integrations list (used for prefilling when config is null) */ @@ -170,6 +177,35 @@ export interface IntegrationWithStatus { * Type from the project's integrations list (used for prefilling when config is null) */ integrationType?: ConfigurableDatabaseIntegrationType; + /** + * `.deepnote.env.yaml` supplies this integration's config. The panel only ever edits SecretStorage, and the + * file wins the merge, so anything saved from here would be silently overridden — the row is read-only. + */ + isFileConfigured?: boolean; /** Federated-auth token status; only meaningful for federated integrations (currently BigQuery + `google-oauth`). */ tokenStatus?: FederatedAuthTokenStatus; } + +/** + * Narrows integration metadata to the federated-auth variant. Shared by the file-config provider and the SQL + * env-vars provider (upstream `isFederatedAuthMetadata`'s generic doesn't unify with our + * `DatabaseIntegrationConfig['metadata']` union); delegates to the exported `isFederatedAuthMethod` at runtime. + */ +export function isFederatedAuthMetadata( + metadata: DatabaseIntegrationConfig['metadata'] +): metadata is Extract { + if (typeof metadata !== 'object' || metadata === null) { + return false; + } + if (!('authMethod' in metadata)) { + return false; + } + const authMethod = metadata.authMethod; + + return typeof authMethod === 'string' && isFederatedAuthMethod(authMethod); +} + +/** The only federated combination this extension implements — see `FederatedAuthSqlBlockCodeGenerator`. */ +export function isSupportedFederatedAuth(integration: DatabaseIntegrationConfig): boolean { + return integration.type === 'big-query' && integration.metadata.authMethod === BigQueryAuthMethods.GoogleOauth; +} diff --git a/src/platform/notebooks/deepnote/integrationsEnvFileSettings.ts b/src/platform/notebooks/deepnote/integrationsEnvFileSettings.ts new file mode 100644 index 0000000000..00aad14320 --- /dev/null +++ b/src/platform/notebooks/deepnote/integrationsEnvFileSettings.ts @@ -0,0 +1,6 @@ +import { Uri, workspace } from 'vscode'; + +/** Shared gate for `deepnote.integrations.envFile.enabled` so provider and watcher stay in sync. */ +export function isIntegrationsEnvFileEnabled(deepnoteFileUri: Uri): boolean { + return workspace.getConfiguration('deepnote', deepnoteFileUri).get('integrations.envFile.enabled', true); +} diff --git a/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.ts b/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.ts new file mode 100644 index 0000000000..0398e1f16e --- /dev/null +++ b/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.ts @@ -0,0 +1,227 @@ +import dotenv from 'dotenv'; +import { inject, injectable } from 'inversify'; +import { Diagnostic, DiagnosticCollection, DiagnosticSeverity, languages, Range, Uri, workspace } from 'vscode'; + +import { + BUILTIN_INTEGRATIONS, + DatabaseIntegrationConfig, + DEFAULT_ENV_FILE, + DEFAULT_INTEGRATIONS_FILE, + parseIntegrations, + ValidationIssue +} from '@deepnote/database-integrations'; + +import { IFileSystem } from '../../common/platform/types'; +import { IDisposableRegistry } from '../../common/types'; +import { logger } from '../../logging'; +import { isIntegrationsEnvFileEnabled } from './integrationsEnvFileSettings'; +import { isFederatedAuthMetadata, isSupportedFederatedAuth } from './integrationTypes'; +import { IIntegrationsFileConfigProvider } from './types'; + +/** + * Stateless loader that reads integration configs from a `.deepnote.env.yaml` file (CLI parity), + * resolving `env:` references against a sibling `.env` file and `process.env`. Replicates the Node + * filesystem/dotenv shell that `@deepnote/database-integrations` does not export, delegating parsing + * to the exported, environment-agnostic `parseIntegrations`. + * + * No caching, no watching: a fresh read happens on every call. That was cheap when the only caller was + * kernel/server (re)start, but `getMergedIntegrationConfigs` now reaches here on every SQL cell execution and on every + * integrations-panel refresh, so each call is up to two `exists` + two `readFile` round-trips and re-publishes + * the YAML's diagnostics. Adding a cache means invalidating it on file change — deliberately not done yet. + */ +@injectable() +export class IntegrationsFileConfigProvider implements IIntegrationsFileConfigProvider { + private readonly diagnostics: DiagnosticCollection | undefined; + + constructor( + @inject(IFileSystem) private readonly fileSystem: IFileSystem, + @inject(IDisposableRegistry) disposables: IDisposableRegistry + ) { + this.diagnostics = languages.createDiagnosticCollection('deepnote-integrations'); + if (this.diagnostics) { + disposables.push(this.diagnostics); + } + } + + public async getConfigsForFile( + deepnoteFileUri: Uri + ): Promise<{ configs: DatabaseIntegrationConfig[]; issues: ValidationIssue[] }> { + try { + const candidateDirs = this.getCandidateDirs(deepnoteFileUri); + + // Both early returns clear first: nothing below republishes diagnostics, so a warning from a previous + // read would otherwise stay pinned in Problems after the file is fixed, deleted or the gate turned off. + if (!isIntegrationsEnvFileEnabled(deepnoteFileUri)) { + this.clearDiagnostics(candidateDirs); + + return { configs: [], issues: [] }; + } + + // Locate the integrations YAML (dir-then-root). A missing file is not an error. + const yamlUri = await this.findFirstExisting(candidateDirs, DEFAULT_INTEGRATIONS_FILE); + if (!yamlUri) { + this.clearDiagnostics(candidateDirs); + + return { configs: [], issues: [] }; + } + + const yaml = await this.fileSystem.readFile(yamlUri); + + // Locate the `.env` (dir-then-root) and resolve `env:` refs against it; real env wins over the file. + const envUri = await this.findFirstExisting(candidateDirs, DEFAULT_ENV_FILE); + const fileEnv = envUri ? dotenv.parse(await this.fileSystem.readFile(envUri)) : {}; + const env: Record = { ...fileEnv, ...this.getProcessEnvironment() }; + + const { integrations, issues } = parseIntegrations({ yaml, env }); + + const result = this.filterIntegrations(integrations, issues); + this.updateDiagnostics(yamlUri, result.issues); + + return result; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + const issue: ValidationIssue = { + path: '', + message: `Failed to read integrations file: ${message}`, + code: 'file_read_error' + }; + logger.error(`IntegrationsFileConfigProvider: ${issue.message}`); + + return { configs: [], issues: [issue] }; + } + } + + /** The process environment merged over the `.env` file; a seam tests override so they never touch the real `process.env`. */ + protected getProcessEnvironment(): Record { + return process.env; + } + + /** Drops any diagnostics published against a candidate dir's `.deepnote.env.yaml`, whether or not it still exists. */ + private clearDiagnostics(dirs: Uri[]): void { + dirs.forEach((dir) => this.diagnostics?.delete(Uri.joinPath(dir, DEFAULT_INTEGRATIONS_FILE))); + } + + /** Drops reserved, unsupported, duplicate, and unsupported federated-auth integrations, recording an issue for each. */ + private filterIntegrations( + integrations: DatabaseIntegrationConfig[], + parseIssues: ValidationIssue[] + ): { configs: DatabaseIntegrationConfig[]; issues: ValidationIssue[] } { + const configs: DatabaseIntegrationConfig[] = []; + const issues: ValidationIssue[] = [...parseIssues]; + const seenIds = new Set(); + + integrations.forEach((integration, index) => { + const issuePath = `integrations[${index}]`; + + if (BUILTIN_INTEGRATIONS.has(integration.id)) { + issues.push({ + path: issuePath, + message: `Integration '${integration.id}' uses a reserved id and was ignored.`, + code: 'reserved_integration_id' + }); + + return; + } + + if (integration.type === 'pandas-dataframe') { + issues.push({ + path: issuePath, + message: `Integration '${integration.id}' has unsupported type '${integration.type}' and was ignored.`, + code: 'unsupported_integration_type' + }); + + return; + } + + if (seenIds.has(integration.id)) { + issues.push({ + path: issuePath, + message: `Integration '${integration.id}' has a duplicate id and was ignored.`, + code: 'duplicate_integration_id' + }); + + return; + } + + if (isFederatedAuthMetadata(integration.metadata) && !isSupportedFederatedAuth(integration)) { + issues.push({ + path: issuePath, + message: `Integration '${integration.id}' uses an unsupported federated authentication method and was ignored.`, + code: 'unsupported_federated_integration' + }); + + return; + } + + seenIds.add(integration.id); + configs.push(integration); + }); + + issues.forEach((issue) => { + logger.warn(`IntegrationsFileConfigProvider: ${issue.code} at '${issue.path}': ${issue.message}`); + }); + + return { configs, issues }; + } + + private async findFirstExisting(dirs: Uri[], fileName: string): Promise { + for (const dir of dirs) { + const candidate = Uri.joinPath(dir, fileName); + if (await this.fileSystem.exists(candidate)) { + return candidate; + } + } + + return undefined; + } + + /** + * Candidate directories to look for the integration/env files in priority order: next to the + * `.deepnote` file first, then the workspace-folder root. Undefined entries are skipped and + * duplicates removed. + */ + private getCandidateDirs(deepnoteFileUri: Uri): Uri[] { + const dirs: Uri[] = [Uri.joinPath(deepnoteFileUri, '..')]; + const workspaceFolder = workspace.getWorkspaceFolder(deepnoteFileUri); + if (workspaceFolder) { + dirs.push(workspaceFolder.uri); + } + + const seen = new Set(); + + return dirs.filter((dir) => { + const key = dir.toString(); + if (seen.has(key)) { + return false; + } + seen.add(key); + + return true; + }); + } + + /** Surfaces validation issues in the Problems panel against the located `.deepnote.env.yaml` so a typo/missing key isn't silent; a clean parse clears them. No-op when diagnostics are unavailable (e.g. web/tests). */ + private updateDiagnostics(yamlUri: Uri, issues: ValidationIssue[]): void { + if (!this.diagnostics) { + return; + } + + if (issues.length === 0) { + this.diagnostics.delete(yamlUri); + + return; + } + + const diagnostics = issues.map((issue) => { + const detail = issue.path + ? `${issue.code} at '${issue.path}': ${issue.message}` + : `${issue.code}: ${issue.message}`; + const diagnostic = new Diagnostic(new Range(0, 0, 0, 0), detail, DiagnosticSeverity.Warning); + diagnostic.source = 'Deepnote integrations'; + + return diagnostic; + }); + + this.diagnostics.set(yamlUri, diagnostics); + } +} diff --git a/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.unit.test.ts b/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.unit.test.ts new file mode 100644 index 0000000000..4c60927684 --- /dev/null +++ b/src/platform/notebooks/deepnote/integrationsFileConfigProvider.node.unit.test.ts @@ -0,0 +1,362 @@ +import assert from 'assert'; + +import { + DatabaseIntegrationConfig, + DEFAULT_ENV_FILE, + DEFAULT_INTEGRATIONS_FILE +} from '@deepnote/database-integrations'; +import dedent from 'dedent'; +import { dump } from 'js-yaml'; +import { anything, instance, mock, verify, when } from 'ts-mockito'; +import { DiagnosticCollection, Uri, WorkspaceConfiguration, WorkspaceFolder } from 'vscode'; + +import { IFileSystem } from '../../common/platform/types'; +import { IntegrationsFileConfigProvider } from './integrationsFileConfigProvider.node'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; + +/** A pgsql entry for `pgIntegrationYaml`: `id`/`name` identify it, every other key overrides a metadata default. */ +interface PgIntegrationSpec { + id: string; + name: string; + [key: string]: string; +} + +/** A file present in the virtual filesystem: either readable `content`, or a `readError` that rejects. */ +interface VirtualFile { + path: string; + content?: string; + readError?: Error; +} + +/** Provider whose process environment is controllable, so tests never read or mutate the real `process.env`. */ +class TestableIntegrationsFileConfigProvider extends IntegrationsFileConfigProvider { + public processEnvironment: Record = {}; + + protected override getProcessEnvironment(): Record { + return this.processEnvironment; + } +} + +suite('IntegrationsFileConfigProvider', () => { + const deepnoteFileUri = Uri.file('/workspace/project/notebook.deepnote'); + const deepnoteDirUri = Uri.joinPath(deepnoteFileUri, '..'); + // Build the expected file paths exactly as the loader does (dir of the `.deepnote` file). + const yamlPath = Uri.joinPath(deepnoteDirUri, DEFAULT_INTEGRATIONS_FILE).fsPath; + const envPath = Uri.joinPath(deepnoteDirUri, DEFAULT_ENV_FILE).fsPath; + + let fileSystem: IFileSystem; + let provider: TestableIntegrationsFileConfigProvider; + let featureEnabled: boolean; + let workspaceFolder: WorkspaceFolder | undefined; + + setup(() => { + resetVSCodeMocks(); + + featureEnabled = true; + workspaceFolder = undefined; + + fileSystem = mock(); + provider = new TestableIntegrationsFileConfigProvider(instance(fileSystem), []); + + // The gate reads `deepnote.integrations.envFile.enabled`; return the current `featureEnabled` value. + when(mockedVSCodeNamespaces.workspace.getConfiguration(anything(), anything())).thenReturn({ + get: () => featureEnabled + } as unknown as WorkspaceConfiguration); + when(mockedVSCodeNamespaces.workspace.getWorkspaceFolder(anything())).thenCall(() => workspaceFolder); + }); + + /** Wires `IFileSystem.exists`/`readFile` to a small in-memory set of files keyed by fsPath. */ + function configureFileSystem(files: VirtualFile[]): void { + const byPath = new Map(files.map((file) => [file.path, file])); + + when(fileSystem.exists(anything())).thenCall((uri: Uri) => Promise.resolve(byPath.has(uri.fsPath))); + when(fileSystem.readFile(anything())).thenCall((uri: Uri) => { + const file = byPath.get(uri.fsPath); + if (!file) { + return Promise.reject(new Error(`ENOENT: ${uri.fsPath}`)); + } + if (file.readError) { + return Promise.reject(file.readError); + } + + return Promise.resolve(file.content ?? ''); + }); + } + + /** Reads a metadata field off a parsed config without narrowing the metadata union. */ + function metadataField(config: DatabaseIntegrationConfig, key: string): unknown { + return (config.metadata as unknown as Record)[key]; + } + + /** + * Rebuilds `provider` with a recording `DiagnosticCollection`. The default `languages` mock returns + * undefined, so a test that cares about diagnostics has to opt in before the constructor runs. + */ + function recordDiagnostics(): { deleted: string[]; set: string[] } { + const recorded: { deleted: string[]; set: string[] } = { deleted: [], set: [] }; + + when(mockedVSCodeNamespaces.languages.createDiagnosticCollection(anything())).thenReturn({ + delete: (uri: Uri) => recorded.deleted.push(uri.fsPath), + dispose: () => undefined, + set: (uri: Uri) => recorded.set.push(uri.fsPath) + } as unknown as DiagnosticCollection); + + provider = new TestableIntegrationsFileConfigProvider(instance(fileSystem), []); + + return recorded; + } + + /** An integrations document holding the given pgsql entries, each with valid default connection metadata. */ + function pgIntegrationYaml(...entries: PgIntegrationSpec[]): string { + return dump({ + integrations: entries.map(({ id, name, ...overrides }) => ({ + id, + name, + type: 'pgsql', + metadata: { + host: 'localhost', + port: '5432', + database: 'mydb', + user: 'root', + password: 'my-secret', + ...overrides + } + })) + }); + } + + test('returns configs for a valid integrations file', async () => { + configureFileSystem([ + { path: yamlPath, content: pgIntegrationYaml({ id: 'my-postgres', name: 'My Postgres' }) } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(configs[0].id, 'my-postgres'); + assert.strictEqual(configs[0].type, 'pgsql'); + assert.deepStrictEqual(issues, []); + }); + + test('lets the process environment override values from the .env file', async () => { + provider.processEnvironment = { DEEPNOTE_TEST_OVERRIDE_PASSWORD: 'secret-from-process-env' }; + configureFileSystem([ + { + path: yamlPath, + content: pgIntegrationYaml({ + id: 'override-postgres', + name: 'Override Postgres', + password: 'env:DEEPNOTE_TEST_OVERRIDE_PASSWORD' + }) + }, + { path: envPath, content: 'DEEPNOTE_TEST_OVERRIDE_PASSWORD=stale-from-dotenv\n' } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(metadataField(configs[0], 'password'), 'secret-from-process-env'); + assert.deepStrictEqual(issues, []); + }); + + test('returns an empty result when the YAML file is missing', async () => { + configureFileSystem([]); + + const result = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(result, { configs: [], issues: [] }); + }); + + test('reports a yaml_parse_error for malformed YAML', async () => { + configureFileSystem([{ path: yamlPath, content: 'integrations:\n - id: "unclosed string' }]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'yaml_parse_error'); + }); + + test('finds the YAML at the workspace-folder root when absent next to the .deepnote file', async () => { + const nestedDeepnoteUri = Uri.file('/workspace/project/sub/notebook.deepnote'); + const rootFolder: WorkspaceFolder = { uri: Uri.file('/workspace/project'), name: 'project', index: 0 }; + const rootYamlPath = Uri.joinPath(rootFolder.uri, DEFAULT_INTEGRATIONS_FILE).fsPath; + + workspaceFolder = rootFolder; + configureFileSystem([ + { path: rootYamlPath, content: pgIntegrationYaml({ id: 'root-postgres', name: 'Root Postgres' }) } + ]); + + const { configs, issues } = await provider.getConfigsForFile(nestedDeepnoteUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(configs[0].id, 'root-postgres'); + assert.deepStrictEqual(issues, []); + }); + + test('drops an integration whose id is reserved (reserved_integration_id)', async () => { + configureFileSystem([ + { path: yamlPath, content: pgIntegrationYaml({ id: 'deepnote-dataframe-sql', name: 'Reserved' }) } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'reserved_integration_id'); + assert.strictEqual(issues[0].path, 'integrations[0]'); + }); + + test('drops an integration with an unsupported pandas-dataframe type (unsupported_integration_type)', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: my-dataframe + name: My Dataframe + type: pandas-dataframe + metadata: {} + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'unsupported_integration_type'); + assert.strictEqual(issues[0].path, 'integrations[0]'); + }); + + test('drops a duplicate id, keeping the first occurrence (duplicate_integration_id)', async () => { + configureFileSystem([ + { + path: yamlPath, + content: pgIntegrationYaml( + { id: 'dup-postgres', name: 'First', host: 'first-host' }, + { id: 'dup-postgres', name: 'Second', host: 'second-host' } + ) + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(configs[0].id, 'dup-postgres'); + assert.strictEqual(metadataField(configs[0], 'host'), 'first-host'); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'duplicate_integration_id'); + assert.strictEqual(issues[0].path, 'integrations[1]'); + }); + + test('returns a file_read_error issue (and does not throw) when reading the file fails', async () => { + configureFileSystem([{ path: yamlPath, readError: new Error('disk failure') }]); + + // Must resolve, never reject: a read failure degrades to an issue, not a thrown error. + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'file_read_error'); + assert.strictEqual(issues[0].path, ''); + assert.ok(issues[0].message.includes('Failed to read integrations file')); + }); + + test('keeps a BigQuery google-oauth integration, resolving env: references inside its metadata', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: bq-oauth + name: BigQuery OAuth + type: big-query + metadata: + authMethod: google-oauth + project: my-project + clientId: my-client-id + clientSecret: "env:DEEPNOTE_TEST_BQ_CLIENT_SECRET" + ` + }, + { path: envPath, content: 'DEEPNOTE_TEST_BQ_CLIENT_SECRET=oauth-secret-from-dotenv\n' } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.strictEqual(configs.length, 1); + assert.strictEqual(configs[0].id, 'bq-oauth'); + assert.strictEqual(configs[0].type, 'big-query'); + assert.strictEqual(metadataField(configs[0], 'authMethod'), 'google-oauth'); + assert.strictEqual(metadataField(configs[0], 'clientSecret'), 'oauth-secret-from-dotenv'); + assert.deepStrictEqual(issues, []); + }); + + test('drops a federated integration using an unsupported method (unsupported_federated_integration)', async () => { + configureFileSystem([ + { + path: yamlPath, + content: dedent` + integrations: + - id: sf-okta + name: Snowflake Okta + type: snowflake + metadata: + authMethod: okta + accountName: my-account + clientId: my-client-id + clientSecret: my-client-secret + oktaSubdomain: my-subdomain + identityProvider: okta + authorizationServer: default + ` + } + ]); + + const { configs, issues } = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(configs, []); + assert.strictEqual(issues.length, 1); + assert.strictEqual(issues[0].code, 'unsupported_federated_integration'); + assert.strictEqual(issues[0].path, 'integrations[0]'); + }); + + test('returns an empty result without touching the filesystem when the feature is disabled', async () => { + featureEnabled = false; + + const result = await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(result, { configs: [], issues: [] }); + verify(fileSystem.exists(anything())).never(); + verify(fileSystem.readFile(anything())).never(); + }); + + test('clears a stale diagnostic once the YAML file is deleted', async () => { + const diagnostics = recordDiagnostics(); + + configureFileSystem([{ path: yamlPath, content: 'integrations:\n - id: "unclosed string' }]); + await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(diagnostics.set, [yamlPath], 'the malformed file should publish a diagnostic'); + + // Nothing below the early return republishes, so the warning would otherwise stay pinned in Problems. + configureFileSystem([]); + await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(diagnostics.deleted, [yamlPath]); + }); + + test('clears a stale diagnostic once the feature is disabled', async () => { + const diagnostics = recordDiagnostics(); + + configureFileSystem([{ path: yamlPath, content: 'integrations:\n - id: "unclosed string' }]); + await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(diagnostics.set, [yamlPath], 'the malformed file should publish a diagnostic'); + + featureEnabled = false; + await provider.getConfigsForFile(deepnoteFileUri); + + assert.deepStrictEqual(diagnostics.deleted, [yamlPath]); + }); +}); diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts index 7a612e927e..90b62410dc 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts @@ -1,37 +1,24 @@ import { inject, injectable } from 'inversify'; -import { CancellationToken, Event, EventEmitter } from 'vscode'; +import { CancellationToken, Event, EventEmitter, Uri } from 'vscode'; -import { - DatabaseIntegrationConfig, - FederatedAuthMethod, - getEnvironmentVariablesForIntegrations, - isFederatedAuthMethod -} from '@deepnote/database-integrations'; +import { DeepnoteFile } from '@deepnote/blocks'; +import { DatabaseIntegrationConfig, getEnvironmentVariablesForIntegrations } from '@deepnote/database-integrations'; import { IDisposableRegistry, Resource } from '../../common/types'; import { EnvironmentVariables } from '../../common/variables/types'; +import { notebookPathToDeepnoteProjectFilePath } from '../../deepnote/deepnoteProjectUtils'; import { logger } from '../../logging'; import { + IIntegrationsFileConfigProvider, IIntegrationStorage, ISqlIntegrationEnvVarsProvider, IPlatformNotebookEditorProvider, IPlatformDeepnoteNotebookManager } from './types'; -import { DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes'; - -/** Narrows metadata to the federated-auth variant; upstream `isFederatedAuthMetadata` can't be reused because its generic doesn't unify with our union. Delegates to upstream `isFederatedAuthMethod` at runtime. */ -function isFederatedAuthMetadata( - metadata: DatabaseIntegrationConfig['metadata'] -): metadata is Extract { - if (typeof metadata !== 'object' || metadata === null) { - return false; - } - if (!('authMethod' in metadata)) { - return false; - } - const authMethod = metadata.authMethod; - return typeof authMethod === 'string' && isFederatedAuthMethod(authMethod); -} +import { DATAFRAME_SQL_INTEGRATION_ID, isFederatedAuthMetadata, isSupportedFederatedAuth } from './integrationTypes'; + +/** One entry of a Deepnote project's `integrations` list. */ +type ProjectIntegration = NonNullable[number]; /** * Provides environment variables for SQL integrations. @@ -49,7 +36,9 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati @inject(IPlatformNotebookEditorProvider) private readonly notebookEditorProvider: IPlatformNotebookEditorProvider, @inject(IPlatformDeepnoteNotebookManager) private readonly notebookManager: IPlatformDeepnoteNotebookManager, - @inject(IDisposableRegistry) disposables: IDisposableRegistry + @inject(IDisposableRegistry) disposables: IDisposableRegistry, + @inject(IIntegrationsFileConfigProvider) + private readonly fileConfigProvider: IIntegrationsFileConfigProvider ) { logger.info('SqlIntegrationEnvironmentVariablesProvider: Constructor called - provider is being instantiated'); // Dispose emitter when extension deactivates @@ -111,21 +100,10 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati `SqlIntegrationEnvironmentVariablesProvider: Found ${projectIntegrations.length} integrations in project` ); - const configResults = await Promise.allSettled( - projectIntegrations.map((integration) => this.integrationStorage.getIntegrationConfig(integration.id)) - ); - const allConfigs: Array = configResults.flatMap((result, index) => { - if (result.status === 'fulfilled') { - return result.value ? [result.value] : []; - } - logger.error( - `SqlIntegrationEnvironmentVariablesProvider: Failed to load integration config ${projectIntegrations[index].id}`, - result.reason - ); - return []; - }); + const fileConfigs = await this.loadFileIntegrationConfigs(notebook.uri); + const allConfigs = await this.mergeIntegrationConfigs(projectIntegrations, fileConfigs); - // Skip federated-auth integrations: tokens are fetched per-cell via per-cell codegen in `FederatedAuthSqlBlockCodeGenerator`, not baked into kernel env. + // Skip federated-auth configs so OAuth client credentials are never injected into the kernel environment. const projectIntegrationConfigs: Array = []; for (const config of allConfigs) { if (isFederatedAuthMetadata(config.metadata)) { @@ -159,4 +137,145 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati return envVars; } + + /** Ids of merged configs that support federated auth (BigQuery + `google-oauth`); no credentials exposed. */ + public async getFederatedAuthCandidates( + resource: Resource, + token?: CancellationToken + ): Promise> { + const configs = await this.getMergedIntegrationConfigs(resource, token); + + return new Set(configs.filter(isSupportedFederatedAuth).map((config) => config.id)); + } + + /** Ids `.deepnote.env.yaml` configures for this notebook; ids only, so no file config or credentials leak out. */ + public async getFileConfiguredIntegrationIds( + resource: Resource, + token?: CancellationToken + ): Promise> { + if (!resource || token?.isCancellationRequested) { + return new Set(); + } + + // Unlike the merge, this needs no project/notebook id: the file is located from the notebook's own path. + const notebook = this.notebookEditorProvider.findAssociatedNotebookDocument(resource); + if (!notebook) { + return new Set(); + } + + const fileConfigs = await this.loadFileIntegrationConfigs(notebook.uri); + + return new Set(fileConfigs.map((config) => config.id)); + } + + /** SecretStorage integrations merged with `.deepnote.env.yaml` (file wins); excludes DuckDB; never pass to `save`. */ + public async getMergedIntegrationConfigs( + resource: Resource, + token?: CancellationToken + ): Promise { + if (!resource || token?.isCancellationRequested) { + return []; + } + + const notebook = this.notebookEditorProvider.findAssociatedNotebookDocument(resource); + if (!notebook) { + return []; + } + + const projectId = notebook.metadata?.deepnoteProjectId as string | undefined; + const notebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; + if (!projectId || !notebookId) { + return []; + } + + const project = this.notebookManager.getProjectForNotebook(projectId, notebookId); + if (!project) { + return []; + } + + const projectIntegrations = project.project.integrations?.slice() ?? []; + const fileConfigs = await this.loadFileIntegrationConfigs(notebook.uri); + + return this.mergeIntegrationConfigs(projectIntegrations, fileConfigs); + } + + /** Loads `.deepnote.env.yaml` configs (CLI parity); failures degrade to [] so SecretStorage still applies. */ + private async loadFileIntegrationConfigs(notebookUri: Uri): Promise { + try { + const result = await this.fileConfigProvider.getConfigsForFile( + notebookPathToDeepnoteProjectFilePath(notebookUri) + ); + result.issues.forEach((issue) => { + logger.warn( + `SqlIntegrationEnvironmentVariablesProvider: integrations file issue ${issue.code} at '${issue.path}': ${issue.message}` + ); + }); + + return result.configs; + } catch (error) { + logger.error( + 'SqlIntegrationEnvironmentVariablesProvider: file integrations source failed; falling back to SecretStorage', + error + ); + + return []; + } + } + + /** File config wins on id conflict; SecretStorage is the fallback for project ids the file lacks; file-only ids are appended additively (CLI parity). */ + private async mergeIntegrationConfigs( + projectIntegrations: ProjectIntegration[], + fileConfigs: DatabaseIntegrationConfig[] + ): Promise { + const fileConfigsById = new Map(fileConfigs.map((config) => [config.id, config])); + const consumedFileIds = new Set(); + + // Read from SecretStorage only the project integrations the file did not provide. + const secretStorageIds = projectIntegrations + .map((integration) => integration.id) + .filter((id) => !fileConfigsById.has(id)); + const secretStorageResults = await Promise.allSettled( + secretStorageIds.map((id) => this.integrationStorage.getIntegrationConfig(id)) + ); + const secretStorageConfigsById = new Map(); + secretStorageResults.forEach((result, index) => { + const id = secretStorageIds[index]; + if (result.status === 'fulfilled') { + if (result.value) { + secretStorageConfigsById.set(id, result.value); + } + + return; + } + logger.error( + `SqlIntegrationEnvironmentVariablesProvider: Failed to load integration config ${id}`, + result.reason + ); + }); + + // Resolve each project integration in declared order: file config wins, else the SecretStorage fallback. + const allConfigs: Array = []; + for (const integration of projectIntegrations) { + const fileConfig = fileConfigsById.get(integration.id); + if (fileConfig) { + consumedFileIds.add(integration.id); + allConfigs.push(fileConfig); + + continue; + } + const secretStorageConfig = secretStorageConfigsById.get(integration.id); + if (secretStorageConfig) { + allConfigs.push(secretStorageConfig); + } + } + + // Append file-only integrations (not declared in project.integrations) additively, deduped by the map. + for (const fileConfig of fileConfigsById.values()) { + if (!consumedFileIds.has(fileConfig.id)) { + allConfigs.push(fileConfig); + } + } + + return allConfigs; + } } diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts index 11b5e57eaa..b6b6230ff2 100644 --- a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.unit.test.ts @@ -1,12 +1,17 @@ import assert from 'assert'; import type { DeepnoteFile } from '@deepnote/blocks'; -import { instance, mock, when } from 'ts-mockito'; +import { anything, instance, mock, verify, when } from 'ts-mockito'; import { CancellationTokenSource, EventEmitter, NotebookDocument, Uri } from 'vscode'; import { IDisposableRegistry } from '../../common/types'; import { SqlIntegrationEnvironmentVariablesProvider } from './sqlIntegrationEnvironmentVariablesProvider'; -import { IIntegrationStorage, IPlatformDeepnoteNotebookManager, IPlatformNotebookEditorProvider } from './types'; -import { DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes'; +import { + IIntegrationsFileConfigProvider, + IIntegrationStorage, + IPlatformDeepnoteNotebookManager, + IPlatformNotebookEditorProvider +} from './types'; +import { ConfigurableDatabaseIntegrationConfig, DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes'; import { DatabaseIntegrationConfig } from '@deepnote/database-integrations'; /** Create a minimal `DeepnoteFile` for tests. */ @@ -29,7 +34,14 @@ function createMockProject( }; } +/** A file-config source that yields nothing — what callers see when no `.deepnote.env.yaml` exists. */ +function emptyFileConfigProvider(): IIntegrationsFileConfigProvider { + return { getConfigsForFile: async () => ({ configs: [], issues: [] }) }; +} + suite('SqlIntegrationEnvironmentVariablesProvider', () => { + const notebookUri = Uri.file('/ws/project.deepnote'); + const duckDbEnvVar = `SQL_${DATAFRAME_SQL_INTEGRATION_ID.toUpperCase().replace(/-/g, '_')}`; let provider: SqlIntegrationEnvironmentVariablesProvider; let integrationStorage: IIntegrationStorage; let notebookEditorProvider: IPlatformNotebookEditorProvider; @@ -37,6 +49,35 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { let disposables: IDisposableRegistry; let onDidChangeIntegrationsEmitter: EventEmitter; + /** A non-federated, non-reserved pgsql config whose host is embedded in the generated connection URL. */ + function pgConfig(id: string, host: string): ConfigurableDatabaseIntegrationConfig { + return { + id, + name: id, + type: 'pgsql', + metadata: { + host, + port: '5432', + database: 'db', + user: 'u', + password: 'p', + sslEnabled: false + } + }; + } + + /** Stubs the resource -> notebook -> project chain that every public method walks for `notebookUri`. */ + function stubNotebookWithProject(project: DeepnoteFile): void { + const notebook = mock(); + when(notebook.uri).thenReturn(notebookUri); + when(notebook.metadata).thenReturn({ + deepnoteProjectId: 'project-123', + deepnoteNotebookId: 'notebook-123' + }); + when(notebookEditorProvider.findAssociatedNotebookDocument(notebookUri)).thenReturn(instance(notebook)); + when(notebookManager.getProjectForNotebook('project-123', 'notebook-123')).thenReturn(project); + } + setup(() => { integrationStorage = mock(); notebookEditorProvider = mock(); @@ -50,7 +91,8 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { instance(integrationStorage), instance(notebookEditorProvider), instance(notebookManager), - disposables + disposables, + emptyFileConfigProvider() ); }); @@ -437,52 +479,271 @@ suite('SqlIntegrationEnvironmentVariablesProvider', () => { }); }); - suite('Federated-auth integrations are skipped', () => { - test('Mixed project: federated integration is skipped, non-federated is included', async () => { - const resource = Uri.file('/test/notebook.deepnote'); - const notebook = mock(); - const postgresConfig: DatabaseIntegrationConfig = { - id: 'pg-1', - name: 'Postgres', - type: 'pgsql', - metadata: { - host: 'localhost', - port: '5432', - database: 'db', - user: 'u', - password: 'p', - sslEnabled: false - } - }; - const federatedConfig: DatabaseIntegrationConfig = { - id: 'bq-oauth', - name: 'OAuth BQ', + suite('File config source (.deepnote.env.yaml) merge', () => { + let fileConfigProvider: IIntegrationsFileConfigProvider; + let providerWithFile: SqlIntegrationEnvironmentVariablesProvider; + + setup(() => { + fileConfigProvider = mock(); + providerWithFile = new SqlIntegrationEnvironmentVariablesProvider( + instance(integrationStorage), + instance(notebookEditorProvider), + instance(notebookManager), + disposables, + instance(fileConfigProvider) + ); + }); + + test('File wins on id conflict: file config used and SecretStorage is not queried for that id', async () => { + stubNotebookWithProject( + createMockProject('project-123', [ + { id: 'shared-db', name: 'shared-db', type: 'pgsql' }, + { id: 'secret-only', name: 'secret-only', type: 'pgsql' } + ]) + ); + when(fileConfigProvider.getConfigsForFile(anything())).thenResolve({ + configs: [pgConfig('shared-db', 'from-file.example.com')], + issues: [] + }); + // Stubbed with a different host to prove the file wins; the provider must never consult it for `shared-db`. + when(integrationStorage.getIntegrationConfig('shared-db')).thenResolve( + pgConfig('shared-db', 'from-secret.example.com') + ); + when(integrationStorage.getIntegrationConfig('secret-only')).thenResolve( + pgConfig('secret-only', 'secret-only.example.com') + ); + + const result = await providerWithFile.getEnvironmentVariables(notebookUri); + + const sharedUrl = JSON.parse(result['SQL_SHARED_DB']!).url as string; + assert.ok(sharedUrl.includes('from-file.example.com'), 'File config host should win the conflict'); + assert.ok(!sharedUrl.includes('from-secret.example.com'), 'SecretStorage host must not be used'); + assert.ok(result['SQL_SECRET_ONLY'], 'SecretStorage-only integration should still be resolved'); + + // The conflicting id must never hit SecretStorage; the SecretStorage-only id must be queried exactly once. + verify(integrationStorage.getIntegrationConfig('shared-db')).never(); + verify(integrationStorage.getIntegrationConfig('secret-only')).once(); + }); + + test('getMergedIntegrationConfigs returns the merged config list (file wins, SecretStorage fallback, file-only additive)', async () => { + stubNotebookWithProject( + createMockProject('project-123', [ + { id: 'shared-db', name: 'shared-db', type: 'pgsql' }, + { id: 'secret-only', name: 'secret-only', type: 'pgsql' } + ]) + ); + when(fileConfigProvider.getConfigsForFile(anything())).thenResolve({ + configs: [ + pgConfig('shared-db', 'from-file.example.com'), + pgConfig('file-only', 'file-only.example.com') + ], + issues: [] + }); + when(integrationStorage.getIntegrationConfig('secret-only')).thenResolve( + pgConfig('secret-only', 'secret-only.example.com') + ); + + const merged = await providerWithFile.getMergedIntegrationConfigs(notebookUri); + const byId = new Map(merged.map((config) => [config.id, config])); + + assert.deepStrictEqual( + [...byId.keys()].sort(), + ['file-only', 'secret-only', 'shared-db'], + 'merged configs must include the file-won, SecretStorage-fallback, and file-only integrations' + ); + const sharedDb = byId.get('shared-db'); + assert.ok( + sharedDb && JSON.stringify(sharedDb.metadata).includes('from-file.example.com'), + 'file config must win the id conflict in the merged list' + ); + assert.ok( + !byId.has(DATAFRAME_SQL_INTEGRATION_ID), + 'the internal DuckDB integration is not part of the merged list' + ); + verify(integrationStorage.getIntegrationConfig('shared-db')).never(); + }); + + test('getMergedIntegrationConfigs returns [] when the resource resolves to no project', async () => { + const merged = await providerWithFile.getMergedIntegrationConfigs(undefined); + + assert.deepStrictEqual(merged, []); + }); + + test('getFileConfiguredIntegrationIds returns the file config ids only', async () => { + stubNotebookWithProject( + createMockProject('project-123', [{ id: 'secret-only', name: 'secret-only', type: 'pgsql' }]) + ); + when(fileConfigProvider.getConfigsForFile(anything())).thenResolve({ + configs: [pgConfig('shared-db', 'from-file.example.com'), pgConfig('file-only', 'file-only.test')], + issues: [] + }); + + const ids = await providerWithFile.getFileConfiguredIntegrationIds(notebookUri); + + assert.deepStrictEqual( + ids, + new Set(['shared-db', 'file-only']), + 'SecretStorage-only ids must not be reported as file-configured' + ); + assert.deepStrictEqual( + await providerWithFile.getFileConfiguredIntegrationIds(undefined), + new Set(), + 'no resource means nothing to look up' + ); + assert.deepStrictEqual( + await providerWithFile.getFileConfiguredIntegrationIds(Uri.file('/ws/not-open.deepnote')), + new Set(), + 'a resource with no associated notebook resolves to no file configs' + ); + }); + + test('File source yields nothing: behavior is SecretStorage-only (unchanged)', async () => { + const providerWithoutFile = new SqlIntegrationEnvironmentVariablesProvider( + instance(integrationStorage), + instance(notebookEditorProvider), + instance(notebookManager), + disposables, + emptyFileConfigProvider() + ); + stubNotebookWithProject( + createMockProject('project-123', [{ id: 'secret-db', name: 'secret-db', type: 'pgsql' }]) + ); + when(integrationStorage.getIntegrationConfig('secret-db')).thenResolve( + pgConfig('secret-db', 'from-secret.example.com') + ); + + const result = await providerWithoutFile.getEnvironmentVariables(notebookUri); + + assert.ok(result['SQL_SECRET_DB'], 'SecretStorage integration should be resolved without a file provider'); + assert.ok( + JSON.parse(result['SQL_SECRET_DB']!).url.includes('from-secret.example.com'), + 'SecretStorage config should be used' + ); + assert.ok(result[duckDbEnvVar], 'DuckDB integration should always be included'); + verify(integrationStorage.getIntegrationConfig('secret-db')).once(); + }); + + test('File source throws: degrades to SecretStorage + DuckDB without rejecting', async () => { + stubNotebookWithProject( + createMockProject('project-123', [{ id: 'secret-db', name: 'secret-db', type: 'pgsql' }]) + ); + when(fileConfigProvider.getConfigsForFile(anything())).thenReject(new Error('boom')); + when(integrationStorage.getIntegrationConfig('secret-db')).thenResolve( + pgConfig('secret-db', 'from-secret.example.com') + ); + + const result = await providerWithFile.getEnvironmentVariables(notebookUri); + + assert.ok( + result['SQL_SECRET_DB'], + 'SecretStorage integration should still be resolved when the file source throws' + ); + assert.ok(result[duckDbEnvVar], 'DuckDB integration should still be included when the file source throws'); + }); + }); + + suite('Federated-auth candidates and env-var exclusion', () => { + let fileConfigProvider: IIntegrationsFileConfigProvider; + let providerWithFile: SqlIntegrationEnvironmentVariablesProvider; + + /** BigQuery + `google-oauth` — the one federated combination this extension implements. */ + function bigQueryOauthConfig(id: string, name: string): ConfigurableDatabaseIntegrationConfig { + return { + id, + name, type: 'big-query', metadata: { authMethod: 'google-oauth', project: 'oauth-project', - clientId: 'client', - clientSecret: 'secret' + clientId: `${id}-client-id`, + clientSecret: `${id}-client-secret` } }; - const project = createMockProject('project-123', [ - { id: 'pg-1', name: 'Postgres', type: 'pgsql' }, - { id: 'bq-oauth', name: 'OAuth BQ', type: 'big-query' } - ]); + } + + setup(() => { + fileConfigProvider = mock(); + providerWithFile = new SqlIntegrationEnvironmentVariablesProvider( + instance(integrationStorage), + instance(notebookEditorProvider), + instance(notebookManager), + disposables, + instance(fileConfigProvider) + ); + }); - when(notebook.metadata).thenReturn({ - deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'notebook-123' + test('File-sourced federated config reaches getMergedIntegrationConfigs but contributes no env vars', async () => { + stubNotebookWithProject(createMockProject('project-123', [])); + when(fileConfigProvider.getConfigsForFile(anything())).thenResolve({ + configs: [bigQueryOauthConfig('bq-file', 'File BQ'), pgConfig('pg-file', 'pg-file.example.com')], + issues: [] }); - when(notebookEditorProvider.findAssociatedNotebookDocument(resource)).thenReturn(instance(notebook)); - when(notebookManager.getProjectForNotebook('project-123', 'notebook-123')).thenReturn(project); - when(integrationStorage.getIntegrationConfig('pg-1')).thenResolve(postgresConfig); - when(integrationStorage.getIntegrationConfig('bq-oauth')).thenResolve(federatedConfig); - const result = await provider.getEnvironmentVariables(resource); + const merged = await providerWithFile.getMergedIntegrationConfigs(notebookUri); + const envVars = await providerWithFile.getEnvironmentVariables(notebookUri); + + assert.deepStrictEqual( + merged.map((config) => config.id), + ['bq-file', 'pg-file'], + 'the file-declared federated config must survive the merge; the SQL LSP and status bar need it' + ); + // Without the skip in `getEnvironmentVariables`, upstream emits every metadata key for the federated + // config (`FILE_BQ_CLIENTID` / `FILE_BQ_CLIENTSECRET`) and no usable `SQL_*` connection var for it. + assert.deepStrictEqual( + Object.keys(envVars).filter((name) => /_CLIENTID$|_CLIENTSECRET$/.test(name)), + [], + 'OAuth client credentials must never reach the kernel environment' + ); + assert.strictEqual( + envVars['SQL_BQ_FILE'], + undefined, + 'no connection var is emitted for a federated config' + ); + assert.ok(envVars['SQL_PG_FILE'], 'a federated config must not suppress its non-federated siblings'); + assert.ok(envVars[duckDbEnvVar], 'DuckDB integration should still be included'); + }); + + test('getFederatedAuthCandidates returns the supported federated ids from either source', async () => { + stubNotebookWithProject( + createMockProject('project-123', [ + { id: 'bq-secret-oauth', name: 'Secret BQ', type: 'big-query' }, + { id: 'bq-service-account', name: 'Service Account BQ', type: 'big-query' }, + { id: 'sf-native-oauth', name: 'Snowflake OAuth', type: 'snowflake' } + ]) + ); + when(fileConfigProvider.getConfigsForFile(anything())).thenResolve({ + configs: [bigQueryOauthConfig('bq-file-oauth', 'File BQ')], + issues: [] + }); + when(integrationStorage.getIntegrationConfig('bq-secret-oauth')).thenResolve( + bigQueryOauthConfig('bq-secret-oauth', 'Secret BQ') + ); + // BigQuery, but not federated at all. + when(integrationStorage.getIntegrationConfig('bq-service-account')).thenResolve({ + id: 'bq-service-account', + name: 'Service Account BQ', + type: 'big-query', + metadata: { + authMethod: 'service-account', + service_account: '{"type":"service_account","project_id":"test"}' + } + }); + // Federated, but not the combination `FederatedAuthSqlBlockCodeGenerator` implements. + when(integrationStorage.getIntegrationConfig('sf-native-oauth')).thenResolve({ + id: 'sf-native-oauth', + name: 'Snowflake OAuth', + type: 'snowflake', + metadata: { + authMethod: 'snowflake', + accountName: 'test-account', + clientId: 'sf-client-id', + clientSecret: 'sf-client-secret' + } + }); + + const candidates = await providerWithFile.getFederatedAuthCandidates(notebookUri); - assert.ok(result['SQL_PG_1'], 'Non-federated postgres env var should be present'); - assert.strictEqual(result['SQL_BQ_OAUTH'], undefined, 'Federated integration env var should be omitted'); + assert.deepStrictEqual(candidates, new Set(['bq-file-oauth', 'bq-secret-oauth'])); }); }); diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.web.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.web.ts new file mode 100644 index 0000000000..1659e28421 --- /dev/null +++ b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.web.ts @@ -0,0 +1,42 @@ +import { injectable } from 'inversify'; +import { EventEmitter } from 'vscode'; + +import { DatabaseIntegrationConfig } from '@deepnote/database-integrations'; + +import { Resource } from '../../common/types'; +import { DisposableBase } from '../../common/utils/lifecycle'; +import { EnvironmentVariables } from '../../common/variables/types'; +import { ISqlIntegrationEnvVarsProvider } from './types'; + +/** + * Web stub for `ISqlIntegrationEnvVarsProvider`. Integration credentials are only resolved for locally + * launched kernels, which web does not have, so this yields no env vars and no configs. Binding it keeps + * the dependency non-optional for every consumer instead of making them branch on a missing provider. + */ +@injectable() +export class SqlIntegrationEnvironmentVariablesProviderWeb + extends DisposableBase + implements ISqlIntegrationEnvVarsProvider +{ + private readonly _onDidChangeEnvironmentVariables = this._register(new EventEmitter()); + + public readonly onDidChangeEnvironmentVariables = this._onDidChangeEnvironmentVariables.event; + + public async getEnvironmentVariables(): Promise { + return {}; + } + + /** Federated auth is node-only (no codegen is bound on web), so nothing is ever a candidate here. */ + public async getFederatedAuthCandidates(): Promise> { + return new Set(); + } + + /** No filesystem on web, so `.deepnote.env.yaml` never contributes configs here. */ + public async getFileConfiguredIntegrationIds(): Promise> { + return new Set(); + } + + public async getMergedIntegrationConfigs(): Promise { + return []; + } +} diff --git a/src/platform/notebooks/deepnote/types.ts b/src/platform/notebooks/deepnote/types.ts index 23be88f5da..6dd16ae811 100644 --- a/src/platform/notebooks/deepnote/types.ts +++ b/src/platform/notebooks/deepnote/types.ts @@ -3,6 +3,7 @@ import { IDisposable, Resource } from '../../common/types'; import { EnvironmentVariables } from '../../common/variables/types'; import { ConfigurableDatabaseIntegrationConfig } from './integrationTypes'; import { DeepnoteFile } from '@deepnote/blocks'; +import { DatabaseIntegrationConfig, ValidationIssue } from '@deepnote/database-integrations'; /** * Settings for select input blocks @@ -96,6 +97,53 @@ export interface ISqlIntegrationEnvVarsProvider { * Get environment variables for SQL integrations used in the given notebook. */ getEnvironmentVariables(resource: Resource, token?: CancellationToken): Promise; + + /** + * Ids that can be federated-authenticated for this notebook — BigQuery + `google-oauth`, from either + * source. Derived state only: it exposes **no config**, so the integrations panel (an editor) can offer an + * Authenticate action without receiving credentials it cannot write back. + */ + getFederatedAuthCandidates(resource: Resource, token?: CancellationToken): Promise>; + + /** + * Ids that `.deepnote.env.yaml` configures for this notebook. Derived state only — it exposes **no + * config and no credentials** — so the integrations panel (a SecretStorage editor) can render them + * read-only without receiving values it could never write back. + */ + getFileConfiguredIntegrationIds(resource: Resource, token?: CancellationToken): Promise>; + + /** + * What actually applies for this notebook: project SecretStorage integrations merged with + * `.deepnote.env.yaml` file configs (file wins, additive file-only), so integration detection, the SQL + * status bar, and the SQL LSP agree with kernel execution. **Read-only** — the file layer cannot be + * written back, so these must never reach `IIntegrationStorage.save`, which only ever edits SecretStorage. + */ + getMergedIntegrationConfigs(resource: Resource, token?: CancellationToken): Promise; +} + +export const IUserpodApiEndpoints = Symbol('IUserpodApiEndpoints'); +export interface IUserpodApiEndpoints { + /** Loopback base URL the toolkit fetches integration env vars from; `undefined` until the server is listening. */ + readonly baseUrl: string | undefined; + + /** Settles (never rejects) once the initial bind attempt completes, so callers can await readiness before reading `baseUrl`. */ + readonly ready: Promise; + + /** Per-project bearer token the endpoint requires (it serves credentials); injected into the toolkit as `DEEPNOTE_RUNTIME__PROJECT_SECRET`. */ + getAuthToken(projectId: string): string; +} + +export const IIntegrationsFileConfigProvider = Symbol('IIntegrationsFileConfigProvider'); +export interface IIntegrationsFileConfigProvider { + /** + * Loads integration configs from a `.deepnote.env.yaml` file located next to the given `.deepnote` + * project file (or at the workspace-folder root), resolving `env:` references against a sibling + * `.env` file and `process.env`. Returns the accepted configs plus any validation issues; a missing + * file yields an empty result and this never throws. + */ + getConfigsForFile( + deepnoteFileUri: Uri + ): Promise<{ configs: DatabaseIntegrationConfig[]; issues: ValidationIssue[] }>; } /** diff --git a/src/test/mocks/vscodeFs.ts b/src/test/mocks/vscodeFs.ts new file mode 100644 index 0000000000..4341b60f06 --- /dev/null +++ b/src/test/mocks/vscodeFs.ts @@ -0,0 +1,23 @@ +import { anything, instance, mock, when } from 'ts-mockito'; +import { Uri } from 'vscode'; + +import { mockedVSCodeNamespaces } from '../vscode-mock'; + +/** + * Stubs `workspace.fs.readFile` to yield `contents` as UTF-8 bytes, and points the mocked + * `workspace.fs` namespace at it. Pass a function to serve different bytes per URI. + * + * Returns the ts-mockito mock so callers can `verify(mockFs.readFile(anything()))`. + * Call after `resetVSCodeMocks()` — the reset replaces the namespace mocks this stubs. + */ +export function stubReadFile(contents: string | ((uri: Uri) => string)): typeof import('vscode').workspace.fs { + const resolveContents = typeof contents === 'function' ? contents : () => contents; + const mockFs = mock(); + + when(mockFs.readFile(anything())).thenCall((uri: Uri) => + Promise.resolve(new TextEncoder().encode(resolveContents(uri))) + ); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + return mockFs; +} diff --git a/src/webviews/webview-side/integrations/IntegrationItem.tsx b/src/webviews/webview-side/integrations/IntegrationItem.tsx index 4c3a206727..48ed75ec7d 100644 --- a/src/webviews/webview-side/integrations/IntegrationItem.tsx +++ b/src/webviews/webview-side/integrations/IntegrationItem.tsx @@ -1,12 +1,11 @@ import * as React from 'react'; -import { BigQueryAuthMethods } from '@deepnote/database-integrations'; import { getLocString } from '../react-common/locReactSide'; -import { ConfigurableDatabaseIntegrationType, IntegrationWithStatus } from './types'; +import { ConfigurableDatabaseIntegrationType, DetectedIntegration } from './types'; import { integrationTypeIcons } from './integrationUtils'; export interface IIntegrationItemProps { - integration: IntegrationWithStatus; + integration: DetectedIntegration; onConfigure: (integrationId: string) => void; onReset: (integrationId: string) => void; onDelete: (integrationId: string) => void; @@ -63,11 +62,17 @@ export const IntegrationItem: React.FC = ({ onDelete, onAuthenticate }) => { - const statusClass = integration.status === 'connected' ? 'status-connected' : 'status-disconnected'; - const statusText = - integration.status === 'connected' - ? getLocString('integrationsConnected', 'Connected') - : getLocString('integrationsNotConfigured', 'Not Configured'); + // Credentials the panel can edit live in SecretStorage, which is exactly what `config` holds. A + // `.deepnote.env.yaml`-configured integration is read-only here: the panel writes SecretStorage only and the + // file wins the merge, so configuring, resetting or deleting it would be a silent no-op at runtime. It still + // counts as configured — it works — so it gets the connected styling. + const isFileConfigured = Boolean(integration.isFileConfigured); + const statusClass = integration.config || isFileConfigured ? 'status-connected' : 'status-disconnected'; + const statusText = isFileConfigured + ? getLocString('integrationsConfiguredInFile', 'Configured in file') + : integration.config + ? getLocString('integrationsConnected', 'Connected') + : getLocString('integrationsNotConfigured', 'Not Configured'); const configureText = integration.config ? getLocString('integrationsReconfigure', 'Reconfigure') : getLocString('integrationsConfigure', 'Configure'); @@ -82,13 +87,11 @@ export const IntegrationItem: React.FC = ({ const typeLabel = type ? getIntegrationTypeLabel(type) : undefined; const typeIcon = type ? integrationTypeIcons[type] : undefined; - // Federated-auth UI: only for BigQuery + google-oauth; hidden for service-account BigQuery and other types. - const isFederatedOauth = - integration.config?.type === 'big-query' && - integration.config.metadata.authMethod === BigQueryAuthMethods.GoogleOauth; + // Federated-auth UI: `tokenStatus` alone decides. The extension gates on its candidate set, which also + // covers integrations declared in `.deepnote.env.yaml` — those have no `config` here, so re-deriving + // eligibility from `config` would hide the action for exactly the case that needs it. const tokenStatus = integration.tokenStatus; - const showFederatedPill = isFederatedOauth && tokenStatus && tokenStatus !== 'unsupported'; - const showFederatedAuthButton = isFederatedOauth && tokenStatus && tokenStatus !== 'unsupported'; + const showFederatedAuth = tokenStatus && tokenStatus !== 'unsupported'; const tokenStatusText = tokenStatus === 'authenticated' @@ -113,7 +116,7 @@ export const IntegrationItem: React.FC = ({ {typeLabel && {typeLabel}} {typeLabel && • } {statusText} - {showFederatedPill && ( + {showFederatedAuth && ( <> • {tokenStatusText} @@ -122,20 +125,23 @@ export const IntegrationItem: React.FC = ({
- - {showFederatedAuthButton && ( + {!isFileConfigured && ( + + )} + {showFederatedAuth && ( )} - {integration.config && ( + {/* An id can live in both SecretStorage and the file; the file wins at runtime, so hide these there too. */} + {integration.config && !isFileConfigured && ( )} - {integration.config && ( + {integration.config && !isFileConfigured && (