feat(cosmos): integrate DbtTaskGroup into dag_factory and fix cems model ref
This commit is contained in:
+40
-5
@@ -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 []
|
||||
|
||||
Reference in New Issue
Block a user