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/contracts.ts

@@ -49,6 +49,69 @@ export namespace SecurityTokenRequest {
 }
 // ------------------------------- </ Security Token Request > ------------------------------------------
 
+// ------------------------------- < Refresh Token Notification > ---------------------------------
+
+/**
+ * Parameters for a refresh token notification sent from STS to ADS
+ */
+export interface RefreshTokenParams {
+	/**
+	 * The tenant ID
+	 */
+	tenantId: string;
+	/**
+	 * The provider that indicates the type of linked account to query
+	 */
+	provider: string;
+	/**
+	 * The identifier of the target resource of the requested token
+	 */
+	resource: string;
+	/**
+	 * The account ID
+	 */
+	accountId: string;
+	/**
+	 * The URI for the editor that needs a token refresh
+	 */
+	uri: string;
+}
+
+export namespace RefreshTokenNotification {
+	export const type = new NotificationType<RefreshTokenParams, void>('account/refreshToken');
+}
+
+
+
+// ------------------------------- </ Refresh Token Notification > -------------------------------
+
+// ------------------------------- < Token Refreshed Notification > ---------------------------------
+
+/**
+ * Parameters for a new refresh token sent from ADS to STS
+ */
+export interface TokenRefreshedParams {
+	/**
+	 * The refresh token
+	 */
+	token: string;
+	/**
+	 * The token expiration, a Unix epoch
+	 */
+	expiresOn: Number;
+	/**
+	 * The URI for the editor that needs a token refresh
+	 */
+	uri: string;
+}
+
+export namespace TokenRefreshedNotification {
+	export const type = new NotificationType<TokenRefreshedParams, void>('account/tokenRefreshed');
+}
+
+// ------------------------------- </ Token Refreshed Notification > -------------------------------
+
+
 // ------------------------------- < Agent Management > ------------------------------------
 // Job management parameters
 export interface AgentJobsParams {

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;

extensions/mssql/src/localizedConstants.ts

@@ -55,3 +55,10 @@ export function sparkJobSubmissionGetApplicationIdFailed(err: string): string {
 export function sparkJobSubmissionLocalFileNotExisted(path: string): string { return localize('sparkJobSubmission.LocalFileNotExisted', "Local file {0} does not existed. ", path); }
 export const sparkJobSubmissionNoSqlBigDataClusterFound = localize('sparkJobSubmission.NoSqlBigDataClusterFound', "No SQL Server Big Data Cluster found.");
 export function sparkConnectionRequired(name: string): string { return localize('sparkConnectionRequired', "Please connect to the Spark cluster before View {0} History.", name); }
+
+
+export function failedToFindTenants(tenantId: string, accountName: string): string { return localize('mssql.failedToFindTenants', "Failed to find tenant '{0}' in account '{1}' when refreshing security token", tenantId, accountName); }
+export function tokenRefreshFailed(name: string): string { return localize('mssql.tokenRefreshFailed', "{0} AAD token refresh failed, please reconnect to enable {0}", name); }
+export const tokenRefreshFailedNoSecurityToken = localize('mssql.tokenRefreshFailedNoSecurityToken', "Editor token refresh failed, autocompletion will be disabled until the editor is disconnected and reconnected");
+export function failedToFindAccount(accountName: string) { return localize('mssql.failedToFindAccount', "Failed to find azure account {0} when executing token refresh", accountName); }
+