zombie-squirrel 0.7.4__py3-none-any.whl → 0.8.1__py3-none-any.whl

zombie_squirrel/__init__.py

@@ -3,12 +3,10 @@
 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.7.4"
+__version__ = "0.8.1"
 
-from zombie_squirrel.squirrels import (  # noqa: F401
-    asset_basics,
-    raw_to_derived,
-    source_data,
-    unique_project_names,
-    unique_subject_ids,
-)
+from zombie_squirrel.acorn_contents.asset_basics import asset_basics  # noqa: F401
+from zombie_squirrel.acorn_contents.raw_to_derived import raw_to_derived  # noqa: F401
+from zombie_squirrel.acorn_contents.source_data import source_data  # noqa: F401
+from zombie_squirrel.acorn_contents.unique_project_names import unique_project_names  # noqa: F401
+from zombie_squirrel.acorn_contents.unique_subject_ids import unique_subject_ids  # noqa: F401

zombie_squirrel/acorn_contents/__init__.py

@@ -0,0 +1,10 @@
+"""Acorns module: individual data fetching functions."""
+
+# Import the acorn modules to trigger decorator registration
+from zombie_squirrel.acorn_contents import (  # noqa: F401
+    asset_basics,
+    raw_to_derived,
+    source_data,
+    unique_project_names,
+    unique_subject_ids,
+)

zombie_squirrel/acorn_contents/asset_basics.py

@@ -0,0 +1,135 @@
+"""Asset basics acorn."""
+
+import logging
+
+import pandas as pd
+from aind_data_access_api.document_db import MetadataDbClient
+
+import zombie_squirrel.acorns as acorns
+
+
+@acorns.register_acorn(acorns.NAMES["basics"])
+def asset_basics(force_update: bool = False) -> pd.DataFrame:
+    """Fetch basic asset metadata including modalities, projects, and subject info.
+
+    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 = acorns.TREE.scurry(acorns.NAMES["basics"])
+
+    FIELDS = [
+        "data_description.modalities",
+        "data_description.project_name",
+        "data_description.data_level",
+        "subject.subject_id",
+        "acquisition.acquisition_start_time",
+        "acquisition.acquisition_end_time",
+        "processing.data_processes.start_date_time",
+        "subject.subject_details.genotype",
+        "other_identifiers",
+        "location",
+    ]
+
+    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",
+                "code_ocean",
+                "process_date",
+                "genotype",
+                "location",
+            ]
+        )
+        client = MetadataDbClient(
+            host=acorns.API_GATEWAY_HOST,
+            version="v2",
+        )
+        # It's a bit complex to get multiple fields that aren't indexed in a database
+        # 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,
+        )
+        keep_ids = []
+        # 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"]:
+                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):
+            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"]},
+                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)
+
+            # Get the process date, convert to YYYY-MM-DD if present
+            data_processes = record.get("processing", {}).get("data_processes", [])
+            if data_processes:
+                latest_process = data_processes[-1]
+                process_datetime = latest_process.get("start_date_time", None)
+                process_date = process_datetime.split("T")[0]
+            else:
+                process_date = None
+
+            # Get the CO asset ID
+            other_identifiers = record.get("other_identifiers", {})
+            if other_identifiers:
+                code_ocean = other_identifiers.get("Code Ocean", None)
+            else:
+                code_ocean = None
+
+            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),
+                "code_ocean": code_ocean,
+                "process_date": process_date,
+                "genotype": record.get("subject", {}).get("subject_details", {}).get("genotype", None),
+                "location": record.get("location", None),
+            }
+            records.append(flat_record)
+
+        # Combine new records with the old df and store in cache
+        new_df = pd.DataFrame(records)
+        df = pd.concat([df[~df["_id"].isin(keep_ids)], new_df], ignore_index=True)
+
+        acorns.TREE.hide(acorns.NAMES["basics"], df)
+
+    return df

zombie_squirrel/acorn_contents/raw_to_derived.py

@@ -0,0 +1,71 @@
+"""Raw to derived mapping acorn."""
+
+import logging
+
+import pandas as pd
+from aind_data_access_api.document_db import MetadataDbClient
+
+import zombie_squirrel.acorns as acorns
+
+
+@acorns.register_acorn(acorns.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 = acorns.TREE.scurry(acorns.NAMES["r2d"])
+
+    if df.empty or force_update:
+        logging.info("Updating cache for raw to derived mapping")
+        client = MetadataDbClient(
+            host=acorns.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)
+        acorns.TREE.hide(acorns.NAMES["r2d"], df)
+
+    return df

zombie_squirrel/acorn_contents/source_data.py

@@ -0,0 +1,50 @@
+"""Source data acorn."""
+
+import logging
+
+import pandas as pd
+from aind_data_access_api.document_db import MetadataDbClient
+
+import zombie_squirrel.acorns as acorns
+
+
+@acorns.register_acorn(acorns.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 = acorns.TREE.scurry(acorns.NAMES["d2r"])
+
+    if df.empty or force_update:
+        logging.info("Updating cache for source data")
+        client = MetadataDbClient(
+            host=acorns.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)
+        acorns.TREE.hide(acorns.NAMES["d2r"], df)
+
+    return df

zombie_squirrel/acorn_contents/unique_project_names.py

@@ -0,0 +1,41 @@
+"""Unique project names acorn."""
+
+import logging
+
+import pandas as pd
+from aind_data_access_api.document_db import MetadataDbClient
+
+import zombie_squirrel.acorns as acorns
+
+
+@acorns.register_acorn(acorns.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 = acorns.TREE.scurry(acorns.NAMES["upn"])
+
+    if df.empty or force_update:
+        # If cache is missing, fetch data
+        logging.info("Updating cache for unique project names")
+        client = MetadataDbClient(
+            host=acorns.API_GATEWAY_HOST,
+            version="v2",
+        )
+        unique_project_names = client.aggregate_docdb_records(
+            pipeline=[
+                {"$group": {"_id": "$data_description.project_name"}},
+                {"$project": {"project_name": "$_id", "_id": 0}},
+            ]
+        )
+        df = pd.DataFrame(unique_project_names)
+        acorns.TREE.hide(acorns.NAMES["upn"], df)
+
+    return df["project_name"].tolist()

zombie_squirrel/acorn_contents/unique_subject_ids.py

@@ -0,0 +1,41 @@
+"""Unique subject IDs acorn."""
+
+import logging
+
+import pandas as pd
+from aind_data_access_api.document_db import MetadataDbClient
+
+import zombie_squirrel.acorns as acorns
+
+
+@acorns.register_acorn(acorns.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 = acorns.TREE.scurry(acorns.NAMES["usi"])
+
+    if df.empty or force_update:
+        # If cache is missing, fetch data
+        logging.info("Updating cache for unique subject IDs")
+        client = MetadataDbClient(
+            host=acorns.API_GATEWAY_HOST,
+            version="v2",
+        )
+        unique_subject_ids = client.aggregate_docdb_records(
+            pipeline=[
+                {"$group": {"_id": "$subject.subject_id"}},
+                {"$project": {"subject_id": "$_id", "_id": 0}},
+            ]
+        )
+        df = pd.DataFrame(unique_subject_ids)
+        acorns.TREE.hide(acorns.NAMES["usi"], df)
+
+    return df["subject_id"].tolist()

zombie_squirrel/acorns.py

@@ -1,96 +1,47 @@
-"""Storage backend interfaces for caching data."""
+"""Acorns: functions to fetch and cache data from MongoDB."""
 
-import io
 import logging
-from abc import ABC, abstractmethod
+import os
+from collections.abc import Callable
+from typing import Any
 
-import boto3
-import duckdb
-import pandas as pd
+from zombie_squirrel.forest import (
+    MemoryTree,
+    S3Tree,
+)
 
-from zombie_squirrel.utils import get_s3_cache_path, prefix_table_name
+# --- Backend setup ---------------------------------------------------
 
+API_GATEWAY_HOST = "api.allenneuraldynamics.org"
 
-class Acorn(ABC):
-    """Base class for a storage backend (the cache)."""
+forest_type = os.getenv("FOREST_TYPE", "memory").lower()
 
-    def __init__(self) -> None:
-        """Initialize the Acorn."""
-        super().__init__()
+if forest_type == "S3":  # pragma: no cover
+    logging.info("Using S3 forest for caching")
+    TREE = S3Tree()
+else:
+    logging.info("Using in-memory forest for caching")
+    TREE = MemoryTree()
 
-    @abstractmethod
-    def hide(self, table_name: str, data: pd.DataFrame) -> None:
-        """Store records in the cache."""
-        pass  # pragma: no cover
+# --- Acorn registry and names -----------------------------------------------------
 
-    @abstractmethod
-    def scurry(self, table_name: str) -> pd.DataFrame:
-        """Fetch records from the cache."""
-        pass  # pragma: no cover
+NAMES = {
+    "upn": "unique_project_names",
+    "usi": "unique_subject_ids",
+    "basics": "asset_basics",
+    "d2r": "source_data",
+    "r2d": "raw_to_derived",
+}
 
+ACORN_REGISTRY: dict[str, Callable[[], Any]] = {}
 
-class S3Acorn(Acorn):
-    """Stores and retrieves caches using AWS S3 with parquet files."""
 
-    def __init__(self) -> None:
-        """Initialize S3Acorn with S3 client."""
-        self.bucket = "aind-scratch-data"
-        self.s3_client = boto3.client("s3")
+def register_acorn(name: str):
+    """Decorator for registering new acorns."""
 
-    def hide(self, table_name: str, data: pd.DataFrame) -> None:
-        """Store DataFrame as parquet file in S3."""
-        filename = prefix_table_name(table_name)
-        s3_key = get_s3_cache_path(filename)
+    def decorator(func):
+        """Register function in acorn registry."""
+        ACORN_REGISTRY[name] = func
+        return func
 
-        # Convert DataFrame to parquet bytes
-        parquet_buffer = io.BytesIO()
-        data.to_parquet(parquet_buffer, index=False)
-        parquet_buffer.seek(0)
-
-        # Upload to S3
-        self.s3_client.put_object(
-            Bucket=self.bucket,
-            Key=s3_key,
-            Body=parquet_buffer.getvalue(),
-        )
-        logging.info(f"Stored cache to S3: s3://{self.bucket}/{s3_key}")
-
-    def scurry(self, table_name: str) -> pd.DataFrame:
-        """Fetch DataFrame from S3 parquet file."""
-        filename = prefix_table_name(table_name)
-        s3_key = get_s3_cache_path(filename)
-
-        try:
-            # Read directly from S3 using DuckDB
-            query = f"""
-                SELECT * FROM read_parquet(
-                    's3://{self.bucket}/{s3_key}'
-                )
-            """
-            result = duckdb.query(query).to_df()
-            logging.info(
-                f"Retrieved cache from S3: s3://{self.bucket}/{s3_key}"
-            )
-            return result
-        except Exception as e:
-            logging.warning(
-                f"Error fetching from cache {s3_key}: {e}"
-            )
-            return pd.DataFrame()
-
-
-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())
+    return decorator

zombie_squirrel/forest.py

@@ -0,0 +1,96 @@
+"""Storage backend interfaces for caching data."""
+
+import io
+import logging
+from abc import ABC, abstractmethod
+
+import boto3
+import duckdb
+import pandas as pd
+
+from zombie_squirrel.utils import get_s3_cache_path, prefix_table_name
+
+
+class Tree(ABC):
+    """Base class for a storage backend (the cache)."""
+
+    def __init__(self) -> None:
+        """Initialize the Tree."""
+        super().__init__()
+
+    @abstractmethod
+    def hide(self, table_name: str, data: pd.DataFrame) -> None:
+        """Store records in the cache."""
+        pass  # pragma: no cover
+
+    @abstractmethod
+    def scurry(self, table_name: str) -> pd.DataFrame:
+        """Fetch records from the cache."""
+        pass  # pragma: no cover
+
+
+class S3Tree(Tree):
+    """Stores and retrieves caches using AWS S3 with parquet files."""
+
+    def __init__(self) -> None:
+        """Initialize S3Acorn with S3 client."""
+        self.bucket = "aind-scratch-data"
+        self.s3_client = boto3.client("s3")
+
+    def hide(self, table_name: str, data: pd.DataFrame) -> None:
+        """Store DataFrame as parquet file in S3."""
+        filename = prefix_table_name(table_name)
+        s3_key = get_s3_cache_path(filename)
+
+        # Convert DataFrame to parquet bytes
+        parquet_buffer = io.BytesIO()
+        data.to_parquet(parquet_buffer, index=False)
+        parquet_buffer.seek(0)
+
+        # Upload to S3
+        self.s3_client.put_object(
+            Bucket=self.bucket,
+            Key=s3_key,
+            Body=parquet_buffer.getvalue(),
+        )
+        logging.info(f"Stored cache to S3: s3://{self.bucket}/{s3_key}")
+
+    def scurry(self, table_name: str) -> pd.DataFrame:
+        """Fetch DataFrame from S3 parquet file."""
+        filename = prefix_table_name(table_name)
+        s3_key = get_s3_cache_path(filename)
+
+        try:
+            # Read directly from S3 using DuckDB
+            query = f"""
+                SELECT * FROM read_parquet(
+                    's3://{self.bucket}/{s3_key}'
+                )
+            """
+            result = duckdb.query(query).to_df()
+            logging.info(
+                f"Retrieved cache from S3: s3://{self.bucket}/{s3_key}"
+            )
+            return result
+        except Exception as e:
+            logging.warning(
+                f"Error fetching from cache {s3_key}: {e}"
+            )
+            return pd.DataFrame()
+
+
+class MemoryTree(Tree):
+    """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())

zombie_squirrel/sync.py

@@ -2,17 +2,17 @@
 
 import logging
 
-from .squirrels import SQUIRREL_REGISTRY
+from .acorns import ACORN_REGISTRY
 
 
 def hide_acorns():
-    """Trigger force update of all registered squirrel functions.
+    """Trigger force update of all registered acorn functions.
 
-    Calls each squirrel function with force_update=True to refresh
-    all cached data in the acorn backend."""
+    Calls each acorn function with force_update=True to refresh
+    all cached data in the tree backend."""
     logging.basicConfig(
         level=logging.INFO,
         format="%(asctime)s %(levelname)s %(message)s"
     )
-    for squirrel in SQUIRREL_REGISTRY.values():
-        squirrel(force_update=True)
+    for acorn in ACORN_REGISTRY.values():
+        acorn(force_update=True)

{zombie_squirrel-0.7.4.dist-info → zombie_squirrel-0.8.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: zombie-squirrel
-Version: 0.7.4
+Version: 0.8.1
 Summary: Generated from aind-library-template
 Author: Allen Institute for Neural Dynamics
 License: MIT
@@ -21,7 +21,7 @@ Dynamic: license-file
 ![Code Style](https://img.shields.io/badge/code%20style-black-black)
 [![semantic-release: angular](https://img.shields.io/badge/semantic--release-angular-e10079?logo=semantic-release)](https://github.com/semantic-release/semantic-release)
 ![Interrogate](https://img.shields.io/badge/interrogate-100.0%25-brightgreen)
-![Coverage](https://img.shields.io/badge/coverage-99%25-brightgreen)
+![Coverage](https://img.shields.io/badge/coverage-100%25-brightgreen)
 ![Python](https://img.shields.io/badge/python->=3.10-blue?logo=python)
 
 <img src="zombie-squirrel_logo.png" width="400" alt="Logo (image from ChatGPT)">
@@ -37,10 +37,10 @@ pip install zombie-squirrel
 ### Set backend
 
 ```bash
-export TREE_SPECIES='s3'
+export FOREST_TYPE='S3'
 ```
 
-Options are 's3', 'MEMORY'.
+Options are 'S3', 'MEMORY'.
 
 ### Scurry (fetch) data
 

zombie_squirrel-0.8.1.dist-info/RECORD

@@ -0,0 +1,16 @@
+zombie_squirrel/__init__.py,sha256=sPD8B2awlREi8c33mhKoGw0khdE5u1XQmU_3hWiVH9A,693
+zombie_squirrel/acorns.py,sha256=BLhySoLfb7D6IfbNMOhxbLJba-8Wp-L9e1w_2KozjjM,1134
+zombie_squirrel/forest.py,sha256=v0K1u0EA0OptzxocFC-fPEi6xYcnJ9SoWJ6aiPF4jLg,2939
+zombie_squirrel/sync.py,sha256=9cpfSzTj0cQz4-d3glMAOejCZgekMirLc-dwEFFQhlg,496
+zombie_squirrel/utils.py,sha256=kojQpHUKlRJD7WEZDfcpQIZTj9iUrtX5_6F-gWWzJW0,628
+zombie_squirrel/acorn_contents/__init__.py,sha256=LsNy5xjlJS981SVxRLD4yPhUg__fuY8-HbPDapOhH-o,285
+zombie_squirrel/acorn_contents/asset_basics.py,sha256=cTg016sZbrJH5_As7CqnirumR1hSMn4GiYX0IReoou0,5513
+zombie_squirrel/acorn_contents/raw_to_derived.py,sha256=EFjLP9szzk3Q6jNhyDwxx8rqhG1bFTVmxeCpCAsv6yg,2445
+zombie_squirrel/acorn_contents/source_data.py,sha256=zbw3DnTwXo6xDTZ1uCoIbcYrVBH2r2xZN3trKtxT8cg,1527
+zombie_squirrel/acorn_contents/unique_project_names.py,sha256=Fu0P2DyI91W7xhx1uRk2YVMyFG8WPihDCj7n_A7zL2E,1289
+zombie_squirrel/acorn_contents/unique_subject_ids.py,sha256=dOVn1ObDF86p8S8US__hlyjTGxT0vz01oeGzjoWGFIc,1260
+zombie_squirrel-0.8.1.dist-info/licenses/LICENSE,sha256=U0Y7B3gZJHXpjJVLgTQjM8e_c8w4JJpLgGhIdsoFR1Y,1092
+zombie_squirrel-0.8.1.dist-info/METADATA,sha256=xEQ04ahuoz0XhbxLpAUCYFLG1kj6VSV2RFMfRtAq1PY,1898
+zombie_squirrel-0.8.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+zombie_squirrel-0.8.1.dist-info/top_level.txt,sha256=FmM0coe4AangURZLjM4JwwRv2B8H6oINYCoZLKLDCKA,16
+zombie_squirrel-0.8.1.dist-info/RECORD,,

{zombie_squirrel-0.7.4.dist-info → zombie_squirrel-0.8.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

zombie_squirrel/squirrels.py

@@ -1,355 +0,0 @@
-"""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 aind_data_access_api.document_db import MetadataDbClient
-
-from zombie_squirrel.acorns import (
-    MemoryAcorn,
-    S3Acorn,
-)
-
-# --- Backend setup ---------------------------------------------------
-
-API_GATEWAY_HOST = "api.allenneuraldynamics.org"
-
-tree_type = os.getenv("TREE_SPECIES", "memory").lower()
-
-if tree_type == "s3":  # pragma: no cover
-    logging.info("Using S3 acorn for caching")
-    ACORN = S3Acorn()
-else:
-    logging.info("Using in-memory acorn for caching")
-    ACORN = MemoryAcorn()
-
-# --- Squirrel registry -----------------------------------------------------
-
-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
-
-
-# --- Squirrels -----------------------------------------------------
-
-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 = ACORN.scurry(NAMES["upn"])
-
-    if df.empty or force_update:
-        # If cache is missing, fetch data
-        logging.info("Updating cache for unique project names")
-        client = MetadataDbClient(
-            host=API_GATEWAY_HOST,
-            version="v2",
-        )
-        unique_project_names = client.aggregate_docdb_records(
-            pipeline=[
-                {"$group": {"_id": "$data_description.project_name"}},
-                {"$project": {"project_name": "$_id", "_id": 0}},
-            ]
-        )
-        df = pd.DataFrame(unique_project_names)
-        ACORN.hide(NAMES["upn"], df)
-
-    return df["project_name"].tolist()
-
-
-@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 = ACORN.scurry(NAMES["usi"])
-
-    if df.empty or force_update:
-        # If cache is missing, fetch data
-        logging.info("Updating cache for unique subject IDs")
-        client = MetadataDbClient(
-            host=API_GATEWAY_HOST,
-            version="v2",
-        )
-        unique_subject_ids = client.aggregate_docdb_records(
-            pipeline=[
-                {"$group": {"_id": "$subject.subject_id"}},
-                {"$project": {"subject_id": "$_id", "_id": 0}},
-            ]
-        )
-        df = pd.DataFrame(unique_subject_ids)
-        ACORN.hide(NAMES["usi"], df)
-
-    return df["subject_id"].tolist()
-
-
-@register_squirrel(NAMES["basics"])
-def asset_basics(force_update: bool = False) -> pd.DataFrame:
-    """Fetch basic asset metadata including modalities, projects, and subject info.
-
-    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 = ACORN.scurry(NAMES["basics"])
-
-    FIELDS = [
-        "data_description.modalities",
-        "data_description.project_name",
-        "data_description.data_level",
-        "subject.subject_id",
-        "acquisition.acquisition_start_time",
-        "acquisition.acquisition_end_time",
-        "processing.data_processes.start_date_time",
-        "subject.subject_details.genotype",
-        "other_identifiers",
-        "location",
-        "name",
-    ]
-
-    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",
-                "code_ocean",
-                "process_date",
-                "genotype",
-                "location",
-                "name",
-            ]
-        )
-        client = MetadataDbClient(
-            host=API_GATEWAY_HOST,
-            version="v2",
-        )
-        # It's a bit complex to get multiple fields that aren't indexed in a database
-        # 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,
-        )
-        keep_ids = []
-        # 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"]:
-                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):
-            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"]},
-                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)
-
-            # Get the process date, convert to YYYY-MM-DD if present
-            data_processes = record.get("processing", {}).get("data_processes", [])
-            if data_processes:
-                latest_process = data_processes[-1]
-                process_datetime = latest_process.get("start_date_time", None)
-                process_date = process_datetime.split("T")[0]
-            else:
-                process_date = None
-
-            # Get the CO asset ID
-            other_identifiers = record.get("other_identifiers", {})
-            code_ocean = None
-            if other_identifiers:
-                co_list = other_identifiers.get("Code Ocean", None)
-                if co_list:
-                    code_ocean = co_list[0]
-
-            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),
-                "code_ocean": code_ocean,
-                "process_date": process_date,
-                "genotype": record.get("subject", {}).get("subject_details", {}).get("genotype", None),
-                "location": record.get("location", None),
-                "name": record.get("name", None),
-            }
-            records.append(flat_record)
-
-        # Combine new records with the old df and store in cache
-        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 = ACORN.scurry(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 = ACORN.scurry(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-0.7.4.dist-info/RECORD

@@ -1,10 +0,0 @@
-zombie_squirrel/__init__.py,sha256=NEzSvrO6WxDaetYvZ7EoSWRQqnTnNxe5XKoKPHt8XGk,409
-zombie_squirrel/acorns.py,sha256=mpinFacaN9BM6CvRy0M76JMb6n3oVPZLJxn8O4J9Wlw,2945
-zombie_squirrel/squirrels.py,sha256=1leLr5gA3gPa39NSLNmTOVFjxC29caciPwEcTNCgj6Y,12399
-zombie_squirrel/sync.py,sha256=84Ta5beHiPuGBVzp9SCo7G1b4McTUohcUIf_TJbNIV8,518
-zombie_squirrel/utils.py,sha256=kojQpHUKlRJD7WEZDfcpQIZTj9iUrtX5_6F-gWWzJW0,628
-zombie_squirrel-0.7.4.dist-info/licenses/LICENSE,sha256=U0Y7B3gZJHXpjJVLgTQjM8e_c8w4JJpLgGhIdsoFR1Y,1092
-zombie_squirrel-0.7.4.dist-info/METADATA,sha256=kVo8V2VH1YpNwJOh9TqQJ2vN780BwQ8Jk6QgjTy3y1o,1898
-zombie_squirrel-0.7.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-zombie_squirrel-0.7.4.dist-info/top_level.txt,sha256=FmM0coe4AangURZLjM4JwwRv2B8H6oINYCoZLKLDCKA,16
-zombie_squirrel-0.7.4.dist-info/RECORD,,