merica-pipelines/dags/helpers/dbt_trino_operator.py

"""
dbt-Trino Iceberg Transformation Operator for MERICA Data Lakehouse Studio
Executes SQL-first transformation models via dbt with Trino/PostgreSQL target.
Emits OpenMetadata 2.0 lineage automatically via BaseLakehouseOperator.
"""
import os
import subprocess
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 DbtTrinoIcebergOperator(BaseLakehouseOperator):
    """
    Executes dbt transformation models targeting Trino/Iceberg or Lakehouse PostgreSQL.
    Extracts model dependencies and emits lineage to OpenMetadata.
    """

    template_fields: Sequence[str] = (
        "select_models",
        "project_dir",
        "profiles_dir",
        "target"
    )

    def __init__(
        self,
        *,
        select_models: str,
        project_dir: Optional[str] = None,
        profiles_dir: Optional[str] = None,
        target: Optional[str] = None,
        layer: str = "silver",
        inlets: Optional[List[str]] = None,
        outlets: Optional[List[str]] = None,
        vars: Optional[Dict[str, Any]] = None,
        **kwargs
    ):
        model_name = select_models.replace(".sql", "").strip()
        custom_outlets = outlets or [f"{layer}.{model_name}"]
        custom_inlets = inlets or [f"bronze.{model_name.replace('silver_', '').replace('gold_', '')}"]

        super().__init__(
            layer=layer,
            engine="dbt_trino",
            inlets=custom_inlets,
            outlets=custom_outlets,
            **kwargs
        )

        resolved_proj = self._resolve_dbt_project_dir(project_dir)
        self.select_models = select_models
        self.project_dir = resolved_proj
        self.profiles_dir = profiles_dir or resolved_proj
        self.target = target or os.getenv("DBT_DEFAULT_TARGET", "dev")
        self.vars = vars or {}

    @staticmethod
    def _resolve_dbt_project_dir(preferred: Optional[str] = None) -> str:
        """Resolve dbt project directory dynamically from Git DAG Bundle or system path."""
        if preferred and os.path.exists(os.path.join(preferred, "dbt_project.yml")):
            return preferred
        
        candidates = [
            os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "dbt")),
            "/opt/airflow/dag_bundles/merica-pipeline/tracking_repo/dbt",
            "/opt/airflow/dag_bundles/merica-pipeline/dbt",
            "/opt/airflow/pipeline-bundles/dbt",
            "/opt/airflow/dag_bundles/dlh-pipeline-bundle/dbt",
            "/opt/airflow/dag_bundles/merica-pipeline-bundles/dbt",
            os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "merica-pipeline", "dbt")),
            os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "merica-pipeline-bundles", "dbt")),
            "/opt/airflow/dbt/dlh_transformation",
            "/opt/airflow/dbt"
        ]
        for c in candidates:
            if os.path.exists(os.path.join(c, "dbt_project.yml")):
                return c
        return preferred or candidates[0]

    def execute(self, context: Context) -> Dict[str, Any]:
        """
        Execute dbt run command using subprocess and capture execution status.
        """
        logger.info(
            "Executing dbt model '%s' (Project: %s, Target: %s, Layer: %s)",
            self.select_models, self.project_dir, self.target, self.layer
        )

        # Ensure writable log and target directories for read-only bundle mounts
        log_dir = "/tmp/dbt_logs"
        target_dir = "/tmp/dbt_target"
        os.makedirs(log_dir, exist_ok=True)
        os.makedirs(target_dir, exist_ok=True)

        cmd = [
            "dbt", "run",
            "--select", self.select_models,
            "--project-dir", self.project_dir,
            "--profiles-dir", self.profiles_dir,
            "--log-path", log_dir,
            "--target-path", target_dir,
        ]

        if self.target:
            cmd.extend(["--target", self.target])

        if self.vars:
            import json
            cmd.extend(["--vars", json.dumps(self.vars)])

        cmd_str = " ".join(cmd)
        logger.info("Running dbt command: %s", cmd_str)

        env = os.environ.copy()
        if not env.get("POSTGRES_HOST") or env.get("POSTGRES_HOST") in ("localhost", "127.0.0.1"):
            env["POSTGRES_HOST"] = "merica-postgres"

        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                env=env,
                check=False
            )

            stdout = result.stdout
            stderr = result.stderr

            if result.returncode != 0:
                logger.error("dbt run failed with returncode %d:\n%s\n%s", result.returncode, stdout, stderr)
                raise RuntimeError(f"dbt run failed for model {self.select_models}: {stderr or stdout}")

            logger.info("✅ dbt model '%s' executed successfully!\n%s", self.select_models, stdout[-500:])

            return {
                "status": "success",
                "engine": "dbt_trino",
                "model": self.select_models,
                "layer": self.layer,
                "returncode": result.returncode
            }
        except Exception as e:
            logger.error("DbtTrinoIcebergOperator execution failed: %s", e)
            raise