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.

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));
 			}
 		}

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<nb.IKernel> {
+	async changeKernel(options: nb.IKernelSpec, oldValue?: nb.IKernel): Promise<nb.IKernel> {
 		this._kernelChangeCompleted = new Deferred<void>();
 		this._isReady = false;
-		let oldKernel = this.kernel;
+		let oldKernel = oldValue? oldValue: this.kernel;
 		let newKernel = this.kernel;
 
 		let kernel = await this.doChangeKernel(options);

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<nb.IKernel>;
 
 	/**

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<void>;
 	private _onClientSessionReady = new Emitter<IClientSession>();
 	private _onProviderIdChanged = new Emitter<string>();
@@ -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<void> {
 		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;