From 035643d07be44ede36be4d841d586699c1985f88 Mon Sep 17 00:00:00 2001
From: ketiltrout <dvw@phas.ubc.ca>
Date: Wed, 18 Oct 2023 10:59:44 -0700
Subject: [PATCH] fix(get): Fix from_id cache-misses

---
 chimedb/dataset/get.py | 18 +++++++++---------
 tests/test_dataset.py  | 14 ++++++++++++++
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/chimedb/dataset/get.py b/chimedb/dataset/get.py
index 42b8e03..9dada8d 100644
--- 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
index 3b41a6a..b12fc1a 100644
--- 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"