Disconnect Object Explorer node when session is disconnected from SQL Tools Service (#3249)

* WIP

* WIP

* Send disconnect event to OE

* Bump dataprotocol to 0.2.9

* Cleanupps

* Address a couple feedback

extensions/mssql/package.json

@@ -18,7 +18,7 @@
     "update-grammar": "node ../../build/npm/update-grammar.js Microsoft/vscode-mssql syntaxes/SQL.plist ./syntaxes/sql.tmLanguage.json"
   },
   "dependencies": {
-    "dataprotocol-client": "github:Microsoft/sqlops-dataprotocolclient#0.2.8",
+    "dataprotocol-client": "github:Microsoft/sqlops-dataprotocolclient#0.2.9",
     "opener": "^1.4.3",
     "service-downloader": "github:anthonydresser/service-downloader#0.1.5",
     "vscode-extension-telemetry": "^0.0.15"

extensions/mssql/src/config.json

@@ -1,6 +1,6 @@
 {
 	"downloadUrl": "https://github.com/Microsoft/sqltoolsservice/releases/download/v{#version#}/microsoft.sqltools.servicelayer-{#fileName#}",
-	"version": "1.5.0-alpha.53",
+	"version": "1.5.0-alpha.55",
 	"downloadFileNames": {
 		"Windows_86": "win-x86-netcoreapp2.2.zip",
 		"Windows_64": "win-x64-netcoreapp2.2.zip",

src/sql/parts/objectExplorer/common/objectExplorerService.ts

@@ -20,6 +20,7 @@ import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
 import { warn, error } from 'sql/base/common/log';
 import { ServerTreeView } from 'sql/parts/objectExplorer/viewlet/serverTreeView';
 import { ICapabilitiesService } from 'sql/services/capabilities/capabilitiesService';
+import * as Utils from 'sql/parts/connection/common/utils';
 
 export const SERVICE_ID = 'ObjectExplorerService';
 
@@ -42,6 +43,8 @@ export interface IObjectExplorerService {
 
 	onSessionCreated(handle: number, sessionResponse: sqlops.ObjectExplorerSession);
 
+	onSessionDisconnected(handle: number, sessionResponse: sqlops.ObjectExplorerSession);
+
 	onNodeExpanded(handle: number, sessionResponse: sqlops.ObjectExplorerExpandInfo);
 
 	/**
@@ -205,6 +208,29 @@ export class ObjectExplorerService implements IObjectExplorerService {
 		this.sendUpdateNodeEvent(connection, errorMessage);
 	}
 
+	/**
+	 * Gets called when session is disconnected
+	 */
+	public onSessionDisconnected(handle: number, session: sqlops.ObjectExplorerSession) {
+		if (this._sessions[session.sessionId]) {
+			let connection: ConnectionProfile = this._sessions[session.sessionId].connection;
+			if (connection && this._connectionManagementService.isProfileConnected(connection)) {
+				let uri: string = Utils.generateUri(connection);
+				if (this._serverTreeView.isObjectExplorerConnectionUri(uri)) {
+					this._serverTreeView.deleteObjectExplorerNodeAndRefreshTree(connection).then(() => {
+						this.sendUpdateNodeEvent(connection, session.errorMessage);
+						connection.isDisconnecting = true;
+						this._connectionManagementService.disconnect(connection).then((value) => {
+							connection.isDisconnecting = false;
+						});
+					});
+				}
+			}
+		} else {
+			warn(`Cannot find session ${session.sessionId}`);
+		}
+	}
+
 	private sendUpdateNodeEvent(connection: ConnectionProfile, errorMessage: string = undefined) {
 		let eventArgs: ObjectExplorerNodeEventArgs = {
 			connection: <IConnectionProfile>connection,

src/sql/parts/objectExplorer/viewlet/serverTreeView.ts

@@ -150,7 +150,7 @@ export class ServerTreeView {
 		});
 	}
 
-	private isObjectExplorerConnectionUri(uri: string): boolean {
+	public isObjectExplorerConnectionUri(uri: string): boolean {
 		let isBackupRestoreUri: boolean = uri.indexOf(ConnectionUtils.ConnectionUriBackupIdAttributeName) >= 0 ||
 			uri.indexOf(ConnectionUtils.ConnectionUriRestoreIdAttributeName) >= 0;
 		return uri && uri.startsWith(ConnectionUtils.uriPrefixes.default) && !isBackupRestoreUri;
@@ -227,16 +227,17 @@ export class ServerTreeView {
 		}
 	}
 
-	public deleteObjectExplorerNodeAndRefreshTree(connection: IConnectionProfile): void {
+	public deleteObjectExplorerNodeAndRefreshTree(connection: IConnectionProfile): Thenable<void> {
 		if (connection) {
 			var conn = this.getConnectionInTreeInput(connection.id);
 			if (conn) {
-				this._objectExplorerService.deleteObjectExplorerNode(conn).then(() => {
+				return this._objectExplorerService.deleteObjectExplorerNode(conn).then(() => {
 					this._tree.collapse(conn);
 					this._tree.refresh(conn);
 				});
 			}
 		}
+		return Promise.resolve();
 	}
 
 	public refreshTree(): void {

src/sql/sqlops.d.ts

@@ -1170,8 +1170,9 @@ declare module 'sqlops' {
 
 		registerOnSessionCreated(handler: (response: ObjectExplorerSession) => any): void;
 
-		registerOnExpandCompleted(handler: (response: ObjectExplorerExpandInfo) => any): void;
+		registerOnSessionDisconnected?(handler: (response: ObjectExplorerSession) => any): void;
 
+		registerOnExpandCompleted(handler: (response: ObjectExplorerExpandInfo) => any): void;
 	}
 
 	// Admin Services interfaces  -----------------------------------------------------------------------

src/sql/workbench/api/node/extHostDataProtocol.ts

@@ -353,6 +353,10 @@ export class ExtHostDataProtocol extends ExtHostDataProtocolShape {
 		this._proxy.$onObjectExplorerSessionCreated(handle, response);
 	}
 
+	public $onObjectExplorerSessionDisconnected(handle: number, response: sqlops.ObjectExplorerSession): void {
+		this._proxy.$onObjectExplorerSessionDisconnected(handle, response);
+	}
+
 	public $onObjectExplorerNodeExpanded(handle: number, response: sqlops.ObjectExplorerExpandInfo): void {
 		this._proxy.$onObjectExplorerNodeExpanded(handle, response);
 	}

src/sql/workbench/api/node/mainThreadDataProtocol.ts

@@ -442,6 +442,10 @@ export class MainThreadDataProtocol implements MainThreadDataProtocolShape {
 		this._objectExplorerService.onSessionCreated(handle, sessionResponse);
 	}
 
+	public $onObjectExplorerSessionDisconnected(handle: number, sessionResponse: sqlops.ObjectExplorerSession): void {
+		this._objectExplorerService.onSessionDisconnected(handle, sessionResponse);
+	}
+
 	public $onObjectExplorerNodeExpanded(handle: number, expandResponse: sqlops.ObjectExplorerExpandInfo): void {
 		this._objectExplorerService.onNodeExpanded(handle, expandResponse);
 	}

src/sql/workbench/api/node/sqlExtHost.api.impl.ts

@@ -217,6 +217,12 @@ export function createApiFactory(
 					extHostDataProvider.$onObjectExplorerSessionCreated(provider.handle, response);
 				});
 
+				if (provider.registerOnSessionDisconnected) {
+					provider.registerOnSessionDisconnected((response: sqlops.ObjectExplorerSession) => {
+						extHostDataProvider.$onObjectExplorerSessionDisconnected(provider.handle, response);
+					});
+				}
+
 				provider.registerOnExpandCompleted((response: sqlops.ObjectExplorerExpandInfo) => {
 					extHostDataProvider.$onObjectExplorerNodeExpanded(provider.handle, response);
 				});

src/sql/workbench/api/node/sqlExtHost.protocol.ts

@@ -486,6 +486,7 @@ export interface MainThreadDataProtocolShape extends IDisposable {
 	$onResultSetComplete(handle: number, resultSetInfo: sqlops.QueryExecuteResultSetCompleteNotificationParams): void;
 	$onQueryMessage(handle: number, message: sqlops.QueryExecuteMessageParams): void;
 	$onObjectExplorerSessionCreated(handle: number, message: sqlops.ObjectExplorerSession): void;
+	$onObjectExplorerSessionDisconnected(handle: number, message: sqlops.ObjectExplorerSession): void;
 	$onObjectExplorerNodeExpanded(handle: number, message: sqlops.ObjectExplorerExpandInfo): void;
 	$onTaskCreated(handle: number, sessionResponse: sqlops.TaskInfo): void;
 	$onTaskStatusChanged(handle: number, sessionResponse: sqlops.TaskProgressInfo): void;

src/sqltest/stubs/objectExplorerProviderTestService.ts

@@ -30,6 +30,10 @@ export class ObjectExplorerProviderTestService implements sqlops.ObjectExplorerP
 
 	}
 
+	public registerOnSessionDisconnected(handler: (response: sqlops.ObjectExplorerSession) => any): void {
+
+	}
+
 	public registerOnExpandCompleted(handler: (response: sqlops.ObjectExplorerExpandInfo) => any): void {
 
 	}