Fix ExecutionMode.AIRFLOW_ASYNC TaskGroup XCom issue

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,
         )

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")
 

dev/dags/simple_dag_async.py

@@ -47,3 +47,41 @@
     },
 )
 # [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,
+):
+    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]