bring Azure context menu to the new tree and enhance start cloud shell command (#13101)

* bring context menu to the new tree

* enhance start cloud shell command

* text

extensions/azurecore/package.json

@@ -282,6 +282,16 @@
           "command": "azure.resource.refresh",
           "when": "contextValue == azure.resource.itemType.account",
           "group": "navigation"
+        },
+        {
+          "command": "azure.resource.signin",
+          "when": "treeId == connectionDialog/azureResourceExplorer",
+          "group": "navigation"
+        },
+        {
+          "command": "azure.resource.refreshall",
+          "when": "treeId == connectionDialog/azureResourceExplorer",
+          "group": "navigation"
         }
       ]
     },

extensions/azurecore/src/azureResource/commands.ts

@@ -30,13 +30,34 @@ export function registerAzureResourceCommands(appContext: AppContext, trees: (Az
 				vscode.window.showInformationMessage(msg);
 				return;
 			}
-			if (!node || !(node instanceof AzureResourceAccountTreeNode)) {
+			let azureAccount: AzureAccount | undefined;
+			if (node instanceof AzureResourceAccountTreeNode) {
+				azureAccount = node.account as AzureAccount;
+			} else {
+				let accounts = await azdata.accounts.getAllAccounts();
+				accounts = accounts.filter(a => a.key.providerId.startsWith('azure'));
+				if (accounts.length === 0) {
+					const signin = localize('azure.signIn', "Sign in");
+					const action = await vscode.window.showErrorMessage(localize('azure.noAccountError', "You are not currently signed into any Azure accounts, Please sign in and then try again."),
+						signin);
+					if (action === signin) {
+						vscode.commands.executeCommand('azure.resource.signin');
+					}
+					return;
+				} else if (accounts.length === 1) {
+					azureAccount = accounts[0];
+				} else {
+					const pickedAccount = await vscode.window.showQuickPick(accounts.map(account => account.displayInfo.displayName), {
+						canPickMany: false,
+						placeHolder: localize('azure.pickAnAzureAccount', "Select an Azure account")
+					});
+					if (!pickedAccount) {
+						vscode.window.showErrorMessage(localize('azure.accountNotSelectedError', "You must select an Azure account for this feature to work."));
 						return;
 					}
-
+					azureAccount = accounts.find(acct => acct.displayInfo.displayName === pickedAccount);
-			const accountNode = node as AzureResourceAccountTreeNode;
+				}
-			const azureAccount = accountNode.account as AzureAccount;
+			}
-
 
 			const terminalService = appContext.getService<IAzureTerminalService>(AzureResourceServiceNames.terminalService);
 

src/sql/workbench/services/connection/browser/connectionBrowseTab.ts

@@ -47,7 +47,7 @@ import { IAction } from 'vs/base/common/actions';
 import { createAndFillInContextMenuActions } from 'vs/platform/actions/browser/menuEntryActionViewItem';
 import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 
-export type TreeElement = ConnectionProviderElement | ITreeItemFromProvider | SavedConnectionNode | ServerTreeElement;
+export type TreeElement = ConnectionDialogTreeProviderElement | ITreeItemFromProvider | SavedConnectionNode | ServerTreeElement;
 
 export class ConnectionBrowseTab implements IPanelTab {
 	public readonly title = localize('connectionDialog.browser', "Browse");
@@ -156,9 +156,17 @@ export class ConnectionBrowserView extends Disposable implements IPanelView {
 				accessibilityProvider: new ListAccessibilityProvider()
 			}) as WorkbenchAsyncDataTree<TreeModel, TreeElement>);
 		this._register(this.tree.onContextMenu(e => {
-			const context = e.element as ITreeItemFromProvider;
+			let context: ITreeItem | ConnectionDialogTreeProviderElement | undefined;
-			if (context?.element) {
+			let actionContext: TreeViewItemHandleArg | ConnectionDialogTreeProviderElement | undefined;
-				this.contextKey.set(context.element);
+			if (instanceOfITreeItemFromProvider(e.element)) {
+				context = e.element.element;
+				actionContext = <TreeViewItemHandleArg>{ $treeViewId: e.element.treeId, $treeItemHandle: context.handle, $treeItem: context };
+			} else if (e.element instanceof ConnectionDialogTreeProviderElement) {
+				context = e.element;
+				actionContext = e.element;
+			}
+			if (context) {
+				this.contextKey.set(context);
 				const menu = this.menuService.createMenu(MenuId.ConnectionDialogBrowseTreeContext, this.contextKeyService);
 				const primary: IAction[] = [];
 				const secondary: IAction[] = [];
@@ -168,7 +176,7 @@ export class ConnectionBrowserView extends Disposable implements IPanelView {
 				this.contextMenuService.showContextMenu({
 					getAnchor: () => e.anchor,
 					getActions: () => result.primary,
-					getActionsContext: () => (<TreeViewItemHandleArg>{ $treeViewId: context.treeId, $treeItemHandle: context.element.handle, $treeItem: context.element })
+					getActionsContext: () => actionContext
 				});
 			}
 		}));
@@ -239,7 +247,11 @@ export interface ITreeItemFromProvider {
 	getChildren?(): Promise<ITreeItemFromProvider[]>
 }
 
-class ConnectionProviderElement {
+export function instanceOfITreeItemFromProvider(obj: any): obj is ITreeItemFromProvider {
+	return !!(<ITreeItemFromProvider>obj)?.element;
+}
+
+class ConnectionDialogTreeProviderElement {
 	public readonly id = this.descriptor.id;
 	public readonly name = this.descriptor.name;
 
@@ -262,7 +274,7 @@ class ListDelegate implements IListVirtualDelegate<TreeElement> {
 	}
 
 	getTemplateId(element: TreeElement): string {
-		if (element instanceof ConnectionProviderElement) {
+		if (element instanceof ConnectionDialogTreeProviderElement) {
 			return ProviderElementRenderer.TEMPLATE_ID;
 		} else if (element instanceof ConnectionProfile) {
 			return ServerTreeRenderer.CONNECTION_TEMPLATE_ID;
@@ -283,7 +295,7 @@ interface ProviderElementTemplate {
 	readonly name: HTMLElement;
 }
 
-class ProviderElementRenderer implements ITreeRenderer<ConnectionProviderElement, void, ProviderElementTemplate> {
+class ProviderElementRenderer implements ITreeRenderer<ConnectionDialogTreeProviderElement, void, ProviderElementTemplate> {
 	public static readonly TEMPLATE_ID = 'ProviderElementTemplate';
 	public readonly templateId = ProviderElementRenderer.TEMPLATE_ID;
 
@@ -293,7 +305,7 @@ class ProviderElementRenderer implements ITreeRenderer<ConnectionProviderElement
 		return { name, icon };
 	}
 
-	renderElement(element: ITreeNode<ConnectionProviderElement, void>, index: number, templateData: ProviderElementTemplate, height: number): void {
+	renderElement(element: ITreeNode<ConnectionDialogTreeProviderElement, void>, index: number, templateData: ProviderElementTemplate, height: number): void {
 		templateData.name.innerText = element.element.name;
 	}
 
@@ -306,7 +318,7 @@ interface SavedConnectionNodeElementTemplate {
 	readonly name: HTMLElement;
 }
 
-class SavedConnectionsNodeRenderer implements ITreeRenderer<ConnectionProviderElement, void, SavedConnectionNodeElementTemplate> {
+class SavedConnectionsNodeRenderer implements ITreeRenderer<ConnectionDialogTreeProviderElement, void, SavedConnectionNodeElementTemplate> {
 	public static readonly TEMPLATE_ID = 'savedConnectionNode';
 	public readonly templateId = SavedConnectionsNodeRenderer.TEMPLATE_ID;
 
@@ -316,7 +328,7 @@ class SavedConnectionsNodeRenderer implements ITreeRenderer<ConnectionProviderEl
 		return { name, icon };
 	}
 
-	renderElement(element: ITreeNode<ConnectionProviderElement, void>, index: number, templateData: SavedConnectionNodeElementTemplate, height: number): void {
+	renderElement(element: ITreeNode<ConnectionDialogTreeProviderElement, void>, index: number, templateData: SavedConnectionNodeElementTemplate, height: number): void {
 		templateData.name.innerText = localize('savedConnections', "Saved Connections");
 	}
 
@@ -326,7 +338,7 @@ class SavedConnectionsNodeRenderer implements ITreeRenderer<ConnectionProviderEl
 
 class IdentityProvider implements IIdentityProvider<TreeElement> {
 	getId(element: TreeElement): string {
-		if (element instanceof ConnectionProviderElement) {
+		if (element instanceof ConnectionDialogTreeProviderElement) {
 			return element.id;
 		} else if (element instanceof ConnectionProfile) {
 			return element.id;
@@ -353,7 +365,7 @@ class TreeModel {
 	getChildren(): TreeElement[] {
 		this._savedConnectionNode = this.instantiationService.createInstance(SavedConnectionNode);
 		const descriptors = Array.from(this.connectionTreeService.descriptors);
-		return [this._savedConnectionNode, ...Iterable.map(this.connectionTreeService.providers, ([id, provider]) => new ConnectionProviderElement(provider, descriptors.find(i => i.id === id)))];
+		return [this._savedConnectionNode, ...Iterable.map(this.connectionTreeService.providers, ([id, provider]) => new ConnectionDialogTreeProviderElement(provider, descriptors.find(i => i.id === id)))];
 	}
 
 	public get savedConnectionNode(): SavedConnectionNode | undefined {
@@ -363,7 +375,7 @@ class TreeModel {
 
 class ListAccessibilityProvider implements IListAccessibilityProvider<TreeElement> {
 	getAriaLabel(element: TreeElement): string {
-		if (element instanceof ConnectionProviderElement) {
+		if (element instanceof ConnectionDialogTreeProviderElement) {
 			return element.name;
 		} else if (element instanceof ConnectionProfile) {
 			return element.serverName;
@@ -394,7 +406,7 @@ class DataSource implements IAsyncDataSource<TreeModel, TreeElement> {
 	hasChildren(element: TreeModel | TreeElement): boolean {
 		if (element instanceof TreeModel) {
 			return true;
-		} else if (element instanceof ConnectionProviderElement) {
+		} else if (element instanceof ConnectionDialogTreeProviderElement) {
 			return true;
 		} else if (element instanceof ConnectionProfile) {
 			return false;
@@ -413,12 +425,7 @@ class DataSource implements IAsyncDataSource<TreeModel, TreeElement> {
 
 	public get expandableTreeNodes(): TreeElement[] {
 		return this.treeNodes.filter(node => {
-			return node instanceof TreeModel
+			return instanceOfITreeItemFromProvider(node) && node.element.collapsibleState !== TreeItemCollapsibleState.None;
-				|| node instanceof SavedConnectionNode
-				|| node instanceof ConnectionProfileGroup
-				|| node instanceof TreeNode
-				|| node instanceof ConnectionProviderElement
-				|| (!(node instanceof ConnectionProfile) && node.element.collapsibleState !== TreeItemCollapsibleState.None);
 		});
 	}
 
@@ -436,7 +443,7 @@ class DataSource implements IAsyncDataSource<TreeModel, TreeElement> {
 				} else if (
 					!(element instanceof TreeModel) &&
 					!(element instanceof TreeNode) &&
-					!(element instanceof ConnectionProviderElement)
+					!(element instanceof ConnectionDialogTreeProviderElement)
 				) {
 					children = (children as ITreeItemFromProvider[]).filter(item => {
 						return item.element.collapsibleState !== TreeItemCollapsibleState.None || this._filterRegex.test(item.element.label.label);
@@ -599,26 +606,36 @@ class TreeItemRenderer extends Disposable implements ITreeRenderer<ITreeItemFrom
 	}
 }
 
-class ContextKey extends Disposable implements IContextKey<ITreeItem> {
+type ContextValueType = ITreeItem | ConnectionDialogTreeProviderElement;
-	static readonly ContextValue = new RawContextKey<string>('contextValue', undefined);
+
-	static readonly Item = new RawContextKey<ITreeItem>('item', undefined);
+class ContextKey extends Disposable implements IContextKey<ContextValueType> {
-	private _contextValueKey: IContextKey<string>;
+	static readonly ContextValue = new RawContextKey<string | undefined>('contextValue', undefined);
-	private _itemKey: IContextKey<ITreeItem>;
+	static readonly TreeId = new RawContextKey<string | undefined>('treeId', undefined);
+	private _contextValueKey: IContextKey<string | undefined>;
+	private _treeIdKey: IContextKey<string | undefined>;
+	private _item: ContextValueType;
 
 	constructor(
 		@IContextKeyService contextKeyService: IContextKeyService
 	) {
 		super();
 		this._contextValueKey = ContextKey.ContextValue.bindTo(contextKeyService);
-		this._itemKey = ContextKey.Item.bindTo(contextKeyService);
+		this._treeIdKey = ContextKey.TreeId.bindTo(contextKeyService);
 	}
-	set(value: ITreeItem): void {
+	set(value: ContextValueType): void {
+		this.reset();
+		this._item = value;
+		if (value instanceof ConnectionDialogTreeProviderElement) {
+			this._treeIdKey.set(value.id);
+		} else {
 			this._contextValueKey.set(value.contextValue);
 		}
+	}
 	reset(): void {
 		this._contextValueKey.reset();
+		this._treeIdKey.reset();
 	}
-	get(): ITreeItem {
+	get(): ContextValueType {
-		return this._itemKey.get();
+		return this._item;
 	}
 }