Accounts: Enable notification for accounts change (#2432)

* Enable notification for accounts change

* Fixed bugs in extHostAccountManagement.test.ts

* Fixed as commented:
1. Removed AccountWithProviderHandle
2. Use a private dictionary _accounts in ExtHostAccountManagement to cache all accounts and corresponding provider handles.
3. getSecurityToken gets provider handle from _accounts for specified account.
4. Added / changed unit tests for getAllAccounts & getSecurityToken

src/sql/sqlops.d.ts

@@ -1855,14 +1855,19 @@ declare module 'sqlops' {
 		 * Gets all added accounts.
 		 * @returns {Thenable<Account>} Promise to return the accounts
 		 */
-		export function getAllAccounts(): Thenable<AccountWithProviderHandle[]>;
+		export function getAllAccounts(): Thenable<Account[]>;
 
 		/**
 		 * Generates a security token by asking the account's provider
 		 * @param {Account} account Account to generate security token for
 		 * @return {Thenable<{}>} Promise to return the security token
 		 */
-		export function getSecurityToken(account: AccountWithProviderHandle): Thenable<{}>;
+		export function getSecurityToken(account: Account): Thenable<{}>;
+
+		/**
+		 * An [event](#Event) which fires when the accounts have changed.
+		 */
+		export const onDidChangeAccounts: vscode.Event<DidChangeAccountsParams>;
 	}
 
 	/**
@@ -1930,19 +1935,9 @@ declare module 'sqlops' {
 		isStale: boolean;
 	}
 
-	/**
+	export interface DidChangeAccountsParams {
-	 * Represents an account with account provider's handle
+		// Updated accounts
-	 */
+		accounts: Account[];
-	export interface AccountWithProviderHandle {
-		/**
-		 * Account
-		 */
-		account: Account;
-
-		/**
-		 * Account's provider handle
-		 */
-		providerHandle: number;
 	}
 
 	// - ACCOUNT PROVIDER //////////////////////////////////////////////////

src/sql/workbench/api/node/extHostAccountManagement.ts

@@ -14,11 +14,14 @@ import {
 	SqlMainContext,
 } from 'sql/workbench/api/node/sqlExtHost.protocol';
 import { IMainContext } from 'vs/workbench/api/node/extHost.protocol';
+import { Event, Emitter } from 'vs/base/common/event';
 
 export class ExtHostAccountManagement extends ExtHostAccountManagementShape {
 	private _handlePool: number = 0;
 	private _proxy: MainThreadAccountManagementShape;
 	private _providers: { [handle: number]: AccountProviderWithMetadata } = {};
+	private _accounts: { [handle: number]: sqlops.Account[] } = {};
+	private readonly _onDidChangeAccounts = new Emitter<sqlops.DidChangeAccountsParams>();
 
 	constructor(mainContext: IMainContext) {
 		super();
@@ -31,10 +34,6 @@ export class ExtHostAccountManagement extends ExtHostAccountManagementShape {
 		return this._withProvider(handle, (provider: sqlops.AccountProvider) => provider.clear(accountKey));
 	}
 
-	public $getSecurityToken(handle: number, account: sqlops.Account): Thenable<{}> {
-		return this._withProvider(handle, (provider: sqlops.AccountProvider) => provider.getSecurityToken(account));
-	}
-
 	public $initialize(handle: number, restoredAccounts: sqlops.Account[]): Thenable<sqlops.Account[]> {
 		return this._withProvider(handle, (provider: sqlops.AccountProvider) => provider.initialize(restoredAccounts));
 	}
@@ -64,31 +63,49 @@ export class ExtHostAccountManagement extends ExtHostAccountManagementShape {
 		this._proxy.$accountUpdated(updatedAccount);
 	}
 
-	public $getAllAccounts(): Thenable<sqlops.AccountWithProviderHandle[]> {
+	public $getAllAccounts(): Thenable<sqlops.Account[]> {
 		if (Object.keys(this._providers).length === 0) {
 			throw new Error('No account providers registered.');
 		}
 
-		let accountWithProviderHandles: sqlops.AccountWithProviderHandle[] = [];
+		this._accounts = {};
-		let promises: Thenable<void>[] = [];
 
-		for (let providerKey in this._providers) {
+		const resultAccounts: sqlops.Account[] = [];
-			let providerHandle = parseInt(providerKey);
-			let provider = this._providers[providerHandle];
 
+		const promises: Thenable<void>[] = [];
+
+		for (const providerKey in this._providers) {
+			const providerHandle = parseInt(providerKey);
+
+			const provider = this._providers[providerHandle];
 			promises.push(this._proxy.$getAccountsForProvider(provider.metadata.id).then(
 				(accounts) => {
-					accounts.forEach((account) => {
+					this._accounts[providerHandle] = accounts;
-						accountWithProviderHandles.push({
+					resultAccounts.push(...accounts);
-							account: account,
-							providerHandle: providerHandle
-						});
-					});
 				}
 			));
 		}
 
-		return Promise.all(promises).then(() => accountWithProviderHandles);
+		return Promise.all(promises).then(() => resultAccounts);
+	}
+
+	public $getSecurityToken(account: sqlops.Account): Thenable<{}> {
+		for (const handle in this._accounts) {
+			const providerHandle = parseInt(handle);
+			if (this._accounts[handle].findIndex((acct) => acct.key.accountId === account.key.accountId) !== -1) {
+				return this._withProvider(providerHandle, (provider: sqlops.AccountProvider) => provider.getSecurityToken(account));
+			}
+		}
+
+		throw new Error(`Account ${account.key.accountId} not found.`);
+	}
+
+	public get onDidChangeAccounts(): Event<sqlops.DidChangeAccountsParams> {
+		return this._onDidChangeAccounts.event;
+	}
+
+	public $accountsChanged(handle: number, accounts: sqlops.Account[]): Thenable<void> {
+		return this._onDidChangeAccounts.fire({ accounts: accounts });
 	}
 
 	public $registerAccountProvider(providerMetadata: sqlops.AccountProviderMetadata, provider: sqlops.AccountProvider): Disposable {

src/sql/workbench/api/node/mainThreadAccountManagement.ts

@@ -16,6 +16,7 @@ import {
 } from 'sql/workbench/api/node/sqlExtHost.protocol';
 import { IExtHostContext } from 'vs/workbench/api/node/extHost.protocol';
 import { extHostNamedCustomer } from 'vs/workbench/api/electron-browser/extHostCustomers';
+import { UpdateAccountListEventParams } from 'sql/services/accountManagement/eventTypes';
 
 @extHostNamedCustomer(SqlMainContext.MainThreadAccountManagement)
 export class MainThreadAccountManagement implements MainThreadAccountManagementShape {
@@ -32,6 +33,20 @@ export class MainThreadAccountManagement implements MainThreadAccountManagementS
 			this._proxy = extHostContext.getProxy(SqlExtHostContext.ExtHostAccountManagement);
 		}
 		this._toDispose = [];
+
+		this._accountManagementService.updateAccountListEvent((e: UpdateAccountListEventParams) => {
+			if (!e) {
+				return;
+			}
+
+			const providerMetadataIndex = Object.values(this._providerMetadata).findIndex((providerMetadata: sqlops.AccountProviderMetadata) => providerMetadata.id === e.providerId);
+			if (providerMetadataIndex === -1) {
+				return;
+			}
+
+			const providerHandle = parseInt(Object.keys(this._providerMetadata)[providerMetadataIndex]);
+			this._proxy.$accountsChanged(providerHandle, e.accountList);
+		});
 	}
 
 	public $beginAutoOAuthDeviceCode(providerId: string, title: string, message: string, userCode: string, uri: string): Thenable<void> {
@@ -62,7 +77,7 @@ export class MainThreadAccountManagement implements MainThreadAccountManagementS
 				return self._proxy.$clear(handle, accountKey);
 			},
 			getSecurityToken(account: sqlops.Account): Thenable<{}> {
-				return self._proxy.$getSecurityToken(handle, account);
+				return self._proxy.$getSecurityToken(account);
 			},
 			initialize(restoredAccounts: sqlops.Account[]): Thenable<sqlops.Account[]> {
 				return self._proxy.$initialize(handle, restoredAccounts);

src/sql/workbench/api/node/sqlExtHost.api.impl.ts

@@ -92,11 +92,14 @@ export function createApiFactory(
 				accountUpdated(updatedAccount: sqlops.Account): void {
 					return extHostAccountManagement.$accountUpdated(updatedAccount);
 				},
-				getAllAccounts(): Thenable<sqlops.AccountWithProviderHandle[]> {
+				getAllAccounts(): Thenable<sqlops.Account[]> {
 					return extHostAccountManagement.$getAllAccounts();
 				},
-				getSecurityToken(account: sqlops.AccountWithProviderHandle): Thenable<{}> {
+				getSecurityToken(account: sqlops.Account): Thenable<{}> {
-					return extHostAccountManagement.$getSecurityToken(account.providerHandle, account.account);
+					return extHostAccountManagement.$getSecurityToken(account);
+				},
+				onDidChangeAccounts(listener: (e: sqlops.DidChangeAccountsParams) => void, thisArgs?: any, disposables?: extHostTypes.Disposable[]) {
+					return extHostAccountManagement.onDidChangeAccounts(listener, thisArgs, disposables);
 				}
 			};
 

src/sql/workbench/api/node/sqlExtHost.protocol.ts

@@ -27,10 +27,11 @@ import {
 export abstract class ExtHostAccountManagementShape {
 	$autoOAuthCancelled(handle: number): Thenable<void> { throw ni(); }
 	$clear(handle: number, accountKey: sqlops.AccountKey): Thenable<void> { throw ni(); }
-	$getSecurityToken(handle: number, account: sqlops.Account): Thenable<{}> { throw ni(); }
+	$getSecurityToken(account: sqlops.Account): Thenable<{}> { throw ni(); }
 	$initialize(handle: number, restoredAccounts: sqlops.Account[]): Thenable<sqlops.Account[]> { throw ni(); }
 	$prompt(handle: number): Thenable<sqlops.Account> { throw ni(); }
 	$refresh(handle: number, account: sqlops.Account): Thenable<sqlops.Account> { throw ni(); }
+	$accountsChanged(handle: number, accounts: sqlops.Account[]): Thenable<void> { throw ni(); }
 }
 
 export abstract class ExtHostConnectionManagementShape {

src/sqltest/workbench/api/extHostAccountManagement.test.ts

@@ -292,16 +292,7 @@ suite('ExtHostAccountManagement', () => {
 		};
 		let mockAccounts = [mockAccount1, mockAccount2];
 
-		let expectedAccounts = [
+		let expectedAccounts = [mockAccount1, mockAccount2];
-			{
-				account: mockAccount1,
-				providerHandle: 0
-			},
-			{
-				account: mockAccount2,
-				providerHandle: 0
-			}
-		];
 
 		let mockAccountManagementService = getMockAccountManagementService(mockAccounts);
 		instantiationService.stub(IAccountManagementService, mockAccountManagementService.object);
@@ -341,6 +332,110 @@ suite('ExtHostAccountManagement', () => {
 			});
 		done();
 	});
+
+	test('GetSecurityToken - Success', (done) => {
+		let mockAccountProviderMetadata = {
+			id: 'azure',
+			displayName: 'Azure'
+		};
+
+		let mockAccount1 = {
+			key: {
+				providerId: mockAccountProviderMetadata.id,
+				accountId: 'azure_account_1'
+			},
+			displayInfo: {
+				contextualDisplayName: 'Microsoft Account',
+				accountType: 'microsoft',
+				displayName: 'Azure Account 1'
+			},
+			properties: [],
+			isStale: false
+		};
+		let mockAccounts = [mockAccount1];
+
+		let mockAccountManagementService = getMockAccountManagementService(mockAccounts);
+		instantiationService.stub(IAccountManagementService, mockAccountManagementService.object);
+		let accountManagementService = instantiationService.createInstance(MainThreadAccountManagement);
+		threadService.set(SqlMainContext.MainThreadAccountManagement, accountManagementService);
+
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		extHost.$registerAccountProvider(mockAccountProviderMetadata, new AccountProviderStub());
+
+		extHost.$getAllAccounts()
+			.then((accounts) => {
+		    		// If: I get security token
+					extHost.$getSecurityToken(mockAccount1)
+						.then((securityToken) => {
+							// Then: The call should have been passed to the account management service
+							mockAccountManagementService.verify(
+								(obj) => obj.getSecurityToken(TypeMoq.It.isAny()),
+								TypeMoq.Times.once()
+							);
+						}
+					);
+				}
+			).then(() => done(), (err) => done(err));
+	});
+
+	test('GetSecurityToken - Account not found', (done) => {
+		let mockAccountProviderMetadata = {
+			id: 'azure',
+			displayName: 'Azure'
+		};
+
+		let mockAccount1 = {
+			key: {
+				providerId: mockAccountProviderMetadata.id,
+				accountId: 'azure_account_1'
+			},
+			displayInfo: {
+				contextualDisplayName: 'Microsoft Account',
+				accountType: 'microsoft',
+				displayName: 'Azure Account 1'
+			},
+			properties: [],
+			isStale: false
+		};
+		let mockAccounts = [mockAccount1];
+
+		let mockAccountManagementService = getMockAccountManagementService(mockAccounts);
+		instantiationService.stub(IAccountManagementService, mockAccountManagementService.object);
+		let accountManagementService = instantiationService.createInstance(MainThreadAccountManagement);
+		threadService.set(SqlMainContext.MainThreadAccountManagement, accountManagementService);
+
+		// Setup: Create ext host account management with registered account provider
+		let extHost = new ExtHostAccountManagement(threadService);
+		extHost.$registerAccountProvider(mockAccountProviderMetadata, new AccountProviderStub());
+
+		let mockAccount2 = {
+			key: {
+				providerId: mockAccountProviderMetadata.id,
+				accountId: 'azure_account_2'
+			},
+			displayInfo: {
+				contextualDisplayName: 'Work/School Account',
+				accountType: 'microsoft',
+				displayName: 'Azure Account 2'
+			},
+			properties: [],
+			isStale: false
+		};
+
+		extHost.$getAllAccounts().then((accounts) => {
+			// If: I get security token for mockAccount2
+			// Then: It should throw
+			assert.throws(
+				() => extHost.$getSecurityToken(mockAccount2),
+				(error) => {
+					return error.message === `Account ${mockAccount2.key.accountId} not found.`;
+				}
+			);
+		});
+
+		done();
+	});
 });
 
 function getMockAccountProvider(): TypeMoq.Mock<sqlops.AccountProvider> {
@@ -362,6 +457,10 @@ function getMockAccountManagementService(accounts: sqlops.Account[]): TypeMoq.Mo
 
 	mockAccountManagementService.setup(x => x.getAccountsForProvider(TypeMoq.It.isAny()))
 		.returns(() => Promise.resolve(accounts));
+	mockAccountManagementService.setup(x => x.getSecurityToken(TypeMoq.It.isValue(accounts[0])))
+	    .returns(() => Promise.resolve({}));
+	mockAccountManagementService.setup(x => x.updateAccountListEvent)
+	    .returns(() => () => { return undefined; } );
 
 	return mockAccountManagementService;
 }