From 2730aa9d31f458cd5c7cbf79160befe690b220fe Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Fri, 7 Nov 2025 14:16:40 +0000 Subject: [PATCH 1/4] Fix ExecutionMode.AIRFLOW_ASYNC TaskGroup XCom issue --- cosmos/airflow/graph.py | 15 +++- cosmos/operators/_asynchronous/bigquery.py | 6 +- dev/dags/simple_dag_async.py | 85 ++++++++++++++++------ 3 files changed, 77 insertions(+), 29 deletions(-) diff --git a/cosmos/airflow/graph.py b/cosmos/airflow/graph.py index 6a9c14a390..fa9f213cb2 100644 --- a/cosmos/airflow/graph.py +++ b/cosmos/airflow/graph.py @@ -657,9 +657,16 @@ def _add_dbt_setup_async_task( ) setup_airflow_task = create_airflow_task(setup_task_metadata, dag, task_group=task_group) - for task_id, task in tasks_map.items(): - if not task.upstream_list: - setup_airflow_task >> task + for node_id, task_or_taskgroup in tasks_map.items(): + node_tasks = ( + list(task_or_taskgroup.children.values()) + if isinstance(task_or_taskgroup, TaskGroup) + else [task_or_taskgroup] + ) + for task in node_tasks: + task.producer_task_id = setup_airflow_task.task_id # type: ignore[attr-defined] + if not task.upstream_list: + setup_airflow_task >> task tasks_map[DBT_SETUP_ASYNC_TASK_ID] = setup_airflow_task @@ -979,7 +986,7 @@ def build_airflow_graph( # noqa: C901 TODO: https://github.com/astronomer/astro execution_mode, {**task_args, "virtualenv_dir": virtualenv_dir}, tasks_map, - task_group, + task_group=task_group, render_config=render_config, async_py_requirements=async_py_requirements, ) diff --git a/cosmos/operators/_asynchronous/bigquery.py b/cosmos/operators/_asynchronous/bigquery.py index 2e9bae55f4..ea4a9ade67 100644 --- a/cosmos/operators/_asynchronous/bigquery.py +++ b/cosmos/operators/_asynchronous/bigquery.py @@ -69,6 +69,7 @@ class DbtRunAirflowAsyncBigqueryOperator(BigQueryInsertJobOperator, AbstractDbtL template_fields_renderers = { "compiled_sql": "sql", } + producer_task_id: str = "dbt_setup_async" def __init__( self, @@ -86,6 +87,7 @@ def __init__( self.dbt_kwargs = dbt_kwargs or {} task_id = self.dbt_kwargs.pop("task_id") self.full_refresh = self.dbt_kwargs.pop("full_refresh", False) + AbstractDbtLocalBase.__init__( self, task_id=task_id, project_dir=project_dir, profile_config=profile_config, **self.dbt_kwargs ) @@ -137,7 +139,9 @@ def get_sql_from_xcom(self, context: Context) -> str: file_path = self.async_context["dbt_node_config"]["file_path"] project_dir_parent = str(Path(self.project_dir).parent) sql_model_path = str(file_path).replace(project_dir_parent, "").lstrip("/") - compressed_b64_sql = context["ti"].xcom_pull(task_ids="dbt_setup_async", key=_sanitize_xcom_key(sql_model_path)) + compressed_b64_sql = context["ti"].xcom_pull( + task_ids=self.producer_task_id, key=_sanitize_xcom_key(sql_model_path) + ) compressed_b64_sql = base64.b64decode(compressed_b64_sql) sql_query = zlib.decompress(compressed_b64_sql).decode("utf-8") diff --git a/dev/dags/simple_dag_async.py b/dev/dags/simple_dag_async.py index 0cbf629cd1..3a505a5757 100644 --- a/dev/dags/simple_dag_async.py +++ b/dev/dags/simple_dag_async.py @@ -2,8 +2,7 @@ from datetime import datetime from pathlib import Path -from cosmos import DbtDag, ExecutionConfig, ExecutionMode, ProfileConfig, ProjectConfig, RenderConfig -from cosmos.constants import TestBehavior +from cosmos import ExecutionConfig, ExecutionMode, ProfileConfig, ProjectConfig, RenderConfig from cosmos.profiles import GoogleCloudServiceAccountDictProfileMapping DEFAULT_DBT_ROOT_PATH = Path(__file__).resolve().parent / "dbt" @@ -23,27 +22,65 @@ # [START airflow_async_execution_mode_example] -simple_dag_async = DbtDag( - # dbt/cosmos-specific parameters - project_config=ProjectConfig( - DBT_PROJECT_PATH, - ), - profile_config=profile_config, - execution_config=ExecutionConfig( - execution_mode=ExecutionMode.AIRFLOW_ASYNC, - async_py_requirements=[f"dbt-bigquery=={DBT_ADAPTER_VERSION}"], - ), - render_config=RenderConfig(select=["path:models"], test_behavior=TestBehavior.NONE), - # normal dag parameters - schedule=None, +# simple_dag_async = DbtDag( +# # dbt/cosmos-specific parameters +# project_config=ProjectConfig( +# DBT_PROJECT_PATH, +# ), +# profile_config=profile_config, +# execution_config=ExecutionConfig( +# execution_mode=ExecutionMode.AIRFLOW_ASYNC, +# async_py_requirements=[f"dbt-bigquery=={DBT_ADAPTER_VERSION}"], +# ), +# render_config=RenderConfig(select=["path:models"], test_behavior=TestBehavior.NONE), +# # normal dag parameters +# schedule=None, +# start_date=datetime(2023, 1, 1), +# catchup=False, +# dag_id="simple_dag_async", +# tags=["simple"], +# operator_args={ +# "location": "US", +# "install_deps": True, +# "full_refresh": True, +# }, +# ) +# [END airflow_async_execution_mode_example] + + +from airflow.models import DAG + +try: + from airflow.providers.standard.operators.empty import EmptyOperator +except ImportError: + from airflow.operators.empty import EmptyOperator + +from cosmos import DbtTaskGroup + +# [START simple_dag_async_taskgroup] +with DAG( + dag_id="simple_dag_async_taskgroup", + schedule="@daily", start_date=datetime(2023, 1, 1), catchup=False, - dag_id="simple_dag_async", - tags=["simple"], - operator_args={ - "location": "US", - "install_deps": True, - "full_refresh": True, - }, -) -# [END airflow_async_execution_mode_example] +): + pre_dbt = EmptyOperator(task_id="pre_dbt") + + first_dbt_task_group = DbtTaskGroup( + group_id="first_dbt_task_group", + execution_config=ExecutionConfig( + execution_mode=ExecutionMode.AIRFLOW_ASYNC, + async_py_requirements=[f"dbt-bigquery=={DBT_ADAPTER_VERSION}"], + ), + render_config=RenderConfig(select=["*customers*"], exclude=["path:seeds"]), + project_config=ProjectConfig(DBT_PROJECT_PATH), + profile_config=profile_config, + operator_args={ + "location": "US", + "install_deps": True, + "full_refresh": True, + }, + ) + + pre_dbt >> first_dbt_task_group +# [END simple_dag_async_taskgroup] From afdbc86d2cdadb60b6beea420aa2779d5cf12f8a Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Fri, 7 Nov 2025 14:35:16 +0000 Subject: [PATCH 2/4] Restore async DAG --- dev/dags/simple_dag_async.py | 49 ++++++++++++++++++------------------ 1 file changed, 25 insertions(+), 24 deletions(-) diff --git a/dev/dags/simple_dag_async.py b/dev/dags/simple_dag_async.py index 3a505a5757..3ef1642299 100644 --- a/dev/dags/simple_dag_async.py +++ b/dev/dags/simple_dag_async.py @@ -2,7 +2,8 @@ from datetime import datetime from pathlib import Path -from cosmos import ExecutionConfig, ExecutionMode, ProfileConfig, ProjectConfig, RenderConfig +from cosmos import DbtDag, ExecutionConfig, ExecutionMode, ProfileConfig, ProjectConfig, RenderConfig +from cosmos.constants import TestBehavior from cosmos.profiles import GoogleCloudServiceAccountDictProfileMapping DEFAULT_DBT_ROOT_PATH = Path(__file__).resolve().parent / "dbt" @@ -22,29 +23,29 @@ # [START airflow_async_execution_mode_example] -# simple_dag_async = DbtDag( -# # dbt/cosmos-specific parameters -# project_config=ProjectConfig( -# DBT_PROJECT_PATH, -# ), -# profile_config=profile_config, -# execution_config=ExecutionConfig( -# execution_mode=ExecutionMode.AIRFLOW_ASYNC, -# async_py_requirements=[f"dbt-bigquery=={DBT_ADAPTER_VERSION}"], -# ), -# render_config=RenderConfig(select=["path:models"], test_behavior=TestBehavior.NONE), -# # normal dag parameters -# schedule=None, -# start_date=datetime(2023, 1, 1), -# catchup=False, -# dag_id="simple_dag_async", -# tags=["simple"], -# operator_args={ -# "location": "US", -# "install_deps": True, -# "full_refresh": True, -# }, -# ) +simple_dag_async = DbtDag( + # dbt/cosmos-specific parameters + project_config=ProjectConfig( + DBT_PROJECT_PATH, + ), + profile_config=profile_config, + execution_config=ExecutionConfig( + execution_mode=ExecutionMode.AIRFLOW_ASYNC, + async_py_requirements=[f"dbt-bigquery=={DBT_ADAPTER_VERSION}"], + ), + render_config=RenderConfig(select=["path:models"], test_behavior=TestBehavior.NONE), + # normal dag parameters + schedule=None, + start_date=datetime(2023, 1, 1), + catchup=False, + dag_id="simple_dag_async", + tags=["simple"], + operator_args={ + "location": "US", + "install_deps": True, + "full_refresh": True, + }, +) # [END airflow_async_execution_mode_example] From 51442327eceeb117fe3a4b5fb57b0eae80c78ebd Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Fri, 7 Nov 2025 14:52:17 +0000 Subject: [PATCH 3/4] Enable running simple_dag_async_taskgroup as a test DAG in the CI --- scripts/test/integration-dbt-async.sh | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/scripts/test/integration-dbt-async.sh b/scripts/test/integration-dbt-async.sh index 7b1c404f19..2e663ca58b 100644 --- a/scripts/test/integration-dbt-async.sh +++ b/scripts/test/integration-dbt-async.sh @@ -68,3 +68,9 @@ pytest -vv \ --cov-report=term-missing \ --cov-report=xml \ "tests/test_async_example_dag.py::test_example_dag[simple_dag_async]" + +pytest -vv \ + --cov=cosmos \ + --cov-report=term-missing \ + --cov-report=xml \ + "tests/test_async_example_dag.py::test_example_dag[simple_dag_async_taskgroup]" From 3bddd23a71f6c2e55da1cbb503004d14aa423aa8 Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Fri, 7 Nov 2025 14:55:43 +0000 Subject: [PATCH 4/4] Address copilot feedback --- cosmos/operators/_asynchronous/bigquery.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cosmos/operators/_asynchronous/bigquery.py b/cosmos/operators/_asynchronous/bigquery.py index ea4a9ade67..56c15b7fcb 100644 --- a/cosmos/operators/_asynchronous/bigquery.py +++ b/cosmos/operators/_asynchronous/bigquery.py @@ -29,6 +29,7 @@ from cosmos.settings import remote_target_path, remote_target_path_conn_id AIRFLOW_VERSION = Version(airflow.__version__) +DEFAULT_PRODUCER_ASYNC_TASK_ID = "dbt_setup_async" def _mock_bigquery_adapter() -> None: @@ -69,7 +70,7 @@ class DbtRunAirflowAsyncBigqueryOperator(BigQueryInsertJobOperator, AbstractDbtL template_fields_renderers = { "compiled_sql": "sql", } - producer_task_id: str = "dbt_setup_async" + producer_task_id: str = DEFAULT_PRODUCER_ASYNC_TASK_ID def __init__( self,