feat: Additional tasks for batchReplayer to enable reactive pipelines (#629)

* batchReplayer: additional tasks for reactive pipelines

* updated e2e test with additionalStepFunctionTasks

* Update unit tests for batch replayer additional tasks

---------

Co-authored-by: Vincent Gromakowski <[email protected]>

Author: rubenhgaws

core/API.md

@@ -7,7 +7,8 @@ import { SfnStateMachine } from 'aws-cdk-lib/aws-events-targets';
 import { PolicyStatement } from 'aws-cdk-lib/aws-iam';
 import { LayerVersion, Runtime } from 'aws-cdk-lib/aws-lambda';
 import { LogGroup, RetentionDays } from 'aws-cdk-lib/aws-logs';
-import { JsonPath, LogLevel, Map, StateMachine, TaskInput } from 'aws-cdk-lib/aws-stepfunctions';
+import { JsonPath, LogLevel, Map, StateMachine, 
+         TaskInput, INextable, IChainable } from 'aws-cdk-lib/aws-stepfunctions';
 import { LambdaInvoke } from 'aws-cdk-lib/aws-stepfunctions-tasks';
 import { Construct } from 'constructs';
 import { PreBundledFunction } from '../common/pre-bundled-function';
@@ -66,6 +67,34 @@ export interface BatchReplayerProps {
    * VPC for the WriteInBatch Lambda function
    */
   readonly vpc?: IVpc;
+  /**
+   * Additional StupFunction Tasks to run sequentially after the BatchReplayer finishes
+   * @default - The BatchReplayer do not have additional Tasks
+   * 
+   * The expected input for the first Task in this sequence is:
+   * 
+   * input = [
+   *  {
+   *    "processedRecords": Int,
+   *    "outputPaths": String [],
+   *    "startTimeinIso": String,
+   *    "endTimeinIso": String
+   *  }
+   * ]
+   * 
+   * Each element in input represents the output of each lambda iterator that replays the data.
+   * 
+   * param: processedRecods -> Number of records processed
+   * param: ouputPaths -> List of files created in S3 
+   *  **  eg. "s3://<sinkBucket name>/<s3ObjectKeySink prefix, if any>/<dataset name>/ingestion_start=<timestamp>/ingestion_end=<timestamp>/<s3 filename>.csv",
+
+   * param: startTimeinIso -> Start Timestamp on original dataset
+   * param: endTimeinIso -> End Timestamp on original dataset
+   * 
+   * *outputPaths* can be used to extract and aggregate new partitions on data and 
+   * trigger additional Tasks.
+   */
+  readonly additionalStepFunctionTasks?: IChainable [];
 }
 
 /**
@@ -142,6 +171,13 @@ export class BatchReplayer extends Construct {
    */
   public readonly vpc?: IVpc;
 
+  /**
+   * Optional Sequence of additional Tasks to append at the end of the Step Function
+   * that replays data that will execute after data has been replayed
+   */
+  private readonly additionalStepFunctionTasks?: IChainable [];
+
+
   /**
    * Constructs a new instance of the BatchReplayer construct
    * @param {Construct} scope the Scope of the CDK Construct
@@ -153,6 +189,7 @@ export class BatchReplayer extends Construct {
 
     this.dataset = props.dataset;
     this.frequency = props.frequency?.toSeconds() || 60;
+    this.additionalStepFunctionTasks = props.additionalStepFunctionTasks;
 
     // Properties for S3 target
     if (props.s3Props) {
@@ -337,7 +374,9 @@ export class BatchReplayer extends Construct {
 
     // Overarching Step Function StateMachine
     const batchReplayStepFn = new StateMachine(this, 'BatchReplayStepFn', {
-      definition: findFilePathsFnTask.next(writeInBatchMapTask),
+      definition: this.chainStepFunctionTasks(
+          findFilePathsFnTask.next(writeInBatchMapTask)
+        ),
       timeout: Duration.minutes(20),
       logs: {
         destination: new LogGroup(this, 'LogGroup', {
@@ -354,4 +393,18 @@ export class BatchReplayer extends Construct {
       targets: [new SfnStateMachine(batchReplayStepFn, {})],
     });
   }
+
+  private chainStepFunctionTasks(requiredTasks: IChainable & INextable) {
+  
+    let base = requiredTasks;
+
+    if (this.additionalStepFunctionTasks) {
+    
+      this.additionalStepFunctionTasks.forEach(newTask => {
+        base = base.next(newTask)
+      });
+    }
+    return base;
+  }
+
 }

core/package.json

@@ -8,13 +8,17 @@
 */
 
 import { Bucket } from 'aws-cdk-lib/aws-s3';
-import { App, Stack, aws_ec2, aws_rds, aws_dynamodb, RemovalPolicy, CfnOutput } from 'aws-cdk-lib';
+import { App, Stack, aws_ec2, aws_rds, aws_stepfunctions_tasks, 
+        aws_dynamodb, RemovalPolicy, CfnOutput, Duration } from 'aws-cdk-lib';
 import { Cluster } from '@aws-cdk/aws-redshift-alpha';
 import { deployStack, destroyStack } from './utils';
 
 import { BatchReplayer } from '../../src/data-generator/batch-replayer';
 import { IS3Sink, DynamoDbSink, DbSink } from '../../src/data-generator/batch-replayer-helpers';
 import { PreparedDataset } from '../../src/data-generator/prepared-dataset';
+import { PreBundledFunction } from '../../src/common/pre-bundled-function';
+import { Runtime } from 'aws-cdk-lib/aws-lambda';
+import { RetentionDays } from 'aws-cdk-lib/aws-logs';
 
 jest.setTimeout(3000000);
 // GIVEN
@@ -66,6 +70,29 @@ const rdsPostgres = new aws_rds.DatabaseInstance(stack, 'PostgreSQL', {
 const rdsPostgresCreds = rdsPostgres.secret ? rdsPostgres.secret.secretArn : '';
 let rdsProps: DbSink = { table: defaultName, connection: rdsPostgresCreds, schema: defaultName, type: 'postgresql' };
 
+/**
+ * Lambda to be used as additional task
+ */
+const summaryOutputState = new PreBundledFunction(stack, 'WriteInBatch', {
+  memorySize: 128,
+  codePath: 'data-generator/resources/lambdas/summary-additional-task',
+  runtime: Runtime.PYTHON_3_9,
+  handler: 'summary-additional-task.handler',
+  logRetention: RetentionDays.ONE_WEEK,
+  timeout: Duration.minutes(1),
+});
+
+/** 
+ * Additional Task
+*/
+const additionalTask = new aws_stepfunctions_tasks
+        .LambdaInvoke(stack, "SummarizeOutput",  {
+            lambdaFunction: summaryOutputState,
+            outputPath: "$",
+            retryOnServiceExceptions: true
+            });
+
+
 const batchReplayer = new BatchReplayer(stack, 'BatchReplay', {
   dataset: PreparedDataset.RETAIL_1_GB_STORE_SALE,
   s3Props: s3Props,
@@ -75,6 +102,7 @@ const batchReplayer = new BatchReplayer(stack, 'BatchReplay', {
   rdsProps: rdsProps,
   vpc: vpc,
   secGroup: secGroup,
+  additionalStepFunctionTasks: [additionalTask]
 });
 
 new BatchReplayer(stack, 'BatchReplay2', {
@@ -86,6 +114,7 @@ new BatchReplayer(stack, 'BatchReplay2', {
   rdsProps: rdsProps,
   vpc: vpc,
   secGroup: secGroup,
+  additionalStepFunctionTasks: [additionalTask]
 });
 
 new CfnOutput(stack, 'DatasetName', {

core/src/data-generator/batch-replayer.ts

@@ -0,0 +1,4 @@
+os
+json
+logging
+typing

core/test/e2e/batch-replayer.test.ts

@@ -0,0 +1,75 @@
+
+import os
+import json
+import logging
+from typing import Set, Any, Dict
+
+logging.basicConfig(level=logging.INFO)
+log = logging.getLogger()
+log.setLevel(os.getenv('LOG_LEVEL', logging.INFO))
+
+default_value = 'NOTFOUND'
+
+def get_dict_value(dictionary: dict, key: str, nullable:bool = False):
+    val = dictionary.get(key, default_value)
+    if not nullable and val == default_value:
+        raise Exception(f"Key: {key} not found in {dictionary}")
+    
+    log.info(f"Got parameter: {key}:{val}")
+
+    return val
+
+def parseLambdaIterationOutput(event: dict)->Dict[str,Any]:
+    processedRecords: int = event.get('processedRecords', 0)
+    startTimeinIso: str = event.get('startTimeinIso', "")
+    endTimeinIso: str = event.get('endTimeinIso', "")
+
+    return   {
+        "processedRecords": processedRecords,
+        "startTimeinIso": startTimeinIso,
+        "endTimeinIso": endTimeinIso,
+    }
+
+
+
+def handler(outputList, ctx):
+    processedRecords: int = 0
+    startTimeinIso: Set[str] = set()
+    endTimeinIso: Set[str] = set()
+    
+    for event in outputList:
+
+        resp = parseLambdaIterationOutput(event)
+        
+        processedRecords += resp.get("processedRecords", 0)
+        startTimeinIso.add(resp.get('startTimeinIso', ""))
+        endTimeinIso  .add(resp.get('endTimeinIso',   ""))
+
+
+    return   {
+        "processedRecords": processedRecords,
+        "startTimeinIso": min(list(startTimeinIso)),
+        "endTimeinIso": max(list(endTimeinIso)),
+        }
+
+# Test from local machine
+if __name__ == '__main__':
+    event=[
+        {
+            "processedRecords": 122,
+            "startTimeinIso": "2021-01-01T00:03:34",
+            "endTimeinIso": "2021-01-01T00:13:34"
+        },
+        {
+            "processedRecords": 389,
+            "startTimeinIso": "2021-01-01T00:03:34",
+            "endTimeinIso": "2021-01-01T00:13:34"
+        },
+        {
+            "processedRecords": 221,
+            "startTimeinIso": "2021-01-01T00:03:34",
+            "endTimeinIso": "2021-01-01T00:13:34"
+        }
+    ]
+    result = handler(event,None)
+    log.info(f'result = {json.dumps(result, indent=2)}')

core/test/e2e/resources/lambdas/summary-additional-task/requirements.txt

@@ -11,7 +11,7 @@ import { Cluster } from '@aws-cdk/aws-redshift-alpha';
 import {
   aws_dynamodb,
   aws_ec2,
-  aws_rds, Duration,
+  aws_rds, aws_stepfunctions_tasks, Duration,
   RemovalPolicy,
   Stack,
 } from 'aws-cdk-lib';
@@ -20,6 +20,9 @@ import { Template } from 'aws-cdk-lib/assertions';
 import { Bucket } from 'aws-cdk-lib/aws-s3';
 import { BatchReplayer, DbSink, DynamoDbSink, IS3Sink, PreparedDataset } from '../../../src';
 import { IVpc, SecurityGroup } from 'aws-cdk-lib/aws-ec2';
+import { PreBundledFunction } from '../../../src/common/pre-bundled-function';
+import { Runtime } from 'aws-cdk-lib/aws-lambda';
+import { RetentionDays } from 'aws-cdk-lib/aws-logs';
 
 let testStack: Stack;
 let s3Props: IS3Sink;
@@ -32,6 +35,7 @@ let rdsProps: DbSink;
 let defaultName = 'test';
 let vpc: IVpc;
 let secGroup: SecurityGroup;
+const expectedAdditionalTaskName = "SummarizeOutput"
 
 beforeEach(() => {
   testStack = new Stack();
@@ -138,4 +142,65 @@ test("BatchReplayer should create a step function", () => {
   template.resourceCountIs("AWS::StepFunctions::StateMachine", 1);
 });
 
+test("BatchReplayer should not have additionalStepFunctionTasks", ()=> {
+
+  const resources = template.toJSON().Resources
+  const stepFnDefinition:any[] = resources.TestBatchReplayerBatchReplayStepFnBB59B3E9.Properties
+  .DefinitionString["Fn::Join"][1]
+
+  // filter strings that contains : expectedAdditionalTaskName
+  const found = stepFnDefinition.map<String>((a) => (typeof a === 'string' || a instanceof String) ? a : "")
+                    .filter((a:String)=>  a.includes(expectedAdditionalTaskName))
+                    
+  expect(found.length).toBe(0)
+
+});
+
+test("BatchReplayer should have one additionalStepFunctionTask : expectedAdditionalTaskName", ()=> {
+  const testStack2 = new Stack();
+
+  /**
+   * Lambda to be used as additional task
+   */
+  const summaryOutputState = new PreBundledFunction(testStack2, 'SummaryFn', {
+    memorySize: 128,
+    codePath: '../test/unit/data-generator/resources/lambdas/summary-additional-task',
+    runtime: Runtime.PYTHON_3_9,
+    handler: 'summary-additional-task.handler',
+    logRetention: RetentionDays.ONE_WEEK,
+    timeout: Duration.minutes(1),
+  });
+
+  /** 
+   * Additional Task
+  */
+  const additionalTask = new aws_stepfunctions_tasks
+          .LambdaInvoke(testStack2, expectedAdditionalTaskName,  {
+              lambdaFunction: summaryOutputState,
+              outputPath: "$",
+              retryOnServiceExceptions: true
+              });
+
+  const bucket2 = new Bucket(testStack2, 'Bucket');
+  const s3Props2 = { sinkBucket: bucket2 };
+
+  new BatchReplayer(testStack2, 'TestBatchReplayer', {
+    dataset: PreparedDataset.RETAIL_1_GB_WEB_SALE,
+    frequency: Duration.seconds(120),
+    s3Props: s3Props2,
+    additionalStepFunctionTasks: [additionalTask]
+  });
+
+  const template2 = Template.fromStack(testStack2);
+
+  const resources = template2.toJSON().Resources
+  const stepFnDefinition: any[] = resources.TestBatchReplayerBatchReplayStepFnBB59B3E9.Properties
+                          .DefinitionString["Fn::Join"][1]
+
+  // filter strings that contains : expectedAdditionalTaskName
+  const found = stepFnDefinition.map<String>((a) => (typeof a === 'string' || a instanceof String) ? a : "")
+                    .filter((a:String)=>  a.includes(expectedAdditionalTaskName))
+
+  expect(found.length > 0).toBeTruthy()
+});
 

core/test/e2e/resources/lambdas/summary-additional-task/summary-additional-task.py

@@ -0,0 +1,4 @@
+os
+json
+logging
+typing