Dag test backcompat

.github/workflows/test.yml

@@ -2,7 +2,7 @@ name: test
 
 on:
   push: # Run on pushes to the default branch
-    branches: [main]
+    branches: [main, af-31, dag-test-backcompat]
   # Also run on pull requests originating from forks. Although this is insecure by default, we need it to run
   # integration tests on forked PRs. As a guardrail, we’ve added an Authorize step to each job, which requires manually
   # approving the workflow run for each pushed commit. Approval only happens after a careful code review of the changes.
@@ -51,8 +51,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
-        airflow-version: ["2.4", "2.5", "2.6", "2.7", "2.8", "2.9", "2.10", "2.11", "3.0"]
+        python-version: ["3.10", "3.11", "3.12"]
+        airflow-version: ["2.4", "2.5", "2.6", "2.7", "2.8", "2.9", "2.10", "2.11", "3.0", "3.1"]
         dbt-version: ["1.10"]
         exclude:
           - python-version: "3.11"
@@ -120,8 +120,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9", "3.10", "3.11"]
-        airflow-version: ["2.4", "2.5", "2.6", "2.7", "2.8", "2.9", "2.10", "2.11", "3.0"]
+        python-version: ["3.10", "3.11"]
+        airflow-version: ["2.4", "2.5", "2.6", "2.7", "2.8", "2.9", "2.10", "2.11", "3.0", "3.1"]
         dbt-version: ["1.10"]
         exclude:
           - python-version: "3.11"
@@ -548,7 +548,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ["3.11"]
-        airflow-version: ["2.10", "3.0" ]
+        airflow-version: ["2.10", "3.0"]
         dbt-version: ["1.9"]
         num-models: [1, 10, 50, 100, 500]
     services:

cosmos/__init__.py

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

cosmos/airflow/dag.py

@@ -6,7 +6,7 @@
 
 from typing import Any
 
-from airflow.models.dag import DAG
+from airflow.models.dag import DAG  # type: ignore[assignment]
 
 from cosmos.converter import DbtToAirflowConverter, airflow_kwargs, specific_kwargs
 

cosmos/airflow/graph.py

@@ -11,7 +11,12 @@
 
 from airflow.models.base import ID_LEN as AIRFLOW_MAX_ID_LENGTH
 from airflow.models.dag import DAG
-from airflow.utils.task_group import TaskGroup
+
+try:
+    # Airflow 3.1 onwards
+    from airflow.sdk import TaskGroup
+except ImportError:
+    from airflow.utils.task_group import TaskGroup
 
 from cosmos import settings
 from cosmos.config import RenderConfig
@@ -541,14 +546,14 @@ def _add_dbt_setup_async_task(
     tasks_map[DBT_SETUP_ASYNC_TASK_ID] = setup_airflow_task
 
 
-def _add_producer_watcher(
+def _add_producer_watcher_and_dependencies(
     dag: DAG,
     task_args: dict[str, Any],
     tasks_map: dict[str, Any],
     task_group: TaskGroup | None,
     render_config: RenderConfig | None = None,
+    nodes: dict[str, DbtNode] | None = None,
 ) -> str:
-
     producer_task_args = task_args.copy()
 
     if render_config is not None:
@@ -562,11 +567,26 @@ def _add_producer_watcher(
         arguments=producer_task_args,
     )
     producer_airflow_task = create_airflow_task(producer_task_metadata, dag, task_group=task_group)
-    for task_id, task in tasks_map.items():
+
+    # Consumer tasks will need to be updated to use the producer task as a dependency
+    for node_id, task_or_taskgroup in tasks_map.items():
         # we want to make the producer task to be the parent of the root dbt nodes, without blocking them from sensing XCom
-        if not task.upstream_list:
-            producer_airflow_task >> task
-            task.trigger_rule = task_args.get("trigger_rule", "always")
+        node_tasks = (
+            list(task_or_taskgroup.children.values())
+            if isinstance(task_or_taskgroup, TaskGroup)
+            else [task_or_taskgroup]
+        )
+
+        # First, we tackle dbt graph nodes that are root nodes
+        if nodes and node_id in nodes and not nodes[node_id].depends_on:
+            producer_airflow_task >> task_or_taskgroup
+            for root_task in node_tasks:
+                if hasattr(root_task, "trigger_rule"):
+                    root_task.trigger_rule = task_args.get("trigger_rule", "always")
+
+        # We also need to set the producer task id too all consumer tasks, regardless if they are root or not
+        for task in node_tasks:
+            task.producer_task_id = producer_airflow_task.task_id  # type: ignore[attr-defined]
 
     tasks_map[PRODUCER_WATCHER_TASK_ID] = producer_airflow_task
     return producer_airflow_task.task_id
@@ -745,16 +765,6 @@ def build_airflow_graph(  # noqa: C901 TODO: https://github.com/astronomer/astro
             logger.debug(f"Conversion of <{node.unique_id}> was successful!")
             tasks_map[node_id] = task_or_group
 
-    if execution_mode == ExecutionMode.WATCHER:
-        producer_watcher_task_id = _add_producer_watcher(
-            dag,
-            task_args,
-            tasks_map,
-            task_group,
-            render_config=render_config,
-        )
-        task_args["producer_watcher_task_id"] = producer_watcher_task_id
-
     # If test_behaviour=="after_all", there will be one test task, run by the end of the DAG
     # The end of a DAG is defined by the DAG leaf tasks (tasks which do not have downstream tasks)
     if test_behavior == TestBehavior.AFTER_ALL:
@@ -790,6 +800,16 @@ def build_airflow_graph(  # noqa: C901 TODO: https://github.com/astronomer/astro
 
     create_airflow_task_dependencies(nodes, tasks_map)
 
+    if execution_mode == ExecutionMode.WATCHER:
+        _add_producer_watcher_and_dependencies(
+            dag=dag,
+            task_args=task_args,
+            tasks_map=tasks_map,
+            task_group=task_group,
+            render_config=render_config,
+            nodes=nodes,
+        )
+
     if settings.enable_setup_async_task:
         _add_dbt_setup_async_task(
             dag,

cosmos/airflow/task_group.py

@@ -6,7 +6,11 @@
 
 from typing import Any
 
-from airflow.utils.task_group import TaskGroup
+try:
+    # Airflow 3.1 onwards
+    from airflow.sdk import TaskGroup
+except ImportError:
+    from airflow.utils.task_group import TaskGroup
 
 from cosmos.converter import DbtToAirflowConverter, airflow_kwargs, specific_kwargs
 

cosmos/cache.py

@@ -17,7 +17,6 @@
 from airflow.models import DagRun, Variable
 from airflow.models.dag import DAG
 from airflow.utils.session import provide_session
-from airflow.utils.task_group import TaskGroup
 from airflow.version import version as airflow_version
 from sqlalchemy import select
 from sqlalchemy.orm import Session
@@ -26,14 +25,16 @@
 
 if TYPE_CHECKING:
     try:
+        # Airflow 3 onwards
         from airflow.sdk import ObjectStoragePath
+        from airflow.utils.task_group import TaskGroup
     except ImportError:
         try:
             from airflow.io.path import ObjectStoragePath
         except ImportError:
             pass
-    except ImportError:
-        pass
+        from airflow.utils.task_group import TaskGroup
+
 from cosmos.constants import (
     DBT_MANIFEST_FILE_NAME,
     DBT_TARGET_DIR_NAME,
@@ -83,7 +84,13 @@ def _configure_remote_cache_dir() -> Path | ObjectStoragePath | None:
             "Airflow 2.8 or later."
         )
 
-    from airflow.io.path import ObjectStoragePath
+    try:
+        from airflow.sdk import ObjectStoragePath
+    except ImportError:
+        try:
+            from airflow.io.path import ObjectStoragePath
+        except ImportError:
+            pass
 
     _configured_cache_dir = ObjectStoragePath(cache_dir_str, conn_id=remote_cache_conn_id)
 

cosmos/config.py

@@ -15,11 +15,12 @@
 
 from cosmos import settings
 
-if TYPE_CHECKING:
+if settings.AIRFLOW_IO_AVAILABLE or TYPE_CHECKING:
     try:
-        from airflow.io.path import ObjectStoragePath
+        from airflow.sdk import ObjectStoragePath
     except ImportError:
-        pass
+        from airflow.io.path import ObjectStoragePath
+
 from cosmos.cache import create_cache_profile, get_cached_profile, is_profile_cache_enabled
 from cosmos.constants import (
     DEFAULT_PROFILES_FILE_NAME,
@@ -236,8 +237,6 @@ def __init__(
                 )
 
             if settings.AIRFLOW_IO_AVAILABLE:
-                from airflow.io.path import ObjectStoragePath
-
                 self.manifest_path = ObjectStoragePath(manifest_path_str, conn_id=manifest_conn_id)
             else:
                 self.manifest_path = Path(manifest_path_str)

cosmos/constants.py

@@ -79,6 +79,8 @@ class TestBehavior(Enum):
     Behavior of the tests.
     """
 
+    __test__ = False
+
     BUILD = "build"
     NONE = "none"
     AFTER_EACH = "after_each"
@@ -116,6 +118,8 @@ class TestIndirectSelection(Enum):
     Modes to configure the test behavior when performing indirect selection.
     """
 
+    __test__ = False
+
     EAGER = "eager"
     CAUTIOUS = "cautious"
     BUILDABLE = "buildable"

cosmos/converter.py

@@ -12,7 +12,12 @@
 from warnings import warn
 
 from airflow.models.dag import DAG
-from airflow.utils.task_group import TaskGroup
+
+try:
+    # Airflow 3.1 onwards
+    from airflow.sdk import TaskGroup
+except ImportError:
+    from airflow.utils.task_group import TaskGroup
 
 from cosmos import cache, settings
 from cosmos.airflow.graph import build_airflow_graph

cosmos/core/airflow.py

@@ -8,7 +8,12 @@
 except ImportError:  # Airflow 2
     from airflow.models import BaseOperator
 from airflow.models.dag import DAG
-from airflow.utils.task_group import TaskGroup
+
+try:
+    # Airflow 3.1 onwards
+    from airflow.sdk import TaskGroup
+except ImportError:
+    from airflow.utils.task_group import TaskGroup
 
 from cosmos.core.graph.entities import Task
 from cosmos.log import get_logger

cosmos/core/graph/entities.py

@@ -3,11 +3,17 @@
 from dataclasses import dataclass, field
 from typing import Any, Dict, List
 
+import airflow
+from packaging.version import Version
+
 from cosmos.log import get_logger
 
 logger = get_logger(__name__)
 
 
+AIRFLOW_VERSION = Version(airflow.__version__)
+
+
 @dataclass
 class CosmosEntity:
     """
@@ -58,6 +64,10 @@ class Task(CosmosEntity):
     """
 
     owner: str = ""
-    operator_class: str = "airflow.operators.empty.EmptyOperator"
+    operator_class: str = (
+        "airflow.operators.empty.EmptyOperator"
+        if AIRFLOW_VERSION < Version("3.0")
+        else "airflow.providers.standard.operators.empty.EmptyOperator"
+    )
     arguments: Dict[str, Any] = field(default_factory=dict)
     extra_context: Dict[str, Any] = field(default_factory=dict)

cosmos/dataset.py

@@ -1,7 +1,15 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from airflow import DAG
-from airflow.utils.task_group import TaskGroup
+
+if TYPE_CHECKING:
+    try:
+        # Airflow 3.1 onwards
+        from airflow.utils.task_group import TaskGroup
+    except ImportError:
+        from airflow.utils.task_group import TaskGroup
 
 
 def get_dataset_alias_name(dag: DAG | None, task_group: TaskGroup | None, task_id: str) -> str:

cosmos/dbt/graph.py

@@ -20,9 +20,13 @@
 
 if TYPE_CHECKING:
     try:
-        from airflow.io.path import ObjectStoragePath
+        # Airflow 3 onwards
+        from airflow.sdk import ObjectStoragePath
     except ImportError:
-        pass
+        try:
+            from airflow.io.path import ObjectStoragePath
+        except ImportError:
+            pass
 
 import cosmos.dbt.runner as dbt_runner
 from cosmos import cache, settings

cosmos/hooks/subprocess.py

@@ -11,7 +11,11 @@
 from tempfile import TemporaryDirectory, gettempdir
 from typing import NamedTuple
 
-from airflow.hooks.base import BaseHook
+try:
+    # Airflow 3.1 onwards
+    from airflow.sdk.bases.hook import BaseHook
+except ImportError:
+    from airflow.hooks.base import BaseHook
 
 
 class FullOutputSubprocessResult(NamedTuple):
@@ -20,7 +24,7 @@ class FullOutputSubprocessResult(NamedTuple):
     full_output: list[str]
 
 
-class FullOutputSubprocessHook(BaseHook):
+class FullOutputSubprocessHook(BaseHook):  # type: ignore[misc]
     """Hook for running processes with the ``subprocess`` module."""
 
     def __init__(self) -> None: