Fix #4089 Linked account cancel (#5347)

VSCode serialization changed rejection handling to only serialize errors, which caused things to break
- Changed to return either account info or a cancel message in the resolve
- Rewrote to use promises. Tracking how to return canceled through 4+ thenables was way trickier than just using a promise
- Updated unit tests to handle new scenario
- Tested integration tests, realized they a) didn't run and b) didn't passed. 
  - Added vscode dev dependency to fix run issue
  - Fixed tests to account for behavior changes in tree state.

extensions/azurecore/package.json

@@ -6,7 +6,7 @@
 	"publisher": "Microsoft",
 	"preview": true,
 	"engines": {
-		"vscode": "^1.25.0",
+		"vscode": "^1.30.1",
 		"azdata": "*"
 	},
 	"activationEvents": [
@@ -169,6 +169,7 @@
 		"@types/request": "^2.48.1",
 		"mocha": "^5.2.0",
 		"should": "^13.2.1",
-		"typemoq": "^2.1.0"
+		"typemoq": "^2.1.0",
+		"vscode": "1.1.26"
 	}
 }

extensions/azurecore/src/account-provider/azureAccountProvider.ts

@@ -103,16 +103,16 @@ export class AzureAccountProvider implements azdata.AccountProvider {
 			});
 	}
 
-	public prompt(): Thenable<AzureAccount> {
+	public prompt(): Thenable<AzureAccount | azdata.PromptFailedResult> {
 		return this.doIfInitialized(() => this.signIn(true));
 	}
 
-	public refresh(account: AzureAccount): Thenable<AzureAccount> {
+	public refresh(account: AzureAccount): Thenable<AzureAccount | azdata.PromptFailedResult> {
 		return this.doIfInitialized(() => this.signIn(false));
 	}
 
 	// PRIVATE METHODS /////////////////////////////////////////////////////
-	private cancelAutoOAuth(): Thenable<void> {
+	private cancelAutoOAuth(): Promise<void> {
 		let self = this;
 
 		if (!this._inProgressAutoOAuth) {
@@ -137,23 +137,23 @@ export class AzureAccountProvider implements azdata.AccountProvider {
 		return Promise.resolve();
 	}
 
-	private clearAccountTokens(accountKey: azdata.AccountKey): Thenable<void> {
+	private async clearAccountTokens(accountKey: azdata.AccountKey): Promise<void> {
 		// Put together a query to look up any tokens associated with the account key
 		let query = <adal.TokenResponse>{ userId: accountKey.accountId };
 
 		// 1) Look up the tokens associated with the query
 		// 2) Remove them
-		return this._tokenCache.findThenable(query)
+		let results = await this._tokenCache.findThenable(query);
-			.then(results => this._tokenCache.removeThenable(results));
+		this._tokenCache.removeThenable(results);
 	}
 
-	private doIfInitialized<T>(op: () => Thenable<T>): Thenable<T> {
+	private doIfInitialized<T>(op: () => Promise<T>): Promise<T> {
 		return this._isInitialized
 			? op()
 			: Promise.reject(localize('accountProviderNotInitialized', 'Account provider not initialized, cannot perform action'));
 	}
 
-	private getAccessTokens(account: AzureAccount, resource: azdata.AzureResource): Thenable<AzureAccountSecurityTokenCollection> {
+	private getAccessTokens(account: AzureAccount, resource: azdata.AzureResource): Promise<AzureAccountSecurityTokenCollection> {
 		let self = this;
 
 		const resourceIdMap = new Map<azdata.AzureResource, string>([
@@ -228,7 +228,7 @@ export class AzureAccountProvider implements azdata.AccountProvider {
 		});
 	}
 
-	private getDeviceLoginToken(oAuth: InProgressAutoOAuth, isAddAccount: boolean): Thenable<adal.TokenResponse> {
+	private getDeviceLoginToken(oAuth: InProgressAutoOAuth, isAddAccount: boolean): Thenable<adal.TokenResponse | azdata.PromptFailedResult> {
 		let self = this;
 
 		// 1) Open the auto OAuth dialog
@@ -239,14 +239,15 @@ export class AzureAccountProvider implements azdata.AccountProvider {
 			localize('refreshAccount', 'Refresh {0} account', self._metadata.displayName);
 		return azdata.accounts.beginAutoOAuthDeviceCode(self._metadata.id, title, oAuth.userCodeInfo.message, oAuth.userCodeInfo.userCode, oAuth.userCodeInfo.verificationUrl)
 			.then(() => {
-				return new Promise<adal.TokenResponse>((resolve, reject) => {
+				return new Promise<adal.TokenResponse | azdata.PromptFailedResult>((resolve, reject) => {
 					let context = oAuth.context;
 					context.acquireTokenWithDeviceCode(self._metadata.settings.signInResourceId, self._metadata.settings.clientId, oAuth.userCodeInfo,
 						(err, response) => {
 							if (err) {
 								if (self._autoOAuthCancelled) {
+									let result: azdata.PromptFailedResult = { canceled: true };
 									// Auto OAuth was cancelled by the user, indicate this with the error we return
-									reject(<azdata.UserCancelledSignInError>{ userCancelledSignIn: true });
+									resolve(result);
 								} else {
 									// Auto OAuth failed for some other reason
 									azdata.accounts.endAutoOAuthDeviceCode();
@@ -368,27 +369,27 @@ export class AzureAccountProvider implements azdata.AccountProvider {
 		});
 	}
 
-	private signIn(isAddAccount: boolean): Thenable<AzureAccount> {
+	private isPromptFailed(value: adal.TokenResponse | azdata.PromptFailedResult): value is azdata.PromptFailedResult {
-		let self = this;
+		return value && (<azdata.PromptFailedResult>value).canceled;
+	}
 
+	private async signIn(isAddAccount: boolean): Promise<AzureAccount | azdata.PromptFailedResult> {
 		// 1) Get the user code for this login
 		// 2) Get an access token from the device code
 		// 3) Get the list of tenants
 		// 4) Generate the AzureAccount object and return it
 		let tokenResponse: adal.TokenResponse = null;
-		return this.getDeviceLoginUserCode()
+		let result: InProgressAutoOAuth = await this.getDeviceLoginUserCode();
-			.then((result: InProgressAutoOAuth) => {
+		this._autoOAuthCancelled = false;
-				self._autoOAuthCancelled = false;
+		this._inProgressAutoOAuth = result;
-				self._inProgressAutoOAuth = result;
+		let response: adal.TokenResponse | azdata.PromptFailedResult = await this.getDeviceLoginToken(this._inProgressAutoOAuth, isAddAccount);
-				return self.getDeviceLoginToken(self._inProgressAutoOAuth, isAddAccount);
+		if (this.isPromptFailed(response)) {
-			})
+			return response;
-			.then((response: adal.TokenResponse) => {
+		}
 		tokenResponse = response;
-				self._autoOAuthCancelled = false;
+		this._autoOAuthCancelled = false;
-				self._inProgressAutoOAuth = null;
+		this._inProgressAutoOAuth = null;
-				return self.getTenants(tokenResponse.userId, tokenResponse.userId);
+		let tenants: Tenant[] = await this.getTenants(tokenResponse.userId, tokenResponse.userId);
-			})
-			.then((tenants: Tenant[]) => {
 		// Figure out where we're getting the identity from
 		let identityProvider = tokenResponse.identityProvider;
 		if (identityProvider) {
@@ -419,7 +420,7 @@ export class AzureAccountProvider implements azdata.AccountProvider {
 
 		return <AzureAccount>{
 			key: {
-						providerId: self._metadata.id,
+				providerId: this._metadata.id,
 				accountId: tokenResponse.userId
 			},
 			name: tokenResponse.userId,
@@ -435,7 +436,6 @@ export class AzureAccountProvider implements azdata.AccountProvider {
 			},
 			isStale: false
 		};
-			});
 	}
 }
 

extensions/azurecore/src/test/azureResource/providers/database/databaseTreeDataProvider.test.ts

@@ -142,7 +142,7 @@ describe('AzureResourceDatabaseTreeDataProvider.getChildren', function (): void
 			should(child.tenantId).equal(mockTenantId);
 			should(child.treeItem.id).equal(`databaseServer_${database.serverFullName}.database_${database.name}`);
 			should(child.treeItem.label).equal(`${database.name} (${database.serverName})`);
-			should(child.treeItem.collapsibleState).equal(vscode.TreeItemCollapsibleState.None);
+			should(child.treeItem.collapsibleState).equal(vscode.TreeItemCollapsibleState.Collapsed);
 			should(child.treeItem.contextValue).equal(AzureResourceItemType.database);
 		}
 	});

extensions/azurecore/src/test/azureResource/providers/databaseServer/databaseServerTreeDataProvider.test.ts

@@ -142,7 +142,7 @@ describe('AzureResourceDatabaseServerTreeDataProvider.getChildren', function ():
 			should(child.tenantId).equal(mockTenantId);
 			should(child.treeItem.id).equal(`databaseServer_${databaseServer.name}`);
 			should(child.treeItem.label).equal(databaseServer.name);
-			should(child.treeItem.collapsibleState).equal(vscode.TreeItemCollapsibleState.None);
+			should(child.treeItem.collapsibleState).equal(vscode.TreeItemCollapsibleState.Collapsed);
 			should(child.treeItem.contextValue).equal(AzureResourceItemType.databaseServer);
 		}
 	});

src/sql/azdata.proposed.d.ts

@@ -2327,11 +2327,11 @@ declare module 'azdata' {
 	 * AccountProvider.refresh or AccountProvider.prompt are rejected with this error, the error
 	 * will not be reported to the user.
 	 */
-	export interface UserCancelledSignInError extends Error {
+	export interface PromptFailedResult {
 		/**
 		 * Type guard for differentiating user cancelled sign in errors from other errors
 		 */
-		userCancelledSignIn: boolean;
+		canceled: boolean;
 	}
 
 	/**
@@ -2382,7 +2382,7 @@ declare module 'azdata' {
 		 * Prompts the user to enter account information.
 		 * Returns an error if the user canceled the operation.
 		 */
-		prompt(): Thenable<Account>;
+		prompt(): Thenable<Account | PromptFailedResult>;
 
 		/**
 		 * Refreshes a stale account.
@@ -2390,7 +2390,7 @@ declare module 'azdata' {
 		 * Otherwise, returns a new updated account instance.
 		 * @param account - An account.
 		 */
-		refresh(account: Account): Thenable<Account>;
+		refresh(account: Account): Thenable<Account | PromptFailedResult>;
 
 		/**
 		 * Clears sensitive information for an account. To be called when account is removed

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

@@ -36,11 +36,11 @@ export class ExtHostAccountManagement extends ExtHostAccountManagementShape {
 		return this._withProvider(handle, (provider: azdata.AccountProvider) => provider.initialize(restoredAccounts));
 	}
 
-	public $prompt(handle: number): Thenable<azdata.Account> {
+	public $prompt(handle: number): Thenable<azdata.Account | azdata.PromptFailedResult> {
 		return this._withProvider(handle, (provider: azdata.AccountProvider) => provider.prompt());
 	}
 
-	public $refresh(handle: number, account: azdata.Account): Thenable<azdata.Account> {
+	public $refresh(handle: number, account: azdata.Account): Thenable<azdata.Account | azdata.PromptFailedResult> {
 		return this._withProvider(handle, (provider: azdata.AccountProvider) => provider.refresh(account));
 	}
 

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

@@ -81,10 +81,10 @@ export class MainThreadAccountManagement implements MainThreadAccountManagementS
 			initialize(restoredAccounts: azdata.Account[]): Thenable<azdata.Account[]> {
 				return self._proxy.$initialize(handle, restoredAccounts);
 			},
-			prompt(): Thenable<azdata.Account> {
+			prompt(): Thenable<azdata.Account | azdata.PromptFailedResult> {
 				return self._proxy.$prompt(handle);
 			},
-			refresh(account: azdata.Account): Thenable<azdata.Account> {
+			refresh(account: azdata.Account): Thenable<azdata.Account | azdata.PromptFailedResult> {
 				return self._proxy.$refresh(handle, account);
 			}
 		};

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

@@ -30,8 +30,8 @@ export abstract class ExtHostAccountManagementShape {
 	$clear(handle: number, accountKey: azdata.AccountKey): Thenable<void> { throw ni(); }
 	$getSecurityToken(account: azdata.Account, resource?: azdata.AzureResource): Thenable<{}> { throw ni(); }
 	$initialize(handle: number, restoredAccounts: azdata.Account[]): Thenable<azdata.Account[]> { throw ni(); }
-	$prompt(handle: number): Thenable<azdata.Account> { throw ni(); }
+	$prompt(handle: number): Thenable<azdata.Account | azdata.PromptFailedResult> { throw ni(); }
-	$refresh(handle: number, account: azdata.Account): Thenable<azdata.Account> { throw ni(); }
+	$refresh(handle: number, account: azdata.Account): Thenable<azdata.Account | azdata.PromptFailedResult> { throw ni(); }
 	$accountsChanged(handle: number, accounts: azdata.Account[]): Thenable<void> { throw ni(); }
 }
 

src/sql/workbench/services/accountManagement/browser/accountManagementService.ts

@@ -21,6 +21,7 @@ import { AccountListStatusbarItem } from 'sql/platform/accounts/browser/accountL
 import { AccountProviderAddedEventParams, UpdateAccountListEventParams } from 'sql/platform/accounts/common/eventTypes';
 import { IAccountManagementService } from 'sql/platform/accounts/common/interfaces';
 import { Deferred } from 'sql/base/common/promise';
+import { localize } from 'vs/nls';
 
 export class AccountManagementService implements IAccountManagementService {
 	// CONSTANTS ///////////////////////////////////////////////////////////
@@ -127,10 +128,13 @@ export class AccountManagementService implements IAccountManagementService {
 	public addAccount(providerId: string): Thenable<void> {
 		let self = this;
 
-		return this.doWithProvider(providerId, (provider) => {
+		return this.doWithProvider(providerId, async (provider) => {
-			return provider.provider.prompt()
+			let account = await provider.provider.prompt();
-				.then(account => self._accountStore.addOrUpdate(account))
+			if (self.isCanceledResult(account)) {
-				.then(result => {
+				return;
+			}
+
+			let result = await self._accountStore.addOrUpdate(account);
 			if (result.accountAdded) {
 				// Add the account to the list
 				provider.accounts.push(result.changedAccount);
@@ -140,15 +144,11 @@ export class AccountManagementService implements IAccountManagementService {
 			}
 
 			self.fireAccountListUpdate(provider, result.accountAdded);
-				})
+		});
-				.then(null, err => {
-					// On error, check to see if the error is because the user cancelled. If so, just ignore
-					if (err && 'userCancelledSignIn' in err) {
-						return Promise.resolve();
 	}
-					return Promise.reject(err);
+
-				});
+	private isCanceledResult(result: azdata.Account | azdata.PromptFailedResult): result is azdata.PromptFailedResult {
-		});
+		return (<azdata.PromptFailedResult>result).canceled;
 	}
 
 	/**
@@ -159,10 +159,16 @@ export class AccountManagementService implements IAccountManagementService {
 	public refreshAccount(account: azdata.Account): Thenable<azdata.Account> {
 		let self = this;
 
-		return this.doWithProvider(account.key.providerId, (provider) => {
+		return this.doWithProvider(account.key.providerId, async (provider) => {
-			return provider.provider.refresh(account)
+			let refreshedAccount = await provider.provider.refresh(account);
-				.then(account => self._accountStore.addOrUpdate(account))
+			if (self.isCanceledResult(refreshedAccount)) {
-				.then(result => {
+				// Pattern here is to throw if this fails. Handled upstream.
+				throw new Error(localize('refreshFailed', 'Refresh account was canceled by the user'));
+			} else {
+				account = refreshedAccount;
+			}
+
+			let result = await self._accountStore.addOrUpdate(account);
 			if (result.accountAdded) {
 				// Add the account to the list
 				provider.accounts.push(result.changedAccount);
@@ -180,7 +186,6 @@ export class AccountManagementService implements IAccountManagementService {
 			self.fireAccountListUpdate(provider, result.accountAdded);
 			return result.changedAccount;
 		});
-		});
 	}
 
 	/**

src/sqltest/services/accountManagement/accountManagementService.test.ts

@@ -608,13 +608,13 @@ function getFailingMockAccountProvider(cancel: boolean): TypeMoq.Mock<azdata.Acc
 	mockProvider.setup(x => x.prompt())
 		.returns(() => {
 			return cancel
-				? Promise.reject(<azdata.UserCancelledSignInError>{ userCancelledSignIn: true }).then()
+				? Promise.resolve(<azdata.PromptFailedResult>{ canceled: true }).then()
 				: Promise.reject(new Error()).then();
 		});
 	mockProvider.setup(x => x.refresh(TypeMoq.It.isAny()))
 		.returns(() => {
 			return cancel
-				? Promise.reject(<azdata.UserCancelledSignInError>{ userCancelledSignIn: true }).then()
+				? Promise.resolve(<azdata.PromptFailedResult>{ canceled: true }).then()
 				: Promise.reject(new Error()).then();
 		});
 	return mockProvider;