improve transactions typing, chaining & merge multiple update item operations

Author: Moicky

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@moicky/dynamodb",
-  "version": "3.0.5",
+  "version": "3.1.0",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "description": "Contains a collection of convenience functions for working with AWS DynamoDB",

src/transactions/index.ts

@@ -50,6 +50,7 @@ export class Transaction {
   private updatedAt: any;
   private operations: { [itemKey: string]: ItemOperation } = {};
   private shouldSplitTransactions: boolean;
+  private hasBeenExecuted: boolean = false;
 
   constructor({
     tableName,
@@ -79,6 +80,12 @@ export class Transaction {
     item: WithoutReferences<T>,
     args?: CreateOperation["args"]
   ) {
+    if (this.hasBeenExecuted) {
+      throw new Error(
+        "[@moicky/dynamodb]: Transaction has already been executed"
+      );
+    }
+
     const itemKey = this.getItemKey(item, args?.TableName);
 
     const createOperation: CreateOperation = {
@@ -92,15 +99,30 @@ export class Transaction {
     return new CreateOperations<T>(createOperation, this);
   }
   update<T extends DynamoDBItemKey>(
-    item: OnlyKey<T>,
+    item: T extends any ? OnlyKey<T> : T,
     args?: UpdateOperation["args"]
   ) {
+    if (this.hasBeenExecuted) {
+      throw new Error(
+        "[@moicky/dynamodb]: Transaction has already been executed"
+      );
+    }
+
     const itemKey = this.getItemKey(item, args?.TableName);
 
+    const actions: UpdateAction[] = [
+      { _type: "set", values: { updatedAt: this.updatedAt } },
+    ];
+
+    const existingOperation = this.operations[itemKey];
+    if (existingOperation && existingOperation._type === "update") {
+      actions.push(...existingOperation.actions);
+    }
+
     const updateOperation: UpdateOperation = {
       _type: "update",
       item,
-      actions: [{ _type: "set", values: { updatedAt: this.updatedAt } }],
+      actions,
       args: { TableName: this.tableName, ...args },
     };
 
@@ -109,19 +131,33 @@ export class Transaction {
     return new UpdateOperations<T>(updateOperation, this);
   }
   delete(item: DynamoDBItemKey, args?: DeleteOperation["args"]) {
+    if (this.hasBeenExecuted) {
+      throw new Error(
+        "[@moicky/dynamodb]: Transaction has already been executed"
+      );
+    }
+
     const itemKey = this.getItemKey(item, args?.TableName);
 
     this.operations[itemKey] = {
       _type: "delete",
       item,
       args: { TableName: this.tableName, ...args },
     };
+
+    return this;
   }
   addConditionFor<T extends DynamoDBItemKey>(
     item: T,
     args?: Partial<ConditionOperation["args"]>
   ) {
-    return new ConditionOperations<T>(this.operations, item, {
+    if (this.hasBeenExecuted) {
+      throw new Error(
+        "[@moicky/dynamodb]: Transaction has already been executed"
+      );
+    }
+
+    return new ConditionOperations<T>(this, this.operations, item, {
       TableName: this.tableName,
       ...args,
     });
@@ -234,6 +270,16 @@ export class Transaction {
 
     return new Promise<Record<string, ResponseItem[]>>(
       async (resolve, reject) => {
+        if (this.hasBeenExecuted) {
+          reject(
+            new Error(
+              "[@moicky/dynamodb]: Transaction has already been executed"
+            )
+          );
+        }
+
+        this.hasBeenExecuted = true;
+
         const operations = Object.values(this.operations);
 
         if (

src/transactions/operations.ts

@@ -1,3 +1,4 @@
+import { TransactWriteItemsCommandInput } from "@aws-sdk/client-dynamodb";
 import { Transaction } from ".";
 import {
   getAttributeNames,
@@ -39,7 +40,7 @@ export class CreateOperations<U extends ItemWithKey> {
 
       const refData =
         ref instanceof Set
-          ? [...ref].map((references) =>
+          ? Array.from(ref).map((references) =>
               createReference({ references, ...refArgs }, this.transaction)
             )
           : createReference({ references: ref, ...refArgs }, this.transaction);
@@ -55,6 +56,12 @@ export class CreateOperations<U extends ItemWithKey> {
       }, this.operation.item);
     });
   }
+
+  execute(
+    args?: Partial<Omit<TransactWriteItemsCommandInput, "TransactItems">>
+  ) {
+    return this.transaction.execute(args);
+  }
 }
 
 export class UpdateOperations<U extends DynamoDBItem> {
@@ -77,7 +84,7 @@ export class UpdateOperations<U extends DynamoDBItem> {
         values: {
           [attributeName]:
             ref instanceof Set
-              ? [...ref].map((references) =>
+              ? Array.from(ref).map((references) =>
                   createReference({ references, ...refArgs }, this.transaction)
                 )
               : createReference(
@@ -147,10 +154,17 @@ export class UpdateOperations<U extends DynamoDBItem> {
       }, {}),
     };
   }
+
+  execute(
+    args?: Partial<Omit<TransactWriteItemsCommandInput, "TransactItems">>
+  ) {
+    return this.transaction.execute(args);
+  }
 }
 
 export class ConditionOperations<U extends DynamoDBItem> {
   constructor(
+    private transaction: Transaction,
     private operations: Transaction["operations"],
     private item: DynamoDBItemKey,
     private args: Partial<ConditionOperation["args"]> & { TableName: string }
@@ -181,6 +195,12 @@ export class ConditionOperations<U extends DynamoDBItem> {
       },
     };
   }
+
+  execute(
+    args?: Partial<Omit<TransactWriteItemsCommandInput, "TransactItems">>
+  ) {
+    return this.transaction.execute(args);
+  }
 }
 
 const arraysToSets = <T extends DynamoDBItemKey>(

src/transactions/references/index.ts

@@ -48,11 +48,15 @@ const getRefsToResolve = (item: any) => {
   for (const value of Object.values<DynamoDBReference>(item)) {
     if (typeof value === "object") {
       if (value instanceof Set || Array.isArray(value)) {
-        refs.push(...[...value].map((v) => getRefsToResolve(v)).flat());
+        refs.push(
+          ...Array.from(value)
+            .map((v) => getRefsToResolve(v))
+            .flat()
+        );
       } else if (value?._type === "dynamodb:reference") {
         refs.push(value?._target);
       } else {
-        refs.push(...getRefsToResolve(value));
+        refs.push(...getRefsToResolve(value).flat());
       }
     }
   }
@@ -69,7 +73,7 @@ const injectRefs = (item: any, refs: Record<string, ItemWithKey>) => {
   for (const [key, value] of Object.entries<DynamoDBReference>(item)) {
     if (typeof value === "object") {
       if (value instanceof Set || Array.isArray(value)) {
-        item[key] = new Set([...value].map((v) => injectRefs(v, refs)));
+        item[key] = new Set(Array.from(value).map((v) => injectRefs(v, refs)));
       } else if (value?._type === "dynamodb:reference") {
         item[key] = refs[itemToStringKey(value._target)];
       } else {

test/transactions/index.test.js

@@ -1,6 +1,5 @@
 require("dotenv/config");
 
-const { ConditionalCheckFailedException } = require("@aws-sdk/client-dynamodb");
 const {
   Transaction,
   getItems,
@@ -525,4 +524,23 @@ describe("Transaction operations", () => {
     const newItem = await getItem(item, { ConsistentRead: true });
     expect(newItem.nested.number).toEqual(3);
   });
+
+  it("should merge update operations for the same item", async () => {
+    const item = generateItem("9003");
+    await putItem(item);
+
+    const transaction = new Transaction();
+    transaction.update(item).set({ title: "Test" });
+    transaction
+      .update(item)
+      .set({ author: "Test Author" })
+      .adjustNumber({ stars: 1 });
+
+    await transaction.execute();
+
+    const newItem = await getItem(item, { ConsistentRead: true });
+    expect(newItem.title).toEqual("Test");
+    expect(newItem.author).toEqual("Test Author");
+    expect(newItem.stars).toEqual(item.stars + 1);
+  });
 });

test/types.ts

@@ -1,3 +1,5 @@
+import "dotenv/config";
+
 import { PutItemCommandOutput } from "@aws-sdk/client-dynamodb";
 import {
   DynamoDBItem,
@@ -10,8 +12,9 @@ import {
   queryPaginatedItems,
   removeAttributes,
   transactGetItems,
+  Transaction,
   updateItem,
-} from "../dist";
+} from "../src";
 
 type DemoItem = {
   PK: `User/${string}`;
@@ -142,3 +145,18 @@ async function playground() {
   const simpleTypedItems = await getItems<DemoItem>([a, a]);
   const [simpleTypedItem1, simpleTypedItem2] = simpleTypedItems;
 }
+
+async function transactionPlayground() {
+  type TxnItem = {
+    PK: `User/${string}`;
+    SK: string;
+    stars: number;
+    otherProp: string;
+  };
+
+  const res = await new Transaction()
+    .update<TxnItem>({ PK: "User/1", SK: "Book/1" })
+    .set({ otherProp: "test" })
+    .adjustNumber({ stars: 1 })
+    .execute();
+}