Merge pull request #1019 from Yelp/u/emanelsabban/capture-json_val

Capturing exception when json_val doesnt exist
Yelp · Jan 8, 2025 · 773708a · 773708a
2 parents 5dbd969 + ea037c6
commit 773708a
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/tron/serialize/runstate/dynamodb_state_store.py b/tron/serialize/runstate/dynamodb_state_store.py
@@ -169,7 +169,13 @@ def _merge_items(self, first_items, remaining_items, read_json=False) -> dict:
                 raw_items[key] += bytes(val["val"]["B"])
             if read_json:
                 for json_val in item:
-                    json_items[key] = json_val["json_val"]["S"]
+                    try:
+                        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
+                        # TODO: it would be nice if we can read the pickled data only for this failed key instead of all keys
+                        read_json = False
         if read_json:
             try:
                 log.info("read_json is enabled. Deserializing JSON items to restore them instead of pickled data.")