Fix 14 api retry (#15)

* fix: adds retry logic to api calls

* chore: removes commented out code

* chore: runs linters

Author: jtyoung84

.flake8

@@ -5,5 +5,6 @@ exclude =
     build,
     .venv,
     venv, 
-    .conda
+    .conda,
+    env
 max-complexity = 10

pyproject.toml

@@ -23,7 +23,8 @@ dependencies = [
     "pandas",
     "awswrangler",
     "jinja2",
-    "requests"
+    "requests",
+    "tenacity"
 ]
 
 [dependency-groups]
@@ -84,5 +85,5 @@ profile = "black"
 skip = [".conda"]
 
 [tool.interrogate]
-exclude = ["setup.py", "docs", "build", ".conda"]
+exclude = ["setup.py", "docs", "build", ".conda", "env"]
 fail-under = 100

src/aind_vast_utils/compile_metrics_job.py

@@ -12,6 +12,12 @@
 from aind_settings_utils.aws import SecretsManagerBaseSettings
 from pydantic import Field, SecretStr
 from pydantic_settings import SettingsConfigDict
+from tenacity import (
+    before_sleep_log,
+    retry,
+    stop_after_attempt,
+    wait_exponential_jitter,
+)
 from vastpy import VASTClient
 
 from aind_vast_utils.models import (
@@ -21,7 +27,8 @@
     QuotaTableRow,
 )
 
-logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
 
 
 class JobSettings(
@@ -72,6 +79,12 @@ def __init__(self, job_settings: JobSettings):
             address=job_settings.address,
         )
 
+    @retry(
+        wait=wait_exponential_jitter(initial=1, max=60, exp_base=2, jitter=1),
+        stop=stop_after_attempt(5),
+        reraise=True,
+        before_sleep=before_sleep_log(logger, logging.WARNING),  # type: ignore
+    )
     def _get_capacity(self, path: str, sort_key: str = "logical") -> Capacity:
         """Get capacity info for VAST cluster."""
 
@@ -81,6 +94,12 @@ def _get_capacity(self, path: str, sort_key: str = "logical") -> Capacity:
         )
         return Capacity.model_validate(response)
 
+    @retry(
+        wait=wait_exponential_jitter(initial=1, max=60, exp_base=2, jitter=1),
+        stop=stop_after_attempt(5),
+        reraise=True,
+        before_sleep=before_sleep_log(logger, logging.WARNING),  # type: ignore
+    )
     def _get_quota(self, path: str) -> Quota:
         """Get quota info for VAST cluster."""
 

src/aind_vast_utils/send_notification_job.py

@@ -15,7 +15,8 @@
 from pydantic import Field
 from pydantic_settings import SettingsConfigDict
 
-logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
 
 
 class JobSettings(
@@ -61,9 +62,6 @@ def _get_table(self, table_name: str) -> pd.DataFrame:
             path = f"{self.job_settings.tables_location}/{table_name}"
             report_date = self.job_settings.report_date
             report_year = report_date.year
-            # my_filter = lambda x: x["report_year"] == str(report_year) and x[
-            #     "report_date"
-            # ] == str(report_date)
             df = wr.s3.read_parquet(
                 path,
                 dataset=True,
@@ -222,8 +220,8 @@ def run_job(self):
             )
             self.send_notification(html_str)
         else:
-            logging.info("All quotas good.")
-            logging.info(quota_df.to_string(index=False))
+            logger.info("All quotas good.")
+            logger.info(quota_df.to_string(index=False))
 
 
 if __name__ == "__main__":