Notebooks: Adding Change Kernel API, 3 Integration Tests (#5287)

* Notebook change kernel

* Fix notifying of k change too many times add tests

* Fix broken unit test

* Deal with comment

extensions/integration-tests/src/notebook.test.ts

@@ -10,7 +10,7 @@ import * as assert from 'assert';
 import * as azdata from 'azdata';
 import * as vscode from 'vscode';
 import { context } from './testContext';
-import { sqlNotebookContent, writeNotebookToFile, sqlKernelMetadata, getFileName, pySparkNotebookContent, pySpark3KernelMetadata, pythonKernelMetadata, sqlNotebookMultipleCellsContent, notebookContentForCellLanguageTest } from './notebook.util';
+import { sqlNotebookContent, writeNotebookToFile, sqlKernelMetadata, getFileName, pySparkNotebookContent, pySpark3KernelMetadata, pythonKernelMetadata, sqlNotebookMultipleCellsContent, notebookContentForCellLanguageTest, sqlKernelSpec, pythonKernelSpec, pySpark3KernelSpec } from './notebook.util';
 import { getBdcServer, getConfigValue, EnvironmentVariable_PYTHON_PATH } from './testConfig';
 import { connectToServer } from './utils';
 import * as fs from 'fs';
@@ -56,6 +56,18 @@ if (context.RunTest) {
 			test('python language test', async function () {
 				await (new NotebookTester()).pythonLanguageTest(this.test.title);
 			});
+
+			test('Change kernel different provider SQL to Python to SQL', async function () {
+				await (new NotebookTester()).sqlNbChangeKernelDifferentProviderTest(this.test.title);
+			});
+
+			test('Change kernel different provider Python to SQL to Python', async function () {
+				await (new NotebookTester()).pythonChangeKernelDifferentProviderTest(this.test.title);
+			});
+
+			test('Change kernel same provider Python to PySpark3 to Python', async function () {
+				await (new NotebookTester()).pythonChangeKernelSameProviderTest(this.test.title);
+			});
 		}
 
 		if (process.env['RUN_PYSPARK_TEST'] === '1') {
@@ -154,6 +166,48 @@ class NotebookTester {
 		assert(actualOutput2[0] === '1', `Expected result: 1, Actual: '${actualOutput2[0]}'`);
 	}
 
+	async sqlNbChangeKernelDifferentProviderTest(title: string): Promise<void> {
+		let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title);
+		assert(notebook.document.providerId === 'sql', `Expected providerId to be sql, Actual: ${notebook.document.providerId}`);
+		assert(notebook.document.kernelSpec.name === 'SQL', `Expected first kernel name: SQL, Actual: ${notebook.document.kernelSpec.name}`);
+
+		let kernelChanged = await notebook.changeKernel(pythonKernelSpec);
+		assert(notebook.document.providerId === 'jupyter', `Expected providerId to be jupyter, Actual: ${notebook.document.providerId}`);
+		assert(kernelChanged && notebook.document.kernelSpec.name === 'python3', `Expected second kernel name: python3, Actual: ${notebook.document.kernelSpec.name}`);
+
+		kernelChanged = await notebook.changeKernel(sqlKernelSpec);
+		assert(notebook.document.providerId === 'sql', `Expected providerId to be sql, Actual: ${notebook.document.providerId}`);
+		assert(kernelChanged && notebook.document.kernelSpec.name === 'SQL', `Expected third kernel name: SQL, Actual: ${notebook.document.kernelSpec.name}`);
+	}
+
+	async pythonChangeKernelDifferentProviderTest(title: string): Promise<void> {
+		let notebook = await this.openNotebook(pySparkNotebookContent, pythonKernelMetadata, title);
+		assert(notebook.document.providerId === 'jupyter', `Expected providerId to be jupyter, Actual: ${notebook.document.providerId}`);
+		assert(notebook.document.kernelSpec.name === 'python3', `Expected first kernel name: python3, Actual: ${notebook.document.kernelSpec.name}`);
+
+		let kernelChanged = await notebook.changeKernel(sqlKernelSpec);
+		assert(notebook.document.providerId === 'sql', `Expected providerId to be sql, Actual: ${notebook.document.providerId}`);
+		assert(kernelChanged && notebook.document.kernelSpec.name === 'SQL', `Expected second kernel name: SQL, Actual: ${notebook.document.kernelSpec.name}`);
+
+		kernelChanged = await notebook.changeKernel(pythonKernelSpec);
+		assert(notebook.document.providerId === 'jupyter', `Expected providerId to be jupyter, Actual: ${notebook.document.providerId}`);
+		assert(kernelChanged && notebook.document.kernelSpec.name === 'python3', `Expected third kernel name: python3, Actual: ${notebook.document.kernelSpec.name}`);
+	}
+
+	async pythonChangeKernelSameProviderTest(title: string): Promise<void> {
+		let notebook = await this.openNotebook(pySparkNotebookContent, pythonKernelMetadata, title);
+		assert(notebook.document.providerId === 'jupyter', `Expected providerId to be jupyter, Actual: ${notebook.document.providerId}`);
+		assert(notebook.document.kernelSpec.name === 'python3', `Expected first kernel name: python3, Actual: ${notebook.document.kernelSpec.name}`);
+
+		let kernelChanged = await notebook.changeKernel(pySpark3KernelSpec);
+		assert(notebook.document.providerId === 'jupyter', `Expected providerId to be jupyter, Actual: ${notebook.document.providerId}`);
+		assert(kernelChanged && notebook.document.kernelSpec.name === 'pyspark3kernel', `Expected second kernel name: pyspark3kernel, Actual: ${notebook.document.kernelSpec.name}`);
+
+		kernelChanged = await notebook.changeKernel(pythonKernelSpec);
+		assert(notebook.document.providerId === 'jupyter', `Expected providerId to be jupyter, Actual: ${notebook.document.providerId}`);
+		assert(kernelChanged && notebook.document.kernelSpec.name === 'python3', `Expected third kernel name: python3, Actual: ${notebook.document.kernelSpec.name}`);
+	}
+
 	async scalaLanguageTest(title: string): Promise<void> {
 		let language = 'scala';
 		await this.cellLanguageTest(notebookContentForCellLanguageTest, title + this.invocationCount++, language, {

extensions/integration-tests/src/notebook.util.ts

@@ -134,6 +134,11 @@ export const pySpark3KernelMetadata = {
 	}
 };
 
+export const pySpark3KernelSpec = {
+	name: 'pyspark3',
+	display_name: 'PySpark3'
+};
+
 export const sqlKernelMetadata = {
 	'kernelspec': {
 		'name': 'SQL',
@@ -141,6 +146,11 @@ export const sqlKernelMetadata = {
 	}
 };
 
+export const sqlKernelSpec: azdata.nb.IKernelSpec = {
+	name: 'SQL',
+	display_name: 'SQL'
+};
+
 export const pythonKernelMetadata = {
 	'kernelspec': {
 		'name': 'python3',
@@ -148,6 +158,11 @@ export const pythonKernelMetadata = {
 	}
 };
 
+export const pythonKernelSpec: azdata.nb.IKernelSpec = {
+	name: 'python3',
+	display_name: 'Python 3'
+};
+
 export function writeNotebookToFile(pythonNotebook: azdata.nb.INotebookContents, testName: string): vscode.Uri {
 	let fileName = getFileName(testName);
 	let notebookContentString = JSON.stringify(pythonNotebook);

src/sql/azdata.proposed.d.ts

@@ -4363,6 +4363,11 @@ declare module 'azdata' {
 			* @return A promise that resolves with a value indicating if the outputs are cleared or not.
 			*/
 			clearAllOutputs(): Thenable<boolean>;
+
+			/**
+			 * Changes the Notebook's kernel. Thenable will resolve only after kernel change is complete.
+			 */
+			changeKernel(kernel: IKernelSpec): Thenable<boolean>;
 		}
 
 		export interface NotebookCell {

src/sql/workbench/api/node/extHostNotebookEditor.ts

@@ -163,6 +163,10 @@ export class ExtHostNotebookEditor implements azdata.nb.NotebookEditor, IDisposa
 		return this._proxy.$clearAllOutputs(this._id);
 	}
 
+	public changeKernel(kernel: azdata.nb.IKernelSpec): Thenable<boolean> {
+		return this._proxy.$changeKernel(this._id, kernel);
+	}
+
 	public edit(callback: (editBuilder: azdata.nb.NotebookEditorEdit) => void, options?: { undoStopBefore: boolean; undoStopAfter: boolean; }): Thenable<boolean> {
 		if (this._disposed) {
 			return Promise.reject(new Error('NotebookEditor#edit not possible on closed editors'));

src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts

@@ -36,11 +36,14 @@ import { localize } from 'vs/nls';
 class MainThreadNotebookEditor extends Disposable {
 	private _contentChangedEmitter = new Emitter<NotebookContentChange>();
 	public readonly contentChanged: Event<NotebookContentChange> = this._contentChangedEmitter.event;
-	private _providerInfo: IProviderInfo;
+	private _providerId: string = '';
+	private _providers: string[] = [];
 
 	constructor(public readonly editor: INotebookEditor) {
 		super();
 		editor.modelReady.then(model => {
+			this._providerId = model.providerId;
+
 			this._register(model.contentChanged((e) => this._contentChangedEmitter.fire(e)));
 			this._register(model.kernelChanged((e) => {
 				let changeEvent: NotebookContentChange = {
@@ -48,9 +51,12 @@ class MainThreadNotebookEditor extends Disposable {
 				};
 				this._contentChangedEmitter.fire(changeEvent);
 			}));
+			this._register(model.onProviderIdChange((provider) => {
+				this._providerId = provider;
+			}));
 		});
 		editor.notebookParams.providerInfo.then(info => {
-			this._providerInfo = info;
+			this._providers = info.providers;
 		});
 	}
 
@@ -67,11 +73,11 @@ class MainThreadNotebookEditor extends Disposable {
 	}
 
 	public get providerId(): string {
-		return this._providerInfo ? this._providerInfo.providerId : undefined;
+		return this._providerId;
 	}
 
 	public get providers(): string[] {
-		return this._providerInfo ? this._providerInfo.providers : [];
+		return this._providers;
 	}
 
 	public get cells(): ICellModel[] {
@@ -306,7 +312,7 @@ class MainThreadNotebookDocumentAndEditorStateComputer extends Disposable {
 export class MainThreadNotebookDocumentsAndEditors extends Disposable implements MainThreadNotebookDocumentsAndEditorsShape {
 	private _proxy: ExtHostNotebookDocumentsAndEditorsShape;
 	private _notebookEditors = new Map<string, MainThreadNotebookEditor>();
-	private _modelToDisposeMap = new Map<string, IDisposable>();
+	private _modelToDisposeMap = new Map<string, IDisposable[]>();
 	constructor(
 		extHostContext: IExtHostContext,
 		@IUntitledEditorService private _untitledEditorService: IUntitledEditorService,
@@ -385,6 +391,14 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		return editor.clearAllOutputs();
 	}
 
+	$changeKernel(id: string, kernel: azdata.nb.IKernelSpec): Promise<boolean> {
+		let editor = this.getEditor(id);
+		if (!editor) {
+			return Promise.reject(disposed(`TextEditor(${id})`));
+		}
+		return this.doChangeKernel(editor, kernel.display_name);
+	}
+
 	//#endregion
 
 	private async doOpenEditor(resource: UriComponents, options: INotebookShowOptions): Promise<string> {
@@ -500,9 +514,11 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 			return;
 		}
 		removedDocuments.forEach(removedDoc => {
-			let listener = this._modelToDisposeMap.get(removedDoc.toString());
-			if (listener) {
-				listener.dispose();
+			let listeners = this._modelToDisposeMap.get(removedDoc.toString());
+			if (listeners && listeners.length) {
+				listeners.forEach(listener => {
+					listener.dispose();
+				});
 				this._modelToDisposeMap.delete(removedDoc.toString());
 			}
 		});
@@ -514,9 +530,9 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		}
 		addedEditors.forEach(editor => {
 			let modelUrl = editor.uri;
-			this._modelToDisposeMap.set(editor.uri.toString(), editor.contentChanged((e) => {
+			this._modelToDisposeMap.set(editor.uri.toString(), [editor.contentChanged((e) => {
 				this._proxy.$acceptModelChanged(modelUrl, this._toNotebookChangeData(e, editor));
-			}));
+			})]);
 		});
 	}
 
@@ -591,4 +607,15 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements
 		}
 		return notebookCells;
 	}
+
+	private async doChangeKernel(editor: MainThreadNotebookEditor, displayName: string): Promise<boolean> {
+		let listeners = this._modelToDisposeMap.get(editor.id);
+		editor.model.changeKernel(displayName);
+		return new Promise((resolve) => {
+			listeners.push(editor.model.kernelChanged((kernel) => {
+				resolve(true);
+			}));
+			this._modelToDisposeMap.set(editor.id, listeners);
+		});
+	}
 }

src/sql/workbench/api/node/sqlExtHost.protocol.ts

@@ -900,6 +900,7 @@ export interface MainThreadNotebookDocumentsAndEditorsShape extends IDisposable
 	$runCell(id: string, cellUri: UriComponents): Promise<boolean>;
 	$runAllCells(id: string): Promise<boolean>;
 	$clearAllOutputs(id: string): Promise<boolean>;
+	$changeKernel(id: string, kernel: azdata.nb.IKernelInfo): Promise<boolean>;
 }
 
 export interface ExtHostExtensionManagementShape {

src/sql/workbench/parts/notebook/models/clientSession.ts

@@ -264,8 +264,11 @@ export class ClientSession implements IClientSession {
 	private async updateCachedKernelSpec(): Promise<void> {
 		this._cachedKernelSpec = undefined;
 		let kernel = this.kernel;
-		if (kernel && kernel.isReady) {
-			this._cachedKernelSpec = await this.kernel.getSpec();
+		if (kernel) {
+			await kernel.ready;
+			if (kernel.isReady) {
+				this._cachedKernelSpec = await kernel.getSpec();
+			}
 		}
 	}
 

src/sql/workbench/parts/notebook/models/modelInterfaces.ts

@@ -332,6 +332,11 @@ export interface INotebookModel {
 	 */
 	readonly contentChanged: Event<NotebookContentChange>;
 
+	/**
+	 * Event fired on notebook provider change
+	 */
+	readonly onProviderIdChange: Event<string>;
+
 	/**
 	 * The trusted mode of the Notebook
 	 */

src/sql/workbench/parts/notebook/models/notebookModel.ts

@@ -457,10 +457,6 @@ export class NotebookModel extends Disposable implements INotebookModel {
 				}
 			}
 			this._onClientSessionReady.fire(clientSession);
-			this._kernelChangedEmitter.fire({
-				oldValue: undefined,
-				newValue: clientSession.kernel
-			});
 		}
 	}
 

src/sqltest/parts/notebook/common.ts

@@ -96,6 +96,9 @@ export class NotebookModelStub implements INotebookModel {
 	get onValidConnectionSelected(): Event<boolean> {
 		throw new Error('method not implemented.');
 	}
+	get onProviderIdChange(): Event<string> {
+		throw new Error('method not impelemented.');
+	}
 	toJSON(): nb.INotebookContents {
 		throw new Error('Method not implemented.');
 	}