Data Explorer Disconnect and Error Handling (#4243)

* adding context

* apply extension changes

* shimming disconnect

* add data explorer context menu and add disconnect to it

* clean up shim code; better handle errors

* remove tpromise

* simplify code

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

@@ -18,26 +18,24 @@ import { IConnectionProfile } from 'azdata';
 import { TreeItemCollapsibleState } from 'vs/workbench/common/views';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import { TPromise } from 'vs/base/common/winjs.base';
-import { equalsIgnoreCase } from 'vs/base/common/strings';
 import { hash } from 'vs/base/common/hash';
 import { generateUuid } from 'vs/base/common/uuid';
-import { URI } from 'vs/base/common/uri';
 
 export const SERVICE_ID = 'oeShimService';
 export const IOEShimService = createDecorator<IOEShimService>(SERVICE_ID);
 
 export interface IOEShimService {
 	_serviceBrand: any;
-	getChildren(node: ITreeItem, identifier: any): Promise<ITreeItem[]>;
+	getChildren(node: ITreeItem, viewId: string): Promise<ITreeItem[]>;
+	disconnectNode(viewId: string, node: ITreeItem): Promise<boolean>;
 	providerExists(providerId: string): boolean;
 }
 
 export class OEShimService implements IOEShimService {
 	_serviceBrand: any;
 
-	// maps a datasource to a provider handle to a session
-	private sessionMap = new Map<any, Map<number, string>>();
-	private nodeIdMap = new Map<string, string>();
+	private sessionMap = new Map<number, string>();
+	private nodeHandleMap = new Map<number, string>();
 
 	constructor(
 		@IObjectExplorerService private oe: IObjectExplorerService,
@@ -47,73 +45,82 @@ export class OEShimService implements IOEShimService {
 	) {
 	}
 
-	private async createSession(providerId: string, node: ITreeItem): TPromise<string> {
+	private async createSession(viewId: string, providerId: string, node: ITreeItem): Promise<string> {
 		let deferred = new Deferred<string>();
 		let connProfile = new ConnectionProfile(this.capabilities, node.payload);
 		connProfile.saveProfile = false;
 		if (this.cm.providerRegistered(providerId)) {
-			connProfile = new ConnectionProfile(this.capabilities, await this.cd.openDialogAndWait(this.cm, { connectionType: ConnectionType.default, showDashboard: false }, connProfile, undefined, false));
+			let userProfile = await this.cd.openDialogAndWait(this.cm, { connectionType: ConnectionType.default, showDashboard: false }, connProfile, undefined, false);
+			if (userProfile) {
+				connProfile = new ConnectionProfile(this.capabilities, userProfile);
+			} else {
+				return Promise.reject('User canceled');
+			}
 		}
 		let sessionResp = await this.oe.createNewSession(providerId, connProfile);
 		let disp = this.oe.onUpdateObjectExplorerNodes(e => {
 			if (e.connection.id === connProfile.id) {
+				if (e.errorMessage) {
+					deferred.reject();
+					return;
+				}
 				let rootNode = this.oe.getSession(sessionResp.sessionId).rootNode;
 				// this is how we know it was shimmed
 				if (rootNode.nodePath) {
-					node.handle = this.oe.getSession(sessionResp.sessionId).rootNode.nodePath;
+					this.nodeHandleMap.set(generateNodeMapKey(viewId, node), rootNode.nodePath);
 				}
 			}
 			disp.dispose();
 			deferred.resolve(sessionResp.sessionId);
 		});
-		return TPromise.wrap(deferred.promise);
+		return deferred.promise;
 	}
 
-	public async getChildren(node: ITreeItem, identifier: any): Promise<ITreeItem[]> {
-		try {
-			if (!this.sessionMap.has(identifier)) {
-				this.sessionMap.set(identifier, new Map<number, string>());
+	public async disconnectNode(viewId: string, node: ITreeItem): Promise<boolean> {
+		// we assume only nodes with payloads can be connected
+		// check to make sure we have an existing connection
+		let key = generateSessionMapKey(viewId, node);
+		let session = this.sessionMap.get(key);
+		if (session) {
+			let closed = (await this.oe.closeSession(node.childProvider, this.oe.getSession(session))).success;
+			if (closed) {
+				this.sessionMap.delete(key);
 			}
-			if (!this.sessionMap.get(identifier).has(hash(node.payload || node.childProvider))) {
-				this.sessionMap.get(identifier).set(hash(node.payload || node.childProvider), await this.createSession(node.childProvider, node));
-			}
-			if (this.nodeIdMap.has(node.handle)) {
-				node.handle = this.nodeIdMap.get(node.handle);
-			}
-			let sessionId = this.sessionMap.get(identifier).get(hash(node.payload || node.childProvider));
-			let treeNode = new TreeNode(undefined, undefined, undefined, node.handle, undefined, undefined, undefined, undefined, undefined, undefined);
-			let profile: IConnectionProfile = node.payload || {
-				providerName: node.childProvider,
-				authenticationType: undefined,
-				azureTenantId: undefined,
-				connectionName: undefined,
-				databaseName: undefined,
-				groupFullName: undefined,
-				groupId: undefined,
-				id: undefined,
-				options: undefined,
-				password: undefined,
-				savePassword: undefined,
-				saveProfile: undefined,
-				serverName: undefined,
-				userName: undefined,
-			};
-			treeNode.connection = new ConnectionProfile(this.capabilities, profile);
-			return TPromise.wrap(this.oe.resolveTreeNodeChildren({
+			return closed;
+		}
+		return Promise.resolve(false);
+	}
+
+	private async getOrCreateSession(viewId: string, node: ITreeItem): Promise<string> {
+		// verify the map is correct
+		let key = generateSessionMapKey(viewId, node);
+		if (!this.sessionMap.has(key)) {
+			this.sessionMap.set(key, await this.createSession(viewId, node.childProvider, node));
+		}
+		return this.sessionMap.get(key);
+	}
+
+	public async getChildren(node: ITreeItem, viewId: string): Promise<ITreeItem[]> {
+		if (node.payload) {
+			let sessionId = await this.getOrCreateSession(viewId, node);
+			let requestHandle = this.nodeHandleMap.get(generateNodeMapKey(viewId, node)) || node.handle;
+			let treeNode = new TreeNode(undefined, undefined, undefined, requestHandle, undefined, undefined, undefined, undefined, undefined, undefined);
+			treeNode.connection = new ConnectionProfile(this.capabilities, node.payload);
+			return this.oe.resolveTreeNodeChildren({
 				success: undefined,
 				sessionId,
 				rootNode: undefined,
 				errorMessage: undefined
-			}, treeNode).then(e => e.map(n => this.mapNodeToITreeItem(n, node))));
-		} catch (e) {
-			return TPromise.as([]);
+			}, treeNode).then(e => e.map(n => this.treeNodeToITreeItem(viewId, n, node)));
+		} else {
+			return Promise.resolve([]);
 		}
 	}
 
-	private mapNodeToITreeItem(node: TreeNode, parentNode: ITreeItem): ITreeItem {
+	private treeNodeToITreeItem(viewId: string, node: TreeNode, parentNode: ITreeItem): ITreeItem {
 		let handle = generateUuid();
-		this.nodeIdMap.set(handle, node.nodePath);
-		return {
+		let nodePath = node.nodePath;
+		let newTreeItem = {
 			parentHandle: node.parent.id,
 			handle,
 			collapsibleState: node.isAlwaysLeaf ? TreeItemCollapsibleState.None : TreeItemCollapsibleState.Collapsed,
@@ -125,9 +132,19 @@ export class OEShimService implements IOEShimService {
 			payload: node.payload || parentNode.payload,
 			contextValue: node.nodeTypeId
 		};
+		this.nodeHandleMap.set(generateNodeMapKey(viewId, newTreeItem), nodePath);
+		return newTreeItem;
 	}
 
 	public providerExists(providerId: string): boolean {
 		return this.oe.providerRegistered(providerId);
 	}
 }
+
+function generateSessionMapKey(viewId: string, node: ITreeItem): number {
+	return hash([viewId, node.childProvider, node.payload]);
+}
+
+function generateNodeMapKey(viewId: string, node: ITreeItem): number {
+	return hash([viewId, node.handle]);
+}

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

@@ -24,6 +24,13 @@ import { ObjectExplorerActionsContext } from 'sql/parts/objectExplorer/viewlet/o
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { IErrorMessageService } from 'sql/platform/errorMessage/common/errorMessageService';
 import { ConnectionViewletPanel } from 'sql/parts/dataExplorer/objectExplorer/connectionViewlet/connectionViewletPanel';
+import { ConnectionManagementService } from 'sql/platform/connection/common/connectionManagementService';
+import { CommandsRegistry } from 'vs/platform/commands/common/commands';
+import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
+import { MenuRegistry, MenuId } from 'vs/platform/actions/common/actions';
+import { ViewsRegistry } from 'vs/workbench/common/views';
+import { ICustomViewDescriptor, TreeViewItemHandleArg } from 'sql/workbench/common/views';
+import { IOEShimService } from 'sql/parts/objectExplorer/common/objectExplorerViewTreeShim';
 
 export class RefreshAction extends Action {
 
@@ -379,6 +386,41 @@ export class DeleteConnectionAction extends Action {
 	}
 }
 
+class DisconnectProfileAction extends Action {
+
+	constructor(
+		@IOEShimService private objectExplorerService: IOEShimService
+	) {
+		super(DisconnectConnectionAction.ID);
+	}
+	run(args: TreeViewItemHandleArg): Promise<boolean> {
+		if (args.$treeItem) {
+			return this.objectExplorerService.disconnectNode(args.$treeViewId, args.$treeItem).then(() => {
+				const { treeView } = (<ICustomViewDescriptor>ViewsRegistry.getView(args.$treeViewId));
+				// we need to collapse it then refresh it so that the tree doesn't try and use it's cache next time the user expands the node
+				return treeView.collapse(args.$treeItem).then(() => treeView.refresh([args.$treeItem]).then(() => true));
+			});
+		}
+		return Promise.resolve(true);
+	}
+}
+
+CommandsRegistry.registerCommand({
+	id: DisconnectConnectionAction.ID,
+	handler: (accessor, args: TreeViewItemHandleArg) => {
+		return accessor.get(IInstantiationService).createInstance(DisconnectProfileAction).run(args);
+	}
+});
+
+MenuRegistry.appendMenuItem(MenuId.DataExplorerContext, {
+	group: 'connection',
+	order: 4,
+	command: {
+		id: DisconnectConnectionAction.ID,
+		title: DisconnectConnectionAction.LABEL
+	}
+});
+
 /**
  * Action to clear search results
  */