zombie-squirrel 0.1.0__py3-none-any.whl

zombie_squirrel/__init__.py

@@ -0,0 +1,4 @@
+"""Init package"""
+__version__ = "0.1.0"
+
+from zombie_squirrel.squirrels import unique_project_names

zombie_squirrel/acorns.py

@@ -0,0 +1,60 @@
+# forest_cache/acorns.py
+from abc import ABC, abstractmethod
+import pandas as pd
+import os
+
+from aind_data_access_api.rds_tables import Client, RDSCredentials
+
+
+def prefix_table_name(table_name: str) -> str:
+    return "zs_" + table_name
+
+
+class Acorn(ABC):
+    """Base class for a storage backend (the cache)."""
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    @abstractmethod
+    def hide(self, table_name: str, data: pd.DataFrame) -> None:
+        """Store records in the cache."""
+        pass
+
+    @abstractmethod
+    def scurry(self, table_name: str) -> pd.DataFrame:
+        """Fetch records from the cache."""
+        pass
+
+
+class RedshiftAcorn(Acorn):
+    """Stores and retrieves caches using aind-data-access-api
+    Redshift Client"""
+
+    def __init__(self) -> None:
+        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:
+        self.rds_client.overwrite_table_with_df(
+            df=data,
+            table_name=prefix_table_name(table_name),
+        )
+
+    def scurry(self, table_name: str) -> pd.DataFrame:
+        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:
+        super().__init__()
+        self._store: dict[str, pd.DataFrame] = {}
+
+    def hide(self, table_name: str, data: pd.DataFrame) -> None:
+        self._store[table_name] = data
+
+    def scurry(self, table_name: str) -> pd.DataFrame:
+        return self._store.get(table_name, pd.DataFrame())

zombie_squirrel/squirrels.py

@@ -0,0 +1,76 @@
+"""Squirrels: functions to fetch and cache data from MongoDB."""
+import pandas as pd
+from typing import Any, Callable, Optional
+from zombie_squirrel.acorns import RedshiftAcorn, MemoryAcorn
+from aind_data_access_api.document_db import MetadataDbClient
+import os
+
+# --- Backend setup ---------------------------------------------------
+
+API_GATEWAY_HOST = "api.allenneuraldynamics.org"
+
+tree_type = os.getenv("TREE_SPECIES", "memory").lower()
+
+if tree_type == "redshift":
+    ACORN = RedshiftAcorn()
+else:
+    ACORN = MemoryAcorn()
+
+# --- Squirrel registry -----------------------------------------------------
+
+SQUIRREL_REGISTRY: dict[str, Callable[[], Any]] = {}
+
+
+def register_squirrel(name: str):
+    """Decorator for registering new squirrels."""
+    def decorator(func):
+        SQUIRREL_REGISTRY[name] = func
+        return func
+    return decorator
+
+
+# --- Squirrels -----------------------------------------------------
+
+
+@register_squirrel("unique-project-names")
+def unique_project_names(force_update: bool = False) -> list[str]:
+    df = ACORN.scurry("unique-project-names")
+
+    if df.empty or force_update:
+        # If cache is missing, fetch data
+        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("unique-project-names", df)
+
+    return df["project_name"].tolist()
+
+
+@register_squirrel("unique-subject-ids")
+def unique_subject_ids(force_update: bool = False) -> list[str]:
+    df = ACORN.scurry("unique-subject-ids")
+
+    if df.empty or force_update:
+        # If cache is missing, fetch data
+        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("unique-subject-ids", df)
+
+    return df["subject_id"].tolist()

zombie_squirrel/sync.py

@@ -0,0 +1,7 @@
+"""Sync all acorns"""
+from .squirrels import SQUIRREL_REGISTRY
+
+
+def hide_acorns():
+    for squirrel in SQUIRREL_REGISTRY.values():
+        squirrel(force_update=True)

zombie_squirrel-0.1.0.dist-info/METADATA

@@ -0,0 +1,56 @@
+Metadata-Version: 2.4
+Name: zombie-squirrel
+Version: 0.1.0
+Summary: Generated from aind-library-template
+Author: Allen Institute for Neural Dynamics
+License: MIT
+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]
+Dynamic: license-file
+
+# zombie-squirrel
+
+[![License](https://img.shields.io/badge/license-MIT-brightgreen)](LICENSE)
+![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-100%25-brightgreen)
+![Python](https://img.shields.io/badge/python->=3.10-blue?logo=python)
+
+## Installation
+
+```bash
+pip install zombie-squirrel
+
+```bash
+uv sync
+```
+
+## Usage
+
+### Set backend
+
+```bash
+export REDSHIFT_SECRETS='/aind/prod/redshift/credentials/readwrite'
+export TREE_SPECIES='REDSHIFT'
+```
+
+Options are 'REDSHIFT', 'MEMORY'.
+
+### Scurry (fetch) data
+
+```python
+import zombie_squirrel as zs
+
+unique_project_names = zs.scurry_project_names()
+```
+
+### Hide the acorns
+
+```python
+from zombie_squirrel.sync import hide_acorns
+hide_acorns()
+```

zombie_squirrel-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+zombie_squirrel/__init__.py,sha256=H8ylm8tAcxCbApkNWTsAt5kGs1AgPxrxrob-4n2ADpg,100
+zombie_squirrel/acorns.py,sha256=3wiZE5RYzoSgUdnan6lCVlE7Xn5VeNDmDiIXHl0YQfA,1785
+zombie_squirrel/squirrels.py,sha256=PTDw5DeRoLhdDwuKHAZsLnIAQulYc-A5b8021WaM0BI,2345
+zombie_squirrel/sync.py,sha256=jslTVIend5Z-sLJuNXKkhn-nqmKK_P0FAiRuFFYRnto,168
+zombie_squirrel-0.1.0.dist-info/licenses/LICENSE,sha256=U0Y7B3gZJHXpjJVLgTQjM8e_c8w4JJpLgGhIdsoFR1Y,1092
+zombie_squirrel-0.1.0.dist-info/METADATA,sha256=5LYlpZBrXyjcCPvZ3J62_E_SVzNeAQ94nHEDXAslxzk,1382
+zombie_squirrel-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+zombie_squirrel-0.1.0.dist-info/top_level.txt,sha256=FmM0coe4AangURZLjM4JwwRv2B8H6oINYCoZLKLDCKA,16
+zombie_squirrel-0.1.0.dist-info/RECORD,,

zombie_squirrel-0.1.0.dist-info/WHEEL

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

zombie_squirrel-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Allen Institute for Neural Dynamics
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

zombie_squirrel-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+zombie_squirrel