Notebooks: Use new Python installation after configuration change (#14765)

* start new jupyter server

* restart session working (removed extra code)

* only restart server once

* shutdown session first then stop server

* add comments remove extra lines

* add comment

* fix test

* only restart jupyter sessions

* Dispose jupytersessionmanager and create new one

* move restart server logic out of notebookmodel

* move methods to azdata proposed

* pr comment

extensions/notebook/src/jupyter/jupyterNotebookManager.ts

@@ -15,8 +15,7 @@ export class JupyterNotebookManager implements nb.NotebookManager, vscode.Dispos
 	private _sessionManager: JupyterSessionManager;
 
 	constructor(private _serverManager: LocalJupyterServerManager, sessionManager?: JupyterSessionManager) {
-		let pythonEnvVarPath = this._serverManager && this._serverManager.jupyterServerInstallation && this._serverManager.jupyterServerInstallation.pythonEnvVarPath;
+		this._sessionManager = sessionManager || new JupyterSessionManager();
-		this._sessionManager = sessionManager || new JupyterSessionManager(pythonEnvVarPath);
 		this._serverManager.onServerStarted(() => {
 			this.setServerSettings(this._serverManager.serverSettings);
 			this._sessionManager.installation = this._serverManager.instanceOptions.install;

extensions/notebook/src/jupyter/jupyterServerInstallation.ts

@@ -438,6 +438,7 @@ export class JupyterServerInstallation implements IJupyterServerInstallation {
 
 						this._installCompletion.resolve();
 						this._installInProgress = false;
+						await vscode.commands.executeCommand('notebook.action.restartJupyterNotebookSessions');
 					})
 					.catch(err => {
 						let errorMsg = msgDependenciesInstallationFailed(utils.getErrorMessage(err));

extensions/notebook/src/jupyter/jupyterSessionManager.ts

@@ -65,7 +65,7 @@ export class JupyterSessionManager implements nb.SessionManager {
 	private static _sessions: JupyterSession[] = [];
 	private _installation: JupyterServerInstallation;
 
-	constructor(private _pythonEnvVarPath?: string) {
+	constructor() {
 		this._isReady = false;
 		this._ready = new Deferred<void>();
 	}
@@ -130,7 +130,7 @@ export class JupyterSessionManager implements nb.SessionManager {
 			return Promise.reject(new Error(localize('errorStartBeforeReady', "Cannot start a session, the manager is not yet initialized")));
 		}
 		let sessionImpl = await this._sessionManager.startNew(options);
-		let jupyterSession = new JupyterSession(sessionImpl, this._installation, skipSettingEnvironmentVars, this._pythonEnvVarPath);
+		let jupyterSession = new JupyterSession(sessionImpl, this._installation, skipSettingEnvironmentVars, this._installation?.pythonEnvVarPath);
 		await jupyterSession.messagesComplete;
 		let index = JupyterSessionManager._sessions.findIndex(session => session.path === options.path);
 		if (index > -1) {

src/sql/azdata.proposed.d.ts

@@ -51,6 +51,16 @@ declare module 'azdata' {
 		export type ICellAttachments = { [key: string]: ICellAttachment };
 		export type ICellAttachment = { [key: string]: string };
 
+		export interface SessionManager {
+			/**
+			 * Shutdown all sessions.
+			 */
+			shutdownAll(): Thenable<void>;
+			/**
+			 * Disposes the session manager.
+			 */
+			dispose(): void;
+		}
 	}
 
 	export type SqlDbType = 'BigInt' | 'Binary' | 'Bit' | 'Char' | 'DateTime' | 'Decimal'

src/sql/workbench/api/browser/mainThreadNotebook.ts

@@ -255,6 +255,14 @@ class SessionManagerWrapper implements azdata.nb.SessionManager {
 			this._specs = specs;
 		}
 	}
+
+	shutdownAll(): Thenable<void> {
+		return this._proxy.ext.$shutdownAll(this.managerHandle);
+	}
+
+	dispose(): void {
+		return this._proxy.ext.$dispose(this.managerHandle);
+	}
 }
 
 class SessionWrapper implements azdata.nb.ISession {

src/sql/workbench/api/common/extHostNotebook.ts

@@ -126,6 +126,12 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 		});
 	}
 
+	$shutdownAll(managerHandle: number): Thenable<void> {
+		return this._withSessionManager(managerHandle, async (sessionManager) => {
+			return sessionManager.shutdownAll();
+		});
+	}
+
 	$changeKernel(sessionId: number, kernelInfo: azdata.nb.IKernelSpec): Thenable<INotebookKernelDetails> {
 		let session = this._getAdapter<azdata.nb.ISession>(sessionId);
 		return session.changeKernel(kernelInfo).then(kernel => this.saveKernel(kernel));
@@ -207,6 +213,12 @@ export class ExtHostNotebook implements ExtHostNotebookShape {
 		future.dispose();
 	}
 
+	$dispose(managerHandle: number): Thenable<void> {
+		return this._withSessionManager(managerHandle, async (sessionManager) => {
+			return sessionManager.dispose();
+		});
+	}
+
 	//#endregion
 
 	//#region APIs called by extensions

src/sql/workbench/api/common/sqlExtHost.protocol.ts

@@ -870,6 +870,8 @@ export interface ExtHostNotebookShape {
 	$refreshSpecs(managerHandle: number): Thenable<azdata.nb.IAllKernels>;
 	$startNewSession(managerHandle: number, options: azdata.nb.ISessionOptions): Thenable<INotebookSessionDetails>;
 	$shutdownSession(managerHandle: number, sessionId: string): Thenable<void>;
+	$shutdownAll(managerHandle: number): Thenable<void>;
+	$dispose(managerHandle: number): void;
 
 	// Session APIs
 	$changeKernel(sessionId: number, kernelInfo: azdata.nb.IKernelSpec): Thenable<INotebookKernelDetails>;

src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts

@@ -254,6 +254,10 @@ export abstract class NotebookInput extends EditorInput {
 		return this.resource;
 	}
 
+	public get notebookModel(): INotebookModel | undefined {
+		return this._model.getNotebookModel();
+	}
+
 	public get notebookFindModel(): NotebookFindModel {
 		if (!this._notebookFindModel) {
 			this._notebookFindModel = new NotebookFindModel(this._model.getNotebookModel());

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

@@ -7,7 +7,7 @@ import { EditorDescriptor, IEditorRegistry, Extensions as EditorExtensions } fro
 import { SyncDescriptor } from 'vs/platform/instantiation/common/descriptors';
 import { IInstantiationService, ServicesAccessor } from 'vs/platform/instantiation/common/instantiation';
 import { localize } from 'vs/nls';
-import { IEditorInputFactoryRegistry, Extensions as EditorInputFactoryExtensions, ActiveEditorContext } from 'vs/workbench/common/editor';
+import { IEditorInputFactoryRegistry, Extensions as EditorInputFactoryExtensions, ActiveEditorContext, IEditorInput } from 'vs/workbench/common/editor';
 
 import { ILanguageAssociationRegistry, Extensions as LanguageAssociationExtensions } from 'sql/workbench/services/languageAssociation/common/languageAssociation';
 import { UntitledNotebookInput } from 'sql/workbench/contrib/notebook/browser/models/untitledNotebookInput';
@@ -33,7 +33,7 @@ import { MssqlNodeContext } from 'sql/workbench/services/objectExplorer/browser/
 import { mssqlProviderName } from 'sql/platform/connection/common/constants';
 import { CommandsRegistry } from 'vs/platform/commands/common/commands';
 import { TreeViewItemHandleArg } from 'sql/workbench/common/views';
-import { ConnectedContext } from 'azdata';
+import { ConnectedContext, nb } from 'azdata';
 import { TreeNodeContextKey } from 'sql/workbench/services/objectExplorer/common/treeNodeContextKey';
 import { ObjectExplorerActionsContext } from 'sql/workbench/services/objectExplorer/browser/objectExplorerActions';
 import { ItemContextKey } from 'sql/workbench/contrib/dashboard/browser/widgets/explorer/explorerContext';
@@ -52,6 +52,9 @@ import { isMacintosh } from 'vs/base/common/platform';
 import { SearchSortOrder } from 'vs/workbench/services/search/common/search';
 import { ImageMimeTypes } from 'sql/workbench/services/notebook/common/contracts';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
+import { NotebookInput } from 'sql/workbench/contrib/notebook/browser/models/notebookInput';
+import { INotebookModel } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
+import { INotebookManager } from 'sql/workbench/services/notebook/browser/notebookService';
 
 Registry.as<IEditorInputFactoryRegistry>(EditorInputFactoryExtensions.EditorInputFactories)
 	.registerEditorInputFactory(FileNotebookInput.ID, FileNoteBookEditorInputFactory);
@@ -167,6 +170,40 @@ CommandsRegistry.registerCommand({
 	}
 });
 
+const RESTART_JUPYTER_NOTEBOOK_SESSIONS = 'notebook.action.restartJupyterNotebookSessions';
+
+CommandsRegistry.registerCommand({
+	id: RESTART_JUPYTER_NOTEBOOK_SESSIONS,
+	handler: async (accessor: ServicesAccessor) => {
+		const editorService: IEditorService = accessor.get(IEditorService);
+		const editors: readonly IEditorInput[] = editorService.editors;
+		let jupyterServerRestarted: boolean = false;
+
+		for (let editor of editors) {
+			if (editor instanceof NotebookInput) {
+				let model: INotebookModel = editor.notebookModel;
+				if (model.providerId === 'jupyter') {
+					// Jupyter server needs to be restarted so that the correct Python installation is used
+					if (!jupyterServerRestarted) {
+						let jupyterNotebookManager: INotebookManager = model.notebookManagers.find(x => x.providerId === 'jupyter');
+						// Shutdown all current Jupyter sessions before stopping the server
+						await jupyterNotebookManager.sessionManager.shutdownAll();
+						// Jupyter session manager needs to be disposed so that a new one is created with the new server info
+						jupyterNotebookManager.sessionManager.dispose();
+						await jupyterNotebookManager.serverManager.stopServer();
+						let spec: nb.IKernelSpec = model.defaultKernel;
+						await jupyterNotebookManager.serverManager.startServer(spec);
+						jupyterServerRestarted = true;
+					}
+
+					// Start a new session for each Jupyter notebook
+					await model.restartSession();
+				}
+			}
+		}
+	}
+});
+
 MenuRegistry.appendMenuItem(MenuId.CommandPalette, {
 	command: {
 		id: TOGGLE_TAB_FOCUS_COMMAND_ID,

src/sql/workbench/contrib/notebook/test/emptySessionClasses.ts

@@ -42,6 +42,13 @@ export class SessionManager implements nb.SessionManager {
 	shutdown(id: string): Thenable<void> {
 		return Promise.resolve();
 	}
+
+	shutdownAll(): Thenable<void> {
+		return Promise.resolve();
+	}
+
+	dispose(): void {
+	}
 }
 
 export class EmptySession implements nb.ISession {

src/sql/workbench/contrib/notebook/test/stubs.ts

@@ -88,6 +88,9 @@ export class NotebookModelStub implements INotebookModel {
 	getStandardKernelFromName(name: string): IStandardKernelWithProvider {
 		throw new Error('Method not implemented.');
 	}
+	restartSession(): Promise<void> {
+		throw new Error('Method not implemented.');
+	}
 	changeKernel(displayName: string): void {
 		throw new Error('Method not implemented.');
 	}

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

@@ -361,6 +361,11 @@ export interface INotebookModel {
 	 */
 	getMetaValue(key: string): any;
 
+	/**
+	 * Restart current active session if it exists
+	 */
+	restartSession(): Promise<void>;
+
 	/**
 	 * Change the current kernel from the Kernel dropdown
 	 * @param displayName kernel name (as displayed in Kernel dropdown)
@@ -387,7 +392,6 @@ export interface INotebookModel {
 	 */
 	moveCell(cellModel: ICellModel, direction: MoveDirection): void;
 
-
 	/**
 	 * Deletes a cell
 	 */
@@ -403,7 +407,6 @@ export interface INotebookModel {
 	 */
 	onCellChange(cell: ICellModel, change: NotebookChangeType): void;
 
-
 	/**
 	 * Push edit operations, basically editing the model. This is the preferred way of
 	 * editing the model. Long-term, this will ensure edit operations can be added to the undo stack

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

@@ -741,6 +741,14 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		}
 	}
 
+	public async restartSession(): Promise<void> {
+		if (this._activeClientSession) {
+			// Old active client sessions have already been shutdown by RESTART_JUPYTER_NOTEBOOK_SESSIONS command
+			this._activeClientSession = undefined;
+			await this.startSession(this.notebookManager, this._selectedKernelDisplayName, true);
+		}
+	}
+
 	// When changing kernel, update the active session
 	private updateActiveClientSession(clientSession: IClientSession) {
 		this._activeClientSession = clientSession;
@@ -1114,7 +1122,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		}
 	}
 
-	private async shutdownActiveSession() {
+	private async shutdownActiveSession(): Promise<void> {
 		if (this._activeClientSession) {
 			try {
 				await this._activeClientSession.ready;

src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts

@@ -133,6 +133,16 @@ export class SqlSessionManager implements nb.SessionManager {
 		}
 		return Promise.resolve();
 	}
+
+	shutdownAll(): Thenable<void> {
+		return Promise.all(SqlSessionManager._sessions.map(session => {
+			return this.shutdown(session.id);
+		})).then();
+	}
+
+	dispose(): void {
+		// no-op
+	}
 }
 
 export class SqlSession implements nb.ISession {

src/sql/workbench/test/electron-browser/api/mainThreadNotebook.test.ts

@@ -161,6 +161,9 @@ class ExtHostNotebookStub implements ExtHostNotebookShape {
 	$shutdownSession(managerHandle: number, sessionId: string): Thenable<void> {
 		throw new Error('Method not implemented.');
 	}
+	$shutdownAll(managerHandle: number): Thenable<void> {
+		throw new Error('Method not implemented.');
+	}
 	$changeKernel(sessionId: number, kernelInfo: azdata.nb.IKernelSpec): Thenable<INotebookKernelDetails> {
 		throw new Error('Method not implemented.');
 	}
@@ -191,4 +194,7 @@ class ExtHostNotebookStub implements ExtHostNotebookShape {
 	$disposeFuture(futureId: number): void {
 		throw new Error('Method not implemented.');
 	}
+	$dispose(): void {
+		throw new Error('Method not implemented.');
+	}
 }