refactor: extract shared integration scanning logic to utility function

jankuca · jankuca · commit 6ce699c71d38 · 2025-10-07T17:42:06.000+02:00
diff --git a/src/notebooks/deepnote/integrations/integrationDetector.ts b/src/notebooks/deepnote/integrations/integrationDetector.ts
@@ -2,8 +2,9 @@ import { inject, injectable } from 'inversify';
 
 import { logger } from '../../../platform/logging';
 import { IDeepnoteNotebookManager } from '../../types';
-import { DATAFRAME_SQL_INTEGRATION_ID, IntegrationStatus, IntegrationWithStatus } from './integrationTypes';
+import { IntegrationStatus, IntegrationWithStatus } from './integrationTypes';
 import { IIntegrationDetector, IIntegrationStorage } from './types';
+import { BlockWithIntegration, scanBlocksForIntegrations } from './integrationUtils';
 
 /**
  * Service for detecting integrations used in Deepnote notebooks
@@ -19,55 +20,31 @@ export class IntegrationDetector implements IIntegrationDetector {
      * Detect all integrations used in the given project
      */
     async detectIntegrations(projectId: string): Promise<Map<string, IntegrationWithStatus>> {
-        const integrations = new Map<string, IntegrationWithStatus>();
-
         // Get the project
         const project = this.notebookManager.getOriginalProject(projectId);
         if (!project) {
             logger.warn(
                 `IntegrationDetector: No project found for ID: ${projectId}. The project may not have been loaded yet.`
             );
-            return integrations;
+            return new Map();
         }
 
         logger.debug(
             `IntegrationDetector: Scanning project ${projectId} with ${project.project.notebooks.length} notebooks`
         );
 
-        // Scan all notebooks in the project
+        // Collect all blocks with SQL integration metadata from all notebooks
+        const blocksWithIntegrations: BlockWithIntegration[] = [];
         for (const notebook of project.project.notebooks) {
             logger.trace(`IntegrationDetector: Scanning notebook ${notebook.id} with ${notebook.blocks.length} blocks`);
 
-            // Scan all blocks in the notebook
             for (const block of notebook.blocks) {
                 // Check if this is a code block with SQL integration metadata
                 if (block.type === 'code' && block.metadata?.sql_integration_id) {
-                    const integrationId = block.metadata.sql_integration_id;
-
-                    // Skip excluded integrations (e.g., internal DuckDB integration)
-                    if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) {
-                        logger.trace(
-                            `IntegrationDetector: Skipping excluded integration: ${integrationId} in block ${block.id}`
-                        );
-                        continue;
-                    }
-
-                    logger.debug(`IntegrationDetector: Found integration: ${integrationId} in block ${block.id}`);
-
-                    // Skip if we've already detected this integration
-                    if (integrations.has(integrationId)) {
-                        continue;
-                    }
-
-                    // Check if the integration is configured
-                    const config = await this.integrationStorage.get(integrationId);
-
-                    const status: IntegrationWithStatus = {
-                        config: config || null,
-                        status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected
-                    };
-
-                    integrations.set(integrationId, status);
+                    blocksWithIntegrations.push({
+                        id: block.id,
+                        sql_integration_id: block.metadata.sql_integration_id
+                    });
                 } else if (block.type === 'code') {
                     logger.trace(
                         `IntegrationDetector: Block ${block.id} has no sql_integration_id. Metadata:`,
@@ -77,8 +54,8 @@ export class IntegrationDetector implements IIntegrationDetector {
             }
         }
 
-        logger.debug(`IntegrationDetector: Found ${integrations.size} integrations`);
-        return integrations;
+        // Use the shared utility to scan blocks and build the status map
+        return scanBlocksForIntegrations(blocksWithIntegrations, this.integrationStorage, 'IntegrationDetector');
     }
 
     /**
diff --git a/src/notebooks/deepnote/integrations/integrationManager.ts b/src/notebooks/deepnote/integrations/integrationManager.ts
@@ -6,7 +6,8 @@ import { Commands } from '../../../platform/common/constants';
 import { logger } from '../../../platform/logging';
 import { IntegrationWebviewProvider } from './integrationWebview';
 import { IIntegrationDetector, IIntegrationStorage } from './types';
-import { DATAFRAME_SQL_INTEGRATION_ID, IntegrationStatus, IntegrationWithStatus } from './integrationTypes';
+import { IntegrationStatus, IntegrationWithStatus } from './integrationTypes';
+import { BlockWithIntegration, scanBlocksForIntegrations } from './integrationUtils';
 
 /**
  * Manages integration UI and commands for Deepnote notebooks
@@ -139,38 +140,22 @@ export class IntegrationManager {
      * This is used when the project isn't stored in the notebook manager
      */
     private async detectIntegrationsFromCells(notebook: NotebookDocument): Promise<Map<string, IntegrationWithStatus>> {
-        const integrations = new Map<string, IntegrationWithStatus>();
+        // Collect all cells with SQL integration metadata
+        const blocksWithIntegrations: BlockWithIntegration[] = [];
 
         for (const cell of notebook.getCells()) {
             const deepnoteMetadata = cell.metadata?.deepnoteMetadata;
             logger.trace(`IntegrationManager: Cell ${cell.index} metadata:`, deepnoteMetadata);
 
             if (deepnoteMetadata?.sql_integration_id) {
-                const integrationId = deepnoteMetadata.sql_integration_id;
-
-                // Skip excluded integrations (e.g., internal DuckDB integration)
-                if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) {
-                    logger.trace(`IntegrationManager: Skipping excluded integration: ${integrationId}`);
-                    continue;
-                }
-
-                // Skip if we've already detected this integration
-                if (integrations.has(integrationId)) {
-                    continue;
-                }
-
-                // Check if the integration is configured
-                const config = await this.integrationStorage.get(integrationId);
-
-                integrations.set(integrationId, {
-                    config: config || null,
-                    status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected
+                blocksWithIntegrations.push({
+                    id: `cell-${cell.index}`,
+                    sql_integration_id: deepnoteMetadata.sql_integration_id
                 });
-
-                logger.debug(`IntegrationManager: Found integration in cell: ${integrationId}`);
             }
         }
 
-        return integrations;
+        // Use the shared utility to scan blocks and build the status map
+        return scanBlocksForIntegrations(blocksWithIntegrations, this.integrationStorage, 'IntegrationManager');
     }
 }
diff --git a/src/notebooks/deepnote/integrations/integrationUtils.ts b/src/notebooks/deepnote/integrations/integrationUtils.ts
@@ -0,0 +1,59 @@
+import { logger } from '../../../platform/logging';
+import { IIntegrationStorage } from './types';
+import { DATAFRAME_SQL_INTEGRATION_ID, IntegrationStatus, IntegrationWithStatus } from './integrationTypes';
+
+/**
+ * Represents a block with SQL integration metadata
+ */
+export interface BlockWithIntegration {
+    id: string;
+    sql_integration_id: string;
+}
+
+/**
+ * Scans blocks for SQL integrations and builds a status map.
+ * This is the core logic shared between IntegrationDetector and IntegrationManager.
+ *
+ * @param blocks - Iterator of blocks to scan (can be from Deepnote project or VSCode notebook cells)
+ * @param integrationStorage - Storage service to check configuration status
+ * @param logContext - Context string for logging (e.g., "IntegrationDetector", "IntegrationManager")
+ * @returns Map of integration IDs to their status
+ */
+export async function scanBlocksForIntegrations(
+    blocks: Iterable<BlockWithIntegration>,
+    integrationStorage: IIntegrationStorage,
+    logContext: string
+): Promise<Map<string, IntegrationWithStatus>> {
+    const integrations = new Map<string, IntegrationWithStatus>();
+
+    for (const block of blocks) {
+        const integrationId = block.sql_integration_id;
+
+        // Skip excluded integrations (e.g., internal DuckDB integration)
+        if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) {
+            logger.trace(`${logContext}: Skipping excluded integration: ${integrationId} in block ${block.id}`);
+            continue;
+        }
+
+        // Skip if we've already detected this integration
+        if (integrations.has(integrationId)) {
+            continue;
+        }
+
+        logger.debug(`${logContext}: Found integration: ${integrationId} in block ${block.id}`);
+
+        // Check if the integration is configured
+        const config = await integrationStorage.get(integrationId);
+
+        const status: IntegrationWithStatus = {
+            config: config || null,
+            status: config ? IntegrationStatus.Connected : IntegrationStatus.Disconnected
+        };
+
+        integrations.set(integrationId, status);
+    }
+
+    logger.debug(`${logContext}: Found ${integrations.size} integrations`);
+    return integrations;
+}
+
