diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.ts index c6439f2d8e..4637b18cbb 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.ts +++ b/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 { + 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 { - 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).deepnote_variable_name; + if (typeof variableName === 'string' && variableName) { + return variableName; } + } + + return 'df'; + } + + private async updateVariableName(cell: NotebookCell): Promise { + 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 { + 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); + } } } diff --git a/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts index ac7cf16386..209ecb3c36 100644 --- a/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/sqlCellStatusBarProvider.unit.test.ts @@ -1,8 +1,9 @@ import { assert } from 'chai'; -import { anything, instance, mock, when } from 'ts-mockito'; +import { anything, instance, mock, verify, when } from 'ts-mockito'; import { CancellationToken, CancellationTokenSource, + EventEmitter, NotebookCell, NotebookCellKind, NotebookDocument, @@ -14,6 +15,9 @@ import { IDisposableRegistry } from '../../platform/common/types'; import { IIntegrationStorage } from './integrations/types'; import { SqlCellStatusBarProvider } from './sqlCellStatusBarProvider'; import { DATAFRAME_SQL_INTEGRATION_ID, IntegrationType } from '../../platform/notebooks/deepnote/integrationTypes'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; +import { createEventHandler } from '../../test/common'; +import { Commands } from '../../platform/common/constants'; suite('SqlCellStatusBarProvider', () => { let provider: SqlCellStatusBarProvider; @@ -38,25 +42,71 @@ suite('SqlCellStatusBarProvider', () => { assert.isUndefined(result); }); - test('returns undefined for SQL cells without integration ID', async () => { + test('returns status bar items for SQL cells without integration ID', async () => { const cell = createMockCell('sql', {}); const result = await provider.provideCellStatusBarItems(cell, cancellationToken); - assert.isUndefined(result); + assert.isDefined(result); + assert.isArray(result); + const items = result as any[]; + assert.strictEqual(items.length, 2); + + // Check "No integration connected" status bar item + const integrationItem = items[0]; + assert.strictEqual(integrationItem.text, '$(database) No integration connected'); + assert.strictEqual(integrationItem.alignment, 1); + assert.isDefined(integrationItem.command); + assert.strictEqual(integrationItem.command.command, 'deepnote.switchSqlIntegration'); + assert.deepStrictEqual(integrationItem.command.arguments, [cell]); + assert.strictEqual(integrationItem.priority, 100); + + // Check variable status bar item + const variableItem = items[1]; + assert.strictEqual(variableItem.text, 'Variable: df'); + assert.strictEqual(variableItem.alignment, 1); + assert.isDefined(variableItem.command); + assert.strictEqual(variableItem.command.command, 'deepnote.updateSqlVariableName'); + assert.deepStrictEqual(variableItem.command.arguments, [cell]); + assert.strictEqual(variableItem.priority, 90); }); - test('returns undefined for SQL cells with dataframe integration ID', async () => { + test('returns status bar items for SQL cells with dataframe integration ID', async () => { const cell = createMockCell('sql', { sql_integration_id: DATAFRAME_SQL_INTEGRATION_ID }); const result = await provider.provideCellStatusBarItems(cell, cancellationToken); - assert.isUndefined(result); + assert.isDefined(result); + assert.isArray(result); + const items = result as any[]; + assert.strictEqual(items.length, 2); + + // Check integration status bar item + const integrationItem = items[0]; + assert.strictEqual(integrationItem.text, '$(database) DataFrame SQL (DuckDB)'); + assert.strictEqual(integrationItem.alignment, 1); + assert.strictEqual( + integrationItem.tooltip, + 'Internal DuckDB integration for querying DataFrames\nClick to switch' + ); + assert.isDefined(integrationItem.command); + assert.strictEqual(integrationItem.command.command, 'deepnote.switchSqlIntegration'); + assert.deepStrictEqual(integrationItem.command.arguments, [cell]); + assert.strictEqual(integrationItem.priority, 100); + + // Check variable status bar item + const variableItem = items[1]; + assert.strictEqual(variableItem.text, 'Variable: df'); + assert.strictEqual(variableItem.alignment, 1); + assert.isDefined(variableItem.command); + assert.strictEqual(variableItem.command.command, 'deepnote.updateSqlVariableName'); + assert.deepStrictEqual(variableItem.command.arguments, [cell]); + assert.strictEqual(variableItem.priority, 90); }); - test('returns status bar item for SQL cell with integration ID', async () => { + test('returns status bar items for SQL cell with integration ID', async () => { const integrationId = 'postgres-123'; const cell = createMockCell( 'sql', @@ -82,11 +132,27 @@ suite('SqlCellStatusBarProvider', () => { const result = await provider.provideCellStatusBarItems(cell, cancellationToken); assert.isDefined(result); - assert.strictEqual((result as any).text, '$(database) My Postgres DB'); - assert.strictEqual((result as any).alignment, 1); // NotebookCellStatusBarAlignment.Left - assert.isDefined((result as any).command); - assert.strictEqual((result as any).command.command, 'deepnote.manageIntegrations'); - assert.deepStrictEqual((result as any).command.arguments, [integrationId]); + assert.isArray(result); + const items = result as any[]; + assert.strictEqual(items.length, 2); + + // Check integration status bar item + const integrationItem = items[0]; + assert.strictEqual(integrationItem.text, '$(database) My Postgres DB'); + assert.strictEqual(integrationItem.alignment, 1); + assert.isDefined(integrationItem.command); + assert.strictEqual(integrationItem.command.command, 'deepnote.switchSqlIntegration'); + assert.deepStrictEqual(integrationItem.command.arguments, [cell]); + assert.strictEqual(integrationItem.priority, 100); + + // Check variable status bar item + const variableItem = items[1]; + assert.strictEqual(variableItem.text, 'Variable: df'); + assert.strictEqual(variableItem.alignment, 1); + assert.isDefined(variableItem.command); + assert.strictEqual(variableItem.command.command, 'deepnote.updateSqlVariableName'); + assert.deepStrictEqual(variableItem.command.arguments, [cell]); + assert.strictEqual(variableItem.priority, 90); }); test('shows "Unknown integration (configure)" when config not found', async () => { @@ -106,10 +172,20 @@ suite('SqlCellStatusBarProvider', () => { const result = await provider.provideCellStatusBarItems(cell, cancellationToken); assert.isDefined(result); - assert.strictEqual((result as any).text, '$(database) Unknown integration (configure)'); + assert.isArray(result); + const items = result as any[]; + assert.strictEqual(items.length, 2); + assert.strictEqual(items[0].text, '$(database) Unknown integration (configure)'); + assert.strictEqual(items[0].alignment, 1); + assert.strictEqual(items[0].command.command, 'deepnote.switchSqlIntegration'); + assert.deepStrictEqual(items[0].command.arguments, [cell]); + assert.strictEqual(items[1].text, 'Variable: df'); + assert.strictEqual(items[1].alignment, 1); + assert.strictEqual(items[1].command.command, 'deepnote.updateSqlVariableName'); + assert.strictEqual(items[1].priority, 90); }); - test('returns undefined when notebook has no project ID', async () => { + test('returns only variable item when notebook has no project ID', async () => { const integrationId = 'postgres-123'; const cell = createMockCell('sql', { sql_integration_id: integrationId @@ -117,7 +193,459 @@ suite('SqlCellStatusBarProvider', () => { const result = await provider.provideCellStatusBarItems(cell, cancellationToken); - assert.isUndefined(result); + assert.isDefined(result); + assert.isArray(result); + const items = result as any[]; + assert.strictEqual(items.length, 1); + + // Check variable status bar item is still shown + const variableItem = items[0]; + assert.strictEqual(variableItem.text, 'Variable: df'); + assert.strictEqual(variableItem.alignment, 1); + assert.strictEqual(variableItem.command.command, 'deepnote.updateSqlVariableName'); + assert.deepStrictEqual(variableItem.command.arguments, [cell]); + assert.strictEqual(variableItem.priority, 90); + }); + + test('shows custom variable name when set in metadata', async () => { + const integrationId = 'postgres-123'; + const cell = createMockCell( + 'sql', + { + sql_integration_id: integrationId, + deepnote_variable_name: 'my_results' + }, + { + deepnoteProjectId: 'project-1' + } + ); + + when(integrationStorage.getProjectIntegrationConfig(anything(), anything())).thenResolve({ + id: integrationId, + name: 'My Postgres DB', + type: IntegrationType.Postgres, + host: 'localhost', + port: 5432, + database: 'test', + username: 'user', + password: 'pass' + }); + + const result = await provider.provideCellStatusBarItems(cell, cancellationToken); + + assert.isDefined(result); + assert.isArray(result); + const items = result as any[]; + assert.strictEqual(items.length, 2); + + // Check variable status bar item shows custom name + const variableItem = items[1]; + assert.strictEqual(variableItem.text, 'Variable: my_results'); + assert.strictEqual(variableItem.alignment, 1); + assert.strictEqual(variableItem.command.command, 'deepnote.updateSqlVariableName'); + }); + + suite('activate', () => { + let activateDisposables: IDisposableRegistry; + let activateProvider: SqlCellStatusBarProvider; + let activateIntegrationStorage: IIntegrationStorage; + + setup(() => { + resetVSCodeMocks(); + activateDisposables = []; + activateIntegrationStorage = mock(); + activateProvider = new SqlCellStatusBarProvider(activateDisposables, instance(activateIntegrationStorage)); + }); + + teardown(() => { + resetVSCodeMocks(); + }); + + test('registers notebook cell status bar provider for deepnote notebooks', () => { + activateProvider.activate(); + + verify( + mockedVSCodeNamespaces.notebooks.registerNotebookCellStatusBarItemProvider('deepnote', activateProvider) + ).once(); + }); + + test('registers deepnote.updateSqlVariableName command', () => { + activateProvider.activate(); + + verify( + mockedVSCodeNamespaces.commands.registerCommand('deepnote.updateSqlVariableName', anything()) + ).once(); + }); + + test('registers deepnote.switchSqlIntegration command', () => { + activateProvider.activate(); + + verify(mockedVSCodeNamespaces.commands.registerCommand('deepnote.switchSqlIntegration', anything())).once(); + }); + + test('listens to integration storage changes', () => { + const onDidChangeIntegrations = new EventEmitter(); + when(activateIntegrationStorage.onDidChangeIntegrations).thenReturn(onDidChangeIntegrations.event); + + activateProvider.activate(); + + // Verify the listener was registered by checking disposables + assert.isTrue(activateDisposables.length > 0); + }); + }); + + suite('event listeners', () => { + let eventDisposables: IDisposableRegistry; + let eventProvider: SqlCellStatusBarProvider; + let eventIntegrationStorage: IIntegrationStorage; + + setup(() => { + eventDisposables = []; + eventIntegrationStorage = mock(); + eventProvider = new SqlCellStatusBarProvider(eventDisposables, instance(eventIntegrationStorage)); + }); + + test('fires onDidChangeCellStatusBarItems when integration storage changes', () => { + const onDidChangeIntegrations = new EventEmitter(); + when(eventIntegrationStorage.onDidChangeIntegrations).thenReturn(onDidChangeIntegrations.event); + + eventProvider.activate(); + + const statusBarChangeHandler = createEventHandler( + eventProvider, + 'onDidChangeCellStatusBarItems', + eventDisposables + ); + + // Fire integration storage change event + onDidChangeIntegrations.fire(); + + assert.strictEqual(statusBarChangeHandler.count, 1, 'onDidChangeCellStatusBarItems should fire once'); + }); + + test('fires onDidChangeCellStatusBarItems multiple times for multiple integration changes', () => { + const onDidChangeIntegrations = new EventEmitter(); + when(eventIntegrationStorage.onDidChangeIntegrations).thenReturn(onDidChangeIntegrations.event); + + eventProvider.activate(); + + const statusBarChangeHandler = createEventHandler( + eventProvider, + 'onDidChangeCellStatusBarItems', + eventDisposables + ); + + // Fire integration storage change event multiple times + onDidChangeIntegrations.fire(); + onDidChangeIntegrations.fire(); + onDidChangeIntegrations.fire(); + + assert.strictEqual( + statusBarChangeHandler.count, + 3, + 'onDidChangeCellStatusBarItems should fire three times' + ); + }); + }); + + suite('updateSqlVariableName command handler', () => { + let commandDisposables: IDisposableRegistry; + let commandProvider: SqlCellStatusBarProvider; + let commandIntegrationStorage: IIntegrationStorage; + let updateVariableNameHandler: Function; + + setup(() => { + resetVSCodeMocks(); + commandDisposables = []; + commandIntegrationStorage = mock(); + commandProvider = new SqlCellStatusBarProvider(commandDisposables, instance(commandIntegrationStorage)); + + // Capture the command handler + when( + mockedVSCodeNamespaces.commands.registerCommand('deepnote.updateSqlVariableName', anything()) + ).thenCall((_, handler) => { + updateVariableNameHandler = handler; + return { + dispose: () => { + return; + } + }; + }); + + commandProvider.activate(); + }); + + teardown(() => { + resetVSCodeMocks(); + }); + + test('updates cell metadata with new variable name', async () => { + const cell = createMockCell('sql', { deepnote_variable_name: 'old_name' }); + const newVariableName = 'new_name'; + + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(newVariableName)); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(true)); + + await updateVariableNameHandler(cell); + + verify(mockedVSCodeNamespaces.window.showInputBox(anything())).once(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).once(); + }); + + test('does not update if user cancels input box', async () => { + const cell = createMockCell('sql', { deepnote_variable_name: 'old_name' }); + + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + await updateVariableNameHandler(cell); + + verify(mockedVSCodeNamespaces.window.showInputBox(anything())).once(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('does not update if new name is same as current name', async () => { + const cell = createMockCell('sql', { deepnote_variable_name: 'same_name' }); + + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('same_name')); + + await updateVariableNameHandler(cell); + + verify(mockedVSCodeNamespaces.window.showInputBox(anything())).once(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('shows error message if workspace edit fails', async () => { + const cell = createMockCell('sql', { deepnote_variable_name: 'old_name' }); + const newVariableName = 'new_name'; + + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(newVariableName)); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(false)); + + await updateVariableNameHandler(cell); + + verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); + }); + + test('fires onDidChangeCellStatusBarItems after successful update', async () => { + const cell = createMockCell('sql', { deepnote_variable_name: 'old_name' }); + const newVariableName = 'new_name'; + + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(newVariableName)); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(true)); + + const statusBarChangeHandler = createEventHandler( + commandProvider, + 'onDidChangeCellStatusBarItems', + commandDisposables + ); + + await updateVariableNameHandler(cell); + + assert.strictEqual(statusBarChangeHandler.count, 1, 'onDidChangeCellStatusBarItems should fire once'); + }); + + test('validates input - rejects empty variable name', async () => { + const cell = createMockCell('sql', {}); + + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenCall((options) => { + const validationResult = options.validateInput(''); + assert.strictEqual(validationResult, 'Variable name cannot be empty'); + return Promise.resolve(undefined); + }); + + await updateVariableNameHandler(cell); + }); + + test('validates input - rejects invalid Python identifier', async () => { + const cell = createMockCell('sql', {}); + + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenCall((options) => { + const validationResult = options.validateInput('123invalid'); + assert.strictEqual(validationResult, 'Variable name must be a valid Python identifier'); + return Promise.resolve(undefined); + }); + + await updateVariableNameHandler(cell); + }); + }); + + suite('switchSqlIntegration command handler', () => { + let commandDisposables: IDisposableRegistry; + let commandProvider: SqlCellStatusBarProvider; + let commandIntegrationStorage: IIntegrationStorage; + let switchIntegrationHandler: Function; + + setup(() => { + resetVSCodeMocks(); + commandDisposables = []; + commandIntegrationStorage = mock(); + commandProvider = new SqlCellStatusBarProvider(commandDisposables, instance(commandIntegrationStorage)); + + // Capture the command handler + when(mockedVSCodeNamespaces.commands.registerCommand('deepnote.switchSqlIntegration', anything())).thenCall( + (_, handler) => { + switchIntegrationHandler = handler; + return { + dispose: () => { + return; + } + }; + } + ); + + commandProvider.activate(); + }); + + teardown(() => { + resetVSCodeMocks(); + }); + + test('updates cell metadata with selected integration', async () => { + const cell = createMockCell('sql', { sql_integration_id: 'old-integration' }); + const newIntegrationId = 'new-integration'; + + when(commandIntegrationStorage.getAll()).thenReturn( + Promise.resolve([ + { + id: newIntegrationId, + name: 'New Integration', + type: IntegrationType.Postgres, + host: 'localhost', + port: 5432, + database: 'test', + username: 'user', + password: 'pass' + } + ]) + ); + + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( + Promise.resolve({ id: newIntegrationId, label: 'New Integration' } as any) + ); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(true)); + + await switchIntegrationHandler(cell); + + verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).once(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).once(); + }); + + test('does not update if user cancels quick pick', async () => { + const cell = createMockCell('sql', { sql_integration_id: 'old-integration' }); + + when(commandIntegrationStorage.getAll()).thenReturn(Promise.resolve([])); + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( + Promise.resolve(undefined) + ); + + await switchIntegrationHandler(cell); + + verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).once(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('shows error message if workspace edit fails', async () => { + const cell = createMockCell('sql', { sql_integration_id: 'old-integration' }); + const newIntegrationId = 'new-integration'; + + when(commandIntegrationStorage.getAll()).thenReturn(Promise.resolve([])); + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( + Promise.resolve({ id: newIntegrationId, label: 'New Integration' } as any) + ); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(false)); + + await switchIntegrationHandler(cell); + + verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); + }); + + test('fires onDidChangeCellStatusBarItems after successful update', async () => { + const cell = createMockCell('sql', { sql_integration_id: 'old-integration' }); + const newIntegrationId = 'new-integration'; + + when(commandIntegrationStorage.getAll()).thenReturn(Promise.resolve([])); + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( + Promise.resolve({ id: newIntegrationId, label: 'New Integration' } as any) + ); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(true)); + + const statusBarChangeHandler = createEventHandler( + commandProvider, + 'onDidChangeCellStatusBarItems', + commandDisposables + ); + + await switchIntegrationHandler(cell); + + assert.strictEqual(statusBarChangeHandler.count, 1, 'onDidChangeCellStatusBarItems should fire once'); + }); + + test('executes manage integrations command when configure option is selected', async () => { + const cell = createMockCell('sql', { sql_integration_id: 'current-integration' }); + + when(commandIntegrationStorage.getAll()).thenReturn(Promise.resolve([])); + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( + Promise.resolve({ id: '__configure__', label: 'Configure current integration' } as any) + ); + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenReturn( + Promise.resolve(undefined) + ); + + await switchIntegrationHandler(cell); + + verify( + mockedVSCodeNamespaces.commands.executeCommand(Commands.ManageIntegrations, 'current-integration') + ).once(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('includes DuckDB integration in quick pick items', async () => { + const cell = createMockCell('sql', {}); + let quickPickItems: any[] = []; + + when(commandIntegrationStorage.getAll()).thenReturn(Promise.resolve([])); + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenCall((items) => { + quickPickItems = items; + return Promise.resolve(undefined); + }); + + await switchIntegrationHandler(cell); + + const duckDbItem = quickPickItems.find((item) => item.id === DATAFRAME_SQL_INTEGRATION_ID); + assert.isDefined(duckDbItem, 'DuckDB integration should be in quick pick items'); + assert.strictEqual(duckDbItem.label, 'DataFrame SQL (DuckDB)'); + }); + + test('marks current integration as selected in quick pick', async () => { + const currentIntegrationId = 'current-integration'; + const cell = createMockCell('sql', { sql_integration_id: currentIntegrationId }); + let quickPickItems: any[] = []; + + when(commandIntegrationStorage.getAll()).thenReturn( + Promise.resolve([ + { + id: currentIntegrationId, + name: 'Current Integration', + type: IntegrationType.Postgres, + host: 'localhost', + port: 5432, + database: 'test', + username: 'user', + password: 'pass' + } + ]) + ); + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenCall((items) => { + quickPickItems = items; + return Promise.resolve(undefined); + }); + + await switchIntegrationHandler(cell); + + const currentItem = quickPickItems.find((item) => item.id === currentIntegrationId); + assert.isDefined(currentItem, 'Current integration should be in quick pick items'); + assert.strictEqual(currentItem.detail, 'Currently selected'); + }); }); function createMockCell(