feat(cosmos): integrate DbtTaskGroup into dag_factory and fix cems model ref

dags/dag_factory.py

@@ -27,6 +27,13 @@ try:
 except ImportError:
     MericaAlertHook = None
 
+try:
+    from cosmos import DbtTaskGroup, ProjectConfig, ProfileConfig, RenderConfig, ExecutionConfig
+    from cosmos.constants import TestBehavior, ExecutionMode
+    COSMOS_AVAILABLE = True
+except ImportError:
+    COSMOS_AVAILABLE = False
+
 import re
 import time
 import json
@@ -521,11 +528,39 @@ def build_dag_from_yaml(config_file: str) -> DAG:
                 elif op_type in ('dbt', 'dbt_transform', 'dbt_trino'):
                     select_m = config.get('select_models') or config.get('dbt_model', '')
                     layer = config.get('dbt_layer') or config.get('layer', 'silver')
-                    task_obj = DbtTrinoIcebergOperator(
+                    use_cosmos = config.get('use_cosmos', False) or (config.get('transformation_engine') == 'cosmos')
-                        task_id=t_id,
+                    if use_cosmos and COSMOS_AVAILABLE:
-                        select_models=select_m,
+                        dbt_candidate = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "dbt"))
-                        layer=layer,
+                        dbt_proj_dir = os.getenv("DBT_PROJECT_DIR") or (dbt_candidate if os.path.exists(dbt_candidate) else "/opt/airflow/dbt")
-                    )
+                        try:
+                            select_args = [select_m] if select_m else ([f"path:models/{layer}"] if layer else None)
+                            task_obj = DbtTaskGroup(
+                                group_id=t_id,
+                                project_config=ProjectConfig(dbt_project_path=dbt_proj_dir),
+                                profile_config=ProfileConfig(
+                                    profile_name=os.getenv("DBT_PROFILE_NAME", "dlh_postgres"),
+                                    target_name=os.getenv("DBT_DEFAULT_TARGET", "dev"),
+                                    profiles_yml_filepath=f"{dbt_proj_dir}/profiles.yml"
+                                ),
+                                render_config=RenderConfig(
+                                    test_behavior=TestBehavior.AFTER_EACH,
+                                    select=select_args
+                                ),
+                                execution_config=ExecutionConfig(execution_mode=ExecutionMode.LOCAL)
+                            )
+                        except Exception as e:
+                            print(f"[COSMOS_FALLBACK] Error creating DbtTaskGroup for {t_id}: {e}")
+                            task_obj = DbtTrinoIcebergOperator(
+                                task_id=t_id,
+                                select_models=select_m,
+                                layer=layer,
+                            )
+                    else:
+                        task_obj = DbtTrinoIcebergOperator(
+                            task_id=t_id,
+                            select_models=select_m,
+                            layer=layer,
+                        )
                 elif op_type in ('great_expectations', 'ge_quality', 'data_quality'):
                     rules = config.get('rules') or []
                     queries = config.get('queries') or []

dbt/models/silver/silver_cems_emission_enriched.sql

@@ -25,7 +25,7 @@ SELECT
     ELSE 'COMPLIANT'
   END AS compliance_status,
   t.sensor_status
-FROM {{ ref('bronze_cems_telemetry_logs') }} t
+FROM {{ ref('bronze_raw_cems_telemetry_logs') }} t
 JOIN {{ ref('bronze_ref_emission_sources') }} s
   ON t.source_id = s.source_id
 WHERE t.sensor_status = 'VALID'