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

This commit is contained in:
2026-10-19 17:58:27 +07:00
parent 56a2c8916b
commit fa0d562eb1
2 changed files with 41 additions and 6 deletions
+40 -5
View File
@@ -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 []