Adding send action for kinesis destination

packages/actions-shared/src/aws-config/index.ts

@@ -0,0 +1,19 @@
+export const APP_AWS_REGION = process.env['AWS_REGION'] || `us-west-2`
+
+export const AWS_REGIONS = [
+  { label: 'us-east-1', value: 'us-east-1' },
+  { label: 'us-east-2', value: 'us-east-2' },
+  { label: 'us-west-1', value: 'us-west-1' },
+  { label: 'us-west-2', value: 'us-west-2' },
+  { label: 'eu-west-1', value: 'eu-west-1' },
+  { label: 'eu-west-2', value: 'eu-west-2' },
+  { label: 'eu-west-3', value: 'eu-west-3' },
+  { label: 'ap-southeast-1', value: 'ap-southeast-1' },
+  { label: 'ap-southeast-2', value: 'ap-southeast-2' },
+  { label: 'sa-east-1', value: 'sa-east-1' },
+  { label: 'ap-northeast-1', value: 'ap-northeast-1' },
+  { label: 'ap-northeast-2', value: 'ap-northeast-2' },
+  { label: 'ap-south-1', value: 'ap-south-1' },
+  { label: 'ca-central-1', value: 'ca-central-1' },
+  { label: 'eu-central-1', value: 'eu-central-1' }
+]

packages/actions-shared/src/index.ts

@@ -7,3 +7,4 @@ export * from './friendbuy/sharedPurchase'
 export * from './friendbuy/sharedSignUp'
 export * from './friendbuy/util'
 export * from './engage/utils'
+export * from './aws-config/index'

packages/destination-actions/package.json

@@ -44,6 +44,7 @@
     "@aws-sdk/client-eventbridge": "^3.741.0",
     "@aws-sdk/client-s3": "^3.600.0",
     "@aws-sdk/client-sts": "3.614.0",
+    "@aws-sdk/client-kinesis": "3.917.0",
     "@bufbuild/protobuf": "^2.2.3",
     "@segment/a1-notation": "^2.1.4",
     "@segment/actions-core": "^3.161.0",

packages/destination-actions/src/destinations/aws-kinesis/__test__/index.test.ts → packages/destination-actions/src/destinations/aws-kinesis/__tests__/index.test.ts

@@ -1,7 +1,7 @@
 import destination from '../index'
 import { assumeRole } from '../../../lib/AWS/sts'
 import { validateIamRoleArnFormat } from '../utils'
-import { APP_AWS_REGION } from '../../../lib/AWS/utils'
+import { APP_AWS_REGION } from '@segment/actions-shared'
 import type { Settings } from '../generated-types'
 import { createTestIntegration } from '../../../../../core/src/create-test-integration'
 
@@ -23,10 +23,6 @@ jest.mock('@segment/actions-core', () => ({
   }))
 }))
 
-jest.mock('../../../lib/AWS/utils', () => ({
-  APP_AWS_REGION: 'us-east-1'
-}))
-
 const testDestination = createTestIntegration(destination)
 
 describe('AWS Kinesis Destination - testAuthentication', () => {

packages/destination-actions/src/destinations/aws-kinesis/__tests__/utils.test.ts

@@ -0,0 +1,141 @@
+import { validateIamRoleArnFormat, send } from '../utils'
+import { Payload } from '../send/generated-types'
+import { KinesisClient, PutRecordsCommand } from '@aws-sdk/client-kinesis'
+import { assumeRole } from '../../../lib/AWS/sts'
+import { IntegrationError } from '@segment/actions-core'
+import { Logger } from '@segment/actions-core/destination-kit'
+
+describe('validateIamRoleArnFormat', () => {
+  it('should return true for a valid IAM Role ARN', () => {
+    const validArns = [
+      'arn:aws:iam::123456789012:role/MyRole',
+      'arn:aws:iam::000000000000:role/service-role/My-Service_Role',
+      'arn:aws:iam::987654321098:role/path/to/MyRole',
+      'arn:aws:iam::111122223333:role/MyRole-With.Special@Chars_+=,.'
+    ]
+
+    for (const arn of validArns) {
+      expect(validateIamRoleArnFormat(arn)).toBe(true)
+    }
+  })
+
+  it('should return false for an ARN with invalid prefix', () => {
+    const invalidArn = 'arn:aws:s3::123456789012:role/MyRole'
+    expect(validateIamRoleArnFormat(invalidArn)).toBe(false)
+  })
+
+  it('should return false if missing account ID', () => {
+    const invalidArn = 'arn:aws:iam:::role/MyRole'
+    expect(validateIamRoleArnFormat(invalidArn)).toBe(false)
+  })
+
+  it('should return false if account ID is not 12 digits', () => {
+    const invalidArns = ['arn:aws:iam::12345:role/MyRole', 'arn:aws:iam::1234567890123:role/MyRole']
+    for (const arn of invalidArns) {
+      expect(validateIamRoleArnFormat(arn)).toBe(false)
+    }
+  })
+
+  it('should return false if missing "role/" segment', () => {
+    const invalidArn = 'arn:aws:iam::123456789012:MyRole'
+    expect(validateIamRoleArnFormat(invalidArn)).toBe(false)
+  })
+
+  it('should return false if role name contains invalid characters', () => {
+    const invalidArns = [
+      'arn:aws:iam::123456789012:role/My Role', // space
+      'arn:aws:iam::123456789012:role/MyRole#InvalidChar'
+    ]
+    for (const arn of invalidArns) {
+      expect(validateIamRoleArnFormat(arn)).toBe(false)
+    }
+  })
+
+  it('should return false for empty or null values', () => {
+    expect(validateIamRoleArnFormat('')).toBe(false)
+    // @ts-expect-error testing invalid input type
+    expect(validateIamRoleArnFormat(null)).toBe(false)
+    // @ts-expect-error testing invalid input type
+    expect(validateIamRoleArnFormat(undefined)).toBe(false)
+  })
+})
+
+jest.mock('@aws-sdk/client-kinesis')
+jest.mock('../../../lib/AWS/sts')
+
+const mockSend = jest.fn()
+const mockLogger: Partial<Logger> = {
+  crit: jest.fn(),
+  info: jest.fn(),
+  warn: jest.fn()
+}
+
+describe('Kinesis send', () => {
+  const mockSettings = {
+    iamRoleArn: 'arn:aws:iam::123456789012:role/TestRole',
+    iamExternalId: 'ext-id'
+  }
+
+  const mockPayloads: Payload[] = [
+    {
+      streamName: 'test-stream',
+      awsRegion: 'us-east-1',
+      partitionKey: 'pk-1',
+      payload: { data: 'test message' },
+      max_batch_size: 500,
+      batch_keys: ['awsRegion']
+    }
+  ]
+
+  beforeEach(() => {
+    jest.clearAllMocks()
+    ;(assumeRole as jest.Mock).mockResolvedValue({
+      accessKeyId: 'mockAccess',
+      secretAccessKey: 'mockSecret',
+      sessionToken: 'mockToken'
+    })
+    ;(KinesisClient as unknown as jest.Mock).mockImplementation(() => ({
+      send: mockSend
+    }))
+  })
+
+  it('should throw IntegrationError if partitionKey is missing', async () => {
+    const invalidPayload = [
+      { ...mockPayloads[0], partitionKey: '' } // missing partitionKey
+    ]
+
+    await expect(send(mockSettings, invalidPayload, undefined, mockLogger as Logger)).rejects.toThrow(IntegrationError)
+
+    expect(mockLogger.crit).not.toHaveBeenCalled()
+  })
+
+  it('should create Kinesis client and send records successfully', async () => {
+    mockSend.mockResolvedValueOnce({ Records: [] })
+
+    await send(mockSettings, mockPayloads, undefined, mockLogger as Logger)
+
+    expect(assumeRole).toHaveBeenCalledWith(
+      mockSettings.iamRoleArn,
+      mockSettings.iamExternalId,
+      expect.any(String) // region
+    )
+
+    expect(KinesisClient).toHaveBeenCalledWith(
+      expect.objectContaining({
+        region: 'us-east-1',
+        credentials: expect.any(Object)
+      })
+    )
+
+    expect(mockSend).toHaveBeenCalledWith(expect.any(PutRecordsCommand))
+  })
+
+  it('should log and rethrow error when Kinesis send fails', async () => {
+    const error = new Error('Kinesis failure')
+    mockSend.mockRejectedValueOnce(error)
+
+    await expect(send(mockSettings, mockPayloads, undefined, mockLogger as Logger)).rejects.toThrow('Kinesis failure')
+
+    expect(mockLogger.crit).toHaveBeenCalledWith('Failed to send batch to Kinesis:', error)
+  })
+})

packages/destination-actions/src/destinations/aws-kinesis/index.ts

@@ -3,7 +3,7 @@ import type { Settings } from './generated-types'
 import { IntegrationError } from '@segment/actions-core'
 import { assumeRole } from '../../lib/AWS/sts'
 import { validateIamRoleArnFormat } from './utils'
-import { APP_AWS_REGION } from '../../lib/AWS/utils'
+import { APP_AWS_REGION } from '@segment/actions-shared'
 
 import send from './send'
 

packages/destination-actions/src/destinations/aws-kinesis/send/index.ts

@@ -1,6 +1,8 @@
 import type { ActionDefinition } from '@segment/actions-core'
 import type { Settings } from '../generated-types'
 import type { Payload } from './generated-types'
+import { AWS_REGIONS } from '@segment/actions-shared'
+import { send } from '../utils'
 
 const action: ActionDefinition<Settings, Payload> = {
   title: 'Send',
@@ -26,18 +28,39 @@ const action: ActionDefinition<Settings, Payload> = {
       description: 'The name of the Kinesis stream to send records to',
       type: 'string',
       required: true,
-      default: { '@path': '$.properties.streamName' }
+      disabledInputMethods: ['variable', 'function', 'freeform', 'enrichment']
     },
     awsRegion: {
       label: 'AWS Region',
       description: 'The AWS region where the Kinesis stream is located',
       type: 'string',
       required: true,
-      default: { '@path': '$.properties.awsRegion' }
+      choices: AWS_REGIONS,
+      disabledInputMethods: ['variable', 'function', 'freeform', 'enrichment']
+    },
+    batch_keys: {
+      label: 'Batch Keys',
+      description: 'The keys to use for batching the events.',
+      type: 'string',
+      unsafe_hidden: true,
+      default: ['awsRegion', 'streamName', 'partitionKey'],
+      multiple: true
+    },
+    max_batch_size: {
+      label: 'Max Batch Size',
+      description: 'The maximum number of payloads to include in a batch.',
+      type: 'number',
+      required: true,
+      minimum: 1,
+      maximum: 500,
+      default: 500
     }
   },
-  perform: async (_request, _data) => {
-    // Todo implement functionality
+  perform: async (_requests, { settings, payload, statsContext, logger }) => {
+    await send(settings, [payload], statsContext, logger)
+  },
+  performBatch: async (_requests, { settings, payload, statsContext, logger }) => {
+    await send(settings, payload, statsContext, logger)
   }
 }