From e8b4c0377014a01f94e5ddc2cb4f38cb85e29655 Mon Sep 17 00:00:00 2001
From: Yurong He <43652751+YurongHe@users.noreply.github.com>
Date: Tue, 23 Apr 2019 16:42:28 -0700
Subject: [PATCH] =?UTF-8?q?Fixed=20#4772=20create=20unique=20connection=20?=
 =?UTF-8?q?for=20each=20notebook=20by=20using=20noteb=E2=80=A6=20(#5163)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Fixed #4772 create unique connection for each notebook by using notebook path as uri
Disconnect sqlConnection from SqlKernel to ensure no connection left after the notebook is closed.

* SqlSessionManager is ADS level manager

* Moved path to SqlKernel constructor
---
 .../notebook/sql/sqlSessionManager.ts         | 40 ++++++++++++++-----
 1 file changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/sql/workbench/services/notebook/sql/sqlSessionManager.ts b/src/sql/workbench/services/notebook/sql/sqlSessionManager.ts
index c49e0fdaf6..c6fa3fa442 100644
--- a/src/sql/workbench/services/notebook/sql/sqlSessionManager.ts
+++ b/src/sql/workbench/services/notebook/sql/sqlSessionManager.ts
@@ -12,7 +12,6 @@ import QueryRunner, { EventType } from 'sql/platform/query/common/queryRunner';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import Severity from 'vs/base/common/severity';
-import * as Utils from 'sql/platform/connection/common/utils';
 import { Deferred } from 'sql/base/common/promise';
 import { Disposable } from 'vs/base/common/lifecycle';
 import { IErrorMessageService } from 'sql/platform/errorMessage/common/errorMessageService';
@@ -46,6 +45,8 @@ export interface SQLData {
 }
 
 export class SqlSessionManager implements nb.SessionManager {
+	private static _sessions: nb.ISession[] = [];
+
 	constructor(private _instantiationService: IInstantiationService) { }
 
 	public get isReady(): boolean {
@@ -65,11 +66,25 @@ export class SqlSessionManager implements nb.SessionManager {
 	}
 
 	startNew(options: nb.ISessionOptions): Thenable<nb.ISession> {
-		let session = new SqlSession(options, this._instantiationService);
-		return Promise.resolve(session);
+		let sqlSession = new SqlSession(options, this._instantiationService);
+		let index = SqlSessionManager._sessions.findIndex(session => session.path === options.path);
+		if (index > -1) {
+			SqlSessionManager._sessions.splice(index);
+		}
+		SqlSessionManager._sessions.push(sqlSession);
+		return Promise.resolve(sqlSession);
 	}
 
 	shutdown(id: string): Thenable<void> {
+		let index = SqlSessionManager._sessions.findIndex(session => session.id === id);
+		if (index > -1) {
+			let sessionManager = SqlSessionManager._sessions[index];
+			SqlSessionManager._sessions.splice(index);
+			if (sessionManager && sessionManager.kernel) {
+				let sqlKernel = sessionManager.kernel as SqlKernel;
+				return sqlKernel.disconnect();
+			}
+		}
 		return Promise.resolve();
 	}
 }
@@ -88,7 +103,7 @@ export class SqlSession implements nb.ISession {
 	}
 
 	constructor(private options: nb.ISessionOptions, private _instantiationService: IInstantiationService) {
-		this._kernel = this._instantiationService.createInstance(SqlKernel);
+		this._kernel = this._instantiationService.createInstance(SqlKernel, options.path);
 	}
 
 	public get canChangeKernels(): boolean {
@@ -96,7 +111,7 @@ export class SqlSession implements nb.ISession {
 	}
 
 	public get id(): string {
-		return this.options.kernelId || '';
+		return this.options.kernelId || this.kernel ? this._kernel.id : '';
 	}
 
 	public get path(): string {
@@ -146,7 +161,8 @@ class SqlKernel extends Disposable implements nb.IKernel {
 	private _executionCount: number = 0;
 	private _magicToExecutorMap = new Map<string, ExternalScriptMagic>();
 
-	constructor(@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
+	constructor(private _path: string,
+		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
 		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService,
 		@IInstantiationService private _instantiationService: IInstantiationService,
 		@IErrorMessageService private _errorMessageService: IErrorMessageService,
@@ -232,9 +248,8 @@ class SqlKernel extends Disposable implements nb.IKernel {
 			}
 			this._queryRunner.runQuery(code);
 		} else if (this._currentConnection && this._currentConnectionProfile) {
-			let connectionUri = Utils.generateUri(this._currentConnectionProfile, 'notebook');
-			this._queryRunner = this._instantiationService.createInstance(QueryRunner, connectionUri);
-			this._connectionManagementService.connect(this._currentConnectionProfile, connectionUri).then((result) => {
+			this._queryRunner = this._instantiationService.createInstance(QueryRunner, this._path);
+			this._connectionManagementService.connect(this._currentConnectionProfile, this._path).then((result) => {
 				this.addQueryEventListeners(this._queryRunner);
 				this._queryRunner.runQuery(code);
 			});
@@ -311,6 +326,13 @@ class SqlKernel extends Disposable implements nb.IKernel {
 		}
 		// TODO issue #2746 should ideally show a warning inside the dialog if have no data
 	}
+
+	public async disconnect(): Promise<void> {
+		if (this._path) {
+			await this._connectionManagementService.disconnect(this._path);
+		}
+		return;
+	}
 }
 
 export class SQLFuture extends Disposable implements FutureInternal {