[STRAT-3865] | Re-authentication Flow for CreateAudience and GetAudience to handle 401

packages/core/src/destination-kit/index.ts

@@ -141,8 +141,8 @@ const instanceOfAudienceDestinationSettingsWithCreateGet = (
 export interface AudienceDestinationDefinition<Settings = unknown, AudienceSettings = unknown>
   extends DestinationDefinition<Settings> {
   audienceConfig:
-    | AudienceDestinationConfiguration
     | AudienceDestinationConfigurationWithCreateGet<Settings, AudienceSettings>
+    | AudienceDestinationConfiguration
 
   audienceFields: Record<string, GlobalSetting>
 
@@ -422,41 +422,94 @@ export class Destination<Settings = JSONObject, AudienceSettings = JSONObject> {
   }
 
   async createAudience(createAudienceInput: CreateAudienceInput<Settings, AudienceSettings>) {
-    const audienceDefinition = this.definition as AudienceDestinationDefinition
-    if (!instanceOfAudienceDestinationSettingsWithCreateGet(audienceDefinition.audienceConfig)) {
+    let settings: JSONObject = createAudienceInput.settings as unknown as JSONObject
+    const { audienceConfig } = this.definition as AudienceDestinationDefinition
+    if (!instanceOfAudienceDestinationSettingsWithCreateGet(audienceConfig)) {
       throw new Error('Unexpected call to createAudience')
     }
-    const destinationSettings = this.getDestinationSettings(createAudienceInput.settings as unknown as JSONObject)
-    const auth = getAuthData(createAudienceInput.settings as unknown as JSONObject)
-    const context: ExecuteInput<Settings, any, AudienceSettings> = {
-      audienceSettings: createAudienceInput.audienceSettings,
-      settings: destinationSettings,
-      payload: undefined,
-      auth
+    const destinationSettings = this.getDestinationSettings(settings)
+    const run = async () => {
+      const auth = getAuthData(settings)
+      const context: ExecuteInput<Settings, any, AudienceSettings> = {
+        audienceSettings: createAudienceInput.audienceSettings,
+        settings: destinationSettings,
+        payload: undefined,
+        auth
+      }
+      const opts = this.extendRequest?.(context) ?? {}
+      const requestClient = createRequestClient({ ...opts, statsContext: context.statsContext })
+      return await audienceConfig?.createAudience(requestClient, createAudienceInput)
     }
-    const options = this.extendRequest?.(context) ?? {}
-    const requestClient = createRequestClient({ ...options, statsContext: context.statsContext })
 
-    return audienceDefinition.audienceConfig?.createAudience(requestClient, createAudienceInput)
+    const onFailedAttempt = async (error: ResponseError & HTTPError) => {
+      const statusCode = error?.status ?? error?.response?.status ?? 500
+
+      // Throw original error if it is unrelated to invalid access tokens and not an oauth2 scheme
+      if (
+        !(
+          statusCode === 401 &&
+          (this.authentication?.scheme === 'oauth2' || this.authentication?.scheme === 'oauth-managed')
+        )
+      ) {
+        throw error
+      }
+
+      const oauthSettings = getOAuth2Data(settings)
+      const newTokens = await this.refreshAccessToken(destinationSettings, oauthSettings)
+      if (!newTokens) {
+        throw new InvalidAuthenticationError('Failed to refresh access token', ErrorCodes.OAUTH_REFRESH_FAILED)
+      }
+
+      // Update `settings` with new tokens
+      settings = updateOAuthSettings(settings, newTokens)
+    }
+    return await retry(run, { retries: 2, onFailedAttempt })
   }
 
   async getAudience(getAudienceInput: GetAudienceInput<Settings, AudienceSettings>) {
-    const audienceDefinition = this.definition as AudienceDestinationDefinition
-    if (!instanceOfAudienceDestinationSettingsWithCreateGet(audienceDefinition.audienceConfig)) {
+    const { audienceConfig } = this.definition as AudienceDestinationDefinition
+    let settings: JSONObject = getAudienceInput.settings as unknown as JSONObject
+    if (!instanceOfAudienceDestinationSettingsWithCreateGet(audienceConfig)) {
       throw new Error('Unexpected call to getAudience')
     }
-    const destinationSettings = this.getDestinationSettings(getAudienceInput.settings as unknown as JSONObject)
-    const auth = getAuthData(getAudienceInput.settings as unknown as JSONObject)
-    const context: ExecuteInput<Settings, any, AudienceSettings> = {
-      audienceSettings: getAudienceInput.audienceSettings,
-      settings: destinationSettings,
-      payload: undefined,
-      auth
+    const destinationSettings = this.getDestinationSettings(settings)
+    const run = async () => {
+      const auth = getAuthData(settings)
+      const context: ExecuteInput<Settings, any, AudienceSettings> = {
+        audienceSettings: getAudienceInput.audienceSettings,
+        settings: destinationSettings,
+        payload: undefined,
+        auth
+      }
+      const opts = this.extendRequest?.(context) ?? {}
+      const requestClient = createRequestClient({ ...opts, statsContext: context.statsContext })
+      return await audienceConfig?.getAudience(requestClient, getAudienceInput)
     }
-    const options = this.extendRequest?.(context) ?? {}
-    const requestClient = createRequestClient({ ...options, statsContext: context.statsContext })
 
-    return audienceDefinition.audienceConfig?.getAudience(requestClient, getAudienceInput)
+    const onFailedAttempt = async (error: ResponseError & HTTPError) => {
+      const statusCode = error?.status ?? error?.response?.status ?? 500
+
+      // Throw original error if it is unrelated to invalid access tokens and not an oauth2 scheme
+      if (
+        !(
+          statusCode === 401 &&
+          (this.authentication?.scheme === 'oauth2' || this.authentication?.scheme === 'oauth-managed')
+        )
+      ) {
+        throw error
+      }
+
+      const oauthSettings = getOAuth2Data(settings)
+      const newTokens = await this.refreshAccessToken(destinationSettings, oauthSettings)
+      if (!newTokens) {
+        throw new InvalidAuthenticationError('Failed to refresh access token', ErrorCodes.OAUTH_REFRESH_FAILED)
+      }
+
+      // Update `settings` with new tokens
+      settings = updateOAuthSettings(settings, newTokens)
+    }
+
+    return await retry(run, { retries: 2, onFailedAttempt })
   }
 
   async testAuthentication(settings: Settings): Promise<void> {

packages/destination-actions/src/destinations/amazon-amc/__tests__/index.test.ts

@@ -1,5 +1,5 @@
 import nock from 'nock'
-import { createTestIntegration, InvalidAuthenticationError } from '@segment/actions-core'
+import { createTestIntegration } from '@segment/actions-core'
 import Definition from '../index'
 import { HTTPError } from '@segment/actions-core/*'
 import { AUTHORIZATION_URL } from '../utils'
@@ -116,15 +116,6 @@ describe('Amazon-Ads (actions)', () => {
       )
     })
 
-    it('should fail if refresh token API gets failed', async () => {
-      const endpoint = AUTHORIZATION_URL[`${settings.region}`]
-      nock(`${endpoint}`).post('/auth/o2/token').reply(401)
-
-      await expect(testDestination.createAudience(createAudienceInputTemp)).rejects.toThrowError(
-        InvalidAuthenticationError
-      )
-    })
-
     it('should throw an HTTPError when createAudience API response is not ok', async () => {
       const endpoint = AUTHORIZATION_URL[`${settings.region}`]
       nock(`${endpoint}`).post('/auth/o2/token').reply(200)
@@ -138,9 +129,6 @@ describe('Amazon-Ads (actions)', () => {
     })
 
     it('creates an audience', async () => {
-      const endpoint = AUTHORIZATION_URL[`${settings.region}`]
-      nock(`${endpoint}`).post('/auth/o2/token').reply(200)
-
       nock(`${settings.region}`)
         .post('/amc/audiences/metadata')
         .matchHeader('content-type', 'application/vnd.amcaudiences.v1+json')
@@ -205,13 +193,6 @@ describe('Amazon-Ads (actions)', () => {
       await expect(audiencePromise).rejects.toHaveProperty('response.statusText', 'Not Found')
       await expect(audiencePromise).rejects.toHaveProperty('response.status', 404)
     })
-    it('should fail if refresh token API gets failed ', async () => {
-      const endpoint = AUTHORIZATION_URL[`${settings.region}`]
-      nock(`${endpoint}`).post('/auth/o2/token').reply(401)
-
-      const audiencePromise = testDestination.getAudience(getAudienceInput)
-      await expect(audiencePromise).rejects.toThrow(InvalidAuthenticationError)
-    })
 
     it('should throw an IntegrationError when the audienceId is not provided', async () => {
       getAudienceInput.externalId = ''

packages/destination-actions/src/destinations/amazon-amc/index.ts

@@ -4,7 +4,6 @@ import type { Settings, AudienceSettings } from './generated-types'
 import {
   AudiencePayload,
   extractNumberAndSubstituteWithStringValue,
-  getAuthSettings,
   getAuthToken,
   REGEX_ADVERTISERID,
   REGEX_AUDIENCEID
@@ -198,10 +197,6 @@ const destination: AudienceDestinationDefinition<Settings, AudienceSettings> = {
         })
       }
 
-      // @ts-ignore - TS doesn't know about the oauth property
-      const authSettings = getAuthSettings(settings)
-      const authToken = await getAuthToken(request, createAudienceInput.settings, authSettings)
-
       let payloadString = JSON.stringify(payload)
       // Regular expression to find a advertiserId numeric string and replace the quoted advertiserId string with an unquoted number
       // AdvertiserId is very big number string and can not be assigned or converted to number directly as it changes the value due to integer overflow.
@@ -211,8 +206,7 @@ const destination: AudienceDestinationDefinition<Settings, AudienceSettings> = {
         method: 'POST',
         body: payloadString,
         headers: {
-          'Content-Type': 'application/vnd.amcaudiences.v1+json',
-          authorization: `Bearer ${authToken}`
+          'Content-Type': 'application/vnd.amcaudiences.v1+json'
         }
       })
 
@@ -229,18 +223,11 @@ const destination: AudienceDestinationDefinition<Settings, AudienceSettings> = {
       const audience_id = getAudienceInput.externalId
       const { settings } = getAudienceInput
       const endpoint = settings.region
-
       if (!audience_id) {
         throw new IntegrationError('Missing audienceId value', 'MISSING_REQUIRED_FIELD', 400)
       }
-      // @ts-ignore - TS doesn't know about the oauth property
-      const authSettings = getAuthSettings(settings)
-      const authToken = await getAuthToken(request, settings, authSettings)
       const response = await request(`${endpoint}/amc/audiences/metadata/${audience_id}`, {
-        method: 'GET',
-        headers: {
-          authorization: `Bearer ${authToken}`
-        }
+        method: 'GET'
       })
       const res = await response.text()
       // Regular expression to find a audienceId number and replace the audienceId with quoted string