apache-airflow-providers-standard 1.9.1rc1__py3-none-any.whl → 1.10.3__py3-none-any.whl

airflow/providers/standard/sensors/external_task.py

@@ -18,14 +18,18 @@ from __future__ import annotations
 
 import datetime
 import os
+import typing
 import warnings
-from collections.abc import Callable, Collection, Iterable
-from typing import TYPE_CHECKING, Any, ClassVar
+from collections.abc import Callable, Collection, Iterable, Sequence
+from typing import TYPE_CHECKING, ClassVar
 
-from airflow.configuration import conf
-from airflow.exceptions import AirflowSkipException
 from airflow.models.dag import DagModel
-from airflow.providers.common.compat.sdk import BaseOperatorLink, BaseSensorOperator
+from airflow.providers.common.compat.sdk import (
+    AirflowSkipException,
+    BaseOperatorLink,
+    BaseSensorOperator,
+    conf,
+)
 from airflow.providers.standard.exceptions import (
     DuplicateStateError,
     ExternalDagDeletedError,
@@ -250,18 +254,21 @@ class ExternalTaskSensor(BaseSensorOperator):
         self._has_checked_existence = False
         self.deferrable = deferrable
         self.poll_interval = poll_interval
+        self.external_dates_filter: str | None = None
 
-    def _get_dttm_filter(self, context):
+    def _get_dttm_filter(self, context: Context) -> Sequence[datetime.datetime]:
         logical_date = self._get_logical_date(context)
 
         if self.execution_delta:
-            dttm = logical_date - self.execution_delta
-        elif self.execution_date_fn:
-            dttm = self._handle_execution_date_fn(context=context)
-        else:
-            dttm = logical_date
+            return [logical_date - self.execution_delta]
+        if self.execution_date_fn:
+            result = self._handle_execution_date_fn(context=context)
+            return result if isinstance(result, list) else [result]
+        return [logical_date]
 
-        return dttm if isinstance(dttm, list) else [dttm]
+    @staticmethod
+    def _serialize_dttm_filter(dttm_filter: Sequence[datetime.datetime]) -> str:
+        return ",".join(dt.isoformat() for dt in dttm_filter)
 
     def poke(self, context: Context) -> bool:
         # delay check to poke rather than __init__ in case it was supplied as XComArgs
@@ -269,7 +276,9 @@ class ExternalTaskSensor(BaseSensorOperator):
             raise ValueError("Duplicate task_ids passed in external_task_ids parameter")
 
         dttm_filter = self._get_dttm_filter(context)
-        serialized_dttm_filter = ",".join(dt.isoformat() for dt in dttm_filter)
+        serialized_dttm_filter = self._serialize_dttm_filter(dttm_filter)
+        # Save as attribute - to be used by listeners
+        self.external_dates_filter = serialized_dttm_filter
 
         if self.external_task_ids:
             self.log.info(
@@ -298,7 +307,7 @@ class ExternalTaskSensor(BaseSensorOperator):
             return self._poke_af3(context, dttm_filter)
         return self._poke_af2(dttm_filter)
 
-    def _poke_af3(self, context: Context, dttm_filter: list[datetime.datetime]) -> bool:
+    def _poke_af3(self, context: Context, dttm_filter: Sequence[datetime.datetime]) -> bool:
         from airflow.providers.standard.utils.sensor_helper import _get_count_by_matched_states
 
         self._has_checked_existence = True
@@ -308,20 +317,20 @@ class ExternalTaskSensor(BaseSensorOperator):
             if self.external_task_ids:
                 return ti.get_ti_count(
                     dag_id=self.external_dag_id,
-                    task_ids=self.external_task_ids,  # type: ignore[arg-type]
-                    logical_dates=dttm_filter,
+                    task_ids=list(self.external_task_ids),
+                    logical_dates=list(dttm_filter),
                     states=states,
                 )
             if self.external_task_group_id:
                 run_id_task_state_map = ti.get_task_states(
                     dag_id=self.external_dag_id,
                     task_group_id=self.external_task_group_id,
-                    logical_dates=dttm_filter,
+                    logical_dates=list(dttm_filter),
                 )
                 return _get_count_by_matched_states(run_id_task_state_map, states)
             return ti.get_dr_count(
                 dag_id=self.external_dag_id,
-                logical_dates=dttm_filter,
+                logical_dates=list(dttm_filter),
                 states=states,
             )
 
@@ -339,7 +348,7 @@ class ExternalTaskSensor(BaseSensorOperator):
         count_allowed = self._calculate_count(count, dttm_filter)
         return count_allowed == len(dttm_filter)
 
-    def _calculate_count(self, count: int, dttm_filter: list[datetime.datetime]) -> float | int:
+    def _calculate_count(self, count: int, dttm_filter: Sequence[datetime.datetime]) -> float | int:
         """Calculate the normalized count based on the type of check."""
         if self.external_task_ids:
             return count / len(self.external_task_ids)
@@ -395,7 +404,7 @@ class ExternalTaskSensor(BaseSensorOperator):
     if not AIRFLOW_V_3_0_PLUS:
 
         @provide_session
-        def _poke_af2(self, dttm_filter: list[datetime.datetime], session: Session = NEW_SESSION) -> bool:
+        def _poke_af2(self, dttm_filter: Sequence[datetime.datetime], session: Session = NEW_SESSION) -> bool:
             if self.check_existence and not self._has_checked_existence:
                 self._check_for_existence(session=session)
 
@@ -416,27 +425,51 @@ class ExternalTaskSensor(BaseSensorOperator):
             super().execute(context)
         else:
             dttm_filter = self._get_dttm_filter(context)
-            logical_or_execution_dates = (
-                {"logical_dates": dttm_filter} if AIRFLOW_V_3_0_PLUS else {"execution_dates": dttm_filter}
-            )
-            self.defer(
-                timeout=self.execution_timeout,
-                trigger=WorkflowTrigger(
-                    external_dag_id=self.external_dag_id,
-                    external_task_group_id=self.external_task_group_id,
-                    external_task_ids=self.external_task_ids,
-                    allowed_states=self.allowed_states,
-                    failed_states=self.failed_states,
-                    skipped_states=self.skipped_states,
-                    poke_interval=self.poll_interval,
-                    soft_fail=self.soft_fail,
-                    **logical_or_execution_dates,
-                ),
-                method_name="execute_complete",
-            )
+            if AIRFLOW_V_3_0_PLUS:
+                self.defer(
+                    timeout=self.execution_timeout,
+                    trigger=WorkflowTrigger(
+                        external_dag_id=self.external_dag_id,
+                        external_task_group_id=self.external_task_group_id,
+                        external_task_ids=self.external_task_ids,
+                        allowed_states=self.allowed_states,
+                        failed_states=self.failed_states,
+                        skipped_states=self.skipped_states,
+                        poke_interval=self.poll_interval,
+                        soft_fail=self.soft_fail,
+                        logical_dates=list(dttm_filter),
+                        run_ids=None,
+                        execution_dates=None,
+                    ),
+                    method_name="execute_complete",
+                )
+            else:
+                self.defer(
+                    timeout=self.execution_timeout,
+                    trigger=WorkflowTrigger(
+                        external_dag_id=self.external_dag_id,
+                        external_task_group_id=self.external_task_group_id,
+                        external_task_ids=self.external_task_ids,
+                        allowed_states=self.allowed_states,
+                        failed_states=self.failed_states,
+                        skipped_states=self.skipped_states,
+                        poke_interval=self.poll_interval,
+                        soft_fail=self.soft_fail,
+                        execution_dates=list(dttm_filter),
+                        logical_dates=None,
+                        run_ids=None,
+                    ),
+                    method_name="execute_complete",
+                )
 
-    def execute_complete(self, context, event=None):
+    def execute_complete(self, context: Context, event: dict[str, typing.Any] | None = None) -> None:
         """Execute when the trigger fires - return immediately."""
+        if event is None:
+            raise ExternalTaskNotFoundError("No event received from trigger")
+
+        # Re-set as attribute after coming back from deferral - to be used by listeners
+        self.external_dates_filter = self._serialize_dttm_filter(self._get_dttm_filter(context))
+
         if event["status"] == "success":
             self.log.info("External tasks %s has executed successfully.", self.external_task_ids)
         elif event["status"] == "skipped":
@@ -453,13 +486,14 @@ class ExternalTaskSensor(BaseSensorOperator):
                 "name of executed task and Dag."
             )
 
-    def _check_for_existence(self, session) -> None:
+    def _check_for_existence(self, session: Session) -> None:
         dag_to_wait = DagModel.get_current(self.external_dag_id, session)
 
         if not dag_to_wait:
             raise ExternalDagNotFoundError(f"The external DAG {self.external_dag_id} does not exist.")
 
-        if not os.path.exists(correct_maybe_zipped(dag_to_wait.fileloc)):
+        path = correct_maybe_zipped(dag_to_wait.fileloc)
+        if not path or not os.path.exists(path):
             raise ExternalDagDeletedError(f"The external DAG {self.external_dag_id} was deleted.")
 
         if self.external_task_ids:
@@ -488,7 +522,7 @@ class ExternalTaskSensor(BaseSensorOperator):
 
         self._has_checked_existence = True
 
-    def get_count(self, dttm_filter, session, states) -> int:
+    def get_count(self, dttm_filter: Sequence[datetime.datetime], session: Session, states: list[str]) -> int:
         """
         Get the count of records against dttm filter and states.
 
@@ -509,15 +543,19 @@ class ExternalTaskSensor(BaseSensorOperator):
             session,
         )
 
-    def get_external_task_group_task_ids(self, session, dttm_filter):
+    def get_external_task_group_task_ids(
+        self, session: Session, dttm_filter: Sequence[datetime.datetime]
+    ) -> list[tuple[str, int]]:
         warnings.warn(
             "This method is deprecated and will be removed in future.", DeprecationWarning, stacklevel=2
         )
+        if self.external_task_group_id is None:
+            return []
         return _get_external_task_group_task_ids(
-            dttm_filter, self.external_task_group_id, self.external_dag_id, session
+            list(dttm_filter), self.external_task_group_id, self.external_dag_id, session
         )
 
-    def _get_logical_date(self, context) -> datetime.datetime:
+    def _get_logical_date(self, context: Context) -> datetime.datetime:
         """
         Handle backwards- and forwards-compatible retrieval of the date.
 
@@ -527,19 +565,21 @@ class ExternalTaskSensor(BaseSensorOperator):
         if AIRFLOW_V_3_0_PLUS:
             logical_date = context.get("logical_date")
             dag_run = context.get("dag_run")
-            if not (logical_date or (dag_run and dag_run.run_after)):
-                raise ValueError(
-                    "Either `logical_date` or `dag_run.run_after` must be provided in the context"
-                )
-            return logical_date or dag_run.run_after
+            if logical_date:
+                return logical_date
+            if dag_run and hasattr(dag_run, "run_after") and dag_run.run_after:
+                return dag_run.run_after
+            raise ValueError("Either `logical_date` or `dag_run.run_after` must be provided in the context")
 
         # Airflow 2.x and earlier: contexts used "execution_date"
         execution_date = context.get("execution_date")
         if not execution_date:
             raise ValueError("Either `execution_date` must be provided in the context`")
+        if not isinstance(execution_date, datetime.datetime):
+            raise ValueError("execution_date must be a datetime object")
         return execution_date
 
-    def _handle_execution_date_fn(self, context) -> Any:
+    def _handle_execution_date_fn(self, context: Context) -> datetime.datetime | list[datetime.datetime]:
         """
         Handle backward compatibility.
 

airflow/providers/standard/sensors/filesystem.py

@@ -20,31 +20,14 @@ from __future__ import annotations
 import datetime
 import os
 from collections.abc import Sequence
-from dataclasses import dataclass
 from functools import cached_property
 from glob import glob
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
-from airflow.exceptions import AirflowException
-from airflow.providers.common.compat.sdk import BaseSensorOperator
+from airflow.providers.common.compat.sdk import AirflowException, BaseSensorOperator, conf
 from airflow.providers.standard.hooks.filesystem import FSHook
 from airflow.providers.standard.triggers.file import FileTrigger
-
-try:
-    from airflow.triggers.base import StartTriggerArgs  # type: ignore[no-redef]
-except ImportError:  # TODO: Remove this when min airflow version is 2.10.0 for standard provider
-
-    @dataclass
-    class StartTriggerArgs:  # type: ignore[no-redef]
-        """Arguments required for start task execution from triggerer."""
-
-        trigger_cls: str
-        next_method: str
-        trigger_kwargs: dict[str, Any] | None = None
-        next_kwargs: dict[str, Any] | None = None
-        timeout: datetime.timedelta | None = None
-
+from airflow.triggers.base import StartTriggerArgs
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

airflow/providers/standard/sensors/time.py

@@ -19,28 +19,12 @@ from __future__ import annotations
 
 import datetime
 import warnings
-from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
-from airflow.configuration import conf
 from airflow.exceptions import AirflowProviderDeprecationWarning
-from airflow.providers.common.compat.sdk import BaseSensorOperator, timezone
+from airflow.providers.common.compat.sdk import BaseSensorOperator, conf, timezone
 from airflow.providers.standard.triggers.temporal import DateTimeTrigger
-
-try:
-    from airflow.triggers.base import StartTriggerArgs  # type: ignore[no-redef]
-except ImportError:  # TODO: Remove this when min airflow version is 2.10.0 for standard provider
-
-    @dataclass
-    class StartTriggerArgs:  # type: ignore[no-redef]
-        """Arguments required for start task execution from triggerer."""
-
-        trigger_cls: str
-        next_method: str
-        trigger_kwargs: dict[str, Any] | None = None
-        next_kwargs: dict[str, Any] | None = None
-        timeout: datetime.timedelta | None = None
-
+from airflow.triggers.base import StartTriggerArgs
 
 if TYPE_CHECKING:
     from airflow.sdk import Context

airflow/providers/standard/sensors/time_delta.py

@@ -25,9 +25,8 @@ from typing import TYPE_CHECKING, Any
 from deprecated.classic import deprecated
 from packaging.version import Version
 
-from airflow.configuration import conf
-from airflow.exceptions import AirflowProviderDeprecationWarning, AirflowSkipException
-from airflow.providers.common.compat.sdk import BaseSensorOperator, timezone
+from airflow.exceptions import AirflowProviderDeprecationWarning
+from airflow.providers.common.compat.sdk import AirflowSkipException, BaseSensorOperator, conf, timezone
 from airflow.providers.standard.triggers.temporal import DateTimeTrigger, TimeDeltaTrigger
 from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
 
@@ -194,9 +193,11 @@ class WaitSensor(BaseSensorOperator):
     def execute(self, context: Context) -> None:
         if self.deferrable:
             self.defer(
-                trigger=TimeDeltaTrigger(self.time_to_wait, end_from_trigger=True)
-                if AIRFLOW_V_3_0_PLUS
-                else TimeDeltaTrigger(self.time_to_wait),
+                trigger=(
+                    TimeDeltaTrigger(self.time_to_wait, end_from_trigger=True)
+                    if AIRFLOW_V_3_0_PLUS
+                    else TimeDeltaTrigger(self.time_to_wait)
+                ),
                 method_name="execute_complete",
             )
         else:

airflow/providers/standard/triggers/external_task.py

@@ -18,10 +18,11 @@ from __future__ import annotations
 
 import asyncio
 import typing
+from collections.abc import Collection
 from typing import Any
 
 from asgiref.sync import sync_to_async
-from sqlalchemy import func
+from sqlalchemy import func, select
 
 from airflow.models import DagRun
 from airflow.providers.standard.utils.sensor_helper import _get_count
@@ -60,9 +61,9 @@ class WorkflowTrigger(BaseTrigger):
         logical_dates: list[datetime] | None = None,
         external_task_ids: typing.Collection[str] | None = None,
         external_task_group_id: str | None = None,
-        failed_states: typing.Iterable[str] | None = None,
-        skipped_states: typing.Iterable[str] | None = None,
-        allowed_states: typing.Iterable[str] | None = None,
+        failed_states: Collection[str] | None = None,
+        skipped_states: Collection[str] | None = None,
+        allowed_states: Collection[str] | None = None,
         poke_interval: float = 2.0,
         soft_fail: bool = False,
         **kwargs,
@@ -129,43 +130,41 @@ class WorkflowTrigger(BaseTrigger):
             self.log.info("Sleeping for %s seconds", self.poke_interval)
             await asyncio.sleep(self.poke_interval)
 
-    async def _get_count_af_3(self, states):
+    async def _get_count_af_3(self, states: Collection[str] | None) -> int:
         from airflow.providers.standard.utils.sensor_helper import _get_count_by_matched_states
         from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
 
-        params = {
-            "dag_id": self.external_dag_id,
-            "logical_dates": self.logical_dates,
-            "run_ids": self.run_ids,
-        }
         if self.external_task_ids:
             count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
-                task_ids=self.external_task_ids,
-                states=states,
-                **params,
+                dag_id=self.external_dag_id,
+                task_ids=list(self.external_task_ids),
+                logical_dates=self.logical_dates,
+                run_ids=self.run_ids,
+                states=list(states) if states else None,
             )
-        elif self.external_task_group_id:
+            return int(count / len(self.external_task_ids))
+        if self.external_task_group_id:
             run_id_task_state_map = await sync_to_async(RuntimeTaskInstance.get_task_states)(
+                dag_id=self.external_dag_id,
                 task_group_id=self.external_task_group_id,
-                **params,
+                logical_dates=self.logical_dates,
+                run_ids=self.run_ids,
             )
             count = await sync_to_async(_get_count_by_matched_states)(
                 run_id_task_state_map=run_id_task_state_map,
-                states=states,
-            )
-        else:
-            count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
-                dag_id=self.external_dag_id,
-                logical_dates=self.logical_dates,
-                run_ids=self.run_ids,
-                states=states,
+                states=states or [],
             )
-        if self.external_task_ids:
-            return count / len(self.external_task_ids)
+            return count
+        count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+            dag_id=self.external_dag_id,
+            logical_dates=self.logical_dates,
+            run_ids=self.run_ids,
+            states=list(states) if states else None,
+        )
         return count
 
     @sync_to_async
-    def _get_count(self, states: typing.Iterable[str] | None) -> int:
+    def _get_count(self, states: Collection[str] | None) -> int:
         """
         Get the count of records against dttm filter and states. Async wrapper for _get_count.
 
@@ -227,23 +226,26 @@ class DagStateTrigger(BaseTrigger):
         elif self.execution_dates:
             runs_ids_or_dates = len(self.execution_dates)
 
+        cls_path, data = self.serialize()
+
         if AIRFLOW_V_3_0_PLUS:
-            event = await self.validate_count_dags_af_3(runs_ids_or_dates_len=runs_ids_or_dates)
-            yield TriggerEvent(event)
+            data.update(  # update with {run_id: run_state} dict
+                await self.validate_count_dags_af_3(runs_ids_or_dates_len=runs_ids_or_dates)
+            )
+            yield TriggerEvent((cls_path, data))
             return
         else:
             while True:
                 num_dags = await self.count_dags()
                 if num_dags == runs_ids_or_dates:
-                    yield TriggerEvent(self.serialize())
+                    yield TriggerEvent((cls_path, data))
                     return
                 await asyncio.sleep(self.poll_interval)
 
-    async def validate_count_dags_af_3(self, runs_ids_or_dates_len: int = 0) -> tuple[str, dict[str, Any]]:
+    async def validate_count_dags_af_3(self, runs_ids_or_dates_len: int = 0) -> dict[str, str]:
         from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
 
-        cls_path, data = self.serialize()
-
+        run_states: dict[str, str] = {}  # {run_id: run_state}
         while True:
             num_dags = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
                 dag_id=self.dag_id,
@@ -258,8 +260,8 @@ class DagStateTrigger(BaseTrigger):
                             dag_id=self.dag_id,
                             run_id=run_id,
                         )
-                        data[run_id] = state
-                        return cls_path, data
+                        run_states[run_id] = state
+                        return run_states
             await asyncio.sleep(self.poll_interval)
 
     if not AIRFLOW_V_3_0_PLUS:
@@ -270,17 +272,18 @@ class DagStateTrigger(BaseTrigger):
         def count_dags(self, *, session: Session = NEW_SESSION) -> int:
             """Count how many dag runs in the database match our criteria."""
             _dag_run_date_condition = (
-                DagRun.run_id.in_(self.run_ids)
+                DagRun.run_id.in_(self.run_ids or [])
                 if AIRFLOW_V_3_0_PLUS
                 else DagRun.execution_date.in_(self.execution_dates)
             )
-            count = (
-                session.query(func.count("*"))  # .count() is inefficient
-                .filter(
+            stmt = (
+                select(func.count())
+                .select_from(DagRun)
+                .where(
                     DagRun.dag_id == self.dag_id,
                     DagRun.state.in_(self.states),
                     _dag_run_date_condition,
                 )
-                .scalar()
             )
-            return typing.cast("int", count)
+            result = session.execute(stmt).scalar()
+            return result or 0

airflow/providers/standard/triggers/file.py

@@ -79,7 +79,7 @@ class FileTrigger(BaseTrigger):
                     self.log.info("Found File %s last modified: %s", path, mod_time)
                     yield TriggerEvent(True)
                     return
-                for _, _, files in os.walk(self.filepath):
+                for _, _, files in os.walk(path):
                     if files:
                         yield TriggerEvent(True)
                         return