diff --git a/packages/extension/icons/info-dark.svg b/packages/extension/icons/info-dark.svg new file mode 100644 index 000000000..2952b397b --- /dev/null +++ b/packages/extension/icons/info-dark.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/packages/extension/icons/info-light.svg b/packages/extension/icons/info-light.svg new file mode 100644 index 000000000..b7c11fe5c --- /dev/null +++ b/packages/extension/icons/info-light.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/packages/extension/icons/limit-dark.svg b/packages/extension/icons/limit-dark.svg new file mode 100644 index 000000000..01aa8d49c --- /dev/null +++ b/packages/extension/icons/limit-dark.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/packages/extension/icons/limit-light.svg b/packages/extension/icons/limit-light.svg new file mode 100644 index 000000000..44949e660 --- /dev/null +++ b/packages/extension/icons/limit-light.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/packages/extension/package.json b/packages/extension/package.json index dd4d5d4db..ac7c19e14 100644 --- a/packages/extension/package.json +++ b/packages/extension/package.json @@ -113,6 +113,17 @@ "configuration": "./language/language-configuration.json" } ], + "notebooks": [ + { + "type": "sqltools-notebook", + "displayName": "SQLTools Notebook", + "selector": [ + { + "filenamePattern": "*.sqlnb" + } + ] + } + ], "grammars": [ { "language": "sql", @@ -121,6 +132,47 @@ } ], "commands": [ + { + "title": "List Available SQLTools Commands", + "command": "sqltools.notebook.listCommands", + "category": "SQLTools Notebook Debug" + }, + { + "title": "Select Connection for Cell", + "command": "sqltools.notebook.selectCellConnection", + "category": "SQLTools Notebook", + "icon": { + "light": "icons/connect-light.svg", + "dark": "icons/connect-dark.svg" + } + }, + { + "title": "Clear Connection for Cell", + "command": "sqltools.notebook.clearCellConnection", + "category": "SQLTools Notebook", + "icon": { + "light": "icons/disconnect-light.svg", + "dark": "icons/disconnect-dark.svg" + } + }, + { + "title": "Show SQL Notebook Connection Help", + "command": "sqltools.notebook.showConnectionHelp", + "category": "SQLTools Notebook", + "icon": { + "light": "icons/help-light.svg", + "dark": "icons/help-dark.svg" + } + }, + { + "title": "Limit Query Results", + "command": "sqltools.notebook.limitResults", + "category": "SQLTools Notebook", + "icon": { + "light": "icons/limit-light.svg", + "dark": "icons/limit-dark.svg" + } + }, { "title": "Connect", "command": "sqltools.selectConnection", @@ -411,6 +463,54 @@ "key": "ctrl+e ctrl+q", "mac": "cmd+e q", "when": "!config.sqltools.disableChordKeybindings && editorTextFocus && editorHasSelection" + }, + { + "command": "sqltools.executeQuery", + "key": "ctrl+alt+e", + "mac": "cmd+alt+e", + "when": "editorTextFocus && editorLangId == 'sql'" + }, + { + "command": "sqltools.executeCurrentQuery", + "key": "ctrl+alt+q", + "mac": "cmd+alt+q", + "when": "editorTextFocus && editorLangId == 'sql'" + }, + { + "command": "sqltools.executeQueryFromFile", + "key": "ctrl+alt+f", + "mac": "cmd+alt+f", + "when": "editorTextFocus && editorLangId == 'sql'" + }, + { + "command": "sqltools.formatSql", + "key": "shift+alt+f", + "mac": "shift+alt+f", + "when": "editorHasSelection && editorLangId == 'sql'" + }, + { + "command": "sqltools.formatSql", + "key": "shift+alt+f", + "mac": "shift+alt+f", + "when": "editorHasSelection" + }, + { + "command": "sqltools.notebook.selectCellConnection", + "key": "ctrl+k ctrl+c", + "mac": "cmd+k cmd+c", + "when": "notebookEditorFocused && notebookType == 'sqltools-notebook'" + }, + { + "command": "sqltools.notebook.clearCellConnection", + "key": "ctrl+k ctrl+r", + "mac": "cmd+k cmd+r", + "when": "notebookEditorFocused && notebookType == 'sqltools-notebook'" + }, + { + "command": "sqltools.notebook.showConnectionHelp", + "key": "ctrl+k ctrl+h", + "mac": "cmd+k cmd+h", + "when": "notebookType == 'sqltools-notebook'" } ], "configuration": { @@ -1069,6 +1169,22 @@ { "command": "sqltools.bookmarkSelection", "when": "editorHasSelection" + }, + { + "command": "sqltools.notebook.selectConnection", + "when": "notebookType == 'sqltools-notebook'" + } + ], + "notebook/title": [ + { + "command": "sqltools.notebook.selectCellConnection", + "when": "notebookType == 'sqltools-notebook'", + "group": "navigation" + }, + { + "command": "sqltools.notebook.showConnectionHelp", + "when": "notebookType == 'sqltools-notebook'", + "group": "navigation" } ], "view/title": [ diff --git a/packages/extension/src/index.ts b/packages/extension/src/index.ts index b99d825c6..94b111e41 100644 --- a/packages/extension/src/index.ts +++ b/packages/extension/src/index.ts @@ -13,6 +13,7 @@ import PluginResourcesMap from '@sqltools/util/plugin-resources'; import SQLToolsLanguageClient from './language-client'; import Timer from '@sqltools/util/timer'; import Utils from './api/utils'; +import { registerSQLNotebook, deactivateSQLNotebook } from './notebook'; const log = createLogger(); @@ -78,6 +79,9 @@ export class SQLToolsExtension implements IExtension { } public deactivate = (): void => { + // Clean up notebook resources + deactivateSQLNotebook(); + return Context.subscriptions.forEach((sub) => void sub.dispose()); } @@ -295,30 +299,49 @@ export class SQLToolsExtension implements IExtension { } } -let instance: SQLToolsExtension; export function activate(ctx: ExtensionContext) { + Context.set(ctx); + migrateFilesToNewPaths(); + + // Register SQL notebook functionality + registerSQLNotebook(ctx); + try { - Context.set(ctx); - if (instance) return; - migrateFilesToNewPaths(); - instance = new SQLToolsExtension(); - instance.registerPlugin([ - FormatterPlugin, + const ext = new SQLToolsExtension(); + + // Register plugins - use the correct approach for each plugin + ext.registerPlugin([ + // These are already plugin objects, don't need 'new' ConnectionManagerPlugin, - new HistoryManagerPlugin, - new BookmarksManagerPlugin, - new AuthenticationProviderPlugin, - new ObjectDropProviderPlugin, - ]) - return instance.activate(); - - } catch (err) { - log.fatal('failed to activate: %O', err); + FormatterPlugin, + // These need instantiation as they're exported as classes + new ObjectDropProviderPlugin(), + new HistoryManagerPlugin(), + new BookmarksManagerPlugin(), + new AuthenticationProviderPlugin(), + ]); + + return ext.activate(); + } catch (e) { + // Create an error handler with a string message + const handler = ErrorHandler.create('Failed to activate extension'); + // Then pass the actual error to the handler + handler(e); } } export function deactivate() { - if (!instance) return; - instance.deactivate(); - instance = undefined; + // Clean up notebook resources + deactivateSQLNotebook(); + + // Just call dispose on each subscription + if (Context.subscriptions) { + Context.subscriptions.forEach(sub => { + try { + sub.dispose(); + } catch (e) { + // Ignore dispose errors + } + }); + } } diff --git a/packages/extension/src/notebook/controller.ts b/packages/extension/src/notebook/controller.ts new file mode 100644 index 000000000..209265c62 --- /dev/null +++ b/packages/extension/src/notebook/controller.ts @@ -0,0 +1,1285 @@ +import * as vscode from 'vscode'; +import { NSDatabase } from '@sqltools/types'; +import { EXT_NAMESPACE } from '@sqltools/util/constants'; + +interface ConnectionDetails { + name?: string; + driver?: string; + [key: string]: any; +} + +export class SQLNotebookController { + readonly controllerId = 'sqltools-notebook-controller'; + readonly notebookType = 'sqltools-notebook'; + readonly label = 'SQLTools Notebook'; + readonly supportedLanguages = ['sql']; + + private readonly _controller: vscode.NotebookController; + private _executionOrder = 0; + // Store last used connection for setting new cell defaults + private _lastUsedConnection: string | null = null; + + constructor() { + this._controller = vscode.notebooks.createNotebookController( + this.controllerId, + this.notebookType, + this.label + ); + + this._controller.supportedLanguages = this.supportedLanguages; + this._controller.supportsExecutionOrder = true; + this._controller.executeHandler = this._execute.bind(this); + + // Register selection command + this._registerCommands(); + + // Listen for notebook changes + vscode.window.onDidChangeActiveNotebookEditor(this._onActiveNotebookChanged, this); + + // Listen for notebook document changes to set default connection for new cells + vscode.workspace.onDidChangeNotebookDocument(this._onNotebookDocumentChanged, this); + + // Register message handler for export + this._registerMessageHandler(); + } + + private _onActiveNotebookChanged(editor: vscode.NotebookEditor | undefined): void { + if (!editor || editor.notebook.notebookType !== this.notebookType) { + // Disable controller when no notebook is active + this._controller.updateNotebookAffinity(editor?.notebook, vscode.NotebookControllerAffinity.Default); + return; + } + + if (this._lastUsedConnection) { + // Enable controller when last connection is available + this._controller.updateNotebookAffinity(editor.notebook, vscode.NotebookControllerAffinity.Preferred); + } else { + // Disable controller when no connection is selected + this._controller.updateNotebookAffinity(editor.notebook, vscode.NotebookControllerAffinity.Default); + } + } + + private _onNotebookDocumentChanged(event: vscode.NotebookDocumentChangeEvent): void { + // Only handle our notebook type + if (event.notebook.notebookType !== this.notebookType) { + return; + } + + // Check for cell additions + const cellChanges = event.contentChanges; + for (const change of cellChanges) { + // Process added cells + if (change.addedCells && change.addedCells.length > 0) { + for (const cell of change.addedCells) { + // Set default connection for new cells + this._setDefaultConnectionForNewCell(cell); + } + } + } + } + + // Helper method to get connection details + private async _getConnectionDetails(connectionId: string): Promise { + try { + const connections = await vscode.commands.executeCommand( + `${EXT_NAMESPACE}.getConnections`, + { connectedOnly: true } + ) || []; + + const connection = connections.find(conn => { + const id = typeof conn === 'string' ? conn : conn.id || conn.name; + return id === connectionId; + }); + + if (!connection) return null; + + // Extract and return connection details + return typeof connection === 'string' + ? { name: connection } + : connection; + } catch (error) { + console.error('Error fetching connection details:', error); + return null; + } + } + + private _registerCommands() { + // Register a command to list all available commands (for debugging) + vscode.commands.registerCommand('sqltools.notebook.listCommands', async () => { + try { + // Get all registered commands that have sqltools in their name + const allCommands = await vscode.commands.getCommands(true); + const sqlCommands = allCommands.filter(cmd => cmd.includes('sqltools')); + + console.log('Available SQLTools commands:', sqlCommands); + + // Show commands in the output channel + vscode.window.showInformationMessage(`Found ${sqlCommands.length} SQLTools commands.`); + + // Return the list for programmatic use + return sqlCommands; + } catch (error) { + console.error('Error listing commands:', error); + return []; + } + }); + + // Register a command to select a connection for a specific cell + vscode.commands.registerCommand('sqltools.notebook.selectCellConnection', async (cell: vscode.NotebookCell) => { + if (!cell || cell.kind !== vscode.NotebookCellKind.Code || cell.document.languageId !== 'sql') { + vscode.window.showErrorMessage('Not a valid SQL cell'); + return; + } + + const connectionId = await this._selectConnection(); + if (connectionId) { + // Store the connection for this cell in cell metadata + this._setCellConnection(cell, connectionId); + + // Remember this as the last used connection + this._lastUsedConnection = connectionId; + + // Show confirmation + const connDetails = await this._getConnectionDetails(connectionId); + const connectionName = connDetails?.name || connectionId; + const driverName = connDetails?.driver ? ` (${connDetails.driver})` : ''; + vscode.window.showInformationMessage(`Cell will use connection: ${connectionName}${driverName}. Run the cell to see results.`); + + return connectionId; + } + return null; + }); + + // Register a command to clear the connection for a specific cell + vscode.commands.registerCommand('sqltools.notebook.clearCellConnection', async (cell: vscode.NotebookCell) => { + if (!cell || cell.kind !== vscode.NotebookCellKind.Code || cell.document.languageId !== 'sql') { + vscode.window.showErrorMessage('Not a valid SQL cell'); + return; + } + + // Clear the connection from this cell's metadata + if (cell.metadata?.sqltools_connection) { + const edit = new vscode.WorkspaceEdit(); + const cellMetadata = { ...cell.metadata }; + delete cellMetadata.sqltools_connection; + + const nbEdit = vscode.NotebookEdit.updateCellMetadata(cell.index, cellMetadata); + edit.set(cell.notebook.uri, [nbEdit]); + + await vscode.workspace.applyEdit(edit); + vscode.window.showInformationMessage('Cell connection cleared. You need to select a connection before running this cell.'); + } else { + vscode.window.showInformationMessage('This cell has no specific connection set.'); + } + }); + + // Register a command to show SQL Notebook connection help + vscode.commands.registerCommand('sqltools.notebook.showConnectionHelp', async () => { + if (!vscode.window.activeNotebookEditor || vscode.window.activeNotebookEditor.notebook.notebookType !== this.notebookType) { + vscode.window.showErrorMessage('No active SQL notebook found'); + return; + } + + const helpContent = `## SQLTools Notebook Connection Commands + +### Cell-Level Connections + +| Command | Keybinding (Windows/Linux) | Keybinding (Mac) | Description | +| ------- | -------------------------- | ---------------- | ----------- | +| Select Cell Connection | Ctrl+K Ctrl+E | Cmd+K Cmd+E | Choose a specific connection for the current cell | +| Clear Cell Connection | Ctrl+K Ctrl+R | Cmd+K Cmd+R | Remove the cell-specific connection | + +#### About Cell Connections + +- Each cell can have its own database connection +- When you create a new cell, it will automatically use the last connection you selected +- The connection for each cell is displayed in the cell status bar +- You must select a connection for each cell before running it + +You can also click on the connection indicator in the status bar to select a connection for the current cell. +`; + + const cell = new vscode.NotebookCellData( + vscode.NotebookCellKind.Markup, + helpContent, + 'markdown' + ); + + const edit = new vscode.WorkspaceEdit(); + const nbEdit = vscode.NotebookEdit.insertCells(0, [cell]); + edit.set(vscode.window.activeNotebookEditor.notebook.uri, [nbEdit]); + + await vscode.workspace.applyEdit(edit); + vscode.window.showInformationMessage('Connection help added to notebook'); + }); + + // Register a command to export notebook results + vscode.commands.registerCommand('sqltools.notebook.exportResults', async (cell: vscode.NotebookCell, format: 'csv' | 'json') => { + if (!cell || !cell.outputs || cell.outputs.length === 0) { + vscode.window.showWarningMessage('No results to export. Run the query first.'); + return; + } + + try { + // Get the table data from the cell's HTML output + const htmlOutput = cell.outputs + .flatMap(output => output.items) + .find(item => item.mime === 'text/html'); + + if (!htmlOutput) { + vscode.window.showErrorMessage('No results table found in the output.'); + return; + } + + // Parse the HTML content to extract table data + const htmlContent = Buffer.from(htmlOutput.data).toString('utf8'); + const tableData = this._extractTableDataFromHtml(htmlContent); + + if (!tableData || tableData.rows.length === 0) { + vscode.window.showErrorMessage('Could not extract data from the results.'); + return; + } + + // Format the data according to the requested format + let formattedContent = ''; + if (format === 'csv') { + formattedContent = this._formatAsCsv(tableData); + } else { // json + formattedContent = this._formatAsJson(tableData); + } + + // Show save dialog + const filters = format === 'csv' + ? { 'CSV Files': ['csv'] } + : { 'JSON Files': ['json'] }; + + const uri = await vscode.window.showSaveDialog({ + filters, + saveLabel: `Export as ${format.toUpperCase()}`, + title: `Export Query Results as ${format.toUpperCase()}` + }); + + if (!uri) return; // User cancelled + + // Write the file + await vscode.workspace.fs.writeFile( + uri, + new Uint8Array(Buffer.from(formattedContent)) + ); + + // Show success message + const openAction = 'Open File'; + const action = await vscode.window.showInformationMessage( + `Results exported to ${uri.fsPath}`, + openAction + ); + + if (action === openAction) { + await vscode.commands.executeCommand('vscode.open', uri); + } + } catch (error) { + console.error('Error exporting results:', error); + vscode.window.showErrorMessage(`Failed to export results: ${error instanceof Error ? error.message : String(error)}`); + } + }); + + // Register a command to limit query results + vscode.commands.registerCommand('sqltools.notebook.limitResults', async (cell: vscode.NotebookCell) => { + const quickPickItems = [ + { label: 'No limit', description: 'Run query without a LIMIT clause', value: 0 }, + { label: 'LIMIT 10', description: 'Return a maximum of 10 rows', value: 10 }, + { label: 'LIMIT 50', description: 'Return a maximum of 50 rows', value: 50 }, + { label: 'LIMIT 100', description: 'Return a maximum of 100 rows', value: 100 }, + { label: 'LIMIT 1000', description: 'Return a maximum of 1000 rows', value: 1000 }, + { label: 'Custom...', description: 'Specify a custom limit', value: -1 } + ]; + + const selectedOption = await vscode.window.showQuickPick(quickPickItems, { + placeHolder: 'Select a limit for query results', + title: 'Limit Query Results' + }); + + if (!selectedOption) return; // User cancelled + + if (selectedOption.value === -1) { + // Handle custom limit input + const customLimit = await vscode.window.showInputBox({ + prompt: 'Enter a custom row limit', + placeHolder: 'e.g., 500', + validateInput: (value) => { + const num = parseInt(value); + return (!isNaN(num) && num > 0) ? null : 'Please enter a positive number'; + } + }); + + if (customLimit) { + const limitValue = parseInt(customLimit); + if (!isNaN(limitValue) && limitValue > 0) { + // Store the limit in notebook metadata or cell metadata + this._setNotebookCellLimit(cell, limitValue); + vscode.window.showInformationMessage(`Query will run with LIMIT ${limitValue}. Run the cell to see results.`); + } + } + } else { + // Store the predefined limit + this._setNotebookCellLimit(cell, selectedOption.value); + if (selectedOption.value === 0) { + vscode.window.showInformationMessage('Query will run without a LIMIT clause. Run the cell to see results.'); + } else { + vscode.window.showInformationMessage(`Query will run with ${selectedOption.label}. Run the cell to see results.`); + } + } + }); + } + + private async _execute( + cells: vscode.NotebookCell[], + notebook: vscode.NotebookDocument, + _controller: vscode.NotebookController + ): Promise { + for (let cell of cells) { + // Only execute cells that have a connection specified + const cellConnectionId = cell.metadata?.sqltools_connection as string; + + // If the cell doesn't have a connection, prompt to select one + if (!cellConnectionId) { + const selectedConnection = await this._selectConnection(); + if (selectedConnection) { + // Set the connection for this cell + this._setCellConnection(cell, selectedConnection); + // Remember this as the last used connection + this._lastUsedConnection = selectedConnection; + // Execute with the newly selected connection + await this._doExecution(cell, selectedConnection); + } else { + // No connection selected, show error in cell output + const execution = this._controller.createNotebookCellExecution(cell); + execution.executionOrder = ++this._executionOrder; + execution.start(Date.now()); + execution.replaceOutput([ + new vscode.NotebookCellOutput([ + vscode.NotebookCellOutputItem.text('No connection selected. Please select a database connection for this cell.') + ]) + ]); + execution.end(false, Date.now()); + } + } else { + // Execute with the cell's connection + await this._doExecution(cell, cellConnectionId); + } + } + } + + private async _doExecution(cell: vscode.NotebookCell, connectionId: string): Promise { + if (cell.document.languageId !== 'sql') { + return; // Only execute SQL cells + } + + const execution = this._controller.createNotebookCellExecution(cell); + execution.executionOrder = ++this._executionOrder; + execution.start(Date.now()); + + try { + if (!connectionId) { + execution.replaceOutput([ + new vscode.NotebookCellOutput([ + vscode.NotebookCellOutputItem.text('No connection selected. Please select a database connection first.') + ]) + ]); + execution.end(false, Date.now()); + return; + } + + // Get connection details for display + const connDetails = await this._getConnectionDetails(connectionId); + const connectionName = connDetails?.name || connectionId; + const driverName = connDetails?.driver ? ` (${connDetails.driver})` : ''; + + // Get the original SQL query from the cell + const originalQuery = cell.document.getText(); + + // Check if we need to apply a LIMIT to this query + const limitValue = cell.metadata?.sqltools_limit as number; + const query = this._applyLimitToQuery(originalQuery, limitValue, connDetails?.driver); + + console.log(`SQLTools Notebook: Executing query with connection ID: ${connectionId}`); + if (limitValue && query !== originalQuery) { + console.log(`SQLTools Notebook: Applied LIMIT ${limitValue} to query`); + } + + try { + // Ensure connectionId is a string before passing it to executeQuery + const connId = typeof connectionId === 'string' ? connectionId : String(connectionId); + + // We need to directly execute the query but prevent the default behavior of showing results in a new window + const results = await vscode.commands.executeCommand( + `${EXT_NAMESPACE}.executeQuery`, + query, + { + connNameOrId: connId, // Use the exact parameter name expected by connection manager + showOutput: false, + runInNotebook: true + } + ); + + if (!results || results.length === 0) { + console.log('SQLTools Notebook: No results returned from query'); + execution.replaceOutput([ + new vscode.NotebookCellOutput([ + vscode.NotebookCellOutputItem.text(`Query executed on ${connectionName}${driverName}. No results returned.`) + ]) + ]); + execution.end(true, Date.now()); + return; + } + + console.log(`SQLTools Notebook: Query returned ${results.length} result sets`); + + // Process and display the results only in the notebook cell + for (const result of results) { + // Add connection info to the result for display + const connectionInfo = `
${connectionName}${driverName}
`; + + // Convert the result to a renderable format + const tableHtml = this._createHtmlTable(result); + + // Use appropriate success/error message based on result status + const statusText = result.error + ? 'Query execution failed.' + : 'Query executed successfully.'; + + execution.replaceOutput([ + new vscode.NotebookCellOutput([ + vscode.NotebookCellOutputItem.text(statusText), + vscode.NotebookCellOutputItem.text(`${connectionInfo}${tableHtml}`, 'text/html') + ]) + ]); + } + } catch (queryError) { + console.error('SQLTools Notebook: Error executing query:', queryError); + throw queryError; // Re-throw to be caught by the outer try/catch + } + + execution.end(true, Date.now()); + } catch (error) { + console.error('SQLTools Notebook: Execution error:', error); + const errorMessage = error instanceof Error + ? `${error.message}\n${error.stack}` + : String(error); + + execution.replaceOutput([ + new vscode.NotebookCellOutput([ + vscode.NotebookCellOutputItem.error(error instanceof Error ? error : new Error(String(error))), + vscode.NotebookCellOutputItem.text(`Error details: ${errorMessage}`, 'text/plain') + ]) + ]); + execution.end(false, Date.now()); + } + } + + private _createHtmlTable(result: NSDatabase.IResult): string { + if (!result.results || result.results.length === 0) { + // Properly format messages from error objects + const errorMessages = result.messages ? result.messages.map(msg => + typeof msg === 'string' ? msg : msg.message || JSON.stringify(msg) + ).join(' ') : ''; + + // Check if this is an error result and style accordingly + if (result.error) { + // Create error styled message with details if available + let errorDetails = ''; + if (result.rawError) { + const rawErrorStr = typeof result.rawError === 'string' + ? result.rawError + : result.rawError.message || JSON.stringify(result.rawError); + errorDetails = `
${rawErrorStr}
`; + } + + return ` +
+
Error executing query: ${errorMessages}
+ ${errorDetails} +
+ `; + } + + return `
No data returned. ${errorMessages}
`; + } + + const columns = result.cols || Object.keys(result.results[0]); + const totalRows = result.results.length; + + let html = ` + +
+ +
+ +
+ + + + + ${columns.map((col, index) => ``).join('')} + + + + `; + + // Add initial rows (first page with default page size of 10) + const initialPageSize = Math.min(10, totalRows); + result.results.slice(0, initialPageSize).forEach(row => { + html += ''; + columns.forEach(col => { + const value = row[col]; + html += ``; + }); + html += ''; + }); + + html += ` + +
${col}
${value === null || value === undefined ? 'NULL' : String(value)}
+ + +
+
+ Rows per page: + + +
+
+ + Page 1 of 1 | + Records 1-${Math.min(initialPageSize, totalRows)} of ${totalRows} + + + +
+
+ `; + + if (result.messages && result.messages.length > 0) { + html += ` +
+ ${result.messages.map(msg => typeof msg === 'string' ? msg : (msg.message || JSON.stringify(msg))).join('
')} +
+ `; + } + + // Add JavaScript for pagination, sorting and filtering + html += ` + + `; + + html += '
'; + return html; + } + + private _registerMessageHandler() { + // Register handler for notebook cell output messages (for export buttons) + vscode.notebooks.registerNotebookCellStatusBarItemProvider(this.notebookType, { + provideCellStatusBarItems: (cell, token) => { + if (cell.kind !== vscode.NotebookCellKind.Code || cell.document.languageId !== 'sql') { + return []; + } + + // Get the current limit setting to display in the button + const currentLimit = cell.metadata?.sqltools_limit as number; + const limitLabel = currentLimit && currentLimit > 0 + ? `Limit: ${currentLimit}` + : 'Limit: none'; + + const limitQueryItem = new vscode.NotebookCellStatusBarItem( + limitLabel, + vscode.NotebookCellStatusBarAlignment.Right + ); + limitQueryItem.command = { + title: 'Limit Results', + command: 'sqltools.notebook.limitResults', + arguments: [cell] + }; + + // Get the current cell's connection to display in the button + const cellConnectionId = cell.metadata?.sqltools_connection as string; + let connectionLabel = 'No Connection Selected'; + let connectionTooltip = 'Select a connection for this cell'; + + if (cellConnectionId) { + // If cell has its own connection, show it + connectionLabel = `Connection: ${cellConnectionId.split('.').pop() || cellConnectionId}`; + connectionTooltip = 'Cell connection'; + } + + const connectionItem = new vscode.NotebookCellStatusBarItem( + connectionLabel, + vscode.NotebookCellStatusBarAlignment.Right + ); + connectionItem.command = { + title: 'Select Cell Connection', + command: 'sqltools.notebook.selectCellConnection', + arguments: [cell] + }; + connectionItem.tooltip = connectionTooltip; + + // Add clear connection item if the cell has a specific connection + const items = [limitQueryItem, connectionItem]; + + if (cellConnectionId) { + const clearConnectionItem = new vscode.NotebookCellStatusBarItem( + 'Clear Connection', + vscode.NotebookCellStatusBarAlignment.Right + ); + clearConnectionItem.command = { + title: 'Clear Cell Connection', + command: 'sqltools.notebook.clearCellConnection', + arguments: [cell] + }; + items.push(clearConnectionItem); + } + + // Export items + const exportCsvItem = new vscode.NotebookCellStatusBarItem( + 'Export CSV', + vscode.NotebookCellStatusBarAlignment.Right + ); + exportCsvItem.command = { + title: 'Export CSV', + command: 'sqltools.notebook.exportResults', + arguments: [cell, 'csv'] + }; + + const exportJsonItem = new vscode.NotebookCellStatusBarItem( + 'Export JSON', + vscode.NotebookCellStatusBarAlignment.Right + ); + exportJsonItem.command = { + title: 'Export JSON', + command: 'sqltools.notebook.exportResults', + arguments: [cell, 'json'] + }; + + items.push(exportCsvItem, exportJsonItem); + return items; + } + }); + } + + private _extractTableDataFromHtml(htmlContent: string): { headers: string[], rows: any[][] } | null { + try { + // Simple regex-based extraction - considers a basic HTML table structure + const tableRegex = /]*>([\s\S]*?)<\/table>/i; + const tableMatch = htmlContent.match(tableRegex); + + if (!tableMatch) return null; + + const tableHtml = tableMatch[0]; + + // Extract headers + const headerRegex = /]*>(.*?)<\/th>/gi; + const headers: string[] = []; + let headerMatch; + while ((headerMatch = headerRegex.exec(tableHtml)) !== null) { + headers.push(headerMatch[1].trim()); + } + + // Extract rows + const rowRegex = /]*>([\s\S]*?)<\/tr>/gi; + const rows: any[][] = []; + + let rowMatch; + let startFromSecondRow = false; // Skip header row + + while ((rowMatch = rowRegex.exec(tableHtml)) !== null) { + if (!startFromSecondRow) { + startFromSecondRow = true; + continue; // Skip header row + } + + const rowHtml = rowMatch[1]; + const cellRegex = /]*>(.*?)<\/td>/gi; + const rowData: any[] = []; + + let cellMatch; + while ((cellMatch = cellRegex.exec(rowHtml)) !== null) { + let value = cellMatch[1].trim(); + // Convert NULL string to null + if (value === 'NULL') { + rowData.push(null); + } else { + // Try to convert to number if applicable + const numValue = Number(value); + rowData.push(isNaN(numValue) ? value : numValue); + } + } + + rows.push(rowData); + } + + return { headers, rows }; + } catch (error) { + console.error('Error extracting table data from HTML:', error); + return null; + } + } + + private _formatAsCsv(data: { headers: string[], rows: any[][] }): string { + const { headers, rows } = data; + + // Escape and quote function for CSV + const escapeCSV = (value: any): string => { + if (value === null || value === undefined) return ''; + const strValue = String(value); + // If the value contains a comma, quote, or newline, wrap it in quotes and escape existing quotes + if (strValue.includes(',') || strValue.includes('"') || strValue.includes('\n')) { + return `"${strValue.replace(/"/g, '""')}"`; + } + return strValue; + }; + + // Format the headers + const csvRows = [ + headers.map(escapeCSV).join(',') + ]; + + // Format each row + for (const row of rows) { + csvRows.push(row.map(escapeCSV).join(',')); + } + + return csvRows.join('\n'); + } + + private _formatAsJson(data: { headers: string[], rows: any[][] }): string { + const { headers, rows } = data; + + // Convert to array of objects with named properties + const jsonData = rows.map(row => { + const obj: Record = {}; + headers.forEach((header, index) => { + obj[header] = row[index]; + }); + return obj; + }); + + return JSON.stringify(jsonData, null, 2); + } + + private _setNotebookCellLimit(cell: vscode.NotebookCell, limitValue: number): void { + // We use VSCode's cell metadata to store our limit value + const edit = new vscode.WorkspaceEdit(); + const cellMetadata = { ...cell.metadata, sqltools_limit: limitValue }; + + const nbEdit = vscode.NotebookEdit.updateCellMetadata(cell.index, cellMetadata); + edit.set(cell.notebook.uri, [nbEdit]); + + vscode.workspace.applyEdit(edit); + } + + private _setCellConnection(cell: vscode.NotebookCell, connectionId: string): void { + // Store the connection ID in cell metadata + const edit = new vscode.WorkspaceEdit(); + const cellMetadata = { ...cell.metadata, sqltools_connection: connectionId }; + + const nbEdit = vscode.NotebookEdit.updateCellMetadata(cell.index, cellMetadata); + edit.set(cell.notebook.uri, [nbEdit]); + + vscode.workspace.applyEdit(edit); + } + + // Method to apply limit to a query if needed + private _applyLimitToQuery(query: string, limitValue: number, driver?: string): string { + if (!limitValue || limitValue <= 0) return query; + + // Skip adding LIMIT for non-SELECT queries or if there's already a LIMIT clause + const trimmedQuery = query.trim().toLowerCase(); + + // Check if it's an action query (non-SELECT) + if (!trimmedQuery.startsWith('select')) return query; + + // Check if query already has a LIMIT clause + if (trimmedQuery.includes(' limit ')) return query; + + // Remove trailing semicolon before adding LIMIT + let processedQuery = query.trim(); + if (processedQuery.endsWith(';')) { + processedQuery = processedQuery.slice(0, -1); + } + + // Add LIMIT based on database driver + switch (driver?.toLowerCase()) { + case 'mssql': + // For SQL Server, we use a different approach + if (trimmedQuery.includes(' top ')) return query; + // Add TOP clause for SQL Server + return processedQuery.replace(/select\s+/i, `SELECT TOP ${limitValue} `); + default: + // Most databases use a LIMIT clause + return `${processedQuery} LIMIT ${limitValue}`; + } + } + + // Add a method to set cell connection based on last used connection + private _setDefaultConnectionForNewCell(cell: vscode.NotebookCell): void { + // Only set default connection for SQL code cells that don't already have a connection + if (cell.kind !== vscode.NotebookCellKind.Code || + cell.document.languageId !== 'sql' || + cell.metadata?.sqltools_connection) { + return; + } + + // If there's a last used connection, apply it to this cell + if (this._lastUsedConnection) { + this._setCellConnection(cell, this._lastUsedConnection); + } + } + + private async _selectConnection(): Promise { + try { + // Get available connections with proper parameters + // Changed to connectedOnly: false to show all available connections + const connections = await vscode.commands.executeCommand( + `${EXT_NAMESPACE}.getConnections`, + { connectedOnly: false, sort: 'connectedFirst' } + ) || []; + + if (connections.length === 0) { + vscode.window.showErrorMessage('No database connections available. Please add a database connection first.'); + return null; + } + + // Create connection items with proper ID extraction and more details + const connectionItems = connections.map(conn => { + // Extract connection ID and name + const id = typeof conn === 'string' ? conn : conn.id || conn.name; + const name = typeof conn === 'string' ? conn : conn.name || conn.id; + const driver = typeof conn === 'string' ? '' : conn.driver || ''; + const server = typeof conn === 'string' ? '' : conn.server || conn.host || ''; + const database = typeof conn === 'string' ? '' : conn.database || ''; + const isConnected = typeof conn === 'string' ? false : conn.isConnected || false; + + // Create a description that shows more connection details + let description = ''; + if (driver) description += driver; + if (server) description += description ? ` - ${server}` : server; + if (database) description += description ? ` - ${database}` : database; + if (isConnected) description += ' (Connected)'; + + return { + label: name, + description: description, + id: id, + isConnected: isConnected + }; + }); + + // Show connection picker with improved descriptions + const selectedConn = await vscode.window.showQuickPick(connectionItems, { + placeHolder: 'Select a database connection for this notebook', + title: 'SQLTools: Select Connection for Notebook', + }); + + if (selectedConn) { + console.log(`SQLTools Notebook: Selected connection ID: ${selectedConn.id}`); + + // If the connection is not already connected, connect to it + if (!selectedConn.isConnected) { + vscode.window.showInformationMessage(`Connecting to ${selectedConn.label}...`); + // Use the selectConnection command to connect to the database + await vscode.commands.executeCommand(`${EXT_NAMESPACE}.selectConnection`, selectedConn.id); + } + + vscode.window.showInformationMessage( + `Using connection: ${selectedConn.label}`, + { modal: false } + ); + return selectedConn.id; + } + + return null; + } catch (error) { + console.error('SQLTools Notebook: Error selecting connection:', error); + vscode.window.showErrorMessage(`Error selecting connection: ${error}`); + return null; + } + } + + dispose() { + console.log('SQLTools Notebook: Disposing controller'); + this._controller.dispose(); + } +} \ No newline at end of file diff --git a/packages/extension/src/notebook/index.ts b/packages/extension/src/notebook/index.ts new file mode 100644 index 000000000..0ac607e00 --- /dev/null +++ b/packages/extension/src/notebook/index.ts @@ -0,0 +1,37 @@ +import * as vscode from 'vscode'; +import { SQLNotebookSerializer } from './serializer'; +import { SQLNotebookController } from './controller'; + +let notebookController: SQLNotebookController | undefined; + +export function registerSQLNotebook(context: vscode.ExtensionContext): void { + // Register the notebook serializer + const notebookType = 'sqltools-notebook'; + const serializer = new SQLNotebookSerializer(); + + context.subscriptions.push( + vscode.workspace.registerNotebookSerializer(notebookType, serializer) + ); + + // Register the notebook controller + notebookController = new SQLNotebookController(); + + // Add to subscriptions for proper cleanup + context.subscriptions.push({ + dispose: () => { + if (notebookController) { + notebookController.dispose(); + notebookController = undefined; + } + } + }); + + console.log('SQLTools Notebook feature registered successfully'); +} + +export function deactivateSQLNotebook(): void { + if (notebookController) { + notebookController.dispose(); + notebookController = undefined; + } +} \ No newline at end of file diff --git a/packages/extension/src/notebook/serializer.ts b/packages/extension/src/notebook/serializer.ts new file mode 100644 index 000000000..1fae7eabd --- /dev/null +++ b/packages/extension/src/notebook/serializer.ts @@ -0,0 +1,72 @@ +import { TextDecoder, TextEncoder } from 'util'; +import * as vscode from 'vscode'; + +interface SQLNotebookCell { + kind: vscode.NotebookCellKind; + value: string; + language: string; +} + +interface SQLNotebook { + cells: SQLNotebookCell[]; +} + +export class SQLNotebookSerializer implements vscode.NotebookSerializer { + async deserializeNotebook( + content: Uint8Array, + _token: vscode.CancellationToken + ): Promise { + try { + const contents = new TextDecoder().decode(content); + + let raw: SQLNotebookCell[]; + try { + const parsed = JSON.parse(contents); + raw = parsed.cells || []; + console.log(`SQLTools Notebook: Deserialized notebook with ${raw.length} cells`); + } catch (parseError) { + console.error('SQLTools Notebook: Error parsing notebook content:', parseError); + // Return an empty notebook if we can't parse the content + raw = []; + } + + const cells = raw.map( + item => new vscode.NotebookCellData( + item.kind, + item.value, + item.language + ) + ); + + return new vscode.NotebookData(cells); + } catch (error) { + console.error('SQLTools Notebook: Error deserializing notebook:', error); + // Return an empty notebook on error + return new vscode.NotebookData([]); + } + } + + async serializeNotebook( + data: vscode.NotebookData, + _token: vscode.CancellationToken + ): Promise { + try { + const contents: SQLNotebookCell[] = []; + + for (const cell of data.cells) { + contents.push({ + kind: cell.kind, + value: cell.value, + language: cell.languageId + }); + } + + console.log(`SQLTools Notebook: Serialized notebook with ${contents.length} cells`); + return new TextEncoder().encode(JSON.stringify({ cells: contents })); + } catch (error) { + console.error('SQLTools Notebook: Error serializing notebook:', error); + // Return empty notebook data on error + return new TextEncoder().encode(JSON.stringify({ cells: [] })); + } + } +} \ No newline at end of file diff --git a/packages/plugins/connection-manager/extension.ts b/packages/plugins/connection-manager/extension.ts index 279473e80..6a154756d 100644 --- a/packages/plugins/connection-manager/extension.ts +++ b/packages/plugins/connection-manager/extension.ts @@ -255,7 +255,7 @@ export class ConnectionManagerPlugin implements IExtensionPlugin { return query; } - private ext_executeQuery = async (query?: string, { connNameOrId, connId, ...opt }: IQueryOptions = {}) => { + private ext_executeQuery = async (query?: string, { connNameOrId, connId, runInNotebook, ...opt }: IQueryOptions & { runInNotebook?: boolean } = {}) => { try { query = typeof query === 'string' ? query : await getSelectedText('execute query'); connNameOrId = connId || connNameOrId; @@ -267,7 +267,14 @@ export class ConnectionManagerPlugin implements IExtensionPlugin { connNameOrId = getAttachedConnection(window.activeTextEditor.document.uri); } - if (connNameOrId && connNameOrId.trim()) { + // Convert connNameOrId to string if it's an object + if (connNameOrId && typeof connNameOrId === 'object') { + connNameOrId = connNameOrId.id || connNameOrId.name; + log.info(`Connection identifier converted from object to: ${connNameOrId}`); + } + + // Now check if we have a valid string connection identifier + if (connNameOrId && typeof connNameOrId === 'string' && connNameOrId.trim()) { connNameOrId = connNameOrId.trim(); const conn = (await this.ext_getConnections({ connectedOnly: false, sort: 'connectedFirst' })).find(c => getConnectionId(c) === connNameOrId || c.name === connNameOrId); if (!conn) { @@ -281,10 +288,17 @@ export class ConnectionManagerPlugin implements IExtensionPlugin { const conn = await this.explorer.getActive() query = await this.replaceParams(query, conn); - const view = await this._openResultsWebview(conn && conn.id, opt.requestId); - const payload = await this._runConnectionCommandWithArgs('query', query, { ...opt, requestId: view.requestId }); - this.updateViewResults(view, payload); - return payload; + // If runInNotebook is true, skip opening the results webview + if (runInNotebook) { + // Just run the query and return the results directly + return await this._runConnectionCommandWithArgs('query', query, { ...opt }); + } else { + // Normal execution with results webview + const view = await this._openResultsWebview(conn && conn.id, opt.requestId); + const payload = await this._runConnectionCommandWithArgs('query', query, { ...opt, requestId: view.requestId }); + this.updateViewResults(view, payload); + return payload; + } } catch (e) { this.errorHandler('Error fetching records.', e); }