add support for ADD statements in update expr

mindler-olli · Mar 28, 2024 · 163db48 · 163db48
1 parent aad28b2
commit 163db48
Show file tree

Hide file tree

Showing 8 changed files with 139 additions and 12 deletions.
diff --git a/src/nodes/addUpdateExpression.ts b/src/nodes/addUpdateExpression.ts
@@ -0,0 +1,5 @@
+export type AddUpdateExpression = {
+  readonly kind: "AddUpdateExpression";
+  readonly key: string;
+  readonly value: unknown;
+};
diff --git a/src/nodes/updateExpression.ts b/src/nodes/updateExpression.ts
@@ -1,8 +1,10 @@
+import { AddUpdateExpression } from "./addUpdateExpression";
 import { RemoveUpdateExpression } from "./removeUpdateExpression";
 import { SetUpdateExpression } from "./setUpdateExpression";
 
 export type UpdateExpression = {
   readonly kind: "UpdateExpression";
   readonly setUpdateExpressions: SetUpdateExpression[];
   readonly removeUpdateExpressions: RemoveUpdateExpression[];
+  readonly addUpdateExpressions: AddUpdateExpression[];
 };
diff --git a/src/queryBuilders/__snapshots__/updateItemQueryBuilder.integration.test.ts.snap b/src/queryBuilders/__snapshots__/updateItemQueryBuilder.integration.test.ts.snap
@@ -1,5 +1,18 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
+exports[`UpdateItemQueryBuilder > handles update item query with ADD statements 1`] = `
+{
+  "dataTimestamp": 200,
+  "someBoolean": true,
+  "someSet": Set {
+    "item1",
+    "item2",
+  },
+  "somethingElse": 7,
+  "userId": "1010",
+}
+`;
+
 exports[`UpdateItemQueryBuilder > handles update item query with REMOVE statements 1`] = `
 {
   "dataTimestamp": 200,

diff --git a/src/queryBuilders/updateItemQueryBuilder.integration.test.ts b/src/queryBuilders/updateItemQueryBuilder.integration.test.ts
@@ -1,15 +1,18 @@
+import { DynamoDBDocumentClient } from "@aws-sdk/lib-dynamodb";
 import { DDB } from "../../test/testFixture";
 import { getDDBClientFor, startDDBTestContainer } from "../../test/testUtil";
 import { Tsynamo } from "./../index";
 
 describe("UpdateItemQueryBuilder", () => {
   let tsynamoClient: Tsynamo<DDB>;
+  let ddbClient: DynamoDBDocumentClient;
 
   beforeAll(async () => {
     const testContainer = await startDDBTestContainer();
+    ddbClient = await getDDBClientFor(testContainer);
 
     tsynamoClient = new Tsynamo<DDB>({
-      ddbClient: await getDDBClientFor(testContainer),
+      ddbClient,
     });
   });
 
@@ -62,4 +65,32 @@ describe("UpdateItemQueryBuilder", () => {
 
     expect(foundItem).toMatchSnapshot();
   });
+
+  it("handles update item query with ADD statements", async () => {
+    await tsynamoClient
+      .putItem("myTable")
+      .item({
+        userId: "1010",
+        dataTimestamp: 200,
+        someBoolean: true,
+      })
+      .execute();
+
+    await tsynamoClient
+      .updateItem("myTable")
+      .keys({ userId: "1010", dataTimestamp: 200 })
+      .add("somethingElse", 7)
+      .add("someSet", new Set(["item1", "item2"]))
+      .execute();
+
+    const foundItem = await tsynamoClient
+      .getItem("myTable")
+      .keys({
+        userId: "1010",
+        dataTimestamp: 200,
+      })
+      .execute();
+
+    expect(foundItem).toMatchSnapshot();
+  });
 });
diff --git a/src/queryBuilders/updateItemQueryBuilder.ts b/src/queryBuilders/updateItemQueryBuilder.ts
@@ -120,8 +120,13 @@ export interface UpdateItemQueryBuilderInterface<
     attribute: Key
   ): UpdateItemQueryBuilderInterface<DDB, Table, O>;
 
-  // TODO: add
-  // TODO: delete?
+  // TODO: This should only be supported for keys that are of type NUMBER or ARRAY
+  add<Key extends ObjectKeyPaths<PickNonKeys<DDB[Table]>>>(
+    attribute: Key,
+    value: StripKeys<GetFromPath<DDB[Table], Key>>
+  ): UpdateItemQueryBuilderInterface<DDB, Table, O>;
+
+  // TODO: Add support for DELETE action
 
   returnValues(
     option: ReturnValuesOptions
@@ -312,6 +317,27 @@ export class UpdateItemQueryBuilder<
     });
   }
 
+  add<Key extends ObjectKeyPaths<PickNonKeys<DDB[Table]>>>(
+    attribute: Key,
+    value: StripKeys<GetFromPath<DDB[Table], Key>>
+  ): UpdateItemQueryBuilderInterface<DDB, Table, O> {
+    return new UpdateItemQueryBuilder<DDB, Table, O>({
+      ...this.#props,
+      node: {
+        ...this.#props.node,
+        updateExpression: {
+          ...this.#props.node.updateExpression,
+          addUpdateExpressions:
+            this.#props.node.updateExpression.addUpdateExpressions.concat({
+              kind: "AddUpdateExpression",
+              key: attribute,
+              value,
+            }),
+        },
+      },
+    });
+  }
+
   returnValues(
     option: ReturnValuesOptions
   ): UpdateItemQueryBuilderInterface<DDB, Table, O> {

diff --git a/src/queryCompiler/queryCompiler.ts b/src/queryCompiler/queryCompiler.ts
@@ -22,6 +22,8 @@ import {
   getExpressionAttributeNameFrom,
   mergeObjectIntoMap,
 } from "./compilerUtil";
+import { RemoveUpdateExpression } from "../nodes/removeUpdateExpression";
+import { AddUpdateExpression } from "../nodes/addUpdateExpression";
 
 export class QueryCompiler {
   compile(rootNode: QueryNode): QueryCommand;
@@ -462,11 +464,29 @@ export class QueryCompiler {
       res += " REMOVE ";
       res += node.removeUpdateExpressions
         .map((removeUpdateExpression) => {
-          return removeUpdateExpression.attribute;
+          return this.compileRemoveUpdateExpression(
+            removeUpdateExpression,
+            attributeNames
+          );
+        })
+        .join(", ");
+    }
+
+    if (node.addUpdateExpressions.length > 0) {
+      res += " ADD ";
+      res += node.addUpdateExpressions
+        .map((addUpdateExpression) => {
+          return this.compileAddUpdateExpression(
+            addUpdateExpression,
+            updateExpressionAttributeValues,
+            attributeNames
+          );
         })
         .join(", ");
     }
 
+    // TODO: Compile ADD actions
+
     return res;
   }
 
@@ -534,7 +554,7 @@ export class QueryCompiler {
 
   compileSetUpdateExpressionFunction(
     functionExpression: SetUpdateExpressionFunction,
-    filterExpressionAttributeValues: Map<string, unknown>,
+    updateExpressionAttributeValues: Map<string, unknown>,
     attributeNames: Map<string, string>
   ) {
     const { function: functionNode } = functionExpression;
@@ -543,20 +563,20 @@ export class QueryCompiler {
     switch (functionNode.kind) {
       case "SetUpdateExpressionIfNotExistsFunction": {
         let rightValue = "";
-        const offset = filterExpressionAttributeValues.size;
+        const offset = updateExpressionAttributeValues.size;
         const attributeValue = `:setUpdateExpressionValue${offset}`;
 
         if (functionNode.right.kind === "SetUpdateExpressionValue") {
           rightValue = attributeValue;
 
-          filterExpressionAttributeValues.set(
+          updateExpressionAttributeValues.set(
             attributeValue,
             functionNode.right.value
           );
         } else {
           rightValue = this.compileSetUpdateExpressionFunction(
             functionNode.right,
-            filterExpressionAttributeValues,
+            updateExpressionAttributeValues,
             attributeNames
           );
         }
@@ -584,25 +604,25 @@ export class QueryCompiler {
         } else {
           leftValue = this.compileSetUpdateExpressionFunction(
             functionNode.left,
-            filterExpressionAttributeValues,
+            updateExpressionAttributeValues,
             attributeNames
           );
         }
 
-        const offset = filterExpressionAttributeValues.size;
+        const offset = updateExpressionAttributeValues.size;
         const attributeValue = `:setUpdateExpressionValue${offset}`;
 
         if (functionNode.right.kind === "SetUpdateExpressionValue") {
           rightValue = attributeValue;
 
-          filterExpressionAttributeValues.set(
+          updateExpressionAttributeValues.set(
             attributeValue,
             functionNode.right.value
           );
         } else {
           rightValue = this.compileSetUpdateExpressionFunction(
             functionNode.right,
-            filterExpressionAttributeValues,
+            updateExpressionAttributeValues,
             attributeNames
           );
         }
@@ -612,4 +632,32 @@ export class QueryCompiler {
       }
     }
   }
+
+  compileRemoveUpdateExpression(
+    node: RemoveUpdateExpression,
+    attributeNames: Map<string, string>
+  ) {
+    const { expressionAttributeName, attributeNameMap } =
+      this.compileAttributeName(node.attribute);
+    const attributeName = expressionAttributeName;
+    mergeObjectIntoMap(attributeNames, attributeNameMap);
+    return attributeName;
+  }
+
+  compileAddUpdateExpression(
+    node: AddUpdateExpression,
+    updateExpressionAttributeValues: Map<string, unknown>,
+    attributeNames: Map<string, string>
+  ) {
+    const { expressionAttributeName, attributeNameMap } =
+      this.compileAttributeName(node.key);
+    const attributeName = expressionAttributeName;
+    mergeObjectIntoMap(attributeNames, attributeNameMap);
+
+    const offset = updateExpressionAttributeValues.size;
+    const attributeValue = `:addUpdateExpressionValue${offset}`;
+    updateExpressionAttributeValues.set(attributeValue, node.value);
+
+    return `${attributeName} ${attributeValue}`;
+  }
 }
diff --git a/src/queryCreator.ts b/src/queryCreator.ts
@@ -145,6 +145,7 @@ export class QueryCreator<DDB> {
           kind: "UpdateExpression",
           setUpdateExpressions: [],
           removeUpdateExpressions: [],
+          addUpdateExpressions: [],
         },
       },
       ddbClient: this.#props.ddbClient,

diff --git a/test/testFixture.ts b/test/testFixture.ts
@@ -23,6 +23,7 @@ export interface DDB {
       };
     };
     tags: string[];
+    someSet: Set<string>;
   };
   myOtherTable: {
     userId: PartitionKey<string>;