zombie-squirrel 0.5.1__py3-none-any.whl → 0.5.3__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.1"
+__version__ = "0.5.3"
 
 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",
@@ -249,7 +248,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")
@@ -291,7 +290,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")
@@ -319,8 +318,6 @@ def raw_to_derived(force_update: bool = False) -> pd.DataFrame:
         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", [])
-            if not source_data_list:
-                continue
             derived_id = derived_record["_id"]
             # Add this derived record to each raw record it depends on
             for source_id in source_data_list:

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.1.dist-info → zombie_squirrel-0.5.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: zombie-squirrel
-Version: 0.5.1
+Version: 0.5.3
 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
@@ -17,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)">
@@ -36,11 +40,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.5.3.dist-info/RECORD

@@ -0,0 +1,10 @@
+zombie_squirrel/__init__.py,sha256=Kc5hE79CUDqA0D2_ISxTgPYz9QyK1PaMwEvUbYh2wEE,409
+zombie_squirrel/acorns.py,sha256=mpinFacaN9BM6CvRy0M76JMb6n3oVPZLJxn8O4J9Wlw,2945
+zombie_squirrel/squirrels.py,sha256=PwLhJdpNZZCWIdqb3-MA3VRzCMf1tWwKcCxpTMbIjn0,11901
+zombie_squirrel/sync.py,sha256=84Ta5beHiPuGBVzp9SCo7G1b4McTUohcUIf_TJbNIV8,518
+zombie_squirrel/utils.py,sha256=kojQpHUKlRJD7WEZDfcpQIZTj9iUrtX5_6F-gWWzJW0,628
+zombie_squirrel-0.5.3.dist-info/licenses/LICENSE,sha256=U0Y7B3gZJHXpjJVLgTQjM8e_c8w4JJpLgGhIdsoFR1Y,1092
+zombie_squirrel-0.5.3.dist-info/METADATA,sha256=T3oVLz6gY4adCdyUdgJT0JA17ziOCU17AuCsrKHSVNw,1911
+zombie_squirrel-0.5.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+zombie_squirrel-0.5.3.dist-info/top_level.txt,sha256=FmM0coe4AangURZLjM4JwwRv2B8H6oINYCoZLKLDCKA,16
+zombie_squirrel-0.5.3.dist-info/RECORD,,

zombie_squirrel-0.5.1.dist-info/RECORD

@@ -1,10 +0,0 @@
-zombie_squirrel/__init__.py,sha256=ihoDdnYVX9vzrIOTRGU1EFj_vGu_6OgJHeOC_6bVyLQ,409
-zombie_squirrel/acorns.py,sha256=4uBzYtYgW2oD5sOohNQUw4qfjmNjmAIK2RlL1Ge1Udo,2597
-zombie_squirrel/squirrels.py,sha256=xFGckcm4hcmFVxQneqeloSZjoY6Qc2wvHT9Wws9L-Rk,12088
-zombie_squirrel/sync.py,sha256=84Ta5beHiPuGBVzp9SCo7G1b4McTUohcUIf_TJbNIV8,518
-zombie_squirrel/utils.py,sha256=woPxU4vYMUv-T0XOjV5ieViksU_q7It_n_5Ll4zpocA,289
-zombie_squirrel-0.5.1.dist-info/licenses/LICENSE,sha256=U0Y7B3gZJHXpjJVLgTQjM8e_c8w4JJpLgGhIdsoFR1Y,1092
-zombie_squirrel-0.5.1.dist-info/METADATA,sha256=CgLCs4WxISN0rMMXSCw0H-E90W8V87bXErZbRbZ26XU,1902
-zombie_squirrel-0.5.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-zombie_squirrel-0.5.1.dist-info/top_level.txt,sha256=FmM0coe4AangURZLjM4JwwRv2B8H6oINYCoZLKLDCKA,16
-zombie_squirrel-0.5.1.dist-info/RECORD,,