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

zombie_squirrel/__init__.py

@@ -1,5 +1,5 @@
 """Init package"""
-__version__ = "0.2.3"
+__version__ = "0.3.0"
 
 from zombie_squirrel.squirrels import (
     unique_project_names,

zombie_squirrel/acorns.py

@@ -59,7 +59,8 @@ class MemoryAcorn(Acorn):
 
 
 def rds_get_handle_empty(acorn: Acorn, table_name: str) -> pd.DataFrame:
-    """Utility function for testing purposes."""
+    """Helper for handling errors when loading from redshift, because
+    there's no helper function """
     try:
         logging.info(f"Fetching from cache: {table_name}")
         df = acorn.scurry(table_name)

zombie_squirrel/squirrels.py

@@ -37,6 +37,7 @@ def register_squirrel(name: str):
 NAMES = {
     "upn": "unique_project_names",
     "usi": "unique_subject_ids",
+    "basics": "asset_basics",
 }
 
 
@@ -84,3 +85,61 @@ def unique_subject_ids(force_update: bool = False) -> list[str]:
         ACORN.hide(NAMES["usi"], df)
 
     return df["subject_id"].tolist()
+
+
+@register_squirrel(NAMES["basics"])
+def asset_basics(force_update: bool = False) -> pd.DataFrame:
+    """Basic asset metadata.
+
+    _id, last_modified,
+    modalities, project names, data_level, subject_id, acquisition_start and _end
+    """
+    df = rds_get_handle_empty(ACORN, 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",
+    ]
+
+    if df.empty or force_update:
+        logging.info("Updating cache for asset basics")
+        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):
+            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)
+
+        # 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)
+
+        ACORN.hide(NAMES["basics"], df)
+
+    return df

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

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

zombie_squirrel-0.3.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+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,,

zombie_squirrel-0.2.3.dist-info/RECORD

@@ -1,10 +0,0 @@
-zombie_squirrel/__init__.py,sha256=i7J51pQP-ID_TrBZtNQmxjgzlqJ8WGQIObfHAzospx8,134
-zombie_squirrel/acorns.py,sha256=MZaScwDpnuuGnrx8a1vRmfv5-fr6h4Idw1_rQ2FWdB0,2132
-zombie_squirrel/squirrels.py,sha256=3ybJQpuNsoM8gBkHSWOKBZ_zOfsnzq35TKh0Aig2voc,2662
-zombie_squirrel/sync.py,sha256=jslTVIend5Z-sLJuNXKkhn-nqmKK_P0FAiRuFFYRnto,168
-zombie_squirrel/utils.py,sha256=74DSFK1Qbp8yQeUXpnli4kqx_QcAc8v4_6FZut0xZ8g,103
-zombie_squirrel-0.2.3.dist-info/licenses/LICENSE,sha256=U0Y7B3gZJHXpjJVLgTQjM8e_c8w4JJpLgGhIdsoFR1Y,1092
-zombie_squirrel-0.2.3.dist-info/METADATA,sha256=o82tnKDGfBAcDwt3fiSpF14octxGxKym5BmluyTyTg0,1382
-zombie_squirrel-0.2.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-zombie_squirrel-0.2.3.dist-info/top_level.txt,sha256=FmM0coe4AangURZLjM4JwwRv2B8H6oINYCoZLKLDCKA,16
-zombie_squirrel-0.2.3.dist-info/RECORD,,