from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from config import Config
from util.base import BaseUtil
from typing import TYPE_CHECKING
if TYPE_CHECKING:
import pandas as pd
[docs]
class EsriToAthenaOperator(BaseOperator):
"""
To convert an ESRI endpoint to Athena table
:param sri_endpoint: the endpoint of the ESRI service
:type esri_endpoint: str
:param fields: list of fields to return, defaults to None which loads all fields by passing `*` to the API
:type fields: list, optional
:param where: filter criteria, defaults to None which means no filter and it sends `1=1` to the API
:type where: str, optional
:param limit: max number of records to get, defaults to all the values in the feature layer
:type limit: int, optional
:param output_table_name: the name to use for the output table
:type output_table_name: str
:param output_path: the path to the output table in S3
:type output_path: str
"""
# dark green
ui_color = "#006400"
@apply_defaults
def __init__(
self,
esri_endpoint,
output_table_name,
output_path,
fields=None,
where=None,
limit=None,
*args,
**kwargs,
):
super(EsriToAthenaOperator, self).__init__(
*args,
**kwargs,
)
self.esri_endpoint = esri_endpoint
self.fields = fields
self.where = where
self.limit = limit
self.output_table_name = output_table_name
self.output_path = output_path
[docs]
def download_from_esri(
self: object,
url: str,
fields=None,
where=None,
limit=None,
**kwargs,
) -> "pd.DataFrame":
"""
Download data from an ESRI endpoint
:param url: the endpoint of the ESRI service
:type url: str
:param fields: list of fields to return, defaults to None which loads all fields by passing `*` to the API
:type fields: list, optional
:param where: filter criteria, defaults to None which means no filter and it sends `1=1` to the API
:type where: str, optional
:param limit: max number of records to get, defaults to all the values in the feature layer
:type limit: int, optional
:return: A pandas DataFrame with the data
:rtype: pd.DataFrame
sample url: https://pcrdgis.agriculture.purdue.edu/arcgis/rest/services/DDI/Digital_Divide_Index_COUNTY_USA_2020/MapServer/0
"""
import geopandas as gpd
import pandas as pd
import json
from tempfile import NamedTemporaryFile
# from util.base import BaseUtil.get_json_safely
print(f"Attempting to download data from ESRI endpoint: {url}")
gpd.options.io_engine = "pyogrio"
# Get the max record count
metadata = BaseUtil.get_request(url=url, params=dict(f="pjson")).json()
max_record_count = metadata["maxRecordCount"]
# default behavior matches all features
if where is None:
where = "1=1"
if fields is None:
fields = "*"
else:
fields = ", ".join(fields)
# extract object IDs of features
queryURL = f"{url}/query"
# Get the total record count
params = dict(
where=where,
returnCountOnly="true",
f="json",
)
response = BaseUtil.get_request(
url=queryURL,
params=params,
)
total_size = BaseUtil.get_json_safely(response)["count"]
# Check the limit
if limit is not None:
total_size = limit
# params for this request
resultOffset = 0
params = dict(
f="json",
outSR="4326",
outFields=fields,
resultOffset=resultOffset,
where=where,
**kwargs,
)
calls = total_size // max_record_count
if calls > 10:
print(
f"Long download time — total download will require {calls} separate requests"
)
out = []
while params["resultOffset"] < total_size:
remaining = total_size - params["resultOffset"]
if remaining < max_record_count:
params["resultRecordCount"] = remaining
# get raw features
response = BaseUtil.get_request(
url=queryURL,
params=params,
)
json_content = BaseUtil.get_json_safely(response)
with NamedTemporaryFile(
suffix=".json",
mode="w+",
) as json_file:
# write to file
json.dump(json_content, json_file)
json_file.flush()
gdf = gpd.read_file(
json_file.name,
driver="ESRIJSON",
crs="EPSG:4326",
)
out.append(gdf)
params["resultOffset"] += len(out[-1])
return pd.concat(out, axis=0).reset_index(drop=True)
[docs]
def write_to_athena(
self: object,
gdf: "pd.DataFrame",
output_table_name: str,
output_path: str,
) -> None:
"""
Write a GeoDataFrame to Athena
:param gdf: the GeoDataFrame to write
:type gdf: 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}")
# project to 4326
gdf.to_crs(
epsg=4326,
inplace=True,
)
# all column types to string
col_types_dict = {col: "string" for col in gdf.columns}
# convert geometry to WKT
gdf["geometry"] = gdf["geometry"].apply(lambda geom: geom.wkt)
wr.s3.to_parquet(
df=gdf,
path=output_path,
dataset=True,
database=Config.DATA_LAKE_DB,
table=f"r_{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: object,
context: dict,
) -> None:
try:
gdf = self.download_from_esri(
url=self.esri_endpoint,
fields=self.fields,
where=self.where,
limit=self.limit,
)
self.write_to_athena(
gdf=gdf,
output_table_name=self.output_table_name,
output_path=self.output_path,
)
except Exception as e:
print(f"Error: {e}")
raise e