Refactor airflow/graph.py to simplify code-base

CHANGELOG.rst

@@ -1,6 +1,30 @@
 Changelog
 =========
 
+1.12.0a2 (2025-11-04)
+-------------------
+
+Breaking changes
+
+* Introduced in the PR #2080. The following functions are expected to be used internally only to Cosmos, so we hope these won't impact end-users, but we are documenting the changes just in case:
+  - ``generate_task_or_group`` receives ``render_config`` instead of its individual configurations, such as ``test_behavior``, ``source_rendering_behavior`` and ``enable_owner_inheritance``
+  - ``create_task_metadata`` receives ``render_config`` instead of its individual configurations, such as ``test_behavior``, ``source_rendering_behavior`` and ``enable_owner_inheritance``
+  - ``create_task_metadata`` now expects the ``node_converters`` argument
+
+Features
+
+* Support applying ``node_converter`` at a task level instead of task group level by @anyapriya in #1759
+
+Enhancements
+
+* Remove usage of contextmanager in plugins for accessing connections in Airflow >= 3.1.2 by @pankajkoti in #2073
+* Refactor ``airflow/graph.py`` to simplify code-base by @tatiana in #2080
+
+Others
+
+* Fix broken CI due to fastapi incompatibility with cadwyn for Airflow 3 by @pankajkoti in #2076
+* Pre-commit updates: #2078
+
 1.11.0 (2025-10-29)
 ---------------------
 

cosmos/__init__.py

@@ -9,7 +9,7 @@
 
 from cosmos import settings
 
-__version__ = "1.11.0"
+__version__ = "1.12.0a2"
 
 if not settings.enable_memory_optimised_imports:
     from cosmos.airflow.dag import DbtDag

cosmos/config.py

@@ -64,6 +64,7 @@ class RenderConfig:
     :param selector: Name of a dbt YAML selector to use for parsing. Only supported when using ``load_method=LoadMode.DBT_LS``.
     :param dbt_deps: (deprecated) Configure to run dbt deps when using dbt ls for dag parsing
     :param node_converters: a dictionary mapping a ``DbtResourceType`` into a callable. Users can control how to render dbt nodes in Airflow. Only supported when using ``load_method=LoadMode.DBT_MANIFEST`` or ``LoadMode.DBT_LS``.
+    :param node_conversion_by_task_group: A boolean that allows users to do node conversion at the task group level instead of task level.  Defaults to True.
     :param dbt_executable_path: The path to the dbt executable for dag generation. Defaults to dbt if available on the path.
     :param env_vars: (Deprecated since Cosmos 1.3 use ProjectConfig.env_vars) A dictionary of environment variables for rendering. Only supported when using ``LoadMode.DBT_LS``.
     :param dbt_project_path: Configures the DBT project location accessible on the airflow controller for DAG rendering. Mutually Exclusive with ProjectConfig.dbt_project_path. Required when using ``load_method=LoadMode.DBT_LS`` or ``load_method=LoadMode.CUSTOM``.
@@ -76,7 +77,6 @@ class RenderConfig:
     :param normalize_task_display_name: A callable that takes a dbt node as input and returns the task display name. This allows users to assign a custom task display name separate from the node ID.
     :param should_detach_multiple_parents_tests: A boolean that allows users to decide whether to run tests with multiple parent dependencies in separate tasks.
     :param enable_owner_inheritance: A boolean that allows users to enable the owner inheritance from dbt models to airflow tasks. Defaults to True.
-    :param node_conversion_by_task_group: A boolean that allows users to do node conversion at the task group level instead of task level.  Defaults to True.
     """
 
     emit_datasets: bool = True
@@ -88,6 +88,7 @@ class RenderConfig:
     selector: str | None = None
     dbt_deps: bool | None = None
     node_converters: dict[DbtResourceType, Callable[..., Any]] | None = None
+    node_conversion_by_task_group: bool | None = True
     dbt_executable_path: str | Path = get_system_dbt()
     env_vars: dict[str, str] | None = None
     dbt_project_path: InitVar[str | Path | None] = None
@@ -101,7 +102,6 @@ class RenderConfig:
     normalize_task_display_name: Callable[..., Any] | None = None
     should_detach_multiple_parents_tests: bool = False
     enable_owner_inheritance: bool | None = True
-    node_conversion_by_task_group: bool | None = True
 
     def __post_init__(self, dbt_project_path: str | Path | None) -> None:
         if self.env_vars:

docs/configuration/render-config.rst

@@ -15,6 +15,7 @@ The ``RenderConfig`` class takes the following arguments:
 - ``selector``: (new in v1.3) name of a dbt YAML selector to use for DAG parsing. Only supported when using ``load_method=LoadMode.DBT_LS``. See `Selecting & Excluding <selecting-excluding.html>`_ for more information.
 - ``dbt_deps``: (deprecated in v1.9, use ``ProjectConfig.install_dbt_deps`` onwards) A Boolean to run dbt deps when using dbt ls for dag parsing. Default True
 - ``node_converters``: a dictionary mapping a ``DbtResourceType`` into a callable. Users can control how to render dbt nodes in Airflow. Only supported when using ``load_method=LoadMode.DBT_MANIFEST`` or ``LoadMode.DBT_LS``. Find more information below.
+- ``node_conversion_by_task_group``: (new in v1.12.0) A boolean to control if node_converters are used at the task group level (ex. converting models with test_behavior=AFTER_EACH means the entire task group is converted including the run task and the test task), or the individual task level (gives more granularity for converting just the run tasks or just the test tasks).  Defaults to True.
 - ``dbt_executable_path``: The path to the dbt executable for dag generation. Defaults to dbt if available on the path.
 - ``dbt_ls_path``:  Should be set when using ``load_method=LoadMode.DBT_LS_OUTPUT``. Path of the user-managed output of ``dbt ls``.
 - ``enable_mock_profile``: When using ``LoadMode.DBT_LS`` with a ``ProfileMapping`` class, by default, Cosmos mocks the values of the profile. Defaults to True. In order to leverage partial parsing, this argument should be set to ``False``. Read `Partial parsing <./partial-parsing.html#profile-configuration.html>`_ for more information.
@@ -50,8 +51,10 @@ Users may opt to use ``InvocationMode.SUBPROCESS`` when they have multiple Pytho
 and do not want Cosmos to use the dbt version installed in the same Python Virtualenv as Airflow to parse the DAG.
 
 
-Customizing how nodes are rendered (experimental)
--------------------------------------------------
+Customizing how nodes are rendered
+----------------------------------
+
+.. versionadded:: 1.2.0
 
 There are circumstances when choosing specific Airflow operators to represent a dbt node is helpful.
 An example could be to use an S3 sensor to represent dbt sources or to create custom operators to handle exposures.
@@ -65,3 +68,7 @@ The following example illustrates how it is possible to tell Cosmos how to conve
     :end-before: [END custom_dbt_nodes]
 
 When defining the mapping for a new type that is not part of Cosmos' ``DbtResourceType`` enumeration, users should use the syntax ``DbtResourceType("new-node-type")`` as opposed to ``DbtResourceType.EXISTING_TYPE``. It will dynamically add the new type to the enumeration ``DbtResourceType`` so that Cosmos can parse these dbt nodes and convert them into the Airflow DAG.
+
+In Cosmos 1.12.0, this feature was further improved by adding the ``RenderConfig.node_conversion_by_task_group`` parameter.
+This parameter allows users to control if node_converters are used at the task group level (ex. converting models with test_behavior=AFTER_EACH means the entire task group is converted including the run task and the test task),
+or the individual task level (gives more granularity for converting just the run tasks or just the test tasks).  Defaults to True.

tests/airflow/test_graph.py

@@ -231,9 +231,11 @@ def test_create_task_group_for_after_each_supported_nodes(node_type: DbtResource
         },
         dbt_project_name="astro_shop",
         node_converters={},
-        test_behavior=TestBehavior.AFTER_EACH,
+        render_config=RenderConfig(
+            test_behavior=TestBehavior.AFTER_EACH,
+            source_rendering_behavior=SOURCE_RENDERING_BEHAVIOR,
+        ),
         on_warning_callback=None,
-        source_rendering_behavior=SOURCE_RENDERING_BEHAVIOR,
     )
     assert isinstance(output, TaskGroup)
     assert list(output.children.keys()) == [f"dbt_node.{task_suffix}", "dbt_node.test"]
@@ -615,7 +617,9 @@ def test_create_task_metadata_source_with_rendering_options(
     metadata = create_task_metadata(
         child_node,
         execution_mode=ExecutionMode.LOCAL,
-        source_rendering_behavior=source_rendering_behavior,
+        render_config=RenderConfig(
+            source_rendering_behavior=source_rendering_behavior,
+        ),
         args={},
         dbt_dag_task_group_identifier="",
     )
@@ -922,10 +926,12 @@ def test_create_task_metadata_normalize_task_id(
         args=args,
         dbt_dag_task_group_identifier="",
         use_task_group=use_task_group,
-        normalize_task_id=normalize_task_id,
-        normalize_task_display_name=normalize_task_display_name,
-        source_rendering_behavior=SourceRenderingBehavior.ALL,
-        test_behavior=test_behavior,
+        render_config=RenderConfig(
+            normalize_task_id=normalize_task_id,
+            normalize_task_display_name=normalize_task_display_name,
+            source_rendering_behavior=SourceRenderingBehavior.ALL,
+            test_behavior=test_behavior,
+        ),
     )
     assert metadata.id == expected_node_id
     if expected_display_name:
@@ -1105,10 +1111,12 @@ def test_owner(dbt_extra_config, expected_owner):
         },
         dbt_project_name="astro_shop",
         node_converters={},
-        test_behavior=TestBehavior.AFTER_EACH,
+        render_config=RenderConfig(
+            test_behavior=TestBehavior.AFTER_EACH,
+            source_rendering_behavior=SOURCE_RENDERING_BEHAVIOR,
+            enable_owner_inheritance=True,
+        ),
         on_warning_callback=None,
-        source_rendering_behavior=SOURCE_RENDERING_BEHAVIOR,
-        enable_owner_inheritance=True,
     )
 
     assert len(output.leaves) == 1
@@ -1232,7 +1240,9 @@ def test_create_task_metadata_disable_owner_inheritance(enable_owner_inheritance
         execution_mode=ExecutionMode.LOCAL,
         args={"project_dir": SAMPLE_PROJ_PATH},
         dbt_dag_task_group_identifier="test_dag",
-        enable_owner_inheritance=enable_owner_inheritance,
+        render_config=RenderConfig(
+            enable_owner_inheritance=enable_owner_inheritance,
+        ),
     )
 
     assert task_metadata is not None
@@ -1322,11 +1332,14 @@ def test_generate_task_or_group_disable_owner_inheritance(enable_owner_inheritan
                     ),
                 ),
             },
-            test_behavior=TestBehavior.NONE,
-            source_rendering_behavior=SOURCE_RENDERING_BEHAVIOR,
+            render_config=RenderConfig(
+                test_behavior=TestBehavior.NONE,
+                source_rendering_behavior=SOURCE_RENDERING_BEHAVIOR,
+                enable_owner_inheritance=enable_owner_inheritance,
+            ),
             test_indirect_selection=TestIndirectSelection.EAGER,
             on_warning_callback=None,
-            enable_owner_inheritance=enable_owner_inheritance,
+            node_converters={},
         )
 
         assert task_or_group is not None

tests/dbt/test_pruning.py

@@ -13,7 +13,7 @@
     create_task_metadata,
     generate_task_or_group,
 )
-from cosmos.config import ProfileConfig
+from cosmos.config import ProfileConfig, RenderConfig
 from cosmos.constants import (
     DbtResourceType,
     ExecutionMode,
@@ -224,9 +224,11 @@ def test_create_task_metadata_source_pruning_disabled(self):
             node=self.source_with_downstream,
             execution_mode=ExecutionMode.LOCAL,
             args=self.task_args,
-            dbt_dag_task_group_identifier="test",
-            source_rendering_behavior=SourceRenderingBehavior.ALL,
-            source_pruning=False,  # Disabled
+            dbt_dag_task_group_identifier="tes  t",
+            render_config=RenderConfig(
+                source_rendering_behavior=SourceRenderingBehavior.ALL,
+                source_pruning=False,  # Disabled
+            ),
             filtered_nodes=self.filtered_nodes_with_model,
         )
 
@@ -235,8 +237,10 @@ def test_create_task_metadata_source_pruning_disabled(self):
             execution_mode=ExecutionMode.LOCAL,
             args=self.task_args,
             dbt_dag_task_group_identifier="test",
-            source_rendering_behavior=SourceRenderingBehavior.ALL,
-            source_pruning=False,  # Disabled
+            render_config=RenderConfig(
+                source_rendering_behavior=SourceRenderingBehavior.ALL,
+                source_pruning=False,  # Disabled
+            ),
             filtered_nodes=self.filtered_nodes_empty,
         )
 
@@ -254,8 +258,10 @@ def test_create_task_metadata_source_pruning_enabled(self):
             execution_mode=ExecutionMode.LOCAL,
             args=self.task_args,
             dbt_dag_task_group_identifier="test",
-            source_rendering_behavior=SourceRenderingBehavior.ALL,
-            source_pruning=True,  # Enabled
+            render_config=RenderConfig(
+                source_rendering_behavior=SourceRenderingBehavior.ALL,
+                source_pruning=True,  # Enabled
+            ),
             filtered_nodes=self.filtered_nodes_with_model,
         )
 
@@ -265,8 +271,10 @@ def test_create_task_metadata_source_pruning_enabled(self):
             execution_mode=ExecutionMode.LOCAL,
             args=self.task_args,
             dbt_dag_task_group_identifier="test",
-            source_rendering_behavior=SourceRenderingBehavior.ALL,
-            source_pruning=True,  # Enabled
+            render_config=RenderConfig(
+                source_rendering_behavior=SourceRenderingBehavior.ALL,
+                source_pruning=True,  # Enabled
+            ),
             filtered_nodes=self.filtered_nodes_with_model,
         )
 
@@ -282,8 +290,10 @@ def test_create_task_metadata_source_pruning_edge_cases(self):
             execution_mode=ExecutionMode.LOCAL,
             args=self.task_args,
             dbt_dag_task_group_identifier="test",
-            source_rendering_behavior=SourceRenderingBehavior.ALL,
-            source_pruning=True,
+            render_config=RenderConfig(
+                source_rendering_behavior=SourceRenderingBehavior.ALL,
+                source_pruning=True,
+            ),
             filtered_nodes=None,  # None means skip pruning check
         )
 
@@ -293,8 +303,10 @@ def test_create_task_metadata_source_pruning_edge_cases(self):
             execution_mode=ExecutionMode.LOCAL,
             args=self.task_args,
             dbt_dag_task_group_identifier="test",
-            source_rendering_behavior=SourceRenderingBehavior.ALL,
-            source_pruning=True,
+            render_config=RenderConfig(
+                source_rendering_behavior=SourceRenderingBehavior.ALL,
+                source_pruning=True,
+            ),
             filtered_nodes={},  # Empty dict means skip pruning check
         )
 
@@ -311,8 +323,10 @@ def test_create_task_metadata_source_pruning_edge_cases(self):
             execution_mode=ExecutionMode.LOCAL,
             args=self.task_args,
             dbt_dag_task_group_identifier="test",
-            source_rendering_behavior=SourceRenderingBehavior.ALL,
-            source_pruning=True,
+            render_config=RenderConfig(
+                source_rendering_behavior=SourceRenderingBehavior.ALL,
+                source_pruning=True,
+            ),
             filtered_nodes={unrelated_model.unique_id: unrelated_model},  # Non-empty but doesn't use this source
         )
 
@@ -333,11 +347,14 @@ def test_generate_task_or_group_source_pruning(self):
             node=self.source_with_downstream,
             execution_mode=ExecutionMode.LOCAL,
             task_args=self.task_args,
-            test_behavior=TestBehavior.NONE,
-            source_rendering_behavior=SourceRenderingBehavior.ALL,
-            source_pruning=True,
+            render_config=RenderConfig(
+                test_behavior=TestBehavior.NONE,
+                source_rendering_behavior=SourceRenderingBehavior.ALL,
+                source_pruning=True,
+            ),
             test_indirect_selection=TestIndirectSelection.EAGER,
             filtered_nodes=self.filtered_nodes_with_model,
+            node_converters={},
         )
 
         # Orphaned source should NOT generate a task
@@ -348,10 +365,13 @@ def test_generate_task_or_group_source_pruning(self):
             execution_mode=ExecutionMode.LOCAL,
             task_args=self.task_args,
             test_behavior=TestBehavior.NONE,
-            source_rendering_behavior=SourceRenderingBehavior.ALL,
-            source_pruning=True,
+            render_config=RenderConfig(
+                source_rendering_behavior=SourceRenderingBehavior.ALL,
+                source_pruning=True,
+            ),
             test_indirect_selection=TestIndirectSelection.EAGER,
             filtered_nodes=self.filtered_nodes_with_model,
+            node_converters={},
         )
 
         assert task_with_downstream is not None
@@ -375,8 +395,10 @@ def test_source_pruning_with_different_source_rendering_behaviors(self):
             execution_mode=ExecutionMode.LOCAL,
             args=self.task_args,
             dbt_dag_task_group_identifier="test",
-            source_rendering_behavior=SourceRenderingBehavior.NONE,
-            source_pruning=True,
+            render_config=RenderConfig(
+                source_rendering_behavior=SourceRenderingBehavior.NONE,
+                source_pruning=True,
+            ),
             filtered_nodes=self.filtered_nodes_with_model,
         )
 
@@ -386,8 +408,10 @@ def test_source_pruning_with_different_source_rendering_behaviors(self):
             execution_mode=ExecutionMode.LOCAL,
             args=self.task_args,
             dbt_dag_task_group_identifier="test",
-            source_rendering_behavior=SourceRenderingBehavior.WITH_TESTS_OR_FRESHNESS,
-            source_pruning=True,
+            render_config=RenderConfig(
+                source_rendering_behavior=SourceRenderingBehavior.WITH_TESTS_OR_FRESHNESS,
+                source_pruning=True,
+            ),
             filtered_nodes=self.filtered_nodes_with_model,
         )
 
@@ -504,11 +528,14 @@ def test_source_pruning_parameter_defaults(self):
             node=source,
             execution_mode=ExecutionMode.LOCAL,
             task_args=task_args,
-            test_behavior=TestBehavior.NONE,
-            source_rendering_behavior=SourceRenderingBehavior.ALL,
+            render_config=RenderConfig(
+                test_behavior=TestBehavior.NONE,
+                source_rendering_behavior=SourceRenderingBehavior.ALL,
+                # source_pruning not specified - should default to False
+            ),
             test_indirect_selection=TestIndirectSelection.EAGER,
-            # source_pruning not specified - should default to False
             # filtered_nodes not specified - should default to None
+            node_converters={},
         )
 
         # Should create a task because pruning is disabled by default