Support Airflow 3.1

.github/workflows/test.yml

@@ -54,8 +54,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"
@@ -123,8 +123,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"
@@ -551,7 +551,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.0a10"
 
 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

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/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:

cosmos/io.py

@@ -5,6 +5,14 @@
 from typing import Any
 from urllib.parse import urlparse
 
+try:
+    from airflow.sdk import ObjectStoragePath
+except ImportError:
+    try:
+        from airflow.io.path import ObjectStoragePath
+    except ImportError:
+        pass
+
 from cosmos import settings
 from cosmos.constants import DEFAULT_TARGET_PATH, FILE_SCHEME_AIRFLOW_DEFAULT_CONN_ID_MAP
 from cosmos.exceptions import CosmosValueError
@@ -158,9 +166,6 @@ def _configure_remote_target_path() -> tuple[Path, str] | tuple[None, None]:
             f"Object Storage feature is unavailable in Airflow version {airflow_version}. Please upgrade to "
             "Airflow 2.8 or later."
         )
-
-    from airflow.io.path import ObjectStoragePath
-
     _configured_target_path = ObjectStoragePath(target_path_str, conn_id=remote_conn_id)
 
     if not _configured_target_path.exists():  # type: ignore[no-untyped-call]
@@ -207,8 +212,6 @@ def upload_to_cloud_storage(project_dir: str, source_subpath: str = DEFAULT_TARG
     if not dest_target_dir:
         raise CosmosValueError("You're trying to upload artifact files, but the remote target path is not configured.")
 
-    from airflow.io.path import ObjectStoragePath
-
     source_target_dir = Path(project_dir) / f"{source_subpath}"
     files = [str(file) for file in source_target_dir.rglob("*") if file.is_file()]
     for file_path in files:

cosmos/operators/_asynchronous/__init__.py

@@ -5,6 +5,14 @@
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
+try:
+    from airflow.sdk import ObjectStoragePath
+except ImportError:
+    try:
+        from airflow.io.path import ObjectStoragePath
+    except ImportError:
+        pass
+
 from cosmos.operators.local import DbtRunLocalOperator
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -59,8 +67,6 @@ def execute(self, context: Context, **kwargs: Any) -> Any:
 
         dest_target_dir, dest_conn_id = self._configure_remote_target_path()
 
-        from airflow.io.path import ObjectStoragePath
-
         dag_task_group_identifier = self.extra_context["dbt_dag_task_group_identifier"]
         run_id = context["run_id"]
         run_dir_path_str = f"{str(dest_target_dir).rstrip('/')}/{dag_task_group_identifier}/{run_id}"

cosmos/operators/_asynchronous/bigquery.py

@@ -150,7 +150,10 @@ def get_remote_sql(self) -> str:
 
         if not settings.AIRFLOW_IO_AVAILABLE:  # pragma: no cover
             raise CosmosValueError(f"Cosmos async support is only available starting in Airflow 2.8 or later.")
-        from airflow.io.path import ObjectStoragePath
+        try:
+            from airflow.sdk import ObjectStoragePath
+        except ImportError:
+            from airflow.io.path import ObjectStoragePath
 
         file_path = self.async_context["dbt_node_config"]["file_path"]  # type: ignore
         dbt_dag_task_group_identifier = self.async_context["dbt_dag_task_group_identifier"]

cosmos/operators/_asynchronous/databricks.py

@@ -3,11 +3,11 @@
 
 from typing import Any
 
-from airflow.utils.context import Context
+from airflow.utils.context import Context  # type: ignore[attr-defined]
 
 try:
     from airflow.sdk.bases.operator import BaseOperator  # Airflow 3
-except ImportError:
+except (ImportError, AttributeError):
     from airflow.models import BaseOperator  # Airflow 2