[DEV-12163] - Office Loader

brus_backend_common/helpers/scripts.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import logging
 import requests
@@ -6,6 +7,7 @@
 from enum import Enum
 from typing import Any, Callable, Literal, Protocol
 
+import aiohttp
 import numpy as np
 import pandas as pd
 from requests.exceptions import HTTPError, RequestException, ConnectionError, ReadTimeout
@@ -34,6 +36,7 @@ def get_with_exception_hand(
     resp_type: Literal["json", "xml"] = "json",
     namespaces: dict | None = None,
     validate_response: Callable[[dict], bool] | None = None,
+    params: dict | None = None,
 ) -> dict[Any, Any] | requests.Response | None:
     """Retrieve data from a feed, allow for multiple retries, checks, and timeouts
 
@@ -68,7 +71,7 @@ def get_with_exception_hand(
     while current_retries < len(retry_sleep_times):
         resp = None
         try:
-            resp = requests.get(url_string, timeout=request_timeout)
+            resp = requests.get(url_string, params=params, timeout=request_timeout)
             resp.raise_for_status()
 
             if not decode:
@@ -124,6 +127,111 @@ def get_with_exception_hand(
     return return_val
 
 
+async def async_get_with_exception_hand(
+    session: aiohttp.ClientSession,
+    url_string: str,
+    max_retries: int = 12,
+    decode: bool = True,
+    resp_type: Literal["json", "xml"] = "json",
+    namespaces: dict | None = None,
+    validate_response: Callable[[dict], bool] | None = None,
+    params: dict | None = None,
+) -> dict[Any, Any] | str | None:
+    """Async version of get_with_exception_hand
+
+    Args:
+        session: aiohttp.ClientSession to reuse for connection pooling
+        url_string: string path to the feed we are getting data from
+        max_retries: maximum number of retries to save time
+        decode: whether to decode the response into a dict
+        resp_type: the format of the response (either 'json' or 'xml')
+        namespaces: dict of namespaces to clean up for the xml parsing
+        validate_response: function to determine if the response is valid
+        params: query parameters for the request
+
+    Returns:
+        The response from the url provided as a dict or string
+
+    Raises:
+        Various exceptions if max retries exceeded
+    """
+    if resp_type not in ["json", "xml"]:
+        raise ValueError("resp_type must be 'xml' or 'json'.")
+    if not 0 <= max_retries <= 12:
+        raise ValueError("max_retries must be 0-12.")
+
+    current_retries = 0
+    retry_sleep_times = [0.5, 1, 5, 30, 60, 180, 300, 360, 420, 480, 540, 600]
+    request_timeout = 60
+    return_val = None
+
+    while current_retries <= max_retries and current_retries < len(retry_sleep_times):
+        resp_text = None
+        try:
+            timeout = aiohttp.ClientTimeout(total=request_timeout)
+            async with session.get(url_string, params=params, timeout=timeout) as resp:
+                resp.raise_for_status()
+
+                if not decode:
+                    return_val = await resp.text()
+                else:
+                    if resp_type == "xml":
+                        resp_text = await resp.text()
+                        resp_dict = xmltodict.parse(resp_text, process_namespaces=True, namespaces=namespaces)
+                    else:
+                        resp_dict = await resp.json()
+
+                    resp_err = resp_dict.get("error")
+                    if resp_err:
+                        message = resp_dict.get("message")
+                        raise ValueError(f"Error processing response: {resp_err} {message}")
+
+                    if validate_response is not None and validate_response(resp_dict) is False:
+                        raise ValueError(f"Invalid response provided: {resp_dict}")
+
+                    return_val = resp_dict
+                break
+
+        except (
+            aiohttp.ClientResponseError,
+            aiohttp.ClientConnectionError,
+            aiohttp.ClientError,
+            asyncio.TimeoutError,
+            json.decoder.JSONDecodeError,
+            ValueError,
+        ) as e:
+            if resp_text:
+                logger.exception(f"Response text: {resp_text}")
+            else:
+                logger.exception(e)
+
+            if current_retries < len(retry_sleep_times) and current_retries < max_retries:
+                logger.info(
+                    f"Sleeping {retry_sleep_times[current_retries]}s and then retrying"
+                    f" with a max wait of {request_timeout}s... (attempt {current_retries + 1}/{max_retries})"
+                )
+                await asyncio.sleep(retry_sleep_times[current_retries])
+                current_retries += 1
+                request_timeout += 60
+            else:
+                logger.error(f"Maximum retry attempts exceeded for {url_string}")
+                raise e
+
+        except xmltodict.expat.ExpatError as e:
+            logger.exception(f"XML parsing error: {e}")
+            if current_retries < len(retry_sleep_times) and current_retries < max_retries:
+                await asyncio.sleep(retry_sleep_times[current_retries])
+                current_retries += 1
+            else:
+                raise e
+
+        except Exception as e:
+            logger.exception(f"An unexpected error occurred: {e}")
+            raise e
+
+    return return_val
+
+
 def trim_nested_obj(obj: Any) -> Any:
     """A recursive version to trim all the values in a nested object
 

brus_backend_common/models/reference.py

@@ -272,3 +272,29 @@ class ProgramActivityParkGold(CSVModel):
             SchemaField("park_name", SchemaType.STRING, False),
         ]
     )
+
+
+class OfficeGold(CSVModel):
+    BUCKET_NAME = CONFIG.LAKEHOUSE_REFERENCE_BUCKET
+    DATABASE_NAME = LakeHouseDatabase.GOLD
+    TABLE_NAME = "office"
+    DESCRIPTION = "Federal Hierarchy Office Data"
+    CSV_NAME = "office.csv"
+    PK = "office_id"
+    STRUCTURE = BaseSchema(
+        [
+            SchemaField("created_at", SchemaType.TIMESTAMP, True),
+            SchemaField("updated_at", SchemaType.TIMESTAMP, True),
+            SchemaField("office_id", SchemaType.INTEGER, False),
+            SchemaField("office_code", SchemaType.STRING, False),
+            SchemaField("office_name", SchemaType.STRING, True),
+            SchemaField("sub_tier_code", SchemaType.STRING, False),
+            SchemaField("agency_code", SchemaType.STRING, False),
+            SchemaField("effective_start_date", SchemaType.TIMESTAMP, True),
+            SchemaField("effective_end_date", SchemaType.TIMESTAMP, True),
+            SchemaField("contract_awards_office", SchemaType.BOOLEAN, False),
+            SchemaField("contract_funding_office", SchemaType.BOOLEAN, False),
+            SchemaField("financial_assistance_awards_office", SchemaType.BOOLEAN, False),
+            SchemaField("financial_assistance_funding_office", SchemaType.BOOLEAN, False),
+        ]
+    )