From 34273907f71473b90418de146485de55b2df7e91 Mon Sep 17 00:00:00 2001
From: Gene Lee <gelee@microsoft.com>
Date: Tue, 28 May 2019 14:45:28 -0700
Subject: [PATCH] Made connection icon to be remembered by memento in
 ConnectionManagementService (#5523)

---
 .../common/connectionManagementService.ts     | 30 +++++++++++++++----
 .../connectionManagementService.test.ts       |  1 +
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/src/sql/platform/connection/common/connectionManagementService.ts b/src/sql/platform/connection/common/connectionManagementService.ts
index e59ceff49c..bcd7a50f48 100644
--- a/src/sql/platform/connection/common/connectionManagementService.ts
+++ b/src/sql/platform/connection/common/connectionManagementService.ts
@@ -49,6 +49,8 @@ import { IConnectionDialogService } from 'sql/workbench/services/connection/comm
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { ILogService } from 'vs/platform/log/common/log';
 import * as interfaces from './interfaces';
+import { IStorageService, StorageScope } from 'vs/platform/storage/common/storage';
+import { Memento } from 'vs/workbench/common/memento';
 
 export class ConnectionManagementService extends Disposable implements IConnectionManagementService {
 
@@ -56,7 +58,6 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 
 	private _providers = new Map<string, { onReady: Thenable<azdata.ConnectionProvider>, properties: ConnectionProviderProperties }>();
 	private _iconProviders = new Map<string, azdata.IconProvider>();
-	private _connectionIconIdCache = new Map<string, string>();
 
 	private _uriToProvider: { [uri: string]: string; } = Object.create(null);
 
@@ -71,6 +72,10 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 	private _onLanguageFlavorChanged = new Emitter<azdata.DidChangeLanguageFlavorParams>();
 	private _connectionGlobalStatus = new ConnectionGlobalStatus(this._statusBarService);
 
+	private _mementoContext: Memento;
+	private _mementoObj: any;
+	private static readonly CONNECTION_MEMENTO = 'ConnectionManagement';
+
 	constructor(
 		private _connectionStore: ConnectionStore,
 		@IConnectionDialogService private _connectionDialogService: IConnectionDialogService,
@@ -85,7 +90,8 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 		@IResourceProviderService private _resourceProviderService: IResourceProviderService,
 		@IAngularEventingService private _angularEventing: IAngularEventingService,
 		@IAccountManagementService private _accountManagementService: IAccountManagementService,
-		@ILogService private logService: ILogService
+		@ILogService private logService: ILogService,
+		@IStorageService private _storageService: IStorageService
 	) {
 		super();
 
@@ -93,6 +99,11 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 			this._connectionStore = _instantiationService.createInstance(ConnectionStore);
 		}
 
+		if (this._storageService) {
+			this._mementoContext = new Memento(ConnectionManagementService.CONNECTION_MEMENTO, this._storageService);
+			this._mementoObj = this._mementoContext.getMemento(StorageScope.GLOBAL);
+		}
+
 		const registry = platform.Registry.as<IConnectionProviderRegistry>(ConnectionProviderExtensions.ConnectionProviderContributions);
 
 		let providerRegistration = (p: { id: string, properties: ConnectionProviderProperties }) => {
@@ -551,15 +562,24 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 			let serverInfo: azdata.ServerInfo = this.getServerInfo(connectionProfile.id);
 			let profile: interfaces.IConnectionProfile = connectionProfile.toIConnectionProfile();
 			iconProvider.getConnectionIconId(profile, serverInfo).then(iconId => {
-				if (iconId) {
-					this._connectionIconIdCache.set(connectionProfile.id, iconId);
+				if (iconId && this._mementoObj && this._mementoContext) {
+					if (!this._mementoObj.CONNECTION_ICON_ID) {
+						this._mementoObj.CONNECTION_ICON_ID = <any>{};
+					}
+					if (this._mementoObj.CONNECTION_ICON_ID[connectionProfile.id] !== iconId) {
+						this._mementoObj.CONNECTION_ICON_ID[connectionProfile.id] = iconId;
+						this._mementoContext.saveMemento();
+					}
 				}
 			});
 		}
 	}
 
 	public getConnectionIconId(connectionId: string): string {
-		return this._connectionIconIdCache.get(connectionId);
+		if (!connectionId || !this._mementoObj || !this._mementoObj.CONNECTION_ICON_ID) {
+			return undefined;
+		}
+		return this._mementoObj.CONNECTION_ICON_ID[connectionId];
 	}
 
 	public showDashboard(connection: IConnectionProfile): Thenable<boolean> {
diff --git a/src/sqltest/parts/connection/connectionManagementService.test.ts b/src/sqltest/parts/connection/connectionManagementService.test.ts
index 9ab1b494a7..30660840ce 100644
--- a/src/sqltest/parts/connection/connectionManagementService.test.ts
+++ b/src/sqltest/parts/connection/connectionManagementService.test.ts
@@ -162,6 +162,7 @@ suite('SQL ConnectionManagementService tests', () => {
 			resourceProviderStubMock.object,
 			undefined,
 			accountManagementService.object,
+			undefined,
 			undefined
 		);
 		return connectionManagementService;