add support for compiling set update expressions

mindler-olli · Mar 27, 2024 · b376baf · b376baf
1 parent 4565d4c
commit b376baf
Show file tree

Hide file tree

Showing 7 changed files with 399 additions and 24 deletions.
diff --git a/src/nodes/setUpdateExpressionFunction.ts b/src/nodes/setUpdateExpressionFunction.ts
@@ -15,6 +15,6 @@ export type SetUpdateExpressionIfNotExistsFunction = {
 
 export type SetUpdateExpressionListAppendFunction = {
   readonly kind: "SetUpdateExpressionListAppendFunction";
-  readonly left: SetUpdateExpressionFunction;
+  readonly left: SetUpdateExpressionFunction | string;
   readonly right: SetUpdateExpressionFunction | SetUpdateExpressionValue;
 };
diff --git a/src/nodes/updateNode.ts b/src/nodes/updateNode.ts
@@ -1,5 +1,6 @@
 import { ExpressionNode } from "./expressionNode";
 import { ItemNode } from "./itemNode";
+import { KeysNode } from "./keysNode";
 import { ReturnValuesNode } from "./returnValuesNode";
 import { TableNode } from "./tableNode";
 import { UpdateExpression } from "./updateExpression";
@@ -9,6 +10,6 @@ export type UpdateNode = {
   readonly table: TableNode;
   readonly conditionExpression: ExpressionNode;
   readonly updateExpression: UpdateExpression;
-  readonly item?: ItemNode;
+  readonly keys?: KeysNode;
   readonly returnValues?: ReturnValuesNode;
 };
diff --git a/src/queryBuilders/__snapshots__/updateItemQueryBuilder.integration.test.ts.snap b/src/queryBuilders/__snapshots__/updateItemQueryBuilder.integration.test.ts.snap
@@ -0,0 +1,12 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`UpdateItemQueryBuilder > handles update item query with SET statements 1`] = `
+{
+  "dataTimestamp": 2,
+  "someBoolean": true,
+  "tags": [
+    "test_tag",
+  ],
+  "userId": "1",
+}
+`;
diff --git a/src/queryBuilders/setUpdateExpressionFunctionQueryBuilder.ts b/src/queryBuilders/setUpdateExpressionFunctionQueryBuilder.ts
@@ -1,3 +1,4 @@
+import exp from "constants";
 import { SetUpdateExpressionFunction } from "../nodes/setUpdateExpressionFunction";
 import { GetFromPath, ObjectKeyPaths, StripKeys } from "../typeHelpers";
 
@@ -18,8 +19,7 @@ export interface SetUpdateExpressionFunctionQueryBuilderInterface<
     ) => SetUpdateExpressionFunction
   ): SetUpdateExpressionFunction;
 
-  // NOTE: List append can also be func func or val val or func val or val func :-)
-  /* listAppend<Key extends ObjectKeyPaths<DDB[Table]>>(
+  listAppend<Key extends ObjectKeyPaths<DDB[Table]>>(
     key: Key,
     value: StripKeys<GetFromPath<DDB[Table], Key>>
   ): SetUpdateExpressionFunction;
@@ -29,7 +29,23 @@ export interface SetUpdateExpressionFunctionQueryBuilderInterface<
     value: (
       builder: SetUpdateExpressionFunctionQueryBuilder<DDB, Table, DDB[Table]>
     ) => SetUpdateExpressionFunction
-  ): SetUpdateExpressionFunction; */
+  ): SetUpdateExpressionFunction;
+
+  listAppend<Key extends ObjectKeyPaths<DDB[Table]>>(
+    key: (
+      builder: SetUpdateExpressionFunctionQueryBuilder<DDB, Table, DDB[Table]>
+    ) => SetUpdateExpressionFunction,
+    value: StripKeys<GetFromPath<DDB[Table], Key>>
+  ): SetUpdateExpressionFunction;
+
+  listAppend(
+    key: (
+      builder: SetUpdateExpressionFunctionQueryBuilder<DDB, Table, DDB[Table]>
+    ) => SetUpdateExpressionFunction,
+    value: (
+      builder: SetUpdateExpressionFunctionQueryBuilder<DDB, Table, DDB[Table]>
+    ) => SetUpdateExpressionFunction
+  ): SetUpdateExpressionFunction;
 }
 
 export class SetUpdateExpressionFunctionQueryBuilder<
@@ -66,7 +82,6 @@ export class SetUpdateExpressionFunctionQueryBuilder<
 
       const expression = builder(setUpdateExpressionBuilder);
 
-
       this.node = {
         kind: "SetUpdateExpressionFunction",
         function: {
@@ -91,4 +106,128 @@ export class SetUpdateExpressionFunctionQueryBuilder<
 
     return this.node;
   }
+
+  listAppend<Key extends ObjectKeyPaths<DDB[Table]>>(
+    ...args:
+      | [key: Key, value: StripKeys<GetFromPath<DDB[Table], Key>>]
+      | [
+          key: Key,
+          value: (
+            builder: SetUpdateExpressionFunctionQueryBuilder<
+              DDB,
+              Table,
+              DDB[Table]
+            >
+          ) => SetUpdateExpressionFunction
+        ]
+      | [
+          key: (
+            builder: SetUpdateExpressionFunctionQueryBuilder<
+              DDB,
+              Table,
+              DDB[Table]
+            >
+          ) => SetUpdateExpressionFunction,
+          value: StripKeys<GetFromPath<DDB[Table], Key>>
+        ]
+      | [
+          key: (
+            builder: SetUpdateExpressionFunctionQueryBuilder<
+              DDB,
+              Table,
+              DDB[Table]
+            >
+          ) => SetUpdateExpressionFunction,
+          value: (
+            builder: SetUpdateExpressionFunctionQueryBuilder<
+              DDB,
+              Table,
+              DDB[Table]
+            >
+          ) => SetUpdateExpressionFunction
+        ]
+  ) {
+    const [left, right] = args;
+
+    if (typeof left === "function" && typeof right === "function") {
+      const setUpdateExpressionBuilderA =
+        new SetUpdateExpressionFunctionQueryBuilder<DDB, Table, O>();
+
+      const setUpdateExpressionBuilderB =
+        new SetUpdateExpressionFunctionQueryBuilder<DDB, Table, O>();
+
+      const builderLeft = left as (
+        builder: SetUpdateExpressionFunctionQueryBuilder<DDB, Table, DDB[Table]>
+      ) => SetUpdateExpressionFunction;
+
+      const builderRight = right as (
+        builder: SetUpdateExpressionFunctionQueryBuilder<DDB, Table, DDB[Table]>
+      ) => SetUpdateExpressionFunction;
+
+      const exprLeft = builderLeft(setUpdateExpressionBuilderA);
+      const exprRight = builderRight(setUpdateExpressionBuilderB);
+
+      this.node = {
+        kind: "SetUpdateExpressionFunction",
+        function: {
+          kind: "SetUpdateExpressionListAppendFunction",
+          left: exprLeft,
+          right: exprRight,
+        },
+      };
+    } else if (typeof left === "function") {
+      const setUpdateExpressionBuilder =
+        new SetUpdateExpressionFunctionQueryBuilder<DDB, Table, O>();
+
+      const builder = left as (
+        builder: SetUpdateExpressionFunctionQueryBuilder<DDB, Table, DDB[Table]>
+      ) => SetUpdateExpressionFunction;
+
+      const expr = builder(setUpdateExpressionBuilder);
+
+      this.node = {
+        kind: "SetUpdateExpressionFunction",
+        function: {
+          kind: "SetUpdateExpressionListAppendFunction",
+          left: expr,
+          right: {
+            kind: "SetUpdateExpressionValue",
+            value: right,
+          },
+        },
+      };
+    } else if (typeof right === "function") {
+      const setUpdateExpressionBuilder =
+        new SetUpdateExpressionFunctionQueryBuilder<DDB, Table, O>();
+
+      const builder = right as (
+        builder: SetUpdateExpressionFunctionQueryBuilder<DDB, Table, DDB[Table]>
+      ) => SetUpdateExpressionFunction;
+
+      const expr = builder(setUpdateExpressionBuilder);
+
+      this.node = {
+        kind: "SetUpdateExpressionFunction",
+        function: {
+          kind: "SetUpdateExpressionListAppendFunction",
+          left,
+          right: expr,
+        },
+      };
+    } else {
+      this.node = {
+        kind: "SetUpdateExpressionFunction",
+        function: {
+          kind: "SetUpdateExpressionListAppendFunction",
+          left,
+          right: {
+            kind: "SetUpdateExpressionValue",
+            value: right,
+          },
+        },
+      };
+    }
+
+    return this.node;
+  }
 }
diff --git a/src/queryBuilders/updateItemQueryBuilder.integration.test.ts b/src/queryBuilders/updateItemQueryBuilder.integration.test.ts
@@ -14,12 +14,21 @@ describe("UpdateItemQueryBuilder", () => {
   });
 
   it("handles update item query with SET statements", async () => {
-    await tsynamoClient
+    const res = await tsynamoClient
       .updateItem("myTable")
+      .keys({ userId: "1", dataTimestamp: 2 })
       .set("someBoolean", "=", (qb) => {
         return qb.ifNotExists("someBoolean", true);
       })
-      .set("dataTimestamp", "+=", 1)
+      .set("tags", "=", (qb) => {
+        return qb.listAppend(
+          (qbb) => qbb.ifNotExists("tags", []),
+          ["test_tag"]
+        );
+      })
+      .returnValues("ALL_NEW")
       .execute();
+
+    expect(res).toMatchSnapshot();
   });
 });
diff --git a/src/queryBuilders/updateItemQueryBuilder.ts b/src/queryBuilders/updateItemQueryBuilder.ts
@@ -8,7 +8,10 @@ import {
   ExecuteOutput,
   GetFromPath,
   ObjectKeyPaths,
+  PickPk,
+  PickSkRequired,
   StripKeys,
+  PickNonKeys,
 } from "../typeHelpers";
 import { preventAwait } from "../util/preventAwait";
 import {
@@ -87,27 +90,23 @@ export interface UpdateItemQueryBuilderInterface<
     ...args: BuilderExprArg<DDB, Table, Key>
   ): UpdateItemQueryBuilderInterface<DDB, Table, O>;
 
-  set<Key extends ObjectKeyPaths<DDB[Table]>>(
+  set<Key extends ObjectKeyPaths<PickNonKeys<DDB[Table]>>>(
     key: Key,
     operand: UpdateExpressionOperands,
     value: StripKeys<GetFromPath<DDB[Table], Key>>
   ): UpdateItemQueryBuilderInterface<DDB, Table, O>;
 
-  set<Key extends ObjectKeyPaths<DDB[Table]>>(
+  set<Key extends ObjectKeyPaths<PickNonKeys<DDB[Table]>>>(
     key: Key,
     operand: UpdateExpressionOperands,
     value: (
       builder: SetUpdateExpressionFunctionQueryBuilder<DDB, Table, DDB[Table]>
     ) => SetUpdateExpressionFunction
   ): UpdateItemQueryBuilderInterface<DDB, Table, O>;
 
-  /* set("key1", "=", ({ ifNotExists }) => {
-    return ifNotExists("key2", ({ ifNotExists }) => {
-      return ifNotExists("key3", ({ listAppend }) => {
-        return listAppend("key3", "dada")
-      })
-    })
-  }) */
+  keys<Keys extends PickPk<DDB[Table]> & PickSkRequired<DDB[Table]>>(
+    pk: Keys
+  ): UpdateItemQueryBuilderInterface<DDB, Table, O>;
 
   // TODO: remove
   // TODO: add
@@ -169,7 +168,7 @@ export class UpdateItemQueryBuilder<
     });
   }
 
-  set<Key extends ObjectKeyPaths<DDB[Table]>, XD extends DDB>(
+  set<Key extends ObjectKeyPaths<PickNonKeys<DDB[Table]>>>(
     ...args:
       | [
           key: Key,
@@ -255,6 +254,21 @@ export class UpdateItemQueryBuilder<
     });
   }
 
+  keys<Keys extends PickPk<DDB[Table]> & PickSkRequired<DDB[Table]>>(
+    keys: Keys
+  ) {
+    return new UpdateItemQueryBuilder<DDB, Table, O>({
+      ...this.#props,
+      node: {
+        ...this.#props.node,
+        keys: {
+          kind: "KeysNode",
+          keys,
+        },
+      },
+    });
+  }
+
   compile = (): UpdateCommand => {
     return this.#props.queryCompiler.compile(this.#props.node);
   };