fix(get): Fix from_id cache-misses (#15)

chime-experiment · Oct 18, 2023 · 5006d22 · 5006d22
2 parents c58ab2d + 035643d
commit 5006d22
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 9 deletions.
diff --git a/chimedb/dataset/get.py b/chimedb/dataset/get.py
@@ -105,14 +105,14 @@ def from_id(cls, state_id, load_data=True):
             try:
                 _logger.debug(f"Loading state {state_id} from database...")
                 if load_data:
-                    _state_cache[state_id] = DatasetState.get(
-                        DatasetState.id == state_id
-                    )
+                    _state_cache[state_id] = DatasetState.get_by_id(state_id)
                 else:
-                    _state_cache[state_id] = DatasetState.select(
-                        DatasetState.type, DatasetState.time
-                    ).get()
-            except orm.DatasetState.DoesNotExist:
+                    _state_cache[state_id] = (
+                        DatasetState.select(DatasetState.type, DatasetState.time)
+                        .where(DatasetState.id == state_id)
+                        .get()
+                    )
+            except DatasetState.DoesNotExist:
                 _logger.warning(f"Could not find state {state_id}.")
                 return None
         return _state_cache[state_id]
@@ -191,8 +191,8 @@ def from_id(cls, ds_id):
             )
         if ds_id not in _dataset_cache:
             try:
-                _dataset_cache[ds_id] = Dataset.get(Dataset.id == ds_id)
-            except orm.Dataset.DoesNotExist:
+                _dataset_cache[ds_id] = Dataset.get_by_id(ds_id)
+            except Dataset.DoesNotExist:
                 _logger.warning(f"Could not find dataset {ds_id}.")
                 return None
         return _dataset_cache[ds_id]

diff --git a/tests/test_dataset.py b/tests/test_dataset.py
@@ -106,3 +106,17 @@ def tests():
         # pre-fetch
         dget.index()
         tests()
+
+    def test_get_missing(self):
+        # These don't exist
+        assert dget.Dataset.from_id("42") is None
+        assert dget.DatasetState.from_id("42") is None
+
+        # But these do
+        ds = dget.Dataset.from_id("1337")
+        assert ds is not None
+        assert ds.id == "1337"
+
+        state = dget.DatasetState.from_id("23")
+        assert state is not None
+        assert state.id == "23"