Add connection profile persistence to MIAA dashboard (#11061)

* wip

* fixes

* fix pg model

* more updates

* Add resourceType check

extensions/arc/src/localizedConstants.ts

@@ -29,7 +29,7 @@ export const newSupportRequest = localize('arc.newSupportRequest', "New support
 export const diagnoseAndSolveProblems = localize('arc.diagnoseAndSolveProblems', "Diagnose and solve problems");
 export const supportAndTroubleshooting = localize('arc.supportAndTroubleshooting', "Support + troubleshooting");
 
-export const createNew = localize('arc.createNew', "Create New");
+export const newInstance = localize('arc.createNew', "New Instance");
 export const deleteText = localize('arc.delete', "Delete");
 export const resetPassword = localize('arc.resetPassword', "Reset Password");
 export const openInAzurePortal = localize('arc.openInAzurePortal', "Open in Azure Portal");
@@ -116,10 +116,9 @@ export function resourceDeleted(name: string): string { return localize('arc.res
 export function couldNotFindAzureResource(name: string): string { return localize('arc.couldNotFindAzureResource', "Could not find Azure resource for {0}", name); }
 export function copiedToClipboard(name: string): string { return localize('arc.copiedToClipboard', "{0} copied to clipboard", name); }
 export function clickTheTroubleshootButton(resourceType: string): string { return localize('arc.clickTheTroubleshootButton', "Click the troubleshoot button to open the Azure Arc {0} troubleshooting notebook.", resourceType); }
-export function numVCores(vCores: string): string {
-	const numCores = +vCores;
-	if (numCores && numCores > 0) {
-		return localize('arc.numVCores', "{0} vCores", numCores);
+export function numVCores(vCores: string | undefined): string {
+	if (vCores && +vCores > 0) {
+		return localize('arc.numVCores', "{0} vCores", vCores);
 	} else {
 		return '-';
 	}

extensions/arc/src/models/controllerModel.ts

@@ -14,7 +14,15 @@ import { AzureArcTreeDataProvider } from '../ui/tree/azureArcTreeDataProvider';
 export type ControllerInfo = {
 	url: string,
 	username: string,
-	rememberPassword: boolean
+	rememberPassword: boolean,
+	resources: ResourceInfo[]
+};
+
+export type ResourceInfo = {
+	namespace: string,
+	name: string,
+	resourceType: ResourceType | string,
+	connectionId?: string
 };
 
 export interface Registration extends RegistrationResponse {

extensions/arc/src/models/miaaModel.ts

@@ -8,15 +8,22 @@ import * as vscode from 'vscode';
 import { SqlInstanceRouterApi } from '../controller/generated/v1/api/sqlInstanceRouterApi';
 import { HybridSqlNsNameGetResponse } from '../controller/generated/v1/model/hybridSqlNsNameGetResponse';
 import { Authentication } from '../controller/generated/v1/api';
+import { ResourceModel } from './resourceModel';
+import { ResourceInfo, Registration } from './controllerModel';
+import { AzureArcTreeDataProvider } from '../ui/tree/azureArcTreeDataProvider';
+import { Deferred } from '../common/promise';
 
 export type DatabaseModel = { name: string, status: string };
 
-export class MiaaModel {
+export class MiaaModel extends ResourceModel {
 
 	private _sqlInstanceRouter: SqlInstanceRouterApi;
 	private _status: HybridSqlNsNameGetResponse | undefined;
 	private _databases: DatabaseModel[] = [];
+	// The saved connection information
 	private _connectionProfile: azdata.IConnectionProfile | undefined = undefined;
+	// The ID of the active connection used to query the server
+	private _activeConnectionId: string | undefined = undefined;
 
 	private readonly _onPasswordUpdated = new vscode.EventEmitter<string>();
 	private readonly _onStatusUpdated = new vscode.EventEmitter<HybridSqlNsNameGetResponse>();
@@ -26,23 +33,12 @@ export class MiaaModel {
 	public onDatabasesUpdated = this._onDatabasesUpdated.event;
 	public passwordLastUpdated?: Date;
 
-	constructor(controllerUrl: string, auth: Authentication, private _namespace: string, private _name: string) {
+	private _refreshPromise: Deferred<void> | undefined = undefined;
+
+	constructor(controllerUrl: string, controllerAuth: Authentication, info: ResourceInfo, registration: Registration, private _treeDataProvider: AzureArcTreeDataProvider) {
+		super(info, registration);
 		this._sqlInstanceRouter = new SqlInstanceRouterApi(controllerUrl);
-		this._sqlInstanceRouter.setDefaultAuthentication(auth);
-	}
-
-	/**
-	 * The name of this instance
-	 */
-	public get name(): string {
-		return this._name;
-	}
-
-	/**
-	 * The namespace of this instance
-	 */
-	public get namespace(): string {
-		return this._namespace;
+		this._sqlInstanceRouter.setDefaultAuthentication(controllerAuth);
 	}
 
 	/**
@@ -72,48 +68,118 @@ export class MiaaModel {
 
 	/** Refreshes the model */
 	public async refresh(): Promise<void> {
-		const instanceRefresh = this._sqlInstanceRouter.apiV1HybridSqlNsNameGet(this._namespace, this._name).then(response => {
-			this._status = response.body;
-			this._onStatusUpdated.fire(this._status);
-		});
-		const promises: Thenable<any>[] = [instanceRefresh];
-		await this.getConnection();
-		if (this._connectionProfile) {
-			const provider = azdata.dataprotocol.getProvider<azdata.MetadataProvider>(this._connectionProfile.providerName, azdata.DataProviderType.MetadataProvider);
-			const databasesRefresh = azdata.connection.getUriForConnection(this._connectionProfile.id).then(ownerUri => {
-				provider.getDatabases(ownerUri).then(databases => {
-					if (databases.length > 0 && typeof (databases[0]) === 'object') {
-						this._databases = (<azdata.DatabaseInfo[]>databases).map(db => { return { name: db.options['name'], status: db.options['state'] }; });
-					} else {
-						this._databases = (<string[]>databases).map(db => { return { name: db, status: '-' }; });
-					}
-					this._onDatabasesUpdated.fire(this._databases);
-				});
-			});
-			promises.push(databasesRefresh);
+		// Only allow one refresh to be happening at a time
+		if (this._refreshPromise) {
+			return this._refreshPromise.promise;
+		}
+		this._refreshPromise = new Deferred();
+		try {
+			const instanceRefresh = this._sqlInstanceRouter.apiV1HybridSqlNsNameGet(this.info.namespace, this.info.name).then(response => {
+				this._status = response.body;
+				this._onStatusUpdated.fire(this._status);
+			});
+			const promises: Thenable<any>[] = [instanceRefresh];
+			await this.getConnectionProfile();
+			if (this._connectionProfile) {
+				// We haven't connected yet so do so now and then store the ID for the active connection
+				if (!this._activeConnectionId) {
+					const result = await azdata.connection.connect(this._connectionProfile, false, false);
+					if (!result.connected) {
+						throw new Error(result.errorMessage);
+					}
+					this._activeConnectionId = result.connectionId;
+				}
+
+				const provider = azdata.dataprotocol.getProvider<azdata.MetadataProvider>(this._connectionProfile.providerName, azdata.DataProviderType.MetadataProvider);
+				const databasesRefresh = azdata.connection.getUriForConnection(this._activeConnectionId).then(ownerUri => {
+					provider.getDatabases(ownerUri).then(databases => {
+						if (!databases) {
+							throw new Error('Could not fetch databases');
+						}
+						if (databases.length > 0 && typeof (databases[0]) === 'object') {
+							this._databases = (<azdata.DatabaseInfo[]>databases).map(db => { return { name: db.options['name'], status: db.options['state'] }; });
+						} else {
+							this._databases = (<string[]>databases).map(db => { return { name: db, status: '-' }; });
+						}
+						this._onDatabasesUpdated.fire(this._databases);
+					});
+				});
+				promises.push(databasesRefresh);
+			}
+			await Promise.all(promises);
+		} finally {
+			this._refreshPromise = undefined;
 		}
-		await Promise.all(promises);
 	}
 
-	private async getConnection(): Promise<void> {
+	/**
+	 * Loads the saved connection profile associated with this model. Will prompt for one if
+	 * we don't have one or can't find it (it was deleted)
+	 */
+	private async getConnectionProfile(): Promise<void> {
 		if (this._connectionProfile) {
 			return;
 		}
-		const connection = await azdata.connection.openConnectionDialog(['MSSQL']);
-		this._connectionProfile = {
-			serverName: connection.options['serverName'],
-			databaseName: connection.options['databaseName'],
-			authenticationType: connection.options['authenticationType'],
-			providerName: 'MSSQL',
-			connectionName: '',
-			userName: connection.options['user'],
-			password: connection.options['password'],
-			savePassword: false,
-			groupFullName: undefined,
-			saveProfile: true,
-			id: connection.connectionId,
-			groupId: undefined,
-			options: connection.options
-		};
+		let connection: azdata.connection.ConnectionProfile | azdata.connection.Connection | undefined;
+
+		if (this.info.connectionId) {
+			try {
+				const connections = await azdata.connection.getConnections();
+				const existingConnection = connections.find(conn => conn.connectionId === this.info.connectionId);
+				if (existingConnection) {
+					const credentials = await azdata.connection.getCredentials(this.info.connectionId);
+					if (credentials) {
+						existingConnection.options['password'] = credentials.password;
+						connection = existingConnection;
+					} else {
+						// We need the password so prompt the user for it
+						const connectionProfile = {
+							serverName: existingConnection.options['serverName'],
+							databaseName: existingConnection.options['databaseName'],
+							authenticationType: existingConnection.options['authenticationType'],
+							providerName: 'MSSQL',
+							connectionName: '',
+							userName: existingConnection.options['user'],
+							password: '',
+							savePassword: false,
+							groupFullName: undefined,
+							saveProfile: true,
+							id: '',
+							groupId: undefined,
+							options: existingConnection.options
+						};
+						connection = await azdata.connection.openConnectionDialog(['MSSQL'], connectionProfile);
+					}
+				}
+			} catch (err) {
+				// ignore - the connection may not necessarily exist anymore and in that case we'll just reprompt for a connection
+			}
+		}
+
+		if (!connection) {
+			// Weren't able to load the existing connection so prompt user for new one
+			connection = await azdata.connection.openConnectionDialog(['MSSQL']);
+		}
+
+		if (connection) {
+			this._connectionProfile = {
+				serverName: connection.options['serverName'] || connection.options['server'],
+				databaseName: connection.options['databaseName'] || connection.options['database'],
+				authenticationType: connection.options['authenticationType'],
+				providerName: 'MSSQL',
+				connectionName: '',
+				userName: connection.options['user'],
+				password: connection.options['password'],
+				savePassword: false,
+				groupFullName: undefined,
+				saveProfile: true,
+				id: connection.connectionId,
+				groupId: undefined,
+				options: connection.options
+			};
+			this.info.connectionId = connection.connectionId;
+			await this._treeDataProvider.saveControllers();
+		}
+
 	}
 }

extensions/arc/src/models/postgresModel.ts

@@ -7,6 +7,8 @@ import * as vscode from 'vscode';
 import * as loc from '../localizedConstants';
 import { DuskyObjectModelsDatabaseService, DatabaseRouterApi, DuskyObjectModelsDatabase, V1Status, V1Pod } from '../controller/generated/dusky/api';
 import { Authentication } from '../controller/auth';
+import { ResourceInfo, Registration } from './controllerModel';
+import { ResourceModel } from './resourceModel';
 
 export enum PodRole {
 	Monitor,
@@ -14,7 +16,7 @@ export enum PodRole {
 	Shard
 }
 
-export class PostgresModel {
+export class PostgresModel extends ResourceModel {
 	private _databaseRouter: DatabaseRouterApi;
 	private _service?: DuskyObjectModelsDatabaseService;
 	private _pods?: V1Pod[];
@@ -25,24 +27,25 @@ export class PostgresModel {
 	public serviceLastUpdated?: Date;
 	public podsLastUpdated?: Date;
 
-	constructor(controllerUrl: string, auth: Authentication, private _namespace: string, private _name: string) {
+	constructor(controllerUrl: string, auth: Authentication, info: ResourceInfo, registration: Registration) {
+		super(info, registration);
 		this._databaseRouter = new DatabaseRouterApi(controllerUrl);
 		this._databaseRouter.setDefaultAuthentication(auth);
 	}
 
 	/** Returns the service's Kubernetes namespace */
 	public get namespace(): string {
-		return this._namespace;
+		return this.info.namespace;
 	}
 
 	/** Returns the service's name */
 	public get name(): string {
-		return this._name;
+		return this.info.name;
 	}
 
 	/** Returns the service's fully qualified name in the format namespace.name */
 	public get fullName(): string {
-		return `${this._namespace}.${this._name}`;
+		return `${this.info.namespace}.${this.info.name}`;
 	}
 
 	/** Returns the service's spec */
@@ -58,12 +61,12 @@ export class PostgresModel {
 	/** Refreshes the model */
 	public async refresh() {
 		await Promise.all([
-			this._databaseRouter.getDuskyDatabaseService(this._namespace, this._name).then(response => {
+			this._databaseRouter.getDuskyDatabaseService(this.info.namespace, this.info.name).then(response => {
 				this._service = response.body;
 				this.serviceLastUpdated = new Date();
 				this._onServiceUpdated.fire(this._service);
 			}),
-			this._databaseRouter.getDuskyPods(this._namespace, this._name).then(response => {
+			this._databaseRouter.getDuskyPods(this.info.namespace, this.info.name).then(response => {
 				this._pods = response.body;
 				this.podsLastUpdated = new Date();
 				this._onPodsUpdated.fire(this._pods!);
@@ -77,7 +80,7 @@ export class PostgresModel {
 	 */
 	public async update(func: (service: DuskyObjectModelsDatabaseService) => void): Promise<DuskyObjectModelsDatabaseService> {
 		// Get the latest spec of the service in case it has changed
-		const service = (await this._databaseRouter.getDuskyDatabaseService(this._namespace, this._name)).body;
+		const service = (await this._databaseRouter.getDuskyDatabaseService(this.info.namespace, this.info.name)).body;
 		service.status = undefined; // can't update the status
 		func(service);
 
@@ -89,7 +92,7 @@ export class PostgresModel {
 
 	/** Deletes the service */
 	public async delete(): Promise<V1Status> {
-		return (await this._databaseRouter.deleteDuskyDatabaseService(this._namespace, this._name)).body;
+		return (await this._databaseRouter.deleteDuskyDatabaseService(this.info.namespace, this.info.name)).body;
 	}
 
 	/** Creates a SQL database in the service */

extensions/arc/src/models/resourceModel.ts

@@ -0,0 +1,26 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+import { ResourceInfo, Registration } from './controllerModel';
+
+export abstract class ResourceModel {
+
+	private readonly _onRegistrationUpdated = new vscode.EventEmitter<Registration>();
+	public onRegistrationUpdated = this._onRegistrationUpdated.event;
+
+	constructor(public info: ResourceInfo, private _registration: Registration) { }
+
+	public get registration(): Registration {
+		return this._registration;
+	}
+
+	public set registration(newValue: Registration) {
+		this._registration = newValue;
+		this._onRegistrationUpdated.fire(this._registration);
+	}
+
+	public abstract refresh(): Promise<void>;
+}

extensions/arc/src/test/ui/tree/azureArcTreeDataProvider.test.ts

@@ -44,7 +44,7 @@ describe('AzureArcTreeDataProvider tests', function (): void {
 			treeDataProvider['_loading'] = false;
 			let children = await treeDataProvider.getChildren();
 			should(children.length).equal(0, 'There initially shouldn\'t be any children');
-			const controllerModel = new ControllerModel(treeDataProvider, { url: '127.0.0.1', username: 'sa', rememberPassword: true });
+			const controllerModel = new ControllerModel(treeDataProvider, { url: '127.0.0.1', username: 'sa', rememberPassword: true, resources: [] });
 			await treeDataProvider.addOrUpdateController(controllerModel, '');
 			should(children.length).equal(1, 'Controller node should be added correctly');
 			await treeDataProvider.addOrUpdateController(controllerModel, '');
@@ -55,11 +55,11 @@ describe('AzureArcTreeDataProvider tests', function (): void {
 			treeDataProvider['_loading'] = false;
 			let children = await treeDataProvider.getChildren();
 			should(children.length).equal(0, 'There initially shouldn\'t be any children');
-			const controllerModel = new ControllerModel(treeDataProvider, { url: '127.0.0.1', username: 'sa', rememberPassword: true });
+			const controllerModel = new ControllerModel(treeDataProvider, { url: '127.0.0.1', username: 'sa', rememberPassword: true, resources: [] });
 			await treeDataProvider.addOrUpdateController(controllerModel, '');
 			should(children.length).equal(1, 'Controller node should be added correctly');
 			should((<ControllerTreeNode>children[0]).model.info.rememberPassword).be.true('Info was not set correctly initially');
-			const controllerModel2 = new ControllerModel(treeDataProvider, { url: '127.0.0.1', username: 'sa', rememberPassword: false });
+			const controllerModel2 = new ControllerModel(treeDataProvider, { url: '127.0.0.1', username: 'sa', rememberPassword: false, resources: [] });
 			await treeDataProvider.addOrUpdateController(controllerModel2, '');
 			should(children.length).equal(1, 'Shouldn\'t add duplicate controller node');
 			should((<ControllerTreeNode>children[0]).model.info.rememberPassword).be.false('Info was not updated correctly');
@@ -84,8 +84,8 @@ describe('AzureArcTreeDataProvider tests', function (): void {
 	describe('removeController', function (): void {
 		it('removing a controller should work as expected', async function (): Promise<void> {
 			treeDataProvider['_loading'] = false;
-			const controllerModel = new ControllerModel(treeDataProvider, { url: '127.0.0.1', username: 'sa', rememberPassword: true });
-			const controllerModel2 = new ControllerModel(treeDataProvider, { url: '127.0.0.2', username: 'cloudsa', rememberPassword: true });
+			const controllerModel = new ControllerModel(treeDataProvider, { url: '127.0.0.1', username: 'sa', rememberPassword: true, resources: [] });
+			const controllerModel2 = new ControllerModel(treeDataProvider, { url: '127.0.0.2', username: 'cloudsa', rememberPassword: true, resources: [] });
 			await treeDataProvider.addOrUpdateController(controllerModel, '');
 			await treeDataProvider.addOrUpdateController(controllerModel2, '');
 			const children = <ControllerTreeNode[]>(await treeDataProvider.getChildren());

extensions/arc/src/ui/dashboards/controller/controllerDashboardOverviewPage.ts

@@ -115,7 +115,7 @@ export class ControllerDashboardOverviewPage extends DashboardPage {
 					headerCssStyles: cssStyles.tableHeader,
 					rowCssStyles: cssStyles.tableRow
 				}, {
-					displayName: loc.computeAndStorage,
+					displayName: loc.compute,
 					valueType: azdata.DeclarativeDataType.string,
 					width: '34%',
 					isReadOnly: true,
@@ -140,12 +140,12 @@ export class ControllerDashboardOverviewPage extends DashboardPage {
 
 	public get toolbarContainer(): azdata.ToolbarContainer {
 
-		const createNewButton = this.modelView.modelBuilder.button().withProperties<azdata.ButtonProperties>({
-			label: loc.createNew,
+		const newInstance = this.modelView.modelBuilder.button().withProperties<azdata.ButtonProperties>({
+			label: loc.newInstance,
 			iconPath: IconPathHelper.add
 		}).component();
 
-		createNewButton.onDidClick(async () => {
+		newInstance.onDidClick(async () => {
 			await vscode.commands.executeCommand('azdata.resource.deploy');
 		});
 
@@ -166,7 +166,7 @@ export class ControllerDashboardOverviewPage extends DashboardPage {
 
 		return this.modelView.modelBuilder.toolbarContainer().withToolbarItems(
 			[
-				{ component: createNewButton, toolbarSeparatorAfter: true },
+				{ component: newInstance, toolbarSeparatorAfter: true },
 				{ component: openInAzurePortalButton }
 			]
 		).component();
@@ -196,7 +196,7 @@ export class ControllerDashboardOverviewPage extends DashboardPage {
 						iconWidth: iconSize
 					})
 					.component();
-				return [imageComponent, r.instanceName, resourceTypeToDisplayName(r.instanceType), r.vCores];
+				return [imageComponent, r.instanceName, resourceTypeToDisplayName(r.instanceType), loc.numVCores(r.vCores)];
 			});
 		this._arcResourcesLoadingComponent.loading = false;
 	}

extensions/arc/src/ui/dashboards/miaa/miaaConnectionStringsPage.ts

@@ -18,8 +18,8 @@ export class MiaaConnectionStringsPage extends DashboardPage {
 
 	constructor(modelView: azdata.ModelView, private _controllerModel: ControllerModel, private _miaaModel: MiaaModel) {
 		super(modelView);
-		this._controllerModel.onRegistrationsUpdated(registrations => {
-			this._instanceRegistration = registrations.find(reg => reg.instanceType === ResourceType.sqlManagedInstances && reg.instanceName === this._miaaModel.name);
+		this._controllerModel.onRegistrationsUpdated(_ => {
+			this._instanceRegistration = this._controllerModel.getRegistration(ResourceType.sqlManagedInstances, this._miaaModel.info.namespace, this._miaaModel.info.name);
 			this.eventuallyRunOnInitialized(() => this.updateConnectionStrings());
 		});
 		this.refresh().catch(err => console.error(err));

extensions/arc/src/ui/dashboards/miaa/miaaDashboardOverviewPage.ts

@@ -34,7 +34,7 @@ export class MiaaDashboardOverviewPage extends DashboardPage {
 		subscriptionId: '-',
 		miaaAdmin: '-',
 		host: '-',
-		vCores: '-'
+		vCores: ''
 	};
 
 	constructor(modelView: azdata.ModelView, private _controllerModel: ControllerModel, private _miaaModel: MiaaModel) {
@@ -166,11 +166,6 @@ export class MiaaDashboardOverviewPage extends DashboardPage {
 
 	public get toolbarContainer(): azdata.ToolbarContainer {
 
-		const createDatabaseButton = this.modelView.modelBuilder.button().withProperties<azdata.ButtonProperties>({
-			label: loc.newDatabase,
-			iconPath: IconPathHelper.add
-		}).component();
-
 		const deleteButton = this.modelView.modelBuilder.button().withProperties<azdata.ButtonProperties>({
 			label: loc.deleteText,
 			iconPath: IconPathHelper.delete
@@ -179,55 +174,48 @@ export class MiaaDashboardOverviewPage extends DashboardPage {
 		deleteButton.onDidClick(async () => {
 			deleteButton.enabled = false;
 			try {
-				if (await promptForResourceDeletion(this._miaaModel.namespace, this._miaaModel.name)) {
-					await this._controllerModel.miaaDelete(this._miaaModel.namespace, this._miaaModel.name);
-					vscode.window.showInformationMessage(loc.resourceDeleted(this._miaaModel.name));
+				if (await promptForResourceDeletion(this._miaaModel.info.namespace, this._miaaModel.info.name)) {
+					await this._controllerModel.miaaDelete(this._miaaModel.info.namespace, this._miaaModel.info.name);
+					vscode.window.showInformationMessage(loc.resourceDeleted(this._miaaModel.info.name));
 				}
 			} catch (error) {
-				vscode.window.showErrorMessage(loc.resourceDeletionFailed(this._miaaModel.name, error));
+				vscode.window.showErrorMessage(loc.resourceDeletionFailed(this._miaaModel.info.name, error));
 			} finally {
 				deleteButton.enabled = true;
 			}
 		});
 
-		const resetPasswordButton = this.modelView.modelBuilder.button().withProperties<azdata.ButtonProperties>({
-			label: loc.resetPassword,
-			iconPath: IconPathHelper.edit
-		}).component();
-
 		const openInAzurePortalButton = this.modelView.modelBuilder.button().withProperties<azdata.ButtonProperties>({
 			label: loc.openInAzurePortal,
 			iconPath: IconPathHelper.openInTab
 		}).component();
 
 		openInAzurePortalButton.onDidClick(async () => {
-			const r = this._controllerModel.getRegistration(ResourceType.sqlManagedInstances, this._miaaModel.namespace, this._miaaModel.name);
+			const r = this._controllerModel.getRegistration(ResourceType.sqlManagedInstances, this._miaaModel.info.namespace, this._miaaModel.info.name);
 			if (r) {
 				vscode.env.openExternal(vscode.Uri.parse(
 					`https://portal.azure.com/#resource/subscriptions/${r.subscriptionId}/resourceGroups/${r.resourceGroupName}/providers/Microsoft.AzureData/${ResourceType.sqlManagedInstances}/${r.instanceName}`));
 			} else {
-				vscode.window.showErrorMessage(loc.couldNotFindRegistration(this._miaaModel.namespace, this._miaaModel.name));
+				vscode.window.showErrorMessage(loc.couldNotFindRegistration(this._miaaModel.info.namespace, this._miaaModel.info.name));
 			}
 		});
 
 		return this.modelView.modelBuilder.toolbarContainer().withToolbarItems(
 			[
-				{ component: createDatabaseButton },
-				{ component: deleteButton },
-				{ component: resetPasswordButton, toolbarSeparatorAfter: true },
+				{ component: deleteButton, toolbarSeparatorAfter: true },
 				{ component: openInAzurePortalButton }
 			]
 		).component();
 	}
 
 	private async handleRegistrationsUpdated(): Promise<void> {
-		const reg = this._controllerModel.getRegistration(ResourceType.sqlManagedInstances, this._miaaModel.namespace, this._miaaModel.name);
+		const reg = this._controllerModel.getRegistration(ResourceType.sqlManagedInstances, this._miaaModel.info.namespace, this._miaaModel.info.name);
 		if (reg) {
 			this._instanceProperties.resourceGroup = reg.resourceGroupName || '-';
 			this._instanceProperties.dataController = this._controllerModel.controllerRegistration?.instanceName || '-';
 			this._instanceProperties.region = (await getAzurecoreApi()).getRegionDisplayName(reg.location);
 			this._instanceProperties.subscriptionId = reg.subscriptionId || '-';
-			this._instanceProperties.vCores = reg.vCores || '-';
+			this._instanceProperties.vCores = reg.vCores || '';
 			this._instanceProperties.host = reg.externalEndpoint || '-';
 			this.refreshDisplayedProperties();
 		}
@@ -239,12 +227,12 @@ export class MiaaDashboardOverviewPage extends DashboardPage {
 	}
 
 	private handleEndpointsUpdated(endpoints: EndpointModel[]): void {
-		const kibanaQuery = `kubernetes_namespace:"${this._miaaModel.namespace}" and instance_name :"${this._miaaModel.name}"`;
+		const kibanaQuery = `kubernetes_namespace:"${this._miaaModel.info.namespace}" and instance_name :"${this._miaaModel.info.name}"`;
 		const kibanaUrl = `${endpoints.find(e => e.name === 'logsui')?.endpoint}/app/kibana#/discover?_a=(query:(language:kuery,query:'${kibanaQuery}'))`;
 		this._kibanaLink.label = kibanaUrl;
 		this._kibanaLink.url = kibanaUrl;
 
-		const grafanaUrl = `${endpoints.find(e => e.name === 'metricsui')?.endpoint}/d/wZx3OUdmz/azure-sql-db-managed-instance-metrics?var-hostname=${this._miaaModel.name}-0`;
+		const grafanaUrl = `${endpoints.find(e => e.name === 'metricsui')?.endpoint}/d/wZx3OUdmz/azure-sql-db-managed-instance-metrics?var-hostname=${this._miaaModel.info.name}-0`;
 		this._grafanaLink.label = grafanaUrl;
 		this._grafanaLink.url = grafanaUrl;
 

extensions/arc/src/ui/dialogs/connectControllerDialog.ts

@@ -87,7 +87,12 @@ export class ConnectToControllerDialog {
 		if (!this.urlInputBox.value || !this.usernameInputBox.value || !this.passwordInputBox.value) {
 			return false;
 		}
-		const controllerInfo: ControllerInfo = { url: this.urlInputBox.value, username: this.usernameInputBox.value, rememberPassword: this.rememberPwCheckBox.checked ?? false };
+		const controllerInfo: ControllerInfo = {
+			url: this.urlInputBox.value,
+			username: this.usernameInputBox.value,
+			rememberPassword: this.rememberPwCheckBox.checked ?? false,
+			resources: []
+		};
 		const controllerModel = new ControllerModel(this._treeDataProvider, controllerInfo, this.passwordInputBox.value);
 		try {
 			// Validate that we can connect to the controller

extensions/arc/src/ui/tree/azureArcTreeDataProvider.ts

@@ -52,7 +52,7 @@ export class AzureArcTreeDataProvider implements vscode.TreeDataProvider<TreeNod
 		if (controllerNode) {
 			controllerNode.model.info = model.info;
 		} else {
-			this._controllerNodes.push(new ControllerTreeNode(model, this._context));
+			this._controllerNodes.push(new ControllerTreeNode(model, this._context, this));
 		}
 		await this.updatePassword(model, password);
 		if (refreshTree) {
@@ -95,7 +95,7 @@ export class AzureArcTreeDataProvider implements vscode.TreeDataProvider<TreeNod
 			const controllerMementos: ControllerInfo[] = this._context.globalState.get(mementoToken) || [];
 			this._controllerNodes = controllerMementos.map(memento => {
 				const controllerModel = new ControllerModel(this, memento);
-				return new ControllerTreeNode(controllerModel, this._context);
+				return new ControllerTreeNode(controllerModel, this._context, this);
 			});
 		} finally {
 			this._loading = false;
@@ -103,8 +103,9 @@ export class AzureArcTreeDataProvider implements vscode.TreeDataProvider<TreeNod
 		}
 	}
 
-	private async saveControllers(): Promise<void> {
-		await this._context.globalState.update(mementoToken, this._controllerNodes.map(node => node.model.info));
+	public async saveControllers(): Promise<void> {
+		const controllerInfo = this._controllerNodes.map(node => node.model.info);
+		await this._context.globalState.update(mementoToken, controllerInfo);
 	}
 }
 

extensions/arc/src/ui/tree/controllerTreeNode.ts

@@ -8,23 +8,25 @@ import { TreeNode } from './treeNode';
 import { MiaaTreeNode } from './miaaTreeNode';
 import { ResourceType } from '../../constants';
 import { PostgresTreeNode } from './postgresTreeNode';
-import { ControllerModel, Registration } from '../../models/controllerModel';
+import { ControllerModel, Registration, ResourceInfo } from '../../models/controllerModel';
 import { ControllerDashboard } from '../dashboards/controller/controllerDashboard';
 import { PostgresModel } from '../../models/postgresModel';
 import { parseInstanceName } from '../../common/utils';
 import { MiaaModel } from '../../models/miaaModel';
 import { Deferred } from '../../common/promise';
 import { RefreshTreeNode } from './refreshTreeNode';
+import { ResourceTreeNode } from './resourceTreeNode';
+import { AzureArcTreeDataProvider } from './azureArcTreeDataProvider';
 
 /**
  * The TreeNode for displaying an Azure Arc Controller
  */
 export class ControllerTreeNode extends TreeNode {
 
-	private _children: TreeNode[] = [];
+	private _children: ResourceTreeNode[] = [];
 	private _childrenRefreshPromise = new Deferred();
 
-	constructor(public model: ControllerModel, private _context: vscode.ExtensionContext) {
+	constructor(public model: ControllerModel, private _context: vscode.ExtensionContext, private _treeDataProvider: AzureArcTreeDataProvider) {
 		super(model.info.url, vscode.TreeItemCollapsibleState.Collapsed, ResourceType.dataControllers);
 		model.onRegistrationsUpdated(registrations => this.refreshChildren(registrations));
 	}
@@ -51,21 +53,52 @@ export class ControllerTreeNode extends TreeNode {
 	}
 
 	private refreshChildren(registrations: Registration[]): void {
-		this._children = <TreeNode[]>registrations.map(registration => {
+		const newChildren: ResourceTreeNode[] = [];
+		registrations.forEach(registration => {
 			if (!registration.instanceNamespace || !registration.instanceName) {
 				console.warn('Registration is missing required namespace and name values, skipping');
-				return undefined;
+				return;
 			}
-			switch (registration.instanceType) {
-				case ResourceType.postgresInstances:
-					const postgresModel = new PostgresModel(this.model.info.url, this.model.auth!, registration.instanceNamespace, parseInstanceName(registration.instanceName));
-					return new PostgresTreeNode(postgresModel, this.model, this._context);
-				case ResourceType.sqlManagedInstances:
-					const miaaModel = new MiaaModel(this.model.info.url, this.model.auth!, registration.instanceNamespace, parseInstanceName(registration.instanceName));
-					return new MiaaTreeNode(miaaModel, this.model);
+
+			const resourceInfo: ResourceInfo = {
+				namespace: registration.instanceNamespace,
+				name: parseInstanceName(registration.instanceName),
+				resourceType: registration.instanceType ?? ''
+			};
+
+			let node = this._children.find(n =>
+				n.model?.info?.name === resourceInfo.name &&
+				n.model?.info?.namespace === resourceInfo.namespace &&
+				n.model?.info?.resourceType === resourceInfo.resourceType);
+
+			// If we don't have this child already then create a new node for it
+			if (!node) {
+				// If we had a stored connectionId copy that over
+				resourceInfo.connectionId = this.model.info.resources.find(info =>
+					info.namespace === resourceInfo.namespace &&
+					info.name === resourceInfo.name &&
+					info.resourceType === resourceInfo.resourceType)?.connectionId;
+
+				switch (registration.instanceType) {
+					case ResourceType.postgresInstances:
+						const postgresModel = new PostgresModel(this.model.info.url, this.model.auth!, resourceInfo, registration);
+						node = new PostgresTreeNode(postgresModel, this.model, this._context);
+						break;
+					case ResourceType.sqlManagedInstances:
+						const miaaModel = new MiaaModel(this.model.info.url, this.model.auth!, resourceInfo, registration, this._treeDataProvider);
+						node = new MiaaTreeNode(miaaModel, this.model);
+						break;
+				}
 			}
-			return undefined;
-		}).filter(item => item); // filter out invalid nodes (controllers or ones without required properties)
+			if (node) {
+				newChildren.push(node);
+			}
+		});
+		this._children = newChildren;
+
+		// Update our model info too
+		this.model.info.resources = <ResourceInfo[]>this._children.map(c => c.model?.info).filter(c => c);
+		this._treeDataProvider.saveControllers();
 		this._childrenRefreshPromise.resolve();
 	}
 }

extensions/arc/src/ui/tree/miaaTreeNode.ts

@@ -15,14 +15,14 @@ import { MiaaDashboard } from '../dashboards/miaa/miaaDashboard';
  */
 export class MiaaTreeNode extends TreeNode {
 
-	constructor(private _model: MiaaModel, private _controllerModel: ControllerModel) {
-		super(_model.name, vscode.TreeItemCollapsibleState.None, ResourceType.sqlManagedInstances);
+	constructor(public model: MiaaModel, private _controllerModel: ControllerModel) {
+		super(model.info.name, vscode.TreeItemCollapsibleState.None, ResourceType.sqlManagedInstances);
 	}
 
 	public async openDashboard(): Promise<void> {
-		const miaaDashboard = new MiaaDashboard(this._controllerModel, this._model);
+		const miaaDashboard = new MiaaDashboard(this._controllerModel, this.model);
 		await Promise.all([
 			miaaDashboard.showDashboard(),
-			this._model.refresh()]);
+			this.model.refresh()]);
 	}
 }

extensions/arc/src/ui/tree/postgresTreeNode.ts

@@ -5,18 +5,18 @@
 
 import * as vscode from 'vscode';
 import { ResourceType } from '../../constants';
-import { TreeNode } from './treeNode';
 import { PostgresModel } from '../../models/postgresModel';
 import { ControllerModel } from '../../models/controllerModel';
 import { PostgresDashboard } from '../dashboards/postgres/postgresDashboard';
+import { ResourceTreeNode } from './resourceTreeNode';
 
 /**
  * The TreeNode for displaying an Postgres Server group
  */
-export class PostgresTreeNode extends TreeNode {
+export class PostgresTreeNode extends ResourceTreeNode {
 
 	constructor(private _model: PostgresModel, private _controllerModel: ControllerModel, private _context: vscode.ExtensionContext) {
-		super(_model.name, vscode.TreeItemCollapsibleState.None, ResourceType.postgresInstances);
+		super(_model.name, vscode.TreeItemCollapsibleState.None, ResourceType.postgresInstances, _model);
 	}
 
 	public async openDashboard(): Promise<void> {

extensions/arc/src/ui/tree/resourceTreeNode.ts

@@ -0,0 +1,17 @@
+/*---------------------------------------------------------------------------------------------
+ *  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 { ResourceModel } from '../../models/resourceModel';
+import { TreeNode } from './treeNode';
+
+/**
+ * A TreeNode belonging to a child of a Controller
+ */
+export abstract class ResourceTreeNode extends TreeNode {
+	constructor(label: string, collapsibleState: vscode.TreeItemCollapsibleState, resourceType?: string, public model?: ResourceModel) {
+		super(label, collapsibleState, resourceType);
+	}
+}