Enable VS Code notebooks with a built-in SQL kernel. (#21995)

extensions/mssql/package.json

@@ -67,6 +67,11 @@
         "category": "MSSQL",
         "title": "%title.designTable%"
       },
+      {
+        "command": "mssql.changeNotebookConnection",
+        "category": "MSSQL",
+        "title": "%title.changeNotebookConnection%"
+      },
       {
         "command": "mssql.newLogin",
         "category": "MSSQL",
@@ -418,6 +423,10 @@
           "command": "mssql.designTable",
           "when": "false"
         },
+        {
+          "command": "mssql.changeNotebookConnection",
+          "when": "false"
+        },
         {
           "command": "mssql.newServerRole",
           "when": "false"

extensions/mssql/package.nls.json

@@ -176,6 +176,7 @@
 	"objectsListProperties.name": "Name",
 	"title.newTable": "New Table",
 	"title.designTable": "Design",
+	"title.changeNotebookConnection": "Change SQL Notebook Connection",
 	"mssql.parallelMessageProcessing" : "[Experimental] Whether the requests to the SQL Tools Service should be handled in parallel. This is introduced to discover the issues there might be when handling all requests in parallel. The default value is false. Relaunch of ADS is required when the value is changed.",
 	"mssql.tableDesigner.preloadDatabaseModel": "Whether to preload the database model when the database node in the object explorer is expanded. When enabled, the loading time of table designer can be reduced. Note: You might see higher than normal memory usage if you need to expand a lot of database nodes.",
 	"mssql.objectExplorer.groupBySchema": "When enabled, the database objects in Object Explorer will be categorized by schema.",

extensions/mssql/src/main.ts

@@ -22,6 +22,7 @@ import { IconPathHelper } from './iconHelper';
 import * as nls from 'vscode-nls';
 import { INotebookConvertService } from './notebookConvert/notebookConvertService';
 import { registerTableDesignerCommands } from './tableDesigner/tableDesigner';
+import { SqlNotebookController } from './sqlNotebook/sqlNotebookController';
 import { registerObjectManagementCommands } from './objectManagement/commands';
 import { TelemetryActions, TelemetryReporter, TelemetryViews } from './telemetry';
 
@@ -114,7 +115,10 @@ export async function activate(context: vscode.ExtensionContext): Promise<IExten
 	registerTableDesignerCommands(appContext);
 	registerObjectManagementCommands(appContext);
 
+	context.subscriptions.push(new SqlNotebookController());
+
 	context.subscriptions.push(TelemetryReporter);
+
 	return createMssqlApi(appContext, server);
 }
 

extensions/mssql/src/sqlNotebook/sqlNotebookController.ts

@@ -0,0 +1,329 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+import * as azdata from 'azdata';
+import * as nls from 'vscode-nls';
+
+const localize = nls.loadMessageBundle();
+
+interface QueryCompletionHandler {
+	ownerUri: string;
+	handler: (results: azdata.BatchSummary[]) => void
+}
+
+interface QueryMessageHandler {
+	ownerUri: string;
+	handler: (results: azdata.QueryExecuteMessageParams) => void
+}
+
+export class SqlNotebookController implements vscode.Disposable {
+	private readonly _cellUriScheme = 'vscode-notebook-cell';
+	private readonly _connectionLabel = (serverName: string) => localize('notebookConnection', 'Connected to: {0}', serverName);
+	private readonly _disconnectedLabel = localize('notebookDisconnected', 'Disconnected');
+
+	private readonly _disposables = new Array<vscode.Disposable>();
+	private readonly _controller: vscode.NotebookController;
+	private readonly _connectionsMap = new Map<vscode.Uri, azdata.connection.Connection>();
+	private readonly _executionOrderMap = new Map<vscode.Uri, number>();
+	private readonly _queryProvider: azdata.QueryProvider;
+	private readonly _connProvider: azdata.ConnectionProvider;
+	private readonly _connectionLabelItem: vscode.StatusBarItem;
+
+	private _queryCompleteHandler: QueryCompletionHandler;
+	private _queryMessageHandler: QueryMessageHandler;
+	private _activeCellUri: string;
+
+	constructor() {
+		this._controller = vscode.notebooks.createNotebookController('sql-controller-id', 'jupyter-notebook', 'SQL');
+
+		this._controller.supportedLanguages = ['sql'];
+		this._controller.supportsExecutionOrder = true;
+		this._controller.executeHandler = this.execute.bind(this);
+
+		const sqlProvider = 'MSSQL';
+		this._queryProvider = azdata.dataprotocol.getProvider<azdata.QueryProvider>(sqlProvider, azdata.DataProviderType.QueryProvider);
+		this._queryProvider.registerOnQueryComplete(result => this.handleQueryComplete(result));
+		this._queryProvider.registerOnMessage(message => this.handleQueryMessage(message));
+
+		this._connProvider = azdata.dataprotocol.getProvider<azdata.ConnectionProvider>(sqlProvider, azdata.DataProviderType.ConnectionProvider);
+
+		const commandName = 'mssql.changeNotebookConnection';
+		let changeConnectionCommand = vscode.commands.registerCommand(commandName, async () => await this.changeConnection());
+		this._disposables.push(changeConnectionCommand);
+
+		this._connectionLabelItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Right);
+		this._connectionLabelItem.text = this._disconnectedLabel;
+		this._connectionLabelItem.tooltip = localize('changeNotebookConnection', 'Change SQL Notebook Connection');
+		this._connectionLabelItem.command = commandName;
+		this._disposables.push(this._connectionLabelItem);
+
+		// Show connection status if there's a notebook already open when ADS starts
+		if (vscode.window.activeTextEditor?.document.notebook) {
+			this._connectionLabelItem.show();
+		}
+
+		let editorChangedEvent = vscode.window.onDidChangeActiveTextEditor(async editor => await this.handleActiveEditorChanged(editor));
+		this._disposables.push(editorChangedEvent);
+
+		let docClosedEvent = vscode.workspace.onDidCloseTextDocument(document => this.handleDocumentClosed(document));
+		this._disposables.push(docClosedEvent);
+	}
+
+	private handleQueryComplete(result: azdata.QueryExecuteCompleteNotificationResult): void {
+		if (this._queryCompleteHandler && this._queryCompleteHandler.ownerUri === result.ownerUri) { // Check if handler is undefined separately in case the result URI is also undefined
+			this._queryCompleteHandler.handler(result.batchSummaries);
+		}
+	}
+
+	private handleQueryMessage(message: azdata.QueryExecuteMessageParams): void {
+		if (this._queryMessageHandler && this._queryMessageHandler.ownerUri === message.ownerUri) { // Check if handler is undefined separately in case the result URI is also undefined
+			this._queryMessageHandler.handler(message);
+		}
+	}
+
+	private async handleActiveEditorChanged(editor: vscode.TextEditor): Promise<void> {
+		let notebook = editor?.document.notebook;
+		if (!notebook) {
+			// Hide status bar item if the current editor isn't a notebook
+			this._connectionLabelItem.hide();
+		} else {
+			let connection = this._connectionsMap.get(notebook.uri);
+			if (connection) {
+				this._connectionLabelItem.text = this._connectionLabel(connection.options['server']);
+
+				// If this editor is for a cell, then update the connection for it
+				this.updateCellConnection(notebook.uri, connection);
+			} else {
+				this._connectionLabelItem.text = this._disconnectedLabel;
+			}
+			this._connectionLabelItem.show();
+		}
+	}
+
+	public getConnectionProfile(connection: azdata.connection.Connection): azdata.IConnectionProfile {
+		let connectionProfile: azdata.IConnectionProfile = {
+			connectionName: connection.options.connectionName,
+			serverName: connection.options.server,
+			databaseName: connection.options.database,
+			userName: connection.options.user,
+			password: connection.options.password,
+			authenticationType: connection.options.authenticationType,
+			savePassword: connection.options.savePassword,
+			groupFullName: undefined,
+			groupId: undefined,
+			providerName: connection.providerName,
+			saveProfile: false,
+			id: connection.connectionId,
+			options: connection.options
+		};
+		return connectionProfile;
+	}
+
+	private handleDocumentClosed(editor: vscode.TextDocument): void {
+		// Have to check isClosed here since this event is also emitted on doc language changes
+		if (editor.notebook && editor.isClosed) {
+			// Remove the connection & execution associations if the doc is closed, but don't close the connection since it might be re-used elsewhere
+			this._connectionsMap.delete(editor.notebook.uri);
+			this._executionOrderMap.delete(editor.notebook.uri);
+		}
+	}
+
+	private updateCellConnection(notebookUri: vscode.Uri, connection: azdata.connection.Connection): void {
+		let docUri = vscode.window.activeTextEditor?.document.uri;
+		if (docUri?.scheme === this._cellUriScheme && docUri?.path === notebookUri.path) {
+			if (this._activeCellUri) {
+				this._connProvider.disconnect(this._activeCellUri).then(() => undefined, error => console.log(error));
+			}
+			this._activeCellUri = docUri.toString();
+			// Delay connecting in case user is clicking between cells a lot
+			setTimeout(() => {
+				if (this._activeCellUri === docUri.toString()) {
+					let profile = this.getConnectionProfile(connection);
+					this._connProvider.connect(docUri.toString(), profile).then(
+						connected => {
+							if (!connected) {
+								console.log(`Failed to update cell connection for cell: ${docUri.toString()}`);
+							}
+						},
+						error => {
+							console.log(error);
+						});
+				}
+			}, 200);
+		}
+	}
+
+	private async changeConnection(notebook?: vscode.NotebookDocument): Promise<azdata.connection.Connection | undefined> {
+		let connection: azdata.connection.Connection;
+		let notebookUri = notebook?.uri ?? vscode.window.activeTextEditor?.document.notebook?.uri;
+		if (notebookUri) {
+			connection = await azdata.connection.openConnectionDialog(['MSSQL']);
+			if (connection) {
+				this._connectionsMap.set(notebookUri, connection);
+				this._connectionLabelItem.text = this._connectionLabel(connection.options['server']);
+
+				// Connect current notebook cell, if there is one
+				this.updateCellConnection(notebookUri, connection);
+			} else {
+				this._connectionLabelItem.text = this._disconnectedLabel;
+			}
+			this._connectionLabelItem.show();
+		}
+		return connection;
+	}
+
+	private async execute(cells: vscode.NotebookCell[], notebook: vscode.NotebookDocument, controller: vscode.NotebookController): Promise<void> {
+		if (this._queryCompleteHandler) {
+			throw new Error(localize('queryInProgressError', 'Another query is currently in progress. Please wait for that query to complete before running these cells.'));
+		}
+
+		let connection = this._connectionsMap.get(notebook.uri);
+		if (!connection) {
+			connection = await this.changeConnection(notebook);
+		}
+
+		let executionOrder = this._executionOrderMap.get(notebook.uri) ?? 0;
+		for (let cell of cells) {
+			await this.doExecution(cell, connection, ++executionOrder);
+		}
+		this._executionOrderMap.set(notebook.uri, executionOrder);
+	}
+
+	private async doExecution(cell: vscode.NotebookCell, connection: azdata.connection.Connection | undefined, executionOrder: number): Promise<void> {
+		const execution = this._controller.createNotebookCellExecution(cell);
+		execution.executionOrder = executionOrder;
+		execution.start(Date.now());
+		await execution.clearOutput();
+		if (!connection) {
+			await execution.appendOutput([
+				new vscode.NotebookCellOutput([
+					vscode.NotebookCellOutputItem.text(localize('noConnectionError', 'No connection provided.'))
+				])
+			]);
+			execution.end(false, Date.now());
+			return;
+		}
+
+		let cancelHandler: vscode.Disposable;
+		try {
+			const ownerUri = await azdata.connection.getUriForConnection(connection.connectionId);
+			await this._queryProvider.runQueryString(ownerUri, cell.document.getText());
+			cancelHandler = execution.token.onCancellationRequested(async () => await this._queryProvider.cancelQuery(ownerUri));
+
+			let queryComplete = new Promise<void>(resolve => {
+				let queryCompleteHandler = async (batchSummaries: azdata.BatchSummary[]) => {
+					let tableHtmlEntries: string[] = [];
+					for (let batchSummary of batchSummaries) {
+						if (execution.token.isCancellationRequested) {
+							break;
+						}
+
+						for (let resultSummary of batchSummary.resultSetSummaries) {
+							if (execution.token.isCancellationRequested) {
+								break;
+							}
+
+							if (resultSummary.rowCount > 0) {
+								// Add column headers
+								let tableHtml =
+									`<style>
+										.output_container .sqlNotebookResults td, .output_container .sqlNotebookResults th {
+											text-align: left;
+										}
+									</style>
+									<table class="sqlNotebookResults"><thead><tr>`;
+								for (let column of resultSummary.columnInfo) {
+									tableHtml += `<th>${htmlEscape(column.columnName)}</th>`;
+								}
+								tableHtml += '</tr></thead>';
+
+								// Add rows and cells
+								let subsetResult = await this._queryProvider.getQueryRows({
+									ownerUri: ownerUri,
+									batchIndex: batchSummary.id,
+									resultSetIndex: resultSummary.id,
+									rowsStartIndex: 0,
+									rowsCount: resultSummary.rowCount
+								});
+								tableHtml += '<tbody>';
+								for (let row of subsetResult.resultSubset.rows) {
+									tableHtml += '<tr>';
+									for (let cell of row) {
+										tableHtml += `<td>${htmlEscape(cell.displayValue)}</td>`;
+									}
+									tableHtml += '</tr>';
+								}
+								tableHtml += '</tbody></table>';
+								tableHtmlEntries.push(tableHtml);
+							}
+						}
+					}
+
+					if (execution.token.isCancellationRequested) {
+						await execution.appendOutput([
+							new vscode.NotebookCellOutput([
+								vscode.NotebookCellOutputItem.text(localize('cellExecutionCancelled', 'Cell execution cancelled.'))
+							])
+						]);
+						execution.end(false, Date.now());
+					} else {
+						await execution.appendOutput([
+							new vscode.NotebookCellOutput([
+								vscode.NotebookCellOutputItem.text(tableHtmlEntries.join('<br><br>'), 'text/html')
+							])
+						]);
+						execution.end(true, Date.now());
+					}
+					resolve();
+				};
+				this._queryCompleteHandler = { ownerUri: ownerUri, handler: queryCompleteHandler };
+			});
+
+			this._queryMessageHandler = {
+				ownerUri: ownerUri,
+				handler: async message => {
+					await execution.appendOutput([
+						new vscode.NotebookCellOutput([
+							vscode.NotebookCellOutputItem.text(message.message.message)
+						])
+					]);
+				}
+			};
+
+			await queryComplete;
+		} catch (error) {
+			await execution.appendOutput([
+				new vscode.NotebookCellOutput([
+					vscode.NotebookCellOutputItem.error(error)
+				])
+			]);
+			execution.end(false, Date.now());
+		} finally {
+			if (cancelHandler) {
+				cancelHandler.dispose();
+			}
+			this._queryCompleteHandler = undefined;
+			this._queryMessageHandler = undefined;
+		}
+	}
+
+	dispose() {
+		this._disposables.forEach(d => d.dispose());
+	}
+}
+
+function htmlEscape(html: string): string {
+	return html.replace(/[<|>|&|"]/g, function (match) {
+		switch (match) {
+			case '<': return '&lt;';
+			case '>': return '&gt;';
+			case '&': return '&amp;';
+			case '"': return '&quot;';
+			case '\'': return '&#39';
+			default: return match;
+		}
+	});
+}

extensions/mssql/src/typings/refs.d.ts

@@ -6,4 +6,5 @@
 /// <reference path='../../../../src/sql/azdata.d.ts'/>
 /// <reference path='../../../../src/sql/azdata.proposed.d.ts'/>
 /// <reference path='../../../../src/vscode-dts/vscode.d.ts'/>
+/// <reference path='../../../../src/vscode-dts/vscode.proposed.textDocumentNotebook.d.ts' />
 /// <reference path='../../../azurecore/src/azurecore.d.ts' />