Kusto Notebook Kernel Changes (#11760)

* Able to get Kernel Alias using extension registery

* Get Kernel Alias through Capability Services

* Notebook Action feature to add kusto to dropdown based on extension - complete

* Fixed indexing issue when Kusto is in kernels Dropdown

* Kusto Kernel listed properly and selected when kernel changes

* Added kernel change when user Attaches To Kusto connection

* Deleted unnecessary code/refactored

* Fix Merge Issues

* Resolving Compile issues - test file error

* Capabilities Provider Changes

* Fixed Notebook Tests

* Rearchitect kernel changes to Notebook Model

* Address minor changes

src/sql/azdata.proposed.d.ts

@@ -40,6 +40,10 @@ declare module 'azdata' {
 		export interface IStandardKernel {
 			readonly blockedOnSAW?: boolean;
 		}
+
+		export interface IKernelChangedArgs {
+			nbKernelAlias?: string
+		}
 	}
 
 	export type SqlDbType = 'BigInt' | 'Binary' | 'Bit' | 'Char' | 'DateTime' | 'Decimal'

src/sql/platform/capabilities/common/capabilitiesService.ts

@@ -21,6 +21,7 @@ export const clientCapabilities = {
 export interface ConnectionProviderProperties {
 	providerId: string;
 	displayName: string;
+	notebookKernelAlias?: string;
 	azureResource?: string;
 	connectionOptions: azdata.ConnectionOption[];
 }

src/sql/workbench/contrib/connection/common/connectionProviderExtension.ts

@@ -26,6 +26,10 @@ const ConnectionProviderContrib: IJSONSchema = {
 			type: 'string',
 			description: localize('schema.displayName', "Display Name for the provider")
 		},
+		notebookKernelAlias: {
+			type: 'string',
+			description: localize('schema.notebookKernelAlias', "Notebook Kernel Alias for the provider")
+		},
 		iconPath: {
 			description: localize('schema.iconPath', "Icon path for the server type"),
 			oneOf: [

src/sql/workbench/contrib/notebook/browser/notebook.component.ts

@@ -313,7 +313,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 			layoutChanged: this._notebookParams.input.layoutChanged,
 			capabilitiesService: this.capabilitiesService,
 			editorLoadedTimestamp: this._notebookParams.input.editorOpenedTimestamp
-		}, this.profile, this.logService, this.notificationService, this.adstelemetryService);
+		}, this.profile, this.logService, this.notificationService, this.adstelemetryService, this.capabilitiesService);
 		let trusted = await this.notebookService.isNotebookTrustCached(this._notebookParams.notebookUri, this.isDirty());
 		this._register(model.onError((errInfo: INotification) => this.handleModelError(errInfo)));
 		this._register(model.contentChanged((change) => this.handleContentChanged(change)));

src/sql/workbench/contrib/notebook/browser/notebookActions.ts

@@ -274,6 +274,7 @@ export class CollapseCellsAction extends ToggleableAction {
 const showAllKernelsConfigName = 'notebook.showAllKernels';
 const workbenchPreviewConfigName = 'workbench.enablePreviewFeatures';
 export const noKernelName = localize('noKernel', "No Kernel");
+
 export class KernelsDropdown extends SelectBox {
 	private model: NotebookModel;
 	private _showAllKernels: boolean = false;
@@ -300,16 +301,21 @@ export class KernelsDropdown extends SelectBox {
 	updateModel(model: INotebookModel): void {
 		this.model = model as NotebookModel;
 		this._register(this.model.kernelChanged((changedArgs: azdata.nb.IKernelChangedArgs) => {
-			this.updateKernel(changedArgs.newValue);
+			this.updateKernel(changedArgs.newValue, changedArgs.nbKernelAlias);
 		}));
 		let kernel = this.model.clientSession && this.model.clientSession.kernel;
 		this.updateKernel(kernel);
 	}
 
 	// Update SelectBox values
-	public updateKernel(kernel: azdata.nb.IKernel) {
+	public updateKernel(kernel: azdata.nb.IKernel, nbKernelAlias?: string) {
 		let kernels: string[] = this._showAllKernels ? [...new Set(this.model.specs.kernels.map(a => a.display_name).concat(this.model.standardKernelsDisplayName()))]
 			: this.model.standardKernelsDisplayName();
+		if (this.model.kernelAliases?.length) {
+			for (let x in this.model.kernelAliases) {
+				kernels.splice(1, 0, this.model.kernelAliases[x]);
+			}
+		}
 		if (kernel && kernel.isReady) {
 			let standardKernel = this.model.getStandardKernelFromName(kernel.name);
 			if (kernels) {
@@ -320,6 +326,9 @@ export class KernelsDropdown extends SelectBox {
 					let kernelSpec = this.model.specs.kernels.find(k => k.name === kernel.name);
 					index = firstIndex(kernels, k => k === kernelSpec?.display_name);
 				}
+				if (nbKernelAlias) {
+					index = kernels.indexOf(nbKernelAlias);
+				}
 				// This is an error case that should never happen
 				// Just in case, setting index to 0
 				if (index < 0) {
@@ -477,6 +486,13 @@ export class AttachToDropdown extends SelectBox {
 			this.model.addAttachToConnectionsToBeDisposed(connectionUri);
 			// Call doChangeContext to set the newly chosen connection in the model
 			this.doChangeContext(connectionProfile);
+
+			//Changes kernel based on connection attached to
+			if (this.model.kernelAliases.includes(connectionProfile.serverCapabilities.notebookKernelAlias)) {
+				this.model.changeKernel(connectionProfile.serverCapabilities.notebookKernelAlias);
+			} else {
+				this.model.changeKernel('SQL');
+			}
 			return true;
 		}
 		catch (error) {

src/sql/workbench/contrib/notebook/test/browser/notebookActions.test.ts

@@ -396,11 +396,12 @@ suite('Notebook Actions', function (): void {
 			const e: azdata.nb.IKernelChangedArgs = <azdata.nb.IKernelChangedArgs>{
 				newValue: <azdata.nb.IKernel>{
 					name: 'StandardKernel2'
-				}
+				},
+				nbKernelAlias: ''
 			};
 			notebookModel.kernelChangedEmitter.fire(e);
 			assert.ok(updateKernelStub.calledOnce, `updateKernel should be called exactly once`);
-			assert.ok(updateKernelStub.calledWithExactly(e.newValue), `updateKernel should be called with the parameter: ${JSON.stringify(e.newValue)}`);
+			assert.ok(updateKernelStub.calledWithExactly(e.newValue, e.nbKernelAlias), `updateKernel should be called with the parameter: ${JSON.stringify(e.newValue), JSON.stringify(e.nbKernelAlias)}`);
 		});
 
 	});

src/sql/workbench/services/notebook/browser/models/notebookModel.ts

@@ -29,6 +29,7 @@ import { startsWith } from 'vs/base/common/strings';
 import { notebookConstants } from 'sql/workbench/services/notebook/browser/interfaces';
 import { IAdsTelemetryService } from 'sql/platform/telemetry/common/telemetry';
 import { Deferred } from 'sql/base/common/promise';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 
 /*
 * Used to control whether a message in a dialog/wizard is displayed as an error,
@@ -81,6 +82,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	private _connectionUrisToDispose: string[] = [];
 	private _textCellsLoading: number = 0;
 	private _standardKernels: notebookUtils.IStandardKernelWithProvider[];
+	private _kernelAliases: string[] = [];
 
 	public requestConnectionHandler: () => Promise<boolean>;
 
@@ -89,7 +91,9 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		public connectionProfile: IConnectionProfile | undefined,
 		@ILogService private readonly logService: ILogService,
 		@INotificationService private readonly notificationService: INotificationService,
-		@IAdsTelemetryService private readonly adstelemetryService: IAdsTelemetryService
+		@IAdsTelemetryService private readonly adstelemetryService: IAdsTelemetryService,
+		@ICapabilitiesService private _capabilitiesService?: ICapabilitiesService
+
 	) {
 		super();
 		if (!_notebookOptions || !_notebookOptions.notebookUri || !_notebookOptions.notebookManagers) {
@@ -232,6 +236,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return this._providerId;
 	}
 
+	public get kernelAliases(): string[] {
+		return this._kernelAliases;
+	}
+
 	public set trustedMode(isTrusted: boolean) {
 		this._trustedMode = isTrusted;
 
@@ -514,7 +522,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		this._onErrorEmitter.fire({ message: error, severity: Severity.Error });
 	}
 
-	public async startSession(manager: INotebookManager, displayName?: string, setErrorStateOnFail?: boolean): Promise<void> {
+	public async startSession(manager: INotebookManager, displayName?: string, setErrorStateOnFail?: boolean, kernelAlias?: string): Promise<void> {
 		if (displayName && this._standardKernels) {
 			let standardKernel = find(this._standardKernels, kernel => kernel.displayName === displayName);
 			if (standardKernel) {
@@ -552,7 +560,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			await clientSession.ready;
 			if (clientSession.kernel) {
 				await clientSession.kernel.ready;
-				await this.updateKernelInfoOnKernelChange(clientSession.kernel);
+				await this.updateKernelInfoOnKernelChange(clientSession.kernel, kernelAlias);
 			}
 			if (clientSession.isInErrorState) {
 				if (setErrorStateOnFail) {
@@ -571,6 +579,15 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	}
 
 	public setDefaultKernelAndProviderId() {
+		if (this._capabilitiesService?.providers) {
+			let providers = this._capabilitiesService.providers;
+			for (const server in providers) {
+				let alias = providers[server].connection.notebookKernelAlias;
+				if (alias && this._kernelAliases.indexOf(alias) === -1) {
+					this._kernelAliases.push(providers[server].connection.notebookKernelAlias);
+				}
+			}
+		}
 		if (this._savedKernelInfo) {
 			this.sanitizeSavedKernelInfo();
 			let provider = this._kernelDisplayNameToNotebookProviderIds.get(this._savedKernelInfo.display_name);
@@ -699,10 +716,15 @@ export class NotebookModel extends Disposable implements INotebookModel {
 			return;
 		}
 		let oldDisplayName = this._activeClientSession && this._activeClientSession.kernel ? this._activeClientSession.kernel.name : undefined;
+		let nbKernelAlias: string;
+		if (this.kernelAliases.includes(displayName)) {
+			displayName = 'SQL';
+			nbKernelAlias = 'Kusto';
+		}
 		try {
 			let changeKernelNeeded = true;
 			if (mustSetProvider) {
-				let providerChanged = await this.tryStartSessionByChangingProviders(displayName);
+				let providerChanged = await this.tryStartSessionByChangingProviders(displayName, nbKernelAlias);
 				// If provider was changed, a new session with new kernel is already created. We can skip calling changeKernel.
 				changeKernelNeeded = !providerChanged;
 			}
@@ -712,7 +734,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 					let kernel = await this._activeClientSession.changeKernel(spec, this._oldKernel);
 					try {
 						await kernel.ready;
-						await this.updateKernelInfoOnKernelChange(kernel);
+						await this.updateKernelInfoOnKernelChange(kernel, nbKernelAlias);
 					} catch (err2) {
 						// TODO should we handle this in any way?
 						this.logService.error(`doChangeKernel: ignoring error ${getErrorMessage(err2)}`);
@@ -742,14 +764,15 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		// Else no need to do anything
 	}
 
-	private async updateKernelInfoOnKernelChange(kernel: nb.IKernel) {
+	private async updateKernelInfoOnKernelChange(kernel: nb.IKernel, kernelAlias?: string) {
 		await this.updateKernelInfo(kernel);
 		if (kernel.info) {
 			this.updateLanguageInfo(kernel.info.language_info);
 		}
 		this._kernelChangedEmitter.fire({
 			newValue: kernel,
-			oldValue: undefined
+			oldValue: undefined,
+			nbKernelAlias: kernelAlias
 		});
 	}
 
@@ -928,7 +951,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	 * Set _providerId and start session if it is new provider
 	 * @param displayName Kernel dispay name
 	 */
-	private async tryStartSessionByChangingProviders(displayName: string): Promise<boolean> {
+	private async tryStartSessionByChangingProviders(displayName: string, kernelAlias?: string): Promise<boolean> {
 		if (displayName) {
 			if (this._activeClientSession && this._activeClientSession.isReady) {
 				this._oldKernel = this._activeClientSession.kernel;
@@ -942,7 +965,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 				await this.shutdownActiveSession();
 				let manager = this.getNotebookManager(providerId);
 				if (manager) {
-					await this.startSession(manager, displayName, false);
+					await this.startSession(manager, displayName, false, kernelAlias);
 				} else {
 					throw new Error(localize('ProviderNoManager', "Can't find notebook manager for provider {0}", providerId));
 				}