Add check for empty val. Merge json correctly

Yelp · Jan 16, 2025 · e351dd1 · e351dd1
1 parent 04418af
commit e351dd1
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tron/serialize/runstate/dynamodb_state_store.py b/tron/serialize/runstate/dynamodb_state_store.py
@@ -165,11 +165,12 @@ def _merge_items(self, first_items, remaining_items, read_json=False) -> dict:
         for key, item in items.items():
             item.sort(key=lambda x: int(x["index"]["N"]))
             for val in item:
-                raw_items[key] += bytes(val["val"]["B"])
+                if "val" in val:
+                    raw_items[key] += bytes(val["val"]["B"])
             if read_json:
                 for json_val in item:
                     try:
-                        json_items[key] = json_val["json_val"]["S"]
+                        json_items[key] += json_val["json_val"]["S"]
                     except Exception:
                         log.exception(f"json_val not found for key {key}")
                         # fallback to pickled data if json_val fails to exist for any key