from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from config import Config
from airflow.models import BaseOperator
from util.base import BaseUtil
[docs]
class CQACollectGeometryOperator(BaseOperator):
ui_color = "#A6E6A6"
@apply_defaults
def __init__(
self,
vintage: str,
fields: list[str],
output_athena_dataset_name: str,
input_sql: str,
# TODO add script for new param skip_cqa_uuid_sql
skip_cqa_uuid_sql: str = "",
datalake_db: str = Config.DATA_LAKE_DB,
*args,
**kwargs,
):
"""
An operator to collect geometry from CostQuest API
:param vintage: the cqa vintage. Get it from https://costquest.azure-api.net/fabric/vintages
:type vintage: str
:param fields: A series of fields. only 5 fields are supported for now https://developer.costquest.com/api-details#api=fabrictiles&operation=fields
:type fields: list[str]
:param output_athena_dataset_name: the output table name
:type output_athena_dataset_name: str
:param input_sql: the sql which is supposed to to have two columns. id and geojson; OPTIMIZATION NOTE: Query should deduplicate records (geographic id e.g. already queried census tracts!) to minimize CQA API calls
:type input_sql: str
:param skip_cqa_uuid_sql: athena sql to select CQA uuids (SELECT ... AS cqa_uuid) that we have already queried once, and exclude from this query
:type skip_cqa_uuid_sql: str, optional
:param datalake_db: the datalake db, defaults to Config.DATA_LAKE_DB
:type datalake_db: str, optional
"""
self.output_athena_dataset_name = output_athena_dataset_name
self.vintage = vintage
self.fields = fields
self.all_df_fields = ["id", "api_call_status", "cqa_uuid"] + self.fields
self.datalake_db = datalake_db
self.input_sql = input_sql
self.base_uri = Config.CQA_BASE_URI
self.cqa_token = ""
self.s3_root = f"s3://{Config.DATA_LAKE_BUCKET}/{Config.RAW_PREFIX}/cqa_collect_geometry/{self.output_athena_dataset_name}/vintage={self.vintage}/"
self.map_name = f"api_calls_{output_athena_dataset_name}"
super().__init__(*args, **kwargs)
[docs]
def collect_cqa(self, geojson):
collect_id_request = BaseUtil.post_request(
url=f"{self.base_uri}fabric/{self.vintage}/collect",
data=geojson,
headers={"apikey": self.cqa_token},
)
collect_ids = collect_id_request.json()
if len(collect_ids) >= 5000:
divide_results_request = BaseUtil.post_request(
f"{self.base_uri}geosvc/divide",
data=geojson,
headers={"apikey": self.cqa_token},
)
divide_results = divide_results_request.json()
for db in divide_results:
divide_calls += 1
print("divide " + str(divide_calls))
yield from self.collect_cqa(db)
else:
yield collect_ids
[docs]
def split_into_chunks(self, array, chunk_size=1000):
idx = -1
for i in range(0, len(array), chunk_size):
idx += 1
yield array[i : i + chunk_size], idx
[docs]
def is_table_available(self):
import awswrangler as wr
return wr.catalog.does_table_exist(
database=self.datalake_db, table=self.output_athena_dataset_name
)
[docs]
def get_api_token(self):
from airflow.hooks.base import BaseHook
self.cqa_token = BaseHook.get_connection(Config.CQA_API_KEY).password
[docs]
def get_error_df(self, id: str, len_chunk: int):
import pandas as pd
output_df = pd.DataFrame(columns=self.all_df_fields)
new_row = {
"id": id,
"api_call_status": "failed",
"cqa_uuid": f"Number of failed uuids in this chunk: {len_chunk}",
}
output_df.loc[len(output_df)] = new_row
return output_df
[docs]
def get_bulk_locations(self, id: str, cqa_location_list: list[str]):
import pandas as pd
for chunk, idx in self.split_into_chunks(cqa_location_list):
try:
collect_id_request = BaseUtil.post_request(
url=f"{self.base_uri}fabric/{self.vintage}/bulk/locations",
# https://developer.costquest.com/api-details#api=fabrictiles&operation=bulk Request body: An array of strings defining the UUIDs of the data records to retrieve
data=chunk,
headers={"apikey": self.cqa_token},
params={"field": ",".join(self.fields)},
)
fabric_data = collect_id_request.json()
output_df = pd.DataFrame(fabric_data)
output_df["id"] = id
output_df["api_call_status"] = "success"
output_df["cqa_uuid"] = chunk
except Exception as ex:
len_chunk = len(chunk)
print(
f"Error processing row {idx+1}- {id}: {ex}, chunk {idx} (number of uuids: {len_chunk})"
)
yield self.get_error_df(id, len_chunk), idx
finally:
yield output_df, idx
[docs]
def execute(self, context):
import awswrangler as wr
import json
import pandas as pd
# add map_index_template
context["params"]["map_name"] = self.map_name
context["map_index_template"] = "{{ params.map_name }}"
self.get_api_token()
print("READING ATHENA INPUT SQL: ", self.input_sql)
input_df = wr.athena.read_sql_query(
sql=self.input_sql,
database=self.datalake_db,
)
if input_df.empty:
raise Exception("Input DataFrame is empty. Please check the input SQL.")
missing = [col for col in ["id", "geojson"] if col not in input_df.columns]
if missing:
raise Exception(
f"Input DataFrame must have columns 'id' and 'geojson'. existing columns are {missing}"
)
if len(self.fields) > 5:
raise Exception("Too many fields. Limit is 5.")
table_exists = self.is_table_available()
if not table_exists:
wr.catalog.create_parquet_table(
database=self.datalake_db,
table=self.output_athena_dataset_name,
path=self.s3_root,
columns_types={f: "string" for f in self.all_df_fields},
description="Table created from DataFrame",
)
print(f"Created table '{self.output_athena_dataset_name}'")
for idx, row in input_df.iterrows():
print(f"Processing row {idx+1}/{len(input_df)}: id={row['id']}")
id = row["id"]
try:
total_records_from_collect = 0
total_records_from_bulk = 0
output_df = pd.DataFrame(columns=self.all_df_fields)
for cqa_location_list in self.collect_cqa(json.loads(row["geojson"])):
df_list = []
total_records_from_collect += len(cqa_location_list)
for chunk_df, idx in self.get_bulk_locations(id, cqa_location_list):
df_list.append(chunk_df)
output_df = pd.concat(df_list)
total_records_from_bulk += len(output_df)
except Exception as ex:
# we want to know which geoids are failing
print(
f"Error processing row {idx+1} of {len(input_df)} i.e. {id=}: {ex}"
)
output_df = self.get_error_df(id, len(input_df))
finally:
print(
f""""
Processed row {idx+1}/{len(input_df)}: id={id}, total_records_from_collect={total_records_from_collect}, total_records_from_bulk={total_records_from_bulk}, output_df={output_df.shape}
"""
)
if not output_df.empty:
output_df = output_df[self.all_df_fields]
wr.s3.to_parquet(
df=output_df,
path=self.s3_root,
database=self.datalake_db,
table=self.output_athena_dataset_name,
dataset=True,
mode="append",
)