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<number, string>();
-	private sessionMap = new Map<any, Map<number, string>>();
+	private nodeHandleMap = new Map<number, string>();
-	private nodeIdMap = new Map<string, 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[]> {
+	public async disconnectNode(viewId: string, node: ITreeItem): Promise<boolean> {
-		try {
+		// we assume only nodes with payloads can be connected
-			if (!this.sessionMap.has(identifier)) {
+		// check to make sure we have an existing connection
-				this.sessionMap.set(identifier, new Map<number, string>());
+		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))) {
+			return closed;
-				this.sessionMap.get(identifier).set(hash(node.payload || node.childProvider), await this.createSession(node.childProvider, node));
 		}
-			if (this.nodeIdMap.has(node.handle)) {
+		return Promise.resolve(false);
-				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);
+	private async getOrCreateSession(viewId: string, node: ITreeItem): Promise<string> {
-			let profile: IConnectionProfile = node.payload || {
+		// verify the map is correct
-				providerName: node.childProvider,
+		let key = generateSessionMapKey(viewId, node);
-				authenticationType: undefined,
+		if (!this.sessionMap.has(key)) {
-				azureTenantId: undefined,
+			this.sessionMap.set(key, await this.createSession(viewId, node.childProvider, node));
-				connectionName: undefined,
+		}
-				databaseName: undefined,
+		return this.sessionMap.get(key);
-				groupFullName: undefined,
+	}
-				groupId: undefined,
+
-				id: undefined,
+	public async getChildren(node: ITreeItem, viewId: string): Promise<ITreeItem[]> {
-				options: undefined,
+		if (node.payload) {
-				password: undefined,
+			let sessionId = await this.getOrCreateSession(viewId, node);
-				savePassword: undefined,
+			let requestHandle = this.nodeHandleMap.get(generateNodeMapKey(viewId, node)) || node.handle;
-				saveProfile: undefined,
+			let treeNode = new TreeNode(undefined, undefined, undefined, requestHandle, undefined, undefined, undefined, undefined, undefined, undefined);
-				serverName: undefined,
+			treeNode.connection = new ConnectionProfile(this.capabilities, node.payload);
-				userName: undefined,
+			return this.oe.resolveTreeNodeChildren({
-			};
-			treeNode.connection = new ConnectionProfile(this.capabilities, profile);
-			return TPromise.wrap(this.oe.resolveTreeNodeChildren({
 				success: undefined,
 				sessionId,
 				rootNode: undefined,
 				errorMessage: undefined
-			}, treeNode).then(e => e.map(n => this.mapNodeToITreeItem(n, node))));
+			}, treeNode).then(e => e.map(n => this.treeNodeToITreeItem(viewId, n, node)));
-		} catch (e) {
+		} else {
-			return TPromise.as([]);
+			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);
+		let nodePath = node.nodePath;
-		return {
+		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
  */

src/sql/workbench/browser/parts/views/customView.ts

@@ -12,7 +12,7 @@ import { IContextMenuService } from 'vs/platform/contextview/browser/contextView
 import { IMenuService, MenuId, MenuItemAction } from 'vs/platform/actions/common/actions';
 import { ContextAwareMenuItemActionItem, fillInActionBarActions, fillInContextMenuActions } from 'vs/platform/actions/browser/menuItemActionItem';
 import { IContextKeyService } from 'vs/platform/contextkey/common/contextkey';
-import { ITreeView, TreeItemCollapsibleState, ITreeViewDataProvider, TreeViewItemHandleArg, ViewContainer, ITreeItemLabel } from 'vs/workbench/common/views';
+import { TreeItemCollapsibleState, ITreeViewDataProvider, TreeViewItemHandleArg, ViewContainer, ITreeItemLabel } from 'vs/workbench/common/views';
 import { FileIconThemableWorkbenchTree } from 'vs/workbench/browser/parts/views/viewsViewlet';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { IProgressService2 } from 'vs/platform/progress/common/progress';
@@ -42,7 +42,7 @@ import { IMarkdownRenderResult } from 'vs/editor/contrib/markdown/markdownRender
 import { ILabelService } from 'vs/platform/label/common/label';
 import { dirname } from 'vs/base/common/resources';
 
-import { ITreeItem } from 'sql/workbench/common/views';
+import { ITreeItem, ITreeView } from 'sql/workbench/common/views';
 import { IOEShimService } from 'sql/parts/objectExplorer/common/objectExplorerViewTreeShim';
 import { equalsIgnoreCase } from 'vs/base/common/strings';
 
@@ -152,6 +152,7 @@ export class CustomTreeView extends Disposable implements ITreeView {
 		private id: string,
 		private container: ViewContainer,
 		@IExtensionService private extensionService: IExtensionService,
+		@IContextKeyService private contextKeyService: IContextKeyService,
 		@IWorkbenchThemeService private themeService: IWorkbenchThemeService,
 		@IInstantiationService private instantiationService: IInstantiationService,
 		@ICommandService private commandService: ICommandService,
@@ -315,7 +316,7 @@ export class CustomTreeView extends Disposable implements ITreeView {
 	private createTree() {
 		const actionItemProvider = (action: IAction) => action instanceof MenuItemAction ? this.instantiationService.createInstance(ContextAwareMenuItemActionItem, action) : undefined;
 		const menus = this.instantiationService.createInstance(TreeMenus, this.id);
-		const dataSource = this.instantiationService.createInstance(TreeDataSource, this, this.container);
+		const dataSource = this.instantiationService.createInstance(TreeDataSource, this, this.container, this.id);
 		const renderer = this.instantiationService.createInstance(TreeRenderer, this.id, menus, actionItemProvider);
 		const controller = this.instantiationService.createInstance(TreeController, this.id, menus);
 		this.tree = this.instantiationService.createInstance(FileIconThemableWorkbenchTree, this.treeContainer, { dataSource, renderer, controller }, {});
@@ -412,6 +413,14 @@ export class CustomTreeView extends Disposable implements ITreeView {
 		return Promise.arguments(null);
 	}
 
+	collapse(itemOrItems: ITreeItem | ITreeItem[]): Thenable<void> {
+		if(this.tree) {
+			itemOrItems = Array.isArray(itemOrItems) ? itemOrItems : [itemOrItems];
+			return this.tree.collapseAll(itemOrItems);
+		}
+		return Promise.arguments(null);
+	}
+
 	setSelection(items: ITreeItem[]): void {
 		if (this.tree) {
 			this.tree.setSelection(items, { source: 'api' });
@@ -496,6 +505,7 @@ class TreeDataSource implements IDataSource {
 	constructor(
 		private treeView: ITreeView,
 		private container: ViewContainer,
+		private id: string,
 		@IProgressService2 private progressService: IProgressService2,
 		@IOEShimService private objectExplorerService: IOEShimService
 	) {
@@ -514,7 +524,15 @@ class TreeDataSource implements IDataSource {
 
 	getChildren(tree: ITree, node: ITreeItem): Promise<any[]> {
 		if (node.childProvider) {
-			return this.progressService.withProgress({ location: this.container.id }, () => this.objectExplorerService.getChildren(node, this));
+			return this.progressService.withProgress({ location: this.container.id }, () => this.objectExplorerService.getChildren(node, this.id)).catch(e => {
+				// if some error is caused we assume something tangently happened
+				// i.e the user could retry if they wanted.
+				// So in order to enable this we need to tell the tree to refresh this node so it will ask us for the data again
+				setTimeout(() => {
+					tree.collapse(node).then(() => tree.refresh(node));
+				});
+				return [];
+			});
 		}
 		if (this.treeView.dataProvider) {
 			return this.progressService.withProgress({ location: this.container.id }, () => this.treeView.dataProvider.getChildren(node));
@@ -752,7 +770,7 @@ class TreeController extends WorkbenchTreeController {
 				}
 			},
 
-			getActionsContext: () => (<TreeViewItemHandleArg>{ $treeViewId: this.treeViewId, $treeItemHandle: node.handle }),
+			getActionsContext: () => (<TreeViewItemHandleArg>{ $treeViewId: this.treeViewId, $treeItemHandle: node.handle, $treeItem: node }),
 
 			actionRunner: new MultipleSelectionActionRunner(() => tree.getSelection())
 		});
@@ -795,17 +813,30 @@ class TreeMenus extends Disposable implements IDisposable {
 	}
 
 	getResourceActions(element: ITreeItem): IAction[] {
-		return this.getActions(MenuId.ViewItemContext, { key: 'viewItem', value: element.contextValue }).primary;
+		return this.mergeActions([
+			this.getActions(MenuId.ViewItemContext, { key: 'viewItem', value: element.contextValue }).primary,
+			this.getActions(MenuId.DataExplorerContext).primary
+		]);
 	}
 
 	getResourceContextActions(element: ITreeItem): IAction[] {
-		return this.getActions(MenuId.ViewItemContext, { key: 'viewItem', value: element.contextValue }).secondary;
+		return this.mergeActions([
+			this.getActions(MenuId.ViewItemContext, { key: 'viewItem', value: element.contextValue }).secondary,
+			this.getActions(MenuId.DataExplorerContext).secondary
+		]);
 	}
 
-	private getActions(menuId: MenuId, context: { key: string, value: string }): { primary: IAction[]; secondary: IAction[]; } {
+	private mergeActions(actions: IAction[][]): IAction[] {
+		return actions.reduce((p, c) => p.concat(...c.filter(a => p.findIndex(x => x.id === a.id) === -1)), [] as IAction[]);
+	}
+
+	private getActions(menuId: MenuId, context?: { key: string, value: string }): { primary: IAction[]; secondary: IAction[]; } {
 		const contextKeyService = this.contextKeyService.createScoped();
 		contextKeyService.createKey('view', this.id);
+
+		if (context) {
 			contextKeyService.createKey(context.key, context.value);
+		}
 
 		const menu = this.menuService.createMenu(menuId, contextKeyService);
 		const primary: IAction[] = [];

src/sql/workbench/common/views.ts

@@ -3,7 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { ITreeViewDataProvider, ITreeItem as vsITreeItem } from 'vs/workbench/common/views';
+import { ITreeViewDataProvider, ITreeItem as vsITreeItem, IViewDescriptor, ITreeView as vsITreeView } from 'vs/workbench/common/views';
 import { IConnectionProfile } from 'azdata';
 
 export interface ITreeComponentItem extends vsITreeItem {
@@ -22,3 +22,21 @@ export interface ITreeItem extends vsITreeItem {
 	childProvider?: string;
 	payload?: IConnectionProfile; // its possible we will want this to be more generic
 }
+
+export interface ITreeView extends vsITreeView {
+
+	collapse(itemOrItems: ITreeItem | ITreeItem[]): Thenable<void>;
+
+}
+
+export type TreeViewItemHandleArg = {
+	$treeViewId: string,
+	$treeItemHandle: string,
+	$treeItem?: ITreeItem
+};
+
+export interface ICustomViewDescriptor extends IViewDescriptor {
+
+	readonly treeView: ITreeView;
+
+}

src/vs/platform/actions/common/actions.ts

@@ -98,7 +98,8 @@ export const enum MenuId {
 	ViewTitle,
 	// {{SQL CARBON EDIT}}
 	ObjectExplorerItemContext,
-	NotebookToolbar
+	NotebookToolbar,
+	DataExplorerContext
 }
 
 export interface IMenuActionOptions {

src/vs/workbench/services/actions/electron-browser/menusExtensionPoint.ts

@@ -43,6 +43,7 @@ namespace schema {
 			// {{SQL CARBON EDIT}}
 			case 'objectExplorer/item/context': return MenuId.ObjectExplorerItemContext;
 			case 'notebook/toolbar': return MenuId.NotebookToolbar;
+			case 'dataExplorer/context': return MenuId.DataExplorerContext;
 		}
 
 		return void 0;