Add persistence and connect dialog to Arc view (#11014)

* Add controller persistence and info prompting

* more stuff

* clean up

* Add arc tests to scripts

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

@@ -0,0 +1,110 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as azdata from 'azdata';
+import * as vscode from 'vscode';
+import * as loc from '../../localizedConstants';
+import { AzureArcTreeDataProvider } from '../tree/azureArcTreeDataProvider';
+import { ControllerModel, ControllerInfo } from '../../models/controllerModel';
+import { Deferred } from '../../common/promise';
+
+export type ConnectToControllerDialogModel = { controllerModel: ControllerModel, password: string };
+
+export class ConnectToControllerDialog {
+	private modelBuilder!: azdata.ModelBuilder;
+
+	private urlInputBox!: azdata.InputBoxComponent;
+	private usernameInputBox!: azdata.InputBoxComponent;
+	private passwordInputBox!: azdata.InputBoxComponent;
+	private rememberPwCheckBox!: azdata.CheckBoxComponent;
+
+	private _completionPromise = new Deferred<ConnectToControllerDialogModel | undefined>();
+
+	constructor(private _treeDataProvider: AzureArcTreeDataProvider) { }
+
+	public showDialog(controllerInfo?: ControllerInfo): void {
+		const dialog = azdata.window.createModelViewDialog(loc.connectToController);
+		dialog.cancelButton.onClick(() => this.handleCancel());
+		dialog.registerContent(async view => {
+			this.modelBuilder = view.modelBuilder;
+
+			this.urlInputBox = this.modelBuilder.inputBox()
+				.withProperties<azdata.InputBoxProperties>({
+					value: controllerInfo?.url,
+					// If we have a model then we're editing an existing connection so don't let them modify the URL
+					readOnly: !!controllerInfo
+				}).component();
+			this.usernameInputBox = this.modelBuilder.inputBox()
+				.withProperties<azdata.InputBoxProperties>({
+					value: controllerInfo?.username
+				}).component();
+			this.passwordInputBox = this.modelBuilder.inputBox()
+				.withProperties<azdata.InputBoxProperties>({
+					inputType: 'password',
+				})
+				.component();
+			this.rememberPwCheckBox = this.modelBuilder.checkBox()
+				.withProperties<azdata.CheckBoxProperties>({
+					label: loc.rememberPassword,
+					checked: controllerInfo?.rememberPassword
+				}).component();
+
+			let formModel = this.modelBuilder.formContainer()
+				.withFormItems([{
+					components: [
+						{
+							component: this.urlInputBox,
+							title: loc.controllerUrl,
+							required: true
+						}, {
+							component: this.usernameInputBox,
+							title: loc.username,
+							required: true
+						}, {
+							component: this.passwordInputBox,
+							title: loc.password,
+							required: true
+						}, {
+							component: this.rememberPwCheckBox,
+							title: ''
+						}
+					],
+					title: ''
+				}]).withLayout({ width: '100%' }).component();
+			await view.initializeModel(formModel);
+			this.urlInputBox.focus();
+		});
+
+		dialog.registerCloseValidator(async () => await this.validate());
+		dialog.okButton.label = loc.connect;
+		dialog.cancelButton.label = loc.cancel;
+		azdata.window.openDialog(dialog);
+	}
+
+	private async validate(): Promise<boolean> {
+		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 controllerModel = new ControllerModel(this._treeDataProvider, controllerInfo, this.passwordInputBox.value);
+		try {
+			// Validate that we can connect to the controller
+			await controllerModel.refresh();
+		} catch (err) {
+			vscode.window.showErrorMessage(loc.connectToControllerFailed(this.urlInputBox.value, err));
+			return false;
+		}
+		this._completionPromise.resolve({ controllerModel: controllerModel, password: this.passwordInputBox.value });
+		return true;
+	}
+
+	private handleCancel(): void {
+		this._completionPromise.resolve(undefined);
+	}
+
+	public waitForClose(): Promise<ConnectToControllerDialogModel | undefined> {
+		return this._completionPromise.promise;
+	}
+}

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

@@ -0,0 +1,113 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the Source EULA. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as azdata from 'azdata';
+import * as vscode from 'vscode';
+import { ControllerTreeNode } from './controllerTreeNode';
+import { TreeNode } from './treeNode';
+import { LoadingControllerNode as LoadingTreeNode } from './loadingTreeNode';
+import { ControllerModel, ControllerInfo } from '../../models/controllerModel';
+
+const mementoToken = 'arcControllers';
+
+/**
+ * The TreeDataProvider for the Azure Arc view, which displays a list of registered
+ * controllers and the resources under them.
+ */
+export class AzureArcTreeDataProvider implements vscode.TreeDataProvider<TreeNode> {
+
+	private _credentialsProvider = azdata.credentials.getProvider('arcControllerPasswords');
+	private _onDidChangeTreeData: vscode.EventEmitter<TreeNode | undefined> = new vscode.EventEmitter<TreeNode | undefined>();
+	readonly onDidChangeTreeData: vscode.Event<TreeNode | undefined> = this._onDidChangeTreeData.event;
+
+	private _loading: boolean = true;
+	private _loadingNode = new LoadingTreeNode();
+
+	private _controllerNodes: ControllerTreeNode[] = [];
+
+	constructor(private _context: vscode.ExtensionContext) {
+		this.loadSavedControllers().catch(err => console.log(`Error loading saved Arc controllers ${err}`));
+	}
+
+	public async getChildren(element?: TreeNode): Promise<TreeNode[]> {
+		if (this._loading) {
+			return [this._loadingNode];
+		}
+
+		if (element) {
+			return element.getChildren();
+		} else {
+			return this._controllerNodes;
+		}
+	}
+
+	public getTreeItem(element: TreeNode): TreeNode | Thenable<TreeNode> {
+		return element;
+	}
+
+	public async addOrUpdateController(model: ControllerModel, password: string, refreshTree = true): Promise<void> {
+		const controllerNode = this._controllerNodes.find(node => model.equals(node.model));
+		if (controllerNode) {
+			controllerNode.model.info = model.info;
+		} else {
+			this._controllerNodes.push(new ControllerTreeNode(model, this._context));
+		}
+		await this.updatePassword(model, password);
+		if (refreshTree) {
+			this._onDidChangeTreeData.fire(undefined);
+		}
+		await this.saveControllers();
+	}
+
+	public async removeController(controllerNode: ControllerTreeNode): Promise<void> {
+		this._controllerNodes = this._controllerNodes.filter(node => node !== controllerNode);
+		this._onDidChangeTreeData.fire(undefined);
+		await this.saveControllers();
+	}
+
+	public async getPassword(info: ControllerInfo): Promise<string> {
+		const provider = await this._credentialsProvider;
+		const credential = await provider.readCredential(getCredentialId(info));
+		return credential.password;
+	}
+
+	/**
+	 * Refreshes the specified node, or the entire tree if node is undefined
+	 * @param node The node to refresh, or undefined for the whole tree
+	 */
+	public refreshNode(node: TreeNode | undefined): void {
+		this._onDidChangeTreeData.fire(node);
+	}
+
+	private async updatePassword(model: ControllerModel, password: string): Promise<void> {
+		const provider = await this._credentialsProvider;
+		if (model.info.rememberPassword) {
+			provider.saveCredential(getCredentialId(model.info), password);
+		} else {
+			provider.deleteCredential(getCredentialId(model.info));
+		}
+	}
+
+	private async loadSavedControllers(): Promise<void> {
+		try {
+			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);
+			});
+		} finally {
+			this._loading = false;
+			this._onDidChangeTreeData.fire(undefined);
+		}
+	}
+
+	private async saveControllers(): Promise<void> {
+		await this._context.globalState.update(mementoToken, this._controllerNodes.map(node => node.model.info));
+	}
+}
+
+function getCredentialId(info: ControllerInfo): string {
+	return `${info.url}::${info.username}`;
+}

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

@@ -1,59 +0,0 @@
-/*---------------------------------------------------------------------------------------------
- *  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 { ControllerTreeNode } from './controllerTreeNode';
-import { TreeNode } from './treeNode';
-import { LoadingControllerNode as LoadingTreeNode } from './loadingTreeNode';
-import { ControllerModel } from '../../models/controllerModel';
-
-/**
- * The TreeDataProvider for the Azure Arc view, which displays a list of registered
- * controllers and the resources under them.
- */
-export class AzureArcTreeDataProvider implements vscode.TreeDataProvider<TreeNode> {
-
-	private _onDidChangeTreeData: vscode.EventEmitter<TreeNode | undefined> = new vscode.EventEmitter<TreeNode | undefined>();
-	readonly onDidChangeTreeData: vscode.Event<TreeNode | undefined> = this._onDidChangeTreeData.event;
-
-	private _loading: boolean = true;
-	private _loadingNode = new LoadingTreeNode();
-
-	private _controllerNodes: ControllerTreeNode[] = [];
-
-	constructor(private _context: vscode.ExtensionContext) {
-		// TODO: 
-		setTimeout(() => {
-			this._loading = false;
-			this._onDidChangeTreeData.fire(undefined);
-		}, 5000);
-	}
-
-	public async getChildren(element?: TreeNode): Promise<TreeNode[]> {
-		if (this._loading) {
-			return [this._loadingNode];
-		}
-
-		if (element) {
-			return element.getChildren();
-		} else {
-			return this._controllerNodes;
-		}
-	}
-
-	public getTreeItem(element: TreeNode): TreeNode | Thenable<TreeNode> {
-		return element;
-	}
-
-	public addController(model: ControllerModel): void {
-		this._controllerNodes.push(new ControllerTreeNode(model, this._context));
-		this._onDidChangeTreeData.fire(undefined);
-	}
-
-	public removeController(controllerNode: ControllerTreeNode): void {
-		this._controllerNodes = this._controllerNodes.filter(node => node !== controllerNode);
-		this._onDidChangeTreeData.fire(undefined);
-	}
-}

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

@@ -13,6 +13,8 @@ import { ControllerDashboard } from '../dashboards/controller/controllerDashboar
 import { PostgresModel } from '../../models/postgresModel';
 import { parseInstanceName } from '../../common/utils';
 import { MiaaModel } from '../../models/miaaModel';
+import { Deferred } from '../../common/promise';
+import { RefreshTreeNode } from './refreshTreeNode';
 
 /**
  * The TreeNode for displaying an Azure Arc Controller
@@ -20,19 +22,31 @@ import { MiaaModel } from '../../models/miaaModel';
 export class ControllerTreeNode extends TreeNode {
 
 	private _children: TreeNode[] = [];
+	private _childrenRefreshPromise = new Deferred();
 
-	constructor(private _model: ControllerModel, private _context: vscode.ExtensionContext) {
-		super(_model.controllerUrl, vscode.TreeItemCollapsibleState.Collapsed, ResourceType.dataControllers);
-		_model.onRegistrationsUpdated(registrations => this.refreshChildren(registrations));
-		_model.refresh().catch(err => console.log(`Error refreshing Arc Controller model for tree node : ${err}`));
+	constructor(public model: ControllerModel, private _context: vscode.ExtensionContext) {
+		super(model.info.url, vscode.TreeItemCollapsibleState.Collapsed, ResourceType.dataControllers);
+		model.onRegistrationsUpdated(registrations => this.refreshChildren(registrations));
 	}
 
 	public async getChildren(): Promise<TreeNode[]> {
+		// First reset our deferred promise so we're sure we'll get the refreshed children
+		this._childrenRefreshPromise = new Deferred();
+		try {
+			await this.model.refresh();
+			await this._childrenRefreshPromise.promise;
+		} catch (err) {
+			// Couldn't get the children and TreeView doesn't have a way to collapse a node
+			// in a way that will refetch its children when expanded again so instead we
+			// display a tempory node that will prompt the user to re-enter credentials
+			return [new RefreshTreeNode(this)];
+		}
+
 		return this._children;
 	}
 
 	public async openDashboard(): Promise<void> {
-		const controllerDashboard = new ControllerDashboard(this._model);
+		const controllerDashboard = new ControllerDashboard(this.model);
 		await controllerDashboard.showDashboard();
 	}
 
@@ -44,13 +58,14 @@ export class ControllerTreeNode extends TreeNode {
 			}
 			switch (registration.instanceType) {
 				case ResourceType.postgresInstances:
-					const postgresModel = new PostgresModel(this._model.controllerUrl, this._model.auth, registration.instanceNamespace, parseInstanceName(registration.instanceName));
-					return new PostgresTreeNode(postgresModel, this._model, this._context);
+					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.controllerUrl, this._model.auth, registration.instanceNamespace, parseInstanceName(registration.instanceName));
-					return new MiaaTreeNode(miaaModel, this._model);
+					const miaaModel = new MiaaModel(this.model.info.url, this.model.auth!, registration.instanceNamespace, parseInstanceName(registration.instanceName));
+					return new MiaaTreeNode(miaaModel, this.model);
 			}
 			return undefined;
 		}).filter(item => item); // filter out invalid nodes (controllers or ones without required properties)
+		this._childrenRefreshPromise.resolve();
 	}
 }

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

@@ -0,0 +1,25 @@
+/*---------------------------------------------------------------------------------------------
+ *  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 loc from '../../localizedConstants';
+import { TreeNode } from './treeNode';
+import { refreshActionId } from '../../constants';
+
+/**
+ * A placeholder TreeNode to display when credentials weren't entered
+ */
+export class RefreshTreeNode extends TreeNode {
+
+	constructor(private _parent: TreeNode) {
+		super(loc.refreshToEnterCredentials, vscode.TreeItemCollapsibleState.None, 'refresh');
+	}
+
+	public command: vscode.Command = {
+		command: refreshActionId,
+		title: loc.refreshToEnterCredentials,
+		arguments: [this._parent]
+	};
+}