merica-pipelines/dags/helpers/base_lakehouse_operator.py

"""
Base Lakehouse Operator for MERICA Data Lakehouse Studio
Provides standard lineage extraction, OpenMetadata 2.0 lineage emission,
and shared catalog resolution across all lakehouse operators.
"""
import os
import json
import logging
import urllib.request
import urllib.parse
from typing import Any, Dict, List, Optional, Sequence, Union

try:
    from airflow.sdk.bases.operator import BaseOperator
    from airflow.sdk.definitions.context import Context
except ImportError:
    try:
        from airflow.models.baseoperator import BaseOperator
        from airflow.utils.context import Context
    except ImportError:
        class BaseOperator:
            def __init__(self, *args, **kwargs):
                self.task_id = kwargs.get("task_id", "lakehouse_operator")
                self.inlets = kwargs.get("inlets", [])
                self.outlets = kwargs.get("outlets", [])
                self.log = logging.getLogger(self.__class__.__name__)
        class Context(dict):
            pass

logger = logging.getLogger(__name__)

ENABLE_OPENMETADATA_SYNC = os.getenv("ENABLE_OPENMETADATA_SYNC", "true").lower() in ("true", "1", "yes")


class BaseLakehouseOperator(BaseOperator):
    """
    Abstract base operator for all MERICA Lakehouse ingestion and transformation tasks.
    Enforces standardized dataset metadata (inlets/outlets) and automated lineage emission
    to OpenMetadata 2.0 REST API upon successful execution.
    """

    template_fields: Sequence[str] = ("layer", "pipeline_id")

    def __init__(
        self,
        *,
        layer: str = "bronze",
        engine: str = "generic",
        pipeline_id: Optional[str] = None,
        inlets: Optional[List[str]] = None,
        outlets: Optional[List[str]] = None,
        lakehouse_catalog: Optional[str] = None,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.layer = layer.lower()
        self.engine = engine.lower()
        self.pipeline_id = pipeline_id
        self.custom_inlets = inlets or []
        self.custom_outlets = outlets or []
        self.lakehouse_catalog = lakehouse_catalog or os.getenv("ICEBERG_CATALOG_NAME", "dlh_lakehouse")

        # Populate Airflow native inlets/outlets if supported
        if inlets:
            self.inlets = inlets
        if outlets:
            self.outlets = outlets

    def post_execute(self, context: Context, result: Any = None):
        """
        Airflow post-execution lifecycle hook.
        Automatically emits lineage edge to OpenMetadata 2.0 when task succeeds.
        """
        super().post_execute(context, result)
        dag_id = self.dag_id or (context.get("dag").dag_id if context.get("dag") else self.pipeline_id)
        task_id = self.task_id

        logger.info(
            "Task '%s' in DAG '%s' (Layer: %s, Engine: %s) completed successfully.",
            task_id, dag_id, self.layer, self.engine
        )

        if ENABLE_OPENMETADATA_SYNC:
            self._emit_openmetadata_lineage(context, dag_id, task_id)

    def _emit_openmetadata_lineage(self, context: Context, dag_id: str, task_id: str):
        """
        Resolve OpenMetadata JWT token and register lineage edges for all inlets -> outlets.
        """
        try:
            token, active_host = self._get_openmetadata_auth()
            if not token or not active_host:
                logger.debug("OpenMetadata host/token unavailable; skipping lineage emission.")
                return

            pipeline_fqn = f"dlh_airflow_service.{dag_id}"
            pipeline_id_uuid = self._resolve_entity_uuid(active_host, token, "pipeline", pipeline_fqn)

            # Emit lineage for each inlet -> outlet pair
            for inlet in self.custom_inlets:
                inlet_fqn = self._normalize_dataset_fqn(inlet)
                inlet_uuid = self._resolve_entity_uuid(active_host, token, "table", inlet_fqn)
                if not inlet_uuid:
                    inlet_uuid = self._ensure_table_entity(active_host, token, inlet_fqn)

                for outlet in self.custom_outlets:
                    outlet_fqn = self._normalize_dataset_fqn(outlet)
                    outlet_uuid = self._resolve_entity_uuid(active_host, token, "table", outlet_fqn)
                    if not outlet_uuid:
                        outlet_uuid = self._ensure_table_entity(active_host, token, outlet_fqn)

                    if inlet_uuid and outlet_uuid:
                        payload = {
                          "edge": {
                            "fromEntity": {"id": inlet_uuid, "type": "table"},
                            "toEntity": {"id": outlet_uuid, "type": "table"},
                          }
                        }
                        if pipeline_id_uuid:
                            payload["edge"]["lineageDetails"] = {
                                "pipeline": {"id": pipeline_id_uuid, "type": "pipeline"}
                            }

                        req = urllib.request.Request(
                            f"{active_host}/api/v1/lineage",
                            data=json.dumps(payload).encode("utf-8"),
                            headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
                            method="PUT"
                        )
                        with urllib.request.urlopen(req, timeout=4) as resp:
                            logger.info(
                                "✅ Lineage emitted: %s -> %s via pipeline %s",
                                inlet_fqn, outlet_fqn, dag_id
                            )
        except Exception as e:
            logger.warning("OpenMetadata Lineage Emission Note: %s", e)

    def _normalize_dataset_fqn(self, dataset_str: str) -> str:
        """
        Convert dataset URI/table into OpenMetadata FQN format.
        e.g. 's3://dlh-bronze/db_cdc/environmental_permits' -> 'dlh_s3_storage.dlh-bronze.db_cdc.environmental_permits'
        e.g. 'silver.environmental_permits_clean' -> 'dlh_postgres_service.dlh_lakehouse.silver.environmental_permits_clean'
        """
        clean = dataset_str.strip()
        if clean.startswith("s3://"):
            parts = clean.replace("s3://", "").strip("/").split("/")
            return f"dlh_minio_s3.{'.'.join(parts)}"
        elif "." in clean:
            return f"dlh_postgres_service.dlh_lakehouse.{clean}"
        return f"dlh_postgres_service.dlh_lakehouse.public.{clean}"

    def _get_openmetadata_auth(self) -> tuple:
        """Authenticate and get access token."""
        hosts = [
            os.getenv("OPENMETADATA_HOST", "http://openmetadata:8585"),
            "http://localhost:8585",
            "http://127.0.0.1:8585"
        ]
        for h in hosts:
            try:
                login_req = urllib.request.Request(
                    f"{h}/api/v1/users/login",
                    data=json.dumps({"email": "admin@open-metadata.org", "password": "YWRtaW4="}).encode("utf-8"),
                    headers={"Content-Type": "application/json"},
                    method="POST"
                )
                with urllib.request.urlopen(login_req, timeout=3) as resp:
                    token = json.loads(resp.read().decode("utf-8")).get("accessToken")
                    if token:
                        return token, h
            except Exception:
                continue
        return None, None

    def _resolve_entity_uuid(self, host: str, token: str, entity_type: str, fqn: str) -> Optional[str]:
        """Fetch entity UUID from OpenMetadata given FQN."""
        try:
            endpoint = "tables" if entity_type == "table" else f"{entity_type}s"
            url = f"{host}/api/v1/{endpoint}/name/{urllib.parse.quote(fqn, safe='')}"
            req = urllib.request.Request(
                url,
                headers={"Authorization": f"Bearer {token}", "Accept": "application/json"},
                method="GET"
            )
            with urllib.request.urlopen(req, timeout=3) as resp:
                data = json.loads(resp.read().decode("utf-8"))
                return data.get("id")
        except Exception:
            return None

    def _ensure_table_entity(self, host: str, token: str, fqn: str) -> Optional[str]:
        """Auto-provision a missing table entity in OpenMetadata."""
        try:
            parts = fqn.split(".")
            table_name = parts[-1]
            schema_fqn = ".".join(parts[:-1]) if len(parts) > 1 else "dlh_postgres_service.dlh_lakehouse.public"
            payload = {
                "name": table_name,
                "displayName": table_name.replace("_", " ").title(),
                "databaseSchema": schema_fqn,
                "description": f"Managed Lakehouse Table '{table_name}'",
                "columns": [
                    {"name": "id", "dataType": "BIGINT", "description": "Primary Key"},
                    {"name": "data", "dataType": "JSON", "description": "Record Payload"},
                    {"name": "updated_at", "dataType": "TIMESTAMP", "description": "Timestamp"}
                ]
            }
            req = urllib.request.Request(
                f"{host}/api/v1/tables",
                data=json.dumps(payload).encode("utf-8"),
                headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"},
                method="PUT"
            )
            with urllib.request.urlopen(req, timeout=4) as resp:
                data = json.loads(resp.read().decode("utf-8"))
                return data.get("id")
        except Exception:
            return None