Refresh token with SqlToolService session update (#21308)

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

@@ -146,10 +146,6 @@ export class AzureAccountProviderService implements vscode.Disposable {
 		const noSystemKeychain = vscode.workspace.getConfiguration(Constants.AzureSection).get<boolean>(Constants.NoSystemKeyChainSection);
 		const platform = os.platform();
 		const tokenCacheKey = `azureTokenCache-${provider.metadata.id}`;
-		const lockOptions = {
-			retryNumber: 100,
-			retryDelay: 50
-		}
 
 		try {
 			if (!this._credentialProvider) {
@@ -176,7 +172,12 @@ export class AzureAccountProviderService implements vscode.Disposable {
 				throw new Error('Unable to intialize persistence for access token cache. Tokens will not persist in system memory for future use.');
 			}
 
-			let persistenceCachePlugin: PersistenceCachePlugin = new PersistenceCachePlugin(this.persistence, lockOptions); // or any of the other ones.
+			let persistenceCachePlugin: PersistenceCachePlugin = new PersistenceCachePlugin(
+				this.persistence, {
+				retryNumber: 500,
+				retryDelay: 150
+			});
+
 			const MSAL_CONFIG = {
 				auth: {
 					clientId: provider.metadata.settings.clientId,

src/sql/azdata.proposed.d.ts

@@ -446,6 +446,22 @@ declare module 'azdata' {
 		showOnConnectionDialog?: boolean;
 	}
 
+	// Object Explorer interfaces  --------------------------------
+	export interface ObjectExplorerSession {
+		/**
+		 * Authentication token for the current session.
+		 */
+		token?: accounts.AccountSecurityToken | undefined;
+	}
+
+	export interface ExpandNodeInfo {
+		/**
+		 * Authentication token for the current session.
+		 */
+		token?: accounts.AccountSecurityToken | undefined;
+	}
+	// End Object Explorer interfaces  ----------------------------
+
 	export interface TaskInfo {
 		targetLocation?: string;
 	}

src/sql/workbench/contrib/objectExplorer/test/browser/connectionTreeActions.test.ts

@@ -403,7 +403,6 @@ suite('SQL Connection Tree Action tests', () => {
 			resolve(connection);
 		}));
 		connectionManagementService.setup(x => x.isConnected(undefined, TypeMoq.It.isAny())).returns(() => isConnectedReturnValue);
-		connectionManagementService.setup(x => x.refreshAzureAccountTokenIfNecessary(TypeMoq.It.isAny())).returns(async () => true);
 
 		let objectExplorerSession = {
 			success: true,
@@ -450,7 +449,6 @@ suite('SQL Connection Tree Action tests', () => {
 
 		return refreshAction.run().then((value) => {
 			connectionManagementService.verify(x => x.isConnected(undefined, TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
-			connectionManagementService.verify(x => x.refreshAzureAccountTokenIfNecessary(TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
 			objectExplorerService.verify(x => x.getObjectExplorerNode(TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());
 			objectExplorerService.verify(x => x.refreshTreeNode(TypeMoq.It.isAny(), TypeMoq.It.isAny()), TypeMoq.Times.exactly(0));
 			tree.verify(x => x.refresh(TypeMoq.It.isAny()), TypeMoq.Times.atLeastOnce());

src/sql/workbench/services/connection/browser/connectionManagementService.ts

@@ -926,7 +926,7 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 	/**
 	 * Refresh Azure access token if it's expired.
 	 * @param uriOrConnectionProfile connection uri or connection profile
-	 * @returns true if no need to refresh or successfully refreshed token
+	 * @returns true if no need to refresh or successfully refreshed token, false if refresh fails or auth mode is not AzureMFA
 	 */
 	public async refreshAzureAccountTokenIfNecessary(uriOrConnectionProfile: string | ConnectionProfile): Promise<boolean> {
 		if (!uriOrConnectionProfile) {
@@ -997,8 +997,10 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 			} else {
 				this._logService.warn(`Invalid expiry time ${expiry} for connection ${connectionProfile.id} with uri ${uri}`);
 			}
+			return true;
 		}
-		return true;
+		else
+			return false;
 	}
 
 	// Request Senders

src/sql/workbench/services/objectExplorer/browser/connectionTreeAction.ts

@@ -48,7 +48,6 @@ export class RefreshAction extends Action {
 		if (this.element instanceof ConnectionProfile) {
 			let connection: ConnectionProfile = this.element;
 			if (this._connectionManagementService.isConnected(undefined, connection)) {
-				await this._connectionManagementService.refreshAzureAccountTokenIfNecessary(connection);
 				treeNode = this._objectExplorerService.getObjectExplorerNode(connection);
 				if (treeNode === undefined) {
 					await this._objectExplorerService.updateObjectExplorerNodes(connection.toIConnectionProfile());
@@ -56,8 +55,6 @@ export class RefreshAction extends Action {
 				}
 			}
 		} else if (this.element instanceof TreeNode) {
-			let connection: ConnectionProfile = this.element.getConnectionProfile();
-			this._connectionManagementService.refreshAzureAccountTokenIfNecessary(connection);
 			treeNode = this.element;
 		}
 

src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts

@@ -440,7 +440,8 @@ export class ObjectExplorerService implements IObjectExplorerService {
 						allProviders.forEach(provider => {
 							self.callExpandOrRefreshFromProvider(provider, {
 								sessionId: session.sessionId!,
-								nodePath: node.nodePath
+								nodePath: node.nodePath,
+								token: session.token
 							}, refresh).then(isExpanding => {
 								if (!isExpanding) {
 									// The provider stated it's not going to expand the node, therefore do not need to track when merging results
@@ -595,7 +596,18 @@ export class ObjectExplorerService implements IObjectExplorerService {
 		session: azdata.ObjectExplorerSession,
 		parentTree: TreeNode,
 		refresh: boolean = false): Promise<TreeNode[]> {
-		const providerName = parentTree.getConnectionProfile()?.providerName;
+		let connection = parentTree.getConnectionProfile();
+		if (connection) {
+			// Refresh access token on connection if needed.
+			let refreshResult = await this._connectionManagementService.refreshAzureAccountTokenIfNecessary(connection);
+			if (refreshResult) {
+				session.token = {
+					token: connection.options['azureAccountToken'],
+					expiresOn: connection.options['expiresOn']
+				};
+			}
+		}
+		const providerName = connection?.providerName;
 		if (!providerName) {
 			throw new Error('Failed to expand node - no provider name');
 		}

src/sql/workbench/services/objectExplorer/test/browser/objectExplorerService.test.ts

@@ -262,6 +262,7 @@ suite('SQL Object Explorer Service tests', () => {
 		connectionManagementService = TypeMoq.Mock.ofType(TestConnectionManagementService, TypeMoq.MockBehavior.Strict);
 		connectionManagementService.setup(x => x.getConnectionGroups()).returns(() => [conProfGroup]);
 		connectionManagementService.setup(x => x.getActiveConnections()).returns(() => [connection]);
+		connectionManagementService.setup(x => x.refreshAzureAccountTokenIfNecessary(TypeMoq.It.isAny())).returns(async () => true);
 		connectionManagementService.setup(x => x.addSavedPassword(TypeMoq.It.isAny())).returns(() => new Promise<ConnectionProfile>((resolve) => {
 			resolve(connection);
 		}));