feat: SQL block variable input and integration switching (#86)

Author: jankuca

src/notebooks/deepnote/sqlCellStatusBarProvider.ts

@@ -4,21 +4,35 @@ import {
     NotebookCell,
     NotebookCellStatusBarItem,
     NotebookCellStatusBarItemProvider,
-    NotebookDocument,
+    NotebookDocumentChangeEvent,
+    NotebookEdit,
     ProviderResult,
+    QuickPickItem,
+    QuickPickItemKind,
+    WorkspaceEdit,
+    commands,
     l10n,
-    notebooks
+    notebooks,
+    window,
+    workspace
 } from 'vscode';
 import { inject, injectable } from 'inversify';
 
 import { IExtensionSyncActivationService } from '../../platform/activation/types';
 import { IDisposableRegistry } from '../../platform/common/types';
-import { Commands } from '../../platform/common/constants';
 import { IIntegrationStorage } from './integrations/types';
-import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes';
+import { Commands } from '../../platform/common/constants';
+import { DATAFRAME_SQL_INTEGRATION_ID, IntegrationType } from '../../platform/notebooks/deepnote/integrationTypes';
 
 /**
- * Provides status bar items for SQL cells showing the integration name
+ * QuickPick item with an integration ID
+ */
+interface LocalQuickPickItem extends QuickPickItem {
+    id: string;
+}
+
+/**
+ * Provides status bar items for SQL cells showing the integration name and variable name
  */
 @injectable()
 export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService {
@@ -42,6 +56,55 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid
             })
         );
 
+        // Refresh when any Deepnote notebook changes (e.g., metadata updated externally)
+        this.disposables.push(
+            workspace.onDidChangeNotebookDocument((e: NotebookDocumentChangeEvent) => {
+                if (e.notebook.notebookType === 'deepnote') {
+                    this._onDidChangeCellStatusBarItems.fire();
+                }
+            })
+        );
+
+        // Register command to update SQL variable name
+        this.disposables.push(
+            commands.registerCommand('deepnote.updateSqlVariableName', async (cell?: NotebookCell) => {
+                if (!cell) {
+                    // Fall back to the active notebook cell
+                    const activeEditor = window.activeNotebookEditor;
+                    if (activeEditor && activeEditor.selection) {
+                        cell = activeEditor.notebook.cellAt(activeEditor.selection.start);
+                    }
+                }
+
+                if (!cell) {
+                    void window.showErrorMessage(l10n.t('No active notebook cell'));
+                    return;
+                }
+
+                await this.updateVariableName(cell);
+            })
+        );
+
+        // Register command to switch SQL integration
+        this.disposables.push(
+            commands.registerCommand('deepnote.switchSqlIntegration', async (cell?: NotebookCell) => {
+                if (!cell) {
+                    // Fall back to the active notebook cell
+                    const activeEditor = window.activeNotebookEditor;
+                    if (activeEditor && activeEditor.selection) {
+                        cell = activeEditor.notebook.cellAt(activeEditor.selection.start);
+                    }
+                }
+
+                if (!cell) {
+                    void window.showErrorMessage(l10n.t('No active notebook cell'));
+                    return;
+                }
+
+                await this.switchIntegration(cell);
+            })
+        );
+
         // Dispose our emitter with the extension
         this.disposables.push(this._onDidChangeCellStatusBarItems);
     }
@@ -59,18 +122,7 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid
             return undefined;
         }
 
-        // Get the integration ID from cell metadata
-        const integrationId = this.getIntegrationId(cell);
-        if (!integrationId) {
-            return undefined;
-        }
-
-        // Don't show status bar for the internal DuckDB integration
-        if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) {
-            return undefined;
-        }
-
-        return this.createStatusBarItem(cell.notebook, integrationId);
+        return this.createStatusBarItems(cell);
     }
 
     private getIntegrationId(cell: NotebookCell): string | undefined {
@@ -86,11 +138,57 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid
         return undefined;
     }
 
-    private async createStatusBarItem(
-        notebook: NotebookDocument,
+    private async createStatusBarItems(cell: NotebookCell): Promise<NotebookCellStatusBarItem[]> {
+        const items: NotebookCellStatusBarItem[] = [];
+
+        // Add integration status bar item
+        const integrationId = this.getIntegrationId(cell);
+        if (integrationId) {
+            const integrationItem = await this.createIntegrationStatusBarItem(cell, integrationId);
+            if (integrationItem) {
+                items.push(integrationItem);
+            }
+        } else {
+            // Show "No integration connected" when no integration is selected
+            items.push({
+                text: `$(database) ${l10n.t('No integration connected')}`,
+                alignment: 1, // NotebookCellStatusBarAlignment.Left
+                priority: 100,
+                tooltip: l10n.t('No SQL integration connected\nClick to select an integration'),
+                command: {
+                    title: l10n.t('Switch Integration'),
+                    command: 'deepnote.switchSqlIntegration',
+                    arguments: [cell]
+                }
+            });
+        }
+
+        // Always add variable status bar item for SQL cells
+        items.push(this.createVariableStatusBarItem(cell));
+
+        return items;
+    }
+
+    private async createIntegrationStatusBarItem(
+        cell: NotebookCell,
         integrationId: string
     ): Promise<NotebookCellStatusBarItem | undefined> {
-        const projectId = notebook.metadata?.deepnoteProjectId;
+        // Handle internal DuckDB integration specially
+        if (integrationId === DATAFRAME_SQL_INTEGRATION_ID) {
+            return {
+                text: `$(database) ${l10n.t('DataFrame SQL (DuckDB)')}`,
+                alignment: 1, // NotebookCellStatusBarAlignment.Left
+                priority: 100,
+                tooltip: l10n.t('Internal DuckDB integration for querying DataFrames\nClick to switch'),
+                command: {
+                    title: l10n.t('Switch Integration'),
+                    command: 'deepnote.switchSqlIntegration',
+                    arguments: [cell]
+                }
+            };
+        }
+
+        const projectId = cell.notebook.metadata?.deepnoteProjectId;
         if (!projectId) {
             return undefined;
         }
@@ -99,16 +197,210 @@ export class SqlCellStatusBarProvider implements NotebookCellStatusBarItemProvid
         const config = await this.integrationStorage.getProjectIntegrationConfig(projectId, integrationId);
         const displayName = config?.name || l10n.t('Unknown integration (configure)');
 
-        // Create a status bar item that opens the integration management UI
+        // Create a status bar item that opens the integration picker
         return {
             text: `$(database) ${displayName}`,
             alignment: 1, // NotebookCellStatusBarAlignment.Left
-            tooltip: l10n.t('SQL Integration: {0}\nClick to configure', displayName),
+            priority: 100,
+            tooltip: l10n.t('SQL Integration: {0}\nClick to switch or configure', displayName),
+            command: {
+                title: l10n.t('Switch Integration'),
+                command: 'deepnote.switchSqlIntegration',
+                arguments: [cell]
+            }
+        };
+    }
+
+    private createVariableStatusBarItem(cell: NotebookCell): NotebookCellStatusBarItem {
+        const variableName = this.getVariableName(cell);
+
+        return {
+            text: l10n.t('Variable: {0}', variableName),
+            alignment: 1, // NotebookCellStatusBarAlignment.Left
+            priority: 90,
+            tooltip: l10n.t('Variable name for SQL query result\nClick to change'),
             command: {
-                title: l10n.t('Configure Integration'),
-                command: Commands.ManageIntegrations,
-                arguments: [integrationId]
+                title: l10n.t('Change Variable Name'),
+                command: 'deepnote.updateSqlVariableName',
+                arguments: [cell]
+            }
+        };
+    }
+
+    private getVariableName(cell: NotebookCell): string {
+        const metadata = cell.metadata;
+        if (metadata && typeof metadata === 'object') {
+            const variableName = (metadata as Record<string, unknown>).deepnote_variable_name;
+            if (typeof variableName === 'string' && variableName) {
+                return variableName;
             }
+        }
+
+        return 'df';
+    }
+
+    private async updateVariableName(cell: NotebookCell): Promise<void> {
+        const currentVariableName = this.getVariableName(cell);
+
+        const newVariableNameInput = await window.showInputBox({
+            prompt: l10n.t('Enter variable name for SQL query result'),
+            value: currentVariableName,
+            ignoreFocusOut: true,
+            validateInput: (value) => {
+                const trimmed = value.trim();
+                if (!trimmed) {
+                    return l10n.t('Variable name cannot be empty');
+                }
+                if (!/^[a-zA-Z_][a-zA-Z0-9_]*$/.test(trimmed)) {
+                    return l10n.t('Variable name must be a valid Python identifier');
+                }
+                return undefined;
+            }
+        });
+
+        const newVariableName = newVariableNameInput?.trim();
+        if (newVariableName === undefined || newVariableName === currentVariableName) {
+            return;
+        }
+
+        // Update cell metadata
+        const edit = new WorkspaceEdit();
+        const updatedMetadata = {
+            ...cell.metadata,
+            deepnote_variable_name: newVariableName
         };
+
+        edit.set(cell.notebook.uri, [NotebookEdit.updateCellMetadata(cell.index, updatedMetadata)]);
+
+        const success = await workspace.applyEdit(edit);
+        if (!success) {
+            void window.showErrorMessage(l10n.t('Failed to update variable name'));
+            return;
+        }
+
+        // Trigger status bar update
+        this._onDidChangeCellStatusBarItems.fire();
+    }
+
+    private async switchIntegration(cell: NotebookCell): Promise<void> {
+        const currentIntegrationId = this.getIntegrationId(cell);
+
+        // Get all available integrations
+        const allIntegrations = await this.integrationStorage.getAll();
+
+        // Build quick pick items
+        const items: (QuickPickItem | LocalQuickPickItem)[] = [];
+
+        // Check if current integration is unknown (not in the list)
+        const isCurrentIntegrationUnknown =
+            currentIntegrationId &&
+            currentIntegrationId !== DATAFRAME_SQL_INTEGRATION_ID &&
+            !allIntegrations.some((i) => i.id === currentIntegrationId);
+
+        // Add current unknown integration first if it exists
+        if (isCurrentIntegrationUnknown && currentIntegrationId) {
+            const item: LocalQuickPickItem = {
+                label: l10n.t('Unknown integration (configure)'),
+                description: currentIntegrationId,
+                detail: l10n.t('Currently selected'),
+                id: currentIntegrationId
+            };
+            items.push(item);
+        }
+
+        // Add all configured integrations
+        for (const integration of allIntegrations) {
+            const typeLabel = this.getIntegrationTypeLabel(integration.type);
+            const item: LocalQuickPickItem = {
+                label: integration.name || integration.id,
+                description: typeLabel,
+                detail: integration.id === currentIntegrationId ? l10n.t('Currently selected') : undefined,
+                // Store the integration ID in a custom property
+                id: integration.id
+            };
+            items.push(item);
+        }
+
+        // Add DuckDB integration
+        const duckDbItem: LocalQuickPickItem = {
+            label: l10n.t('DataFrame SQL (DuckDB)'),
+            description: l10n.t('DuckDB'),
+            detail: currentIntegrationId === DATAFRAME_SQL_INTEGRATION_ID ? l10n.t('Currently selected') : undefined,
+            id: DATAFRAME_SQL_INTEGRATION_ID
+        };
+        items.push(duckDbItem);
+
+        // Add "Configure current integration" option (with separator)
+        if (currentIntegrationId && currentIntegrationId !== DATAFRAME_SQL_INTEGRATION_ID) {
+            // Add separator
+            const separator: QuickPickItem = {
+                label: '',
+                kind: QuickPickItemKind.Separator
+            };
+            items.push(separator);
+
+            const configureItem: LocalQuickPickItem = {
+                label: l10n.t('Configure current integration'),
+                id: '__configure__'
+            };
+            items.push(configureItem);
+        }
+
+        const selected = await window.showQuickPick(items, {
+            placeHolder: l10n.t('Select SQL integration'),
+            matchOnDescription: true
+        });
+
+        if (!selected) {
+            return;
+        }
+
+        // Type guard to check if selected item has an id property
+        if (!('id' in selected)) {
+            return;
+        }
+
+        const selectedItem = selected as LocalQuickPickItem;
+        const selectedId = selectedItem.id;
+
+        // Handle "Configure current integration" option
+        if (selectedId === '__configure__' && currentIntegrationId) {
+            await commands.executeCommand(Commands.ManageIntegrations, currentIntegrationId);
+            return;
+        }
+
+        // No change
+        if (selectedId === currentIntegrationId) {
+            return;
+        }
+
+        // Update cell metadata with new integration ID
+        const edit = new WorkspaceEdit();
+        const updatedMetadata = {
+            ...cell.metadata,
+            sql_integration_id: selectedId
+        };
+
+        edit.set(cell.notebook.uri, [NotebookEdit.updateCellMetadata(cell.index, updatedMetadata)]);
+
+        const success = await workspace.applyEdit(edit);
+        if (!success) {
+            void window.showErrorMessage(l10n.t('Failed to select integration'));
+            return;
+        }
+
+        // Trigger status bar update
+        this._onDidChangeCellStatusBarItems.fire();
+    }
+
+    private getIntegrationTypeLabel(type: IntegrationType): string {
+        switch (type) {
+            case IntegrationType.Postgres:
+                return l10n.t('PostgreSQL');
+            case IntegrationType.BigQuery:
+                return l10n.t('BigQuery');
+            default:
+                return String(type);
+        }
     }
 }