diff --git a/src/sql/platform/connection/common/connectionManagement.ts b/src/sql/platform/connection/common/connectionManagement.ts index 11020422b0..d34cbdd4d7 100644 --- a/src/sql/platform/connection/common/connectionManagement.ts +++ b/src/sql/platform/connection/common/connectionManagement.ts @@ -176,7 +176,7 @@ export interface IConnectionManagementService { isConnected(fileUri: string): boolean; - refreshAzureAccountTokenIfNecessary(uri: string): Promise; + refreshAzureAccountTokenIfNecessary(uriOrConnectionProfile: string | ConnectionProfile): Promise; /** * Returns true if the connection profile is connected */ diff --git a/src/sql/platform/connection/test/common/testConnectionManagementService.ts b/src/sql/platform/connection/test/common/testConnectionManagementService.ts index 19588bfe56..694716c073 100644 --- a/src/sql/platform/connection/test/common/testConnectionManagementService.ts +++ b/src/sql/platform/connection/test/common/testConnectionManagementService.ts @@ -319,7 +319,7 @@ export class TestConnectionManagementService implements IConnectionManagementSer return undefined!; } - refreshAzureAccountTokenIfNecessary(uri: string): Promise { + refreshAzureAccountTokenIfNecessary(uriOrConnectionProfile: string | ConnectionProfile): Promise { return undefined; } 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 489906fe26..c700c8c32a 100644 --- a/src/sql/workbench/contrib/objectExplorer/test/browser/connectionTreeActions.test.ts +++ b/src/sql/workbench/contrib/objectExplorer/test/browser/connectionTreeActions.test.ts @@ -403,6 +403,7 @@ 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, @@ -449,6 +450,7 @@ 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 02a17d8c25..07296b4910 100644 --- a/src/sql/workbench/services/connection/browser/connectionManagementService.ts +++ b/src/sql/workbench/services/connection/browser/connectionManagementService.ts @@ -906,14 +906,15 @@ export class ConnectionManagementService extends Disposable implements IConnecti } const tenantId = connection.azureTenantId; const token = await this._accountManagementService.getAccountSecurityToken(account, tenantId, azureResource); - this._logService.debug(`Got token for tenant ${token}`); if (!token) { - this._logService.info(`No security tokens found for account`); + this._logService.warn(`No security tokens found for account`); + } else { + this._logService.debug(`Got access token for tenant ${tenantId} that expires in ${(token.expiresOn - new Date().getTime()) / 1000} seconds`); + connection.options['azureAccountToken'] = token.token; + connection.options['expiresOn'] = token.expiresOn; + connection.options['password'] = ''; + return true; } - connection.options['azureAccountToken'] = token.token; - connection.options['expiresOn'] = token.expiresOn; - connection.options['password'] = ''; - return true; } else { this._logService.info(`Could not find Azure account with name ${accountId}`); } @@ -925,56 +926,78 @@ export class ConnectionManagementService extends Disposable implements IConnecti /** * Refresh Azure access token if it's expired. - * @param uri connection uri + * @param uriOrConnectionProfile connection uri or connection profile * @returns true if no need to refresh or successfully refreshed token */ - public async refreshAzureAccountTokenIfNecessary(uri: string): Promise { - const profile = this._connectionStatusManager.getConnectionProfile(uri); - if (!profile) { - this._logService.warn(`Connection not found for uri ${uri}`); + public async refreshAzureAccountTokenIfNecessary(uriOrConnectionProfile: string | ConnectionProfile): Promise { + if (!uriOrConnectionProfile) { + this._logService.warn(`refreshAzureAccountTokenIfNecessary: Neither Connection uri nor connection profile received.`); return false; } - //wait for the pending reconnction promise if any + let uri: string; + let connectionProfile: ConnectionProfile; + + if (typeof uriOrConnectionProfile === 'string') { + uri = uriOrConnectionProfile; + connectionProfile = this._connectionStatusManager.getConnectionProfile(uri); + if (!connectionProfile) { + this._logService.warn(`Connection not found for uri ${uri} when refreshing token`); + return false; + } + } else { + connectionProfile = uriOrConnectionProfile; + uri = this.getConnectionUri(connectionProfile); + } + + // Wait for the pending reconnction promise if any + // We expect uri to be defined const previousReconnectPromise = this._uriToReconnectPromiseMap[uri]; if (previousReconnectPromise) { - this._logService.info(`Found pending reconnect promise for uri ${uri}, waiting.`); + this._logService.debug(`Found pending reconnect promise for uri ${uri}, waiting.`); try { const previousConnectionResult = await previousReconnectPromise; if (previousConnectionResult && previousConnectionResult.connected) { - this._logService.info(`Previous pending reconnection for uri ${uri} succeeded.`); + this._logService.debug(`Previous pending reconnection for uri ${uri} succeeded.`); return true; } - this._logService.info(`Previous pending reconnection for uri ${uri} failed.`); + this._logService.debug(`Previous pending reconnection for uri ${uri} failed.`); } catch (err) { - this._logService.info(`Previous pending reconnect promise for uri ${uri} is rejected with error ${err}, will attempt to reconnect if necessary.`); + this._logService.debug(`Previous pending reconnect promise for uri ${uri} is rejected with error ${err}, will attempt to reconnect if necessary.`); } } - const expiry = profile.options.expiresOn; - if (typeof expiry === 'number' && !Number.isNaN(expiry)) { - const currentTime = new Date().getTime() / 1000; - const maxTolerance = 2 * 60; // two minutes - if (expiry - currentTime < maxTolerance) { - this._logService.info(`Access token expired for connection ${profile.id} with uri ${uri}`); - try { - const connectionResultPromise = this.connect(profile, uri); - this._uriToReconnectPromiseMap[uri] = connectionResultPromise; - const connectionResult = await connectionResultPromise; - if (!connectionResult) { - this._logService.error(`Failed to refresh connection ${profile.id} with uri ${uri}, invalid connection result.`); - throw new Error(nls.localize('connection.invalidConnectionResult', "Connection result is invalid")); - } else if (!connectionResult.connected) { - this._logService.error(`Failed to refresh connection ${profile.id} with uri ${uri}, error code: ${connectionResult.errorCode}, error message: ${connectionResult.errorMessage}`); - throw new Error(nls.localize('connection.refreshAzureTokenFailure', "Failed to refresh Azure account token for connection")); + // We expect connectionProfile to be defined + if (connectionProfile && connectionProfile.authenticationType === Constants.AuthenticationType.AzureMFA) { + const expiry = connectionProfile.options.expiresOn; + if (typeof expiry === 'number' && !Number.isNaN(expiry)) { + const currentTime = new Date().getTime() / 1000; + const maxTolerance = 2 * 60; // two minutes + if (expiry - currentTime < maxTolerance) { + this._logService.debug(`Access token expired for connection ${connectionProfile.id} with uri ${uri}`); + try { + const connectionResultPromise = this.connect(connectionProfile, uri); + this._uriToReconnectPromiseMap[uri] = connectionResultPromise; + const connectionResult = await connectionResultPromise; + if (!connectionResult) { + this._logService.error(`Failed to refresh connection ${connectionProfile.id} with uri ${uri}, invalid connection result.`); + throw new Error(nls.localize('connection.invalidConnectionResult', "Connection result is invalid")); + } else if (!connectionResult.connected) { + this._logService.error(`Failed to refresh connection ${connectionProfile.id} with uri ${uri}, error code: ${connectionResult.errorCode}, error message: ${connectionResult.errorMessage}`); + throw new Error(nls.localize('connection.refreshAzureTokenFailure', "Failed to refresh Azure account token for connection")); + } + this._logService.debug(`Successfully refreshed token for connection ${connectionProfile.id} with uri ${uri}, result: ${connectionResult.connected} ${connectionResult.connectionProfile}, ${this._connectionStatusManager.getConnectionProfile(uri)}`); + return true; + } finally { + delete this._uriToReconnectPromiseMap[uri]; } - this._logService.info(`Successfully refreshed token for connection ${profile.id} with uri ${uri}, result: ${connectionResult.connected} ${connectionResult.connectionProfile}, isConnected: ${this.isConnected(uri)}, ${this._connectionStatusManager.getConnectionProfile(uri)}`); - return true; - } finally { - delete this._uriToReconnectPromiseMap[uri]; } + else { + this._logService.debug(`No need to refresh Azure acccount token for connection ${connectionProfile.id} with uri ${uri}`); + } + } else { + this._logService.warn(`Invalid expiry time ${expiry} for connection ${connectionProfile.id} with uri ${uri}`); } - this._logService.info(`No need to refresh Azure acccount token for connection ${profile.id} with uri ${uri}`); } return true; } diff --git a/src/sql/workbench/services/objectExplorer/browser/connectionTreeAction.ts b/src/sql/workbench/services/objectExplorer/browser/connectionTreeAction.ts index 80ebd4cdcd..aab80e6644 100644 --- a/src/sql/workbench/services/objectExplorer/browser/connectionTreeAction.ts +++ b/src/sql/workbench/services/objectExplorer/browser/connectionTreeAction.ts @@ -48,6 +48,7 @@ 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()); @@ -55,6 +56,8 @@ export class RefreshAction extends Action { } } } else if (this.element instanceof TreeNode) { + let connection: ConnectionProfile = this.element.getConnectionProfile(); + this._connectionManagementService.refreshAzureAccountTokenIfNecessary(connection); treeNode = this.element; }