enable filtering, account node context menu and introduce flat account tree node (#13066)

* add search box

* switch back to the traditional azure tree

* Revert "switch back to the traditional azure tree"

This reverts commit 7904b9cd599591e94412ec79da23590068de46b6.

* flat account tree node and filtering

* add comment

* context menu

* fix test

* handle disposable

* add logging

extensions/azurecore/package.json

@@ -37,7 +37,7 @@
             "description": "%azure.resource.config.filter.description%"
           },
           "azure.tenant.config.filter": {
-            "type":"array",
+            "type": "array",
             "default": [],
             "description": "%azure.tenant.config.filter.description%"
           }
@@ -271,6 +271,18 @@
           "when": "viewItem == azure.resource.itemType.account",
           "group": "azurecore"
         }
+      ],
+      "connectionDialog/browseTree": [
+        {
+          "command": "azure.resource.selectsubscriptions",
+          "when": "contextValue == azure.resource.itemType.account",
+          "group": "navigation"
+        },
+        {
+          "command": "azure.resource.refresh",
+          "when": "contextValue == azure.resource.itemType.account",
+          "group": "navigation"
+        }
       ]
     },
     "hasAzureResourceProviders": true

extensions/azurecore/src/azureResource/commands.ts

@@ -18,8 +18,10 @@ import { AzureResourceAccountTreeNode } from './tree/accountTreeNode';
 import { IAzureResourceSubscriptionService, IAzureResourceSubscriptionFilterService, IAzureTerminalService } from '../azureResource/interfaces';
 import { AzureResourceServiceNames } from './constants';
 import { AzureAccount, Tenant } from '../account-provider/interfaces';
+import { FlatAccountTreeNode } from './tree/flatAccountTreeNode';
+import { ConnectionDialogTreeProvider } from './tree/connectionDialogTreeProvider';
 
-export function registerAzureResourceCommands(appContext: AppContext, tree: AzureResourceTreeProvider): void {
+export function registerAzureResourceCommands(appContext: AppContext, trees: (AzureResourceTreeProvider | ConnectionDialogTreeProvider)[]): void {
 	vscode.commands.registerCommand('azure.resource.startterminal', async (node?: TreeNode) => {
 		try {
 			const enablePreviewFeatures = vscode.workspace.getConfiguration('workbench').get('enablePreviewFeatures');
@@ -73,7 +75,7 @@ export function registerAzureResourceCommands(appContext: AppContext, tree: Azur
 	// Resource Tree commands
 
 	vscode.commands.registerCommand('azure.resource.selectsubscriptions', async (node?: TreeNode) => {
-		if (!(node instanceof AzureResourceAccountTreeNode)) {
+		if (!(node instanceof AzureResourceAccountTreeNode) && !(node instanceof FlatAccountTreeNode)) {
 			return;
 		}
 
@@ -85,9 +87,7 @@ export function registerAzureResourceCommands(appContext: AppContext, tree: Azur
 		const subscriptionService = appContext.getService<IAzureResourceSubscriptionService>(AzureResourceServiceNames.subscriptionService);
 		const subscriptionFilterService = appContext.getService<IAzureResourceSubscriptionFilterService>(AzureResourceServiceNames.subscriptionFilterService);
 
-		const accountNode = node as AzureResourceAccountTreeNode;
-
-		const subscriptions = (await accountNode.getCachedSubscriptions()) || <azureResource.AzureResourceSubscription[]>[];
+		const subscriptions = [];
 		if (subscriptions.length === 0) {
 			try {
 
@@ -105,7 +105,7 @@ export function registerAzureResourceCommands(appContext: AppContext, tree: Azur
 			}
 		}
 
-		let selectedSubscriptions = await subscriptionFilterService.getSelectedSubscriptions(accountNode.account);
+		let selectedSubscriptions = await subscriptionFilterService.getSelectedSubscriptions(account);
 		if (!selectedSubscriptions) {
 			selectedSubscriptions = [];
 		}
@@ -132,17 +132,25 @@ export function registerAzureResourceCommands(appContext: AppContext, tree: Azur
 
 		const selectedSubscriptionQuickPickItems = await vscode.window.showQuickPick(subscriptionQuickPickItems, { canPickMany: true });
 		if (selectedSubscriptionQuickPickItems && selectedSubscriptionQuickPickItems.length > 0) {
-			await tree.refresh(node, false);
+			for (const tree of trees) {
+				await tree.refresh(undefined, false);
+			}
 
 			selectedSubscriptions = selectedSubscriptionQuickPickItems.map((subscriptionItem) => subscriptionItem.subscription);
-			await subscriptionFilterService.saveSelectedSubscriptions(accountNode.account, selectedSubscriptions);
+			await subscriptionFilterService.saveSelectedSubscriptions(account, selectedSubscriptions);
 		}
 	});
 
-	vscode.commands.registerCommand('azure.resource.refreshall', () => tree.notifyNodeChanged(undefined));
+	vscode.commands.registerCommand('azure.resource.refreshall', () => {
+		for (const tree of trees) {
+			tree.notifyNodeChanged(undefined);
+		}
+	});
 
 	vscode.commands.registerCommand('azure.resource.refresh', async (node?: TreeNode) => {
-		await tree.refresh(node, true);
+		for (const tree of trees) {
+			await tree.refresh(node, true);
+		}
 	});
 
 	vscode.commands.registerCommand('azure.resource.signin', async (node?: TreeNode) => {

extensions/azurecore/src/azureResource/tree/connectionDialogTreeProvider.ts

@@ -0,0 +1,112 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+import * as azdata from 'azdata';
+import { AppContext } from '../../appContext';
+import * as nls from 'vscode-nls';
+const localize = nls.loadMessageBundle();
+
+import { TreeNode } from '../treeNode';
+import { AzureResourceAccountNotSignedInTreeNode } from './accountNotSignedInTreeNode';
+import { AzureResourceMessageTreeNode } from '../messageTreeNode';
+import { AzureResourceContainerTreeNodeBase } from './baseTreeNodes';
+import { AzureResourceErrorMessageUtil, equals } from '../utils';
+import { IAzureResourceTreeChangeHandler } from './treeChangeHandler';
+import { FlatAccountTreeNode } from './flatAccountTreeNode';
+import { Logger } from '../../utils/Logger';
+
+export class ConnectionDialogTreeProvider implements vscode.TreeDataProvider<TreeNode>, IAzureResourceTreeChangeHandler {
+	public isSystemInitialized: boolean = false;
+
+	private accounts: azdata.Account[];
+	private _onDidChangeTreeData = new vscode.EventEmitter<TreeNode>();
+	private loadingAccountsPromise: Promise<void>;
+
+	public constructor(private readonly appContext: AppContext) {
+		azdata.accounts.onDidChangeAccounts(async (e: azdata.DidChangeAccountsParams) => {
+			// This event sends it per provider, we need to make sure we get all the azure related accounts
+			let accounts = await azdata.accounts.getAllAccounts();
+			accounts = accounts.filter(a => a.key.providerId.startsWith('azure'));
+			// the onDidChangeAccounts event will trigger in many cases where the accounts didn't actually change
+			// the notifyNodeChanged event triggers a refresh which triggers a getChildren which can trigger this callback
+			// this below check short-circuits the infinite callback loop
+			this.setSystemInitialized();
+			if (!equals(accounts, this.accounts)) {
+				this.accounts = accounts;
+				this.notifyNodeChanged(undefined);
+			}
+		});
+	}
+
+	public async getChildren(element?: TreeNode): Promise<TreeNode[]> {
+		if (element) {
+			return element.getChildren(true);
+		}
+
+		if (!this.isSystemInitialized) {
+			if (!this.loadingAccountsPromise) {
+				this.loadingAccountsPromise = this.loadAccounts();
+			}
+			return [AzureResourceMessageTreeNode.create(localize('azure.resource.tree.treeProvider.loadingLabel', "Loading ..."), undefined)];
+		}
+
+		try {
+			if (this.accounts && this.accounts.length > 0) {
+				const accountNodes: FlatAccountTreeNode[] = [];
+				for (const account of this.accounts) {
+					const accountNode = new FlatAccountTreeNode(account, this.appContext, this);
+					await accountNode.updateLabel();
+					accountNodes.push(accountNode);
+				}
+				return accountNodes;
+			} else {
+				return [new AzureResourceAccountNotSignedInTreeNode()];
+			}
+		} catch (error) {
+			return [AzureResourceMessageTreeNode.create(AzureResourceErrorMessageUtil.getErrorMessage(error), undefined)];
+		}
+	}
+
+	private async loadAccounts(): Promise<void> {
+		try {
+			this.accounts = await azdata.accounts.getAllAccounts();
+			// System has been initialized
+			this.setSystemInitialized();
+			this._onDidChangeTreeData.fire(undefined);
+		} catch (err) {
+			// Skip for now, we can assume that the accounts changed event will eventually notify instead
+			Logger.error('loadAccounts failed with the following error: {0}', err.message ?? err);
+			this.isSystemInitialized = false;
+		}
+	}
+
+	private setSystemInitialized(): void {
+		this.isSystemInitialized = true;
+		this.loadingAccountsPromise = undefined;
+	}
+
+	public get onDidChangeTreeData(): vscode.Event<TreeNode> {
+		return this._onDidChangeTreeData.event;
+	}
+
+	public notifyNodeChanged(node: TreeNode): void {
+		this._onDidChangeTreeData.fire(node);
+	}
+
+	public async refresh(node: TreeNode, isClearingCache: boolean): Promise<void> {
+		if (isClearingCache) {
+			if ((node instanceof AzureResourceContainerTreeNodeBase)) {
+				node.clearCache();
+			}
+		}
+
+		this._onDidChangeTreeData.fire(node);
+	}
+
+	public getTreeItem(element: TreeNode): vscode.TreeItem | Thenable<vscode.TreeItem> {
+		return element.getTreeItem();
+	}
+}

extensions/azurecore/src/azureResource/tree/flatAccountTreeNode.ts

@@ -0,0 +1,173 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+import * as azdata from 'azdata';
+import { TokenCredentials } from '@azure/ms-rest-js';
+
+import * as nls from 'vscode-nls';
+const localize = nls.loadMessageBundle();
+
+import { AppContext } from '../../appContext';
+import { azureResource } from 'azureResource';
+import { TreeNode } from '../treeNode';
+import { AzureResourceCredentialError } from '../errors';
+import { AzureResourceContainerTreeNodeBase } from './baseTreeNodes';
+import { AzureResourceItemType, AzureResourceServiceNames } from '../constants';
+import { AzureResourceMessageTreeNode } from '../messageTreeNode';
+import { IAzureResourceTreeChangeHandler } from './treeChangeHandler';
+import { IAzureResourceSubscriptionService, IAzureResourceSubscriptionFilterService, IAzureResourceNodeWithProviderId } from '../../azureResource/interfaces';
+import { AzureAccount } from '../../account-provider/interfaces';
+import { AzureResourceService } from '../resourceService';
+import { AzureResourceResourceTreeNode } from '../resourceTreeNode';
+import { AzureResourceErrorMessageUtil } from '../utils';
+
+export class FlatAccountTreeNode extends AzureResourceContainerTreeNodeBase {
+	public constructor(
+		public readonly account: AzureAccount,
+		appContext: AppContext,
+		treeChangeHandler: IAzureResourceTreeChangeHandler
+	) {
+		super(appContext, treeChangeHandler, undefined);
+
+		this._subscriptionService = this.appContext.getService<IAzureResourceSubscriptionService>(AzureResourceServiceNames.subscriptionService);
+		this._subscriptionFilterService = this.appContext.getService<IAzureResourceSubscriptionFilterService>(AzureResourceServiceNames.subscriptionFilterService);
+		this._resourceService = this.appContext.getService<AzureResourceService>(AzureResourceServiceNames.resourceService);
+
+		this._id = `account_${this.account.key.accountId}`;
+		this.setCacheKey(`${this._id}.dataresources`);
+		this._label = account.displayInfo.displayName;
+	}
+
+	public async updateLabel(): Promise<void> {
+		const subscriptionInfo = await this.getSubscriptionInfo();
+		if (subscriptionInfo.total !== 0) {
+			this._label = localize({
+				key: 'azure.resource.tree.accountTreeNode.title',
+				comment: [
+					'{0} is the display name of the azure account',
+					'{1} is the number of selected subscriptions in this account',
+					'{2} is the number of total subscriptions in this account'
+				]
+			}, "{0} ({1}/{2} subscriptions)", this.account.displayInfo.displayName, subscriptionInfo.selected, subscriptionInfo.total);
+		} else {
+			this._label = this.account.displayInfo.displayName;
+		}
+	}
+
+	private async getSubscriptionInfo(): Promise<{
+		subscriptions: azureResource.AzureResourceSubscription[],
+		total: number,
+		selected: number
+	}> {
+		let subscriptions: azureResource.AzureResourceSubscription[] = [];
+		try {
+			for (const tenant of this.account.properties.tenants) {
+				const token = await azdata.accounts.getAccountSecurityToken(this.account, tenant.id, azdata.AzureResource.ResourceManagement);
+
+				subscriptions.push(...(await this._subscriptionService.getSubscriptions(this.account, new TokenCredentials(token.token, token.tokenType), tenant.id) || <azureResource.AzureResourceSubscription[]>[]));
+			}
+		} catch (error) {
+			throw new AzureResourceCredentialError(localize('azure.resource.tree.accountTreeNode.credentialError', "Failed to get credential for account {0}. Please refresh the account.", this.account.key.accountId), error);
+		}
+		const total = subscriptions.length;
+		let selected = total;
+
+		const selectedSubscriptions = await this._subscriptionFilterService.getSelectedSubscriptions(this.account);
+		const selectedSubscriptionIds = (selectedSubscriptions || <azureResource.AzureResourceSubscription[]>[]).map((subscription) => subscription.id);
+		if (selectedSubscriptionIds.length > 0) {
+			subscriptions = subscriptions.filter((subscription) => selectedSubscriptionIds.indexOf(subscription.id) !== -1);
+			selected = selectedSubscriptionIds.length;
+		}
+		return {
+			subscriptions,
+			total,
+			selected
+		};
+	}
+
+	public async getChildren(): Promise<TreeNode[]> {
+		try {
+			let dataResources: IAzureResourceNodeWithProviderId[] = [];
+			if (this._isClearingCache) {
+				let subscriptions: azureResource.AzureResourceSubscription[] = (await this.getSubscriptionInfo()).subscriptions;
+
+				if (subscriptions.length === 0) {
+					return [AzureResourceMessageTreeNode.create(FlatAccountTreeNode.noSubscriptionsLabel, this)];
+				} else {
+					// Filter out everything that we can't authenticate to.
+					subscriptions = subscriptions.filter(async s => {
+						const token = await azdata.accounts.getAccountSecurityToken(this.account, s.tenant, azdata.AzureResource.ResourceManagement);
+						if (!token) {
+							console.info(`Account does not have permissions to view subscription ${JSON.stringify(s)}.`);
+							return false;
+						}
+						return true;
+					});
+				}
+
+				const resourceProviderIds = await this._resourceService.listResourceProviderIds();
+				for (const subscription of subscriptions) {
+					for (const providerId of resourceProviderIds) {
+						const resourceTypes = await this._resourceService.getRootChildren(providerId, this.account, subscription, subscription.tenant);
+						for (const resourceType of resourceTypes) {
+							dataResources.push(...await this._resourceService.getChildren(providerId, resourceType.resourceNode));
+						}
+					}
+				}
+				dataResources = dataResources.sort((a, b) => { return a.resourceNode.treeItem.label.localeCompare(b.resourceNode.treeItem.label); });
+				this.updateCache(dataResources);
+				this._isClearingCache = false;
+			} else {
+				dataResources = this.getCache<IAzureResourceNodeWithProviderId[]>();
+			}
+
+			return dataResources.map(dr => new AzureResourceResourceTreeNode(dr, this, this.appContext));
+		} catch (error) {
+			if (error instanceof AzureResourceCredentialError) {
+				vscode.commands.executeCommand('azure.resource.signin');
+			}
+			return [AzureResourceMessageTreeNode.create(AzureResourceErrorMessageUtil.getErrorMessage(error), this)];
+		}
+	}
+
+	public getTreeItem(): vscode.TreeItem | Promise<vscode.TreeItem> {
+		const item = new vscode.TreeItem(this._label, vscode.TreeItemCollapsibleState.Collapsed);
+		item.id = this._id;
+		item.contextValue = AzureResourceItemType.account;
+		item.iconPath = {
+			dark: this.appContext.extensionContext.asAbsolutePath('resources/dark/account_inverse.svg'),
+			light: this.appContext.extensionContext.asAbsolutePath('resources/light/account.svg')
+		};
+		return item;
+	}
+
+	public getNodeInfo(): azdata.NodeInfo {
+		return {
+			label: this._label,
+			isLeaf: false,
+			errorMessage: undefined,
+			metadata: undefined,
+			nodePath: this.generateNodePath(),
+			nodeStatus: undefined,
+			nodeType: AzureResourceItemType.account,
+			nodeSubType: undefined,
+			iconType: AzureResourceItemType.account
+		};
+	}
+
+	public get nodePathValue(): string {
+		return this._id;
+	}
+
+	private _subscriptionService: IAzureResourceSubscriptionService = undefined;
+	private _subscriptionFilterService: IAzureResourceSubscriptionFilterService = undefined;
+	private _resourceService: AzureResourceService = undefined;
+
+	private _id: string = undefined;
+	private _label: string = undefined;
+
+	private static readonly noSubscriptionsLabel = localize('azure.resource.tree.accountTreeNode.noSubscriptionsLabel', "No Subscriptions found.");
+}

extensions/azurecore/src/extension.ts

@@ -43,7 +43,7 @@ import * as loc from './localizedConstants';
 import * as constants from './constants';
 import { AzureResourceGroupService } from './azureResource/providers/resourceGroup/resourceGroupService';
 import { Logger } from './utils/Logger';
-import { FlatAzureResourceTreeProvider } from './azureResource/tree/flatTreeProvider';
+import { ConnectionDialogTreeProvider } from './azureResource/tree/connectionDialogTreeProvider';
 import { AzureDataGridProvider } from './azureDataGridProvider';
 
 let extensionContext: vscode.ExtensionContext;
@@ -85,11 +85,11 @@ export async function activate(context: vscode.ExtensionContext): Promise<azurec
 
 	registerAzureServices(appContext);
 	const azureResourceTree = new AzureResourceTreeProvider(appContext);
-	const flatAzureResourceTree = new FlatAzureResourceTreeProvider(appContext);
-	pushDisposable(vscode.window.registerTreeDataProvider('connectionDialog/azureResourceExplorer', flatAzureResourceTree));
+	const connectionDialogTree = new ConnectionDialogTreeProvider(appContext);
+	pushDisposable(vscode.window.registerTreeDataProvider('connectionDialog/azureResourceExplorer', connectionDialogTree));
 	pushDisposable(vscode.window.registerTreeDataProvider('azureResourceExplorer', azureResourceTree));
 	pushDisposable(vscode.workspace.onDidChangeConfiguration(e => onDidChangeConfiguration(e), this));
-	registerAzureResourceCommands(appContext, azureResourceTree);
+	registerAzureResourceCommands(appContext, [azureResourceTree, connectionDialogTree]);
 	azdata.dataprotocol.registerDataGridProvider(new AzureDataGridProvider(appContext));
 
 	return {

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

@@ -24,7 +24,7 @@ import { ILabelService } from 'vs/platform/label/common/label';
 import { WorkbenchAsyncDataTree } from 'vs/platform/list/browser/listService';
 import { FileThemeIcon, FolderThemeIcon, IThemeService, ThemeIcon } from 'vs/platform/theme/common/themeService';
 import { IResourceLabel, ResourceLabels } from 'vs/workbench/browser/labels';
-import { ITreeItemLabel, ITreeViewDataProvider, TreeItemCollapsibleState } from 'vs/workbench/common/views';
+import { ITreeItemLabel, ITreeViewDataProvider, TreeItemCollapsibleState, TreeViewItemHandleArg } from 'vs/workbench/common/views';
 import { Emitter, Event } from 'vs/base/common/event';
 import { AsyncRecentConnectionTreeDataSource } from 'sql/workbench/services/objectExplorer/browser/asyncRecentConnectionTreeDataSource';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
@@ -36,6 +36,15 @@ import { TreeNode } from 'sql/workbench/services/objectExplorer/common/treeNode'
 import { ServerTreeRenderer } from 'sql/workbench/services/objectExplorer/browser/serverTreeRenderer';
 import { ConnectionProfileGroupRenderer, ConnectionProfileRenderer, TreeNodeRenderer } from 'sql/workbench/services/objectExplorer/browser/asyncServerTreeRenderer';
 import { ColorScheme } from 'vs/platform/theme/common/theme';
+import { InputBox } from 'sql/base/browser/ui/inputBox/inputBox';
+import { IContextMenuService, IContextViewService } from 'vs/platform/contextview/browser/contextView';
+import { attachInputBoxStyler } from 'sql/platform/theme/common/styler';
+import { debounce } from 'vs/base/common/decorators';
+import { ICommandService } from 'vs/platform/commands/common/commands';
+import { IMenuService, MenuId } from 'vs/platform/actions/common/actions';
+import { IContextKey, IContextKeyService, RawContextKey } from 'vs/platform/contextkey/common/contextkey';
+import { IAction } from 'vs/base/common/actions';
+import { createAndFillInContextMenuActions } from 'vs/platform/actions/browser/menuEntryActionViewItem';
 
 export type TreeElement = ConnectionProviderElement | ITreeItemFromProvider | SavedConnectionNode | ServerTreeElement;
 
@@ -48,8 +57,13 @@ export class ConnectionBrowseTab implements IPanelTab {
 
 export class ConnectionBrowserView extends Disposable implements IPanelView {
 	private tree: WorkbenchAsyncDataTree<TreeModel, TreeElement> | undefined;
+	private filterInput: InputBox | undefined;
+	private treeContainer: HTMLElement | undefined;
 	private model: TreeModel | undefined;
 	private treeLabels: ResourceLabels | undefined;
+	private treeDataSource: DataSource | undefined;
+	private readonly contextKey = new ContextKey(this.contextKeyService);
+
 	public onDidChangeVisibility = Event.None;
 
 	private readonly _onSelect = this._register(new Emitter<ITreeMouseEvent<TreeElement>>());
@@ -60,14 +74,58 @@ export class ConnectionBrowserView extends Disposable implements IPanelView {
 
 	constructor(
 		@IInstantiationService private readonly instantiationService: IInstantiationService,
-		@IConnectionTreeService private readonly connectionTreeService: IConnectionTreeService
+		@IConnectionTreeService private readonly connectionTreeService: IConnectionTreeService,
+		@IContextViewService private readonly contextViewService: IContextViewService,
+		@IThemeService private readonly themeService: IThemeService,
+		@ICommandService private readonly commandService: ICommandService,
+		@IMenuService private readonly menuService: IMenuService,
+		@IContextKeyService private readonly contextKeyService: IContextKeyService,
+		@IContextMenuService private readonly contextMenuService: IContextMenuService
 	) {
 		super();
 		this.connectionTreeService.setView(this);
 	}
 
 	render(container: HTMLElement): void {
+		this.renderFilterBox(container);
+		this.renderTree(container);
+	}
 
+	renderFilterBox(container: HTMLElement): void {
+		this.filterInput = new InputBox(container, this.contextViewService, {
+			placeholder: localize('connectionDialog.FilterPlaceHolder', "Type here to filter the list"),
+			ariaLabel: localize('connectionDialog.FilterInputTitle', "Filter connections")
+		});
+		this.filterInput.element.style.margin = '5px';
+		this._register(this.filterInput);
+		this._register(attachInputBoxStyler(this.filterInput, this.themeService));
+		this._register(this.filterInput.onDidChange(async () => {
+			await this.applyFilter();
+		}));
+	}
+
+	@debounce(500)
+	async applyFilter(): Promise<void> {
+		this.treeDataSource.setFilter(this.filterInput.value);
+		await this.refresh();
+		await this.expandAll();
+	}
+
+	async expandAll(): Promise<void> {
+		const expandedTreeItems: TreeElement[] = [];
+		let treeItemsToExpand: TreeElement[] = this.treeDataSource.expandableTreeNodes;
+		// expand the nodes one by one here to avoid the possible azure api traffic throttling.
+		while (treeItemsToExpand.length !== 0) {
+			for (const treeItem of treeItemsToExpand) {
+				await this.tree.expand(treeItem);
+			}
+			expandedTreeItems.push(...treeItemsToExpand);
+			treeItemsToExpand = this.treeDataSource.expandableTreeNodes.filter(el => expandedTreeItems.indexOf(el) === -1);
+		}
+	}
+
+	renderTree(container: HTMLElement): void {
+		this.treeContainer = container.appendChild(DOM.$('div'));
 		this.treeLabels = this._register(this.instantiationService.createInstance(ResourceLabels, this));
 		const renderers: ITreeRenderer<TreeElement, any, any>[] = [
 			new ProviderElementRenderer(),
@@ -79,14 +137,14 @@ export class ConnectionBrowserView extends Disposable implements IPanelView {
 		];
 
 		this.model = this.instantiationService.createInstance(TreeModel);
-
+		this.treeDataSource = new DataSource();
 		this.tree = this._register(this.instantiationService.createInstance(
 			WorkbenchAsyncDataTree,
 			'Browser Connections',
-			container,
+			this.treeContainer,
 			new ListDelegate(),
 			renderers,
-			new DataSource(),
+			this.treeDataSource,
 			{
 				identityProvider: new IdentityProvider(),
 				horizontalScrolling: false,
@@ -94,9 +152,37 @@ export class ConnectionBrowserView extends Disposable implements IPanelView {
 				transformOptimization: false,
 				accessibilityProvider: new ListAccessibilityProvider()
 			}) as WorkbenchAsyncDataTree<TreeModel, TreeElement>);
+		this._register(this.tree.onContextMenu(e => {
+			const context = e.element as ITreeItemFromProvider;
+			if (context?.element) {
+				this.contextKey.set(context.element);
+				const menu = this.menuService.createMenu(MenuId.ConnectionDialogBrowseTreeContext, this.contextKeyService);
+				const primary: IAction[] = [];
+				const secondary: IAction[] = [];
+				const result = { primary, secondary };
+				createAndFillInContextMenuActions(menu, { shouldForwardArgs: true }, result, this.contextMenuService);
 
-		this.tree.onMouseDblClick(e => this._onDblClick.fire(e));
-		this.tree.onMouseClick(e => this._onSelect.fire(e));
+				this.contextMenuService.showContextMenu({
+					getAnchor: () => e.anchor,
+					getActions: () => result.primary,
+					getActionsContext: () => (<TreeViewItemHandleArg>{ $treeViewId: context.treeId, $treeItemHandle: context.element.handle, $treeItem: context.element })
+				});
+			}
+		}));
+		this._register(this.tree.onMouseDblClick(e => this._onDblClick.fire(e)));
+		this._register(this.tree.onMouseClick(e => this._onSelect.fire(e)));
+		this._register(this.tree.onDidOpen((e) => {
+			if (!e.browserEvent) {
+				return;
+			}
+			const selection = this.tree.getSelection();
+			if (selection.length === 1) {
+				const selectedNode = selection[0];
+				if ('element' in selectedNode && selectedNode.element.command) {
+					this.commandService.executeCommand(selectedNode.element.command.id, ...(selectedNode.element.command.arguments || []));
+				}
+			}
+		}));
 
 		this.tree.setInput(this.model);
 
@@ -105,18 +191,15 @@ export class ConnectionBrowserView extends Disposable implements IPanelView {
 
 	async refresh(items?: ITreeItem[]): Promise<void> {
 		if (this.tree) {
-			if (items) {
-				for (const item of items) {
-					await this.tree.updateChildren({ element: item });
-				}
-			} else {
-				return this.tree.updateChildren();
-			}
+			return this.tree.updateChildren();
 		}
 	}
 
 	layout(dimension: DOM.Dimension): void {
-		this.tree.layout(dimension.height, dimension.width);
+		const treeHeight = dimension.height - DOM.getTotalHeight(this.filterInput.element);
+		this.treeContainer.style.width = `${dimension.width}px`;
+		this.treeContainer.style.height = `${treeHeight}px`;
+		this.tree.layout(treeHeight, dimension.width);
 	}
 
 	focus(): void {
@@ -126,6 +209,7 @@ export class ConnectionBrowserView extends Disposable implements IPanelView {
 
 export interface ITreeItemFromProvider {
 	readonly element: ITreeItem;
+	readonly treeId?: string;
 	getChildren?(): Promise<ITreeItemFromProvider[]>
 }
 
@@ -140,6 +224,7 @@ class ConnectionProviderElement {
 		const children = await this.provider.getChildren(element);
 		return children.map(v => ({
 			element: v,
+			treeId: this.descriptor.id,
 			getChildren: () => this.getChildren(v)
 		}));
 	}
@@ -267,6 +352,13 @@ class ListAccessibilityProvider implements IListAccessibilityProvider<TreeElemen
 }
 
 class DataSource implements IAsyncDataSource<TreeModel, TreeElement> {
+	private _filter: string | undefined;
+	private _filterRegex: RegExp | undefined;
+	public setFilter(filter: string): void {
+		this._filter = filter;
+		this._filterRegex = new RegExp(filter, 'i');
+	}
+
 	hasChildren(element: TreeModel | TreeElement): boolean {
 		if (element instanceof TreeModel) {
 			return true;
@@ -285,9 +377,42 @@ class DataSource implements IAsyncDataSource<TreeModel, TreeElement> {
 		}
 	}
 
-	getChildren(element: TreeModel | TreeElement): Iterable<TreeElement> | Promise<Iterable<TreeElement>> {
+	public treeNodes: TreeElement[] = [];
+
+	public get expandableTreeNodes(): TreeElement[] {
+		return this.treeNodes.filter(node => {
+			return node instanceof TreeModel
+				|| node instanceof SavedConnectionNode
+				|| node instanceof ConnectionProfileGroup
+				|| node instanceof TreeNode
+				|| node instanceof ConnectionProviderElement
+				|| (!(node instanceof ConnectionProfile) && node.element.collapsibleState !== TreeItemCollapsibleState.None);
+		});
+	}
+
+	async getChildren(element: TreeModel | TreeElement): Promise<Iterable<TreeElement>> {
+		if (element instanceof TreeModel) {
+			this.treeNodes = [];
+		}
 		if (!(element instanceof ConnectionProfile)) {
-			return element.getChildren();
+			let children = await element.getChildren();
+			if (this._filter) {
+				if ((element instanceof SavedConnectionNode) || (element instanceof ConnectionProfileGroup)) {
+					children = (children as (ConnectionProfile | ConnectionProfileGroup)[]).filter(item => {
+						return (item instanceof ConnectionProfileGroup) || this._filterRegex.test(item.title);
+					});
+				} else if (
+					!(element instanceof TreeModel) &&
+					!(element instanceof TreeNode) &&
+					!(element instanceof ConnectionProviderElement)
+				) {
+					children = (children as ITreeItemFromProvider[]).filter(item => {
+						return item.element.collapsibleState !== TreeItemCollapsibleState.None || this._filterRegex.test(item.element.label.label);
+					});
+				}
+			}
+			this.treeNodes.push(...children);
+			return children;
 		}
 		return [];
 	}
@@ -441,3 +566,27 @@ class TreeItemRenderer extends Disposable implements ITreeRenderer<ITreeItemFrom
 		templateData.elementDisposable.dispose();
 	}
 }
+
+class ContextKey extends Disposable implements IContextKey<ITreeItem> {
+	static readonly ContextValue = new RawContextKey<string>('contextValue', undefined);
+	static readonly Item = new RawContextKey<ITreeItem>('item', undefined);
+	private _contextValueKey: IContextKey<string>;
+	private _itemKey: IContextKey<ITreeItem>;
+
+	constructor(
+		@IContextKeyService contextKeyService: IContextKeyService
+	) {
+		super();
+		this._contextValueKey = ContextKey.ContextValue.bindTo(contextKeyService);
+		this._itemKey = ContextKey.Item.bindTo(contextKeyService);
+	}
+	set(value: ITreeItem): void {
+		this._contextValueKey.set(value.contextValue);
+	}
+	reset(): void {
+		this._contextValueKey.reset();
+	}
+	get(): ITreeItem {
+		return this._itemKey.get();
+	}
+}

src/sql/workbench/services/connection/test/browser/connectionDialogWidget.test.ts

@@ -28,6 +28,7 @@ import { Registry } from 'vs/platform/registry/common/platform';
 import { SyncDescriptor } from 'vs/platform/instantiation/common/descriptors';
 import { TestTreeView } from 'sql/workbench/services/connection/test/browser/testTreeView';
 import { ConnectionTreeService, IConnectionTreeService } from 'sql/workbench/services/connection/common/connectionTreeService';
+import { IContextKeyService } from 'vs/platform/contextkey/common/contextkey';
 suite('ConnectionDialogWidget tests', () => {
 	const testTreeViewId = 'testTreeView';
 	const ViewsRegistry = Registry.as<IViewsRegistry>(Extensions.ViewsRegistry);
@@ -53,6 +54,8 @@ suite('ConnectionDialogWidget tests', () => {
 		cmInstantiationService = new TestInstantiationService();
 		cmInstantiationService.stub(IStorageService, new TestStorageService());
 		cmInstantiationService.stub(IConnectionTreeService, new ConnectionTreeService());
+		cmInstantiationService.stub(IContextKeyService, new MockContextKeyService());
+
 		mockConnectionManagementService = TypeMoq.Mock.ofType(ConnectionManagementService, TypeMoq.MockBehavior.Strict,
 			undefined, // connection store
 			undefined, // connection status manager

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

@@ -134,6 +134,7 @@ export class MenuId {
 	static readonly ExplorerWidgetContext = new MenuId('ExplorerWidgetContext'); // {{SQL CARBON EDIT}}
 	static readonly DashboardToolbar = new MenuId('DashboardToolbar'); // {{SQL CARBON EDIT}}
 	static readonly NotebookTitle = new MenuId('NotebookTitle'); // {{SQL CARBON EDIT}}
+	static readonly ConnectionDialogBrowseTreeContext = new MenuId('ConnectionDialogBrowseTreeContext'); // {{SQL CARBON EDIT}}
 	static readonly TimelineItemContext = new MenuId('TimelineItemContext');
 	static readonly TimelineTitle = new MenuId('TimelineTitle');
 	static readonly TimelineTitleContext = new MenuId('TimelineTitleContext');

src/vs/workbench/api/common/menusExtensionPoint.ts

@@ -206,6 +206,11 @@ const apiMenus: IAPIMenu[] = [
 		id: MenuId.ObjectExplorerItemContext,
 		description: localize('objectExplorer.context', "The object explorer item context menu")
 	},
+	{
+		key: 'connectionDialog/browseTree',
+		id: MenuId.ConnectionDialogBrowseTreeContext,
+		description: localize('connectionDialogBrowseTree.context', "The connection dialog's browse tree context menu")
+	}
 	// {{SQL CARBON EDIT}} end menu entries
 ];