From 17901fbf3d33e175bd04dd44be0ed512fd4c9e4f Mon Sep 17 00:00:00 2001 From: Chris LaFreniere <40371649+chlafreniere@users.noreply.github.com> Date: Wed, 13 Mar 2019 11:32:12 -0700 Subject: [PATCH] Fix Issue when OE connection Disconnects on Notebook Close (#4425) * Fix OE connection closing when notebook closes * handle connections created through Add new connection --- .../parts/notebook/models/notebookModel.ts | 32 +++++++++++++++++-- src/sql/parts/notebook/notebookActions.ts | 1 + .../common/connectionManagementService.ts | 4 +-- src/sql/platform/connection/common/utils.ts | 3 +- 4 files changed, 34 insertions(+), 6 deletions(-) diff --git a/src/sql/parts/notebook/models/notebookModel.ts b/src/sql/parts/notebook/models/notebookModel.ts index 3c3988a692..adfeb5f902 100644 --- a/src/sql/parts/notebook/models/notebookModel.ts +++ b/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(); 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 { 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 { + 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 { + this._connectionsToDispose.forEach(async conn => { + await this.notebookOptions.connectionService.disconnect(conn).catch(e => console.log(e)); + }); + this._connectionsToDispose = []; + } + /** * Serialize the model to JSON. */ diff --git a/src/sql/parts/notebook/notebookActions.ts b/src/sql/parts/notebook/notebookActions.ts index 7bb811ee84..117d369f99 100644 --- a/src/sql/parts/notebook/notebookActions.ts +++ b/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); }); diff --git a/src/sql/platform/connection/common/connectionManagementService.ts b/src/sql/platform/connection/common/connectionManagementService.ts index 22ab12db69..18376b86da 100644 --- a/src/sql/platform/connection/common/connectionManagementService.ts +++ b/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 { + public connectIfNotConnected(connection: IConnectionProfile, purpose?: 'dashboard' | 'insights' | 'connection' | 'notebook', saveConnection: boolean = false): Promise { return new Promise((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)) { diff --git a/src/sql/platform/connection/common/utils.ts b/src/sql/platform/connection/common/utils.ts index cb05a5f5f4..8de53e3fba 100644 --- a/src/sql/platform/connection/common/utils.ts +++ b/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://' };