From 014bca031c5b70a29fd59865c26f3374ee64be81 Mon Sep 17 00:00:00 2001
From: Anthony Dresser <andresse@microsoft.com>
Date: Tue, 3 Jul 2018 11:49:57 -0700
Subject: [PATCH] add logic to clean up providers when appropriate (#1824)

---
 .../connectionDialogService.ts                | 23 ++++++-------
 .../capabilities/capabilitiesService.ts       | 32 +++++++++++++++++--
 2 files changed, 40 insertions(+), 15 deletions(-)

diff --git a/src/sql/parts/connection/connectionDialog/connectionDialogService.ts b/src/sql/parts/connection/connectionDialog/connectionDialogService.ts
index 06e66d1eba..f547dd8942 100644
--- a/src/sql/parts/connection/connectionDialog/connectionDialogService.ts
+++ b/src/sql/parts/connection/connectionDialog/connectionDialogService.ts
@@ -280,26 +280,23 @@ export class ConnectionDialogService implements IConnectionDialogService {
 
 		return new Promise<void>((resolve, reject) => {
 			// only create the provider maps first time the dialog gets called
-			let capabilitiesPromise: Promise<void> = Promise.resolve();
 			if (this._providerTypes.length === 0) {
 				entries(this._capabilitiesService.providers).forEach(p => {
 					this._providerTypes.push(p[1].connection.displayName);
 					this._providerNameToDisplayNameMap[p[0]] = p[1].connection.displayName;
 				});
 			}
-			capabilitiesPromise.then(s => {
-				this.updateModelServerCapabilities(model);
-				// If connecting from a query editor set "save connection" to false
-				if (params && params.input && params.connectionType === ConnectionType.editor) {
-					this._model.saveProfile = false;
-				}
+			this.updateModelServerCapabilities(model);
+			// If connecting from a query editor set "save connection" to false
+			if (params && params.input && params.connectionType === ConnectionType.editor) {
+				this._model.saveProfile = false;
+			}
 
-				resolve(this.showDialogWithModel().then(() => {
-					if (connectionResult && connectionResult.errorMessage) {
-						this.showErrorDialog(Severity.Error, this._connectionErrorTitle, connectionResult.errorMessage, connectionResult.callStack);
-					}
-				}));
-			}, e => reject(e));
+			resolve(this.showDialogWithModel().then(() => {
+				if (connectionResult && connectionResult.errorMessage) {
+					this.showErrorDialog(Severity.Error, this._connectionErrorTitle, connectionResult.errorMessage, connectionResult.callStack);
+				}
+			}));
 		});
 	}
 
diff --git a/src/sql/services/capabilities/capabilitiesService.ts b/src/sql/services/capabilities/capabilitiesService.ts
index 7db1917b17..a5298908de 100644
--- a/src/sql/services/capabilities/capabilitiesService.ts
+++ b/src/sql/services/capabilities/capabilitiesService.ts
@@ -20,6 +20,9 @@ import { createDecorator } from 'vs/platform/instantiation/common/instantiation'
 import { IDisposable, dispose, Disposable } from 'vs/base/common/lifecycle';
 import { IStorageService } from 'vs/platform/storage/common/storage';
 import { Registry } from 'vs/platform/registry/common/platform';
+import { IExtensionManagementService } from 'vs/platform/extensionManagement/common/extensionManagement';
+import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
+import { getIdFromLocalExtensionId } from 'vs/platform/extensionManagement/common/extensionManagementUtil';
 
 export const SERVICE_ID = 'capabilitiesService';
 export const HOST_NAME = 'sqlops';
@@ -101,7 +104,9 @@ export class CapabilitiesService extends Disposable implements ICapabilitiesServ
 	public readonly onCapabilitiesRegistered = this._onCapabilitiesRegistered.event;
 
 	constructor(
-		@IStorageService private _storageService: IStorageService
+		@IStorageService private _storageService: IStorageService,
+		@IExtensionService extensionService: IExtensionService,
+		@IExtensionManagementService extentionManagementService: IExtensionManagementService
 	) {
 		super();
 
@@ -114,12 +119,35 @@ export class CapabilitiesService extends Disposable implements ICapabilitiesServ
 			this.handleConnectionProvider({ id: v[0], properties: v[1] });
 		});
 		// register for when new extensions are added
-		connectionRegistry.onNewProvider(this.handleConnectionProvider, this);
+		this._register(connectionRegistry.onNewProvider(this.handleConnectionProvider, this));
 
 		// handle adding already known capabilities (could have caching problems)
 		Object.entries(this.capabilities.connectionProviderCache).map(v => {
 			this.handleConnectionProvider({ id: v[0], properties: v[1] }, false);
 		});
+
+		extensionService.whenInstalledExtensionsRegistered().then(() => {
+			this.cleanupProviders();
+		});
+
+		this._register(extentionManagementService.onDidUninstallExtension(({ identifier }) => {
+			let extensionid = getIdFromLocalExtensionId(identifier.id);
+			extensionService.getExtensions().then(i => {
+				let extension = i.find(c => c.id === extensionid);
+				let id = extension.contributes['connectionProvider'].providerId;
+				delete this.capabilities.connectionProviderCache[id];
+			});
+		}));
+	}
+
+	private cleanupProviders(): void {
+		let knownProviders = Object.keys(connectionRegistry.providers);
+		for (let key in this.capabilities.connectionProviderCache) {
+			if (!knownProviders.includes(key)) {
+				this._providers.delete(key);
+				delete this.capabilities.connectionProviderCache[key];
+			}
+		}
 	}
 
 	private handleConnectionProvider(e: { id: string, properties: ConnectionProviderProperties }, isNew = true): void {