from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from config import Config
from util.base import BaseUtil
from typing import List, TypedDict
# from typing import TYPE_CHECKING
# if TYPE_CHECKING:
# import pandas as pd
RAW = Config.RAW_PREFIX
TRANSFORMED = Config.TRANSFORMED_PREFIX
class OutputDict(TypedDict):
output_table_name: str
df: "pd.DataFrame"
[docs]
class PeeringDBToAthenaOperator(BaseOperator):
"""
To convert a PeeringDB endpoint to Athena table
:param peeringdb_endpoint: the endpoint of the PeeringDB API
:type: peeringdb_endpoint: str
:param fields: list of fields to return, defaults to None which loads all fields
"""
@apply_defaults
def __init__(
self,
peeringdb_endpoint,
output_table_name,
output_path,
where=None,
limit=None,
*args,
**kwargs,
):
super(PeeringDBToAthenaOperator, self).__init__(
*args,
**kwargs,
)
self.peeringdb_endpoint = peeringdb_endpoint
self.where = where or {}
self.limit = limit
self.output_table_name = output_table_name
self.output_path = output_path
[docs]
def get_api_token(self) -> str:
from airflow.hooks.base import BaseHook
print(f"Attempting to fetch PeeringDB credentials from AWS Secret Manager")
return BaseHook.get_connection(Config.PEERINGDB_API_KEY).extra
[docs]
def get_total_count(self, url: str, headers: dict) -> int:
"""
Get total count of records for pagination upon depth expansion
"""
if self.where:
count_params = {
k: v
for k, v in self.where.items()
if k not in ["limit", "skip", "depth"]
}
else:
count_params = {}
count_response = BaseUtil.get_request(url, headers=headers, params=count_params)
count_response.raise_for_status()
response_data = count_response.json()
total_count = len(response_data["data"])
print(f"Total expanded records to fetch: {total_count}")
return total_count
[docs]
def fetch_data(self, url: str, headers: dict) -> list:
"""Fetch data without pagination when depth is not needed"""
params = dict(self.where) if self.where else {}
if self.limit:
limit = min(self.limit, 250)
if limit != self.limit:
print(f"Limit provided exceeds max allowed, limit is 250")
params["limit"] = limit
response = BaseUtil.get_request(url, headers=headers, params=params)
response.raise_for_status()
json_response = response.json()
all_data = json_response["data"]
print(f"Total records fetched: {len(all_data)}")
return all_data
[docs]
def fetch_paginated_data(self, url: str, headers: dict, total_count: int) -> list:
"""
Fetch all data with pagination
"""
all_data = []
skip = 0
if self.limit:
limit = min(self.limit, 250)
if limit != self.limit:
print(f"Limit provided exceeds max allowed, limit is 250")
else:
limit = 250
while skip <= total_count:
params = dict(self.where)
params.update({"limit": limit, "skip": skip})
print(f"Fetching records {skip} to {skip + limit}")
response = BaseUtil.get_request(url, headers=headers, params=params)
response.raise_for_status()
response_data = response.json()
all_data.extend(response_data["data"])
skip += limit
print(f"Total records fetched: {len(all_data)}")
return all_data
[docs]
def process_ix_data(self, ix_data: list) -> List[OutputDict]:
"""
Process IX data to extract both IXPs and facilities with US presence
:param ix_data: Raw IX data from API
:return: List[OutputDict] with 'ixps' and 'facilities' DataFrames
"""
import pandas as pd
import json
from tempfile import NamedTemporaryFile
# Filter IXPs and collect US facilities
ixps = []
us_facs = []
us_fac_ids = set()
for ixp in ix_data:
has_us_facility = False
# Check each ixp facility, and collect US based ones
if 'fac_set' in ixp and len(ixp['fac_set']) > 0:
for facility in ixp['fac_set']:
if facility.get('country') == 'US':
has_us_facility = True
if facility['id'] not in us_fac_ids:
us_facs.append(facility)
us_fac_ids.add(facility['id'])
if ixp['country'] == 'US' or has_us_facility:
ixps.append(ixp)
print(f'Found {len(ixps)} IXPs with US presence (country=US or has US facilities)')
print(f'Found {len(us_facs)} unique US facilities')
result = []
with NamedTemporaryFile(suffix=".json", mode="w+") as json_file:
json.dump(ixps, json_file)
json_file.flush()
result.append({'output_table_name': 'ixps', 'df': pd.read_json(json_file.name)})
with NamedTemporaryFile(suffix=".json", mode="w+") as json_file:
json.dump(us_facs, json_file)
json_file.flush()
result.append({'output_table_name': 'ixp_facilities', 'df': pd.read_json(json_file.name)})
return result
[docs]
def download_from_peeringdb(
self,
url: str,
token: str,
base_output_table_name: str
) -> List[OutputDict]:
"""
Download data from PeeringDB endpoint
:param
# need to include depth= 2 to expand sets into objects
# next / previous - pagination - while next is not None
# count is total count? maybe
"""
import pandas as pd
import json
from tempfile import NamedTemporaryFile
print(f"Attempting to download data from PeeringDB endpoint: {url}")
headers = self.get_headers(token)
# handle queries that need to fetch nested data
if "depth" in self.where and self.where["depth"] > 1:
total_count = self.get_total_count(url, headers)
all_data = self.fetch_paginated_data(
url, headers=headers, total_count=total_count
)
# handle nested data extraction for ixp endpoint
if url.endswith("ix"):
return self.process_ix_data(ix_data=all_data)
else:
all_data = self.fetch_data(url, headers=headers)
# All Data - manual filter to "US" due to API inconsistencies
all_us_data = [
item for item in all_data if "country" in item and item["country"] == "US"
]
print(f"Found {len(all_us_data)} US records for {url}")
with NamedTemporaryFile(
suffix=".json",
mode="w+",
) as json_file:
json.dump(all_us_data, json_file)
json_file.flush()
df = pd.read_json(json_file.name)
return [{'output_table_name': base_output_table_name, 'df': df}]
[docs]
def write_to_athena(
self,
df: "pd.DataFrame",
output_table_name: str,
output_path: str,
table_prefix: str
) -> None:
"""
Write a dataframe to Athena
:param df: the DataFrame to write
:type df: pd.DataFrame
:param output_table_name: name of the output table in Athena. "r_" prefix is added to the table name.
:type output_table_name: str
:param output_path: the path to the output table in S3
:type output_path: str
"""
try:
import awswrangler as wr
print(f"Output table name: {output_table_name}")
print(f"Destination filepath: {output_path}")
col_types_dict = {col: "string" for col in df.columns}
wr.s3.to_parquet(
df=df,
path=output_path,
dataset=True,
database=Config.DATA_LAKE_DB,
table=f"{table_prefix}{output_table_name}",
mode="overwrite",
index=False,
dtype=col_types_dict,
)
print(f"Saved to S3: {output_path}")
print(f"Saved to Athena: {output_table_name}")
except Exception as e:
print(f"Error: {e}")
raise e
[docs]
def execute(self, context) -> None:
try:
token = self.get_api_token()
raw_dfs = self.download_from_peeringdb(url=self.peeringdb_endpoint, base_output_table_name=self.output_table_name, token=token)
for raw_df in raw_dfs:
raw_base_path = self.output_path.replace('CONFIG', RAW)
raw_table_specific_path = f"{raw_base_path}{raw_df['output_table_name']}/"
# Write query data to Athena
self.write_to_athena(
df=raw_df['df'],
output_table_name=raw_df['output_table_name'],
output_path=raw_table_specific_path,
table_prefix='r_'
)
transformed_df = self.add_transformed_fields_to_df(raw_df)
transformed_base_path = self.output_path.replace('CONFIG', TRANSFORMED)
transformed_table_specific_path = f"{transformed_base_path}{transformed_df['output_table_name']}/"
# Write transformed query data to Athena
self.write_to_athena(
df=transformed_df['df'],
output_table_name=transformed_df['output_table_name'],
output_path=transformed_table_specific_path,
table_prefix='t_'
)
except Exception as e:
print(f"Error: {e}")
raise e