From aad28b2911644ec461d8de05ab664a81b95994ec Mon Sep 17 00:00:00 2001
From: Olli Warro <olli.warro@mindler.se>
Date: Thu, 28 Mar 2024 12:58:16 +0200
Subject: [PATCH] add support for remove update expr

---
 src/nodes/removeUpdateExpression.ts           |  4 +++
 src/nodes/updateExpression.ts                 |  2 ++
 ...eItemQueryBuilder.integration.test.ts.snap |  8 +++++
 ...updateItemQueryBuilder.integration.test.ts | 30 ++++++++++++++++-
 src/queryBuilders/updateItemQueryBuilder.ts   | 24 +++++++++++++-
 src/queryCompiler/queryCompiler.ts            | 32 +++++++++++++------
 src/queryCreator.ts                           |  1 +
 7 files changed, 89 insertions(+), 12 deletions(-)
 create mode 100644 src/nodes/removeUpdateExpression.ts

diff --git a/src/nodes/removeUpdateExpression.ts b/src/nodes/removeUpdateExpression.ts
new file mode 100644
index 0000000..1fb9f13
--- /dev/null
+++ b/src/nodes/removeUpdateExpression.ts
@@ -0,0 +1,4 @@
+export type RemoveUpdateExpression = {
+  readonly kind: "RemoveUpdateExpression";
+  readonly attribute: string;
+};
diff --git a/src/nodes/updateExpression.ts b/src/nodes/updateExpression.ts
index 6987e55..1b8c8aa 100644
--- a/src/nodes/updateExpression.ts
+++ b/src/nodes/updateExpression.ts
@@ -1,6 +1,8 @@
+import { RemoveUpdateExpression } from "./removeUpdateExpression";
 import { SetUpdateExpression } from "./setUpdateExpression";
 
 export type UpdateExpression = {
   readonly kind: "UpdateExpression";
   readonly setUpdateExpressions: SetUpdateExpression[];
+  readonly removeUpdateExpressions: RemoveUpdateExpression[];
 };
diff --git a/src/queryBuilders/__snapshots__/updateItemQueryBuilder.integration.test.ts.snap b/src/queryBuilders/__snapshots__/updateItemQueryBuilder.integration.test.ts.snap
index c8cdc04..52184b2 100644
--- a/src/queryBuilders/__snapshots__/updateItemQueryBuilder.integration.test.ts.snap
+++ b/src/queryBuilders/__snapshots__/updateItemQueryBuilder.integration.test.ts.snap
@@ -1,5 +1,13 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
+exports[`UpdateItemQueryBuilder > handles update item query with REMOVE statements 1`] = `
+{
+  "dataTimestamp": 200,
+  "someBoolean": true,
+  "userId": "1010",
+}
+`;
+
 exports[`UpdateItemQueryBuilder > handles update item query with SET statements 1`] = `
 {
   "dataTimestamp": 2,
diff --git a/src/queryBuilders/updateItemQueryBuilder.integration.test.ts b/src/queryBuilders/updateItemQueryBuilder.integration.test.ts
index 6ebafb5..a891618 100644
--- a/src/queryBuilders/updateItemQueryBuilder.integration.test.ts
+++ b/src/queryBuilders/updateItemQueryBuilder.integration.test.ts
@@ -27,11 +27,39 @@ describe("UpdateItemQueryBuilder", () => {
         );
       })
       .set("somethingElse", "+=", (qb) => {
-        return [qb.ifNotExists("somethingElse", 1), 2]
+        return [qb.ifNotExists("somethingElse", 1), 2];
       })
       .returnValues("ALL_NEW")
       .execute();
 
     expect(res).toMatchSnapshot();
   });
+
+  it("handles update item query with REMOVE statements", async () => {
+    await tsynamoClient
+      .putItem("myTable")
+      .item({
+        userId: "1010",
+        dataTimestamp: 200,
+        somethingElse: 313,
+        someBoolean: true,
+      })
+      .execute();
+
+    await tsynamoClient
+      .updateItem("myTable")
+      .keys({ userId: "1010", dataTimestamp: 200 })
+      .remove("somethingElse")
+      .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
index 4755627..d488a3f 100644
--- a/src/queryBuilders/updateItemQueryBuilder.ts
+++ b/src/queryBuilders/updateItemQueryBuilder.ts
@@ -116,7 +116,10 @@ export interface UpdateItemQueryBuilderInterface<
     pk: Keys
   ): UpdateItemQueryBuilderInterface<DDB, Table, O>;
 
-  // TODO: remove
+  remove<Key extends ObjectKeyPaths<PickNonKeys<DDB[Table]>>>(
+    attribute: Key
+  ): UpdateItemQueryBuilderInterface<DDB, Table, O>;
+
   // TODO: add
   // TODO: delete?
 
@@ -290,6 +293,25 @@ export class UpdateItemQueryBuilder<
     }
   }
 
+  remove<Key extends ObjectKeyPaths<PickNonKeys<DDB[Table]>>>(
+    attribute: Key
+  ): UpdateItemQueryBuilderInterface<DDB, Table, O> {
+    return new UpdateItemQueryBuilder<DDB, Table, O>({
+      ...this.#props,
+      node: {
+        ...this.#props.node,
+        updateExpression: {
+          ...this.#props.node.updateExpression,
+          removeUpdateExpressions:
+            this.#props.node.updateExpression.removeUpdateExpressions.concat({
+              kind: "RemoveUpdateExpression",
+              attribute,
+            }),
+        },
+      },
+    });
+  }
+
   returnValues(
     option: ReturnValuesOptions
   ): UpdateItemQueryBuilderInterface<DDB, Table, O> {
diff --git a/src/queryCompiler/queryCompiler.ts b/src/queryCompiler/queryCompiler.ts
index 7df0726..a3e8222 100644
--- a/src/queryCompiler/queryCompiler.ts
+++ b/src/queryCompiler/queryCompiler.ts
@@ -443,17 +443,29 @@ export class QueryCompiler {
     updateExpressionAttributeValues: Map<string, unknown>,
     attributeNames: Map<string, string>
   ) {
-    let res = "SET ";
+    let res = "";
 
-    res += node.setUpdateExpressions
-      .map((setUpdateExpression) => {
-        return this.compileSetUpdateExpression(
-          setUpdateExpression,
-          updateExpressionAttributeValues,
-          attributeNames
-        );
-      })
-      .join(", ");
+    if (node.setUpdateExpressions.length > 0) {
+      res += "SET ";
+      res += node.setUpdateExpressions
+        .map((setUpdateExpression) => {
+          return this.compileSetUpdateExpression(
+            setUpdateExpression,
+            updateExpressionAttributeValues,
+            attributeNames
+          );
+        })
+        .join(", ");
+    }
+
+    if (node.removeUpdateExpressions.length > 0) {
+      res += " REMOVE ";
+      res += node.removeUpdateExpressions
+        .map((removeUpdateExpression) => {
+          return removeUpdateExpression.attribute;
+        })
+        .join(", ");
+    }
 
     return res;
   }
diff --git a/src/queryCreator.ts b/src/queryCreator.ts
index ff91db0..878a9af 100644
--- a/src/queryCreator.ts
+++ b/src/queryCreator.ts
@@ -144,6 +144,7 @@ export class QueryCreator<DDB> {
         updateExpression: {
           kind: "UpdateExpression",
           setUpdateExpressions: [],
+          removeUpdateExpressions: [],
         },
       },
       ddbClient: this.#props.ddbClient,