intial work for supporting transactions

mindler-olli · Apr 18, 2024 · 47c9ce2 · 47c9ce2
1 parent dc22939
commit 47c9ce2
Show file tree

Hide file tree

Showing 8 changed files with 211 additions and 38 deletions.
diff --git a/src/nodes/TransactItemNode.ts b/src/nodes/TransactItemNode.ts
@@ -0,0 +1,6 @@
+import { PutNode } from "./putNode";
+
+export type TransactItemNode = {
+  readonly kind: "TransactItemNode";
+  readonly Put?: PutNode;
+};
diff --git a/src/nodes/transactionNode.ts b/src/nodes/transactionNode.ts
@@ -0,0 +1,6 @@
+import { TransactItemNode } from "./TransactItemNode";
+
+export type TransactionNode = {
+  readonly kind: "TransactionNode";
+  readonly transactItems: TransactItemNode[];
+};
diff --git a/src/queryBuilders/__snapshots__/transactionBuilder.integration.test.ts.snap b/src/queryBuilders/__snapshots__/transactionBuilder.integration.test.ts.snap
@@ -0,0 +1,14 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`TransactionBuilder > handles puts 1`] = `
+[
+  {
+    "dataTimestamp": 1,
+    "userId": "9999",
+  },
+  {
+    "dataTimestamp": 2,
+    "userId": "9999",
+  },
+]
+`;
diff --git a/src/queryBuilders/putItemQueryBuilder.ts b/src/queryBuilders/putItemQueryBuilder.ts
@@ -16,72 +16,76 @@ import {
   NotExprArg,
 } from "./expressionBuilder";
 
-export interface PutItemQueryBuilderInterface<DDB, Table extends keyof DDB, O> {
+export interface PutItemQueryBuilderInterface<
+  DDB,
+  Table extends keyof DDB,
+  O extends DDB[Table]
+> {
   // conditionExpression
   conditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: ComparatorExprArg<DDB, Table, Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   conditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: AttributeFuncExprArg<Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   conditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: AttributeBeginsWithExprArg<Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   conditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: AttributeContainsExprArg<DDB, Table, Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   conditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: AttributeBetweenExprArg<DDB, Table, Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   conditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: NotExprArg<DDB, Table, Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   conditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: BuilderExprArg<DDB, Table, Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   // orConditionExpression
   orConditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: ComparatorExprArg<DDB, Table, Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   orConditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: AttributeFuncExprArg<Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   orConditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: AttributeBeginsWithExprArg<Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   orConditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: AttributeContainsExprArg<DDB, Table, Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   orConditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: AttributeBetweenExprArg<DDB, Table, Key>
   ): PutItemQueryBuilderInterface<DDB, Table, O>;
 
   orConditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: NotExprArg<DDB, Table, Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   orConditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: BuilderExprArg<DDB, Table, Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   returnValues(
     option: Extract<ReturnValuesOptions, "NONE" | "ALL_OLD">
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   item<Item extends ExecuteOutput<O>>(
     item: Item
-  ): PutItemQueryBuilderInterface<DDB, Table, O>;
+  ): PutItemQueryBuilder<DDB, Table, O>;
 
   compile(): PutCommand;
   execute(): Promise<ExecuteOutput<O>[] | undefined>;
@@ -101,7 +105,7 @@ export class PutItemQueryBuilder<
 
   conditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: ExprArgs<DDB, Table, O, Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O> {
+  ): PutItemQueryBuilder<DDB, Table, O> {
     const eB = new ExpressionBuilder<DDB, Table, O>({
       node: { ...this.#props.node.conditionExpression },
     });
@@ -119,7 +123,7 @@ export class PutItemQueryBuilder<
 
   orConditionExpression<Key extends ObjectKeyPaths<DDB[Table]>>(
     ...args: ExprArgs<DDB, Table, O, Key>
-  ): PutItemQueryBuilderInterface<DDB, Table, O> {
+  ): PutItemQueryBuilder<DDB, Table, O> {
     const eB = new ExpressionBuilder<DDB, Table, O>({
       node: { ...this.#props.node.conditionExpression },
     });
@@ -137,7 +141,7 @@ export class PutItemQueryBuilder<
 
   item<Item extends ExecuteOutput<O>>(
     item: Item
-  ): PutItemQueryBuilderInterface<DDB, Table, O> {
+  ): PutItemQueryBuilder<DDB, Table, O> {
     return new PutItemQueryBuilder<DDB, Table, O>({
       ...this.#props,
       node: {
@@ -152,7 +156,7 @@ export class PutItemQueryBuilder<
 
   returnValues(
     option: Extract<ReturnValuesOptions, "NONE" | "ALL_OLD">
-  ): PutItemQueryBuilderInterface<DDB, Table, O> {
+  ): PutItemQueryBuilder<DDB, Table, O> {
     return new PutItemQueryBuilder<DDB, Table, O>({
       ...this.#props,
       node: {
@@ -168,11 +172,16 @@ export class PutItemQueryBuilder<
   compile = (): PutCommand => {
     return this.#props.queryCompiler.compile(this.#props.node);
   };
+
   execute = async (): Promise<ExecuteOutput<O>[] | undefined> => {
     const putCommand = this.compile();
     const data = await this.#props.ddbClient.send(putCommand);
     return data.Attributes as any;
   };
+
+  public get node() {
+    return this.#props.node;
+  }
 }
 
 preventAwait(

diff --git a/src/queryBuilders/transactionBuilder.integration.test.ts b/src/queryBuilders/transactionBuilder.integration.test.ts
@@ -0,0 +1,43 @@
+import { DynamoDBDocumentClient } from "@aws-sdk/lib-dynamodb";
+import { DDB } from "../../test/testFixture";
+import { getDDBClientFor, startDDBTestContainer } from "../../test/testUtil";
+import { Tsynamo } from "../index";
+
+describe("TransactionBuilder", () => {
+  let tsynamoClient: Tsynamo<DDB>;
+  let ddbClient: DynamoDBDocumentClient;
+
+  beforeAll(async () => {
+    const testContainer = await startDDBTestContainer();
+    ddbClient = await getDDBClientFor(testContainer);
+
+    tsynamoClient = new Tsynamo<DDB>({
+      ddbClient,
+    });
+  });
+
+  it("handles puts", async () => {
+    const trx = tsynamoClient.createTransaction();
+
+    trx.addItem({
+      Put: tsynamoClient
+        .putItem("myTable")
+        .item({ userId: "9999", dataTimestamp: 1 }),
+    });
+
+    trx.addItem({
+      Put: tsynamoClient
+        .putItem("myTable")
+        .item({ userId: "9999", dataTimestamp: 2 }),
+    });
+
+    await trx.execute();
+
+    const result = await tsynamoClient
+      .query("myTable")
+      .keyCondition("userId", "=", "9999")
+      .execute();
+
+    expect(result).toMatchSnapshot();
+  });
+});
diff --git a/src/queryBuilders/transactionBuilder.ts b/src/queryBuilders/transactionBuilder.ts
@@ -0,0 +1,44 @@
+import { DynamoDBDocumentClient } from "@aws-sdk/lib-dynamodb";
+import { TransactionNode } from "../nodes/transactionNode";
+import { QueryCompiler } from "../queryCompiler";
+import { PutItemQueryBuilder } from "./putItemQueryBuilder";
+
+export interface TransactionBuilderInterface<DDB> {
+  /**
+   * TODO: Update, Put, Delete, ConditionCheck
+   */
+  addItem(item: { Put?: PutItemQueryBuilder<DDB, any, any> }): void;
+
+  execute(): Promise<void>;
+}
+
+export class TransactionBuilder<DDB>
+  implements TransactionBuilderInterface<DDB>
+{
+  readonly #props: TransactionBuilderProps;
+
+  constructor(props: TransactionBuilderProps) {
+    this.#props = props;
+  }
+
+  addItem(item: { Put?: PutItemQueryBuilder<DDB, any, any> }) {
+    this.#props.node.transactItems.push({
+      kind: "TransactItemNode",
+      Put: item.Put?.node,
+    });
+  }
+
+  async execute() {
+    const transactionCommand = this.#props.queryCompiler.compile(
+      this.#props.node
+    );
+
+    await this.#props.ddbClient.send(transactionCommand);
+  }
+}
+
+interface TransactionBuilderProps {
+  readonly node: TransactionNode;
+  readonly ddbClient: DynamoDBDocumentClient;
+  readonly queryCompiler: QueryCompiler;
+}
diff --git a/src/queryCompiler/queryCompiler.ts b/src/queryCompiler/queryCompiler.ts
@@ -1,38 +1,51 @@
+import { TransactWriteItem } from "@aws-sdk/client-dynamodb";
 import {
   DeleteCommand,
   GetCommand,
   PutCommand,
+  PutCommandInput,
   QueryCommand,
+  TransactWriteCommand,
   UpdateCommand,
 } from "@aws-sdk/lib-dynamodb";
+import { AddUpdateExpression } from "../nodes/addUpdateExpression";
 import { AttributesNode } from "../nodes/attributesNode";
 import { DeleteNode } from "../nodes/deleteNode";
+import { DeleteUpdateExpression } from "../nodes/deleteUpdateExpression";
 import { ExpressionJoinTypeNode } from "../nodes/expressionJoinTypeNode";
 import { ExpressionNode } from "../nodes/expressionNode";
 import { GetNode } from "../nodes/getNode";
 import { KeyConditionNode } from "../nodes/keyConditionNode";
 import { PutNode } from "../nodes/putNode";
 import { QueryNode } from "../nodes/queryNode";
+import { RemoveUpdateExpression } from "../nodes/removeUpdateExpression";
 import { SetUpdateExpression } from "../nodes/setUpdateExpression";
 import { SetUpdateExpressionFunction } from "../nodes/setUpdateExpressionFunction";
+import { TransactionNode } from "../nodes/transactionNode";
 import { UpdateExpression } from "../nodes/updateExpression";
 import { UpdateNode } from "../nodes/updateNode";
 import {
   getAttributeNameFrom,
   getExpressionAttributeNameFrom,
   mergeObjectIntoMap,
 } from "./compilerUtil";
-import { RemoveUpdateExpression } from "../nodes/removeUpdateExpression";
-import { AddUpdateExpression } from "../nodes/addUpdateExpression";
-import { DeleteUpdateExpression } from "../nodes/deleteUpdateExpression";
 
 export class QueryCompiler {
   compile(rootNode: QueryNode): QueryCommand;
   compile(rootNode: GetNode): GetCommand;
   compile(rootNode: PutNode): PutCommand;
   compile(rootNode: DeleteNode): DeleteCommand;
   compile(rootNode: UpdateNode): UpdateCommand;
-  compile(rootNode: QueryNode | GetNode | PutNode | DeleteNode | UpdateNode) {
+  compile(rootNode: TransactionNode): TransactWriteCommand;
+  compile(
+    rootNode:
+      | QueryNode
+      | GetNode
+      | PutNode
+      | DeleteNode
+      | UpdateNode
+      | TransactionNode
+  ) {
     switch (rootNode.kind) {
       case "GetNode":
         return this.compileGetNode(rootNode);
@@ -44,6 +57,8 @@ export class QueryCompiler {
         return this.compileDeleteNode(rootNode);
       case "UpdateNode":
         return this.compileUpdateNode(rootNode);
+      case "TransactionNode":
+        return this.compileTransactionNode(rootNode);
     }
   }
 
@@ -122,6 +137,10 @@ export class QueryCompiler {
   }
 
   compilePutNode(putNode: PutNode) {
+    return new PutCommand(this.compilePutCmdInput(putNode));
+  }
+
+  compilePutCmdInput(putNode: PutNode): PutCommandInput {
     const {
       table: tableNode,
       item: itemNode,
@@ -138,7 +157,7 @@ export class QueryCompiler {
       attributeNames
     );
 
-    return new PutCommand({
+    return {
       TableName: tableNode.table,
       Item: itemNode?.item,
       ReturnValues: returnValuesNode?.option,
@@ -157,7 +176,7 @@ export class QueryCompiler {
               ...Object.fromEntries(attributeNames),
             }
           : undefined,
-    });
+    };
   }
 
   compileDeleteNode(deleteNode: DeleteNode) {
@@ -252,6 +271,26 @@ export class QueryCompiler {
     });
   }
 
+  compileTransactionNode(transactionNode: TransactionNode) {
+    const TransactItems = transactionNode.transactItems.map((item) => {
+      const compiledTransactItem: TransactWriteItem = {};
+
+      if (item.Put) {
+        compiledTransactItem.Put = this.compilePutCmdInput(item.Put);
+      }
+
+      return compiledTransactItem;
+    });
+
+    console.log("WUT", {
+      TransactItems: TransactItems.map((x) => x.Put),
+    });
+
+    return new TransactWriteCommand({
+      TransactItems: TransactItems,
+    });
+  }
+
   compileAttributeNamesNode(node?: AttributesNode) {
     const ProjectionExpression = node?.attributes
       .map((att) => getExpressionAttributeNameFrom(att))