Batch to total request message.

Author: petermyers

package-lock.json

@@ -1,6 +1,6 @@
 {
   "name": "aws-lambda-stream",
-  "version": "1.1.2",
+  "version": "1.1.3",
   "description": "Create stream processors with AWS Lambda functions.",
   "keywords": [
     "aws",

package.json

@@ -2,17 +2,19 @@ import _ from 'highland';
 
 import Connector from '../connectors/sqs';
 
-import { toBatchUow, unBatchUow } from '../utils/batch';
+import { batchWithSize, toBatchUow, unBatchUow } from '../utils/batch';
 import { ratelimit } from '../utils/ratelimit';
 import { rejectWithFault } from '../utils/faults';
 import { debug as d } from '../utils/print';
+import { storeClaimcheck } from './claimcheck';
 
 export const sendToSqs = ({ // eslint-disable-line import/prefer-default-export
   id: pipelineId,
   debug = d('sqs'),
   queueUrl = process.env.QUEUE_URL,
   messageField = 'message',
   batchSize = Number(process.env.SQS_BATCH_SIZE) || Number(process.env.BATCH_SIZE) || 10,
+  maxPublishRequestSize = Number(process.env.PUBLISH_MAX_REQ_SIZE) || Number(process.env.MAX_REQ_SIZE) || 256 * 1024,
   parallel = Number(process.env.SQS_PARALLEL) || Number(process.env.PARALLEL) || 8,
   step = 'send',
   ...opt
@@ -46,7 +48,15 @@ export const sendToSqs = ({ // eslint-disable-line import/prefer-default-export
   return (s) => s
     .through(ratelimit(opt))
 
-    .batch(batchSize)
+    .consume(batchWithSize({
+      ...opt,
+      batchSize,
+      maxRequestSize: maxPublishRequestSize,
+      requestEntryField: messageField,
+      claimcheckEventField: 'MessageBody',
+      debug,
+    }))
+    .through(storeClaimcheck(opt))
     .map(toBatchUow)
 
     .map(toInputParams)

src/sinks/sqs.js

@@ -51,6 +51,8 @@ export const compact = (rule) => {
 export const batchWithSize = ({
   claimCheckBucketName = process.env.CLAIMCHECK_BUCKET_NAME,
   putClaimcheckRequest = 'putClaimcheckRequest',
+  // Detail handles EB, but others like SQS may need something else like 'MessageBody'
+  claimcheckEventField = 'Detail',
   ...opt
 }) => {
   let batched = [];
@@ -77,9 +79,9 @@ export const batchWithSize = ({
           logMetrics([x], [size], opt);
           if (claimCheckBucketName) {
             // setup claim check
-            x[putClaimcheckRequest] = toPutClaimcheckRequest(JSON.parse(x[opt.requestEntryField].Detail), claimCheckBucketName);
-            x[opt.requestEntryField].Detail = JSON.stringify(toClaimcheckEvent(
-              JSON.parse(x[opt.requestEntryField].Detail),
+            x[putClaimcheckRequest] = toPutClaimcheckRequest(JSON.parse(x[opt.requestEntryField][claimcheckEventField]), claimCheckBucketName);
+            x[opt.requestEntryField][claimcheckEventField] = JSON.stringify(toClaimcheckEvent(
+              JSON.parse(x[opt.requestEntryField][claimcheckEventField]),
               claimCheckBucketName,
             ));
             size = Buffer.byteLength(JSON.stringify(x[opt.requestEntryField]));

src/utils/batch.js

@@ -18,6 +18,11 @@ describe('sinks/sqs.js', () => {
         Id: '1',
         MessageBody: JSON.stringify({ f1: 'v1' }),
       },
+    }, {
+      message: {
+        Id: '2',
+        MessageBody: JSON.stringify({ f1: 'v2' }),
+      },
     }];
 
     _(uows)
@@ -26,7 +31,7 @@ describe('sinks/sqs.js', () => {
       .tap((collected) => {
         // console.log(JSON.stringify(collected, null, 2));
 
-        expect(collected.length).to.equal(1);
+        expect(collected.length).to.equal(2);
         expect(collected[0]).to.deep.equal({
           message: {
             Id: '1',
@@ -36,6 +41,79 @@ describe('sinks/sqs.js', () => {
             Entries: [{
               Id: '1',
               MessageBody: JSON.stringify({ f1: 'v1' }),
+            }, {
+              Id: '2',
+              MessageBody: JSON.stringify({ f1: 'v2' }),
+            }],
+          },
+          sendMessageBatchResponse: {},
+        });
+        expect(collected[1]).to.deep.equal({
+          message: {
+            Id: '2',
+            MessageBody: JSON.stringify({ f1: 'v2' }),
+          },
+          inputParams: {
+            Entries: [{
+              Id: '1',
+              MessageBody: JSON.stringify({ f1: 'v1' }),
+            }, {
+              Id: '2',
+              MessageBody: JSON.stringify({ f1: 'v2' }),
+            }],
+          },
+          sendMessageBatchResponse: {},
+        });
+      })
+      .done(done);
+  });
+
+  it('should split a batch due to size constraints', (done) => {
+    sinon.stub(Connector.prototype, 'sendMessageBatch').resolves({});
+
+    const uows = [{
+      message: {
+        Id: '1',
+        MessageBody: JSON.stringify({ f1: 'v1' }),
+      },
+    }, {
+      message: {
+        Id: '2',
+        MessageBody: JSON.stringify({ f1: 'v2' }),
+      },
+    }];
+
+    _(uows)
+      .through(sendToSqs({
+        maxPublishRequestSize: 50,
+      }))
+      .collect()
+      .tap((collected) => {
+        // console.log(JSON.stringify(collected, null, 2));
+
+        expect(collected.length).to.equal(2);
+        expect(collected[0]).to.deep.equal({
+          message: {
+            Id: '1',
+            MessageBody: JSON.stringify({ f1: 'v1' }),
+          },
+          inputParams: {
+            Entries: [{
+              Id: '1',
+              MessageBody: JSON.stringify({ f1: 'v1' }),
+            }],
+          },
+          sendMessageBatchResponse: {},
+        });
+        expect(collected[1]).to.deep.equal({
+          message: {
+            Id: '2',
+            MessageBody: JSON.stringify({ f1: 'v2' }),
+          },
+          inputParams: {
+            Entries: [{
+              Id: '2',
+              MessageBody: JSON.stringify({ f1: 'v2' }),
             }],
           },
           sendMessageBatchResponse: {},