merica-pipelines/dags/helpers/duckdb_operator.py

"""
DuckDB Iceberg Transformation Operator for MERICA Data Lakehouse Studio
Executes in-process SQL transformations from Bronze Parquet to Silver Iceberg tables.
Applies Lazy Import Guardrail to prevent memory bloat on DAG Processor.
"""
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 DuckDBIcebergOperator(BaseLakehouseOperator):
    """
    DuckDB-powered Lakehouse Operator.
    Transforms Bronze Parquet datasets to clean Silver Iceberg/Parquet tables
    with high-throughput, low-memory in-process SQL execution.
    """

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

    def __init__(
        self,
        *,
        source_parquet_path: str,
        target_table_name: str,
        transformation_sql: Optional[str] = None,
        primary_keys: Optional[List[str]] = None,
        order_by_col: Optional[str] = "updated_at",
        s3_endpoint: Optional[str] = None,
        s3_access_key: Optional[str] = None,
        s3_secret_key: Optional[str] = None,
        s3_bucket: Optional[str] = "dlh-silver",
        **kwargs
    ):
        inlets = kwargs.pop("inlets", [source_parquet_path])
        clean_target = target_table_name if "." in target_table_name else f"silver.{target_table_name}"
        outlets = kwargs.pop("outlets", [clean_target])

        super().__init__(
            layer="silver",
            engine="duckdb",
            inlets=inlets,
            outlets=outlets,
            **kwargs
        )

        self.source_parquet_path = source_parquet_path
        self.target_table_name = target_table_name
        self.primary_keys = primary_keys or ["id"]
        self.order_by_col = order_by_col

        if not transformation_sql:
            pk_cols = ", ".join(self.primary_keys)
            order_col = self.order_by_col or self.primary_keys[0]
            read_func = "read_json_auto" if self.source_parquet_path.lower().endswith(".json") else "read_parquet"
            self.transformation_sql = f"""
                WITH raw_source AS (
                    SELECT * FROM {read_func}('{self.source_parquet_path}')
                ),
                ranked AS (
                    SELECT *,
                        ROW_NUMBER() OVER (PARTITION BY {pk_cols} ORDER BY {order_col} DESC) as _row_num
                    FROM raw_source
                )
                SELECT * EXCLUDE (_row_num)
                FROM ranked
                WHERE _row_num = 1
            """.strip()
        else:
            self.transformation_sql = transformation_sql

        self.s3_endpoint = s3_endpoint or os.getenv("MINIO_ENDPOINT", "minio:9000")
        self.s3_access_key = s3_access_key or os.getenv("MINIO_ROOT_USER", "minio_admin")
        self.s3_secret_key = s3_secret_key or os.getenv("MINIO_ROOT_PASSWORD", "minio_password")
        self.s3_bucket = s3_bucket or "dlh-silver"

    def execute(self, context: Context) -> Dict[str, Any]:
        """
        Execute DuckDB SQL pipeline.
        Lazy imports duckdb inside execute() to respect ADR-03 & avoid scheduler OOM.
        """
        logger.info(
            "Starting DuckDB Iceberg transformation: %s -> %s",
            self.source_parquet_path, self.target_table_name
        )

        # Lazy Import Guardrail
        try:
            import duckdb
        except ImportError as e:
            logger.error("DuckDB is not installed in the runtime environment: %s", e)
            raise

        con = duckdb.connect(database=":memory:")
        try:
            # Configure MinIO / S3 connectivity in DuckDB
            clean_endpoint = self.s3_endpoint.replace("http://", "").replace("https://", "")
            con.execute("INSTALL httpfs; LOAD httpfs;")
            con.execute(f"SET s3_endpoint = '{clean_endpoint}';")
            con.execute(f"SET s3_access_key_id = '{self.s3_access_key}';")
            con.execute(f"SET s3_secret_access_key = '{self.s3_secret_key}';")
            con.execute("SET s3_use_ssl = false;")
            con.execute("SET s3_url_style = 'path';")
            con.execute("SET s3_region = 'us-east-1';")

            # Resolve query
            target_s3_path = f"s3://{self.s3_bucket}/{self.target_table_name}/data.parquet"

            if self.transformation_sql.strip().upper().startswith("COPY"):
                query = self.transformation_sql
            else:
                query = f"""
                    COPY (
                        {self.transformation_sql}
                    ) TO '{target_s3_path}' (FORMAT PARQUET, COMPRESSION ZSTD);
                """

            logger.info("Executing DuckDB Query:\n%s", query)
            con.execute(query)

            # Get row count
            count_res = con.execute(f"SELECT count(*) FROM read_parquet('{target_s3_path}')").fetchone()
            rows_written = count_res[0] if count_res else 0

            logger.info(
                "✅ DuckDB Transformation completed successfully! %d rows written to %s",
                rows_written, target_s3_path
            )

            return {
                "status": "success",
                "engine": "duckdb",
                "target_table": self.target_table_name,
                "target_path": target_s3_path,
                "rows_processed": rows_written
            }
        except Exception as e:
            logger.error("DuckDB Transformation Failed: %s", e)
            raise
        finally:
            con.close()