fix: entity remove missing attributes (#441)

sam-goodwin · Sep 8, 2023 · 557c02e · 557c02e
1 parent 8cd6492
commit 557c02e
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 15 deletions.
diff --git a/apps/tests/aws-runtime/test/test-service.ts b/apps/tests/aws-runtime/test/test-service.ts
@@ -557,6 +557,8 @@ export const counter = entity("counter5", {
       z.literal("default"),
       z.literal("another"),
       z.literal("another2"),
+      z.literal("remove"),
+      z.literal("remove2"),
     ]),
     id: z.string(),
     optional: z.string().optional(),
@@ -884,6 +886,24 @@ export const entityWorkflow = workflow(
       { select: ["n", "counterNumber"] }
     );
 
+    /**
+     * Set putting undefined or missing values.
+     */
+    await Promise.all([
+      counter.put({ namespace: "remove", id, n: 1, optional: "someValue" }),
+      counter.put({ namespace: "remove2", id, n: 1, optional: "someValue" }),
+    ]);
+
+    await Promise.all([
+      counter.put({ namespace: "remove", id, n: 2 }),
+      counter.put({ namespace: "remove2", id, n: 2, optional: undefined }),
+    ]);
+
+    const removed = await countersByNamespace.query(
+      { namespace: { $beginsWith: "remove" }, id },
+      { select: ["n", "optional"] }
+    );
+
     return [
       result0,
       result1,
@@ -903,6 +923,7 @@ export const entityWorkflow = workflow(
         // @ts-check - should not error
         return c.value;
       }),
+      removed.entries?.map(({ value }) => value),
     ];
   }
 );

diff --git a/apps/tests/aws-runtime/test/tester.test.ts b/apps/tests/aws-runtime/test/tester.test.ts
@@ -235,6 +235,8 @@ eventualRuntimeTestHarness(
         { counterNumber: 2, n: 2 },
         { counterNumber: 3, n: 1 },
       ]),
+      // testing optional set, optional field should be erased
+      [{ n: 2 }, { n: 2 }],
     ]);
 
     testCompletion("transaction", transactionWorkflow, ([one, two, three]) => {

diff --git a/packages/@eventual/aws-runtime/src/handlers/entity-stream-worker.ts b/packages/@eventual/aws-runtime/src/handlers/entity-stream-worker.ts
@@ -11,6 +11,7 @@ import {
   createEntityStreamWorker,
   getLazy,
   normalizeCompositeKey,
+  removeGeneratedKeyAttributes,
 } from "@eventual/core-runtime";
 import type { EntityStreamOperation } from "@eventual/core/internal";
 import type { DynamoDBStreamHandler } from "aws-lambda";
@@ -88,12 +89,7 @@ export default (async (event) => {
     function removeSyntheticKey(value: Record<string, any> | undefined) {
       if (value) {
         delete value[EntityEntityRecord.VERSION_FIELD];
-        if (normalizedKey.partition.attributes.length > 0) {
-          delete value[normalizedKey.partition.keyAttribute];
-        }
-        if (normalizedKey.sort && normalizedKey.sort.attributes.length > 0) {
-          delete value[normalizedKey.sort.keyAttribute];
-        }
+        removeGeneratedKeyAttributes(entity!, value, false, true);
       }
     }
 

diff --git a/packages/@eventual/aws-runtime/src/stores/entity-store.ts b/packages/@eventual/aws-runtime/src/stores/entity-store.ts
@@ -423,6 +423,13 @@ export class AWSEntityStore extends EntityStore {
       value
     );
 
+    /**
+     * Any attribute not in the input object, but in the schema, will be removed
+     */
+    const missingAttributes = Object.keys(entity.attributes.shape).filter(
+      (a) => !(a in value)
+    );
+
     // remove the entity keys if they are not generated
     const valueToSave = removeKeyAttributes(
       entity,
@@ -437,15 +444,21 @@ export class AWSEntityStore extends EntityStore {
 
     return {
       Key: this.entityKey(key),
-      UpdateExpression: [
-        "SET #__version=if_not_exists(#__version, :__startingVersion) + :__versionIncrement",
-        ...Object.keys(valueRecord).map(
-          (key) =>
-            `${formatAttributeNameMapKey(key)}=${formatAttributeValueMapKey(
-              key
-            )}`
-        ),
-      ].join(","),
+      UpdateExpression:
+        [
+          "SET #__version=if_not_exists(#__version, :__startingVersion) + :__versionIncrement",
+          ...Object.keys(valueRecord).map(
+            (key) =>
+              `${formatAttributeNameMapKey(key)}=${formatAttributeValueMapKey(
+                key
+              )}`
+          ),
+        ].join(",") +
+        (missingAttributes.length > 0
+          ? ` REMOVE ${missingAttributes
+              .map((a) => formatAttributeNameMapKey(a))
+              .join(",")}`
+          : ""),
       ExpressionAttributeNames: {
         "#__version": EntityEntityRecord.VERSION_FIELD,
         ...Object.fromEntries(
@@ -454,6 +467,9 @@ export class AWSEntityStore extends EntityStore {
             key,
           ])
         ),
+        ...Object.fromEntries(
+          missingAttributes.map((m) => [formatAttributeNameMapKey(m), m])
+        ),
       },
       ExpressionAttributeValues: {
         ...(options?.expectedVersion