zombie-squirrel 0.3.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

zombie_squirrel/__init__.py

@@ -1,7 +1,14 @@
-"""Init package"""
-__version__ = "0.3.0"
+"""Zombie-squirrel: caching and synchronization for AIND metadata.
 
-from zombie_squirrel.squirrels import (
+Provides functions to fetch and cache project names, subject IDs, and asset
+metadata from the AIND metadata database with support for multiple backends."""
+
+__version__ = "0.4.1"
+
+from zombie_squirrel.squirrels import (  # noqa: F401
+    asset_basics,
+    raw_to_derived,
+    source_data,
     unique_project_names,
     unique_subject_ids,
 )

zombie_squirrel/acorns.py

@@ -1,10 +1,12 @@
-# forest_cache/acorns.py
-from abc import ABC, abstractmethod
+"""Storage backend interfaces for caching data."""
+
 import logging
-import pandas as pd
 import os
+from abc import ABC, abstractmethod
 
+import pandas as pd
 from aind_data_access_api.rds_tables import Client, RDSCredentials
+
 from zombie_squirrel.utils import prefix_table_name
 
 
@@ -12,17 +14,18 @@ class Acorn(ABC):
     """Base class for a storage backend (the cache)."""
 
     def __init__(self) -> None:
+        """Initialize the Acorn."""
         super().__init__()
 
     @abstractmethod
     def hide(self, table_name: str, data: pd.DataFrame) -> None:
         """Store records in the cache."""
-        pass
+        pass  # pragma: no cover
 
     @abstractmethod
     def scurry(self, table_name: str) -> pd.DataFrame:
         """Fetch records from the cache."""
-        pass
+        pass  # pragma: no cover
 
 
 class RedshiftAcorn(Acorn):
@@ -30,37 +33,44 @@ class RedshiftAcorn(Acorn):
     Redshift Client"""
 
     def __init__(self) -> None:
+        """Initialize RedshiftAcorn with Redshift credentials."""
         REDSHIFT_SECRETS = os.getenv("REDSHIFT_SECRETS", "/aind/prod/redshift/credentials/readwrite")
         self.rds_client = Client(
             credentials=RDSCredentials(aws_secrets_name=REDSHIFT_SECRETS),
         )
 
     def hide(self, table_name: str, data: pd.DataFrame) -> None:
+        """Store DataFrame in Redshift table."""
         self.rds_client.overwrite_table_with_df(
             df=data,
             table_name=prefix_table_name(table_name),
         )
 
     def scurry(self, table_name: str) -> pd.DataFrame:
+        """Fetch DataFrame from Redshift table."""
         return self.rds_client.read_table(table_name=prefix_table_name(table_name))
 
 
 class MemoryAcorn(Acorn):
     """A simple in-memory backend for testing or local development."""
+
     def __init__(self) -> None:
+        """Initialize MemoryAcorn with empty store."""
         super().__init__()
         self._store: dict[str, pd.DataFrame] = {}
 
     def hide(self, table_name: str, data: pd.DataFrame) -> None:
+        """Store DataFrame in memory."""
         self._store[table_name] = data
 
     def scurry(self, table_name: str) -> pd.DataFrame:
+        """Fetch DataFrame from memory."""
         return self._store.get(table_name, pd.DataFrame())
 
 
 def rds_get_handle_empty(acorn: Acorn, table_name: str) -> pd.DataFrame:
     """Helper for handling errors when loading from redshift, because
-    there's no helper function """
+    there's no helper function"""
     try:
         logging.info(f"Fetching from cache: {table_name}")
         df = acorn.scurry(table_name)

zombie_squirrel/squirrels.py

@@ -1,10 +1,18 @@
 """Squirrels: functions to fetch and cache data from MongoDB."""
+
+import logging
+import os
+from collections.abc import Callable
+from typing import Any
+
 import pandas as pd
-from typing import Any, Callable
-from zombie_squirrel.acorns import RedshiftAcorn, MemoryAcorn, rds_get_handle_empty
 from aind_data_access_api.document_db import MetadataDbClient
-import os
-import logging
+
+from zombie_squirrel.acorns import (
+    MemoryAcorn,
+    RedshiftAcorn,
+    rds_get_handle_empty,
+)
 
 # --- Backend setup ---------------------------------------------------
 
@@ -12,7 +20,7 @@ API_GATEWAY_HOST = "api.allenneuraldynamics.org"
 
 tree_type = os.getenv("TREE_SPECIES", "memory").lower()
 
-if tree_type == "redshift":
+if tree_type == "redshift":  # pragma: no cover
     logging.info("Using Redshift acorn for caching")
     ACORN = RedshiftAcorn()
 else:
@@ -26,9 +34,12 @@ SQUIRREL_REGISTRY: dict[str, Callable[[], Any]] = {}
 
 def register_squirrel(name: str):
     """Decorator for registering new squirrels."""
+
     def decorator(func):
+        """Register function in squirrel registry."""
         SQUIRREL_REGISTRY[name] = func
         return func
+
     return decorator
 
 
@@ -38,11 +49,23 @@ NAMES = {
     "upn": "unique_project_names",
     "usi": "unique_subject_ids",
     "basics": "asset_basics",
+    "d2r": "source_data",
+    "r2d": "raw_to_derived",
 }
 
 
 @register_squirrel(NAMES["upn"])
 def unique_project_names(force_update: bool = False) -> list[str]:
+    """Fetch unique project names from metadata database.
+
+    Returns cached results if available, fetches from database if cache is empty
+    or force_update is True.
+
+    Args:
+        force_update: If True, bypass cache and fetch fresh data from database.
+
+    Returns:
+        List of unique project names."""
     df = rds_get_handle_empty(ACORN, NAMES["upn"])
 
     if df.empty or force_update:
@@ -66,6 +89,16 @@ def unique_project_names(force_update: bool = False) -> list[str]:
 
 @register_squirrel(NAMES["usi"])
 def unique_subject_ids(force_update: bool = False) -> list[str]:
+    """Fetch unique subject IDs from metadata database.
+
+    Returns cached results if available, fetches from database if cache is empty
+    or force_update is True.
+
+    Args:
+        force_update: If True, bypass cache and fetch fresh data from database.
+
+    Returns:
+        List of unique subject IDs."""
     df = rds_get_handle_empty(ACORN, NAMES["usi"])
 
     if df.empty or force_update:
@@ -89,13 +122,20 @@ def unique_subject_ids(force_update: bool = False) -> list[str]:
 
 @register_squirrel(NAMES["basics"])
 def asset_basics(force_update: bool = False) -> pd.DataFrame:
-    """Basic asset metadata.
+    """Fetch basic asset metadata including modalities, projects, and subject info.
 
-    _id, last_modified,
-    modalities, project names, data_level, subject_id, acquisition_start and _end
-    """
+    Returns a DataFrame with columns: _id, _last_modified, modalities,
+    project_name, data_level, subject_id, acquisition_start_time, and
+    acquisition_end_time. Uses incremental updates based on _last_modified
+    timestamps to avoid re-fetching unchanged records.
+
+    Args:
+        force_update: If True, bypass cache and fetch fresh data from database.
+
+    Returns:
+        DataFrame with basic asset metadata."""
     df = rds_get_handle_empty(ACORN, NAMES["basics"])
-    
+
     FIELDS = [
         "data_description.modalities",
         "data_description.project_name",
@@ -107,6 +147,18 @@ def asset_basics(force_update: bool = False) -> pd.DataFrame:
 
     if df.empty or force_update:
         logging.info("Updating cache for asset basics")
+        df = pd.DataFrame(
+            columns=[
+                "_id",
+                "_last_modified",
+                "modalities",
+                "project_name",
+                "data_level",
+                "subject_id",
+                "acquisition_start_time",
+                "acquisition_end_time",
+            ]
+        )
         client = MetadataDbClient(
             host=API_GATEWAY_HOST,
             version="v2",
@@ -115,31 +167,157 @@ def asset_basics(force_update: bool = False) -> pd.DataFrame:
         # as large as DocDB. We'll also try to limit ourselves to only updating fields
         # that are necessary
         record_ids = client.retrieve_docdb_records(
-            filter_query={}, projection={"_id": 1, "last_modified": 1}, limit=0,
+            filter_query={},
+            projection={"_id": 1, "_last_modified": 1},
+            limit=0,
         )
         keep_ids = []
-        # Drop all _ids where last_modified matches cache
+        # Drop all _ids where _last_modified matches cache
         for record in record_ids:
             cached_row = df[df["_id"] == record["_id"]]
-            if cached_row.empty or cached_row["last_modified"].values[0] != record["last_modified"]:
+            if cached_row.empty or cached_row["_last_modified"].values[0] != record["_last_modified"]:
                 keep_ids.append(record["_id"])
 
         # Now batch by 100 IDs at a time to avoid overloading server, and fetch all the fields
         BATCH_SIZE = 100
         asset_records = []
         for i in range(0, len(keep_ids), BATCH_SIZE):
-            batch_ids = keep_ids[i:i + BATCH_SIZE]
+            logging.info(f"Fetching asset basics batch {i // BATCH_SIZE + 1}...")
+            batch_ids = keep_ids[i : i + BATCH_SIZE]
             batch_records = client.retrieve_docdb_records(
                 filter_query={"_id": {"$in": batch_ids}},
-                projection={field: 1 for field in FIELDS + ["_id", "last_modified"]},
+                projection={field: 1 for field in FIELDS + ["_id", "_last_modified"]},
                 limit=0,
             )
             asset_records.extend(batch_records)
 
+        # Unwrap nested fields
+        records = []
+        for record in asset_records:
+            modalities = record.get("data_description", {}).get("modalities", [])
+            modality_abbreviations = [modality["abbreviation"] for modality in modalities if "abbreviation" in modality]
+            modality_abbreviations_str = ", ".join(modality_abbreviations)
+            flat_record = {
+                "_id": record["_id"],
+                "_last_modified": record.get("_last_modified", None),
+                "modalities": modality_abbreviations_str,
+                "project_name": record.get("data_description", {}).get("project_name", None),
+                "data_level": record.get("data_description", {}).get("data_level", None),
+                "subject_id": record.get("subject", {}).get("subject_id", None),
+                "acquisition_start_time": record.get("acquisition", {}).get("acquisition_start_time", None),
+                "acquisition_end_time": record.get("acquisition", {}).get("acquisition_end_time", None),
+            }
+            records.append(flat_record)
+
         # Combine new records with the old df and store in cache
-        new_df = pd.DataFrame(asset_records)
-        df = pd.concat([df[df["_id"].isin(keep_ids) == False], new_df], ignore_index=True)
+        new_df = pd.DataFrame(records)
+        df = pd.concat([df[~df["_id"].isin(keep_ids)], new_df], ignore_index=True)
 
         ACORN.hide(NAMES["basics"], df)
 
     return df
+
+
+@register_squirrel(NAMES["d2r"])
+def source_data(force_update: bool = False) -> pd.DataFrame:
+    """Fetch source data references for derived records.
+
+    Returns a DataFrame mapping record IDs to their upstream source data
+    dependencies as comma-separated lists.
+
+    Args:
+        force_update: If True, bypass cache and fetch fresh data from database.
+
+    Returns:
+        DataFrame with _id and source_data columns."""
+    df = rds_get_handle_empty(ACORN, NAMES["d2r"])
+
+    if df.empty or force_update:
+        logging.info("Updating cache for source data")
+        client = MetadataDbClient(
+            host=API_GATEWAY_HOST,
+            version="v2",
+        )
+        records = client.retrieve_docdb_records(
+            filter_query={},
+            projection={"_id": 1, "data_description.source_data": 1},
+            limit=0,
+        )
+        data = []
+        for record in records:
+            source_data_list = record.get("data_description", {}).get("source_data", [])
+            source_data_str = ", ".join(source_data_list) if source_data_list else ""
+            data.append(
+                {
+                    "_id": record["_id"],
+                    "source_data": source_data_str,
+                }
+            )
+
+        df = pd.DataFrame(data)
+        ACORN.hide(NAMES["d2r"], df)
+
+    return df
+
+
+@register_squirrel(NAMES["r2d"])
+def raw_to_derived(force_update: bool = False) -> pd.DataFrame:
+    """Fetch mapping of raw records to their derived records.
+
+    Returns a DataFrame mapping raw record IDs to lists of derived record IDs
+    that depend on them as source data.
+
+    Args:
+        force_update: If True, bypass cache and fetch fresh data from database.
+
+    Returns:
+        DataFrame with _id and derived_records columns."""
+    df = rds_get_handle_empty(ACORN, NAMES["r2d"])
+
+    if df.empty or force_update:
+        logging.info("Updating cache for raw to derived mapping")
+        client = MetadataDbClient(
+            host=API_GATEWAY_HOST,
+            version="v2",
+        )
+
+        # Get all raw record IDs
+        raw_records = client.retrieve_docdb_records(
+            filter_query={"data_description.data_level": "raw"},
+            projection={"_id": 1},
+            limit=0,
+        )
+        raw_ids = {record["_id"] for record in raw_records}
+
+        # Get all derived records with their _id and source_data
+        derived_records = client.retrieve_docdb_records(
+            filter_query={"data_description.data_level": "derived"},
+            projection={"_id": 1, "data_description.source_data": 1},
+            limit=0,
+        )
+
+        # Build mapping: raw_id -> list of derived _ids
+        raw_to_derived_map = {raw_id: [] for raw_id in raw_ids}
+        for derived_record in derived_records:
+            source_data_list = derived_record.get("data_description", {}).get("source_data", [])
+            derived_id = derived_record["_id"]
+            # Add this derived record to each raw record it depends on
+            for source_id in source_data_list:
+                if source_id in raw_to_derived_map:
+                    raw_to_derived_map[source_id].append(derived_id)
+
+        # Convert to DataFrame
+        data = []
+        for raw_id, derived_ids in raw_to_derived_map.items():
+            derived_ids_str = ", ".join(derived_ids)
+            data.append(
+                {
+                    "_id": raw_id,
+                    "derived_records": derived_ids_str,
+                }
+            )
+
+        df = pd.DataFrame(data)
+        ACORN.hide(NAMES["r2d"], df)
+
+    return df

zombie_squirrel/sync.py

@@ -1,7 +1,12 @@
-"""Sync all acorns"""
+"""Synchronization utilities for updating all cached data."""
+
 from .squirrels import SQUIRREL_REGISTRY
 
 
 def hide_acorns():
+    """Trigger force update of all registered squirrel functions.
+
+    Calls each squirrel function with force_update=True to refresh
+    all cached data in the acorn backend."""
     for squirrel in SQUIRREL_REGISTRY.values():
         squirrel(force_update=True)

zombie_squirrel/utils.py

@@ -1,5 +1,12 @@
-"""Utility functions"""
+"""Utility functions for zombie-squirrel package."""
 
 
 def prefix_table_name(table_name: str) -> str:
+    """Add zombie-squirrel prefix to table names.
+
+    Args:
+        table_name: The base table name.
+
+    Returns:
+        Table name with 'zs_' prefix."""
     return "zs_" + table_name

{zombie_squirrel-0.3.0.dist-info → zombie_squirrel-0.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: zombie-squirrel
-Version: 0.3.0
+Version: 0.4.1
 Summary: Generated from aind-library-template
 Author: Allen Institute for Neural Dynamics
 License: MIT

zombie_squirrel-0.4.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+zombie_squirrel/__init__.py,sha256=ymbif5B8eqt_ci5zT41uXpZuu37byq_VabBOnXT7UXk,409
+zombie_squirrel/acorns.py,sha256=4uBzYtYgW2oD5sOohNQUw4qfjmNjmAIK2RlL1Ge1Udo,2597
+zombie_squirrel/squirrels.py,sha256=b1kQ2itTBo4o0e0r8Fg56YcJsiJAIqxzs86CSv0ExXE,11181
+zombie_squirrel/sync.py,sha256=lUTv6d3rEpWR-X_qkRkVxycsLmFjOvMvAKcFuWBCwyg,388
+zombie_squirrel/utils.py,sha256=woPxU4vYMUv-T0XOjV5ieViksU_q7It_n_5Ll4zpocA,289
+zombie_squirrel-0.4.1.dist-info/licenses/LICENSE,sha256=U0Y7B3gZJHXpjJVLgTQjM8e_c8w4JJpLgGhIdsoFR1Y,1092
+zombie_squirrel-0.4.1.dist-info/METADATA,sha256=-xxg9BwS_3ISqxE7sxpcpOVhkFdTxet9kQYjkfdw_n0,1382
+zombie_squirrel-0.4.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+zombie_squirrel-0.4.1.dist-info/top_level.txt,sha256=FmM0coe4AangURZLjM4JwwRv2B8H6oINYCoZLKLDCKA,16
+zombie_squirrel-0.4.1.dist-info/RECORD,,

zombie_squirrel-0.3.0.dist-info/RECORD

@@ -1,10 +0,0 @@
-zombie_squirrel/__init__.py,sha256=SU-iCf1qQAQUKc6VnAdEUcGBlcUwuDZ3CKpJXGSrzb0,134
-zombie_squirrel/acorns.py,sha256=1mCnWCDFRnbHLddCCgiUG3RumuKUjMKVbyTVoYI0FB8,2188
-zombie_squirrel/squirrels.py,sha256=BkS4l9cUB7ZXQV3ySa0HCgSLei_Kc47X28BJY69Bz54,4915
-zombie_squirrel/sync.py,sha256=jslTVIend5Z-sLJuNXKkhn-nqmKK_P0FAiRuFFYRnto,168
-zombie_squirrel/utils.py,sha256=74DSFK1Qbp8yQeUXpnli4kqx_QcAc8v4_6FZut0xZ8g,103
-zombie_squirrel-0.3.0.dist-info/licenses/LICENSE,sha256=U0Y7B3gZJHXpjJVLgTQjM8e_c8w4JJpLgGhIdsoFR1Y,1092
-zombie_squirrel-0.3.0.dist-info/METADATA,sha256=3YrkG4IV-36OekD-DEjPlxAbVOXvuHzjesCarDgJ2eQ,1382
-zombie_squirrel-0.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-zombie_squirrel-0.3.0.dist-info/top_level.txt,sha256=FmM0coe4AangURZLjM4JwwRv2B8H6oINYCoZLKLDCKA,16
-zombie_squirrel-0.3.0.dist-info/RECORD,,