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

src/sql/azdata.proposed.d.ts

@@ -92,6 +92,11 @@ declare module 'azdata' {
 			 */
 			dispose(): void;
 		}
+
+		/**
+		 * An event that is emitted when a [notebook document](#NotebookDocument) is closed.
+		 */
+		export const onDidCloseNotebookDocument: vscode.Event<NotebookDocument>;
 	}
 
 	export type SqlDbType = 'BigInt' | 'Binary' | 'Bit' | 'Char' | 'DateTime' | 'Decimal'

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>{
+			let cell = convertToVSCodeNotebookCell(content.code, content.cellIndex, content.notebookUri, this._kernelSpec.language);
-				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
-				}
-			};
-
 			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
+	};
+}

src/sql/workbench/test/electron-browser/api/vscodeNotebookApi.test.ts

@@ -3,7 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { VSCodeContentManager } from 'sql/workbench/api/common/vscodeSerializationProvider';
+import { convertToVSCodeCellOutput, VSCodeContentManager, convertToADSCellOutput } from 'sql/workbench/api/common/vscodeSerializationProvider';
 import type * as vscode from 'vscode';
 import type * as azdata from 'azdata';
 import * as sinon from 'sinon';
@@ -12,6 +12,9 @@ import { VSBuffer } from 'vs/base/common/buffer';
 import * as assert from 'assert';
 import { OutputTypes } from 'sql/workbench/services/notebook/common/contracts';
 import { NBFORMAT, NBFORMAT_MINOR } from 'sql/workbench/common/constants';
+import { convertToVSCodeNotebookCell } from 'sql/workbench/api/common/vscodeExecuteProvider';
+import { VSCodeNotebookDocument } from 'sql/workbench/api/common/vscodeNotebookDocument';
+import { URI } from 'vs/base/common/uri';
 
 class MockNotebookSerializer implements vscode.NotebookSerializer {
 	deserializeNotebook(content: Uint8Array, token: vscode.CancellationToken): vscode.NotebookData | Thenable<vscode.NotebookData> {
@@ -235,7 +238,7 @@ suite('Notebook Serializer', () => {
 			}
 		];
 
-		let actualOutput = VSCodeContentManager.convertToADSCellOutput(cellOutput, 1);
+		let actualOutput = convertToADSCellOutput(cellOutput, 1);
 		assert.deepStrictEqual(actualOutput, expectedADSOutput);
 	});
 
@@ -260,7 +263,7 @@ suite('Notebook Serializer', () => {
 			id: 'testId',
 			metadata: undefined
 		};
-		let actualOutput = VSCodeContentManager.convertToVSCodeCellOutput(cellOutput);
+		let actualOutput = convertToVSCodeCellOutput(cellOutput);
 		assert.deepStrictEqual(actualOutput, expectedVSCodeOutput);
 	});
 
@@ -281,7 +284,7 @@ suite('Notebook Serializer', () => {
 			id: 'testId',
 			metadata: undefined
 		};
-		let actualOutput = VSCodeContentManager.convertToVSCodeCellOutput(cellOutput);
+		let actualOutput = convertToVSCodeCellOutput(cellOutput);
 		assert.deepStrictEqual(actualOutput, expectedVSCodeOutput);
 	});
 
@@ -301,7 +304,7 @@ suite('Notebook Serializer', () => {
 			id: 'testId',
 			metadata: undefined
 		};
-		let actualOutput = VSCodeContentManager.convertToVSCodeCellOutput(cellOutput);
+		let actualOutput = convertToVSCodeCellOutput(cellOutput);
 		assert.deepStrictEqual(actualOutput, expectedVSCodeOutput);
 	});
 
@@ -320,7 +323,7 @@ suite('Notebook Serializer', () => {
 			id: 'testId',
 			metadata: undefined
 		};
-		let actualOutput = VSCodeContentManager.convertToVSCodeCellOutput(cellOutput);
+		let actualOutput = convertToVSCodeCellOutput(cellOutput);
 		assert.deepStrictEqual(actualOutput, expectedVSCodeOutput);
 	});
 
@@ -334,6 +337,109 @@ suite('Notebook Serializer', () => {
 		assert(serializeSpy.calledOnce);
 		assert.deepStrictEqual(serializeSpy.firstCall.args[0], expectedSerializeArg);
 	});
+
+	const testDoc: azdata.nb.NotebookDocument = {
+		uri: URI.parse('untitled:a/b/c/testNotebook.ipynb'),
+		fileName: 'testFile',
+		providerId: 'testProvider',
+		isUntitled: true,
+		isDirty: true,
+		isClosed: false,
+		cells: [{
+			contents: {
+				cell_type: 'code',
+				source: '1+1'
+			}
+		}, {
+			contents: {
+				cell_type: 'markdown',
+				source: 'abc'
+			}
+		}],
+		kernelSpec: {
+			name: 'testKernel',
+			language: 'testLanguage',
+			display_name: 'testKernelName'
+		},
+		save: () => undefined,
+		setTrusted: () => undefined,
+		validateCellRange: () => undefined
+	};
+
+	test('Convert ADS NotebookDocument into VS Code NotebookDocument', async () => {
+		let expectedDoc: vscode.NotebookDocument = {
+			get uri() { return testDoc.uri; },
+			get notebookType() { return testDoc.providerId; },
+			get version() { return undefined; },
+			get isDirty() { return true; },
+			get isUntitled() { return true; },
+			get isClosed() { return false; },
+			get metadata() { return {}; },
+			get cellCount() { return 2; },
+			cellAt: () => undefined,
+			getCells: () => undefined,
+			save: () => undefined
+		};
+
+		let actualDoc = new VSCodeNotebookDocument(testDoc);
+		assert.deepStrictEqual(actualDoc.uri, expectedDoc.uri);
+		assert.strictEqual(actualDoc.notebookType, expectedDoc.notebookType);
+		assert.strictEqual(actualDoc.version, expectedDoc.version);
+		assert.strictEqual(actualDoc.isDirty, expectedDoc.isDirty);
+		assert.strictEqual(actualDoc.isUntitled, expectedDoc.isUntitled);
+		assert.strictEqual(actualDoc.isClosed, expectedDoc.isClosed);
+		assert.deepStrictEqual(actualDoc.metadata, expectedDoc.metadata);
+		assert.strictEqual(actualDoc.cellCount, expectedDoc.cellCount);
+	});
+
+	// Have to validate cell fields manually since one of the NotebookCell fields is a function pointer,
+	// which throws off the deepEqual assertions.
+	function validateCellMatches(actual: vscode.NotebookCell, expected: vscode.NotebookCell): void {
+		assert.strictEqual(actual.index, expected.index);
+		assert.deepStrictEqual(actual.document.uri, expected.document.uri);
+		assert.strictEqual(actual.document.languageId, expected.document.languageId);
+		assert.deepStrictEqual(actual.notebook.uri, expected.notebook.uri);
+	}
+	function validateCellsMatch(actual: vscode.NotebookCell[], expected: vscode.NotebookCell[]): void {
+		assert.strictEqual(actual.length, expected.length, 'Cell arrays did not have equal lengths.');
+		for (let i = 0; i < actual.length; i++) {
+			validateCellMatches(actual[i], expected[i]);
+		}
+	}
+
+	test('Retrieve range of cells from VS Code NotebookDocument', async () => {
+		let expectedCells: vscode.NotebookCell[] = testDoc.cells.map((cell, index) => convertToVSCodeNotebookCell(cell.contents.source, index, testDoc.uri, testDoc.kernelSpec.language));
+		let vsDoc = new VSCodeNotebookDocument(testDoc);
+
+		let actualCells = vsDoc.getCells();
+		validateCellsMatch(actualCells, expectedCells);
+
+		actualCells = vsDoc.getCells({ start: 0, end: 2, isEmpty: false, with: () => undefined });
+		validateCellsMatch(actualCells, expectedCells);
+
+		actualCells = vsDoc.getCells({ start: 0, end: 1, isEmpty: false, with: () => undefined });
+		validateCellsMatch(actualCells, [expectedCells[0]]);
+
+		actualCells = vsDoc.getCells({ start: 1, end: 2, isEmpty: false, with: () => undefined });
+		validateCellsMatch(actualCells, [expectedCells[1]]);
+	});
+
+	test('Retrieve specific cell from VS Code NotebookDocument', async () => {
+		let expectedCells: vscode.NotebookCell[] = testDoc.cells.map((cell, index) => convertToVSCodeNotebookCell(cell.contents.source, index, testDoc.uri, testDoc.kernelSpec.language));
+		let vsDoc = new VSCodeNotebookDocument(testDoc);
+
+		let firstCell = vsDoc.cellAt(0);
+		validateCellMatches(firstCell, expectedCells[0]);
+
+		firstCell = vsDoc.cellAt(-5);
+		validateCellMatches(firstCell, expectedCells[0]);
+
+		let secondCell = vsDoc.cellAt(1);
+		validateCellMatches(secondCell, expectedCells[1]);
+
+		secondCell = vsDoc.cellAt(10);
+		validateCellMatches(secondCell, expectedCells[1]);
+	});
 });
 
 suite('Notebook Controller', () => {

src/vs/workbench/api/common/extHost.api.impl.ts

@@ -93,6 +93,7 @@ import { matchesScheme } from 'vs/platform/opener/common/opener';
 // import { ExtHostInteractive } from 'vs/workbench/api/common/extHostInteractive'; {{SQL CARBON EDIT}} Remove until we need it
 import { ExtHostNotebook } from 'sql/workbench/api/common/extHostNotebook';
 import { functionalityNotSupportedError } from 'sql/base/common/locConstants';
+import { ExtHostNotebookDocumentsAndEditors } from 'sql/workbench/api/common/extHostNotebookDocumentsAndEditors';
 
 export interface IExtensionApiFactory {
 	(extension: IExtensionDescription, registry: ExtensionDescriptionRegistry, configProvider: ExtHostConfigProvider): typeof vscode;
@@ -101,7 +102,7 @@ export interface IExtensionApiFactory {
 /**
  * This method instantiates and returns the extension API surface
  */
-export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor, extHostNotebook: ExtHostNotebook): IExtensionApiFactory { // {{SQL CARBON EDIT}} Add ExtHostNotebook
+export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor, extHostNotebook: ExtHostNotebook, extHostNotebookDocumentsAndEditors: ExtHostNotebookDocumentsAndEditors): IExtensionApiFactory { // {{SQL CARBON EDIT}} Add ExtHostNotebook
 
 	// services
 	const initData = accessor.get(IExtHostInitDataService);
@@ -902,16 +903,15 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor, ex
 				// return extHostNotebook.getNotebookDocument(uri).apiNotebook;
 			},
 			get onDidOpenNotebookDocument(): Event<vscode.NotebookDocument> {
-				// {{SQL CARBON EDIT}} Disable VS Code notebooks
+				// {{SQL CARBON EDIT}} Use our own notebooks
-				throw new Error(functionalityNotSupportedError);
+				return extHostNotebookDocumentsAndEditors.onDidOpenVSCodeNotebookDocument;
-				// return extHostNotebook.onDidOpenNotebookDocument;
 			},
 			get onDidCloseNotebookDocument(): Event<vscode.NotebookDocument> {
-				// {{SQL CARBON EDIT}} Disable VS Code notebooks
+				// {{SQL CARBON EDIT}} Use our own notebooks
-				throw new Error(functionalityNotSupportedError);
+				return extHostNotebookDocumentsAndEditors.onDidCloseVSCodeNotebookDocument;
-				// return extHostNotebook.onDidCloseNotebookDocument;
 			},
 			registerNotebookSerializer(viewType: string, serializer: vscode.NotebookSerializer, options?: vscode.NotebookDocumentContentOptions, registration?: vscode.NotebookRegistrationData) {
+				// {{SQL CARBON EDIT}} Use our own notebooks
 				return extHostNotebook.registerNotebookSerializer(viewType, serializer, options, extension.enableProposedApi ? registration : undefined);
 			},
 			registerNotebookContentProvider: (viewType: string, provider: vscode.NotebookContentProvider, options?: vscode.NotebookDocumentContentOptions, registration?: vscode.NotebookRegistrationData) => {
@@ -1152,6 +1152,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor, ex
 		// namespace: notebook
 		const notebooks: typeof vscode.notebooks = {
 			createNotebookController(id: string, notebookType: string, label: string, handler?, rendererScripts?: vscode.NotebookRendererScript[]) {
+				// {{SQL CARBON EDIT}} Use our own notebooks
 				return extHostNotebook.createNotebookController(extension, id, notebookType, label, handler, extension.enableProposedApi ? rendererScripts : undefined);
 			},
 			registerNotebookCellStatusBarItemProvider: (notebookType: string, provider: vscode.NotebookCellStatusBarItemProvider) => {