fix(tests): raise key error

ZettaAI · May 17, 2024 · dd2c0ec · dd2c0ec
1 parent 32aeffe
commit dd2c0ec
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/zetta_utils/layer/db_layer/datastore/backend.py b/zetta_utils/layer/db_layer/datastore/backend.py
@@ -91,7 +91,10 @@ def __len__(self) -> int:
         return 0
 
     def _read_single_entity(self, idx: DBIndex):
-        row_key = self.client.key("Row", idx.row_keys[0])
+        row_key_str = idx.row_keys[0]
+        if row_key_str not in self:
+            raise KeyError(row_key_str)
+        row_key = self.client.key("Row", row_key_str)
         _query = self.client.query(kind="Column", ancestor=row_key)
         entities = list(_query.fetch())
         return _get_data_from_entities(idx, entities)
@@ -146,7 +149,7 @@ def query(self, column_filter: dict[str, list] | None = None) -> list[str]:
         return [entity.key.parent.id_or_name for entity in entities]
 
     @retry(
-        retry=retry_if_not_exception_type((RuntimeError, TypeError, ValueError)),
+        retry=retry_if_not_exception_type((KeyError, RuntimeError, TypeError, ValueError)),
         stop=stop_after_attempt(3),
         wait=wait_exponential(multiplier=1, min=4, max=10),
     )
@@ -163,7 +166,7 @@ def read(self, idx: DBIndex) -> DBDataT:
         return _get_data_from_entities(idx, entities)
 
     @retry(
-        retry=retry_if_not_exception_type((RuntimeError, TypeError, ValueError)),
+        retry=retry_if_not_exception_type((KeyError, RuntimeError, TypeError, ValueError)),
         stop=stop_after_attempt(3),
         wait=wait_exponential(multiplier=1, min=4, max=10),
     )