From 23dfd690a63a0064a802b08cddcdb599d13c2717 Mon Sep 17 00:00:00 2001 From: Cheena Malhotra <13396919+cheenamalhotra@users.noreply.github.com> Date: Tue, 29 Nov 2022 15:26:15 -0800 Subject: [PATCH] Refresh token with SqlToolService session update (#21308) --- .../azureAccountProviderService.ts | 11 ++++++----- src/sql/azdata.proposed.d.ts | 16 ++++++++++++++++ .../test/browser/connectionTreeActions.test.ts | 2 -- .../browser/connectionManagementService.ts | 6 ++++-- .../browser/connectionTreeAction.ts | 3 --- .../browser/objectExplorerService.ts | 16 ++++++++++++++-- .../test/browser/objectExplorerService.test.ts | 1 + 7 files changed, 41 insertions(+), 14 deletions(-) diff --git a/extensions/azurecore/src/account-provider/azureAccountProviderService.ts b/extensions/azurecore/src/account-provider/azureAccountProviderService.ts index 0b0e91300d..0f3642557a 100644 --- a/extensions/azurecore/src/account-provider/azureAccountProviderService.ts +++ b/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(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, diff --git a/src/sql/azdata.proposed.d.ts b/src/sql/azdata.proposed.d.ts index 053f934384..d14acc8d4f 100644 --- a/src/sql/azdata.proposed.d.ts +++ b/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; } diff --git a/src/sql/workbench/contrib/objectExplorer/test/browser/connectionTreeActions.test.ts b/src/sql/workbench/contrib/objectExplorer/test/browser/connectionTreeActions.test.ts index c700c8c32a..489906fe26 100644 --- a/src/sql/workbench/contrib/objectExplorer/test/browser/connectionTreeActions.test.ts +++ b/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()); diff --git a/src/sql/workbench/services/connection/browser/connectionManagementService.ts b/src/sql/workbench/services/connection/browser/connectionManagementService.ts index dffbc894ca..aca67d8769 100644 --- a/src/sql/workbench/services/connection/browser/connectionManagementService.ts +++ b/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 { 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 diff --git a/src/sql/workbench/services/objectExplorer/browser/connectionTreeAction.ts b/src/sql/workbench/services/objectExplorer/browser/connectionTreeAction.ts index aab80e6644..80ebd4cdcd 100644 --- a/src/sql/workbench/services/objectExplorer/browser/connectionTreeAction.ts +++ b/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; } diff --git a/src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts b/src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts index 472cbe644c..513b12de73 100644 --- a/src/sql/workbench/services/objectExplorer/browser/objectExplorerService.ts +++ b/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 { - 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'); } diff --git a/src/sql/workbench/services/objectExplorer/test/browser/objectExplorerService.test.ts b/src/sql/workbench/services/objectExplorer/test/browser/objectExplorerService.test.ts index 9fb15eef53..0ca4163c9c 100644 --- a/src/sql/workbench/services/objectExplorer/test/browser/objectExplorerService.test.ts +++ b/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((resolve) => { resolve(connection); }));