Merge branch 'main' into fix-doc-default-values

Author: tatiana

.github/workflows/codeql.yml

@@ -43,7 +43,7 @@ jobs:
 
     steps:
     - name: Checkout repository
-      uses: actions/checkout@v5.0.1
+      uses: actions/checkout@v6.0.0
       with:
         persist-credentials: false
 

.github/workflows/deploy.yml

@@ -12,7 +12,7 @@ jobs:
       contents: read
 
     steps:
-      - uses: actions/checkout@v5.0.1
+      - uses: actions/checkout@v6.0.0
         with:
           persist-credentials: false
 

.github/workflows/test.yml

@@ -33,7 +33,7 @@ jobs:
     permissions:
       contents: read
     steps:
-      - uses: actions/checkout@v5.0.1
+      - uses: actions/checkout@v6.0.0
         with:
           ref: ${{ github.event.pull_request.head.sha || github.ref }}
           persist-credentials: false
@@ -77,7 +77,7 @@ jobs:
           - python-version: "3.12"
             airflow-version: "2.8"
     steps:
-      - uses: actions/checkout@v5.0.1
+      - uses: actions/checkout@v6.0.0
         with:
           ref: ${{ github.event.pull_request.head.sha || github.ref }}
           persist-credentials: false
@@ -140,7 +140,7 @@ jobs:
         ports:
           - 5432:5432
     steps:
-      - uses: actions/checkout@v5.0.1
+      - uses: actions/checkout@v6.0.0
         with:
           ref: ${{ github.event.pull_request.head.sha || github.ref }}
           persist-credentials: false
@@ -223,7 +223,7 @@ jobs:
           - 5432:5432
 
     steps:
-      - uses: actions/checkout@v5.0.1
+      - uses: actions/checkout@v6.0.0
         with:
           ref: ${{ github.event.pull_request.head.sha || github.ref }}
           persist-credentials: false
@@ -309,7 +309,7 @@ jobs:
           - 5432:5432
 
     steps:
-      - uses: actions/checkout@v5.0.1
+      - uses: actions/checkout@v6.0.0
         with:
           ref: ${{ github.event.pull_request.head.sha || github.ref }}
           persist-credentials: false
@@ -385,7 +385,7 @@ jobs:
           - 5432:5432
 
     steps:
-      - uses: actions/checkout@v5.0.1
+      - uses: actions/checkout@v6.0.0
         with:
           ref: ${{ github.event.pull_request.head.sha || github.ref }}
           persist-credentials: false
@@ -455,7 +455,7 @@ jobs:
         dbt-version: ["2.0"]  # dbt Fusion
 
     steps:
-      - uses: actions/checkout@v5.0.1
+      - uses: actions/checkout@v6.0.0
         with:
           ref: ${{ github.event.pull_request.head.sha || github.ref }}
           persist-credentials: false
@@ -552,7 +552,7 @@ jobs:
         ports:
           - 5432:5432
     steps:
-      - uses: actions/checkout@v5.0.1
+      - uses: actions/checkout@v6.0.0
         with:
           ref: ${{ github.event.pull_request.head.sha || github.ref }}
           persist-credentials: false
@@ -614,7 +614,7 @@ jobs:
         airflow-version: [ "2.10", "3.0" ]
         dbt-version: [ "1.10" ]
     steps:
-      - uses: actions/checkout@v5.0.1
+      - uses: actions/checkout@v6.0.0
         with:
           ref: ${{ github.event.pull_request.head.sha || github.ref }}
           persist-credentials: false
@@ -685,7 +685,7 @@ jobs:
     permissions:
       contents: read
     steps:
-      - uses: actions/checkout@v5.0.1
+      - uses: actions/checkout@v6.0.0
         with:
           ref: ${{ github.event.pull_request.head.sha || github.ref }}
           persist-credentials: false

cosmos/_triggers/watcher.py

@@ -11,6 +11,8 @@
 from asgiref.sync import sync_to_async
 from packaging.version import Version
 
+from cosmos._utils.watcher_state import build_producer_state_fetcher
+
 AIRFLOW_VERSION = Version(airflow.__version__)
 
 
@@ -72,8 +74,10 @@ def _get_xcom_val() -> Any | None:
                         task_id=self.producer_task_id,
                         run_id=self.run_id,
                     )
-                    .one()
+                    .one_or_none()
                 )
+                if ti is None:
+                    return None
                 return ti.xcom_pull(task_ids=self.producer_task_id, key=key)
 
         return await sync_to_async(_get_xcom_val)()
@@ -103,11 +107,26 @@ async def _parse_node_status(self) -> str | None:
         )
         return node_result.get("status")
 
+    async def _get_producer_task_status(self) -> str | None:
+        """Retrieve the producer task state for both Airflow 2 and Airflow 3."""
+
+        fetch_state = build_producer_state_fetcher(
+            airflow_version=AIRFLOW_VERSION,
+            dag_id=self.dag_id,
+            run_id=self.run_id,
+            producer_task_id=self.producer_task_id,
+            logger=self.log,
+        )
+        if fetch_state is None:
+            return None
+
+        return await sync_to_async(fetch_state)()
+
     async def run(self) -> AsyncIterator[TriggerEvent]:
         self.log.info("Starting WatcherTrigger for model: %s", self.model_unique_id)
 
         while True:
-            producer_task_state = await self.get_xcom_val("state")
+            producer_task_state = await self._get_producer_task_status()
             node_status = await self._parse_node_status()
             if node_status == "success":
                 self.log.info("Model '%s' succeeded", self.model_unique_id)

cosmos/_utils/watcher_state.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable
+
+from packaging.version import Version
+
+ProducerStateFetcher = Callable[[], str | None]
+
+
+def _load_airflow2_dependencies() -> tuple[Any, Callable[[], Any]]:
+    from airflow.models import TaskInstance
+    from airflow.utils.session import create_session
+
+    return TaskInstance, create_session
+
+
+def _load_airflow3_dependencies() -> Any:
+    from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+    return RuntimeTaskInstance
+
+
+def build_producer_state_fetcher(
+    *,
+    airflow_version: Version,
+    dag_id: str,
+    run_id: str,
+    producer_task_id: str,
+    logger: logging.Logger,
+) -> ProducerStateFetcher | None:
+    """Return a callable that fetches the producer task state for the given Airflow version."""
+
+    if airflow_version < Version("3.0.0"):
+        try:
+            TaskInstance, create_session = _load_airflow2_dependencies()
+        except ImportError as exc:  # pragma: no cover - defensive guard for stripped test envs
+            logger.warning("Could not import Airflow 2 state dependencies: %s", exc)
+            return None
+
+        def fetch_state_airflow2() -> str | None:
+            with create_session() as session:
+                ti = (
+                    session.query(TaskInstance)
+                    .filter_by(
+                        dag_id=dag_id,
+                        task_id=producer_task_id,
+                        run_id=run_id,
+                    )
+                    .one_or_none()
+                )
+                if ti is not None:
+                    return str(ti.state)
+                return None
+
+        return fetch_state_airflow2
+
+    try:
+        RuntimeTaskInstance = _load_airflow3_dependencies()
+    except (ImportError, NameError) as exc:  # pragma: no cover - Airflow 3 libs missing
+        logger.warning("Could not load Airflow 3 RuntimeTaskInstance: %s", exc)
+        return None
+
+    def fetch_state_airflow3() -> str | None:
+        try:
+            task_states = RuntimeTaskInstance.get_task_states(
+                dag_id=dag_id,
+                task_ids=[producer_task_id],
+                run_ids=[run_id],
+            )
+        except NameError as exc:  # pragma: no cover - Airflow 3.0 missing supervisor comms
+            logger.warning("RuntimeTaskInstance.get_task_states unavailable due to NameError: %s", exc)
+            return None
+        state = task_states.get(run_id, {}).get(producer_task_id)
+        if state is not None:
+            return str(state)
+        return None
+
+    return fetch_state_airflow3

cosmos/operators/watcher.py

@@ -7,7 +7,7 @@
 from datetime import timedelta
 from pathlib import Path
 from threading import Lock
-from typing import TYPE_CHECKING, Any, Callable, List, Union
+from typing import TYPE_CHECKING, Any
 
 from cosmos._triggers.watcher import WatcherTrigger, _parse_compressed_xcom
 
@@ -30,8 +30,8 @@
 except ImportError:  # pragma: no cover
     from airflow.operators.empty import EmptyOperator  # type: ignore[no-redef]
 
-from packaging.version import Version
 
+from cosmos._utils.watcher_state import build_producer_state_fetcher
 from cosmos.config import ProfileConfig
 from cosmos.constants import AIRFLOW_VERSION, PRODUCER_WATCHER_TASK_ID, InvocationMode
 from cosmos.operators.base import (
@@ -110,27 +110,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         kwargs["default_args"] = default_args
         kwargs["retries"] = 0
 
-        on_failure_callback = self._set_on_failure_callback(kwargs.pop("on_failure_callback", None))
-        super().__init__(task_id=task_id, *args, on_failure_callback=on_failure_callback, **kwargs)
-
-    def _set_on_failure_callback(
-        self, user_callback: Any
-    ) -> Union[Callable[[Context], None], List[Callable[[Context], None]]]:
-        default_callback = self._store_producer_task_state
-
-        if AIRFLOW_VERSION < Version("2.6.0"):
-            # Older versions only support a single callable
-            return default_callback
-        else:
-            if user_callback is None:
-                # No callback provided — use default in a list
-                return [default_callback]
-            elif isinstance(user_callback, list):
-                # Append to existing list of callbacks (make a copy to avoid side effects)
-                return user_callback + [default_callback]
-            else:
-                # Single callable provided — wrap it in a list and append ours
-                return [user_callback, default_callback]
+        super().__init__(task_id=task_id, *args, **kwargs)
 
     @staticmethod
     def _serialize_event(event_message: EventMsg) -> dict[str, Any]:
@@ -179,9 +159,6 @@ def _finalize(self, context: Context, startup_events: list[dict[str, Any]]) -> N
         if startup_events:
             safe_xcom_push(task_instance=context["ti"], key="dbt_startup_events", value=startup_events)
 
-    def _store_producer_task_state(self, context: Context) -> None:
-        safe_xcom_push(task_instance=context["ti"], key="state", value="failed")
-
     def execute(self, context: Context, **kwargs: Any) -> Any:
         task_instance = context.get("ti")
         if task_instance is None:
@@ -371,8 +348,27 @@ def _get_status_from_run_results(self, ti: Any, context: Context) -> Any:
 
         return node_result.get("status")
 
-    def _get_producer_task_state(self, ti: Any) -> Any:
-        return ti.xcom_pull(task_ids=self.producer_task_id, key="state")
+    def _get_producer_task_status(self, context: Context) -> str | None:
+        """
+        Get the task status of the producer task for both Airflow 2 and Airflow 3.
+
+        Returns the state of the producer task instance, or None if not found.
+        """
+        ti = context["ti"]
+        run_id = context["run_id"]
+        dag_id = ti.dag_id
+
+        fetch_state = build_producer_state_fetcher(
+            airflow_version=AIRFLOW_VERSION,
+            dag_id=dag_id,
+            run_id=run_id,
+            producer_task_id=self.producer_task_id,
+            logger=logger,
+        )
+        if fetch_state is None:
+            return None
+
+        return fetch_state()
 
     def execute(self, context: Context, **kwargs: Any) -> None:
         if not self.deferrable:
@@ -433,7 +429,7 @@ def poke(self, context: Context) -> bool:
             return self._fallback_to_local_run(try_number, context)
 
         # We have assumption here that both the build producer and the sensor task will have same invocation mode
-        producer_task_state = self._get_producer_task_state(ti)
+        producer_task_state = self._get_producer_task_status(context)
         if self._use_event():
             status = self._get_status_from_events(ti, context)
         else:

docs/getting_started/async-execution-mode.rst

@@ -248,5 +248,7 @@ Limitations
 
 9. **TeardownAsyncOperator limitation**: When using a remote object location, in addition to the ``SetupAsyncOperator``, a ``TeardownAsyncOperator`` is also added to the DAG. This task will delete the SQL files from the remote location by the end of the DAG Run. This is can lead to a limitation from a retry perspective, as described in the issue `#2066 <https://github.com/astronomer/astronomer-cosmos/issues/2066>`_. This can be avoided by setting the ``enable_teardown_async_task`` configuration to ``False``, as described in the :ref:`enable_teardown_async_task` section.
 
+10. **Dataset events not emitted**: Dataset events are not currently emitted after dbt models complete when using ``ExecutionMode.AIRFLOW_ASYNC``. This means downstream DAGs scheduled with ``Dataset`` or ``DatasetAlias`` will not trigger automatically. This behaviour is present in ``ExecutionMode.LOCAL`` but is currently missing in async mode. This issue is being tracked in `#2141 <https://github.com/astronomer/astronomer-cosmos/issues/2141>`_.
+
 
 For a comparison between different Cosmos execution modes, please, check the :ref:`execution-modes-comparison` section.