diff --git a/src/sql/azdata.proposed.d.ts b/src/sql/azdata.proposed.d.ts index df62ca72d2..754076a8f5 100644 --- a/src/sql/azdata.proposed.d.ts +++ b/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' diff --git a/src/sql/platform/capabilities/common/capabilitiesService.ts b/src/sql/platform/capabilities/common/capabilitiesService.ts index 2116d7d86b..f9da677ca1 100644 --- a/src/sql/platform/capabilities/common/capabilitiesService.ts +++ b/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[]; } diff --git a/src/sql/workbench/contrib/connection/common/connectionProviderExtension.ts b/src/sql/workbench/contrib/connection/common/connectionProviderExtension.ts index dc23c4f697..faaaf73a60 100644 --- a/src/sql/workbench/contrib/connection/common/connectionProviderExtension.ts +++ b/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: [ diff --git a/src/sql/workbench/contrib/notebook/browser/notebook.component.ts b/src/sql/workbench/contrib/notebook/browser/notebook.component.ts index 0bb919cf09..1cbbc52b68 100644 --- a/src/sql/workbench/contrib/notebook/browser/notebook.component.ts +++ b/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))); diff --git a/src/sql/workbench/contrib/notebook/browser/notebookActions.ts b/src/sql/workbench/contrib/notebook/browser/notebookActions.ts index 3e8d9a2270..4f8eb9062d 100644 --- a/src/sql/workbench/contrib/notebook/browser/notebookActions.ts +++ b/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) { diff --git a/src/sql/workbench/contrib/notebook/test/browser/notebookActions.test.ts b/src/sql/workbench/contrib/notebook/test/browser/notebookActions.test.ts index 471ffd8aa1..fad894d857 100644 --- a/src/sql/workbench/contrib/notebook/test/browser/notebookActions.test.ts +++ b/src/sql/workbench/contrib/notebook/test/browser/notebookActions.test.ts @@ -396,11 +396,12 @@ suite('Notebook Actions', function (): void { const e: azdata.nb.IKernelChangedArgs = { newValue: { 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)}`); }); }); diff --git a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts index 664c999c95..c5c43a4c2b 100644 --- a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts +++ b/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; @@ -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 { + public async startSession(manager: INotebookManager, displayName?: string, setErrorStateOnFail?: boolean, kernelAlias?: string): Promise { 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 { + private async tryStartSessionByChangingProviders(displayName: string, kernelAlias?: string): Promise { 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)); }