Source code for operators.s3_shp_to_athena

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

from config import Config




class S3ShapefileToAthenaOperator(BaseOperator):
    """
    To convert a shapefile in S3 to Athena table

    Ensure that all the relevant files are in the same directory.
    (e.g. .shp, .shx, .dbf, .prj, .cpg)

    Default EPSG is 4326

    :s3_shp_path: the path to the shapefile in S3
    :type s3_shp_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
    """

    ui_color = "#A6E6A6"

    @apply_defaults
    def __init__(
        self,
        s3_shp_path,
        output_table_name,
        output_path,
        s3_shp_layer=None,
        *args,
        **kwargs,
    ):
        super(S3ShapefileToAthenaOperator, self).__init__(
            *args,
            **kwargs,
        )
        self.s3_shp_path = s3_shp_path
        self.output_table_name = output_table_name
        self.output_path = output_path
        self.s3_shp_layer = s3_shp_layer

    def _resolve_zip_shp_path(self, s3_path: str) -> str:
        """
        If the path is a .zip file with a nested directory structure,
        find the .shp file inside and return the proper path for GDAL/Fiona.
        """
        if not s3_path.lower().endswith(".zip"):
            return s3_path

        import io  # Lazy load: standard library
        import zipfile  # Lazy load: standard library

        import boto3  # Lazy load: heavy AWS SDK

        path_without_prefix = s3_path.replace("s3://", "")
        bucket = path_without_prefix.split("/")[0]
        key = "/".join(path_without_prefix.split("/")[1:])

        s3_client = boto3.client("s3")
        zip_obj = s3_client.get_object(Bucket=bucket, Key=key)
        zip_content = io.BytesIO(zip_obj["Body"].read())

        with zipfile.ZipFile(zip_content, "r") as zf:
            shp_files = [f for f in zf.namelist() if f.lower().endswith(".shp")]
            if not shp_files:
                raise ValueError(f"No .shp file found in zip: {s3_path}")
            if len(shp_files) > 1:
                self.log.warning(
                    f"Multiple .shp files found in zip, using first: {shp_files}"
                )
            shp_inside = shp_files[0]

        if "/" in shp_inside:
            resolved_path = f"/vsizip/vsis3/{bucket}/{key}/{shp_inside}"
            self.log.info(f"Resolved nested zip path: {resolved_path}")
            return resolved_path

        return s3_path



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

            self.log.info(f"Source filepath: {self.s3_shp_path}")
            self.log.info(f"Shapefile layer name: {self.s3_shp_layer}")
            self.log.info(f"Output table name: {self.output_table_name}")
            self.log.info(f"Destination filepath: {self.output_path}")

            read_path = self._resolve_zip_shp_path(self.s3_shp_path)

            # Read the shapefile
            if not self.s3_shp_layer:
                gdf = gpd.read_file(read_path)
            else:
                gdf = gpd.read_file(read_path, layer=self.s3_shp_layer)

            # project to 4326
            gdf.to_crs(
                crs=Config.CRS_GPS,
                inplace=True,
            )

            # Create a new column with WKT strings
            gdf["geometry_wkt"] = gdf.geometry.apply(
                lambda geom: geom.wkt if geom is not None else None
            )
            # Convert to regular DataFrame
            gdf = pd.DataFrame(gdf)
            # Drop geometry column and rename
            gdf = gdf.drop(columns=["geometry"]).rename(
                columns={"geometry_wkt": "geometry"}
            )

            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",
                index=False,
                dtype={col: "string" for col in gdf.columns},
            )

            self.log.info(f"Saved to S3: {response['paths']}")
            self.log.info(
                f"Saved table to Athena as: {Config.DATA_LAKE_DB}.{self.output_table_name}"
            )
            self.log.info(f"Partitions, if any: {response['partitions_values']}")
        except Exception as e:
            self.log.error(f"Error: {e}")
            raise e