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..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,6 +70,7 @@ class DbtRunAirflowAsyncBigqueryOperator(BigQueryInsertJobOperator, AbstractDbtL template_fields_renderers = { "compiled_sql": "sql", } + producer_task_id: str = DEFAULT_PRODUCER_ASYNC_TASK_ID def __init__( self, @@ -86,6 +88,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 +140,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..3ef1642299 100644 --- a/dev/dags/simple_dag_async.py +++ b/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] 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]"