diff --git a/extensions/integration-tests/src/notebook.test.ts b/extensions/integration-tests/src/notebook.test.ts index 51104d3a45..40d252acf1 100644 --- a/extensions/integration-tests/src/notebook.test.ts +++ b/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 { + 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 { + 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 { + 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 { let language = 'scala'; await this.cellLanguageTest(notebookContentForCellLanguageTest, title + this.invocationCount++, language, { diff --git a/extensions/integration-tests/src/notebook.util.ts b/extensions/integration-tests/src/notebook.util.ts index 97a7ee1bbd..976d192439 100644 --- a/extensions/integration-tests/src/notebook.util.ts +++ b/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); diff --git a/src/sql/azdata.proposed.d.ts b/src/sql/azdata.proposed.d.ts index da04465f97..1c649c7f9b 100644 --- a/src/sql/azdata.proposed.d.ts +++ b/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; + + /** + * Changes the Notebook's kernel. Thenable will resolve only after kernel change is complete. + */ + changeKernel(kernel: IKernelSpec): Thenable; } export interface NotebookCell { diff --git a/src/sql/workbench/api/node/extHostNotebookEditor.ts b/src/sql/workbench/api/node/extHostNotebookEditor.ts index 5b5083a1b4..f0772d2cf4 100644 --- a/src/sql/workbench/api/node/extHostNotebookEditor.ts +++ b/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 { + return this._proxy.$changeKernel(this._id, kernel); + } + public edit(callback: (editBuilder: azdata.nb.NotebookEditorEdit) => void, options?: { undoStopBefore: boolean; undoStopAfter: boolean; }): Thenable { if (this._disposed) { return Promise.reject(new Error('NotebookEditor#edit not possible on closed editors')); diff --git a/src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts b/src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts index 8a1e7c9273..3a91454e9a 100644 --- a/src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts +++ b/src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts @@ -36,11 +36,14 @@ import { localize } from 'vs/nls'; class MainThreadNotebookEditor extends Disposable { private _contentChangedEmitter = new Emitter(); public readonly contentChanged: Event = 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(); - private _modelToDisposeMap = new Map(); + private _modelToDisposeMap = new Map(); 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 { + 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 { @@ -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 { + 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); + }); + } } diff --git a/src/sql/workbench/api/node/sqlExtHost.protocol.ts b/src/sql/workbench/api/node/sqlExtHost.protocol.ts index 15a696e455..e7bda1b071 100644 --- a/src/sql/workbench/api/node/sqlExtHost.protocol.ts +++ b/src/sql/workbench/api/node/sqlExtHost.protocol.ts @@ -900,6 +900,7 @@ export interface MainThreadNotebookDocumentsAndEditorsShape extends IDisposable $runCell(id: string, cellUri: UriComponents): Promise; $runAllCells(id: string): Promise; $clearAllOutputs(id: string): Promise; + $changeKernel(id: string, kernel: azdata.nb.IKernelInfo): Promise; } export interface ExtHostExtensionManagementShape { diff --git a/src/sql/workbench/parts/notebook/models/clientSession.ts b/src/sql/workbench/parts/notebook/models/clientSession.ts index c23404c3b5..ed6432d651 100644 --- a/src/sql/workbench/parts/notebook/models/clientSession.ts +++ b/src/sql/workbench/parts/notebook/models/clientSession.ts @@ -264,8 +264,11 @@ export class ClientSession implements IClientSession { private async updateCachedKernelSpec(): Promise { 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(); + } } } diff --git a/src/sql/workbench/parts/notebook/models/modelInterfaces.ts b/src/sql/workbench/parts/notebook/models/modelInterfaces.ts index ef9422bdb2..3f3177b311 100644 --- a/src/sql/workbench/parts/notebook/models/modelInterfaces.ts +++ b/src/sql/workbench/parts/notebook/models/modelInterfaces.ts @@ -332,6 +332,11 @@ export interface INotebookModel { */ readonly contentChanged: Event; + /** + * Event fired on notebook provider change + */ + readonly onProviderIdChange: Event; + /** * The trusted mode of the Notebook */ diff --git a/src/sql/workbench/parts/notebook/models/notebookModel.ts b/src/sql/workbench/parts/notebook/models/notebookModel.ts index 4c3f70bda2..3abefd0890 100644 --- a/src/sql/workbench/parts/notebook/models/notebookModel.ts +++ b/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 - }); } } diff --git a/src/sqltest/parts/notebook/common.ts b/src/sqltest/parts/notebook/common.ts index 7c2b21abf6..7f45a7dce0 100644 --- a/src/sqltest/parts/notebook/common.ts +++ b/src/sqltest/parts/notebook/common.ts @@ -96,6 +96,9 @@ export class NotebookModelStub implements INotebookModel { get onValidConnectionSelected(): Event { throw new Error('method not implemented.'); } + get onProviderIdChange(): Event { + throw new Error('method not impelemented.'); + } toJSON(): nb.INotebookContents { throw new Error('Method not implemented.'); }