introduce narrower ConfigurableDatabaseIntegration* types

jankuca · jankuca · commit 5d0b560f6f0f · 2025-11-02T15:32:36.000+01:00
diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts
@@ -2,6 +2,7 @@ import * as assert from 'assert';
 
 import { DeepnoteNotebookManager } from './deepnoteNotebookManager';
 import type { DeepnoteProject } from '../../platform/deepnote/deepnoteTypes';
+import { ProjectIntegration } from '../types';
 
 suite('DeepnoteNotebookManager', () => {
     let manager: DeepnoteNotebookManager;
@@ -187,7 +188,7 @@ suite('DeepnoteNotebookManager', () => {
         test('should update integrations list for existing project and return true', () => {
             manager.storeOriginalProject('project-123', mockProject, 'notebook-456');
 
-            const integrations = [
+            const integrations: ProjectIntegration[] = [
                 { id: 'int-1', name: 'PostgreSQL', type: 'pgsql' },
                 { id: 'int-2', name: 'BigQuery', type: 'big-query' }
             ];
@@ -211,7 +212,7 @@ suite('DeepnoteNotebookManager', () => {
 
             manager.storeOriginalProject('project-123', projectWithIntegrations, 'notebook-456');
 
-            const newIntegrations = [
+            const newIntegrations: ProjectIntegration[] = [
                 { id: 'new-int-1', name: 'New Integration 1', type: 'pgsql' },
                 { id: 'new-int-2', name: 'New Integration 2', type: 'big-query' }
             ];
@@ -257,7 +258,7 @@ suite('DeepnoteNotebookManager', () => {
         test('should preserve other project properties and return true', () => {
             manager.storeOriginalProject('project-123', mockProject, 'notebook-456');
 
-            const integrations = [{ id: 'int-1', name: 'PostgreSQL', type: 'pgsql' }];
+            const integrations: ProjectIntegration[] = [{ id: 'int-1', name: 'PostgreSQL', type: 'pgsql' }];
 
             const result = manager.updateProjectIntegrations('project-123', integrations);
 
@@ -275,7 +276,7 @@ suite('DeepnoteNotebookManager', () => {
             manager.storeOriginalProject('project-123', mockProject, 'notebook-456');
             manager.updateCurrentNotebookId('project-123', undefined as any);
 
-            const integrations = [
+            const integrations: ProjectIntegration[] = [
                 { id: 'int-1', name: 'PostgreSQL', type: 'pgsql' },
                 { id: 'int-2', name: 'BigQuery', type: 'big-query' }
             ];
diff --git a/src/notebooks/deepnote/integrations/integrationDetector.ts b/src/notebooks/deepnote/integrations/integrationDetector.ts
@@ -2,9 +2,13 @@ import { inject, injectable } from 'inversify';
 
 import { logger } from '../../../platform/logging';
 import { IDeepnoteNotebookManager } from '../../types';
-import { IntegrationStatus, IntegrationWithStatus } from '../../../platform/notebooks/deepnote/integrationTypes';
+import {
+    ConfigurableDatabaseIntegrationType,
+    IntegrationStatus,
+    IntegrationWithStatus
+} from '../../../platform/notebooks/deepnote/integrationTypes';
 import { IIntegrationDetector, IIntegrationStorage } from './types';
-import { DatabaseIntegrationType, databaseIntegrationTypes } from '@deepnote/database-integrations';
+import { databaseIntegrationTypes } from '@deepnote/database-integrations';
 
 /**
  * Service for detecting integrations used in Deepnote notebooks
@@ -41,7 +45,10 @@ export class IntegrationDetector implements IIntegrationDetector {
         for (const projectIntegration of projectIntegrations) {
             const integrationId = projectIntegration.id;
             const integrationType = projectIntegration.type;
-            if (!(databaseIntegrationTypes as readonly string[]).includes(integrationType)) {
+            if (
+                !(databaseIntegrationTypes as readonly string[]).includes(integrationType) ||
+                integrationType === 'pandas-dataframe'
+            ) {
                 logger.debug(`IntegrationDetector: Skipping unsupported integration type: ${integrationType}`);
                 continue;
             }
@@ -53,7 +60,7 @@ export class IntegrationDetector implements IIntegrationDetector {
                 status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected,
                 // Include integration metadata from project for prefilling when config is null
                 integrationName: projectIntegration.name,
-                integrationType: integrationType as DatabaseIntegrationType
+                integrationType: integrationType as ConfigurableDatabaseIntegrationType
             };
 
             integrations.set(integrationId, status);
diff --git a/src/notebooks/deepnote/integrations/integrationManager.ts b/src/notebooks/deepnote/integrations/integrationManager.ts
@@ -161,6 +161,11 @@ export class IntegrationManager implements IIntegrationManager {
                 integrationType = projectIntegration.type as DatabaseIntegrationType;
             }
 
+            if (integrationType === 'pandas-dataframe') {
+                logger.debug(`IntegrationManager: Skipping internal DuckDB integration ${selectedIntegrationId}`);
+                return;
+            }
+
             integrations.set(selectedIntegrationId, {
                 config: config || null,
                 status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected,
diff --git a/src/notebooks/deepnote/integrations/integrationWebview.ts b/src/notebooks/deepnote/integrations/integrationWebview.ts
@@ -7,8 +7,11 @@ import { logger } from '../../../platform/logging';
 import { LocalizedMessages, SharedMessages } from '../../../messageTypes';
 import { IDeepnoteNotebookManager, ProjectIntegration } from '../../types';
 import { IIntegrationStorage, IIntegrationWebviewProvider } from './types';
-import { IntegrationStatus, IntegrationWithStatus } from '../../../platform/notebooks/deepnote/integrationTypes';
-import { DatabaseIntegrationConfig } from '@deepnote/database-integrations';
+import {
+    ConfigurableDatabaseIntegrationConfig,
+    IntegrationStatus,
+    IntegrationWithStatus
+} from '../../../platform/notebooks/deepnote/integrationTypes';
 
 /**
  * Manages the webview panel for integration configuration
@@ -217,7 +220,7 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider {
     private async handleMessage(message: {
         type: string;
         integrationId?: string;
-        config?: DatabaseIntegrationConfig;
+        config?: ConfigurableDatabaseIntegrationConfig;
     }): Promise<void> {
         switch (message.type) {
             case 'configure':
@@ -259,7 +262,10 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider {
     /**
      * Save the configuration for an integration
      */
-    private async saveConfiguration(integrationId: string, config: DatabaseIntegrationConfig): Promise<void> {
+    private async saveConfiguration(
+        integrationId: string,
+        config: ConfigurableDatabaseIntegrationConfig
+    ): Promise<void> {
         try {
             await this.integrationStorage.save(config);
 
@@ -345,12 +351,6 @@ export class IntegrationWebviewProvider implements IIntegrationWebviewProvider {
                     return null;
                 }
 
-                // Skip DuckDB integration (internal, not a real Deepnote integration)
-                if (type === 'pandas-dataframe') {
-                    logger.trace(`IntegrationWebviewProvider: Skipping internal DuckDB integration ${id}`);
-                    return null;
-                }
-
                 return {
                     id,
                     name: integration.config?.name || integration.integrationName || id,
diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts
@@ -22,7 +22,10 @@ import { IExtensionSyncActivationService } from '../../platform/activation/types
 import { IDisposableRegistry } from '../../platform/common/types';
 import { IIntegrationStorage } from './integrations/types';
 import { Commands } from '../../platform/common/constants';
-import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes';
+import {
+    ConfigurableDatabaseIntegrationType,
+    DATAFRAME_SQL_INTEGRATION_ID
+} from '../../platform/notebooks/deepnote/integrationTypes';
 import { IDeepnoteNotebookManager } from '../types';
 import { DatabaseIntegrationType, databaseIntegrationTypes } from '@deepnote/database-integrations';
 
@@ -338,15 +341,20 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid
 
         // Add all project integrations
         for (const projectIntegration of projectIntegrations) {
+            const integrationType =
+                projectIntegration.type &&
+                (databaseIntegrationTypes as readonly string[]).includes(projectIntegration.type)
+                    ? (projectIntegration.type as DatabaseIntegrationType)
+                    : null;
+
             // Skip the internal DuckDB integration
-            if (projectIntegration.id === DATAFRAME_SQL_INTEGRATION_ID) {
+            if (projectIntegration.id === DATAFRAME_SQL_INTEGRATION_ID || integrationType === 'pandas-dataframe') {
                 continue;
             }
 
-            const integrationType = projectIntegration.type;
             const typeLabel =
                 integrationType && (databaseIntegrationTypes as readonly string[]).includes(integrationType)
-                    ? this.getIntegrationTypeLabel(integrationType as DatabaseIntegrationType)
+                    ? this.getIntegrationTypeLabel(integrationType)
                     : projectIntegration.type;
 
             const item: LocalQuickPickItem = {
@@ -430,7 +438,7 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid
         this._onDidChangeCellStatusBarItems.fire();
     }
 
-    private getIntegrationTypeLabel(type: DatabaseIntegrationType): string {
+    private getIntegrationTypeLabel(type: ConfigurableDatabaseIntegrationType): string {
         switch (type) {
             case 'pgsql':
                 return l10n.t('PostgreSQL');
diff --git a/src/notebooks/types.ts b/src/notebooks/types.ts
@@ -5,6 +5,7 @@ import { NotebookDocument, NotebookEditor, Uri, type Event } from 'vscode';
 import { Resource } from '../platform/common/types';
 import type { EnvironmentPath } from '@vscode/python-extension';
 import { DeepnoteProject } from '../platform/deepnote/deepnoteTypes';
+import { ConfigurableDatabaseIntegrationType } from '../platform/notebooks/deepnote/integrationTypes';
 
 export interface IEmbedNotebookEditorProvider {
     findNotebookEditor(resource: Resource): NotebookEditor | undefined;
@@ -31,7 +32,7 @@ export interface INotebookPythonEnvironmentService {
 export interface ProjectIntegration {
     id: string;
     name: string;
-    type: string;
+    type: ConfigurableDatabaseIntegrationType;
 }
 
 export const IDeepnoteNotebookManager = Symbol('IDeepnoteNotebookManager');
diff --git a/src/platform/notebooks/deepnote/integrationStorage.ts b/src/platform/notebooks/deepnote/integrationStorage.ts
@@ -6,17 +6,13 @@ import { IAsyncDisposableRegistry } from '../../common/types';
 import { logger } from '../../logging';
 import { IIntegrationStorage } from './types';
 import { upgradeLegacyIntegrationConfig } from './legacyIntegrationConfigUtils';
-import {
-    DatabaseIntegrationConfig,
-    databaseIntegrationTypes,
-    databaseMetadataSchemasByType
-} from '@deepnote/database-integrations';
-import { DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes';
+import { databaseIntegrationTypes, databaseMetadataSchemasByType } from '@deepnote/database-integrations';
+import { ConfigurableDatabaseIntegrationConfig, DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes';
 
 const INTEGRATION_SERVICE_NAME = 'deepnote-integrations';
 
 // NOTE: We need a way to upgrade existing configurations to the new format of deepnote/database-integrations.
-type VersionedDatabaseIntegrationConfig = DatabaseIntegrationConfig & { version: 1 };
+type VersionedDatabaseIntegrationConfig = ConfigurableDatabaseIntegrationConfig & { version: 1 };
 
 function storeEncryptedIntegrationConfig(
     encryptedStorage: IEncryptedStorage,
@@ -33,7 +29,7 @@ function storeEncryptedIntegrationConfig(
  */
 @injectable()
 export class IntegrationStorage implements IIntegrationStorage {
-    private readonly cache: Map<string, DatabaseIntegrationConfig> = new Map();
+    private readonly cache: Map<string, ConfigurableDatabaseIntegrationConfig> = new Map();
 
     private cacheLoaded = false;
 
@@ -52,15 +48,15 @@ export class IntegrationStorage implements IIntegrationStorage {
     /**
      * Get all stored integration configurations.
      */
-    async getAll(): Promise<DatabaseIntegrationConfig[]> {
+    async getAll(): Promise<ConfigurableDatabaseIntegrationConfig[]> {
         await this.ensureCacheLoaded();
         return Array.from(this.cache.values());
     }
 
     /**
      * Get a specific integration configuration by ID
      */
-    async getIntegrationConfig(integrationId: string): Promise<DatabaseIntegrationConfig | undefined> {
+    async getIntegrationConfig(integrationId: string): Promise<ConfigurableDatabaseIntegrationConfig | undefined> {
         await this.ensureCacheLoaded();
         return this.cache.get(integrationId);
     }
@@ -73,15 +69,15 @@ export class IntegrationStorage implements IIntegrationStorage {
     async getProjectIntegrationConfig(
         _projectId: string,
         integrationId: string
-    ): Promise<DatabaseIntegrationConfig | undefined> {
+    ): Promise<ConfigurableDatabaseIntegrationConfig | undefined> {
         return this.getIntegrationConfig(integrationId);
     }
 
     /**
      * Save or update an integration configuration
      */
-    async save(config: DatabaseIntegrationConfig): Promise<void> {
-        if (config.type === 'pandas-dataframe' || config.id === DATAFRAME_SQL_INTEGRATION_ID) {
+    async save(config: ConfigurableDatabaseIntegrationConfig): Promise<void> {
+        if ((config.type as string) === 'pandas-dataframe' || config.id === DATAFRAME_SQL_INTEGRATION_ID) {
             logger.warn(`IntegrationStorage: Skipping save for internal DuckDB integration ${config.id}`);
             return;
         }
@@ -210,7 +206,7 @@ export class IntegrationStorage implements IIntegrationStorage {
                             const config =
                                 databaseIntegrationTypes.includes(rawConfig.type) &&
                                 rawConfig.type !== 'pandas-dataframe'
-                                    ? (rawConfig as DatabaseIntegrationConfig)
+                                    ? (rawConfig as ConfigurableDatabaseIntegrationConfig)
                                     : null;
                             const validMetadata = config
                                 ? databaseMetadataSchemasByType[config.type].safeParse(config.metadata).data
@@ -219,7 +215,7 @@ export class IntegrationStorage implements IIntegrationStorage {
                                 this.cache.set(
                                     id,
                                     // NOTE: We must cast here because there is no union-wide schema parser at the moment.
-                                    { ...config, metadata: validMetadata } as DatabaseIntegrationConfig
+                                    { ...config, metadata: validMetadata } as ConfigurableDatabaseIntegrationConfig
                                 );
                             } else {
                                 logger.warn(`Invalid integration config for ${id}, marking for deletion`);
diff --git a/src/platform/notebooks/deepnote/integrationStorage.unit.test.ts b/src/platform/notebooks/deepnote/integrationStorage.unit.test.ts
@@ -197,7 +197,7 @@ suite('IntegrationStorage', () => {
                 metadata: {}
             };
 
-            await storage.save(config);
+            await storage.save(config as any);
             const result = await storage.getIntegrationConfig('dataframe-1');
 
             assert.strictEqual(result, undefined);
@@ -211,7 +211,7 @@ suite('IntegrationStorage', () => {
                 metadata: {}
             };
 
-            await storage.save(config);
+            await storage.save(config as any);
             const result = await storage.getIntegrationConfig(DATAFRAME_SQL_INTEGRATION_ID);
 
             assert.strictEqual(result, undefined);
diff --git a/src/platform/notebooks/deepnote/integrationTypes.ts b/src/platform/notebooks/deepnote/integrationTypes.ts
@@ -136,6 +136,13 @@ export type LegacyIntegrationConfig =
     | LegacySnowflakeIntegrationConfig
     | LegacyDuckDBIntegrationConfig;
 
+export type ConfigurableDatabaseIntegrationConfig = Extract<
+    DatabaseIntegrationConfig,
+    { type: ConfigurableDatabaseIntegrationType }
+>;
+
+export type ConfigurableDatabaseIntegrationType = Exclude<DatabaseIntegrationType, 'pandas-dataframe'>;
+
 /**
  * Integration connection status
  */
@@ -149,7 +156,7 @@ export enum IntegrationStatus {
  * Integration with its current status
  */
 export interface IntegrationWithStatus {
-    config: DatabaseIntegrationConfig | null;
+    config: ConfigurableDatabaseIntegrationConfig | null;
     status: IntegrationStatus;
     error?: string;
     /**
@@ -159,5 +166,5 @@ export interface IntegrationWithStatus {
     /**
      * Type from the project's integrations list (used for prefilling when config is null)
      */
-    integrationType?: DatabaseIntegrationType;
+    integrationType?: ConfigurableDatabaseIntegrationType;
 }
diff --git a/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts b/src/platform/notebooks/deepnote/sqlIntegrationEnvironmentVariablesProvider.ts
@@ -11,7 +11,7 @@ import {
     IPlatformDeepnoteNotebookManager
 } from './types';
 import { DATAFRAME_SQL_INTEGRATION_ID } from './integrationTypes';
-import { getEnvironmentVariablesForIntegrations } from '@deepnote/database-integrations';
+import { DatabaseIntegrationConfig, getEnvironmentVariablesForIntegrations } from '@deepnote/database-integrations';
 
 /**
  * Provides environment variables for SQL integrations.
@@ -88,7 +88,7 @@ export class SqlIntegrationEnvironmentVariablesProvider implements ISqlIntegrati
             `SqlIntegrationEnvironmentVariablesProvider: Found ${projectIntegrations.length} integrations in project`
         );
 
-        const projectIntegrationConfigs = (
+        const projectIntegrationConfigs: Array<DatabaseIntegrationConfig> = (
             await Promise.all(
                 projectIntegrations.map((integration) => {
                     return this.integrationStorage.getIntegrationConfig(integration.id);
diff --git a/src/platform/notebooks/deepnote/types.ts b/src/platform/notebooks/deepnote/types.ts
@@ -2,7 +2,7 @@ import { CancellationToken, Event, NotebookDocument, Uri } from 'vscode';
 import { IDisposable, Resource } from '../../common/types';
 import { EnvironmentVariables } from '../../common/variables/types';
 import { DeepnoteProject } from '../../deepnote/deepnoteTypes';
-import { DatabaseIntegrationConfig } from '@deepnote/database-integrations';
+import { ConfigurableDatabaseIntegrationConfig } from './integrationTypes';
 
 /**
  * Settings for select input blocks
@@ -31,7 +31,7 @@ export interface IIntegrationStorage extends IDisposable {
      */
     readonly onDidChangeIntegrations: Event<void>;
 
-    getAll(): Promise<DatabaseIntegrationConfig[]>;
+    getAll(): Promise<ConfigurableDatabaseIntegrationConfig[]>;
 
     /**
      * Retrieves the global (non-project-scoped) integration configuration by integration ID.
@@ -49,17 +49,17 @@ export interface IIntegrationStorage extends IDisposable {
      *          - The `IntegrationConfig` object if a global configuration exists for the given ID
      *          - `undefined` if no global configuration exists for the given integration ID
      */
-    getIntegrationConfig(integrationId: string): Promise<DatabaseIntegrationConfig | undefined>;
+    getIntegrationConfig(integrationId: string): Promise<ConfigurableDatabaseIntegrationConfig | undefined>;
 
     /**
      * Get integration configuration for a specific project and integration
      */
     getProjectIntegrationConfig(
         projectId: string,
         integrationId: string
-    ): Promise<DatabaseIntegrationConfig | undefined>;
+    ): Promise<ConfigurableDatabaseIntegrationConfig | undefined>;
 
-    save(config: DatabaseIntegrationConfig): Promise<void>;
+    save(config: ConfigurableDatabaseIntegrationConfig): Promise<void>;
     delete(integrationId: string): Promise<void>;
     exists(integrationId: string): Promise<boolean>;
     clear(): Promise<void>;
diff --git a/src/webviews/webview-side/integrations/ConfigurationForm.tsx b/src/webviews/webview-side/integrations/ConfigurationForm.tsx
diff --git a/src/webviews/webview-side/integrations/IntegrationItem.tsx b/src/webviews/webview-side/integrations/IntegrationItem.tsx
diff --git a/src/webviews/webview-side/integrations/IntegrationPanel.tsx b/src/webviews/webview-side/integrations/IntegrationPanel.tsx
diff --git a/src/webviews/webview-side/integrations/types.ts b/src/webviews/webview-side/integrations/types.ts
