diff --git a/extensions/notebook/src/jupyter/jupyterNotebookManager.ts b/extensions/notebook/src/jupyter/jupyterNotebookManager.ts index d1e326b185..3e79ef1703 100644 --- a/extensions/notebook/src/jupyter/jupyterNotebookManager.ts +++ b/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(pythonEnvVarPath); + this._sessionManager = sessionManager || new JupyterSessionManager(); this._serverManager.onServerStarted(() => { this.setServerSettings(this._serverManager.serverSettings); this._sessionManager.installation = this._serverManager.instanceOptions.install; diff --git a/extensions/notebook/src/jupyter/jupyterServerInstallation.ts b/extensions/notebook/src/jupyter/jupyterServerInstallation.ts index 2a3f3895aa..949ec031ea 100644 --- a/extensions/notebook/src/jupyter/jupyterServerInstallation.ts +++ b/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)); diff --git a/extensions/notebook/src/jupyter/jupyterSessionManager.ts b/extensions/notebook/src/jupyter/jupyterSessionManager.ts index 7056677b63..9be5de2a50 100644 --- a/extensions/notebook/src/jupyter/jupyterSessionManager.ts +++ b/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(); } @@ -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) { diff --git a/src/sql/azdata.proposed.d.ts b/src/sql/azdata.proposed.d.ts index ac6e5d6db9..f99e13973c 100644 --- a/src/sql/azdata.proposed.d.ts +++ b/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; + /** + * Disposes the session manager. + */ + dispose(): void; + } } export type SqlDbType = 'BigInt' | 'Binary' | 'Bit' | 'Char' | 'DateTime' | 'Decimal' diff --git a/src/sql/workbench/api/browser/mainThreadNotebook.ts b/src/sql/workbench/api/browser/mainThreadNotebook.ts index 37c2ca305a..c0d923ba43 100644 --- a/src/sql/workbench/api/browser/mainThreadNotebook.ts +++ b/src/sql/workbench/api/browser/mainThreadNotebook.ts @@ -255,6 +255,14 @@ class SessionManagerWrapper implements azdata.nb.SessionManager { this._specs = specs; } } + + shutdownAll(): Thenable { + return this._proxy.ext.$shutdownAll(this.managerHandle); + } + + dispose(): void { + return this._proxy.ext.$dispose(this.managerHandle); + } } class SessionWrapper implements azdata.nb.ISession { diff --git a/src/sql/workbench/api/common/extHostNotebook.ts b/src/sql/workbench/api/common/extHostNotebook.ts index 8380a9d054..055f1b892c 100644 --- a/src/sql/workbench/api/common/extHostNotebook.ts +++ b/src/sql/workbench/api/common/extHostNotebook.ts @@ -126,6 +126,12 @@ export class ExtHostNotebook implements ExtHostNotebookShape { }); } + $shutdownAll(managerHandle: number): Thenable { + return this._withSessionManager(managerHandle, async (sessionManager) => { + return sessionManager.shutdownAll(); + }); + } + $changeKernel(sessionId: number, kernelInfo: azdata.nb.IKernelSpec): Thenable { let session = this._getAdapter(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 { + return this._withSessionManager(managerHandle, async (sessionManager) => { + return sessionManager.dispose(); + }); + } + //#endregion //#region APIs called by extensions diff --git a/src/sql/workbench/api/common/sqlExtHost.protocol.ts b/src/sql/workbench/api/common/sqlExtHost.protocol.ts index 069c91000e..412febb3cf 100644 --- a/src/sql/workbench/api/common/sqlExtHost.protocol.ts +++ b/src/sql/workbench/api/common/sqlExtHost.protocol.ts @@ -870,6 +870,8 @@ export interface ExtHostNotebookShape { $refreshSpecs(managerHandle: number): Thenable; $startNewSession(managerHandle: number, options: azdata.nb.ISessionOptions): Thenable; $shutdownSession(managerHandle: number, sessionId: string): Thenable; + $shutdownAll(managerHandle: number): Thenable; + $dispose(managerHandle: number): void; // Session APIs $changeKernel(sessionId: number, kernelInfo: azdata.nb.IKernelSpec): Thenable; diff --git a/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts b/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts index a78a5b1f54..fedd335173 100644 --- a/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts +++ b/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()); diff --git a/src/sql/workbench/contrib/notebook/browser/notebook.contribution.ts b/src/sql/workbench/contrib/notebook/browser/notebook.contribution.ts index 9ae4af77c0..0f15ace551 100644 --- a/src/sql/workbench/contrib/notebook/browser/notebook.contribution.ts +++ b/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(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, diff --git a/src/sql/workbench/contrib/notebook/test/emptySessionClasses.ts b/src/sql/workbench/contrib/notebook/test/emptySessionClasses.ts index 48578e87af..c2cb747034 100644 --- a/src/sql/workbench/contrib/notebook/test/emptySessionClasses.ts +++ b/src/sql/workbench/contrib/notebook/test/emptySessionClasses.ts @@ -42,6 +42,13 @@ export class SessionManager implements nb.SessionManager { shutdown(id: string): Thenable { return Promise.resolve(); } + + shutdownAll(): Thenable { + return Promise.resolve(); + } + + dispose(): void { + } } export class EmptySession implements nb.ISession { diff --git a/src/sql/workbench/contrib/notebook/test/stubs.ts b/src/sql/workbench/contrib/notebook/test/stubs.ts index f13dedcd83..2014cbcb6b 100644 --- a/src/sql/workbench/contrib/notebook/test/stubs.ts +++ b/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 { + throw new Error('Method not implemented.'); + } changeKernel(displayName: string): void { throw new Error('Method not implemented.'); } diff --git a/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts b/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts index bd250a8036..48a1d36d4a 100644 --- a/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts +++ b/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; + /** * 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 diff --git a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts index e6ea4f6c6c..f4da200494 100644 --- a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts +++ b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts @@ -741,6 +741,14 @@ export class NotebookModel extends Disposable implements INotebookModel { } } + public async restartSession(): Promise { + 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 { if (this._activeClientSession) { try { await this._activeClientSession.ready; diff --git a/src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts b/src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts index f5baa1d74a..a5d54cf563 100644 --- a/src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts +++ b/src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts @@ -133,6 +133,16 @@ export class SqlSessionManager implements nb.SessionManager { } return Promise.resolve(); } + + shutdownAll(): Thenable { + return Promise.all(SqlSessionManager._sessions.map(session => { + return this.shutdown(session.id); + })).then(); + } + + dispose(): void { + // no-op + } } export class SqlSession implements nb.ISession { diff --git a/src/sql/workbench/test/electron-browser/api/mainThreadNotebook.test.ts b/src/sql/workbench/test/electron-browser/api/mainThreadNotebook.test.ts index 28aeb3f935..0654f91e6b 100644 --- a/src/sql/workbench/test/electron-browser/api/mainThreadNotebook.test.ts +++ b/src/sql/workbench/test/electron-browser/api/mainThreadNotebook.test.ts @@ -161,6 +161,9 @@ class ExtHostNotebookStub implements ExtHostNotebookShape { $shutdownSession(managerHandle: number, sessionId: string): Thenable { throw new Error('Method not implemented.'); } + $shutdownAll(managerHandle: number): Thenable { + throw new Error('Method not implemented.'); + } $changeKernel(sessionId: number, kernelInfo: azdata.nb.IKernelSpec): Thenable { 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.'); + } }