merica-pipelines/dags/helpers/pyiceberg_operator.py

"""
PyIceberg Transformation Operator for MERICA Data Lakehouse Studio
Executes native Apache Iceberg table management using pyiceberg and PyArrow.
Applies Lazy Import Guardrail to protect Airflow Scheduler memory.
"""
import os
import logging
from typing import Any, Dict, List, Optional, Sequence

from helpers.base_lakehouse_operator import BaseLakehouseOperator
try:
    from airflow.sdk.definitions.context import Context
except ImportError:
    try:
        from airflow.utils.context import Context
    except ImportError:
        class Context(dict):
            pass

logger = logging.getLogger(__name__)


class PyIcebergTransformationOperator(BaseLakehouseOperator):
    """
    Apache PyIceberg-powered Lakehouse Operator.
    Loads Bronze Parquet data, creates/evolves Silver Iceberg tables,
    and commits atomic snapshots to the Iceberg Catalog.
    """

    template_fields: Sequence[str] = (
        "source_parquet_path",
        "target_table_name"
    )

    def __init__(
        self,
        *,
        source_parquet_path: str,
        target_table_name: str,
        namespace: str = "silver",
        write_mode: str = "overwrite",  # "append" or "overwrite"
        catalog_name: Optional[str] = "dlh_lakehouse",
        catalog_type: Optional[str] = "sql",  # "sql" (PostgreSQL) or "rest"
        catalog_uri: Optional[str] = None,
        s3_endpoint: Optional[str] = None,
        primary_keys: Optional[List[str]] = None,
        **kwargs
    ):
        inlets = kwargs.pop("inlets", [source_parquet_path])
        clean_target = target_table_name if "." in target_table_name else f"{namespace}.{target_table_name}"
        outlets = kwargs.pop("outlets", [clean_target])

        super().__init__(
            layer=namespace,
            engine="pyiceberg",
            inlets=inlets,
            outlets=outlets,
            **kwargs
        )

        self.source_parquet_path = source_parquet_path
        self.target_table_name = target_table_name
        self.namespace = namespace
        self.write_mode = write_mode.lower()
        self.catalog_name = catalog_name or "dlh_lakehouse"
        self.catalog_type = catalog_type or "sql"
        self.primary_keys = primary_keys or ["id"]

        self.catalog_uri = catalog_uri or os.getenv(
            "ICEBERG_CATALOG_URI",
            "postgresql+psycopg2://dlh_user:dlh_password@postgres:5432/dlh_lakehouse"
        )
        self.s3_endpoint = s3_endpoint or os.getenv("MINIO_ENDPOINT", "http://minio:9000")

    def execute(self, context: Context) -> Dict[str, Any]:
        """
        Execute PyIceberg Table snapshot commit.
        Lazy imports pyiceberg and pyarrow inside execute().
        """
        logger.info(
            "Starting PyIceberg Table transformation: %s -> %s.%s",
            self.source_parquet_path, self.namespace, self.target_table_name
        )

        # Lazy Import Guardrail
        try:
            import pyarrow.parquet as pq
            import pyarrow.fs as pafs
            import pyarrow.csv as pacsv
            from pyiceberg.catalog import load_catalog
        except ImportError as e:
            logger.error("PyIceberg or PyArrow is missing: %s", e)
            raise

        try:
            # 1. Read Bronze dataset into PyArrow Table via S3FileSystem
            logger.info("Reading source dataset from %s", self.source_parquet_path)
            s3_access_key = os.getenv("MINIO_ROOT_USER", "minio_admin")
            s3_secret_key = os.getenv("MINIO_ROOT_PASSWORD", "minio_password")
            
            clean_endpoint = self.s3_endpoint
            if not clean_endpoint.startswith("http"):
                clean_endpoint = f"http://{clean_endpoint}"
                
            s3_fs = pafs.S3FileSystem(
                endpoint_override=clean_endpoint,
                access_key=s3_access_key,
                secret_key=s3_secret_key,
                scheme="http"
            )

            if self.source_parquet_path.startswith("s3://"):
                s3_key = self.source_parquet_path.replace("s3://", "")
                if s3_key.lower().endswith(".csv"):
                    arrow_table = pacsv.read_csv(s3_key, filesystem=s3_fs)
                else:
                    arrow_table = pq.read_table(s3_key, filesystem=s3_fs)
            else:
                if self.source_parquet_path.lower().endswith(".csv"):
                    arrow_table = pacsv.read_csv(self.source_parquet_path)
                else:
                    arrow_table = pq.read_table(self.source_parquet_path)

            row_count = arrow_table.num_rows

            # 2. Connect to Iceberg Catalog
            catalog_props = {
                "type": self.catalog_type,
                "uri": self.catalog_uri,
                "s3.endpoint": clean_endpoint,
                "s3.access-key-id": s3_access_key,
                "s3.secret-access-key": s3_secret_key,
                "s3.path-style-access": "true",
                "s3.region": "us-east-1",
            }
            catalog = load_catalog(self.catalog_name, **catalog_props)

            # Ensure namespace exists in catalog
            clean_table = self.target_table_name.split(".")[-1]
            full_identifier = f"{self.namespace}.{clean_table}"
            try:
                catalog.create_namespace_if_not_exists(self.namespace)
            except Exception as ns_err:
                logger.debug("Namespace check note: %s", ns_err)

            # 3. Load or Create Table
            try:
                table = catalog.load_table(full_identifier)
                logger.info("Loaded existing Iceberg table '%s'", full_identifier)
            except Exception:
                logger.info("Table '%s' not found. Creating new Iceberg table...", full_identifier)
                table = catalog.create_table(
                    identifier=full_identifier,
                    schema=arrow_table.schema,
                    location=f"s3://dlh-silver/{clean_table}"
                )

            # 4. Commit snapshot (Append or Overwrite)
            if self.write_mode == "append":
                table.append(arrow_table)
            else:
                table.overwrite(arrow_table)

            logger.info(
                "✅ PyIceberg snapshot committed! %d records committed to table '%s'",
                row_count, full_identifier
            )

            return {
                "status": "success",
                "engine": "pyiceberg",
                "table": full_identifier,
                "rows_committed": row_count,
                "write_mode": self.write_mode
            }
        except Exception as e:
            logger.error("PyIceberg Transformation Failed: %s", e)
            raise