Data explorer/context menu initial (#6264)

* fixed context for data explorer

* added more files

* initial servers and database context menu actions finished

* added all actions for servers and databases with correct conditions

* added nodetype and nodelabel for subtype actions

* added nodeinfo logic to oe shim

* fixed context for cms

* added all scripting actions to data explorer

* review comments

* fix import

* fix correct context key

* removed unused import

* separate PR for commands and menus added

* rename mssql context menu nodes

* remove command id constants

src/sql/workbench/api/common/sqlExtHostTypes.ts

@@ -191,6 +191,11 @@ export enum StepCompletionAction {
 	GoToStep = 4
 }
 
+export enum ExtensionNodeType {
+	Server = 'Server',
+	Database = 'Database'
+}
+
 export interface IComponentShape {
 	type: ModelComponentTypes;
 	id: string;

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

@@ -21,7 +21,9 @@ export class ExtHostConnectionManagement extends ExtHostConnectionManagementShap
 	public $getCurrentConnection(): Thenable<azdata.connection.ConnectionProfile> {
 		let connection: any = this._proxy.$getCurrentConnection();
 		connection.then((conn) => {
-			conn.providerId = conn.providerName;
+			if (conn && conn.providerId) {
+				conn.providerId = conn.providerName;
+			}
 		});
 		return connection;
 	}

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

@@ -552,7 +552,8 @@ export function createApiFactory(
 				ColumnType: sqlExtHostTypes.ColumnType,
 				ActionOnCellCheckboxCheck: sqlExtHostTypes.ActionOnCellCheckboxCheck,
 				StepCompletionAction: sqlExtHostTypes.StepCompletionAction,
-				AgentSubSystem: sqlExtHostTypes.AgentSubSystem
+				AgentSubSystem: sqlExtHostTypes.AgentSubSystem,
+				ExtensionNodeType: sqlExtHostTypes.ExtensionNodeType
 			};
 		},
 

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

@@ -43,7 +43,6 @@ import { dirname } from 'vs/base/common/resources';
 
 import { ITreeItem, ITreeView } from 'sql/workbench/common/views';
 import { IOEShimService } from 'sql/workbench/parts/objectExplorer/common/objectExplorerViewTreeShim';
-import { equalsIgnoreCase } from 'vs/base/common/strings';
 import { NodeContextKey } from 'sql/workbench/parts/dataExplorer/common/nodeContext';
 import { fillInActionBarActions, fillInContextMenuActions, ContextAwareMenuEntryActionViewItem } from 'vs/platform/actions/browser/menuEntryActionViewItem';
 import { mssqlProviderName } from 'sql/platform/connection/common/constants';

src/sql/workbench/common/views.ts

@@ -5,6 +5,7 @@
 
 import { ITreeViewDataProvider, ITreeItem as vsITreeItem, IViewDescriptor, ITreeView as vsITreeView } from 'vs/workbench/common/views';
 import { IConnectionProfile } from 'azdata';
+import { ExtensionNodeType } from 'sql/workbench/api/common/sqlExtHostTypes';
 
 export interface ITreeComponentItem extends vsITreeItem {
 	checked?: boolean;
@@ -22,6 +23,7 @@ export interface ITreeItem extends vsITreeItem {
 	childProvider?: string;
 	payload?: IConnectionProfile; // its possible we will want this to be more generic
 	sqlIcon?: string;
+	type?: ExtensionNodeType;
 }
 
 export interface ITreeView extends vsITreeView {

src/sql/workbench/parts/dataExplorer/browser/dataExplorerViewlet.ts

@@ -112,6 +112,7 @@ export class DataExplorerViewlet extends ViewContainerViewlet {
 				actions.push(...group[1]);
 			}
 		});
+		menu.dispose();
 		return actions;
 	}
 

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

@@ -0,0 +1,159 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as os from 'os';
+import { INodeContextValue } from 'sql/workbench/parts/dataExplorer/common/nodeContext';
+import { RawContextKey, IContextKeyService, IContextKey } from 'vs/platform/contextkey/common/contextkey';
+import { Disposable } from 'vs/base/common/lifecycle';
+import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
+import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
+import { mssqlProviderName } from 'sql/platform/connection/common/constants';
+import { NodeType } from 'sql/workbench/parts/objectExplorer/common/nodeType';
+import { ExtensionNodeType } from 'sql/workbench/api/common/sqlExtHostTypes';
+
+export class MssqlNodeContext extends Disposable {
+
+	static readonly canSelect = new Set([NodeType.Table, NodeType.View]);
+	static readonly canEditData = new Set([NodeType.Table]);
+	static readonly canCreateOrDelete = new Set([NodeType.AggregateFunction, NodeType.PartitionFunction, NodeType.ScalarValuedFunction,
+	NodeType.Schema, NodeType.StoredProcedure, NodeType.Table, NodeType.TableValuedFunction,
+	NodeType.User, NodeType.UserDefinedTableType, NodeType.View]);
+	static readonly canExecute = new Set([NodeType.StoredProcedure]);
+	static readonly canAlter = new Set([NodeType.AggregateFunction, NodeType.PartitionFunction, NodeType.ScalarValuedFunction,
+	NodeType.StoredProcedure, NodeType.TableValuedFunction, NodeType.View]);
+
+	// General node context keys
+	static NodeProvider = new RawContextKey<string>('nodeProvider', undefined);
+	static IsDatabaseOrServer = new RawContextKey<boolean>('isDatabaseOrServer', false);
+	static IsWindows = new RawContextKey<boolean>('isWindows', os.platform() === 'win32');
+	static IsCloud = new RawContextKey<boolean>('isCloud', false);
+	static NodeType = new RawContextKey<string>('nodeType', undefined);
+	static NodeLabel = new RawContextKey<string>('nodeLabel', undefined);
+
+	// Scripting context keys
+	static CanScriptAsSelect = new RawContextKey<boolean>('canScriptAsSelect', false);
+	static CanEditData = new RawContextKey<boolean>('canEditData', false);
+	static CanScriptAsCreateOrDelete = new RawContextKey<boolean>('canScriptAsCreateOeDelete', false);
+	static CanScriptAsExecute = new RawContextKey<boolean>('canScriptAsExecute', false);
+	static CanScriptAsAlter = new RawContextKey<boolean>('canScriptAsAlter', false);
+
+	private nodeProviderKey: IContextKey<string>;
+	private isCloudKey: IContextKey<boolean>;
+	private nodeTypeKey: IContextKey<string>;
+	private nodeLabelKey: IContextKey<string>;
+	private isDatabaseOrServerKey: IContextKey<boolean>;
+
+	private canScriptAsSelectKey: IContextKey<boolean>;
+	private canEditDataKey: IContextKey<boolean>;
+	private canScriptAsCreateOrDeleteKey: IContextKey<boolean>;
+	private canScriptAsExecuteKey: IContextKey<boolean>;
+	private canScriptAsAlterKey: IContextKey<boolean>;
+
+	constructor(
+		private nodeContextValue: INodeContextValue,
+		@IContextKeyService private contextKeyService: IContextKeyService,
+		@IConnectionManagementService private connectionManagementService: IConnectionManagementService,
+		@ICapabilitiesService private capabilitiesService: ICapabilitiesService
+	) {
+		super();
+		this.bindContextKeys();
+
+		// Set additional node context keys
+		if (this.nodeContextValue.node) {
+			const node = this.nodeContextValue.node;
+			if (node.payload) {
+				this.setNodeProvider();
+				this.setIsCloud();
+				if (node.type) {
+					this.setIsDatabaseOrServer();
+					this.nodeTypeKey.set(node.type);
+				} else if (node.contextValue && node.providerHandle === mssqlProviderName) {
+					this.setIsDatabaseOrServer();
+					this.setScriptingContextKeys();
+					this.nodeTypeKey.set(node.contextValue);
+				}
+			}
+			if (node.label) {
+				this.nodeLabelKey.set(node.label.label);
+			}
+		}
+	}
+
+	private bindContextKeys(): void {
+		this.isCloudKey = MssqlNodeContext.IsCloud.bindTo(this.contextKeyService);
+		this.nodeTypeKey = MssqlNodeContext.NodeType.bindTo(this.contextKeyService);
+		this.nodeLabelKey = MssqlNodeContext.NodeLabel.bindTo(this.contextKeyService);
+		this.isDatabaseOrServerKey = MssqlNodeContext.IsDatabaseOrServer.bindTo(this.contextKeyService);
+		this.canScriptAsSelectKey = MssqlNodeContext.CanScriptAsSelect.bindTo(this.contextKeyService);
+		this.canEditDataKey = MssqlNodeContext.CanEditData.bindTo(this.contextKeyService);
+		this.canScriptAsCreateOrDeleteKey = MssqlNodeContext.CanScriptAsCreateOrDelete.bindTo(this.contextKeyService);
+		this.canScriptAsExecuteKey = MssqlNodeContext.CanScriptAsExecute.bindTo(this.contextKeyService);
+		this.canScriptAsAlterKey = MssqlNodeContext.CanScriptAsAlter.bindTo(this.contextKeyService);
+		this.nodeProviderKey = MssqlNodeContext.NodeProvider.bindTo(this.contextKeyService);
+	}
+
+	/**
+	 * Helper function to get the node provider
+	 */
+	private setNodeProvider(): void {
+		if (this.nodeContextValue.node.payload.providerName) {
+			this.nodeProviderKey.set(this.nodeContextValue.node.payload.providerName);
+		} else if (this.nodeContextValue.node.childProvider) {
+			this.nodeProviderKey.set(this.nodeContextValue.node.childProvider);
+		}
+	}
+
+	/**
+	 * Helper function to tell whether a connected node is cloud or not
+	 */
+	private setIsCloud(): void {
+		const profile = new ConnectionProfile(this.capabilitiesService,
+			this.nodeContextValue.node.payload);
+		const connection = this.connectionManagementService.findExistingConnection(profile);
+		if (connection) {
+			const serverInfo = this.connectionManagementService.getServerInfo(connection.id);
+			if (serverInfo.isCloud) {
+				this.isCloudKey.set(true);
+			}
+		}
+	}
+
+	/**
+	 * Helper function to tell whether a connected node is a database or a
+	 * server or not. Added this key because this is easier to write than
+	 * writing an OR statement in ContextKeyExpr
+	 */
+	private setIsDatabaseOrServer(): void {
+		const isDatabaseOrServer = (this.nodeContextValue.node.contextValue === NodeType.Server ||
+			this.nodeContextValue.node.contextValue === NodeType.Database ||
+			this.nodeContextValue.node.type === ExtensionNodeType.Server ||
+			this.nodeContextValue.node.type === ExtensionNodeType.Database);
+		this.isDatabaseOrServerKey.set(isDatabaseOrServer);
+	}
+
+	/**
+	 * Helper function to get the correct context from node for showing
+	 * scripting context menu actions
+	 */
+	private setScriptingContextKeys(): void {
+		const nodeType = this.nodeContextValue.node.contextValue;
+		if (MssqlNodeContext.canCreateOrDelete.has(nodeType)) {
+			this.canScriptAsCreateOrDeleteKey.set(true);
+		}
+		if (MssqlNodeContext.canEditData.has(nodeType)) {
+			this.canEditDataKey.set(true);
+		}
+		if (MssqlNodeContext.canAlter.has(nodeType)) {
+			this.canScriptAsAlterKey.set(true);
+		}
+		if (MssqlNodeContext.canExecute.has(nodeType)) {
+			this.canScriptAsExecuteKey.set(true);
+		}
+		if (MssqlNodeContext.canSelect.has(nodeType)) {
+			this.canScriptAsSelectKey.set(true);
+		}
+	}
+}

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

@@ -9,6 +9,9 @@ import { ITreeItem } from 'sql/workbench/common/views';
 import { Disposable } from 'vs/base/common/lifecycle';
 import { IContextKey, IContextKeyService, RawContextKey } from 'vs/platform/contextkey/common/contextkey';
 import { IQueryManagementService } from 'sql/platform/query/common/queryManagement';
+import { MssqlNodeContext } from 'sql/workbench/parts/dataExplorer/common/mssqlNodeContext';
+import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
+import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 
 export interface INodeContextValue {
 	node: ITreeItem;
@@ -30,10 +33,14 @@ export class NodeContextKey extends Disposable implements IContextKey<INodeConte
 	private readonly _viewItemKey: IContextKey<string>;
 	private readonly _nodeContextKey: IContextKey<INodeContextValue>;
 
+	private _nodeContextUtils: MssqlNodeContext;
+
 	constructor(
-		@IContextKeyService contextKeyService: IContextKeyService,
+		@IContextKeyService private contextKeyService: IContextKeyService,
 		@IOEShimService private oeService: IOEShimService,
-		@IQueryManagementService queryManagementService: IQueryManagementService
+		@IQueryManagementService queryManagementService: IQueryManagementService,
+		@IConnectionManagementService private connectionManagementService: IConnectionManagementService,
+		@ICapabilitiesService private capabilitiesService: ICapabilitiesService
 	) {
 		super();
 
@@ -62,6 +69,8 @@ export class NodeContextKey extends Disposable implements IContextKey<INodeConte
 		}
 		this._nodeContextKey.set(value);
 		this._viewIdKey.set(value.viewId);
+		this._nodeContextUtils = new MssqlNodeContext(this._nodeContextKey.get(), this.contextKeyService,
+			this.connectionManagementService, this.capabilitiesService);
 	}
 
 	reset(): void {
@@ -71,6 +80,7 @@ export class NodeContextKey extends Disposable implements IContextKey<INodeConte
 		this._connectedKey.reset();
 		this._connectionContextKey.reset();
 		this._nodeContextKey.reset();
+		this._nodeContextUtils.dispose();
 	}
 
 	get(): INodeContextValue | undefined {

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

@@ -4,10 +4,26 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { localize } from 'vs/nls';
+import { mssqlProviderName } from 'sql/platform/connection/common/constants';
 import { MenuId, MenuRegistry } from 'vs/platform/actions/common/actions';
-import { DISCONNECT_COMMAND_ID, MANAGE_COMMAND_ID, NEW_QUERY_COMMAND_ID, REFRESH_COMMAND_ID } from './nodeCommands';
-import { ContextKeyExpr, ContextKeyRegexExpr } from 'vs/platform/contextkey/common/contextkey';
+import {
+	DISCONNECT_COMMAND_ID, MANAGE_COMMAND_ID, NEW_QUERY_COMMAND_ID, REFRESH_COMMAND_ID
+} from './nodeCommands';
+import { ContextKeyExpr } from 'vs/platform/contextkey/common/contextkey';
 import { NodeContextKey } from 'sql/workbench/parts/dataExplorer/common/nodeContext';
+import { MssqlNodeContext } from 'sql/workbench/parts/dataExplorer/common/mssqlNodeContext';
+
+
+// Disconnect
+MenuRegistry.appendMenuItem(MenuId.DataExplorerContext, {
+	group: 'connection',
+	order: 4,
+	command: {
+		id: DISCONNECT_COMMAND_ID,
+		title: localize('disconnect', 'Disconnect')
+	},
+	when: NodeContextKey.IsConnected //add for folders
+});
 
 MenuRegistry.appendMenuItem(MenuId.DataExplorerContext, {
 	group: 'connection',
@@ -16,13 +32,12 @@ MenuRegistry.appendMenuItem(MenuId.DataExplorerContext, {
 		id: DISCONNECT_COMMAND_ID,
 		title: localize('disconnect', 'Disconnect')
 	},
-	when: NodeContextKey.IsConnected
+	when: ContextKeyExpr.and(NodeContextKey.IsConnected,
+		MssqlNodeContext.NodeProvider.isEqualTo(mssqlProviderName),
+		MssqlNodeContext.IsDatabaseOrServer)
 });
 
-// The weird regex is because we want this to generically apply to Database and Server nodes but right now
-// there isn't a consistent standard for the values there. We can't just search for database or server being
-// in the string at all because there's lots of node types which have those values (such as ServerLevelLogin)
-// that we don't want these menu items showing up on.
+// New Query
 MenuRegistry.appendMenuItem(MenuId.DataExplorerContext, {
 	group: 'connection',
 	order: 2,
@@ -30,13 +45,10 @@ MenuRegistry.appendMenuItem(MenuId.DataExplorerContext, {
 		id: NEW_QUERY_COMMAND_ID,
 		title: localize('newQuery', 'New Query')
 	},
-	when: ContextKeyExpr.and(
-		NodeContextKey.IsConnectable,
-		new ContextKeyRegexExpr('viewItem', /.+itemType\.database.*|^database$/i))
+	when: MssqlNodeContext.IsDatabaseOrServer
 });
 
-// Note that we don't show this for Databases under Server nodes (viewItem == Database) because
-// of an issue there where the connection always being master instead of the actual DB
+// Manage
 MenuRegistry.appendMenuItem(MenuId.DataExplorerContext, {
 	group: 'connection',
 	order: 1,
@@ -44,17 +56,17 @@ MenuRegistry.appendMenuItem(MenuId.DataExplorerContext, {
 		id: MANAGE_COMMAND_ID,
 		title: localize('manage', 'Manage')
 	},
-	when: ContextKeyExpr.and(
-		NodeContextKey.IsConnectable,
-		new ContextKeyRegexExpr('viewItem', /.+itemType\.database.*/i))
+	when: MssqlNodeContext.IsDatabaseOrServer
 });
 
+
+// Refresh
 MenuRegistry.appendMenuItem(MenuId.DataExplorerContext, {
 	group: 'connection',
-	order: 4,
+	order: 6,
 	command: {
 		id: REFRESH_COMMAND_ID,
 		title: localize('refresh', 'Refresh')
 	},
-	when: NodeContextKey.IsConnected
-});
+	when: NodeContextKey.IsConnectable
+});

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

@@ -10,7 +10,7 @@ import { ConnectionProfile } from 'sql/platform/connection/common/connectionProf
 import { generateUri } from 'sql/platform/connection/common/utils';
 import { ICustomViewDescriptor, TreeViewItemHandleArg } from 'sql/workbench/common/views';
 import { IQueryEditorService } from 'sql/workbench/services/queryEditor/common/queryEditorService';
-import { CommandsRegistry } from 'vs/platform/commands/common/commands';
+import { CommandsRegistry, ICommandService } from 'vs/platform/commands/common/commands';
 import { IViewsRegistry, Extensions } from 'vs/workbench/common/views';
 import { IProgressService2 } from 'vs/platform/progress/common/progress';
 import { Registry } from 'vs/platform/registry/common/platform';
@@ -20,6 +20,7 @@ export const MANAGE_COMMAND_ID = 'dataExplorer.manage';
 export const NEW_QUERY_COMMAND_ID = 'dataExplorer.newQuery';
 export const REFRESH_COMMAND_ID = 'dataExplorer.refresh';
 
+// Disconnect
 CommandsRegistry.registerCommand({
 	id: DISCONNECT_COMMAND_ID,
 	handler: (accessor, args: TreeViewItemHandleArg) => {
@@ -35,6 +36,7 @@ CommandsRegistry.registerCommand({
 	}
 });
 
+// New Query
 CommandsRegistry.registerCommand({
 	id: NEW_QUERY_COMMAND_ID,
 	handler: (accessor, args: TreeViewItemHandleArg) => {
@@ -58,6 +60,7 @@ CommandsRegistry.registerCommand({
 	}
 });
 
+// Manage
 CommandsRegistry.registerCommand({
 	id: MANAGE_COMMAND_ID,
 	handler: (accessor, args: TreeViewItemHandleArg) => {
@@ -79,6 +82,7 @@ CommandsRegistry.registerCommand({
 	}
 });
 
+// Refresh
 CommandsRegistry.registerCommand({
 	id: REFRESH_COMMAND_ID,
 	handler: (accessor, args: TreeViewItemHandleArg) => {
@@ -94,4 +98,4 @@ CommandsRegistry.registerCommand({
 		}
 		return Promise.resolve(true);
 	}
-});
+});

src/sql/workbench/parts/notebook/notebookActions.ts

@@ -558,7 +558,6 @@ export class NewNotebookAction extends Action {
 
 	run(context?: azdata.ConnectedContext): Promise<void> {
 		return this.commandService.executeCommand(NewNotebookAction.INTERNAL_NEW_NOTEBOOK_CMD_ID, context);
-
 	}
 
 }

src/sql/workbench/parts/objectExplorer/browser/objectExplorerActions.ts

@@ -60,7 +60,6 @@ export class OEAction extends ExecuteCommandAction {
 	public async run(actionContext: any): Promise<boolean> {
 		this._treeSelectionHandler = this._instantiationService.createInstance(TreeSelectionHandler);
 
-
 		let profile: IConnectionProfile;
 		if (actionContext instanceof ObjectExplorerActionsContext) {
 			if (actionContext.isConnectionNode) {
@@ -240,7 +239,7 @@ export class OEScriptCreateAction extends ScriptCreateAction {
 		@IScriptingService protected _scriptingService: IScriptingService,
 		@IObjectExplorerService private _objectExplorerService: IObjectExplorerService,
 		@IInstantiationService private _instantiationService: IInstantiationService,
-		@IErrorMessageService protected _errorMessageService: IErrorMessageService
+		@IErrorMessageService protected _errorMessageService: IErrorMessageService,
 	) {
 		super(id, label, _queryEditorService, _connectionManagementService, _scriptingService, _errorMessageService);
 	}

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

@@ -3,7 +3,7 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Deferred } from 'sql/base/common/promise';
+import * as azdata from 'azdata';
 import { TreeNode } from 'sql/workbench/parts/objectExplorer/common/treeNode';
 import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 import { ConnectionType, IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
@@ -17,6 +17,7 @@ import { generateUuid } from 'vs/base/common/uuid';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import { TreeItemCollapsibleState } from 'vs/workbench/common/views';
 import { localize } from 'vs/nls';
+import { NodeType } from 'sql/workbench/parts/objectExplorer/common/nodeType';
 
 export const SERVICE_ID = 'oeShimService';
 export const IOEShimService = createDecorator<IOEShimService>(SERVICE_ID);
@@ -27,6 +28,7 @@ export interface IOEShimService {
 	disconnectNode(viewId: string, node: ITreeItem): Promise<boolean>;
 	providerExists(providerId: string): boolean;
 	isNodeConnected(viewId: string, node: ITreeItem): boolean;
+	getNodeInfoForTreeItem(treeItem: ITreeItem): azdata.NodeInfo;
 }
 
 export class OEShimService extends Disposable implements IOEShimService {
@@ -34,6 +36,7 @@ export class OEShimService extends Disposable implements IOEShimService {
 
 	private sessionMap = new Map<number, string>();
 	private nodeHandleMap = new Map<number, string>();
+	private nodeInfoMap = new Map<ITreeItem, azdata.NodeInfo>();
 
 	constructor(
 		@IObjectExplorerService private oe: IObjectExplorerService,
@@ -140,8 +143,8 @@ export class OEShimService extends Disposable implements IOEShimService {
 
 	private treeNodeToITreeItem(viewId: string, node: TreeNode, parentNode: ITreeItem): ITreeItem {
 		let handle = generateUuid();
-		let nodePath = node.nodePath;
 		let icon: string = '';
+		let nodePath = node.nodePath;
 		if (node.iconType) {
 			icon = (typeof node.iconType === 'string') ? node.iconType : node.iconType.id;
 		} else {
@@ -154,6 +157,31 @@ export class OEShimService extends Disposable implements IOEShimService {
 			}
 		}
 		icon = icon.toLowerCase();
+		// Change the database if the node has a different database
+		// than its parent
+		let databaseChanged = false;
+		let updatedPayload: azdata.IConnectionProfile | any = {};
+		if (node.nodeTypeId === NodeType.Database) {
+			const database = node.getDatabaseName();
+			if (database) {
+				databaseChanged = true;
+				updatedPayload = Object.assign(updatedPayload, parentNode.payload);
+				updatedPayload.databaseName = node.getDatabaseName();
+			}
+		}
+		const nodeInfo: azdata.NodeInfo = {
+			nodePath: nodePath,
+			nodeType: node.nodeTypeId,
+			nodeSubType: node.nodeSubType,
+			nodeStatus: node.nodeStatus,
+			label: node.label,
+			isLeaf: node.isAlwaysLeaf,
+			metadata: node.metadata,
+			errorMessage: node.errorStateMessage,
+			iconType: icon,
+			childProvider: node.childProvider || parentNode.childProvider,
+			payload: node.payload || (databaseChanged ? updatedPayload : parentNode.payload)
+		};
 		let newTreeItem: ITreeItem = {
 			parentHandle: node.parent.id,
 			handle,
@@ -163,11 +191,12 @@ export class OEShimService extends Disposable implements IOEShimService {
 			},
 			childProvider: node.childProvider || parentNode.childProvider,
 			providerHandle: parentNode.childProvider,
-			payload: node.payload || parentNode.payload,
+			payload: node.payload || (databaseChanged ? updatedPayload : parentNode.payload),
 			contextValue: node.nodeTypeId,
 			sqlIcon: icon
 		};
 		this.nodeHandleMap.set(generateNodeMapKey(viewId, newTreeItem), nodePath);
+		this.nodeInfoMap.set(newTreeItem, nodeInfo);
 		return newTreeItem;
 	}
 
@@ -178,6 +207,13 @@ export class OEShimService extends Disposable implements IOEShimService {
 	public isNodeConnected(viewId: string, node: ITreeItem): boolean {
 		return this.sessionMap.has(generateSessionMapKey(viewId, node));
 	}
+
+	public getNodeInfoForTreeItem(treeItem: ITreeItem): azdata.NodeInfo {
+		if (this.nodeInfoMap.has(treeItem)) {
+			return this.nodeInfoMap.get(treeItem);
+		}
+		return undefined;
+	}
 }
 
 function generateSessionMapKey(viewId: string, node: ITreeItem): number {