zombie-squirrel 0.5.0__py3-none-any.whl → 0.6.1__py3-none-any.whl

zombie_squirrel/__init__.py

@@ -3,7 +3,7 @@
 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.5.0"
+__version__ = "0.6.1"
 
 from zombie_squirrel.squirrels import (  # noqa: F401
     asset_basics,

zombie_squirrel/acorns.py

@@ -1,13 +1,14 @@
 """Storage backend interfaces for caching data."""
 
+import io
 import logging
-import os
 from abc import ABC, abstractmethod
 
+import boto3
+import duckdb
 import pandas as pd
-from aind_data_access_api.rds_tables import Client, RDSCredentials
 
-from zombie_squirrel.utils import prefix_table_name
+from zombie_squirrel.utils import get_s3_cache_path, prefix_table_name
 
 
 class Acorn(ABC):
@@ -28,27 +29,54 @@ class Acorn(ABC):
         pass  # pragma: no cover
 
 
-class RedshiftAcorn(Acorn):
-    """Stores and retrieves caches using aind-data-access-api
-    Redshift Client"""
+class S3Acorn(Acorn):
+    """Stores and retrieves caches using AWS S3 with parquet files."""
 
     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),
-        )
+        """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 in Redshift table."""
-        self.rds_client.overwrite_table_with_df(
-            df=data,
-            table_name=prefix_table_name(table_name),
+        """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 Redshift table."""
-        return self.rds_client.read_table(table_name=prefix_table_name(table_name))
+        """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):
@@ -66,16 +94,3 @@ class MemoryAcorn(Acorn):
     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"""
-    try:
-        logging.info(f"Fetching from cache: {table_name}")
-        df = acorn.scurry(table_name)
-    except Exception as e:
-        logging.warning(f"Error fetching from cache: {e}")
-        df = pd.DataFrame()
-
-    return df

zombie_squirrel/squirrels.py

@@ -10,8 +10,7 @@ from aind_data_access_api.document_db import MetadataDbClient
 
 from zombie_squirrel.acorns import (
     MemoryAcorn,
-    RedshiftAcorn,
-    rds_get_handle_empty,
+    S3Acorn,
 )
 
 # --- Backend setup ---------------------------------------------------
@@ -20,9 +19,9 @@ API_GATEWAY_HOST = "api.allenneuraldynamics.org"
 
 tree_type = os.getenv("TREE_SPECIES", "memory").lower()
 
-if tree_type == "redshift":  # pragma: no cover
-    logging.info("Using Redshift acorn for caching")
-    ACORN = RedshiftAcorn()
+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()
@@ -66,7 +65,7 @@ def unique_project_names(force_update: bool = False) -> list[str]:
 
     Returns:
         List of unique project names."""
-    df = rds_get_handle_empty(ACORN, NAMES["upn"])
+    df = ACORN.scurry(NAMES["upn"])
 
     if df.empty or force_update:
         # If cache is missing, fetch data
@@ -99,7 +98,7 @@ def unique_subject_ids(force_update: bool = False) -> list[str]:
 
     Returns:
         List of unique subject IDs."""
-    df = rds_get_handle_empty(ACORN, NAMES["usi"])
+    df = ACORN.scurry(NAMES["usi"])
 
     if df.empty or force_update:
         # If cache is missing, fetch data
@@ -134,7 +133,7 @@ def asset_basics(force_update: bool = False) -> pd.DataFrame:
 
     Returns:
         DataFrame with basic asset metadata."""
-    df = rds_get_handle_empty(ACORN, NAMES["basics"])
+    df = ACORN.scurry(NAMES["basics"])
 
     FIELDS = [
         "data_description.modalities",
@@ -145,6 +144,7 @@ def asset_basics(force_update: bool = False) -> pd.DataFrame:
         "acquisition.acquisition_end_time",
         "processing.data_processes.start_date_time",
         "subject.subject_details.genotype",
+        "other_identifiers",
         "location",
     ]
 
@@ -160,6 +160,7 @@ def asset_basics(force_update: bool = False) -> pd.DataFrame:
                 "subject_id",
                 "acquisition_start_time",
                 "acquisition_end_time",
+                "code_ocean",
                 "process_date",
                 "genotype",
                 "location",
@@ -205,12 +206,21 @@ def asset_basics(force_update: bool = False) -> pd.DataFrame:
             modality_abbreviations_str = ", ".join(modality_abbreviations)
 
             # Get the process date, convert to YYYY-MM-DD if present
-            process_datetime = record.get("processing", {}).get("data_processes", [{}])[-1].get("start_date_time", None)
-            if process_datetime:
+            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),
@@ -220,6 +230,7 @@ def asset_basics(force_update: bool = False) -> pd.DataFrame:
                 "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),
@@ -247,7 +258,7 @@ def source_data(force_update: bool = False) -> pd.DataFrame:
 
     Returns:
         DataFrame with _id and source_data columns."""
-    df = rds_get_handle_empty(ACORN, NAMES["d2r"])
+    df = ACORN.scurry(NAMES["d2r"])
 
     if df.empty or force_update:
         logging.info("Updating cache for source data")
@@ -289,7 +300,7 @@ def raw_to_derived(force_update: bool = False) -> pd.DataFrame:
 
     Returns:
         DataFrame with _id and derived_records columns."""
-    df = rds_get_handle_empty(ACORN, NAMES["r2d"])
+    df = ACORN.scurry(NAMES["r2d"])
 
     if df.empty or force_update:
         logging.info("Updating cache for raw to derived mapping")

zombie_squirrel/utils.py

@@ -2,11 +2,22 @@
 
 
 def prefix_table_name(table_name: str) -> str:
-    """Add zombie-squirrel prefix to table names.
+    """Add zombie-squirrel prefix and parquet extension to filenames.
 
     Args:
         table_name: The base table name.
 
     Returns:
-        Table name with 'zs_' prefix."""
-    return "zs_" + table_name
+        Filename with 'zs_' prefix and '.pqt' extension."""
+    return "zs_" + table_name + ".pqt"
+
+
+def get_s3_cache_path(filename: str) -> str:
+    """Get the full S3 path for a cache file.
+
+    Args:
+        filename: The cache filename (e.g., "zs_unique_project_names.pqt").
+
+    Returns:
+        Full S3 path: application-caches/filename"""
+    return f"application-caches/{filename}"

{zombie_squirrel-0.5.0.dist-info → zombie_squirrel-0.6.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: zombie-squirrel
-Version: 0.5.0
+Version: 0.6.1
 Summary: Generated from aind-library-template
 Author: Allen Institute for Neural Dynamics
 License: MIT
@@ -8,7 +8,11 @@ Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: aind-data-access-api[docdb,rds]
+Requires-Dist: duckdb
+Requires-Dist: fastparquet
+Requires-Dist: boto3
+Requires-Dist: pandas
+Requires-Dist: aind-data-access-api[docdb]
 Dynamic: license-file
 
 # zombie-squirrel
@@ -26,9 +30,6 @@ Dynamic: license-file
 
 ```bash
 pip install zombie-squirrel
-
-```bash
-uv sync
 ```
 
 ## Usage
@@ -36,11 +37,10 @@ uv sync
 ### Set backend
 
 ```bash
-export REDSHIFT_SECRETS='/aind/prod/redshift/credentials/readwrite'
-export TREE_SPECIES='REDSHIFT'
+export TREE_SPECIES='s3'
 ```
 
-Options are 'REDSHIFT', 'MEMORY'.
+Options are 's3', 'MEMORY'.
 
 ### Scurry (fetch) data
 

zombie_squirrel-0.6.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+zombie_squirrel/__init__.py,sha256=v9cqU6rYDVeeVPFLDyhe6N28mEP2j7vXb8R98nwc--0,409
+zombie_squirrel/acorns.py,sha256=mpinFacaN9BM6CvRy0M76JMb6n3oVPZLJxn8O4J9Wlw,2945
+zombie_squirrel/squirrels.py,sha256=JWO0ihe7J3P9HbrSA4DHG3-o7EhM93WiwtxC2yVD0L0,12262
+zombie_squirrel/sync.py,sha256=84Ta5beHiPuGBVzp9SCo7G1b4McTUohcUIf_TJbNIV8,518
+zombie_squirrel/utils.py,sha256=kojQpHUKlRJD7WEZDfcpQIZTj9iUrtX5_6F-gWWzJW0,628
+zombie_squirrel-0.6.1.dist-info/licenses/LICENSE,sha256=U0Y7B3gZJHXpjJVLgTQjM8e_c8w4JJpLgGhIdsoFR1Y,1092
+zombie_squirrel-0.6.1.dist-info/METADATA,sha256=vywgegMbrLaSk7VdS_L3cT0PT9YD53nRA3o1mKUEVqc,1893
+zombie_squirrel-0.6.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+zombie_squirrel-0.6.1.dist-info/top_level.txt,sha256=FmM0coe4AangURZLjM4JwwRv2B8H6oINYCoZLKLDCKA,16
+zombie_squirrel-0.6.1.dist-info/RECORD,,

zombie_squirrel-0.5.0.dist-info/RECORD

@@ -1,10 +0,0 @@
-zombie_squirrel/__init__.py,sha256=4zjJN8AP-RKQir_Wgvy0KqVUcV2REOE9kTHwrMyRV5U,409
-zombie_squirrel/acorns.py,sha256=4uBzYtYgW2oD5sOohNQUw4qfjmNjmAIK2RlL1Ge1Udo,2597
-zombie_squirrel/squirrels.py,sha256=BjRF7MhSJMK4BsBhHU5PzAdG6Rb7OloTMXUvrA7pZT8,11934
-zombie_squirrel/sync.py,sha256=84Ta5beHiPuGBVzp9SCo7G1b4McTUohcUIf_TJbNIV8,518
-zombie_squirrel/utils.py,sha256=woPxU4vYMUv-T0XOjV5ieViksU_q7It_n_5Ll4zpocA,289
-zombie_squirrel-0.5.0.dist-info/licenses/LICENSE,sha256=U0Y7B3gZJHXpjJVLgTQjM8e_c8w4JJpLgGhIdsoFR1Y,1092
-zombie_squirrel-0.5.0.dist-info/METADATA,sha256=WDnQXat1d6eC1-RkESW5yZ1lmu8vaHPYEM87SvBuB1g,1902
-zombie_squirrel-0.5.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-zombie_squirrel-0.5.0.dist-info/top_level.txt,sha256=FmM0coe4AangURZLjM4JwwRv2B8H6oINYCoZLKLDCKA,16
-zombie_squirrel-0.5.0.dist-info/RECORD,,