From 504d5c91bc9d1781ea9f9eb6c245d43475941fef Mon Sep 17 00:00:00 2001
From: Yurong He <43652751+YurongHe@users.noreply.github.com>
Date: Wed, 3 Apr 2019 19:00:56 -0700
Subject: [PATCH] =?UTF-8?q?Fixed=20#4800=20need=20to=20use=20ConnectionPro?=
 =?UTF-8?q?file=20in=20order=20to=20get=20the=20correct=E2=80=A6=20(#4812)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Fixed #4800 need to use ConnectionProfile in order to get the correct connection

* Go back to create connect in run cell to avoid to fail to run cell or close the connection used by other.
---
 src/sql/parts/notebook/models/notebookModel.ts       | 12 ++++++------
 src/sql/parts/notebook/notebookActions.ts            |  2 +-
 .../services/notebook/sql/sqlSessionManager.ts       | 10 +++++++---
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/sql/parts/notebook/models/notebookModel.ts b/src/sql/parts/notebook/models/notebookModel.ts
index f1c59e3328..30bf4946f2 100644
--- a/src/sql/parts/notebook/models/notebookModel.ts
+++ b/src/sql/parts/notebook/models/notebookModel.ts
@@ -637,19 +637,19 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return spec;
 	}
 
-	public async changeContext(server: string, newConnection?: ConnectionProfile, hideErrorMessage?: boolean): Promise<void> {
+	public async changeContext(title: string, newConnection?: ConnectionProfile, hideErrorMessage?: boolean): Promise<void> {
 		try {
 			if (!newConnection) {
-				newConnection = this._activeContexts.otherConnections.find((connection) => connection.serverName === server);
+				newConnection = this._activeContexts.otherConnections.find((connection) => connection.title === title);
 			}
-			if ((!newConnection) && (this._activeContexts.defaultConnection.serverName === server)) {
+			if ((!newConnection) && (this._activeContexts.defaultConnection.title === title)) {
 				newConnection = this._activeContexts.defaultConnection;
 			}
 
-			if (this._activeConnection) {
-				this._otherConnections.push(this._activeConnection);
-			}
 			if (newConnection) {
+				if (this._activeConnection && this._activeConnection.id !== newConnection.id) {
+					this._otherConnections.push(this._activeConnection);
+				}
 				this._activeConnection = newConnection;
 				this.refreshConnections(newConnection);
 				this._activeClientSession.updateConnection(newConnection.toIConnectionProfile()).then(
diff --git a/src/sql/parts/notebook/notebookActions.ts b/src/sql/parts/notebook/notebookActions.ts
index 4829369837..3934963258 100644
--- a/src/sql/parts/notebook/notebookActions.ts
+++ b/src/sql/parts/notebook/notebookActions.ts
@@ -389,8 +389,8 @@ export class AttachToDropdown extends SelectBox {
 				else {
 					connections.push(msgAddNewConnection);
 				}
+				this.setOptions(connections);
 			}
-			this.setOptions(connections);
 		}
 	}
 
diff --git a/src/sql/workbench/services/notebook/sql/sqlSessionManager.ts b/src/sql/workbench/services/notebook/sql/sqlSessionManager.ts
index 53e5692f45..45e9836187 100644
--- a/src/sql/workbench/services/notebook/sql/sqlSessionManager.ts
+++ b/src/sql/workbench/services/notebook/sql/sqlSessionManager.ts
@@ -23,6 +23,7 @@ import { escape } from 'sql/base/common/strings';
 import { elapsedTimeLabel } from 'sql/parts/query/common/localizedConstants';
 import * as notebookUtils from 'sql/parts/notebook/notebookUtils';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 
 export const sqlKernelError: string = localize("sqlKernelError", "SQL kernel error");
 export const MAX_ROWS = 5000;
@@ -138,6 +139,7 @@ export class SqlSession implements nb.ISession {
 class SqlKernel extends Disposable implements nb.IKernel {
 	private _queryRunner: QueryRunner;
 	private _currentConnection: IConnectionProfile;
+	private _currentConnectionProfile: ConnectionProfile;
 	static kernelId: number = 0;
 
 	private _id: string;
@@ -146,6 +148,7 @@ class SqlKernel extends Disposable implements nb.IKernel {
 	private _magicToExecutorMap = new Map<string, ExternalScriptMagic>();
 
 	constructor(@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
+		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService,
 		@IInstantiationService private _instantiationService: IInstantiationService,
 		@IErrorMessageService private _errorMessageService: IErrorMessageService,
 		@IConfigurationService private _configurationService: IConfigurationService
@@ -206,6 +209,7 @@ class SqlKernel extends Disposable implements nb.IKernel {
 
 	public set connection(conn: IConnectionProfile) {
 		this._currentConnection = conn;
+		this._currentConnectionProfile = new ConnectionProfile(this._capabilitiesService, this._currentConnection);
 		this._queryRunner = undefined;
 	}
 
@@ -224,10 +228,10 @@ class SqlKernel extends Disposable implements nb.IKernel {
 				this._future.handleDone();
 			}
 			this._queryRunner.runQuery(code);
-		} else if (this._currentConnection) {
-			let connectionUri = Utils.generateUri(this._currentConnection, 'notebook');
+		} else if (this._currentConnection && this._currentConnectionProfile) {
+			let connectionUri = Utils.generateUri(this._currentConnectionProfile, 'notebook');
 			this._queryRunner = this._instantiationService.createInstance(QueryRunner, connectionUri);
-			this._connectionManagementService.connect(this._currentConnection, connectionUri).then((result) => {
+			this._connectionManagementService.connect(this._currentConnectionProfile, connectionUri).then((result) => {
 				this.addQueryEventListeners(this._queryRunner);
 				this._queryRunner.runQuery(code);
 			});