diff --git a/src/sql/workbench/contrib/notebook/browser/notebookActions.ts b/src/sql/workbench/contrib/notebook/browser/notebookActions.ts index c731a2219e..a7d0f5f59b 100644 --- a/src/sql/workbench/contrib/notebook/browser/notebookActions.ts +++ b/src/sql/workbench/contrib/notebook/browser/notebookActions.ts @@ -337,6 +337,7 @@ export class KernelsDropdown extends SelectBox { index = 0; } this.setOptions(kernels, index); + this.model.selectedKernelDisplayName = kernels[index]; } } else if (this.model.clientSession?.isInErrorState) { kernels.unshift(noKernelName); diff --git a/src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts b/src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts index 9c9a960c59..78f5bc6695 100644 --- a/src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts +++ b/src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts @@ -69,7 +69,7 @@ let expectedNotebookContentOneCell: nb.INotebookContents = { cells: [{ cell_type: CellTypes.Code, source: ['insert into t1 values (c1, c2)'], - metadata: { language: 'python' }, + metadata: { language: 'sql' }, execution_count: 1 }], metadata: { @@ -83,6 +83,27 @@ let expectedNotebookContentOneCell: nb.INotebookContents = { nbformat_minor: 5 }; +let expectedKernelAliasNotebookContentOneCell: nb.INotebookContents = { + cells: [{ + cell_type: CellTypes.Code, + source: ['StormEvents | summarize Count = count() by State | sort by Count | limit 10'], + execution_count: 1 + }], + metadata: { + kernelspec: { + name: 'mssql', + language: 'sql', + display_name: 'SQL' + }, + language_info: { + name: 'fake', + version: '' + } + }, + nbformat: 4, + nbformat_minor: 5 +}; + let expectedParameterizedNotebookContent: nb.INotebookContents = { cells: [{ cell_type: CellTypes.Code, @@ -601,7 +622,7 @@ suite('notebook model', function (): void { }); test('Should not be in error state if client session initialization succeeds', async function (): Promise { - let model = await loadModelAndStartClientSession(); + let model = await loadModelAndStartClientSession(expectedNotebookContent); assert.equal(model.inErrorState, false); assert.equal(model.notebookManagers.length, 1); @@ -628,7 +649,7 @@ suite('notebook model', function (): void { }); test('Should close active session when closed', async function () { - let model = await loadModelAndStartClientSession(); + let model = await loadModelAndStartClientSession(expectedNotebookContent); // After client session is started, ensure session is ready assert(model.isSessionReady); @@ -641,7 +662,7 @@ suite('notebook model', function (): void { }); test('Should disconnect when connection profile created by notebook', async function () { - let model = await loadModelAndStartClientSession(); + let model = await loadModelAndStartClientSession(expectedNotebookContent); // Ensure notebook prefix is present in the connection URI queryConnectionService.setup(c => c.getConnectionUri(TypeMoq.It.isAny())).returns(() => `${uriPrefixes.notebook}some/path`); await changeContextWithConnectionProfile(model); @@ -657,7 +678,7 @@ suite('notebook model', function (): void { }); test('Should not disconnect when connection profile not created by notebook', async function () { - let model = await loadModelAndStartClientSession(); + let model = await loadModelAndStartClientSession(expectedNotebookContent); // Ensure notebook prefix isn't present in connection URI queryConnectionService.setup(c => c.getConnectionUri(TypeMoq.It.isAny())).returns(() => `${uriPrefixes.default}some/path`); await changeContextWithConnectionProfile(model); @@ -692,7 +713,7 @@ suite('notebook model', function (): void { }); test('Should connect to Fake (kernel alias) connection and set kernelAlias', async function () { - let model = await loadModelAndStartClientSession(); + let model = await loadModelAndStartClientSession(expectedNotebookContent); // Ensure notebook prefix is present in the connection URI queryConnectionService.setup(c => c.getConnectionUri(TypeMoq.It.isAny())).returns(() => `${uriPrefixes.notebook}some/path`); @@ -717,7 +738,7 @@ suite('notebook model', function (): void { }); test('Should change kernel when connecting to a Fake (kernel alias) connection', async function () { - let model = await loadModelAndStartClientSession(); + let model = await loadModelAndStartClientSession(expectedNotebookContent); // Ensure notebook prefix is present in the connection URI queryConnectionService.setup(c => c.getConnectionUri(TypeMoq.It.isAny())).returns(() => `${uriPrefixes.notebook}some/path`); await changeContextWithFakeConnectionProfile(model); @@ -742,7 +763,7 @@ suite('notebook model', function (): void { }); test('Should change kernel from Fake (kernel alias) to SQL kernel when connecting to SQL connection', async function () { - let model = await loadModelAndStartClientSession(); + let model = await loadModelAndStartClientSession(expectedNotebookContent); // Ensure notebook prefix is present in the connection URI queryConnectionService.setup(c => c.getConnectionUri(TypeMoq.It.isAny())).returns(() => `${uriPrefixes.notebook}some/path`); @@ -865,9 +886,28 @@ suite('notebook model', function (): void { assert.equal(output.metadata['multi_connection_mode'], true, 'multi_connection_mode not saved correctly to notebook metadata'); }); - async function loadModelAndStartClientSession(): Promise { + test('Should keep kernel alias as language info kernel alias name even if kernel spec is seralized as SQL', async function () { let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager); - mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent)); + mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedKernelAliasNotebookContentOneCell)); + defaultModelOptions.contentManager = mockContentManager.object; + + queryConnectionService.setup(c => c.getActiveConnections(TypeMoq.It.isAny())).returns(() => null); + + // Given I have a session that fails to start + sessionReady.resolve(); + + let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService); + await model.loadContents(); + + await model.requestModelLoad(); + + // Check to see if language info is set to kernel alias + assert.equal(model.languageInfo.name, 'fake', 'Notebook language info is not set properly'); + }); + + async function loadModelAndStartClientSession(notebookContent: nb.INotebookContents): Promise { + let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager); + mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(notebookContent)); defaultModelOptions.contentManager = mockContentManager.object; queryConnectionService.setup(c => c.getActiveConnections(TypeMoq.It.isAny())).returns(() => null); @@ -880,6 +920,7 @@ suite('notebook model', function (): void { }); let model = new NotebookModel(options, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService, capabilitiesService); model.onClientSessionReady((session) => actualSession = session); + await model.requestModelLoad(); await model.startSession(notebookManagers[0]); diff --git a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts index 5246aedc3d..e6ea4f6c6c 100644 --- a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts +++ b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts @@ -288,6 +288,10 @@ export class NotebookModel extends Disposable implements INotebookModel { return this._selectedKernelDisplayName; } + public set selectedKernelDisplayName(kernel: string) { + this._selectedKernelDisplayName = kernel; + } + public set trustedMode(isTrusted: boolean) { this._trustedMode = isTrusted; @@ -951,6 +955,11 @@ export class NotebookModel extends Disposable implements INotebookModel { private async updateKernelInfoOnKernelChange(kernel: nb.IKernel, kernelAlias?: string) { await this.updateKernelInfo(kernel); + this.kernelAliases.forEach(kernel => { + if (this._defaultLanguageInfo?.name === kernel.toLowerCase()) { + kernelAlias = kernel; + } + }); if (kernel.info) { this.updateLanguageInfo(kernel.info.language_info); }