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/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);
+	}
+});