feat: ensure migrations are run in a correct order

src/custom-resource-migrations-runner.ts

@@ -7,6 +7,7 @@ import {
 } from "aws-lambda";
 import { Construct } from "constructs";
 import { $AWS, Function, Table } from "functionless";
+import sortBy from "lodash.sortby";
 import { marshall } from "typesafe-dynamodb/lib/marshall";
 import { MigrationHistoryItem, MigrationStatus } from "./migrations-manager";
 
@@ -50,17 +51,20 @@ export default class CustomResourceMigrationsRunner extends Construct {
 
           console.log({ storedMigrations, migrationIdStateMachinePairs });
 
-          // todo: Ensure chronological order of migrations.
-          const migrationsToRun = migrationIdStateMachinePairs.filter(
-            (migrationStateMachinePair) =>
-              !(storedMigrations.Items ?? []).find(
-                (storedMigration) =>
-                  storedMigration.id.S === migrationStateMachinePair.migrationId
-              )
+          const migrationsToRun = sortBy(
+            migrationIdStateMachinePairs.filter(
+              (migrationStateMachinePair) =>
+                !(storedMigrations.Items ?? []).find(
+                  (storedMigration) =>
+                    storedMigration.id.S ===
+                    migrationStateMachinePair.migrationId
+                )
+            ),
+            // migrationID starts with date
+            "migrationId"
           );
 
-          console.log({ migrationsToRun });
-
+          // Run migrations sequentially
           for (const migration of migrationsToRun) {
             // todo: Depending on the cloudformation transition (success/rollback) we could either use Up or Down state machine
             const command = new StartExecutionCommand({

src/migration.ts

@@ -12,19 +12,16 @@ import {
   $AWS,
   $SFN,
 } from "functionless";
-import { ScanOutput } from "typesafe-dynamodb/lib/scan";
-
-export type ScanTableOptions = {
-  segments: number;
-};
+import { QueryInput, QueryOutput } from "typesafe-dynamodb/lib/query";
+import { ScanOutput, ScanInput } from "typesafe-dynamodb/lib/scan";
 
 export type MigrationProps = {
   tableArn: string;
   migrationName: string;
 };
 
 export type TransformFunctionType<T extends object> = (input: {
-  result: ScanOutput<T, any, any>;
+  result: ScanOutput<T, any, any> | QueryOutput<T, any, any>;
 }) => Promise<any>;
 
 export class Migration<T extends object> extends NestedStack {
@@ -43,12 +40,68 @@ export class Migration<T extends object> extends NestedStack {
     );
   }
 
+  public query(
+    transformFn: TransformFunctionType<T>,
+    options?: QueryInput<T, string, string, string, keyof T, any>
+  ) {
+    // "this" cannot be referenced in a Function.
+    const table = this.table;
+
+    const transformFunction = new Function(
+      this,
+      "MigrationCallbackFunction",
+      transformFn
+    );
+
+    const stateMachine = new StepFunction(
+      this,
+      "MigrationStateMachine",
+      {
+        stateMachineName: this.migrationName,
+        logs: {
+          destination: new LogGroup(this, "MigrationLogGroup", {
+            retention: RetentionDays.ONE_WEEK,
+          }),
+          level: LogLevel.ALL,
+        },
+      },
+      async () => {
+        let lastEvaluatedKey;
+        let firstRun = true;
+
+        while (firstRun || lastEvaluatedKey) {
+          firstRun = false;
+
+          const result: ScanOutput<any, any, any> = await $AWS.DynamoDB.Scan({
+            Table: table,
+            ...options,
+            ExclusiveStartKey: lastEvaluatedKey,
+          });
+
+          if (result.LastEvaluatedKey) {
+            lastEvaluatedKey = result.LastEvaluatedKey;
+          }
+
+          await transformFunction({ result });
+        }
+      }
+    );
+
+    this.stateMachineArn = new CfnOutput(this, "StateMachineArn", {
+      exportName: `${this.migrationName}StateMachineArn`,
+      value: stateMachine.resource.stateMachineArn,
+    });
+
+    return stateMachine;
+  }
+
   // Creates a state machine scanning whole table in parallel and applying transform function to each item.
   public scan(
-    _transformFn: TransformFunctionType<T>,
-    options?: ScanTableOptions
+    transformFn: TransformFunctionType<T>,
+    options?: ScanInput<T, string, string, keyof T, any>
   ) {
-    const totalSegments = options?.segments ?? 10;
+    // By default, use factor of 10 for parallelism
+    const totalSegments = options?.TotalSegments ?? 10;
     const segments = Array.from({ length: totalSegments }, (_, i) => i);
 
     // "this" cannot be referenced in a Function.
@@ -57,7 +110,7 @@ export class Migration<T extends object> extends NestedStack {
     const transformFunction = new Function(
       this,
       "MigrationCallbackFunction",
-      _transformFn
+      transformFn
     );
 
     const stateMachine = new StepFunction(
@@ -83,6 +136,7 @@ export class Migration<T extends object> extends NestedStack {
             const result: ScanOutput<any, any, any> = await $AWS.DynamoDB.Scan({
               Table: table,
               TotalSegments: totalSegments,
+              ...options,
               Segment: index,
               ExclusiveStartKey: lastEvaluatedKey,
             });
@@ -97,8 +151,6 @@ export class Migration<T extends object> extends NestedStack {
       }
     );
 
-    console.log(this.migrationName);
-
     this.stateMachineArn = new CfnOutput(this, "StateMachineArn", {
       exportName: `${this.migrationName}StateMachineArn`,
       value: stateMachine.resource.stateMachineArn,