Add Data Explorer Context and apply to disconnect (#4265)

* 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

* add node context on data explorer

* formatting

* fix various errors with how the context menus work

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

@@ -4,22 +4,19 @@
  *--------------------------------------------------------------------------------------------*/
 'use strict';
 
-import { ITreeItem } from 'sql/workbench/common/views';
-import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
-import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
-import { TreeNode } from 'sql/parts/objectExplorer/common/treeNode';
-import { IConnectionManagementService, ConnectionType } from 'sql/platform/connection/common/connectionManagement';
 import { Deferred } from 'sql/base/common/promise';
-import { IObjectExplorerService } from 'sql/workbench/services/objectExplorer/common/objectExplorerService';
+import { TreeNode } from 'sql/parts/objectExplorer/common/treeNode';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
+import { ConnectionType, IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
+import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
+import { ITreeItem } from 'sql/workbench/common/views';
 import { IConnectionDialogService } from 'sql/workbench/services/connection/common/connectionDialogService';
-
-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 { IObjectExplorerService } from 'sql/workbench/services/objectExplorer/common/objectExplorerService';
 import { hash } from 'vs/base/common/hash';
+import { Disposable } from 'vs/base/common/lifecycle';
 import { generateUuid } from 'vs/base/common/uuid';
+import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
+import { TreeItemCollapsibleState } from 'vs/workbench/common/views';
 
 export const SERVICE_ID = 'oeShimService';
 export const IOEShimService = createDecorator<IOEShimService>(SERVICE_ID);
@@ -29,9 +26,10 @@ export interface IOEShimService {
 	getChildren(node: ITreeItem, viewId: string): Promise<ITreeItem[]>;
 	disconnectNode(viewId: string, node: ITreeItem): Promise<boolean>;
 	providerExists(providerId: string): boolean;
+	isNodeConnected(viewId: string, node: ITreeItem): boolean;
 }
 
-export class OEShimService implements IOEShimService {
+export class OEShimService extends Disposable implements IOEShimService {
 	_serviceBrand: any;
 
 	private sessionMap = new Map<number, string>();
@@ -43,6 +41,7 @@ export class OEShimService implements IOEShimService {
 		@IConnectionDialogService private cd: IConnectionDialogService,
 		@ICapabilitiesService private capabilities: ICapabilitiesService
 	) {
+		super();
 	}
 
 	private async createSession(viewId: string, providerId: string, node: ITreeItem): Promise<string> {
@@ -139,6 +138,10 @@ export class OEShimService implements IOEShimService {
 	public providerExists(providerId: string): boolean {
 		return this.oe.providerRegistered(providerId);
 	}
+
+	public isNodeConnected(viewId: string, node: ITreeItem): boolean {
+		return this.sessionMap.has(generateSessionMapKey(viewId, node));
+	}
 }
 
 function generateSessionMapKey(viewId: string, node: ITreeItem): number {

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

@@ -386,41 +386,6 @@ 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

@@ -45,6 +45,7 @@ import { dirname } from 'vs/base/common/resources';
 import { ITreeItem, ITreeView } from 'sql/workbench/common/views';
 import { IOEShimService } from 'sql/parts/objectExplorer/common/objectExplorerViewTreeShim';
 import { equalsIgnoreCase } from 'vs/base/common/strings';
+import { NodeContextKey } from 'sql/workbench/parts/dataExplorer/common/nodeContext';
 
 class TitleMenus implements IDisposable {
 
@@ -148,11 +149,12 @@ export class CustomTreeView extends Disposable implements ITreeView {
 	private _onDidChangeActions: Emitter<void> = this._register(new Emitter<void>());
 	readonly onDidChangeActions: Event<void> = this._onDidChangeActions.event;
 
+	private nodeContext: NodeContextKey;
+
 	constructor(
 		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,
@@ -176,6 +178,9 @@ export class CustomTreeView extends Disposable implements ITreeView {
 				this.markdownResult.dispose();
 			}
 		}));
+
+		this.nodeContext = this._register(instantiationService.createInstance(NodeContextKey));
+
 		this.create();
 	}
 
@@ -326,6 +331,10 @@ export class CustomTreeView extends Disposable implements ITreeView {
 		this._register(this.tree.onDidExpandItem(e => this._onDidExpandItem.fire(e.item.getElement())));
 		this._register(this.tree.onDidCollapseItem(e => this._onDidCollapseItem.fire(e.item.getElement())));
 		this._register(this.tree.onDidChangeSelection(e => this._onDidChangeSelection.fire(e.selection)));
+		// Update resource context based on focused element
+		this._register(this.tree.onDidChangeFocus((e: { focus: ITreeItem }) => {
+			this.nodeContext.set({ node: e.focus, viewId: this.id });
+		}));
 		this.tree.setInput(this.root).then(() => this.updateContentAreas());
 	}
 
@@ -414,7 +423,7 @@ export class CustomTreeView extends Disposable implements ITreeView {
 	}
 
 	collapse(itemOrItems: ITreeItem | ITreeItem[]): Thenable<void> {
-		if(this.tree) {
+		if (this.tree) {
 			itemOrItems = Array.isArray(itemOrItems) ? itemOrItems : [itemOrItems];
 			return this.tree.collapseAll(itemOrItems);
 		}
@@ -815,14 +824,14 @@ class TreeMenus extends Disposable implements IDisposable {
 	getResourceActions(element: ITreeItem): IAction[] {
 		return this.mergeActions([
 			this.getActions(MenuId.ViewItemContext, { key: 'viewItem', value: element.contextValue }).primary,
-			this.getActions(MenuId.DataExplorerContext).primary
+			this.getActions(MenuId.DataExplorerContext, { key: 'viewItem', value: element.contextValue }).primary
 		]);
 	}
 
 	getResourceContextActions(element: ITreeItem): IAction[] {
 		return this.mergeActions([
 			this.getActions(MenuId.ViewItemContext, { key: 'viewItem', value: element.contextValue }).secondary,
-			this.getActions(MenuId.DataExplorerContext).secondary
+			this.getActions(MenuId.DataExplorerContext, { key: 'viewItem', value: element.contextValue }).secondary
 		]);
 	}
 
@@ -830,13 +839,10 @@ class TreeMenus extends Disposable implements IDisposable {
 		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[]; } {
+	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);
-		}
+		contextKeyService.createKey(context.key, context.value);
 
 		const menu = this.menuService.createMenu(menuId, contextKeyService);
 		const primary: IAction[] = [];

src/sql/workbench/parts/dataExplorer/common/nodeContext.ts

@@ -0,0 +1,73 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { ConnectionContextKey } from 'sql/parts/connection/common/connectionContextKey';
+import { IOEShimService } from 'sql/parts/objectExplorer/common/objectExplorerViewTreeShim';
+import { ITreeItem } from 'sql/workbench/common/views';
+import { Disposable } from 'vs/base/common/lifecycle';
+import { IContextKey, IContextKeyService, RawContextKey } from 'vs/platform/contextkey/common/contextkey';
+
+export interface INodeContextValue {
+	node: ITreeItem;
+	viewId: string;
+}
+
+export class NodeContextKey extends Disposable implements IContextKey<INodeContextValue> {
+
+	static IsConnectable = new RawContextKey<boolean>('isConnectable', false);
+	static IsConnected = new RawContextKey<boolean>('isConnected', false);
+	static ViewId = new RawContextKey<string>('view', undefined);
+	static ViewItem = new RawContextKey<string>('viewItem', undefined);
+	static Node = new RawContextKey<INodeContextValue>('node', undefined);
+
+	private readonly _connectionContextKey: ConnectionContextKey;
+	private readonly _connectableKey: IContextKey<boolean>;
+	private readonly _connectedKey: IContextKey<boolean>;
+	private readonly _viewIdKey: IContextKey<string>;
+	private readonly _viewItemKey: IContextKey<string>;
+	private readonly _nodeContextKey: IContextKey<INodeContextValue>;
+
+	constructor(
+		@IContextKeyService contextKeyService: IContextKeyService,
+		@IOEShimService private oeService: IOEShimService
+	) {
+		super();
+
+		this._connectableKey = NodeContextKey.IsConnectable.bindTo(contextKeyService);
+		this._connectedKey = NodeContextKey.IsConnected.bindTo(contextKeyService);
+		this._viewIdKey = NodeContextKey.ViewId.bindTo(contextKeyService);
+		this._viewItemKey = NodeContextKey.ViewItem.bindTo(contextKeyService);
+		this._nodeContextKey = NodeContextKey.Node.bindTo(contextKeyService);
+		this._connectionContextKey = new ConnectionContextKey(contextKeyService);
+	}
+
+	set(value: INodeContextValue) {
+		if (value.node && value.node.payload) {
+			this._connectableKey.set(true);
+			this._connectedKey.set(this.oeService.isNodeConnected(value.viewId, value.node));
+			this._connectionContextKey.set(value.node.payload);
+		} else {
+			this._connectableKey.set(false);
+			this._connectedKey.set(false);
+			this._connectionContextKey.reset();
+		}
+		this._nodeContextKey.set(value);
+		this._viewIdKey.set(value.viewId);
+		this._viewItemKey.set(value.node.contextValue);
+	}
+
+	reset(): void {
+		this._viewIdKey.reset();
+		this._viewItemKey.reset();
+		this._connectableKey.reset();
+		this._connectedKey.reset();
+		this._connectionContextKey.reset();
+		this._nodeContextKey.reset();
+	}
+
+	get(): INodeContextValue | undefined {
+		return this._nodeContextKey.get();
+	}
+}

src/sql/workbench/parts/dataExplorer/electron-browser/nodeActions.contribution.ts

@@ -0,0 +1,19 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { NodeContextKey } from 'sql/workbench/parts/dataExplorer/common/nodeContext';
+import { localize } from 'vs/nls';
+import { MenuId, MenuRegistry } from 'vs/platform/actions/common/actions';
+import { DISCONNECT_COMMAND_ID } from './nodeCommands';
+
+MenuRegistry.appendMenuItem(MenuId.DataExplorerContext, {
+	group: 'connection',
+	order: 4,
+	command: {
+		id: DISCONNECT_COMMAND_ID,
+		title: localize('disconnect', 'Disconnect')
+	},
+	when: NodeContextKey.IsConnected
+});

src/sql/workbench/parts/dataExplorer/electron-browser/nodeCommands.ts

@@ -0,0 +1,26 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IOEShimService } from 'sql/parts/objectExplorer/common/objectExplorerViewTreeShim';
+import { ICustomViewDescriptor, TreeViewItemHandleArg } from 'sql/workbench/common/views';
+import { CommandsRegistry } from 'vs/platform/commands/common/commands';
+import { ViewsRegistry } from 'vs/workbench/common/views';
+
+export const DISCONNECT_COMMAND_ID = 'dataExplorer.disconnect';
+
+CommandsRegistry.registerCommand({
+	id: DISCONNECT_COMMAND_ID,
+	handler: (accessor, args: TreeViewItemHandleArg) => {
+		if (args.$treeItem) {
+			const oeService = accessor.get(IOEShimService);
+			return oeService.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);
+	}
+});