Source code for operators.s3_geojson_to_athena

"""Operator that converts a GeoJSON file in S3 into a parquet-backed Athena table with a WKT column."""

from config import Config
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults




class S3GeoJsonToAthenaOperator(BaseOperator):
    """
    Converts a GeoJSON file stored in S3 to an Athena table, stored as parquet
    The output fill will contain a 'wkt' column in addition to the WKB geometry column
    The expected input file must be in the geopandas-parseable .geojson format
    Please follow standard naming and file structure conventions for both the storage of the
        geojson and the output table in Athena
    This operation is performed ON THE AIRFLOW SERVER. Adapt as needed to scale.

    :s3_geojson_path: the path to the GeoJSON in S3
    :type s3_geojson_path: str

    :output_table_name: the name to use for the output table
    :type output_table_name: str

    :output_path: the path to the output table in S3
    :type output_path: str

    :geometry_column: the name of the geometry column in the GeoJSON (default is "geometry")
    :type geometry_column: str

    :custom_columns: a dictionary of custom columns and their data types (default is None, string assigned for all columns)
    :type custom_columns: dict
    """

    ui_color = "420a69"

    @apply_defaults
    def __init__(
        self: object,
        s3_geojson_path: str,
        output_table_name: str,
        output_path: str,
        geometry_column: str = "geometry",
        custom_columns: dict = None,
        *args: list,
        **kwargs: dict,
    ):
        super().__init__(
            *args,
            **kwargs,
        )
        self.s3_geojson_path = s3_geojson_path
        self.output_table_name = output_table_name
        self.output_path = output_path
        self.custom_columns = custom_columns
        self.geometry_column = geometry_column



    def execute(
        self: object,
        context: dict,
    ) -> None:
        try:
            import awswrangler as wr
            import geopandas as gpd

            # Call me a castor canadensis cause I be stacking logs
            self.log.info(f"Source filepath: {self.s3_geojson_path}")
            self.log.info(f"Output table name: {self.output_table_name}")
            self.log.info(f"Destination filepath: {self.output_path}")
            self.log.info(f"Geometry column: {self.geometry_column}")
            self.log.info(f"Custom columns: {self.custom_columns}")

            # Read the GeoJSON directly from S3 into local memory
            gdf = gpd.read_file(self.s3_geojson_path)

            # Adjust projection to 4326
            gdf.to_crs(epsg=4326, inplace=True)

            # Cast geometry column to WKT, add to gdf as column
            gdf["wkt"] = gdf[f"{self.geometry_column}"].to_wkt()

            # Cast original column to WKB, replace geometry column with WKB
            gdf[f"{self.geometry_column}"] = gdf[f"{self.geometry_column}"].to_wkb()

            # Establish column names and types
            if self.custom_columns is None:
                col_types_dict = self.custom_columns
            else:
                # Assign all columns as string by default
                col_types_dict = {col: "string" for col in gdf.columns}
                # Remove geometry column from the dictionary to retain source data type
                col_types_dict.pop(f"{self.geometry_column}")

            # Write parquet file to S3, register table in Athena
            response = wr.s3.to_parquet(
                df=gdf,
                path=self.output_path,
                dataset=True,
                database=Config.DATA_LAKE_DB,
                table=self.output_table_name,
                mode="overwrite",
                dtype=col_types_dict,
            )

            # Logging
            self.log.info(f"Saved to S3: {response['paths']}")
            self.log.info(
                f"Saved to Athena as: {Config.DATA_LAKE_DB}.{self.output_table_name}"
            )

        except Exception as e:
            self.log.error(f"Error: {e}")
            raise e