diff --git a/src/sql/platform/capabilities/test/common/testCapabilitiesService.ts b/src/sql/platform/capabilities/test/common/testCapabilitiesService.ts index 4cb33eab20..b41805bf74 100644 --- a/src/sql/platform/capabilities/test/common/testCapabilitiesService.ts +++ b/src/sql/platform/capabilities/test/common/testCapabilitiesService.ts @@ -16,6 +16,7 @@ import { IDisposable } from 'vs/base/common/lifecycle'; export class TestCapabilitiesService implements ICapabilitiesService { private pgsqlProviderName = 'PGSQL'; + private kustoProviderName = 'KUSTO'; public _serviceBrand: undefined; public capabilities: { [id: string]: ProviderFeatures } = {}; @@ -103,11 +104,18 @@ export class TestCapabilitiesService implements ICapabilitiesService { }; let pgSQLCapabilities = { providerId: this.pgsqlProviderName, - displayName: this.pgsqlProviderName, + displayName: 'PostgreSQL', connectionOptions: connectionProvider, }; + let kustoCapabilities = { + providerId: this.kustoProviderName, + displayName: 'Kusto', + connectionOptions: connectionProvider, + notebookKernelAlias: 'Kusto' + }; this.capabilities[mssqlProviderName] = { connection: msSQLCapabilities }; this.capabilities[this.pgsqlProviderName] = { connection: pgSQLCapabilities }; + this.capabilities[this.kustoProviderName] = { connection: kustoCapabilities }; } registerConnectionProvider(id: string, properties: ConnectionProviderProperties): IDisposable { diff --git a/src/sql/workbench/contrib/notebook/browser/notebookActions.ts b/src/sql/workbench/contrib/notebook/browser/notebookActions.ts index 736098b300..f7013791cf 100644 --- a/src/sql/workbench/contrib/notebook/browser/notebookActions.ts +++ b/src/sql/workbench/contrib/notebook/browser/notebookActions.ts @@ -409,7 +409,12 @@ export class AttachToDropdown extends SelectBox { let currentKernelName = this.model.clientSession.kernel.name.toLowerCase(); let currentKernelSpec = find(this.model.specs.kernels, kernel => kernel.name && kernel.name.toLowerCase() === currentKernelName); if (currentKernelSpec) { - kernelDisplayName = currentKernelSpec.display_name; + //KernelDisplayName should be Kusto when connecting to Kusto connection + if ((this.model.context?.serverCapabilities.notebookKernelAlias && this.model.currentKernelAlias === this.model.context?.serverCapabilities.notebookKernelAlias) || (this.model.kernelAliases.includes(this.model.selectedKernelDisplayName) && this.model.selectedKernelDisplayName)) { + kernelDisplayName = this.model.context?.serverCapabilities.notebookKernelAlias || this.model.selectedKernelDisplayName; + } else { + kernelDisplayName = currentKernelSpec.display_name; + } } } return kernelDisplayName; @@ -420,14 +425,17 @@ export class AttachToDropdown extends SelectBox { let connProviderIds = this.model.getApplicableConnectionProviderIds(currentKernel); if ((connProviderIds && connProviderIds.length === 0) || currentKernel === noKernel) { this.setOptions([msgLocalHost]); - } - else { - let connections: string[] = model.context && model.context.title ? [model.context.title] : [msgSelectConnection]; + } else { + let connections: string[] = model.context && model.context.title && (connProviderIds.includes(this.model.context.providerName)) ? [model.context.title] : [msgSelectConnection]; if (!find(connections, x => x === msgChangeConnection)) { connections.push(msgChangeConnection); } this.setOptions(connections, 0); this.enable(); + + if (this.model.kernelAliases.includes(currentKernel) && this.model.selectedKernelDisplayName !== currentKernel) { + this.model.changeKernel(currentKernel); + } } } 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 15b614f87f..3287e64562 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 @@ -6,6 +6,7 @@ import * as TypeMoq from 'typemoq'; import { nb } from 'azdata'; import * as assert from 'assert'; +import * as sinon from 'sinon'; import { INotificationService } from 'vs/platform/notification/common/notification'; import { TestNotificationService } from 'vs/platform/notification/test/common/testNotificationService'; @@ -485,6 +486,59 @@ suite('notebook model', function (): void { assert(output.metadata['custom-object']['prop2'] === 'value2', 'Custom metadata for object was not preserved'); }); + test('Should get Kusto connection and set kernelAlias', async function () { + let model = await loadModelAndStartClientSession(); + + // 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); + + //Check to see if Kusto is added to kernelAliases + assert(!isUndefinedOrNull(model.kernelAliases)); + let expectedAlias = ['Kusto']; + let kernelAliases = model.kernelAliases; + + assert.equal(kernelAliases.length, 1); + assert(kernelAliases.includes(expectedAlias[0])); + + // // After client session is started, ensure context isn't null/undefined + assert(!isUndefinedOrNull(model.context), 'context should exist after call to change context'); + + // After closing the notebook + await model.handleClosed(); + + // Ensure disconnect is called once + queryConnectionService.verify((c) => c.disconnect(TypeMoq.It.isAny()), TypeMoq.Times.once()); + }); + + test.skip('Should change kernel to Kusto when connecting to Kusto connection', async function () { + let model = await loadModelAndStartClientSession(); + + // 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); + + // // After client session is started, ensure context isn't null/undefined + assert(!isUndefinedOrNull(model.context), 'context should exist after call to change context'); + + let doChangeKernelStub = sinon.spy(model, 'doChangeKernel').withArgs(model.kernelAliases[0]); + + //Change to kusto kernel + //TODO issue with Test not setting serverCapabilities of context + await model.changeKernel(model.kernelAliases[0]); + let notebookKernelAlias = capabilitiesService.instance.providers.KUSTO.connection.notebookKernelAlias; + assert.equal(model.selectedKernelDisplayName, model.kernelAliases[0]); + assert.equal(model.currentKernelAlias, notebookKernelAlias); + sinon.assert.called(doChangeKernelStub); + sinon.restore(doChangeKernelStub); + + // After closing the notebook + await model.handleClosed(); + + // Ensure disconnect is called once + queryConnectionService.verify((c) => c.disconnect(TypeMoq.It.isAny()), TypeMoq.Times.once()); + }); + async function loadModelAndStartClientSession(): Promise { let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager); mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent)); @@ -505,7 +559,8 @@ suite('notebook model', function (): void { let options: INotebookModelOptions = assign({}, defaultModelOptions, >{ factory: mockModelFactory.object }); - let model = new NotebookModel(options, undefined, logService, undefined, new NullAdsTelemetryService()); + let capabilitiesService = new TestCapabilitiesService; + let model = new NotebookModel(options, undefined, logService, undefined, new NullAdsTelemetryService(), capabilitiesService); model.onClientSessionReady((session) => actualSession = session); await model.requestModelLoad(); @@ -534,7 +589,7 @@ suite('notebook model', function (): void { userName: 'testUsername', groupId: undefined, providerName: mssqlProviderName, - options: {}, + options: { serverCapabilities: capabilitiesService.instance.providers.KUSTO.connection }, saveProfile: true, id: 'testID' }); diff --git a/src/sql/workbench/services/connection/test/browser/connectionManagementService.test.ts b/src/sql/workbench/services/connection/test/browser/connectionManagementService.test.ts index 009948aa06..141c29fab3 100644 --- a/src/sql/workbench/services/connection/test/browser/connectionManagementService.test.ts +++ b/src/sql/workbench/services/connection/test/browser/connectionManagementService.test.ts @@ -1098,11 +1098,12 @@ suite('SQL ConnectionManagementService tests', () => { }); test('providerNameToDisplayNameMap should return all providers', () => { - let expectedNames = ['MSSQL', 'PGSQL']; + let expectedNames = ['MSSQL', 'PGSQL', 'KUSTO']; let providerNames = Object.keys(connectionManagementService.providerNameToDisplayNameMap); - assert.equal(providerNames.length, 2); + assert.equal(providerNames.length, 3); assert.equal(providerNames[0], expectedNames[0]); assert.equal(providerNames[1], expectedNames[1]); + assert.equal(providerNames[2], expectedNames[2]); }); test('ensureDefaultLanguageFlavor should not send event if uri is connected', () => { diff --git a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts index c5c43a4c2b..e186898417 100644 --- a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts +++ b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts @@ -83,6 +83,8 @@ export class NotebookModel extends Disposable implements INotebookModel { private _textCellsLoading: number = 0; private _standardKernels: notebookUtils.IStandardKernelWithProvider[]; private _kernelAliases: string[] = []; + private _currentKernelAlias: string; + private _selectedKernelDisplayName: string; public requestConnectionHandler: () => Promise; @@ -240,6 +242,14 @@ export class NotebookModel extends Disposable implements INotebookModel { return this._kernelAliases; } + public get currentKernelAlias(): string { + return this._currentKernelAlias; + } + + public get selectedKernelDisplayName(): string { + return this._selectedKernelDisplayName; + } + public set trustedMode(isTrusted: boolean) { this._trustedMode = isTrusted; @@ -629,7 +639,12 @@ export class NotebookModel extends Disposable implements INotebookModel { if (this._standardKernels) { let standardKernels = find(this._standardKernels, kernel => this._defaultKernel && kernel.displayName === this._defaultKernel.display_name); let connectionProviderIds = standardKernels ? standardKernels.connectionProviderIds : undefined; - return profile && connectionProviderIds && find(connectionProviderIds, provider => provider === profile.providerName) !== undefined; + let providerFeatures = this._capabilitiesService.getCapabilities(profile.providerName); + if (connectionProviderIds.length > 0) { + this._currentKernelAlias = providerFeatures?.connection.notebookKernelAlias; + this._kernelDisplayNameToConnectionProviderIds.set(this._currentKernelAlias, [profile.providerName]); + } + return this._currentKernelAlias || profile && connectionProviderIds && find(connectionProviderIds, provider => provider === profile.providerName) !== undefined; } return false; } @@ -706,8 +721,15 @@ export class NotebookModel extends Disposable implements INotebookModel { } public changeKernel(displayName: string): void { - this._contextsLoadingEmitter.fire(); - this.doChangeKernel(displayName, true).catch(e => this.logService.error(e)); + this._selectedKernelDisplayName = displayName; + this._currentKernelAlias = this.context?.serverCapabilities.notebookKernelAlias; + if (this.kernelAliases.includes(this.currentKernelAlias) && displayName === this.currentKernelAlias) { + this.doChangeKernel(displayName, true).catch(e => this.logService.error(e)); + } else { + this._currentKernelAlias = undefined; + this._contextsLoadingEmitter.fire(); + this.doChangeKernel(displayName, true).catch(e => this.logService.error(e)); + } } private async doChangeKernel(displayName: string, mustSetProvider: boolean = true, restoreOnFail: boolean = true): Promise { @@ -718,8 +740,10 @@ export class NotebookModel extends Disposable implements INotebookModel { let oldDisplayName = this._activeClientSession && this._activeClientSession.kernel ? this._activeClientSession.kernel.name : undefined; let nbKernelAlias: string; if (this.kernelAliases.includes(displayName)) { + this._currentKernelAlias = displayName; displayName = 'SQL'; - nbKernelAlias = 'Kusto'; + nbKernelAlias = this._currentKernelAlias; + this._kernelDisplayNameToConnectionProviderIds.set(this.currentKernelAlias, [this.currentKernelAlias.toUpperCase()]); } try { let changeKernelNeeded = true; @@ -798,6 +822,16 @@ export class NotebookModel extends Disposable implements INotebookModel { } if (newConnection) { + if (newConnection.serverCapabilities?.notebookKernelAlias) { + this._currentKernelAlias = newConnection.serverCapabilities.notebookKernelAlias; + let sqlConnectionProvider = this._kernelDisplayNameToConnectionProviderIds.get('SQL'); + let index = sqlConnectionProvider.indexOf(newConnection.serverCapabilities.notebookKernelAlias.toUpperCase()); + if (index > -1) { + sqlConnectionProvider.splice(index, 1); + } + this._kernelDisplayNameToConnectionProviderIds.set('SQL', sqlConnectionProvider); + this._kernelDisplayNameToConnectionProviderIds.set(newConnection.serverCapabilities.notebookKernelAlias, [newConnection.providerName]); + } this._activeConnection = newConnection; this.setActiveConnectionIfDifferent(newConnection); this._activeClientSession.updateConnection(newConnection.toIConnectionProfile()).then(