Refreshes token for intellisense (#19214)

* wip

* wip

* wip

* working refresh token for intellisense

* pr review comments

* pr review comments

* add link

* pr comments

* change authority -> tenantId

* refactor tenant

* fix build

* add js doc comments, other pr changes

* fix error messaging

* add log

* added logs

* remove expiresOn

* fix error messaging

* pr comments

* remove localized strings from logs

extensions/mssql/src/features.ts

@@ -13,6 +13,7 @@ import * as Utils from './utils';
 import * as UUID from 'vscode-languageclient/lib/utils/uuid';
 import { DataItemCache } from './util/dataCache';
 import * as azurecore from 'azurecore';
+import * as localizedConstants from './localizedConstants';
 
 const localize = nls.loadMessageBundle();
 
@@ -43,7 +44,17 @@ export class AccountFeature implements StaticFeature {
 		let timeToLiveInSeconds = 10;
 		this.tokenCache = new DataItemCache(this.getToken, timeToLiveInSeconds);
 		this._client.onRequest(contracts.SecurityTokenRequest.type, async (request): Promise<contracts.RequestSecurityTokenResponse | undefined> => {
-			return this.tokenCache.getData(request);
+			return await this.tokenCache.getData(request);
+		});
+		this._client.onNotification(contracts.RefreshTokenNotification.type, async (request) => {
+			// Refresh token, then inform client the token has been updated. This is done as separate notification messages due to the synchronous processing nature of STS currently https://github.com/microsoft/azuredatastudio/issues/17179
+			let result = await this.refreshToken(request);
+			if (!result) {
+				void window.showErrorMessage(localizedConstants.tokenRefreshFailed('autocompletion'));
+				console.log(`Token Refresh Failed ${request.toString()}`);
+				throw Error(localizedConstants.tokenRefreshFailed('autocompletion'));
+			}
+			this._client.sendNotification(contracts.TokenRefreshedNotification.type, result);
 		});
 	}
 
@@ -92,6 +103,38 @@ export class AccountFeature implements StaticFeature {
 		return params;
 	}
 
+	protected async refreshToken(request: contracts.RefreshTokenParams): Promise<contracts.TokenRefreshedParams> {
+
+		// find account
+		const accountList = await azdata.accounts.getAllAccounts();
+		const account = accountList.find(a => a.key.accountId === request.accountId);
+		if (account) {
+			console.log(`Failed to find azure account ${request.accountId} when executing token refresh`);
+			throw Error(localizedConstants.failedToFindAccount(request.accountId));
+		}
+
+		// find tenant
+		const tenant = account.properties.tenants.find(tenant => tenant.id === request.tenantId);
+		if (!tenant) {
+			console.log(`Failed to find tenant ${request.tenantId} in account ${account.displayInfo.displayName} when refreshing security token`);
+			throw Error(localizedConstants.failedToFindTenants(request.tenantId, account.displayInfo.displayName));
+		}
+
+		// Get the updated token, which will handle refreshing it if necessary
+		const securityToken = await azdata.accounts.getAccountSecurityToken(account, tenant.id, azdata.AzureResource.ResourceManagement);
+		if (!securityToken) {
+			console.log('Editor token refresh failed, autocompletion will be disabled until the editor is disconnected and reconnected');
+			throw Error(localizedConstants.tokenRefreshFailedNoSecurityToken);
+		}
+		let params: contracts.TokenRefreshedParams = {
+			token: securityToken.token,
+			expiresOn: securityToken.expiresOn,
+			uri: request.uri
+		};
+
+		return params;
+	}
+
 	static AccountQuickPickItem = class implements QuickPickItem {
 		account: azdata.Account;
 		label: string;