diff --git a/extensions/notebook/src/jupyter/jupyterController.ts b/extensions/notebook/src/jupyter/jupyterController.ts index d68767d5c2..483596b7ae 100644 --- a/extensions/notebook/src/jupyter/jupyterController.ts +++ b/extensions/notebook/src/jupyter/jupyterController.ts @@ -19,7 +19,7 @@ import { IPrompter, QuestionTypes, IQuestion } from '../prompts/question'; import { AppContext } from '../common/appContext'; import { ApiWrapper } from '../common/apiWrapper'; -import { LocalJupyterServerManager } from './jupyterServerManager'; +import { LocalJupyterServerManager, ServerInstanceFactory } from './jupyterServerManager'; import { NotebookCompletionItemProvider } from '../intellisense/completionItemProvider'; import { JupyterNotebookProvider } from './jupyterNotebookProvider'; import { ConfigurePythonDialog } from '../dialog/configurePythonDialog'; @@ -31,6 +31,7 @@ let untitledCounter = 0; export class JupyterController implements vscode.Disposable { private _jupyterInstallation: JupyterServerInstallation; private _notebookInstances: IServerInstance[] = []; + private _serverInstanceFactory: ServerInstanceFactory = new ServerInstanceFactory(); private outputChannel: vscode.OutputChannel; private prompter: IPrompter; @@ -92,7 +93,8 @@ export class JupyterController implements vscode.Disposable { documentPath: documentUri.fsPath, jupyterInstallation: this._jupyterInstallation, extensionContext: this.extensionContext, - apiWrapper: this.apiWrapper + apiWrapper: this.apiWrapper, + factory: this._serverInstanceFactory })); azdata.nb.registerNotebookProvider(notebookProvider); return notebookProvider; diff --git a/extensions/notebook/src/jupyter/jupyterNotebookManager.ts b/extensions/notebook/src/jupyter/jupyterNotebookManager.ts index 273008bbb7..a12fbedacc 100644 --- a/extensions/notebook/src/jupyter/jupyterNotebookManager.ts +++ b/extensions/notebook/src/jupyter/jupyterNotebookManager.ts @@ -53,4 +53,4 @@ export class JupyterNotebookManager implements nb.NotebookManager, vscode.Dispos this._serverManager.stopServer().then(success => undefined, error => this.apiWrapper.showErrorMessage(error)); } } -} \ No newline at end of file +} diff --git a/extensions/notebook/src/jupyter/jupyterNotebookProvider.ts b/extensions/notebook/src/jupyter/jupyterNotebookProvider.ts index 6571f3821b..4ddbd017b2 100644 --- a/extensions/notebook/src/jupyter/jupyterNotebookProvider.ts +++ b/extensions/notebook/src/jupyter/jupyterNotebookProvider.ts @@ -3,16 +3,17 @@ * Licensed under the Source EULA. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -'use strict'; - import { nb } from 'azdata'; import * as vscode from 'vscode'; import * as nls from 'vscode-nls'; +import * as path from 'path'; const localize = nls.loadMessageBundle(); import * as constants from '../common/constants'; +import * as utils from '../common/utils'; import { JupyterNotebookManager } from './jupyterNotebookManager'; import { LocalJupyterServerManager } from './jupyterServerManager'; +import { JupyterSessionManager } from './jupyterSessionManager'; export type ServerManagerFactory = (documentUri: vscode.Uri) => LocalJupyterServerManager; @@ -31,12 +32,16 @@ export class JupyterNotebookProvider implements nb.NotebookProvider { } private doGetNotebookManager(notebookUri: vscode.Uri): nb.NotebookManager { - let uriString = notebookUri.toString(); - let manager = this.managerTracker.get(uriString); + let baseFolder = this.transformToBaseFolder(notebookUri.fsPath.toString()); + let manager = this.managerTracker.get(baseFolder); if (!manager) { - let serverManager = this.createServerManager(notebookUri); + let baseFolderUri = vscode.Uri.file(baseFolder); + if (!baseFolderUri) { + baseFolderUri = notebookUri; + } + let serverManager = this.createServerManager(baseFolderUri); manager = new JupyterNotebookManager(serverManager); - this.managerTracker.set(uriString, manager); + this.managerTracker.set(baseFolder, manager); } return manager; } @@ -46,15 +51,52 @@ export class JupyterNotebookProvider implements nb.NotebookProvider { // As this is a notification method, will skip throwing an error here return; } - let uriString = notebookUri.toString(); - let manager = this.managerTracker.get(uriString); + let baseFolder = this.transformToBaseFolder(notebookUri.fsPath.toString()); + let manager = this.managerTracker.get(baseFolder); if (manager) { - this.managerTracker.delete(uriString); - manager.dispose(); + let sessionManager = (manager.sessionManager as JupyterSessionManager); + let session = sessionManager.listRunning().find(e => e.path === notebookUri.fsPath); + if (session) { + manager.sessionManager.shutdown(session.id); + } + if (sessionManager.listRunning().length === 0) { + const FiveMinutesInMs = 5 * 60 * 1000; + setTimeout(() => { + if (sessionManager.listRunning().length === 0) { + this.managerTracker.delete(baseFolder); + manager.dispose(); + } + }, FiveMinutesInMs); + } } } public get standardKernels(): nb.IStandardKernel[] { return []; } + + private transformToBaseFolder(notebookPath: string): string { + let parsedPath = path.parse(notebookPath); + let userHome = utils.getUserHome(); + let relativePathStrUserHome = path.relative(notebookPath, userHome); + if (notebookPath === '.' || relativePathStrUserHome.endsWith('..') || relativePathStrUserHome === '') { + // If you don't match the notebookPath's casing for drive letters, + // a 404 will result when trying to create a new session on Windows + if (notebookPath && userHome && notebookPath[0].toLowerCase() === userHome[0].toLowerCase()) { + userHome = notebookPath[0] + userHome.substr(1); + } + // If the user is using a system version of python, then + // '.' will try to create a notebook in a system directory. + // Since this will fail due to permissions issues, use the user's + // home folder instead. + return userHome; + } else { + let splitDirName: string[] = path.dirname(notebookPath).split(path.sep); + if (splitDirName.length > 1) { + return path.join(parsedPath.root, splitDirName[1]); + } else { + return userHome; + } + } + } } diff --git a/extensions/notebook/src/jupyter/jupyterServerManager.ts b/extensions/notebook/src/jupyter/jupyterServerManager.ts index 4f6907b3e0..733dfa008d 100644 --- a/extensions/notebook/src/jupyter/jupyterServerManager.ts +++ b/extensions/notebook/src/jupyter/jupyterServerManager.ts @@ -16,7 +16,7 @@ import { ApiWrapper } from '../common/apiWrapper'; import { JupyterServerInstallation } from './jupyterServerInstallation'; import * as utils from '../common/utils'; import { IServerInstance } from './common'; -import { PerNotebookServerInstance, IInstanceOptions } from './serverInstance'; +import { PerFolderServerInstance, IInstanceOptions } from './serverInstance'; import { CommandContext } from '../common/constants'; export interface IServerManagerOptions { @@ -30,11 +30,11 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo private _serverSettings: Partial; private _onServerStarted = new vscode.EventEmitter(); private _instanceOptions: IInstanceOptions; - private apiWrapper: ApiWrapper; - private jupyterServer: IServerInstance; + private _apiWrapper: ApiWrapper; + private _jupyterServer: IServerInstance; factory: ServerInstanceFactory; constructor(private options: IServerManagerOptions) { - this.apiWrapper = options.apiWrapper || new ApiWrapper(); + this._apiWrapper = options.apiWrapper || new ApiWrapper(); this.factory = options.factory || new ServerInstanceFactory(); } @@ -43,7 +43,7 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo } public get isStarted(): boolean { - return !!this.jupyterServer; + return !!this._jupyterServer; } public get instanceOptions(): IInstanceOptions { @@ -56,11 +56,13 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo public async startServer(): Promise { try { - this.jupyterServer = await this.doStartServer(); - this.options.extensionContext.subscriptions.push(this); - let partialSettings = LocalJupyterServerManager.getLocalConnectionSettings(this.jupyterServer.uri); - this._serverSettings = partialSettings; - this._onServerStarted.fire(); + if (!this._jupyterServer) { + this._jupyterServer = await this.doStartServer(); + this.options.extensionContext.subscriptions.push(this); + let partialSettings = LocalJupyterServerManager.getLocalConnectionSettings(this._jupyterServer.uri); + this._serverSettings = partialSettings; + this._onServerStarted.fire(); + } } catch (error) { // this is caught and notified up the stack, no longer showing a message here @@ -71,13 +73,13 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo public dispose(): void { this.stopServer().catch(err => { let msg = utils.getErrorMessage(err); - this.apiWrapper.showErrorMessage(localize('shutdownError', 'Shutdown of Notebook server failed: {0}', msg)); + this._apiWrapper.showErrorMessage(localize('shutdownError', 'Shutdown of Notebook server failed: {0}', msg)); }); } public async stopServer(): Promise { - if (this.jupyterServer) { - await this.jupyterServer.stop(); + if (this._jupyterServer) { + await this._jupyterServer.stop(); } } @@ -106,7 +108,7 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo let installation = this.options.jupyterInstallation; await installation.promptForPythonInstall(); await installation.promptForPackageUpgrade(); - this.apiWrapper.setCommandContext(CommandContext.NotebookPythonInstalled, true); + this._apiWrapper.setCommandContext(CommandContext.NotebookPythonInstalled, true); // Calculate the path to use as the notebook-dir for Jupyter based on the path of the uri of the // notebook to open. This will be the workspace folder if the notebook uri is inside a workspace @@ -118,9 +120,11 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo // /path2/nb2.ipynb // /path2/nb3.ipynb // ... will result in 2 notebook servers being started, one for /path1/ and one for /path2/ - let notebookDir = this.apiWrapper.getWorkspacePathFromUri(vscode.Uri.file(this.documentPath)); + let notebookDir = this._apiWrapper.getWorkspacePathFromUri(vscode.Uri.file(this.documentPath)); if (!notebookDir) { - let docDir = path.dirname(this.documentPath); + let docDir; + // If a folder is passed in for documentPath, use the folder instead of calling dirname + docDir = path.extname(this.documentPath) ? path.dirname(this.documentPath) : this.documentPath; if (docDir === '.') { // If the user is using a system version of python, then // '.' will try to create a notebook in a system directory. @@ -156,7 +160,7 @@ export class LocalJupyterServerManager implements nb.ServerManager, vscode.Dispo export class ServerInstanceFactory { createInstance(options: IInstanceOptions): IServerInstance { - return new PerNotebookServerInstance(options); + return new PerFolderServerInstance(options); } } diff --git a/extensions/notebook/src/jupyter/jupyterSessionManager.ts b/extensions/notebook/src/jupyter/jupyterSessionManager.ts index ea023d8b89..6b772f73d4 100644 --- a/extensions/notebook/src/jupyter/jupyterSessionManager.ts +++ b/extensions/notebook/src/jupyter/jupyterSessionManager.ts @@ -117,13 +117,14 @@ export class JupyterSessionManager implements nb.SessionManager { return allKernels; } - public async startNew(options: nb.ISessionOptions): Promise { + public async startNew(options: nb.ISessionOptions, skipSettingEnvironmentVars?: boolean): Promise { if (!this._isReady) { // no-op 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); + let jupyterSession = new JupyterSession(sessionImpl, skipSettingEnvironmentVars); + await jupyterSession.messagesComplete; let index = JupyterSessionManager._sessions.findIndex(session => session.path === options.path); if (index > -1) { JupyterSessionManager._sessions.splice(index); @@ -167,8 +168,10 @@ export class JupyterSessionManager implements nb.SessionManager { export class JupyterSession implements nb.ISession { private _kernel: nb.IKernel; + private _messagesComplete: Deferred = new Deferred(); - constructor(private sessionImpl: Session.ISession) { + constructor(private sessionImpl: Session.ISession, skipSettingEnvironmentVars?: boolean) { + this.setEnvironmentVars(skipSettingEnvironmentVars); } public get canChangeKernels(): boolean { @@ -205,6 +208,11 @@ export class JupyterSession implements nb.ISession { return this._kernel; } + // Sent when startup messages have been sent + public get messagesComplete(): Promise { + return this._messagesComplete.promise; + } + public async changeKernel(kernelInfo: nb.IKernelSpec): Promise { // For now, Jupyter implementation handles disposal etc. so we can just // null out our kernel and let the changeKernel call handle this @@ -321,6 +329,22 @@ export class JupyterSession implements nb.ISession { } return endpoints.find(ep => ep.serviceName.toLowerCase() === serviceName.toLowerCase()); } + + private async setEnvironmentVars(skip: boolean = false): Promise { + if (!skip && this.sessionImpl) { + let allCode: string = ''; + for (let i = 0; i < Object.keys(process.env).length; i++) { + let key = Object.keys(process.env)[i]; + // Jupyter doesn't seem to alow for setting multiple variables at once, so doing it with multiple commands + allCode += `%set_env ${key}=${process.env[key]}${EOL}`; + } + let future = this.sessionImpl.kernel.requestExecute({ + code: allCode + }, true); + await future.done; + } + this._messagesComplete.resolve(); + } } interface ICredentials { diff --git a/extensions/notebook/src/jupyter/serverInstance.ts b/extensions/notebook/src/jupyter/serverInstance.ts index 60bff8e4f8..fc7b995bd1 100644 --- a/extensions/notebook/src/jupyter/serverInstance.ts +++ b/extensions/notebook/src/jupyter/serverInstance.ts @@ -99,7 +99,7 @@ export class ServerInstanceUtils { } } -export class PerNotebookServerInstance implements IServerInstance { +export class PerFolderServerInstance implements IServerInstance { /** * Root of the jupyter directory structure. Config and data roots will be @@ -123,6 +123,7 @@ export class PerNotebookServerInstance implements IServerInstance { private _port: string; private _uri: vscode.Uri; private _isStarted: boolean = false; + private _isStopping: boolean = false; private utils: ServerInstanceUtils; private childProcess: ChildProcess; private errorHandler: ErrorHandler = new ErrorHandler(); @@ -132,7 +133,7 @@ export class PerNotebookServerInstance implements IServerInstance { } public get isStarted(): boolean { - return this._isStarted; + return this._isStarted && !this._isStopping; } public get port(): string { @@ -153,13 +154,14 @@ export class PerNotebookServerInstance implements IServerInstance { public async stop(): Promise { try { + this._isStopping = true; if (this.baseDir) { let exists = await this.utils.pathExists(this.baseDir); if (exists) { await this.utils.removeDir(this.baseDir); } } - if (this.isStarted) { + if (this._isStarted) { let install = this.options.install; let stopCommand = `"${install.pythonExecutable}" -m jupyter notebook stop ${this._port}`; await this.utils.executeBufferedCommand(stopCommand, install.execOptions, install.outputChannel); @@ -171,6 +173,7 @@ export class PerNotebookServerInstance implements IServerInstance { this._isStarted = false; this.utils.ensureProcessEnded(this.childProcess); this.handleConnectionClosed(); + } } diff --git a/extensions/notebook/src/test/model/serverInstance.test.ts b/extensions/notebook/src/test/model/serverInstance.test.ts index e2f282a4bd..05c4a9f926 100644 --- a/extensions/notebook/src/test/model/serverInstance.test.ts +++ b/extensions/notebook/src/test/model/serverInstance.test.ts @@ -11,7 +11,7 @@ import 'mocha'; import { JupyterServerInstallation } from '../../jupyter/jupyterServerInstallation'; import { ApiWrapper } from '../..//common/apiWrapper'; -import { PerNotebookServerInstance, ServerInstanceUtils } from '../../jupyter/serverInstance'; +import { PerFolderServerInstance, ServerInstanceUtils } from '../../jupyter/serverInstance'; import { MockOutputChannel } from '../common/stubs'; import * as testUtils from '../common/testUtils'; import { LocalJupyterServerManager } from '../../jupyter/jupyterServerManager'; @@ -27,7 +27,7 @@ describe('Jupyter server instance', function (): void { let mockOutputChannel: TypeMoq.IMock; let mockApiWrapper: TypeMoq.IMock; let mockUtils: TypeMoq.IMock; - let serverInstance: PerNotebookServerInstance; + let serverInstance: PerFolderServerInstance; beforeEach(() => { mockApiWrapper = TypeMoq.Mock.ofType(ApiWrapper); @@ -40,7 +40,7 @@ describe('Jupyter server instance', function (): void { mockInstall.object.execOptions = { env: Object.assign({}, process.env) }; mockUtils = TypeMoq.Mock.ofType(ServerInstanceUtils); mockUtils.setup(u => u.ensureProcessEnded(TypeMoq.It.isAny())).returns(() => undefined); - serverInstance = new PerNotebookServerInstance({ + serverInstance = new PerFolderServerInstance({ documentPath: expectedPath, install: mockInstall.object }, mockUtils.object); diff --git a/extensions/notebook/src/test/model/sessionManager.test.ts b/extensions/notebook/src/test/model/sessionManager.test.ts index 81a4f01b6c..7432b705b4 100644 --- a/extensions/notebook/src/test/model/sessionManager.test.ts +++ b/extensions/notebook/src/test/model/sessionManager.test.ts @@ -85,7 +85,7 @@ describe('Jupyter Session Manager', function (): void { mockJupyterManager.setup(m => m.startNew(TypeMoq.It.isAny())).returns(() => Promise.resolve(expectedSessionInfo)); // When I call startSession - let session = await sessionManager.startNew(sessionOptions); + let session = await sessionManager.startNew(sessionOptions, true); // Then I expect the parameters passed to be correct should(session.path).equal(sessionOptions.path); should(session.canChangeKernels).be.true(); @@ -147,7 +147,7 @@ describe('Jupyter Session', function (): void { kernel = session.kernel; // Then I expect it to have the ID, and only be called once should(kernel.id).equal('id'); - mockJupyterSession.verify(s => s.kernel, TypeMoq.Times.once()); + mockJupyterSession.verify(s => s.kernel, TypeMoq.Times.exactly(2)); }); it('should send name in changeKernel request', async function (): Promise { diff --git a/src/sql/workbench/api/common/extHostNotebook.ts b/src/sql/workbench/api/common/extHostNotebook.ts index c47b6632ca..c6263f526a 100644 --- a/src/sql/workbench/api/common/extHostNotebook.ts +++ b/src/sql/workbench/api/common/extHostNotebook.ts @@ -34,7 +34,7 @@ export class ExtHostNotebook implements ExtHostNotebookShape { let adapter = this.findManagerForUri(uriString); if (!adapter) { adapter = await this._withProvider(providerHandle, (provider) => { - return this.createManager(provider, uri); + return this.getOrCreateManager(provider, uri); }); } @@ -242,7 +242,7 @@ export class ExtHostNotebook implements ExtHostNotebookShape { return undefined; } - private async createManager(provider: azdata.nb.NotebookProvider, notebookUri: URI): Promise { + private async getOrCreateManager(provider: azdata.nb.NotebookProvider, notebookUri: URI): Promise { let manager = await provider.getNotebookManager(notebookUri); let uriString = notebookUri.toString(); let adapter = new NotebookManagerAdapter(provider, manager, uriString);