ENH: add search term for retrieving Entries that are nested under an ancestor

docs/source/upcoming_release_notes/108-ancestor_search_option.rst

@@ -0,0 +1,22 @@
+108 ancestor search option
+#################
+
+API Breaks
+----------
+- N/A
+
+Features
+--------
+- Add "ancestor" search option, returning entries reachable from the ancestor
+
+Bugfixes
+--------
+- Flatten and cache entries in FilestoreBackend.save_entry
+
+Maintenance
+-----------
+- N/A
+
+Contributors
+------------
+- shilorigins

superscore/backends/filestore.py

@@ -8,14 +8,15 @@
 import os
 import shutil
 from dataclasses import fields, replace
-from typing import Any, Dict, Generator, Optional, Union
+from functools import cache
+from typing import Any, Container, Dict, Generator, Optional, Union
 from uuid import UUID, uuid4
 
 from apischema import deserialize, serialize
 
 from superscore.backends.core import SearchTermType, _Backend
 from superscore.errors import BackendError
-from superscore.model import Entry, Root
+from superscore.model import Entry, Nestable, Root
 from superscore.utils import build_abs_path
 
 logger = logging.getLogger(__name__)
@@ -263,7 +264,7 @@ def save_entry(self, entry: Entry) -> None:
             if db.get(entry.uuid):
                 raise BackendError("Entry already exists, try updating the entry "
                                    "instead of saving it")
-            db[entry.uuid] = entry
+            self.flatten_and_cache(entry)
             self._root.entries.append(entry)
 
     def update_entry(self, entry: Entry) -> None:
@@ -285,13 +286,17 @@ def search(self, *search_terms: SearchTermType) -> Generator[Entry, None, None]:
         Keys are attributes on `Entry` subclasses, or special keywords.
         Values can be a single value or a tuple of values depending on operator.
         """
+        reachable = cache(self._gather_reachable)
         with self._load_and_store_context() as db:
             for entry in db.values():
                 conditions = []
                 for attr, op, target in search_terms:
                     # TODO: search for child pvs?
                     if attr == "entry_type":
                         conditions.append(isinstance(entry, target))
+                    elif attr == "ancestor":
+                        # `target` must be UUID since `reachable` is cached
+                        conditions.append(entry.uuid in reachable(target))
                     else:
                         try:
                             # check entry attribute by name
@@ -302,6 +307,22 @@ def search(self, *search_terms: SearchTermType) -> Generator[Entry, None, None]:
                 if all(conditions):
                     yield entry
 
+    def _gather_reachable(self, ancestor: Union[Entry, UUID]) -> Container[UUID]:
+        """
+        Finds all entries accessible from ancestor, including ancestor, and returns
+        their UUIDs. This makes it easy to check if one entry is hierarchically under another.
+        """
+        reachable = set()
+        q = [ancestor]
+        while len(q) > 0:
+            cur = q.pop()
+            if not isinstance(cur, Entry):
+                cur = self._entry_cache[cur]
+            reachable.add(cur.uuid)
+            if isinstance(cur, Nestable):
+                q.extend(cur.children)
+        return reachable
+
     @contextlib.contextmanager
     def _load_and_store_context(self) -> Generator[Dict[UUID, Any], None, None]:
         """

superscore/tests/test_backend.py

@@ -202,3 +202,17 @@ def test_update_entry(backends: _Backend):
     p1 = Parameter()
     with pytest.raises(BackendError):
         backends.update_entry(p1)
+
+
+@pytest.mark.parametrize("filestore_backend", [("linac_data",)], indirect=True)
+def test_gather_reachable(filestore_backend: _Backend):
+    # top-level snapshot
+    reachable = filestore_backend._gather_reachable(UUID("06282731-33ea-4270-ba14-098872e627dc"))
+    assert len(reachable) == 32
+    assert UUID("927ef6cb-e45f-4175-aa5f-6c6eec1f3ae4") in reachable
+
+    # direct parent snapshot; works with UUID or Entry
+    entry = filestore_backend.get_entry(UUID("2f709b4b-79da-4a8b-8693-eed2c389cb3a"))
+    reachable = filestore_backend._gather_reachable(entry)
+    assert len(reachable) == 3
+    assert UUID("927ef6cb-e45f-4175-aa5f-6c6eec1f3ae4") in reachable

superscore/tests/test_client.py

@@ -211,8 +211,39 @@ def test_fill_depth(fill_depth: int):
 
 
 @pytest.mark.parametrize("filestore_backend", [("linac_with_comparison_snapshot",)], indirect=True)
-def test_parametrized_filestore(sample_client: Client):
-    assert len(list(sample_client.search())) > 0
+def test_search_entries_by_ancestor(sample_client: Client):
+    entries = tuple(sample_client.search(
+        ("entry_type", "eq", Setpoint),
+        ("pv_name", "eq", "LASR:GUNB:TEST1"),
+    ))
+    assert len(entries) == 2
+    entries = tuple(sample_client.search(
+        ("entry_type", "eq", Setpoint),
+        ("pv_name", "eq", "LASR:GUNB:TEST1"),
+        ("ancestor", "eq", UUID("06282731-33ea-4270-ba14-098872e627dc")),  # top-level snapshot
+    ))
+    assert len(entries) == 1
+    entries = tuple(sample_client.search(
+        ("entry_type", "eq", Setpoint),
+        ("pv_name", "eq", "LASR:GUNB:TEST1"),
+        ("ancestor", "eq", UUID("2f709b4b-79da-4a8b-8693-eed2c389cb3a")),  # direct parent
+    ))
+    assert len(entries) == 1
+
+
+@pytest.mark.parametrize("filestore_backend", [("linac_with_comparison_snapshot",)], indirect=True)
+def test_search_caching(sample_client: Client):
+    entry = sample_client.backend.get_entry(UUID("2f709b4b-79da-4a8b-8693-eed2c389cb3a"))
+    result = sample_client.search(
+        ("ancestor", "eq", UUID("2f709b4b-79da-4a8b-8693-eed2c389cb3a")),
+    )
+    assert len(tuple(result)) == 3
+    entry.children = []
+    sample_client.backend.update_entry(entry)
+    result = sample_client.search(
+        ("ancestor", "eq", UUID("2f709b4b-79da-4a8b-8693-eed2c389cb3a")),
+    )
+    assert len(tuple(result)) == 1  # update is picked up in new search
 
 
 def test_parametrized_filestore_empty(sample_client: Client):