From 97a4de411190bf531e2e0221d48354542f7ac72c Mon Sep 17 00:00:00 2001
From: Charles Gagnon <chgagnon@microsoft.com>
Date: Mon, 7 Dec 2020 10:27:37 -0800
Subject: [PATCH] Have resource deployment providers return disposables
 (#13690)

* Add dependent field provider to resource deployment

* Change name to value provider service

* Add error handling

* providerId -> id

* Set dropdown value correctly

* missed id

* back to providerId

* fix updating missed id

* Make resource deployment providers disposable
---
 extensions/arc/src/extension.ts                        |  2 +-
 extensions/azdata/src/extension.ts                     |  2 +-
 extensions/azurecore/src/extension.ts                  |  4 ++--
 .../src/services/optionSourcesService.ts               | 10 +++++++++-
 .../src/services/valueProviderService.ts               | 10 +++++++++-
 .../src/typings/resource-deployment.d.ts               |  5 +++--
 6 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/extensions/arc/src/extension.ts b/extensions/arc/src/extension.ts
index fccd30eb2f..dffb4ba905 100644
--- a/extensions/arc/src/extension.ts
+++ b/extensions/arc/src/extension.ts
@@ -67,7 +67,7 @@ export async function activate(context: vscode.ExtensionContext): Promise<arc.IE
 
 	// register option sources
 	const rdApi = <rd.IExtension>vscode.extensions.getExtension(rd.extension.name)?.exports;
-	rdApi.registerOptionsSourceProvider(new ArcControllersOptionsSourceProvider(treeDataProvider));
+	context.subscriptions.push(rdApi.registerOptionsSourceProvider(new ArcControllersOptionsSourceProvider(treeDataProvider)));
 
 	return arcApi(treeDataProvider);
 }
diff --git a/extensions/azdata/src/extension.ts b/extensions/azdata/src/extension.ts
index 8697e3e334..57bea887ad 100644
--- a/extensions/azdata/src/extension.ts
+++ b/extensions/azdata/src/extension.ts
@@ -65,7 +65,7 @@ export async function activate(context: vscode.ExtensionContext): Promise<azdata
 
 	// register option source(s)
 	const rdApi = <rd.IExtension>vscode.extensions.getExtension(rd.extension.name)?.exports;
-	rdApi.registerOptionsSourceProvider(new ArcControllerConfigProfilesOptionsSource(azdataApi));
+	context.subscriptions.push(rdApi.registerOptionsSourceProvider(new ArcControllerConfigProfilesOptionsSource(azdataApi)));
 
 	return azdataApi;
 }
diff --git a/extensions/azurecore/src/extension.ts b/extensions/azurecore/src/extension.ts
index ebaedf0cd2..55c53cf80b 100644
--- a/extensions/azurecore/src/extension.ts
+++ b/extensions/azurecore/src/extension.ts
@@ -109,7 +109,7 @@ export async function activate(context: vscode.ExtensionContext): Promise<azurec
 	// Don't block on this since there's a bit of a circular dependency here with the extension activation since resource deployment
 	// depends on this extension too. It's fine to wait a bit for that to finish before registering the provider
 	vscode.extensions.getExtension(resourceDeployment.extension.name).activate().then((api: resourceDeployment.IExtension) => {
-		api.registerValueProvider({
+		context.subscriptions.push(api.registerValueProvider({
 			id: 'subscription-id-to-tenant-id',
 			getValue: async (triggerValue: string) => {
 				if (triggerValue === '') {
@@ -137,7 +137,7 @@ export async function activate(context: vscode.ExtensionContext): Promise<azurec
 				console.error(`Unable to find subscription with ID ${triggerValue} when mapping subscription ID to tenant ID`);
 				return '';
 			}
-		});
+		}));
 	});
 
 	return {
diff --git a/extensions/resource-deployment/src/services/optionSourcesService.ts b/extensions/resource-deployment/src/services/optionSourcesService.ts
index dc6cf77b6a..3c6e12d276 100644
--- a/extensions/resource-deployment/src/services/optionSourcesService.ts
+++ b/extensions/resource-deployment/src/services/optionSourcesService.ts
@@ -3,16 +3,24 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import * as vscode from 'vscode';
 import * as rd from 'resource-deployment';
 import * as loc from '../localizedConstants';
 
 class OptionsSourcesService {
 	private _optionsSourceStore = new Map<string, rd.IOptionsSourceProvider>();
-	registerOptionsSourceProvider(provider: rd.IOptionsSourceProvider): void {
+	registerOptionsSourceProvider(provider: rd.IOptionsSourceProvider): vscode.Disposable {
 		if (this._optionsSourceStore.has(provider.id)) {
 			throw new Error(loc.optionsSourceAlreadyDefined(provider.id));
 		}
 		this._optionsSourceStore.set(provider.id, provider);
+		return {
+			dispose: () => this.unregisterOptionsSourceProvider(provider.id)
+		};
+	}
+
+	private unregisterOptionsSourceProvider(providerId: string): void {
+		this._optionsSourceStore.delete(providerId);
 	}
 
 	getOptionsSource(optionsSourceProviderId: string): rd.IOptionsSourceProvider {
diff --git a/extensions/resource-deployment/src/services/valueProviderService.ts b/extensions/resource-deployment/src/services/valueProviderService.ts
index 152b9a89d6..d447200f12 100644
--- a/extensions/resource-deployment/src/services/valueProviderService.ts
+++ b/extensions/resource-deployment/src/services/valueProviderService.ts
@@ -3,16 +3,24 @@
  *  Licensed under the Source EULA. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import * as vscode from 'vscode';
 import * as rd from 'resource-deployment';
 import * as loc from '../localizedConstants';
 
 class ValueProviderService {
 	private _valueProviderStore = new Map<string, rd.IValueProvider>();
-	registerValueProvider(provider: rd.IValueProvider): void {
+	registerValueProvider(provider: rd.IValueProvider): vscode.Disposable {
 		if (this._valueProviderStore.has(provider.id)) {
 			throw new Error(loc.valueProviderAlreadyDefined(provider.id));
 		}
 		this._valueProviderStore.set(provider.id, provider);
+		return {
+			dispose: () => this.unregisterValueProvider(provider.id)
+		};
+	}
+
+	private unregisterValueProvider(providerId: string): void {
+		this._valueProviderStore.delete(providerId);
 	}
 
 	getValueProvider(providerId: string): rd.IValueProvider {
diff --git a/extensions/resource-deployment/src/typings/resource-deployment.d.ts b/extensions/resource-deployment/src/typings/resource-deployment.d.ts
index def6463787..01e4c5a156 100644
--- a/extensions/resource-deployment/src/typings/resource-deployment.d.ts
+++ b/extensions/resource-deployment/src/typings/resource-deployment.d.ts
@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 declare module 'resource-deployment' {
 	import * as azdata from 'azdata';
+	import * as vscode from 'vscode';
 
 	export const enum ErrorType {
 		userCancelled,
@@ -36,7 +37,7 @@ declare module 'resource-deployment' {
 	 */
 
 	export interface IExtension {
-		registerOptionsSourceProvider(provider: IOptionsSourceProvider): void,
-		registerValueProvider(provider: IValueProvider): void
+		registerOptionsSourceProvider(provider: IOptionsSourceProvider): vscode.Disposable,
+		registerValueProvider(provider: IValueProvider): vscode.Disposable
 	}
 }