ckaczor/azuredatastudio
1
0
Fork 0
mirror of https://github.com/ckaczor/azuredatastudio.git synced 2026-02-02 01:25:39 -05:00
Code Issues Packages Projects Releases Wiki Activity

Kusto Kernel New Notebook Changes (#12085)

Browse Source 
* Kusto New Notebook Action Changes

* Kusto Cluster properly switches context when kernel changes

* SQL Connections kernel change to Kusto works properly

* Multiple New Kusto Notebooks open properly and change kernels properly

* Fix SQL Notebook - Change to Kusto Kernel

* Fix unit tests

* Address comments

* Add test, and finalize changes
This commit is contained in:
Vasu Bhog
2020-09-09 09:19:36 -05:00
committed by GitHub
parent a840057cd8
commit 3a0be70783
5 changed files with 119 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/sql/workbench/contrib/notebook/browser/notebookActions.ts
View File

@@ -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);
}
}
}

59
src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts
View File

@@ -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<NotebookModel> {
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, <Partial<INotebookModelOptions>>{
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'
});

5
src/sql/workbench/services/connection/test/browser/connectionManagementService.test.ts
View File

@@ -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', () => {

42
src/sql/workbench/services/notebook/browser/models/notebookModel.ts
View File

@@ -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<boolean>;
@@ -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<void> {
@@ -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(