Tenant list UI for Firewall Rules (#11539)

* Start on the tenant list

* continue work

* Finish up...

* Fix test

* Fix

* Fix tests

* Some PR feedback

* Move responsibilities around

* Fix comment

src/sql/platform/accounts/common/accountPickerViewModel.ts

@@ -18,6 +18,7 @@ export class AccountPickerViewModel {
 	public get updateAccountListEvent(): Event<UpdateAccountListEventParams> { return this._updateAccountListEmitter.event; }
 
 	public selectedAccount: azdata.Account | undefined;
+	public selectedTenantId: string | undefined;
 
 	constructor(
 		_providerId: string,
@@ -35,13 +36,12 @@ export class AccountPickerViewModel {
 	 * Loads an initial list of accounts from the account management service
 	 * @return Promise to return the list of accounts
 	 */
-	public initialize(): Thenable<azdata.Account[]> {
+	public async initialize(): Promise<azdata.Account[]> {
-		// Load a baseline of the accounts for the provider
+		try {
-		return this._accountManagementService.getAccounts()
+			const accounts = await this._accountManagementService.getAccounts();
-			.then(undefined, () => {
+			return accounts;
-				// In the event we failed to lookup accounts for the provider, just send
+		} catch{
-				// back an empty collection
+			return [];
-				return [];
+		}
-			});
 	}
 }

src/sql/platform/accounts/common/firewallRuleViewModel.ts

@@ -11,6 +11,7 @@ import * as azdata from 'azdata';
 export class FirewallRuleViewModel {
 	public isIPAddressSelected: boolean;
 	public selectedAccount: azdata.Account | undefined;
+	public selectedTenantId: string | undefined;
 
 	private _defaultIPAddress?: string;
 	private _defaultFromSubnetIPRange?: string;

src/sql/workbench/services/accountManagement/browser/accountDialog.ts

@@ -44,6 +44,7 @@ import { ILayoutService } from 'vs/platform/layout/browser/layoutService';
 import { Registry } from 'vs/platform/registry/common/platform';
 import { SyncDescriptor } from 'vs/platform/instantiation/common/descriptors';
 import { Iterable } from 'vs/base/common/iterator';
+import { Tenant, TenantListDelegate, TenantListRenderer } from 'sql/workbench/services/accountManagement/browser/tenantListRenderer';
 
 export const VIEWLET_ID = 'workbench.view.accountpanel';
 
@@ -61,6 +62,8 @@ export const ACCOUNT_VIEW_CONTAINER = Registry.as<IViewContainersRegistry>(ViewC
 class AccountPanel extends ViewPane {
 	public index: number;
 	private accountList: List<azdata.Account>;
+	private tenantList: List<Tenant>;
+
 
 	constructor(
 		private options: IViewPaneOptions,
@@ -79,13 +82,14 @@ class AccountPanel extends ViewPane {
 
 	protected renderBody(container: HTMLElement): void {
 		this.accountList = new List<azdata.Account>('AccountList', container, new AccountListDelegate(AccountDialog.ACCOUNTLIST_HEIGHT), [this.instantiationService.createInstance(AccountListRenderer)]);
+		this.tenantList = new List<Tenant>('TenantList', container, new TenantListDelegate(AccountDialog.ACCOUNTLIST_HEIGHT), [this.instantiationService.createInstance(TenantListRenderer)]);
 		this._register(attachListStyler(this.accountList, this.themeService));
+		this._register(attachListStyler(this.tenantList, this.themeService));
 	}
 
 	protected layoutBody(size: number): void {
-		if (this.accountList) {
+		this.accountList?.layout(size);
-			this.accountList.layout(size);
+		this.tenantList?.layout(size);
-		}
 	}
 
 	public get length(): number {
@@ -102,6 +106,11 @@ class AccountPanel extends ViewPane {
 
 	public setSelection(indexes: number[]) {
 		this.accountList.setSelection(indexes);
+		this.updateTenants(this.accountList.getSelection[0]);
+	}
+
+	private updateTenants(account: azdata.Account) {
+		this.tenantList.splice(0, this.tenantList.length, account?.properties?.tenants ?? []);
 	}
 
 	public getActions(): IAction[] {

src/sql/workbench/services/accountManagement/browser/accountListRenderer.ts

@@ -32,17 +32,20 @@ export class AccountListDelegate implements IListVirtualDelegate<azdata.Account>
 	}
 }
 
-export interface AccountPickerListTemplate {
+export interface PickerListTemplate {
 	root: HTMLElement;
+	label: HTMLElement;
+	displayName: HTMLElement;
+	content: HTMLElement;
+}
+
+export interface AccountPickerListTemplate extends PickerListTemplate {
 	icon: HTMLElement;
 	badgeContent: HTMLElement;
 	contextualDisplayName: HTMLElement;
-	label: HTMLElement;
-	displayName: HTMLElement;
 }
 
 export interface AccountListTemplate extends AccountPickerListTemplate {
-	content: HTMLElement;
 	actions: ActionBar;
 }
 

src/sql/workbench/services/accountManagement/browser/accountPicker.ts

@@ -11,10 +11,11 @@ import * as azdata from 'azdata';
 export const IAccountPickerService = createDecorator<IAccountPickerService>('AccountPickerService');
 export interface IAccountPickerService {
 	_serviceBrand: undefined;
-	renderAccountPicker(container: HTMLElement): void;
+	renderAccountPicker(rootContainer: HTMLElement): void;
 	addAccountCompleteEvent: Event<void>;
 	addAccountErrorEvent: Event<string>;
 	addAccountStartEvent: Event<void>;
 	onAccountSelectionChangeEvent: Event<azdata.Account | undefined>;
+	onTenantSelectionChangeEvent: Event<string | undefined>;
 	selectedAccount: azdata.Account | undefined;
 }

src/sql/workbench/services/accountManagement/browser/accountPickerImpl.ts

@@ -5,6 +5,7 @@
 
 import 'vs/css!./media/accountPicker';
 import * as DOM from 'vs/base/browser/dom';
+import { localize } from 'vs/nls';
 import { Event, Emitter } from 'vs/base/common/event';
 import { List } from 'vs/base/browser/ui/list/listWidget';
 import { IDropdownOptions } from 'vs/base/browser/ui/dropdown/dropdown';
@@ -24,15 +25,22 @@ import { AddAccountAction, RefreshAccountAction } from 'sql/platform/accounts/co
 import { AccountPickerListRenderer, AccountListDelegate } from 'sql/workbench/services/accountManagement/browser/accountListRenderer';
 import { AccountPickerViewModel } from 'sql/platform/accounts/common/accountPickerViewModel';
 import { firstIndex } from 'vs/base/common/arrays';
+import { Tenant, TenantListDelegate, TenantPickerListRenderer } from 'sql/workbench/services/accountManagement/browser/tenantListRenderer';
 
 export class AccountPicker extends Disposable {
 	public static ACCOUNTPICKERLIST_HEIGHT = 47;
 	public viewModel: AccountPickerViewModel;
 	private _accountList: List<azdata.Account>;
-	private _rootElement: HTMLElement;
+	private _rootContainer: HTMLElement;
+
+	private _accountContainer: HTMLElement;
 	private _refreshContainer: HTMLElement;
-	private _listContainer: HTMLElement;
+	private _accountListContainer: HTMLElement;
 	private _dropdown: DropdownList;
+	private _tenantContainer: HTMLElement;
+	private _tenantListContainer: HTMLElement;
+	private _tenantList: List<Tenant>;
+	private _tenantDropdown: DropdownList;
 	private _refreshAccountAction: RefreshAccountAction;
 
 	// EVENTING ////////////////////////////////////////////////////////////
@@ -48,6 +56,9 @@ export class AccountPicker extends Disposable {
 	private _onAccountSelectionChangeEvent: Emitter<azdata.Account | undefined>;
 	public get onAccountSelectionChangeEvent(): Event<azdata.Account | undefined> { return this._onAccountSelectionChangeEvent.event; }
 
+	private _onTenantSelectionChangeEvent: Emitter<string | undefined>;
+	public get onTenantSelectionChangeEvent(): Event<string | undefined> { return this._onTenantSelectionChangeEvent.event; }
+
 	constructor(
 		private _providerId: string,
 		@IThemeService private _themeService: IThemeService,
@@ -61,6 +72,7 @@ export class AccountPicker extends Disposable {
 		this._addAccountErrorEmitter = new Emitter<string>();
 		this._addAccountStartEmitter = new Emitter<void>();
 		this._onAccountSelectionChangeEvent = new Emitter<azdata.Account>();
+		this._onTenantSelectionChangeEvent = new Emitter<string | undefined>();
 
 		// Create the view model, wire up the events, and initialize with baseline data
 		this.viewModel = this._instantiationService.createInstance(AccountPickerViewModel, this._providerId);
@@ -75,8 +87,8 @@ export class AccountPicker extends Disposable {
 	/**
 	 * Render account picker
 	 */
-	public render(container: HTMLElement): void {
+	public render(rootContainer: HTMLElement): void {
-		DOM.append(container, this._rootElement);
+		DOM.append(rootContainer, this._rootContainer);
 	}
 
 	// PUBLIC METHODS //////////////////////////////////////////////////////
@@ -85,18 +97,50 @@ export class AccountPicker extends Disposable {
 	 */
 	public createAccountPickerComponent() {
 		// Create an account list
-		const delegate = new AccountListDelegate(AccountPicker.ACCOUNTPICKERLIST_HEIGHT);
+		const accountDelegate = new AccountListDelegate(AccountPicker.ACCOUNTPICKERLIST_HEIGHT);
+		const tenantDelegate = new TenantListDelegate(AccountPicker.ACCOUNTPICKERLIST_HEIGHT);
+
 		const accountRenderer = new AccountPickerListRenderer();
-		this._listContainer = DOM.$('div.account-list-container');
+		const tenantRenderer = new TenantPickerListRenderer();
-		this._accountList = new List<azdata.Account>('AccountPicker', this._listContainer, delegate, [accountRenderer]);
+		this._rootContainer = DOM.$('div.account-picker-container');
+
+		const azureAccountLabel = localize('azureAccount', "Azure account");
+		const azureTenantLabel = localize('azureTenant', "Azure tenant");
+
+		const accountLabel = this.createLabelElement(azureAccountLabel, true);
+		this._accountListContainer = DOM.append(accountLabel, DOM.$('div.account-list-container'));
+
+		const tenantLabel = this.createLabelElement(azureTenantLabel, true);
+		this._tenantListContainer = DOM.append(tenantLabel, DOM.$('div.tenant-list-container'));
+
+
+		this._accountList = new List<azdata.Account>('AccountPicker', this._accountListContainer, accountDelegate, [accountRenderer], {
+			setRowLineHeight: false,
+		});
+		this._tenantList = new List<Tenant>('TenantPicker', this._tenantListContainer, tenantDelegate, [tenantRenderer]);
+
 		this._register(attachListStyler(this._accountList, this._themeService));
+		this._register(attachListStyler(this._tenantList, this._themeService));
 
-		this._rootElement = DOM.$('div.account-picker-container');
+		this._accountContainer = DOM.$('div.account-picker');
+		this._tenantContainer = DOM.$('div.tenant-picker');
 
-		// Create a dropdown for account picker
+		DOM.append(this._accountContainer, accountLabel);
-		const option: IDropdownOptions = {
+		DOM.append(this._tenantContainer, tenantLabel);
+
+
+		DOM.append(this._rootContainer, this._accountContainer);
+		DOM.append(this._rootContainer, this._tenantContainer);
+
+		// Create dropdowns for account and tenant pickers
+		const accountOptions: IDropdownOptions = {
 			contextViewProvider: this._contextViewService,
-			labelRenderer: (container) => this.renderLabel(container)
+			labelRenderer: (container) => this.renderAccountLabel(container)
+		};
+
+		const tenantOption: IDropdownOptions = {
+			contextViewProvider: this._contextViewService,
+			labelRenderer: (container) => this.renderTenantLabel(container)
 		};
 
 		// Create the add account action
@@ -105,8 +149,12 @@ export class AccountPicker extends Disposable {
 		addAccountAction.addAccountErrorEvent((msg) => this._addAccountErrorEmitter.fire(msg));
 		addAccountAction.addAccountStartEvent(() => this._addAccountStartEmitter.fire());
 
-		this._dropdown = this._register(new DropdownList(this._rootElement, option, this._listContainer, this._accountList, addAccountAction));
+		this._dropdown = this._register(new DropdownList(this._accountContainer, accountOptions, this._accountListContainer, this._accountList, addAccountAction));
+		this._tenantDropdown = this._register(new DropdownList(this._tenantContainer, tenantOption, this._tenantListContainer, this._tenantList));
+
 		this._register(attachDropdownStyler(this._dropdown, this._themeService));
+		this._register(attachDropdownStyler(this._tenantDropdown, this._themeService));
+
 		this._register(this._accountList.onDidChangeSelection((e: IListEvent<azdata.Account>) => {
 			if (e.elements.length === 1) {
 				this._dropdown.renderLabel();
@@ -114,8 +162,15 @@ export class AccountPicker extends Disposable {
 			}
 		}));
 
+		this._register(this._tenantList.onDidChangeSelection((e: IListEvent<Tenant>) => {
+			if (e.elements.length === 1) {
+				this._tenantDropdown.renderLabel();
+				this.onTenantSelectionChange(e.elements[0].id);
+			}
+		}));
+
 		// Create refresh account action
-		this._refreshContainer = DOM.append(this._rootElement, DOM.$('div.refresh-container'));
+		this._refreshContainer = DOM.append(this._accountContainer, DOM.$('div.refresh-container'));
 		DOM.append(this._refreshContainer, DOM.$('div.sql codicon warning'));
 		const actionBar = new ActionBar(this._refreshContainer, { animated: false });
 		this._refreshAccountAction = this._instantiationService.createInstance(RefreshAccountAction);
@@ -146,6 +201,17 @@ export class AccountPicker extends Disposable {
 	}
 
 	// PRIVATE HELPERS /////////////////////////////////////////////////////
+
+	private createLabelElement(content: string, isHeader?: boolean) {
+		let className = 'dialog-label';
+		if (isHeader) {
+			className += ' header';
+		}
+		const element = DOM.$(`.${className}`);
+		element.innerText = content;
+		return element;
+	}
+
 	private onAccountSelectionChange(account: azdata.Account | undefined) {
 		this.viewModel.selectedAccount = account;
 		if (account && account.isStale) {
@@ -153,12 +219,25 @@ export class AccountPicker extends Disposable {
 			DOM.show(this._refreshContainer);
 		} else {
 			DOM.hide(this._refreshContainer);
+
+			if (account.properties.tenants?.length > 1) {
+				DOM.show(this._tenantContainer);
+				this.updateTenantList(account);
+			} else {
+				DOM.hide(this._tenantContainer);
+			}
+			this.onTenantSelectionChange(account?.properties?.tenants[0]?.id);
 		}
 
 		this._onAccountSelectionChangeEvent.fire(account);
 	}
 
-	private renderLabel(container: HTMLElement): IDisposable | null {
+	private onTenantSelectionChange(tenantId: string | undefined) {
+		this.viewModel.selectedTenantId = tenantId;
+		this._onTenantSelectionChangeEvent.fire(tenantId);
+	}
+
+	private renderAccountLabel(container: HTMLElement): IDisposable | null {
 		if (container.hasChildNodes()) {
 			for (let i = 0; i < container.childNodes.length; i++) {
 				container.removeChild(container.childNodes.item(i));
@@ -193,6 +272,32 @@ export class AccountPicker extends Disposable {
 		return null;
 	}
 
+	private renderTenantLabel(container: HTMLElement): IDisposable | null {
+		if (container.hasChildNodes()) {
+			for (let i = 0; i < container.childNodes.length; i++) {
+				container.removeChild(container.childNodes.item(i));
+			}
+		}
+
+		const selectedTenants = this._tenantList.getSelectedElements();
+		const tenant = selectedTenants ? selectedTenants[0] : undefined;
+		if (tenant) {
+			const row = DOM.append(container, DOM.$('div.selected-tenant-container'));
+			const label = DOM.append(row, DOM.$('div.label'));
+
+			// TODO: Pick between the light and dark logo
+			label.innerText = tenant.displayName;
+		}
+		return null;
+	}
+
+	private updateTenantList(account: azdata.Account): void {
+		this._tenantList.splice(0, this._tenantList.length, account?.properties?.tenants ?? []);
+		this._tenantList.setSelection([0]);
+		this._tenantDropdown.renderLabel();
+		this._tenantList.layout(this._tenantList.contentHeight);
+	}
+
 	private updateAccountList(accounts: azdata.Account[]): void {
 		// keep the selection to the current one
 		const selectedElements = this._accountList.getSelectedElements();

src/sql/workbench/services/accountManagement/browser/accountPickerService.ts

@@ -28,6 +28,9 @@ export class AccountPickerService implements IAccountPickerService {
 	private _onAccountSelectionChangeEvent: Emitter<azdata.Account | undefined>;
 	public get onAccountSelectionChangeEvent(): Event<azdata.Account | undefined> { return this._onAccountSelectionChangeEvent.event; }
 
+	private _onTenantSelectionChangeEvent: Emitter<string | undefined>;
+	public get onTenantSelectionChangeEvent(): Event<string | undefined> { return this._onTenantSelectionChangeEvent.event; }
+
 	constructor(
 		@IInstantiationService private _instantiationService: IInstantiationService
 	) {
@@ -36,6 +39,7 @@ export class AccountPickerService implements IAccountPickerService {
 		this._addAccountErrorEmitter = new Emitter<string>();
 		this._addAccountStartEmitter = new Emitter<void>();
 		this._onAccountSelectionChangeEvent = new Emitter<azdata.Account>();
+		this._onTenantSelectionChangeEvent = new Emitter<string | undefined>();
 	}
 
 	/**
@@ -48,7 +52,7 @@ export class AccountPickerService implements IAccountPickerService {
 	/**
 	 * Render account picker
 	 */
-	public renderAccountPicker(container: HTMLElement): void {
+	public renderAccountPicker(rootContainer: HTMLElement): void {
 		if (!this._accountPicker) {
 			// TODO: expand support to multiple providers
 			const providerId: string = 'azure_publicCloud';
@@ -60,6 +64,7 @@ export class AccountPickerService implements IAccountPickerService {
 		this._accountPicker.addAccountErrorEvent((msg) => this._addAccountErrorEmitter.fire(msg));
 		this._accountPicker.addAccountStartEvent(() => this._addAccountStartEmitter.fire());
 		this._accountPicker.onAccountSelectionChangeEvent((account) => this._onAccountSelectionChangeEvent.fire(account));
-		this._accountPicker.render(container);
+		this._accountPicker.onTenantSelectionChangeEvent((tenantId) => this._onTenantSelectionChangeEvent.fire(tenantId));
+		this._accountPicker.render(rootContainer);
 	}
 }

src/sql/workbench/services/accountManagement/browser/media/accountListRenderer.css

@@ -3,12 +3,14 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-.list-row.account-picker-list {
+.list-row.account-picker-list,
+.list-row.tenant-picker-list {
 	display: flex;
 	align-items: flex-start;
 }
 
-.list-row.account-picker-list .label {
+.list-row.account-picker-list .label,
+.list-row.tenant-picker-list .label {
 	flex: 1 1 auto;
 	margin-left: 15px;
 	overflow: hidden;
@@ -18,14 +20,16 @@
 	font-size: 15px;
 }
 
-.list-row.account-picker-list .label .display-name {
+.list-row.account-picker-list .label .display-name,
+.list-row.tenant-picker-list .label .display-name {
 	font-size: 13px;
 }
 
-.list-row.account-picker-list .label .content {
+.list-row.account-picker-list .label .content,
+.list-row.tenant-picker-list .label .content{
 	opacity: 0.7;
 }
 
-.account-logo {
+.account-picker-list .account-logo {
 	background: no-repeat center center;
 }

src/sql/workbench/services/accountManagement/browser/media/accountPicker.css

@@ -4,7 +4,8 @@
  *--------------------------------------------------------------------------------------------*/
 
 /* Selected account */
-.selected-account-container {
+.selected-account-container,
+.selected-tenant-container {
 	padding: 6px;
 	display: flex;
 	align-items: flex-start;
@@ -16,7 +17,8 @@
 	width: 25px;
 }
 
-.selected-account-container .label {
+.selected-account-container .label,
+.selected-tenant-container .label {
 	flex: 1 1 auto;
 	padding-left: 10px;
 	align-self: center;
@@ -49,7 +51,7 @@
 }
 
 /* Account list */
-.account-list-container .list-row {
+.account-list-container .list-row, .tenant-list-container .list-row {
 	padding: 6px;
 }
 

src/sql/workbench/services/accountManagement/browser/tenantListRenderer.ts

@@ -0,0 +1,82 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { IListRenderer, IListVirtualDelegate } from 'vs/base/browser/ui/list/list';
+import { PickerListTemplate } from 'sql/workbench/services/accountManagement/browser/accountListRenderer';
+import * as DOM from 'vs/base/browser/dom';
+
+export interface Tenant {
+	id: string;
+	displayName: string;
+}
+
+export interface TenantPickerListTemplate extends PickerListTemplate {
+}
+
+export class TenantListDelegate implements IListVirtualDelegate<Tenant> {
+
+	constructor(
+		private _height: number
+	) {
+	}
+
+	public getHeight(element: Tenant): number {
+		return this._height;
+	}
+
+	public getTemplateId(element: Tenant): string {
+		return 'tenantListRenderer';
+	}
+}
+
+export class TenantPickerListRenderer implements IListRenderer<Tenant, TenantPickerListTemplate> {
+	public static TEMPLATE_ID = 'tenantListRenderer';
+
+	public get templateId(): string {
+		return TenantPickerListRenderer.TEMPLATE_ID;
+	}
+
+	public renderTemplate(container: HTMLElement): TenantPickerListTemplate {
+		const tableTemplate: TenantPickerListTemplate = Object.create(null);
+		tableTemplate.root = DOM.append(container, DOM.$('div.list-row.tenant-picker-list'));
+		tableTemplate.label = DOM.append(tableTemplate.root, DOM.$('div.label'));
+		tableTemplate.displayName = DOM.append(tableTemplate.label, DOM.$('div.display-name'));
+		return tableTemplate;
+	}
+
+	public renderElement(tenant: Tenant, index: number, templateData: PickerListTemplate): void {
+		templateData.displayName.innerText = tenant.displayName;
+	}
+
+	public disposeTemplate(template: PickerListTemplate): void {
+		// noop
+	}
+
+	public disposeElement(element: Tenant, index: number, templateData: PickerListTemplate): void {
+		// noop
+	}
+}
+
+export class TenantListRenderer extends TenantPickerListRenderer {
+	constructor(
+	) {
+		super();
+	}
+
+	public get templateId(): string {
+		return TenantPickerListRenderer.TEMPLATE_ID;
+	}
+
+	public renderTemplate(container: HTMLElement): PickerListTemplate {
+		const tableTemplate = super.renderTemplate(container) as PickerListTemplate;
+		tableTemplate.content = DOM.append(tableTemplate.label, DOM.$('div.content'));
+
+		return tableTemplate;
+	}
+
+	public renderElement(tenant: Tenant, index: number, templateData: PickerListTemplate): void {
+		super.renderElement(tenant, index, templateData);
+	}
+}

src/sql/workbench/services/accountManagement/test/browser/accountPickerService.test.ts

@@ -20,6 +20,7 @@ let mockAddAccountCompleteEmitter: Emitter<void>;
 let mockAddAccountErrorEmitter: Emitter<string>;
 let mockAddAccountStartEmitter: Emitter<void>;
 let mockOnAccountSelectionChangeEvent: Emitter<azdata.Account>;
+let mockOnTenantSelectionChangeEvent: Emitter<string>;
 
 // TESTS ///////////////////////////////////////////////////////////////////
 suite('Account picker service tests', () => {
@@ -29,6 +30,7 @@ suite('Account picker service tests', () => {
 		mockAddAccountErrorEmitter = new Emitter<string>();
 		mockAddAccountStartEmitter = new Emitter<void>();
 		mockOnAccountSelectionChangeEvent = new Emitter<azdata.Account>();
+		mockOnTenantSelectionChangeEvent = new Emitter<string>();
 	});
 
 	test('Construction - Events are properly defined', () => {
@@ -111,6 +113,8 @@ function createInstantiationService(): InstantiationService {
 		.returns(() => mockAddAccountStartEmitter.event);
 	mockAccountDialog.setup(x => x.onAccountSelectionChangeEvent)
 		.returns((account) => mockOnAccountSelectionChangeEvent.event);
+	mockAccountDialog.setup(x => x.onTenantSelectionChangeEvent)
+		.returns((tenant) => mockOnTenantSelectionChangeEvent.event);
 	mockAccountDialog.setup(x => x.render(TypeMoq.It.isAny()))
 		.returns((container) => undefined);
 	mockAccountDialog.setup(x => x.createAccountPickerComponent());

src/sql/workbench/services/resourceProvider/browser/firewallRuleDialog.ts

@@ -138,11 +138,10 @@ export class FirewallRuleDialog extends Modal {
 		});
 		this._accountPickerService.addAccountStartEvent(() => this.spinner = true);
 		this._accountPickerService.onAccountSelectionChangeEvent((account) => this.onAccountSelectionChange(account));
+		this._accountPickerService.onTenantSelectionChangeEvent((tenantId) => this.onTenantSelectionChange(tenantId));
 
 		const azureAccountSection = DOM.append(body, DOM.$('.azure-account-section.new-section'));
-		const azureAccountLabel = localize('azureAccount', "Azure account");
+		this._accountPickerService.renderAccountPicker(azureAccountSection);
-		this.createLabelElement(azureAccountSection, azureAccountLabel, true);
-		this._accountPickerService.renderAccountPicker(DOM.append(azureAccountSection, DOM.$('.dialog-input')));
 
 		const firewallRuleSection = DOM.append(body, DOM.$('.firewall-rule-section.new-section'));
 		const firewallRuleLabel = localize('filewallRule', "Firewall rule");
@@ -215,7 +214,9 @@ export class FirewallRuleDialog extends Modal {
 		if (isHeader) {
 			className += ' header';
 		}
-		DOM.append(container, DOM.$(`.${className}`)).innerText = content;
+		const element = DOM.append(container, DOM.$(`.${className}`));
+		element.innerText = content;
+		return element;
 	}
 
 	// Update theming that is specific to firewall rule flyout body
@@ -289,6 +290,10 @@ export class FirewallRuleDialog extends Modal {
 		}
 	}
 
+	public onTenantSelectionChange(tenantId: string): void {
+		this.viewModel.selectedTenantId = tenantId;
+	}
+
 	public onServiceComplete() {
 		this._createButton.enabled = true;
 		this.spinner = false;

src/sql/workbench/services/resourceProvider/browser/firewallRuleDialogController.ts

@@ -60,7 +60,7 @@ export class FirewallRuleDialogController {
 	private async handleOnCreateFirewallRule(): Promise<void> {
 		const resourceProviderId = this._resourceProviderId;
 		try {
-			const tenantId = this._connection.azureTenantId;
+			const tenantId = this._firewallRuleDialog.viewModel.selectedTenantId;
 			const token = await this._accountManagementService.getAccountSecurityToken(this._firewallRuleDialog.viewModel.selectedAccount!, tenantId, AzureResource.ResourceManagement);
 			const securityTokenMappings = {
 				[tenantId]: token

src/sql/workbench/services/resourceProvider/browser/media/firewallRuleDialog.css

@@ -42,10 +42,6 @@
 	padding-bottom: 0px;
 }
 
-.modal .firewall-rule-dialog .azure-account-section {
-	height: 92px;
-}
-
 /* Firewall rule description section */
 .modal .firewall-rule-dialog a:link {
 	text-decoration: underline;

src/sql/workbench/services/resourceProvider/test/browser/firewallRuleDialogController.test.ts

@@ -52,6 +52,7 @@ suite('Firewall rule dialog controller tests', () => {
 		mockFirewallRuleViewModel.setup(x => x.updateDefaultValues(TypeMoq.It.isAny()))
 			.returns((ipAddress) => undefined);
 		mockFirewallRuleViewModel.object.selectedAccount = account;
+		mockFirewallRuleViewModel.object.selectedTenantId = 'tenantId';
 		mockFirewallRuleViewModel.object.isIPAddressSelected = true;
 
 		// Create a mocked out instantiation service