zombie-squirrel 0.5.1__tar.gz → 0.5.3__tar.gz

{zombie_squirrel-0.5.1/src/zombie_squirrel.egg-info → zombie_squirrel-0.5.3}/PKG-INFO

@@ -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.1 → zombie_squirrel-0.5.3}/README.md

@@ -4,7 +4,7 @@
 ![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)">
@@ -23,11 +23,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.1 → zombie_squirrel-0.5.3}/pyproject.toml

@@ -17,7 +17,11 @@ readme = "README.md"
 dynamic = ["version"]
 
 dependencies = [
-    'aind-data-access-api[docdb,rds]',
+    'duckdb',
+    'fastparquet',
+    'boto3',
+    'pandas',
+    'aind-data-access-api[docdb]',
 ]
 
 [dependency-groups]

{zombie_squirrel-0.5.1 → zombie_squirrel-0.5.3}/src/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-0.5.3/src/zombie_squirrel/acorns.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 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  # pragma: no cover
+
+    @abstractmethod
+    def scurry(self, table_name: str) -> pd.DataFrame:
+        """Fetch records from the cache."""
+        pass  # pragma: no cover
+
+
+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 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 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())

{zombie_squirrel-0.5.1 → zombie_squirrel-0.5.3}/src/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-0.5.3/src/zombie_squirrel/utils.py

@@ -0,0 +1,23 @@
+"""Utility functions for zombie-squirrel package."""
+
+
+def prefix_table_name(table_name: str) -> str:
+    """Add zombie-squirrel prefix and parquet extension to filenames.
+
+    Args:
+        table_name: The base table name.
+
+    Returns:
+        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 → zombie_squirrel-0.5.3/src/zombie_squirrel.egg-info}/PKG-INFO

@@ -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/src/zombie_squirrel.egg-info/requires.txt

@@ -0,0 +1,5 @@
+duckdb
+fastparquet
+boto3
+pandas
+aind-data-access-api[docdb]

zombie_squirrel-0.5.3/tests/test_acorns.py

@@ -0,0 +1,184 @@
+"""Unit tests for zombie_squirrel.acorns module.
+
+Tests for abstract base class, memory backend, and S3 backend
+for caching functionality."""
+
+import unittest
+from unittest.mock import MagicMock, Mock, patch
+
+import pandas as pd
+
+from zombie_squirrel.acorns import (
+    Acorn,
+    MemoryAcorn,
+    S3Acorn,
+)
+
+
+class TestAcornAbstractClass(unittest.TestCase):
+    """Tests for Acorn abstract base class."""
+
+    def test_acorn_cannot_be_instantiated(self):
+        """Test that Acorn abstract class cannot be instantiated."""
+        with self.assertRaises(TypeError):
+            Acorn()
+
+    def test_acorn_subclass_must_implement_hide(self):
+        """Test that subclasses must implement hide method."""
+
+        class IncompleteAcorn(Acorn):
+            """Incomplete Acorn subclass missing hide method."""
+
+            def scurry(self, table_name: str) -> pd.DataFrame:  # pragma: no cover
+                """Fetch records from the cache."""
+                return pd.DataFrame()
+
+        with self.assertRaises(TypeError):
+            IncompleteAcorn()
+
+    def test_acorn_subclass_must_implement_scurry(self):
+        """Test that subclasses must implement scurry method."""
+
+        class IncompleteAcorn(Acorn):
+            """Incomplete Acorn subclass missing scurry method."""
+
+            def hide(self, table_name: str, data: pd.DataFrame) -> None:  # pragma: no cover
+                """Store records in the cache."""
+                pass
+
+        with self.assertRaises(TypeError):
+            IncompleteAcorn()
+
+
+class TestMemoryAcorn(unittest.TestCase):
+    """Tests for MemoryAcorn implementation."""
+
+    def setUp(self):
+        """Initialize a fresh MemoryAcorn for each test."""
+        self.acorn = MemoryAcorn()
+
+    def test_hide_and_scurry_basic(self):
+        """Test basic hide and scurry operations."""
+        df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
+        self.acorn.hide("test_table", df)
+
+        retrieved = self.acorn.scurry("test_table")
+        pd.testing.assert_frame_equal(df, retrieved)
+
+    def test_scurry_empty_table(self):
+        """Test scurrying a table that doesn't exist returns empty DataFrame."""
+        result = self.acorn.scurry("nonexistent_table")
+        self.assertTrue(result.empty)
+        self.assertIsInstance(result, pd.DataFrame)
+
+    def test_hide_overwrites_existing(self):
+        """Test that hiding data overwrites existing data."""
+        df1 = pd.DataFrame({"col1": [1, 2, 3]})
+        df2 = pd.DataFrame({"col1": [4, 5, 6]})
+
+        self.acorn.hide("table", df1)
+        self.acorn.hide("table", df2)
+
+        retrieved = self.acorn.scurry("table")
+        pd.testing.assert_frame_equal(df2, retrieved)
+
+    def test_multiple_tables(self):
+        """Test managing multiple tables."""
+        df1 = pd.DataFrame({"col1": [1, 2]})
+        df2 = pd.DataFrame({"col2": ["a", "b"]})
+
+        self.acorn.hide("table1", df1)
+        self.acorn.hide("table2", df2)
+
+        retrieved1 = self.acorn.scurry("table1")
+        retrieved2 = self.acorn.scurry("table2")
+
+        pd.testing.assert_frame_equal(df1, retrieved1)
+        pd.testing.assert_frame_equal(df2, retrieved2)
+
+    def test_hide_empty_dataframe(self):
+        """Test hiding an empty DataFrame."""
+        df = pd.DataFrame()
+        self.acorn.hide("empty_table", df)
+
+        retrieved = self.acorn.scurry("empty_table")
+        pd.testing.assert_frame_equal(df, retrieved)
+
+
+class TestS3Acorn(unittest.TestCase):
+    """Tests for S3Acorn implementation with mocking."""
+
+    @patch("zombie_squirrel.acorns.boto3.client")
+    def test_s3_acorn_initialization(self, mock_boto3_client):
+        """Test S3Acorn initialization."""
+        mock_s3_client = MagicMock()
+        mock_boto3_client.return_value = mock_s3_client
+
+        acorn = S3Acorn()
+
+        self.assertEqual(acorn.bucket, "aind-scratch-data")
+        self.assertEqual(acorn.s3_client, mock_s3_client)
+        mock_boto3_client.assert_called_once_with("s3")
+
+    @patch("zombie_squirrel.acorns.boto3.client")
+    def test_s3_hide(self, mock_boto3_client):
+        """Test S3Acorn.hide method writes to S3."""
+        mock_s3_client = MagicMock()
+        mock_boto3_client.return_value = mock_s3_client
+
+        acorn = S3Acorn()
+        df = pd.DataFrame({"col1": [1, 2, 3]})
+
+        acorn.hide("test_table", df)
+
+        mock_s3_client.put_object.assert_called_once()
+        call_kwargs = mock_s3_client.put_object.call_args[1]
+        self.assertEqual(call_kwargs["Bucket"], "aind-scratch-data")
+        self.assertEqual(
+            call_kwargs["Key"], "application-caches/zs_test_table.pqt"
+        )
+        self.assertIsInstance(call_kwargs["Body"], bytes)
+
+    @patch("zombie_squirrel.acorns.duckdb.query")
+    @patch("zombie_squirrel.acorns.boto3.client")
+    def test_s3_scurry(self, mock_boto3_client, mock_duckdb_query):
+        """Test S3Acorn.scurry method reads from S3 using DuckDB."""
+        mock_s3_client = MagicMock()
+        mock_boto3_client.return_value = mock_s3_client
+
+        expected_df = pd.DataFrame({"col1": [1, 2, 3]})
+        mock_result = MagicMock()
+        mock_result.to_df.return_value = expected_df
+        mock_duckdb_query.return_value = mock_result
+
+        acorn = S3Acorn()
+        result = acorn.scurry("test_table")
+
+        # Verify DuckDB was called with correct S3 path
+        mock_duckdb_query.assert_called_once()
+        query_call = mock_duckdb_query.call_args[0][0]
+        self.assertIn(
+            "s3://aind-scratch-data/application-caches/zs_test_table.pqt",
+            query_call,
+        )
+        pd.testing.assert_frame_equal(result, expected_df)
+
+    @patch("zombie_squirrel.acorns.duckdb.query")
+    @patch("zombie_squirrel.acorns.boto3.client")
+    def test_s3_scurry_handles_error(
+        self, mock_boto3_client, mock_duckdb_query
+    ):
+        """Test S3Acorn.scurry returns empty DataFrame on error."""
+        mock_s3_client = MagicMock()
+        mock_boto3_client.return_value = mock_s3_client
+        mock_duckdb_query.side_effect = Exception("S3 access error")
+
+        acorn = S3Acorn()
+        result = acorn.scurry("nonexistent_table")
+
+        self.assertTrue(result.empty)
+        self.assertIsInstance(result, pd.DataFrame)
+
+
+if __name__ == "__main__":
+    unittest.main()

{zombie_squirrel-0.5.1 → zombie_squirrel-0.5.3}/tests/test_squirrels.py

@@ -204,7 +204,47 @@ class TestAssetBasics(unittest.TestCase):
 
     @patch("zombie_squirrel.squirrels.ACORN", new_callable=MemoryAcorn)
     @patch("zombie_squirrel.squirrels.MetadataDbClient")
-    def test_asset_basics_incremental_update(self, mock_client_class, mock_acorn):
+    def test_asset_basics_with_data_processes(
+        self, mock_client_class, mock_acorn
+    ):
+        """Test asset_basics includes process_date from data_processes."""
+        mock_client_instance = MagicMock()
+        mock_client_class.return_value = mock_client_instance
+
+        mock_client_instance.retrieve_docdb_records.return_value = [
+            {
+                "_id": "id1",
+                "_last_modified": "2023-01-01",
+                "data_description": {
+                    "modalities": [{"abbreviation": "img"}],
+                    "project_name": "proj1",
+                    "data_level": "raw",
+                },
+                "subject": {"subject_id": "sub001"},
+                "acquisition": {
+                    "acquisition_start_time": "2023-01-01T10:00:00",
+                    "acquisition_end_time": "2023-01-01T11:00:00",
+                },
+                "processing": {
+                    "data_processes": [
+                        {"start_date_time": "2023-01-15T14:30:00"},
+                        {"start_date_time": "2023-01-20T09:15:00"},
+                    ]
+                },
+            }
+        ]
+
+        result = asset_basics()
+
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result.iloc[0]["_id"], "id1")
+        self.assertEqual(result.iloc[0]["process_date"], "2023-01-20")
+
+    @patch("zombie_squirrel.squirrels.ACORN", new_callable=MemoryAcorn)
+    @patch("zombie_squirrel.squirrels.MetadataDbClient")
+    def test_asset_basics_incremental_update(
+        self, mock_client_class, mock_acorn
+    ):
         """Test incremental cache update with partial data refresh."""
         mock_client_instance = MagicMock()
         mock_client_class.return_value = mock_client_instance

zombie_squirrel-0.5.3/tests/test_utils.py

@@ -0,0 +1,54 @@
+"""Unit tests for zombie_squirrel.utils module.
+
+Tests for utility functions."""
+
+import unittest
+
+from zombie_squirrel.utils import get_s3_cache_path, prefix_table_name
+
+
+class TestPrefixTableName(unittest.TestCase):
+    """Tests for the prefix_table_name function."""
+
+    def test_prefix_table_name_basic(self):
+        """Test that prefix_table_name adds 'zs_' prefix and '.pqt' ext."""
+        result = prefix_table_name("my_table")
+        self.assertEqual(result, "zs_my_table.pqt")
+
+    def test_prefix_table_name_empty_string(self):
+        """Test with empty string."""
+        result = prefix_table_name("")
+        self.assertEqual(result, "zs_.pqt")
+
+    def test_prefix_table_name_single_char(self):
+        """Test with single character."""
+        result = prefix_table_name("a")
+        self.assertEqual(result, "zs_a.pqt")
+
+    def test_prefix_table_name_with_underscores(self):
+        """Test with table name containing underscores."""
+        result = prefix_table_name("my_long_table_name")
+        self.assertEqual(result, "zs_my_long_table_name.pqt")
+
+    def test_prefix_table_name_with_numbers(self):
+        """Test with table name containing numbers."""
+        result = prefix_table_name("table123")
+        self.assertEqual(result, "zs_table123.pqt")
+
+
+class TestGetS3CachePath(unittest.TestCase):
+    """Tests for the get_s3_cache_path function."""
+
+    def test_get_s3_cache_path_basic(self):
+        """Test that get_s3_cache_path constructs correct S3 path."""
+        result = get_s3_cache_path("zs_test.pqt")
+        self.assertEqual(result, "application-caches/zs_test.pqt")
+
+    def test_get_s3_cache_path_various_names(self):
+        """Test with various filenames."""
+        result = get_s3_cache_path("zs_my_data.pqt")
+        self.assertEqual(result, "application-caches/zs_my_data.pqt")
+
+
+if __name__ == "__main__":
+    unittest.main()

zombie_squirrel-0.5.1/src/zombie_squirrel/acorns.py

@@ -1,81 +0,0 @@
-"""Storage backend interfaces for caching data."""
-
-import logging
-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
-
-
-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  # pragma: no cover
-
-    @abstractmethod
-    def scurry(self, table_name: str) -> pd.DataFrame:
-        """Fetch records from the cache."""
-        pass  # pragma: no cover
-
-
-class RedshiftAcorn(Acorn):
-    """Stores and retrieves caches using aind-data-access-api
-    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"""
-    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-0.5.1/src/zombie_squirrel/utils.py

@@ -1,12 +0,0 @@
-"""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.5.1/src/zombie_squirrel.egg-info/requires.txt

@@ -1 +0,0 @@
-aind-data-access-api[docdb,rds]

zombie_squirrel-0.5.1/tests/test_acorns.py

@@ -1,217 +0,0 @@
-"""Unit tests for zombie_squirrel.acorns module.
-
-Tests for abstract base class, memory backend, and Redshift backend
-for caching functionality."""
-
-import os
-import unittest
-from unittest.mock import MagicMock, Mock, patch
-
-import pandas as pd
-
-from zombie_squirrel.acorns import (
-    Acorn,
-    MemoryAcorn,
-    RedshiftAcorn,
-    rds_get_handle_empty,
-)
-
-
-class TestAcornAbstractClass(unittest.TestCase):
-    """Tests for Acorn abstract base class."""
-
-    def test_acorn_cannot_be_instantiated(self):
-        """Test that Acorn abstract class cannot be instantiated."""
-        with self.assertRaises(TypeError):
-            Acorn()
-
-    def test_acorn_subclass_must_implement_hide(self):
-        """Test that subclasses must implement hide method."""
-
-        class IncompleteAcorn(Acorn):
-            """Incomplete Acorn subclass missing hide method."""
-
-            def scurry(self, table_name: str) -> pd.DataFrame:  # pragma: no cover
-                """Fetch records from the cache."""
-                return pd.DataFrame()
-
-        with self.assertRaises(TypeError):
-            IncompleteAcorn()
-
-    def test_acorn_subclass_must_implement_scurry(self):
-        """Test that subclasses must implement scurry method."""
-
-        class IncompleteAcorn(Acorn):
-            """Incomplete Acorn subclass missing scurry method."""
-
-            def hide(self, table_name: str, data: pd.DataFrame) -> None:  # pragma: no cover
-                """Store records in the cache."""
-                pass
-
-        with self.assertRaises(TypeError):
-            IncompleteAcorn()
-
-
-class TestMemoryAcorn(unittest.TestCase):
-    """Tests for MemoryAcorn implementation."""
-
-    def setUp(self):
-        """Initialize a fresh MemoryAcorn for each test."""
-        self.acorn = MemoryAcorn()
-
-    def test_hide_and_scurry_basic(self):
-        """Test basic hide and scurry operations."""
-        df = pd.DataFrame({"col1": [1, 2, 3], "col2": ["a", "b", "c"]})
-        self.acorn.hide("test_table", df)
-
-        retrieved = self.acorn.scurry("test_table")
-        pd.testing.assert_frame_equal(df, retrieved)
-
-    def test_scurry_empty_table(self):
-        """Test scurrying a table that doesn't exist returns empty DataFrame."""
-        result = self.acorn.scurry("nonexistent_table")
-        self.assertTrue(result.empty)
-        self.assertIsInstance(result, pd.DataFrame)
-
-    def test_hide_overwrites_existing(self):
-        """Test that hiding data overwrites existing data."""
-        df1 = pd.DataFrame({"col1": [1, 2, 3]})
-        df2 = pd.DataFrame({"col1": [4, 5, 6]})
-
-        self.acorn.hide("table", df1)
-        self.acorn.hide("table", df2)
-
-        retrieved = self.acorn.scurry("table")
-        pd.testing.assert_frame_equal(df2, retrieved)
-
-    def test_multiple_tables(self):
-        """Test managing multiple tables."""
-        df1 = pd.DataFrame({"col1": [1, 2]})
-        df2 = pd.DataFrame({"col2": ["a", "b"]})
-
-        self.acorn.hide("table1", df1)
-        self.acorn.hide("table2", df2)
-
-        retrieved1 = self.acorn.scurry("table1")
-        retrieved2 = self.acorn.scurry("table2")
-
-        pd.testing.assert_frame_equal(df1, retrieved1)
-        pd.testing.assert_frame_equal(df2, retrieved2)
-
-    def test_hide_empty_dataframe(self):
-        """Test hiding an empty DataFrame."""
-        df = pd.DataFrame()
-        self.acorn.hide("empty_table", df)
-
-        retrieved = self.acorn.scurry("empty_table")
-        pd.testing.assert_frame_equal(df, retrieved)
-
-
-class TestRedshiftAcorn(unittest.TestCase):
-    """Tests for RedshiftAcorn implementation with mocking."""
-
-    @patch("zombie_squirrel.acorns.RDSCredentials")
-    @patch("zombie_squirrel.acorns.Client")
-    def test_redshift_acorn_initialization(self, mock_client_class, mock_credentials_class):
-        """Test RedshiftAcorn initialization."""
-        mock_client_instance = MagicMock()
-        mock_client_class.return_value = mock_client_instance
-        mock_credentials_instance = MagicMock()
-        mock_credentials_class.return_value = mock_credentials_instance
-
-        acorn = RedshiftAcorn()
-
-        self.assertEqual(acorn.rds_client, mock_client_instance)
-        mock_client_class.assert_called_once()
-
-    @patch("zombie_squirrel.acorns.RDSCredentials")
-    @patch("zombie_squirrel.acorns.Client")
-    def test_redshift_hide(self, mock_client_class, mock_credentials_class):
-        """Test RedshiftAcorn.hide method."""
-        mock_client_instance = MagicMock()
-        mock_client_class.return_value = mock_client_instance
-        mock_credentials_instance = MagicMock()
-        mock_credentials_class.return_value = mock_credentials_instance
-
-        acorn = RedshiftAcorn()
-        df = pd.DataFrame({"col1": [1, 2, 3]})
-
-        acorn.hide("test_table", df)
-
-        mock_client_instance.overwrite_table_with_df.assert_called_once()
-        call_args = mock_client_instance.overwrite_table_with_df.call_args
-        pd.testing.assert_frame_equal(call_args[1]["df"], df)
-        self.assertEqual(call_args[1]["table_name"], "zs_test_table")
-
-    @patch("zombie_squirrel.acorns.RDSCredentials")
-    @patch("zombie_squirrel.acorns.Client")
-    def test_redshift_scurry(self, mock_client_class, mock_credentials_class):
-        """Test RedshiftAcorn.scurry method."""
-        mock_client_instance = MagicMock()
-        expected_df = pd.DataFrame({"col1": [1, 2, 3]})
-        mock_client_instance.read_table.return_value = expected_df
-        mock_client_class.return_value = mock_client_instance
-        mock_credentials_instance = MagicMock()
-        mock_credentials_class.return_value = mock_credentials_instance
-
-        acorn = RedshiftAcorn()
-        result = acorn.scurry("test_table")
-
-        mock_client_instance.read_table.assert_called_once_with(table_name="zs_test_table")
-        pd.testing.assert_frame_equal(result, expected_df)
-
-    @patch.dict("os.environ", {}, clear=False)
-    @patch("zombie_squirrel.acorns.RDSCredentials")
-    @patch("zombie_squirrel.acorns.Client")
-    def test_redshift_default_secrets_path(self, mock_client_class, mock_credentials_class):
-        """Test RedshiftAcorn uses default secrets path."""
-        if "REDSHIFT_SECRETS" in os.environ:  # pragma: no cover
-            del os.environ["REDSHIFT_SECRETS"]  # pragma: no cover
-
-        mock_client_instance = MagicMock()
-        mock_client_class.return_value = mock_client_instance
-        mock_credentials_instance = MagicMock()
-        mock_credentials_class.return_value = mock_credentials_instance
-
-        RedshiftAcorn()
-
-        mock_client_class.assert_called_once()
-        call_args = mock_client_class.call_args
-        self.assertIsNotNone(call_args)
-
-
-class TestRdsGetHandleEmpty(unittest.TestCase):
-    """Tests for rds_get_handle_empty helper function."""
-
-    def test_rds_get_handle_empty_success(self):
-        """Test successful retrieval from acorn."""
-        acorn = MemoryAcorn()
-        df = pd.DataFrame({"col1": [1, 2, 3]})
-        acorn.hide("test_table", df)
-
-        result = rds_get_handle_empty(acorn, "test_table")
-
-        pd.testing.assert_frame_equal(result, df)
-
-    def test_rds_get_handle_empty_missing_table(self):
-        """Test returns empty DataFrame when table is missing."""
-        acorn = MemoryAcorn()
-
-        result = rds_get_handle_empty(acorn, "nonexistent_table")
-
-        self.assertTrue(result.empty)
-        self.assertIsInstance(result, pd.DataFrame)
-
-    def test_rds_get_handle_empty_exception(self):
-        """Test returns empty DataFrame when acorn raises exception."""
-        acorn = Mock(spec=["scurry"])
-        acorn.scurry.side_effect = Exception("Connection error")
-
-        result = rds_get_handle_empty(acorn, "test_table")
-
-        self.assertTrue(result.empty)
-        self.assertIsInstance(result, pd.DataFrame)
-
-
-if __name__ == "__main__":
-    unittest.main()

zombie_squirrel-0.5.1/tests/test_utils.py

@@ -1,40 +0,0 @@
-"""Unit tests for zombie_squirrel.utils module.
-
-Tests for utility functions."""
-
-import unittest
-
-from zombie_squirrel.utils import prefix_table_name
-
-
-class TestPrefixTableName(unittest.TestCase):
-    """Tests for the prefix_table_name function."""
-
-    def test_prefix_table_name_basic(self):
-        """Test that prefix_table_name correctly adds 'zs_' prefix."""
-        result = prefix_table_name("my_table")
-        self.assertEqual(result, "zs_my_table")
-
-    def test_prefix_table_name_empty_string(self):
-        """Test with empty string."""
-        result = prefix_table_name("")
-        self.assertEqual(result, "zs_")
-
-    def test_prefix_table_name_single_char(self):
-        """Test with single character."""
-        result = prefix_table_name("a")
-        self.assertEqual(result, "zs_a")
-
-    def test_prefix_table_name_with_underscores(self):
-        """Test with table name containing underscores."""
-        result = prefix_table_name("my_long_table_name")
-        self.assertEqual(result, "zs_my_long_table_name")
-
-    def test_prefix_table_name_with_numbers(self):
-        """Test with table name containing numbers."""
-        result = prefix_table_name("table123")
-        self.assertEqual(result, "zs_table123")
-
-
-if __name__ == "__main__":
-    unittest.main()