Fix #4452 Notebook is reloaded with wrong kernel (#4480)

* Fix #4452 Notebook is reloaded with wrong kernel
- Await all extension registration before getting providers and providerId
To do this, we need to await way up in the NotebookInput, and promise the model that we'll have values eventually

src/sql/parts/notebook/notebook.component.ts

@@ -24,7 +24,7 @@ import { VIEWLET_ID, IExtensionsViewlet } from 'vs/workbench/parts/extensions/co
 import { CommonServiceInterface } from 'sql/services/common/commonServiceInterface.service';
 import { AngularDisposable } from 'sql/base/node/lifecycle';
 import { CellTypes, CellType } from 'sql/parts/notebook/models/contracts';
-import { ICellModel, IModelFactory, INotebookModel, NotebookContentChange, notebookConstants } from 'sql/parts/notebook/models/modelInterfaces';
+import { ICellModel, IModelFactory, INotebookModel, NotebookContentChange } from 'sql/parts/notebook/models/modelInterfaces';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { INotebookService, INotebookParams, INotebookManager, INotebookEditor, DEFAULT_NOTEBOOK_PROVIDER, SQL_NOTEBOOK_PROVIDER } from 'sql/workbench/services/notebook/common/notebookService';
 import { IBootstrapParams } from 'sql/services/bootstrap/bootstrapService';
@@ -38,8 +38,6 @@ import { KernelsDropdown, AttachToDropdown, AddCellAction, TrustedAction } from
 import { IObjectExplorerService } from 'sql/workbench/services/objectExplorer/common/objectExplorerService';
 import * as TaskUtilities from 'sql/workbench/common/taskUtilities';
 import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
-import { IUntitledEditorService } from 'vs/workbench/services/untitled/common/untitledEditorService';
-import { IEditorGroupsService } from 'vs/workbench/services/group/common/editorGroupsService';
 import { IConnectionDialogService } from 'sql/workbench/services/connection/common/connectionDialogService';
 import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 import { CellMagicMapper } from 'sql/parts/notebook/models/cellMagicMapper';
@@ -254,8 +252,9 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		this.detectChanges();
 	}
 
-	private async setNotebookManager() {
-		for (let providerId of this._notebookParams.providers) {
+	private async setNotebookManager(): Promise<void> {
+		let providerInfo = await this._notebookParams.providerInfo;
+		for (let providerId of providerInfo.providers) {
 			let notebookManager = await this.notebookService.getOrCreateNotebookManager(providerId, this._notebookParams.notebookUri);
 			this.notebookManagers.push(notebookManager);
 		}
@@ -265,12 +264,14 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		// Wait on registration for now. Long-term would be good to cache and refresh
 		await this.notebookService.registrationComplete;
 		// Refresh the provider if we had been using default
-		if (DEFAULT_NOTEBOOK_PROVIDER === this._notebookParams.providerId) {
+		let providerInfo = await this._notebookParams.providerInfo;
+
+		if (DEFAULT_NOTEBOOK_PROVIDER === providerInfo.providerId) {
 			let providers = notebookUtils.getProvidersForFileName(this._notebookParams.notebookUri.fsPath, this.notebookService);
 			let tsqlProvider = providers.find(provider => provider === SQL_NOTEBOOK_PROVIDER);
-			this._notebookParams.providerId = tsqlProvider ? SQL_NOTEBOOK_PROVIDER : providers[0];
+			providerInfo.providerId = tsqlProvider ? SQL_NOTEBOOK_PROVIDER : providers[0];
 		}
-		if (DEFAULT_NOTEBOOK_PROVIDER === this._notebookParams.providerId) {
+		if (DEFAULT_NOTEBOOK_PROVIDER === providerInfo.providerId) {
 			// If it's still the default, warn them they should install an extension
 			this.notificationService.prompt(Severity.Warning,
 				localize('noKernelInstalled', 'Please install the SQL Server 2019 extension to run cells'),

src/sql/parts/notebook/notebookEditor.ts

@@ -15,7 +15,7 @@ import { CancellationToken } from 'vs/base/common/cancellation';
 import { NotebookInput } from 'sql/parts/notebook/notebookInput';
 import { NotebookModule } from 'sql/parts/notebook/notebook.module';
 import { NOTEBOOK_SELECTOR } from 'sql/parts/notebook/notebook.component';
-import { INotebookParams, DEFAULT_NOTEBOOK_PROVIDER } from 'sql/workbench/services/notebook/common/notebookService';
+import { INotebookParams } from 'sql/workbench/services/notebook/common/notebookService';
 import { IStorageService } from 'vs/platform/storage/common/storage';
 import { $ } from 'sql/base/browser/builder';
 
@@ -92,8 +92,7 @@ export class NotebookEditor extends BaseEditor {
 		let params: INotebookParams = {
 			notebookUri: input.notebookUri,
 			input: input,
-			providerId: input.providerId ? input.providerId : DEFAULT_NOTEBOOK_PROVIDER,
-			providers: input.providers ? input.providers : [DEFAULT_NOTEBOOK_PROVIDER],
+			providerInfo: input.getProviderInfo(),
 			isTrusted: input.isTrusted,
 			profile: input.connectionProfile
 		};

src/sql/parts/notebook/notebookInput.ts

@@ -15,7 +15,7 @@ import * as resources from 'vs/base/common/resources';
 import * as azdata from 'azdata';
 
 import { IStandardKernelWithProvider, getProvidersForFileName, getStandardKernelsForProvider } from 'sql/parts/notebook/notebookUtils';
-import { INotebookService, DEFAULT_NOTEBOOK_PROVIDER } from 'sql/workbench/services/notebook/common/notebookService';
+import { INotebookService, DEFAULT_NOTEBOOK_PROVIDER, IProviderInfo } from 'sql/workbench/services/notebook/common/notebookService';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { ITextModelService } from 'vs/editor/common/services/resolverService';
 import { INotebookModel, IContentManager } from 'sql/parts/notebook/models/modelInterfaces';
@@ -29,6 +29,7 @@ import { ITextFileService, ISaveOptions } from 'vs/workbench/services/textfile/c
 import { LocalContentManager } from 'sql/workbench/services/notebook/node/localContentManager';
 import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorInput';
+import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
 
 export type ModeViewSaveHandler = (handle: number) => Thenable<boolean>;
 
@@ -140,6 +141,7 @@ export class NotebookInput extends EditorInput {
 	private _model: NotebookEditorModel;
 	private _untitledEditorService: IUntitledEditorService;
 	private _contentManager: IContentManager;
+	private _providersLoaded: Promise<void>;
 
 	constructor(private _title: string,
 		private resource: URI,
@@ -147,13 +149,14 @@ export class NotebookInput extends EditorInput {
 		@ITextModelService private textModelService: ITextModelService,
 		@IUntitledEditorService untitledEditorService: IUntitledEditorService,
 		@IInstantiationService private instantiationService: IInstantiationService,
-		@INotebookService private notebookService: INotebookService
+		@INotebookService private notebookService: INotebookService,
+		@IExtensionService private extensionService: IExtensionService
 	) {
 		super();
 		this._untitledEditorService = untitledEditorService;
 		this.resource = resource;
 		this._standardKernels = [];
-		this.assignProviders();
+		this._providersLoaded = this.assignProviders();
 	}
 
 	public get textInput(): UntitledEditorInput {
@@ -186,14 +189,13 @@ export class NotebookInput extends EditorInput {
 		return this._title;
 	}
 
-	public get providerId(): string {
-		return this._providerId;
+	public async getProviderInfo(): Promise<IProviderInfo> {
+		await this._providersLoaded;
+		return {
+			providerId: this._providerId ? this._providerId : DEFAULT_NOTEBOOK_PROVIDER,
+			providers: this._providers ? this._providers : [DEFAULT_NOTEBOOK_PROVIDER]
+		};
 	}
-
-	public set providerId(value: string) {
-		this._providerId = value;
-	}
-
 	public get isTrusted(): boolean {
 		return this._isTrusted;
 	}
@@ -214,14 +216,6 @@ export class NotebookInput extends EditorInput {
 		return this._standardKernels;
 	}
 
-	public get providers(): string[] {
-		return this._providers;
-	}
-
-	public set providers(value: string[]) {
-		this._providers = value;
-	}
-
 	public save(): TPromise<boolean> {
 		let options: ISaveOptions = { force: false };
 		return this._model.save(options);
@@ -280,7 +274,8 @@ export class NotebookInput extends EditorInput {
 		}
 	}
 
-	private assignProviders(): void {
+	private async assignProviders(): Promise<void> {
+		await this.extensionService.whenInstalledExtensionsRegistered();
 		let providerIds: string[] = getProvidersForFileName(this._title, this.notebookService);
 		if (providerIds && providerIds.length > 0) {
 			this._providerId = providerIds.filter(provider => provider !== DEFAULT_NOTEBOOK_PROVIDER)[0];

src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts

@@ -22,8 +22,8 @@ import {
 	SqlMainContext, MainThreadNotebookDocumentsAndEditorsShape, SqlExtHostContext, ExtHostNotebookDocumentsAndEditorsShape,
 	INotebookDocumentsAndEditorsDelta, INotebookEditorAddData, INotebookShowOptions, INotebookModelAddedData, INotebookModelChangedData
 } from 'sql/workbench/api/node/sqlExtHost.protocol';
-import { NotebookInput, NotebookEditorModel } from 'sql/parts/notebook/notebookInput';
-import { INotebookService, INotebookEditor, DEFAULT_NOTEBOOK_PROVIDER } from 'sql/workbench/services/notebook/common/notebookService';
+import { NotebookInput } from 'sql/parts/notebook/notebookInput';
+import { INotebookService, INotebookEditor, IProviderInfo } from 'sql/workbench/services/notebook/common/notebookService';
 import { TPromise } from 'vs/base/common/winjs.base';
 import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
 import { disposed } from 'vs/base/common/errors';
@@ -37,6 +37,7 @@ import { UntitledEditorInput } from 'vs/workbench/common/editor/untitledEditorIn
 class MainThreadNotebookEditor extends Disposable {
 	private _contentChangedEmitter = new Emitter<NotebookContentChange>();
 	public readonly contentChanged: Event<NotebookContentChange> = this._contentChangedEmitter.event;
+	private _providerInfo: IProviderInfo;
 
 	constructor(public readonly editor: INotebookEditor) {
 		super();
@@ -49,6 +50,9 @@ class MainThreadNotebookEditor extends Disposable {
 				this._contentChangedEmitter.fire(changeEvent);
 			}));
 		});
+		editor.notebookParams.providerInfo.then(info => {
+			this._providerInfo = info;
+		});
 	}
 
 	public get uri(): URI {
@@ -64,11 +68,11 @@ class MainThreadNotebookEditor extends Disposable {
 	}
 
 	public get providerId(): string {
-		return this.editor.notebookParams.providerId;
+		return this._providerInfo ? this._providerInfo.providerId : undefined;
 	}
 
 	public get providers(): string[] {
-		return this.editor.notebookParams.providers;
+		return this._providerInfo ? this._providerInfo.providers : [];
 	}
 
 	public get cells(): ICellModel[] {

src/sql/workbench/services/notebook/common/notebookService.ts

@@ -85,11 +85,14 @@ export interface INotebookManager {
 	readonly serverManager: azdata.nb.ServerManager;
 }
 
+export interface IProviderInfo {
+	providerId: string;
+	providers: string[];
+}
 export interface INotebookParams extends IBootstrapParams {
 	notebookUri: URI;
 	input: NotebookInput;
-	providerId: string;
-	providers: string[];
+	providerInfo: Promise<IProviderInfo>;
 	isTrusted: boolean;
 	profile?: IConnectionProfile;
 	modelFactory?: ModelFactory;