Add batching util for SQS sink.

petermyers · petermyers · commit abe53f509ff1 · 2025-07-29T17:21:34.000-04:00
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "aws-lambda-stream",
-  "version": "1.1.3",
+  "version": "1.1.9",
   "description": "Create stream processors with AWS Lambda functions.",
   "keywords": [
     "aws",
diff --git a/src/connectors/s3.js b/src/connectors/s3.js
@@ -3,7 +3,7 @@
 import { Readable } from 'stream';
 import {
   CopyObjectCommand,
-  DeleteObjectCommand, GetObjectCommand, ListObjectsV2Command, PutObjectCommand, S3Client,
+  DeleteObjectCommand, GetObjectCommand, HeadObjectCommand, ListObjectsV2Command, PutObjectCommand, S3Client,
 } from '@aws-sdk/client-s3';
 import { NodeHttpHandler } from '@smithy/node-http-handler';
 import Promise from 'bluebird';
@@ -47,6 +47,15 @@ class Connector {
     return this.clients[pipelineId];
   }
 
+  headObject(inputParams, ctx) {
+    const params = {
+      Bucket: this.bucketName,
+      ...inputParams,
+    };
+
+    return this._sendCommand(new HeadObjectCommand(params), ctx);
+  }
+
   putObject(inputParams, ctx) {
     const params = {
       Bucket: this.bucketName,
@@ -75,6 +84,16 @@ class Connector {
       .then(async (response) => ({ ...response, Body: await response.Body.transformToString() }));
   }
 
+  getObjectAsByteArray(inputParams, ctx) {
+    const params = {
+      Bucket: this.bucketName,
+      ...inputParams,
+    };
+
+    return this._sendCommand(new GetObjectCommand(params), ctx)
+      .then(async (response) => ({ ...response, Body: await response.Body.transformToByteArray() }));
+  }
+
   getObjectStream(inputParams, ctx) {
     const params = {
       Bucket: this.bucketName,
diff --git a/src/flavors/circuitbreaker.js b/src/flavors/circuitbreaker.js
@@ -15,6 +15,9 @@ const toUpdateRequest = (opt) => faulty((uow) => ({
     UUID: process.env.ESM_ID, // Ref: TriggerEventSourceMappingDynamodbEntitiesTable
     Enabled: uow.event.alarmData.state.value !== 'ALARM',
     BatchSize: uow.event.alarmData.state.value === 'INSUFFICIENT_DATA'
-      ? 1 : process.env.BATCH_SIZE || 100,
+      ? 1 : Number(opt.circuitBreakerMaxBatchSize)
+      || Number(process.env.CIRCUIT_BREAKER_MAX_BATCH_SIZE)
+      || Number(process.env.BATCH_SIZE)
+      || 100,
   },
 }));
diff --git a/src/queries/s3.js b/src/queries/s3.js
@@ -80,6 +80,35 @@ export const getObjectFromS3AsStream = ({
     .flatMap(getObject);
 };
 
+export const getObjectFromS3AsByteArray = ({
+  id: pipelineId,
+  debug = d('s3'),
+  bucketName = process.env.BUCKET_NAME,
+  getRequestField = 'getRequest',
+  getResponseField = 'getResponse',
+  parallel = Number(process.env.S3_PARALLEL) || Number(process.env.PARALLEL) || 8,
+  step = 'get',
+  ...opt
+} = {}) => {
+  const connector = new Connector({
+    pipelineId, debug, bucketName, ...opt,
+  });
+
+  const getObject = (uow) => {
+    if (!uow[getRequestField]) return _(Promise.resolve(uow));
+
+    const p = () => connector.getObjectAsByteArray(uow[getRequestField], uow)
+      .then((getResponse) => ({ ...uow, [getResponseField]: getResponse })) // TODO decompress
+      .catch(rejectWithFault(uow));
+
+    return _(uow.metrics?.w(p, step) || p()); // wrap promise in a stream
+  };
+
+  return (s) => s
+    .map(getObject)
+    .parallel(parallel);
+};
+
 export const splitS3Object = ({
   delimiter = '\n',
   getResponseField = 'getResponse',
@@ -193,3 +222,32 @@ export const pageObjectsFromS3 = ({
     .map(listObjects)
     .parallel(parallel);
 };
+
+export const headS3Object = ({
+  id: pipelineId,
+  debug = d('s3'),
+  bucketName = process.env.BUCKET_NAME,
+  headRequestField = 'headRequest',
+  headResponseField = 'headResponse',
+  parallel = Number(process.env.S3_PARALLEL) || Number(process.env.PARALLEL) || 8,
+  step = 'get',
+  ...opt
+} = {}) => {
+  const connector = new Connector({
+    pipelineId, debug, bucketName, ...opt,
+  });
+
+  const headObject = (uow) => {
+    if (!uow[headRequestField]) return _(Promise.resolve(uow));
+
+    const p = () => connector.headObject(uow[headRequestField], uow)
+      .then((headResponse) => ({ ...uow, [headResponseField]: headResponse }))
+      .catch(rejectWithFault(uow));
+
+    return _(uow.metrics?.w(p, step) || p()); // wrap promise in a stream
+  };
+
+  return (s) => s
+    .map(headObject)
+    .parallel(parallel);
+};
diff --git a/src/utils/encryption.js b/src/utils/encryption.js
@@ -128,10 +128,11 @@ export const encryptEvent = ({
       return _(Promise.resolve(uow));
     }
 
+    const eemMetadata = typeof eem === 'function' ? eem(uow[sourceField], uow) : eem;
     const p = encryptObject(uow[sourceField], {
       masterKeyAlias,
       regions,
-      ...eem, // fields and overrides
+      ...eemMetadata, // fields and overrides
       AES,
     })
       // .tap(debug)
@@ -165,11 +166,12 @@ export const encryptData = ({
   masterKeyAlias = process.env.MASTER_KEY_ALIAS,
   regions = (process.env.KMS_REGIONS && process.env.KMS_REGIONS.split(',')),
   AES = true,
-} = {}) => async (data) => {
+} = {}) => async (data, ctx = {}) => {
+  const eemMetadata = typeof eem === 'function' ? eem(data, ctx) : eem;
   const result = await encryptObject(data, {
     masterKeyAlias,
     regions,
-    ...eem, // fields and overrides
+    ...eemMetadata, // fields and overrides
     AES,
   })
     // .tap(debug)
diff --git a/test/unit/connectors/s3.test.js b/test/unit/connectors/s3.test.js
@@ -6,7 +6,7 @@ import { Readable } from 'stream';
 import { mockClient } from 'aws-sdk-client-mock';
 import {
   CopyObjectCommand,
-  DeleteObjectCommand, GetObjectCommand, ListObjectsV2Command, PutObjectCommand, S3Client,
+  DeleteObjectCommand, GetObjectCommand, HeadObjectCommand, ListObjectsV2Command, PutObjectCommand, S3Client,
 } from '@aws-sdk/client-s3';
 import { sdkStreamMixin } from '@smithy/util-stream';
 
@@ -95,6 +95,26 @@ describe('connectors/s3.js', () => {
     });
     expect(data).to.deep.equal({ Body: 'b' });
   });
+  it('should get object as byte array', async () => {
+    const arr = new Uint8Array([104, 101, 108, 108, 111]);
+    const spy = sinon.spy(() => ({ Body: sdkStreamMixin(Readable.from('hello')) }));
+    mockS3.on(GetObjectCommand).callsFake(spy);
+
+    const inputParams = {
+      Key: 'k1',
+    };
+
+    const data = await new Connector({
+      debug: debug('s3'),
+      bucketName: 'b1',
+    }).getObjectAsByteArray(inputParams);
+
+    expect(spy).to.have.been.calledWith({
+      Bucket: 'b1',
+      Key: 'k1',
+    });
+    expect(data).to.deep.equal({ Body: arr });
+  });
 
   it('should get object as stream', (done) => {
     const spy = sinon.spy(() => ({ Body: sdkStreamMixin(Readable.from(Buffer.from('data'))) }));
@@ -199,4 +219,20 @@ describe('connectors/s3.js', () => {
     });
     expect(data).to.deep.equal({});
   });
+  it('should head object', async () => {
+    const spy = sinon.spy(() => ({}));
+    mockS3.on(HeadObjectCommand).callsFake(spy);
+    const inputParams = {
+      Key: 'k1',
+    };
+    const data = await new Connector({
+      debug: debug('s3'),
+      bucketName: 'b1',
+    }).headObject(inputParams);
+    expect(spy).to.have.been.calledWith({
+      Key: 'k1',
+      Bucket: 'b1',
+    });
+    expect(data).to.deep.equal({});
+  });
 });
diff --git a/test/unit/flavors/circuitbreaker.test.js b/test/unit/flavors/circuitbreaker.test.js
@@ -2,7 +2,7 @@ import 'mocha';
 import { expect } from 'chai';
 import sinon from 'sinon';
 
-import { initialize } from '../../../src';
+import { initialize, initializeFrom } from '../../../src';
 
 import { fromAlarm } from '../../../src/from/cw';
 
@@ -125,4 +125,85 @@ describe('flavors/circuitbreaker.js', () => {
       })
       .done(done);
   });
+
+  it('should enable event source mapping with batch size from opt)', (done) => {
+    const rules = [
+      {
+        id: 'circuitBreaker',
+        flavor: circuitBreaker,
+        circuitBreakerMaxBatchSize: '10',
+      },
+    ];
+    const event = {
+      source: 'aws.cloudwatch',
+      alarmArn:
+        'arn:aws:cloudwatch:us-east-1:444455556666:alarm:lambda-demo-metric-alarm',
+      accountId: '444455556666',
+      time: '2023-08-04T12:36:15.490+0000',
+      region: 'us-east-1',
+      alarmData: {
+        alarmName: 'lambda-demo-metric-alarm',
+        state: {
+          value: 'OK',
+        },
+        previousState: {
+          value: 'INSUFFICIENT_DATA',
+        },
+      },
+    };
+
+    initialize({
+      ...initializeFrom(rules),
+    })
+      .assemble(fromAlarm(event), false)
+      .collect()
+      // .tap((collected) => console.log(JSON.stringify(collected, null, 2)))
+      .tap((collected) => {
+        expect(collected.length).to.equal(1);
+        expect(collected[0].pipeline).to.equal('circuitBreaker');
+        expect(collected[0].updateRequest).to.deep.equal({
+          UUID: 'a092f90d-9948-4964-95b5-32c46093f734',
+          Enabled: true,
+          BatchSize: 10,
+        });
+      })
+      .done(done);
+  });
+  it('should enable event source mapping and force Number type for env var of CIRCUIT_BREAKER_MAX_BATCH_SIZE', (done) => {
+    process.env.CIRCUIT_BREAKER_MAX_BATCH_SIZE = '10';
+    const event = {
+      source: 'aws.cloudwatch',
+      alarmArn:
+        'arn:aws:cloudwatch:us-east-1:444455556666:alarm:lambda-demo-metric-alarm',
+      accountId: '444455556666',
+      time: '2023-08-04T12:36:15.490+0000',
+      region: 'us-east-1',
+      alarmData: {
+        alarmName: 'lambda-demo-metric-alarm',
+        state: {
+          value: 'OK',
+        },
+        previousState: {
+          value: 'INSUFFICIENT_DATA',
+        },
+      },
+    };
+
+    initialize({
+      circuitBreaker,
+    })
+      .assemble(fromAlarm(event), false)
+      .collect()
+      // .tap((collected) => console.log(JSON.stringify(collected, null, 2)))
+      .tap((collected) => {
+        expect(collected.length).to.equal(1);
+        expect(collected[0].pipeline).to.equal('circuitBreaker');
+        expect(collected[0].updateRequest).to.deep.equal({
+          UUID: 'a092f90d-9948-4964-95b5-32c46093f734',
+          Enabled: true,
+          BatchSize: 10,
+        });
+      })
+      .done(done);
+  });
 });
diff --git a/test/unit/queries/s3.test.js b/test/unit/queries/s3.test.js
diff --git a/test/unit/utils/encryption.test.js b/test/unit/utils/encryption.test.js

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "aws-lambda-stream",`
`3`		`- "version": "1.1.3",`
	`3`	`+ "version": "1.1.9",`
`4`	`4`	`"description": "Create stream processors with AWS Lambda functions.",`
`5`	`5`	`"keywords": [`
`6`	`6`	`"aws",`