Configure Airflow tasks using dbt model meta

cosmos/airflow/graph.py

@@ -98,6 +98,7 @@ def create_test_task_metadata(
     extra_context = {}
 
     task_owner = ""
+    airflow_task_config = {}
     if test_indirect_selection != TestIndirectSelection.EAGER:
         task_args["indirect_selection"] = test_indirect_selection.value
     if node is not None:
@@ -110,6 +111,7 @@ def create_test_task_metadata(
 
         extra_context = {"dbt_node_config": node.context_dict}
         task_owner = node.owner
+        airflow_task_config = node.airflow_task_config
 
     elif render_config is not None:  # TestBehavior.AFTER_ALL
         task_args["select"] = render_config.select
@@ -119,6 +121,7 @@ def create_test_task_metadata(
     return TaskMetadata(
         id=test_task_name,
         owner=task_owner,
+        airflow_task_config=airflow_task_config,
         operator_class=calculate_operator_class(
             execution_mode=execution_mode,
             dbt_class="DbtTest",
@@ -191,6 +194,7 @@ def create_task_metadata(
         task_metadata = TaskMetadata(
             id=task_id,
             owner=node.owner,
+            airflow_task_config=node.airflow_task_config,
             operator_class=calculate_operator_class(
                 execution_mode=execution_mode, dbt_class=dbt_resource_to_class[node.resource_type]
             ),

cosmos/core/airflow.py

@@ -32,6 +32,9 @@
     if task.owner != "":
         task_kwargs["owner"] = task.owner
 
+    for k, v in task.airflow_task_config.items():
+        task_kwargs[k] = v
+
     airflow_task = Operator(
         task_id=task.id,
         dag=dag,

cosmos/core/graph/entities.py

@@ -58,6 +58,7 @@ class Task(CosmosEntity):
     """
 
     owner: str = ""
+    airflow_task_config: Dict[str, Any] = field(default_factory=dict)
     operator_class: str = "airflow.operators.empty.EmptyOperator"
     arguments: Dict[str, Any] = field(default_factory=dict)
     extra_context: Dict[str, Any] = field(default_factory=dict)

cosmos/dbt/graph.py

@@ -13,7 +13,7 @@
 from functools import cached_property
 from pathlib import Path
 from subprocess import PIPE, Popen
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
 from airflow.models import Variable
 
@@ -67,6 +67,22 @@
     has_freshness: bool = False
     has_test: bool = False
 
+    @property
+    def airflow_task_config(self) -> Dict[str, Any]:
+        """
+        This method is designed to extend the dbt project's functionality by incorporating Airflow-related metadata into the dbt YAML configuration.
+        Since dbt projects are independent of Airflow, adding Airflow-specific information to the `meta` field within the dbt YAML allows Airflow tasks to
+        utilize this information during execution.
+
+        Examples: pool, pool_slots, queue, ...
+        Returns:
+            Dict[str, Any]: A dictionary containing custom metadata configurations for integration with Airflow.
+        """
+        airflow_task_config = self.config.get("meta", {}).get("cosmos")
+        if isinstance(airflow_task_config, dict):
+            return airflow_task_config
+        return {}
+
     @property
     def resource_name(self) -> str:
         """