from typing import Literal
from config import Config
from airflow.hooks.base import BaseHook
import re, time, random, json
[docs]
class BaseUtil:
"""
The BaseUtil contains central airflow utilities. These can be used in any dag, operator, hook, or util.
Many of these functions are wrappers around external libraries, such as requests, or centralized entrypoints into wrangler and pandas.
"""
[docs]
@staticmethod
def get_default_states(fips: bool = True) -> list[str]:
"""
Returns a list of default states, either fips codes or capitalized abbreviations.
:param fips: whether or not to return fips codes
:return: the list of states
"""
return list(Config.STATE_FIPS.values() if fips else Config.STATE_FIPS.keys())[
:-1
]
[docs]
@staticmethod
def get_abbrev_from_fips(fips: str) -> str:
"""
Returns a list of default states, either fips codes or capitalized abbreviations.
:param fips: the fips code that we want to convert to the abbreviation
:return: abbreviation
"""
return {v: k for k, v in Config.STATE_FIPS.items()}[fips]
[docs]
@staticmethod
def post_request(
url: str,
headers: dict = None,
data: dict = None,
params: dict = None,
verify: bool = None,
log: bool = False,
timeout: int = 60 * 60 * Config.REQUEST_TIMEOUT_HOURS,
retry_count: int = 1,
retry_sleep: int = 5,
form_data=None,
):
"""
Make a post request to an endpoint
:param url: the url to call
:type url: str
:param headers: the headers to pass to the request
:type headers: dict, optional
:param data: the data to serialize as JSON in request body
:type data: dict, optional
:param params: the params to pass to the request
:type params: dict, optional
:param verify: whether or not to verify the request
:type verify: bool, optional
:param log: whether or not to log the request
:type log: bool, optional
:param timeout: the timeout for the request in seconds
:type timeout: int, optional, defaults to 60 * 60 * Config.REQUEST_TIMEOUT_HOURS
:param retry_count: the number of times to retry the request
:type retry_count: int, optional
:param retry_sleep: the time
:type retry_sleep: int, optional
:param form_data: the form data to pass to the request
:type form_data: dict, optional
:return: The response.
:rtype: requests.Response
"""
import requests
import time
if log:
print(f"Calling {url}")
print(f"Headers: {headers}")
print(f"Data: {data}")
for attempt in range(retry_count):
if log:
print(f"Attempt: {attempt}")
try:
if form_data:
response = requests.post(
url,
data=form_data,
headers=headers,
verify=verify,
timeout=timeout,
params=params,
)
else:
response = requests.post(
url,
# serializes the data as JSON in the request body
json=data,
headers=headers,
verify=verify,
timeout=timeout,
params=params,
)
if response.status_code == 403:
# Enhanced error message with more details
error_detail = {
"url": url,
"status_code": response.status_code,
"reason": response.reason,
"body": response.text,
"headers": dict(response.headers),
}
raise Exception(
f"Access denied (403). Details: {json.dumps(error_detail, indent=2)}"
)
elif response.status_code == 429:
time.sleep(retry_sleep)
raise Exception(f"Rate limit hit. Waiting.... 429, {response}")
elif response.status_code == 200:
return response
else:
raise Exception(f"Error: {response.status_code} - {response.text}")
except Exception as e:
print(f"Calling {url}")
print(f"Headers: {headers}")
print(f"Data: {data}")
print(f"Params: {params}")
print(f"Error: {e}")
continue
raise Exception("BaseUtil.post_request failed")
[docs]
@staticmethod
def get_request(
url: str,
headers: dict = None,
params: dict = None,
proxies: dict = None,
verify: bool = None,
log: bool = True,
timeout: int = 60 * 60 * Config.REQUEST_TIMEOUT_HOURS,
retry_count: int = 1,
stream: Literal[True, False] = False,
):
"""Layer over requests.get to centralize access to external library
:param url: the url to call
:param headers: the headers to pass to the request
:param params: the params to pass to the request
:param proxies: the proxies to pass to the request
:param verify: whether or not to verify the request
:param log: whether or not to log the request
:param timeout: the timeout for the request
:param retry_count: the number of times to retry the request
:return: The response.
:rtype: requests.Response
"""
from requests import get
if "fcc.gov" in url and not headers:
headers = Config.FCC_HEADERS
if log:
print(f"Calling {url}")
print(f"Headers: {headers}")
print(f"Params: {params}")
print(f"Proxy: {proxies}")
for attempt in range(retry_count):
if log:
print(f"Attempt: {attempt}")
try:
return get(
url,
headers=headers,
params=params,
proxies=proxies,
verify=verify,
timeout=timeout,
stream=stream,
)
except Exception as e:
continue
raise Exception("BaseUtil.get_request failed")
[docs]
@staticmethod
def external_get_request(
url: str,
headers: bool = None,
ultra_premium: bool = False,
api_conn: str = Config.SCRAPERAPI_CONN,
):
"""
Make get request to external scraper via scraperapi
:param url: the url to call
:param headers: the headers to pass to the request
:param ultra_premium: whether or not to use ultra premium proxy
"""
scraper_key = BaseHook.get_connection(api_conn).password
scraper_param = {
"api_key": scraper_key,
"timeout": Config.SCRAPER_TIMEOUT,
"url": url,
"keep_headers": headers is not None,
"headers": headers or {},
"ultra_premium": ultra_premium,
}
return BaseUtil.get_request(Config.SCRAPER_URL, params=scraper_param)
[docs]
@staticmethod
def proxy_get_request(
url: str,
headers: dict = None,
premium: bool = False,
stealth: bool = False,
api_conn: str = Config.SCRAPERBEE_CONN,
retry_count: int = 1,
):
"""
Make get request to via paid proxy via scraperbee
:param url: the url to call
:param headers: the headers to pass to the request
:param premium: whether or not to use premium proxy
:param stealth: whether or not to use stealth proxy
:param api_conn: the connection to use for the proxy
:return: The response.
"""
proxy_key = BaseHook.get_connection(api_conn).password
proxies = {
"http": f"http://{proxy_key}:render_js=False&premium_proxy={premium}&stealth_proxy={stealth}@proxy.scrapingbee.com:8886",
"https": f"https://{proxy_key}:render_js=False&premium_proxy={premium}&stealth_proxy={stealth}@proxy.scrapingbee.com:8887",
}
return BaseUtil.get_request(
url, headers=headers, proxies=proxies, verify=False, retry_count=retry_count
)
[docs]
@staticmethod
def normalize_string(input_string: str) -> str:
"""
Normalize a string to be used as a column name
:param input_string: the string to normalize
:return: the normalized string
"""
return re.sub(
r"\W+",
"",
str(input_string)
.replace(" ", "_")
.replace("-", "_")
.replace(".", "_")
.replace("/", "_")
.lower(),
)
[docs]
@staticmethod
def get_num_unnamed(df, divisor: str = 1) -> int:
"""
Get the number of unnamed columns in a dataframe
:param df: the dataframe to count
:return num_unnamed: the number of unnamed columns
"""
return int(
max(
int(col.split(": ")[1])
for col in df.columns.tolist()
if "Unnamed:" in col
)
/ divisor
)
[docs]
@staticmethod
def normalize_df(df):
"""
Normalize a dataframe's column names
:param df: the dataframe to prep
:return: df
"""
norm_cols = {col: BaseUtil.normalize_string(col) for col in df.columns.tolist()}
df.rename(mapper=norm_cols, axis="columns", inplace=True)
return df
[docs]
@staticmethod
def prep_wrangler_df(df):
"""
Prep a wrangler dataframe for raw ingestion
:param df: the dataframe to prep
:return: df
:rtype: [pandas.Dataframe]
"""
for col in df.columns.tolist():
if df[col].isnull().all():
df[col] = df[col].astype(str)
return df
[docs]
@staticmethod
def lob_address(
street_address: str, city: str = "", state: str = "", zip_code: str = ""
) -> dict:
"""
Send address to the cached lob api
:param street_address: the street address, required
:param city: the city, must be provided with state
:param state: the state
:param zip_code: the zip code, optionally include instead of city and state
"""
import requests
lob_key = BaseHook.get_connection(Config.ROB_AUTH_KEY).password
return requests.post(
f"{Config.ROB_URL}/validated-addresses/lob",
headers={"Authorization": lob_key},
json={
"street_address": street_address,
"city": city,
"state": state,
"zip_code": zip_code,
},
)
[docs]
@staticmethod
def lob_address_bulk(input_df, sleep: float = 0.03, max_retries=5) -> dict:
"""
Send address to the cached lob api
:param input_df: a pandas dataframe with following columns `"id","city","primary_line","state","zip_code"`
returns an array with following columns `query, id, response, created_at`
"""
import requests
import json
import base64
from requests_auth_aws_sigv4 import AWSSigV4
lob_key = BaseHook.get_connection(Config.ROB_AUTH_KEY).password
input_json = input_df.to_json(orient="records")
input_dict = json.loads(input_json)
input_jsonstring = json.dumps({"addresses": input_dict})
input_base64_enncoded = base64.b64encode(input_jsonstring.encode())
data = json.dumps(
{
"body": input_base64_enncoded.decode("utf-8"),
"resource": "/{proxy+}",
"path": "/validated-addresses/lob-bulk",
"httpMethod": "POST",
"isBase64Encoded": "true",
"headers": {
"Authorization": lob_key,
"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8",
},
"requestContext": {},
}
)
for i in range(max_retries):
time.sleep(random.random())
try:
response = requests.request(
method="POST",
url=Config.ROB_API_LAMBDA_FUNCTION_URL,
data=data,
auth=AWSSigV4("lambda"),
headers={"Authorization": lob_key},
)
json_response = response.json()
response_body = json.loads(json_response["body"])
print("API CALL :", json_response["statusCode"])
if json_response["statusCode"] == 200:
return response_body
print("API response ", response_body)
except Exception as ex:
# TODO: need better exception handling
print("error: ", ex, "attempt no: ", i)
print("DEBUG: ALL ATTEMPTS FAILED, returning empty list (of dict)")
return []
[docs]
@staticmethod
def get_aws_secret(secret_id: str) -> str:
"""
Get the secret from the aws secret manager
:param secret_id: the secret id
:return: the secret in string format
:raises: ClientError
"""
import boto3
from botocore.exceptions import ClientError
client = boto3.client("secretsmanager")
try:
return client.get_secret_value(SecretId=secret_id)["SecretString"]
except ClientError as e:
# For a list of exceptions thrown, see
# https://docs.aws.amazon.com/secretsmanager/latest/apireference/API_GetSecretValue.html
print("Error getting secret value. returning None", e)
raise e
[docs]
@staticmethod
def get_json_safely(response) -> dict:
"""Get the JSON from the response, safely
:param response: the response to get the JSON from
:type response: _type_
:raises ValueError: if there is an error in the JSON
:return: the JSON
:rtype: dict
"""
# bad status code
if response.status_code != 200:
response.raise_for_status()
json = response.json() # get the JSON
if "error" in json:
raise ValueError("Error: %s" % json["error"])
return json