Fix Issue when OE connection Disconnects on Notebook Close (#4425)

* Fix OE connection closing when notebook closes

* handle connections created through Add new connection

src/sql/parts/notebook/models/notebookModel.ts

@@ -23,6 +23,7 @@ import { INotification, Severity } from 'vs/platform/notification/common/notific
 import { URI } from 'vs/base/common/uri';
 import { ISingleNotebookEditOperation } from 'sql/workbench/api/common/sqlExtHostTypes';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
+import { uriPrefixes } from 'sql/platform/connection/common/utils';
 
 /*
 * Used to control whether a message in a dialog/wizard is displayed as an error,
@@ -71,6 +72,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 	private _onValidConnectionSelected = new Emitter<boolean>();
 	private _oldKernel: nb.IKernel;
 	private _clientSessionListeners: IDisposable[] = [];
+	private _connectionsToDispose: ConnectionProfile[] = [];
 
 	constructor(private _notebookOptions: INotebookModelOptions, startSessionImmediately?: boolean, private connectionProfile?: IConnectionProfile) {
 		super();
@@ -695,6 +697,10 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return newKernelDisplayName;
 	}
 
+	public addAttachToConnectionsToBeDisposed(conn: ConnectionProfile) {
+		this._connectionsToDispose.push(conn);
+	}
+
 	private setErrorState(errMsg: string): void {
 		this._inErrorState = true;
 		let msg = localize('startSessionFailed', 'Could not start session: {0}', errMsg);
@@ -704,19 +710,21 @@ export class NotebookModel extends Disposable implements INotebookModel {
 
 	public dispose(): void {
 		super.dispose();
+		this.disconnectAttachToConnections();
 		this.handleClosed();
 	}
 
 	public async handleClosed(): Promise<void> {
 		try {
 			if (this.notebookOptions && this.notebookOptions.connectionService) {
-				let connectionService = this.notebookOptions.connectionService;
 				if (this._otherConnections) {
-					notebookUtils.asyncForEach(this._otherConnections, async (conn) => await connectionService.disconnect(conn).catch(e => console.log(e)));
+					notebookUtils.asyncForEach(this._otherConnections, async (conn) => {
+						await this.disconnectNotebookConnection(conn);
+					});
 					this._otherConnections = [];
 				}
 				if (this._activeConnection) {
-					await this.notebookOptions.connectionService.disconnect(this._activeConnection).catch(e => console.log(e));
+					await this.disconnectNotebookConnection(this._activeConnection);
 					this._activeConnection = undefined;
 				}
 			}
@@ -840,6 +848,24 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return [];
 	}
 
+	// Check for and disconnect from any new connections opened while in the notebook
+	// Note: notebooks should always connect with the connection URI in the following format,
+	// so that connections can be tracked accordingly throughout ADS:
+	// let connectionUri = Utils.generateUri(connection, 'notebook');
+	private async disconnectNotebookConnection(conn: ConnectionProfile): Promise<void> {
+		if (this.notebookOptions.connectionService.getConnectionUri(conn).includes(uriPrefixes.notebook)) {
+			await this.notebookOptions.connectionService.disconnect(conn).catch(e => console.log(e));
+		}
+	}
+
+	// Disconnect any connections that were added through the "Add new connection" functionality in the Attach To dropdown
+	private async disconnectAttachToConnections(): Promise<void> {
+		this._connectionsToDispose.forEach(async conn => {
+			await this.notebookOptions.connectionService.disconnect(conn).catch(e => console.log(e));
+		});
+		this._connectionsToDispose = [];
+	}
+
 	/**
 	 * Serialize the model to JSON.
 	 */

src/sql/parts/notebook/notebookActions.ts

@@ -438,6 +438,7 @@ export class AttachToDropdown extends SelectBox {
 				}
 				this.select(index);
 
+				this.model.addAttachToConnectionsToBeDisposed(connectionProfile);
 				// Call doChangeContext to set the newly chosen connection in the model
 				this.doChangeContext(connectionProfile);
 			});

src/sql/platform/connection/common/connectionManagementService.ts

@@ -379,7 +379,7 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 	 * otherwise tries to make a connection and returns the owner uri when connection is complete
 	 * The purpose is connection by default
 	 */
-	public connectIfNotConnected(connection: IConnectionProfile, purpose?: 'dashboard' | 'insights' | 'connection', saveConnection: boolean = false): Promise<string> {
+	public connectIfNotConnected(connection: IConnectionProfile, purpose?: 'dashboard' | 'insights' | 'connection' | 'notebook', saveConnection: boolean = false): Promise<string> {
 		return new Promise<string>((resolve, reject) => {
 			let ownerUri: string = Utils.generateUri(connection, purpose);
 			if (this._connectionStatusManager.isConnected(ownerUri)) {
@@ -1156,7 +1156,7 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 	 * Finds existing connection for given profile and purpose is any exists.
 	 * The purpose is connection by default
 	 */
-	public findExistingConnection(connection: IConnectionProfile, purpose?: 'dashboard' | 'insights' | 'connection'): ConnectionProfile {
+	public findExistingConnection(connection: IConnectionProfile, purpose?: 'dashboard' | 'insights' | 'connection' | 'notebook'): ConnectionProfile {
 		let connectionUri = Utils.generateUri(connection, purpose);
 		let existingConnection = this._connectionStatusManager.findConnection(connectionUri);
 		if (existingConnection && this._connectionStatusManager.isConnected(connectionUri)) {

src/sql/platform/connection/common/utils.ts

@@ -17,7 +17,8 @@ export const uriPrefixes = {
 	default: 'connection://',
 	connection: 'connection://',
 	dashboard: 'dashboard://',
-	insights: 'insights://'
+	insights: 'insights://',
+	notebook: 'notebook://'
 };