diff --git a/dags/dag_factory.py b/dags/dag_factory.py index 1bfabdd..106ad72 100644 --- a/dags/dag_factory.py +++ b/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( - task_id=t_id, - select_models=select_m, - layer=layer, - ) + use_cosmos = config.get('use_cosmos', False) or (config.get('transformation_engine') == 'cosmos') + if use_cosmos and COSMOS_AVAILABLE: + dbt_candidate = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "dbt")) + 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 [] diff --git a/dbt/models/silver/silver_cems_emission_enriched.sql b/dbt/models/silver/silver_cems_emission_enriched.sql index 793f874..69f3eb7 100644 --- a/dbt/models/silver/silver_cems_emission_enriched.sql +++ b/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' \ No newline at end of file