From ebc208cacdcdda792ac9594a537afafba47cda71 Mon Sep 17 00:00:00 2001 From: Yurong He <43652751+YurongHe@users.noreply.github.com> Date: Fri, 1 Mar 2019 22:11:45 -0800 Subject: [PATCH] Fixed #4181 and #4167 change kernel issue between SQL and Sparks #4238 (#4256) * Fixed #4181 and #4167 The problems are: - When change Kernel, setProviderIdForKernel switches Kernel first. So when we switch kernel across provider for exmple from PySpark3 to SQL. The Kernel is set SQL already in ClientSession.changeKernel. Then we lost oldKernel info. The fix is cache the old session in model before switch and use it to get the correct context - SQL Kenerl could make mulitple connects from "Add new connection". While we didn't track them, those connections didn't close properly when close the notebook The fix is saving the connections made from "Add new connection" in model. and close them and activeConnection when close notebook Problem is not solved yet in this PR: - Didn't shutdown Jupytper when swich kernel from spark kernel to SQL. --- .../notebook/cellViews/code.component.ts | 6 ++--- .../parts/notebook/models/clientSession.ts | 4 +-- .../parts/notebook/models/modelInterfaces.ts | 3 ++- .../parts/notebook/models/notebookModel.ts | 25 +++++++++++++++++-- 4 files changed, 30 insertions(+), 8 deletions(-) diff --git a/src/sql/parts/notebook/cellViews/code.component.ts b/src/sql/parts/notebook/cellViews/code.component.ts index 4eccfdff87..c55c94735d 100644 --- a/src/sql/parts/notebook/cellViews/code.component.ts +++ b/src/sql/parts/notebook/cellViews/code.component.ts @@ -137,13 +137,13 @@ export class CodeComponent extends AngularDisposable implements OnInit, OnChange } } - private updateConnectionState(isConnected: boolean) { + private updateConnectionState(shouldConnect: boolean) { if (this.isSqlCodeCell()) { let cellUri = this.cellModel.cellUri.toString(); let connectionService = this.connectionService; - if (!isConnected && connectionService && connectionService.isConnected(cellUri)) { + if (!shouldConnect && connectionService && connectionService.isConnected(cellUri)) { connectionService.disconnect(cellUri).catch(e => console.log(e)); - } else if (this._model.activeConnection && this._model.activeConnection.id !== '-1') { + } else if (shouldConnect && this._model.activeConnection && this._model.activeConnection.id !== '-1') { connectionService.connect(this._model.activeConnection, cellUri).catch(e => console.log(e)); } } diff --git a/src/sql/parts/notebook/models/clientSession.ts b/src/sql/parts/notebook/models/clientSession.ts index ec9ff8c2b1..d572d287a7 100644 --- a/src/sql/parts/notebook/models/clientSession.ts +++ b/src/sql/parts/notebook/models/clientSession.ts @@ -225,10 +225,10 @@ export class ClientSession implements IClientSession { /** * Change the current kernel associated with the document. */ - async changeKernel(options: nb.IKernelSpec): Promise { + async changeKernel(options: nb.IKernelSpec, oldValue?: nb.IKernel): Promise { this._kernelChangeCompleted = new Deferred(); this._isReady = false; - let oldKernel = this.kernel; + let oldKernel = oldValue? oldValue: this.kernel; let newKernel = this.kernel; let kernel = await this.doChangeKernel(options); diff --git a/src/sql/parts/notebook/models/modelInterfaces.ts b/src/sql/parts/notebook/models/modelInterfaces.ts index bec0f3d9b9..728ef89e8d 100644 --- a/src/sql/parts/notebook/models/modelInterfaces.ts +++ b/src/sql/parts/notebook/models/modelInterfaces.ts @@ -141,7 +141,8 @@ export interface IClientSession extends IDisposable { * Change the current kernel associated with the document. */ changeKernel( - options: nb.IKernelSpec + options: nb.IKernelSpec, + oldKernel?: nb.IKernel ): Promise; /** diff --git a/src/sql/parts/notebook/models/notebookModel.ts b/src/sql/parts/notebook/models/notebookModel.ts index 44b50b8d3d..e858194a8e 100644 --- a/src/sql/parts/notebook/models/notebookModel.ts +++ b/src/sql/parts/notebook/models/notebookModel.ts @@ -8,7 +8,7 @@ import { nb, connection } from 'azdata'; import { localize } from 'vs/nls'; -import { Event, Emitter } from 'vs/base/common/event'; +import { Event, Emitter, forEach } from 'vs/base/common/event'; import { Disposable } from 'vs/base/common/lifecycle'; import { CellModel } from './cell'; @@ -48,6 +48,7 @@ export class NotebookModel extends Disposable implements INotebookModel { private _inErrorState: boolean = false; private _clientSessions: IClientSession[] = []; private _activeClientSession: IClientSession; + private _oldClientSession: IClientSession; private _sessionLoadFinished: Promise; private _onClientSessionReady = new Emitter(); private _onProviderIdChanged = new Emitter(); @@ -62,6 +63,7 @@ export class NotebookModel extends Disposable implements INotebookModel { private readonly _nbformat: number = nbversion.MAJOR_VERSION; private readonly _nbformatMinor: number = nbversion.MINOR_VERSION; private _activeConnection: ConnectionProfile; + private _otherConnections: ConnectionProfile[] = []; private _activeCell: ICellModel; private _providerId: string; private _defaultKernel: nb.IKernelSpec; @@ -478,7 +480,8 @@ export class NotebookModel extends Disposable implements INotebookModel { kernelSpec = kernelSpecs.find(spec => spec.name === this.notebookManager.sessionManager.specs.defaultKernel); } if (this._activeClientSession && this._activeClientSession.isReady) { - return this._activeClientSession.changeKernel(kernelSpec) + let oldKernel = this._oldClientSession ? this._oldClientSession.kernel : null; + return this._activeClientSession.changeKernel(kernelSpec, oldKernel) .then((kernel) => { this.updateKernelInfo(kernel); kernel.ready.then(() => { @@ -503,6 +506,9 @@ export class NotebookModel extends Disposable implements INotebookModel { newConnection = this._activeContexts.defaultConnection; } let newConnectionProfile = new ConnectionProfile(this._notebookOptions.capabilitiesService, newConnection); + if (this._activeConnection) { + this._otherConnections.push(this._activeConnection); + } this._activeConnection = newConnectionProfile; this.refreshConnections(newConnectionProfile); this._activeClientSession.updateConnection(this._activeConnection.toIConnectionProfile()).then( @@ -620,6 +626,17 @@ export class NotebookModel extends Disposable implements INotebookModel { public async handleClosed(): Promise { try { + if (this.notebookOptions && this.notebookOptions.connectionService) { + let connectionService = this.notebookOptions.connectionService; + if (this._otherConnections) { + this._otherConnections.forEach(conn => connectionService.disconnect(conn).catch(e => console.log(e))); + this._otherConnections = []; + } + if (this._activeConnection) { + this.notebookOptions.connectionService.disconnect(this._activeConnection).catch(e => console.log(e)); + this._activeConnection = undefined; + } + } if (this._activeClientSession) { try { await this._activeClientSession.ready; @@ -629,6 +646,7 @@ export class NotebookModel extends Disposable implements INotebookModel { await this._activeClientSession.shutdown(); this._clientSessions = undefined; this._activeClientSession = undefined; + } } catch (err) { this.notifyError(localize('shutdownError', 'An error occurred when closing the notebook: {0}', err)); @@ -711,6 +729,9 @@ export class NotebookModel extends Disposable implements INotebookModel { if (this.notebookManagers[i].sessionManager && this.notebookManagers[i].sessionManager.specs && this.notebookManagers[i].sessionManager.specs.kernels) { let index = this.notebookManagers[i].sessionManager.specs.kernels.findIndex(kernel => kernel.name === kernelSpec.name); if (index >= 0) { + if (this._activeClientSession) { + this._oldClientSession = this._activeClientSession; + } this._activeClientSession = this._clientSessions[i]; if (this.notebookManagers[i].providerId !== this._providerId) { this._providerId = this.notebookManagers[i].providerId;