Implement batch support for Attio Actions

packages/destination-actions/src/destinations/attio/api/index.ts

@@ -3,6 +3,17 @@ import { ModifiedResponse } from '@segment/actions-core'
 import get from 'lodash/get'
 import sortBy from 'lodash/sortBy'
 
+export type SimpleValue = string | number | boolean
+
+type BatchAssertion = {
+  object: string
+  mode: 'create-or-update'
+  matching_attribute: string
+  multiselect_values: 'append'
+  values: Record<string, null | SimpleValue | Array<SimpleValue> | BatchAssertion | Array<BatchAssertion>>
+  received_at: string
+}
+
 export type AssertResponse = {
   data: {
     id: {
@@ -59,6 +70,26 @@ export class AttioClient {
     )
   }
 
+  /**
+   * Send a series of (nested) assertions in a single HTTP call
+   *
+   * @param assertions One or more assertions to apply
+   * @param requestOptions Additional options for the request
+   */
+  async batchAssert({
+    assertions,
+    requestOptions
+  }: {
+    assertions: Array<BatchAssertion>
+    requestOptions?: Partial<RequestOptions>
+  }): Promise<ModifiedResponse<AssertResponse>> {
+    return await this.request(`${this.api_url}/v2/batch/records`, {
+      method: 'put',
+      json: { assertions },
+      ...requestOptions
+    })
+  }
+
   /**
    * List all of the available Objects in the Attio workspace.
    */

packages/destination-actions/src/destinations/attio/assertRecord/__tests__/index.test.ts

@@ -13,7 +13,8 @@ const event = createTestEvent({
   traits: {
     name: 'Stair car',
     number_of_wheels: 4
-  }
+  },
+  receivedAt: '2024-05-24T10:00:00.000Z'
 })
 
 const mapping = {
@@ -26,6 +27,9 @@ const mapping = {
     number_of_wheels: {
       '@path': '$.traits.number_of_wheels'
     }
+  },
+  received_at: {
+    '@path': '$.receivedAt'
   }
 }
 
@@ -109,4 +113,32 @@ describe('Attio.assertRecord', () => {
       })
     })
   })
+
+  it('uses the batch assertion endpoint', async () => {
+    nock('https://api.attio.com')
+      .put('/v2/batch/records', {
+        assertions: [
+          {
+            object: 'vehicles',
+            mode: 'create-or-update',
+            matching_attribute: 'name',
+            multiselect_values: 'append',
+            values: {
+              name: 'Stair car',
+              number_of_wheels: 4
+            },
+            received_at: '2024-05-24T10:00:00.000Z'
+          }
+        ]
+      })
+      .reply(202, '')
+
+    const [response] = await testDestination.testBatchAction('assertRecord', {
+      events: [event],
+      mapping,
+      settings: {}
+    })
+
+    expect(response.status).toBe(202)
+  })
 })

packages/destination-actions/src/destinations/attio/assertRecord/index.ts

@@ -2,7 +2,8 @@ import type { ActionDefinition, DynamicFieldResponse, RequestClient } from '@seg
 import type { InputField } from '@segment/actions-core'
 import type { Settings } from '../generated-types'
 import type { Payload } from './generated-types'
-import { AttioClient } from '../api'
+import { AttioClient, SimpleValue } from '../api'
+import { commonFields } from '../common-fields'
 
 const object: InputField = {
   type: 'string',
@@ -68,7 +69,8 @@ const action: ActionDefinition<Settings, Payload> = {
   fields: {
     object,
     matching_attribute,
-    attributes
+    attributes,
+    ...commonFields
   },
 
   dynamicFields: {
@@ -83,6 +85,21 @@ const action: ActionDefinition<Settings, Payload> = {
       matching_attribute: payload.matching_attribute,
       values: payload.attributes ?? {}
     })
+  },
+
+  performBatch: async (request, { payload }) => {
+    const client = new AttioClient(request)
+
+    return await client.batchAssert({
+      assertions: payload.map((item) => ({
+        object: item.object,
+        mode: 'create-or-update',
+        matching_attribute: item.matching_attribute,
+        multiselect_values: 'append',
+        values: (item.attributes as Record<string, SimpleValue>) ?? {},
+        received_at: item.received_at.toString()
+      }))
+    })
   }
 }
 

packages/destination-actions/src/destinations/attio/common-fields.ts

@@ -0,0 +1,32 @@
+import { ActionDefinition } from '@segment/actions-core'
+import { Settings } from './generated-types'
+
+export const commonFields: ActionDefinition<Settings>['fields'] = {
+  enable_batching: {
+    label: 'Send data to Attio in batches',
+    description: 'Send data to Attio in batches for much better performance.',
+    type: 'boolean',
+    required: false,
+    unsafe_hidden: true,
+    default: true
+  },
+
+  batch_size: {
+    label: 'Batch Size',
+    description: 'Max batch size to send to Attio (limit is 10,000)',
+    type: 'number',
+    required: false,
+    unsafe_hidden: true,
+    default: 1_000
+  },
+
+  received_at: {
+    label: 'Received at',
+    description: 'When the event was received.',
+    type: 'datetime',
+    required: true,
+    default: {
+      '@path': '$.receivedAt'
+    }
+  }
+}

packages/destination-actions/src/destinations/attio/groupWorkspace/__tests__/index.test.ts

@@ -12,13 +12,17 @@ const event = createTestEvent({
   traits: {
     id: '42',
     domain
-  }
+  },
+  receivedAt: '2024-05-24T10:00:00.000Z'
 })
 
 const mapping = {
   domain: { '@path': '$.traits.domain' },
   workspace_id: { '@path': '$.traits.id' },
-  user_id: { '@path': '$.userId' }
+  user_id: { '@path': '$.userId' },
+  received_at: {
+    '@path': '$.receivedAt'
+  }
 }
 
 describe('Attio.groupWorkspace', () => {
@@ -180,4 +184,44 @@ describe('Attio.groupWorkspace', () => {
       })
     ).rejects.toThrowError()
   })
+
+  it('uses the batch assertion endpoint', async () => {
+    nock('https://api.attio.com')
+      .put('/v2/batch/records', {
+        assertions: [
+          {
+            object: 'workspaces',
+            mode: 'create-or-update',
+            matching_attribute: 'workspace_id',
+            multiselect_values: 'append',
+            values: {
+              workspace_id: '42',
+              users: ['user1234'],
+
+              company: {
+                object: 'companies',
+                mode: 'create-or-update',
+                matching_attribute: 'domains',
+                multiselect_values: 'append',
+                values: {
+                  domains: domain
+                },
+                received_at: '2024-05-24T10:00:00.000Z'
+              }
+            },
+            received_at: '2024-05-24T10:00:00.000Z'
+          }
+        ]
+      })
+      .reply(202, '')
+
+    const responses = await testDestination.testBatchAction('groupWorkspace', {
+      events: [event],
+      mapping,
+      settings: {}
+    })
+
+    expect(responses.length).toBe(2)
+    expect(responses[1].status).toBe(202)
+  })
 })

packages/destination-actions/src/destinations/attio/groupWorkspace/index.ts

@@ -3,6 +3,7 @@ import type { InputField } from '@segment/actions-core'
 import type { Settings } from '../generated-types'
 import type { Payload } from './generated-types'
 import { AttioClient } from '../api'
+import { commonFields } from '../common-fields'
 
 const domain: InputField = {
   type: 'string',
@@ -76,7 +77,8 @@ const action: ActionDefinition<Settings, Payload> = {
     workspace_id,
     user_id,
     company_attributes,
-    workspace_attributes
+    workspace_attributes,
+    ...commonFields
   },
 
   perform: async (request, { payload }) => {
@@ -101,6 +103,37 @@ const action: ActionDefinition<Settings, Payload> = {
         ...(payload.workspace_attributes ?? {})
       }
     })
+  },
+
+  performBatch: async (request, { payload }) => {
+    const client = new AttioClient(request)
+
+    return await client.batchAssert({
+      assertions: payload.map((item) => ({
+        object: 'workspaces',
+        mode: 'create-or-update',
+        matching_attribute: 'workspace_id',
+        multiselect_values: 'append',
+        values: {
+          workspace_id: item.workspace_id,
+          ...(item.user_id ? { users: [item.user_id] } : {}),
+          ...(item.workspace_attributes ?? {}),
+
+          company: {
+            object: 'companies',
+            mode: 'create-or-update',
+            matching_attribute: 'domains',
+            multiselect_values: 'append',
+            values: {
+              domains: item.domain,
+              ...(item.company_attributes ?? {})
+            },
+            received_at: item.received_at.toString()
+          }
+        },
+        received_at: item.received_at.toString()
+      }))
+    })
   }
 }
 

packages/destination-actions/src/destinations/attio/identifyUser/__tests__/index.test.ts

@@ -12,7 +12,8 @@ const event = createTestEvent({
   traits: {
     name: 'George Oscar Bluth',
     email
-  }
+  },
+  receivedAt: '2024-05-24T10:00:00.000Z'
 })
 
 const mapping = {
@@ -22,6 +23,9 @@ const mapping = {
     name: {
       '@path': '$.traits.name'
     }
+  },
+  received_at: {
+    '@path': '$.receivedAt'
   }
 }
 
@@ -80,4 +84,45 @@ describe('Attio.identifyUser', () => {
       })
     ).rejects.toThrowError()
   })
+
+  it('uses the batch assertion endpoint', async () => {
+    nock('https://api.attio.com')
+      .put('/v2/batch/records', {
+        assertions: [
+          {
+            object: 'users',
+            mode: 'create-or-update',
+            matching_attribute: 'user_id',
+            multiselect_values: 'append',
+            values: {
+              primary_email_address: email,
+              user_id: event.userId,
+              name: event.traits?.name,
+
+              person: {
+                object: 'people',
+                mode: 'create-or-update',
+                matching_attribute: 'email_addresses',
+                multiselect_values: 'append',
+                values: {
+                  email_addresses: email
+                },
+                received_at: '2024-05-24T10:00:00.000Z'
+              }
+            },
+            received_at: '2024-05-24T10:00:00.000Z'
+          }
+        ]
+      })
+      .reply(202, '')
+
+    const responses = await testDestination.testBatchAction('identifyUser', {
+      events: [event],
+      mapping,
+      settings: {}
+    })
+
+    expect(responses.length).toBe(2)
+    expect(responses[1].status).toBe(202)
+  })
 })