feat(cosmos): integrate DbtTaskGroup into dag_factory and fix cems model ref
This commit is contained in:
@@ -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,6 +528,34 @@ 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')
|
||||
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,
|
||||
|
||||
@@ -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'
|
||||
Reference in New Issue
Block a user