Add support for VS Code notebook Open & Close events in extension APIs (#17992)

src/sql/workbench/api/common/adsNotebookController.ts

@@ -11,9 +11,9 @@ import * as extHostTypeConverters from 'vs/workbench/api/common/extHostTypeConve
 import { Deferred } from 'sql/base/common/promise';
 import { ExtHostNotebookDocumentsAndEditors } from 'sql/workbench/api/common/extHostNotebookDocumentsAndEditors';
 import { URI } from 'vs/base/common/uri';
-import { VSCodeContentManager } from 'sql/workbench/api/common/vscodeSerializationProvider';
 import { NotebookCellExecutionTaskState } from 'vs/workbench/api/common/extHostNotebookKernels';
 import { asArray } from 'vs/base/common/arrays';
+import { convertToADSCellOutput } from 'sql/workbench/api/common/vscodeSerializationProvider';
 
 type SelectionChangedEvent = { selected: boolean, notebook: vscode.NotebookDocument; };
 type MessageReceivedEvent = { editor: vscode.NotebookEditor, message: any; };
@@ -224,7 +224,7 @@ class ADSNotebookCellExecution implements vscode.NotebookCellExecution {
 		const targetCell = typeof cell === 'number' ? this._cell.notebook.cellAt(cell) : (cell ?? this._cell);
 		const editor = this._extHostNotebookDocumentsAndEditors.getEditor(URI.from(targetCell.notebook.uri).toString());
 		await editor.edit(builder => {
-			const adsOutputs = VSCodeContentManager.convertToADSCellOutput(outputs);
+			const adsOutputs = convertToADSCellOutput(outputs);
 			builder.updateCell(targetCell.index, { outputs: adsOutputs }, append);
 		});
 	}
@@ -233,7 +233,7 @@ class ADSNotebookCellExecution implements vscode.NotebookCellExecution {
 		this.verifyStateForOutput();
 		const editor = this._extHostNotebookDocumentsAndEditors.getEditor(URI.from(this._cell.notebook.uri).toString());
 		await editor.edit(builder => {
-			const adsOutput = VSCodeContentManager.convertToADSCellOutput({ id: output.id, items: asArray(items) }, undefined);
+			const adsOutput = convertToADSCellOutput({ id: output.id, items: asArray(items) }, undefined);
 			builder.updateCellOutput(this._cell.index, { outputs: adsOutput }, append);
 		});
 	}

src/sql/workbench/api/common/extHostNotebookDocumentsAndEditors.ts

@@ -21,6 +21,7 @@ import {
 } from 'sql/workbench/api/common/sqlExtHost.protocol';
 import { ExtHostNotebookDocumentData } from 'sql/workbench/api/common/extHostNotebookDocumentData';
 import { ExtHostNotebookEditor } from 'sql/workbench/api/common/extHostNotebookEditor';
+import { VSCodeNotebookDocument } from 'sql/workbench/api/common/vscodeNotebookDocument';
 
 type Adapter = azdata.nb.NavigationProvider;
 
@@ -39,13 +40,19 @@ export class ExtHostNotebookDocumentsAndEditors implements ExtHostNotebookDocume
 	private readonly _onDidChangeVisibleNotebookEditors = new Emitter<ExtHostNotebookEditor[]>();
 	private readonly _onDidChangeActiveNotebookEditor = new Emitter<ExtHostNotebookEditor>();
 	private _onDidOpenNotebook = new Emitter<azdata.nb.NotebookDocument>();
+	private _onDidCloseNotebook = new Emitter<azdata.nb.NotebookDocument>();
 	private _onDidChangeNotebookCell = new Emitter<azdata.nb.NotebookCellChangeEvent>();
 
 	readonly onDidChangeVisibleNotebookEditors: Event<ExtHostNotebookEditor[]> = this._onDidChangeVisibleNotebookEditors.event;
 	readonly onDidChangeActiveNotebookEditor: Event<ExtHostNotebookEditor> = this._onDidChangeActiveNotebookEditor.event;
 	readonly onDidOpenNotebookDocument: Event<azdata.nb.NotebookDocument> = this._onDidOpenNotebook.event;
+	readonly onDidCloseNotebookDocument: Event<azdata.nb.NotebookDocument> = this._onDidCloseNotebook.event;
 	readonly onDidChangeNotebookCell: Event<azdata.nb.NotebookCellChangeEvent> = this._onDidChangeNotebookCell.event;
 
+	private _onDidOpenVSCodeNotebook = new Emitter<vscode.NotebookDocument>();
+	private _onDidCloseVSCodeNotebook = new Emitter<vscode.NotebookDocument>();
+	readonly onDidOpenVSCodeNotebookDocument: Event<vscode.NotebookDocument> = this._onDidOpenVSCodeNotebook.event;
+	readonly onDidCloseVSCodeNotebookDocument: Event<vscode.NotebookDocument> = this._onDidCloseVSCodeNotebook.event;
 
 	constructor(
 		private readonly _mainContext: IMainContext,
@@ -53,6 +60,9 @@ export class ExtHostNotebookDocumentsAndEditors implements ExtHostNotebookDocume
 		if (this._mainContext) {
 			this._proxy = this._mainContext.getProxy(SqlMainContext.MainThreadNotebookDocumentsAndEditors);
 		}
+
+		this.onDidOpenNotebookDocument(notebook => this._onDidOpenVSCodeNotebook.fire(new VSCodeNotebookDocument(notebook)));
+		this.onDidCloseNotebookDocument(notebook => this._onDidCloseVSCodeNotebook.fire(new VSCodeNotebookDocument(notebook)));
 	}
 
 	dispose() {
@@ -128,7 +138,7 @@ export class ExtHostNotebookDocumentsAndEditors implements ExtHostNotebookDocume
 
 		// now that the internal state is complete, fire events
 		if (removedDocuments) {
-			// TODO add doc close event
+			removedDocuments.forEach(d => this._onDidCloseNotebook.fire(d.document));
 		}
 		if (addedDocuments) {
 			addedDocuments.forEach(d => this._onDidOpenNotebook.fire(d.document));

src/sql/workbench/api/common/sqlExtHost.api.impl.ts

@@ -49,16 +49,17 @@ export interface IExtensionApiFactory {
 export interface IAdsExtensionApiFactory {
 	azdata: IAzdataExtensionApiFactory;
 	extHostNotebook: ExtHostNotebook;
+	extHostNotebookDocumentsAndEditors: ExtHostNotebookDocumentsAndEditors;
 }
 
 /**
  * This method instantiates and returns the extension API surface
  */
 export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): IExtensionApiFactory {
-	const { azdata, extHostNotebook } = createAdsApiFactory(accessor);
+	const { azdata, extHostNotebook, extHostNotebookDocumentsAndEditors } = createAdsApiFactory(accessor);
 	return {
 		azdata,
-		vscode: vsApiFactory(accessor, extHostNotebook)
+		vscode: vsApiFactory(accessor, extHostNotebook, extHostNotebookDocumentsAndEditors)
 	};
 }
 
@@ -538,6 +539,9 @@ export function createAdsApiFactory(accessor: ServicesAccessor): IAdsExtensionAp
 				get onDidOpenNotebookDocument() {
 					return extHostNotebookDocumentsAndEditors.onDidOpenNotebookDocument;
 				},
+				get onDidCloseNotebookDocument() {
+					return extHostNotebookDocumentsAndEditors.onDidCloseNotebookDocument;
+				},
 				get onDidChangeActiveNotebookEditor() {
 					return extHostNotebookDocumentsAndEditors.onDidChangeActiveNotebookEditor;
 				},
@@ -631,6 +635,7 @@ export function createAdsApiFactory(accessor: ServicesAccessor): IAdsExtensionAp
 				designers: designers
 			};
 		},
-		extHostNotebook: extHostNotebook
+		extHostNotebook: extHostNotebook,
+		extHostNotebookDocumentsAndEditors: extHostNotebookDocumentsAndEditors
 	};
 }

src/sql/workbench/api/common/vscodeExecuteProvider.ts

@@ -7,6 +7,7 @@ import type * as vscode from 'vscode';
 import type * as azdata from 'azdata';
 import { ADSNotebookController } from 'sql/workbench/api/common/adsNotebookController';
 import * as nls from 'vs/nls';
+import { URI } from 'vs/base/common/uri';
 
 class VSCodeFuture implements azdata.nb.IFuture {
 	private _inProgress = true;
@@ -129,18 +130,7 @@ class VSCodeKernel implements azdata.nb.IKernel {
 	requestExecute(content: azdata.nb.IExecuteRequest, disposeOnDone?: boolean): azdata.nb.IFuture {
 		let executePromise: Promise<void>;
 		if (this._controller.executeHandler) {
-			let cell = <vscode.NotebookCell>{
-				index: content.cellIndex,
-				document: <vscode.TextDocument>{
-					uri: content.notebookUri,
-					languageId: this._kernelSpec.language,
-					getText: () => Array.isArray(content.code) ? content.code.join('') : content.code,
-				},
-				notebook: <vscode.NotebookDocument>{
-					uri: content.notebookUri
-				}
-			};
-
+			let cell = convertToVSCodeNotebookCell(content.code, content.cellIndex, content.notebookUri, this._kernelSpec.language);
 			executePromise = Promise.resolve(this._controller.executeHandler([cell], cell.notebook, this._controller));
 		}
 		else {
@@ -315,3 +305,17 @@ export class VSCodeExecuteProvider implements azdata.nb.NotebookExecuteProvider
 		// No-op
 	}
 }
+
+export function convertToVSCodeNotebookCell(cellSource: string | string[], index: number, uri: URI, language: string): vscode.NotebookCell {
+	return <vscode.NotebookCell>{
+		index: index,
+		document: <vscode.TextDocument>{
+			uri: uri,
+			languageId: language,
+			getText: () => Array.isArray(cellSource) ? cellSource.join('') : cellSource,
+		},
+		notebook: <vscode.NotebookDocument>{
+			uri: uri
+		}
+	};
+}

src/sql/workbench/api/common/vscodeNotebookDocument.ts

@@ -0,0 +1,58 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import type * as vscode from 'vscode';
+import type * as azdata from 'azdata';
+import { convertToVSCodeNotebookCell } from 'sql/workbench/api/common/vscodeExecuteProvider';
+
+export class VSCodeNotebookDocument implements vscode.NotebookDocument {
+	private readonly _convertedCells: vscode.NotebookCell[];
+
+	constructor(private readonly _notebookDoc: azdata.nb.NotebookDocument) {
+		this._convertedCells = this._notebookDoc.cells?.map((cell, index) => convertToVSCodeNotebookCell(cell.contents.source, index, this._notebookDoc.uri, this._notebookDoc.kernelSpec?.language));
+	}
+
+	public get uri() { return this._notebookDoc.uri; }
+
+	public get version() { return undefined; }
+
+	public get notebookType() { return this._notebookDoc.providerId; }
+
+	public get isDirty() { return this._notebookDoc.isDirty; }
+
+	public get isUntitled() { return this._notebookDoc.isUntitled; }
+
+	public get isClosed() { return this._notebookDoc.isClosed; }
+
+	public get metadata() { return {}; }
+
+	public get cellCount() { return this._notebookDoc.cells?.length; }
+
+	cellAt(index: number): vscode.NotebookCell {
+		if (this._notebookDoc.cells) {
+			if (index < 0) {
+				index = 0;
+			} else if (index >= this._notebookDoc.cells.length) {
+				index = this._convertedCells.length - 1;
+			}
+			return this._convertedCells[index];
+		}
+		return undefined;
+	}
+
+	getCells(range?: vscode.NotebookRange): vscode.NotebookCell[] {
+		let cells: vscode.NotebookCell[] = [];
+		if (range) {
+			cells = this._convertedCells?.slice(range.start, range.end);
+		} else {
+			cells = this._convertedCells;
+		}
+		return cells;
+	}
+
+	save(): Thenable<boolean> {
+		return this._notebookDoc.save();
+	}
+}

src/sql/workbench/api/common/vscodeSerializationProvider.ts

@@ -16,22 +16,6 @@ export class VSCodeContentManager implements azdata.nb.ContentManager {
 	constructor(private readonly _serializer: vscode.NotebookSerializer) {
 	}
 
-	public static convertToADSCellOutput(outputs: vscode.NotebookCellOutput | vscode.NotebookCellOutput[], executionOrder?: number): azdata.nb.IDisplayResult[] {
-		return asArray(outputs).map(output => {
-			let outputData = {};
-			for (let item of output.items) {
-				outputData[item.mime] = VSBuffer.wrap(item.data).toString();
-			}
-			return {
-				output_type: 'execute_result',
-				data: outputData,
-				execution_count: executionOrder,
-				metadata: output.metadata,
-				id: output.id
-			};
-		});
-	}
-
 	public async deserializeNotebook(contents: string): Promise<azdata.nb.INotebookContents> {
 		let buffer = VSBuffer.fromString(contents);
 		let notebookData = await this._serializer.deserializeNotebook(buffer.buffer, new CancellationTokenSource().token);
@@ -45,7 +29,7 @@ export class VSCodeContentManager implements azdata.nb.ContentManager {
 						language: cell.languageId
 					},
 					execution_count: executionOrder,
-					outputs: cell.outputs ? VSCodeContentManager.convertToADSCellOutput(cell.outputs, executionOrder) : undefined
+					outputs: cell.outputs ? convertToADSCellOutput(cell.outputs, executionOrder) : undefined
 				};
 			}),
 			metadata: notebookData.metadata ?? {},
@@ -59,43 +43,6 @@ export class VSCodeContentManager implements azdata.nb.ContentManager {
 		return result;
 	}
 
-	public static convertToVSCodeCellOutput(output: azdata.nb.ICellOutput): vscode.NotebookCellOutput {
-		let convertedOutputItems: vscode.NotebookCellOutputItem[];
-		switch (output.output_type) {
-			case OutputTypes.ExecuteResult:
-			case OutputTypes.DisplayData:
-			case OutputTypes.UpdateDisplayData:
-				let displayOutput = output as azdata.nb.IDisplayResult;
-				convertedOutputItems = Object.keys(displayOutput.data).map<vscode.NotebookCellOutputItem>(key => {
-					return {
-						mime: key,
-						data: VSBuffer.fromString(displayOutput.data[key]).buffer
-					};
-				});
-				break;
-			case OutputTypes.Stream:
-				let streamOutput = output as azdata.nb.IStreamResult;
-				convertedOutputItems = [{
-					mime: 'text/html',
-					data: VSBuffer.fromString(Array.isArray(streamOutput.text) ? streamOutput.text.join('') : streamOutput.text).buffer
-				}];
-				break;
-			case OutputTypes.Error:
-				let errorOutput = output as azdata.nb.IErrorResult;
-				let errorString = errorOutput.ename + ': ' + errorOutput.evalue + (errorOutput.traceback ? '\n' + errorOutput.traceback?.join('\n') : '');
-				convertedOutputItems = [{
-					mime: 'text/html',
-					data: VSBuffer.fromString(errorString).buffer
-				}];
-				break;
-		}
-		return {
-			items: convertedOutputItems,
-			metadata: output.metadata,
-			id: output.id
-		};
-	}
-
 	public async serializeNotebook(notebook: azdata.nb.INotebookContents): Promise<string> {
 		let notebookData: vscode.NotebookData = {
 			cells: notebook.cells?.map<vscode.NotebookCellData>(cell => {
@@ -103,7 +50,7 @@ export class VSCodeContentManager implements azdata.nb.ContentManager {
 					kind: cell.cell_type === 'code' ? NotebookCellKind.Code : NotebookCellKind.Markup,
 					value: Array.isArray(cell.source) ? cell.source.join('\n') : cell.source,
 					languageId: cell.metadata?.language,
-					outputs: cell.outputs?.map<vscode.NotebookCellOutput>(output => VSCodeContentManager.convertToVSCodeCellOutput(output)),
+					outputs: cell.outputs?.map<vscode.NotebookCellOutput>(output => convertToVSCodeCellOutput(output)),
 					executionSummary: {
 						executionOrder: cell.execution_count
 					}
@@ -148,3 +95,56 @@ export class VSCodeSerializationProvider implements azdata.nb.NotebookSerializat
 		return Promise.resolve(this._manager);
 	}
 }
+
+export function convertToADSCellOutput(outputs: vscode.NotebookCellOutput | vscode.NotebookCellOutput[], executionOrder?: number): azdata.nb.IDisplayResult[] {
+	return asArray(outputs).map(output => {
+		let outputData = {};
+		for (let item of output.items) {
+			outputData[item.mime] = VSBuffer.wrap(item.data).toString();
+		}
+		return {
+			output_type: 'execute_result',
+			data: outputData,
+			execution_count: executionOrder,
+			metadata: output.metadata,
+			id: output.id
+		};
+	});
+}
+
+export function convertToVSCodeCellOutput(output: azdata.nb.ICellOutput): vscode.NotebookCellOutput {
+	let convertedOutputItems: vscode.NotebookCellOutputItem[];
+	switch (output.output_type) {
+		case OutputTypes.ExecuteResult:
+		case OutputTypes.DisplayData:
+		case OutputTypes.UpdateDisplayData:
+			let displayOutput = output as azdata.nb.IDisplayResult;
+			convertedOutputItems = Object.keys(displayOutput.data).map<vscode.NotebookCellOutputItem>(key => {
+				return {
+					mime: key,
+					data: VSBuffer.fromString(displayOutput.data[key]).buffer
+				};
+			});
+			break;
+		case OutputTypes.Stream:
+			let streamOutput = output as azdata.nb.IStreamResult;
+			convertedOutputItems = [{
+				mime: 'text/html',
+				data: VSBuffer.fromString(Array.isArray(streamOutput.text) ? streamOutput.text.join('') : streamOutput.text).buffer
+			}];
+			break;
+		case OutputTypes.Error:
+			let errorOutput = output as azdata.nb.IErrorResult;
+			let errorString = errorOutput.ename + ': ' + errorOutput.evalue + (errorOutput.traceback ? '\n' + errorOutput.traceback?.join('\n') : '');
+			convertedOutputItems = [{
+				mime: 'text/html',
+				data: VSBuffer.fromString(errorString).buffer
+			}];
+			break;
+	}
+	return {
+		items: convertedOutputItems,
+		metadata: output.metadata,
+		id: output.id
+	};
+}