apache-airflow-providers-standard 1.0.0.dev1__py3-none-any.whl → 1.1.0rc1__py3-none-any.whl

airflow/providers/standard/sensors/filesystem.py

@@ -0,0 +1,158 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+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.standard.hooks.filesystem import FSHook
+from airflow.providers.standard.triggers.file import FileTrigger
+from airflow.sensors.base import BaseSensorOperator
+
+try:
+    from airflow.triggers.base import StartTriggerArgs
+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
+
+
+if TYPE_CHECKING:
+    try:
+        from airflow.sdk.definitions.context import Context
+    except ImportError:
+        # TODO: Remove once provider drops support for Airflow 2
+        from airflow.utils.context import Context
+
+
+class FileSensor(BaseSensorOperator):
+    """
+    Waits for a file or folder to land in a filesystem.
+
+    If the path given is a directory then this sensor will only return true if
+    any files exist inside it (either directly, or within a subdirectory)
+
+    :param fs_conn_id: reference to the File (path)
+        connection id
+    :param filepath: File or folder name (relative to
+        the base path set within the connection), can be a glob.
+    :param recursive: when set to ``True``, enables recursive directory matching behavior of
+        ``**`` in glob filepath parameter. Defaults to ``False``.
+    :param deferrable: If waiting for completion, whether to defer the task until done,
+        default is ``False``.
+    :param start_from_trigger: Start the task directly from the triggerer without going into the worker.
+    :param trigger_kwargs: The keyword arguments passed to the trigger when start_from_trigger is set to True
+        during dynamic task mapping. This argument is not used in standard usage.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/operator:FileSensor`
+
+
+    """
+
+    template_fields: Sequence[str] = ("filepath",)
+    ui_color = "#91818a"
+    start_trigger_args = StartTriggerArgs(
+        trigger_cls="airflow.providers.standard.triggers.file.FileTrigger",
+        trigger_kwargs={},
+        next_method="execute_complete",
+        next_kwargs=None,
+        timeout=None,
+    )
+    start_from_trigger = False
+
+    def __init__(
+        self,
+        *,
+        filepath,
+        fs_conn_id="fs_default",
+        recursive=False,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        start_from_trigger: bool = False,
+        trigger_kwargs: dict[str, Any] | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.filepath = filepath
+        self.fs_conn_id = fs_conn_id
+        self.recursive = recursive
+        self.deferrable = deferrable
+
+        self.start_from_trigger = start_from_trigger
+
+        if self.deferrable and self.start_from_trigger:
+            self.start_trigger_args.timeout = datetime.timedelta(seconds=self.timeout)
+            self.start_trigger_args.trigger_kwargs = dict(
+                filepath=self.path,
+                recursive=self.recursive,
+                poke_interval=self.poke_interval,
+            )
+
+    @cached_property
+    def path(self) -> str:
+        hook = FSHook(self.fs_conn_id)
+        basepath = hook.get_path()
+        full_path = os.path.join(basepath, self.filepath)
+        return full_path
+
+    def poke(self, context: Context) -> bool:
+        self.log.info("Poking for file %s", self.path)
+        for path in glob(self.path, recursive=self.recursive):
+            if os.path.isfile(path):
+                mod_time = datetime.datetime.fromtimestamp(os.path.getmtime(path)).strftime("%Y%m%d%H%M%S")
+                self.log.info("Found File %s last modified: %s", path, mod_time)
+                return True
+
+            for _, _, files in os.walk(path):
+                if files:
+                    return True
+        return False
+
+    def execute(self, context: Context) -> None:
+        if not self.deferrable:
+            super().execute(context=context)
+        if not self.poke(context=context):
+            self.defer(
+                timeout=datetime.timedelta(seconds=self.timeout),
+                trigger=FileTrigger(
+                    filepath=self.path,
+                    recursive=self.recursive,
+                    poke_interval=self.poke_interval,
+                ),
+                method_name="execute_complete",
+            )
+
+    def execute_complete(self, context: Context, event: bool | None = None) -> None:
+        if not event:
+            raise AirflowException("%s task failed as %s not found.", self.task_id, self.filepath)
+        self.log.info("%s completed successfully as %s found.", self.task_id, self.filepath)

airflow/providers/standard/sensors/python.py

@@ -0,0 +1,84 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any, Callable
+
+from airflow.sensors.base import BaseSensorOperator, PokeReturnValue
+from airflow.utils.context import context_merge
+from airflow.utils.operator_helpers import determine_kwargs
+
+if TYPE_CHECKING:
+    try:
+        from airflow.sdk.definitions.context import Context
+    except ImportError:
+        # TODO: Remove once provider drops support for Airflow 2
+        from airflow.utils.context import Context
+
+
+class PythonSensor(BaseSensorOperator):
+    """
+    Waits for a Python callable to return True.
+
+    User could put input argument in templates_dict
+    e.g ``templates_dict = {'start_ds': 1970}``
+    and access the argument by calling ``kwargs['templates_dict']['start_ds']``
+    in the callable
+
+    :param python_callable: A reference to an object that is callable
+    :param op_kwargs: a dictionary of keyword arguments that will get unpacked
+        in your function
+    :param op_args: a list of positional arguments that will get unpacked when
+        calling your callable
+    :param templates_dict: a dictionary where the values are templates that
+        will get templated by the Airflow engine sometime between
+        ``__init__`` and ``execute`` takes place and are made available
+        in your callable's context after the template has been applied.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/operator:PythonSensor`
+    """
+
+    template_fields: Sequence[str] = ("templates_dict", "op_args", "op_kwargs")
+
+    def __init__(
+        self,
+        *,
+        python_callable: Callable,
+        op_args: list | None = None,
+        op_kwargs: Mapping[str, Any] | None = None,
+        templates_dict: dict | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.python_callable = python_callable
+        self.op_args = op_args or []
+        self.op_kwargs = op_kwargs or {}
+        self.templates_dict = templates_dict
+
+    def poke(self, context: Context) -> PokeReturnValue | bool:
+        context_merge(context, self.op_kwargs, templates_dict=self.templates_dict)
+        self.op_kwargs = determine_kwargs(self.python_callable, self.op_args, context)
+
+        self.log.info("Poking callable: %s", str(self.python_callable))
+        return_value = self.python_callable(*self.op_args, **self.op_kwargs)
+        if isinstance(return_value, PokeReturnValue):
+            return return_value
+        return PokeReturnValue(bool(return_value))

airflow/providers/standard/sensors/time.py

@@ -18,15 +18,36 @@
 from __future__ import annotations
 
 import datetime
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, NoReturn
 
+from airflow.providers.standard.triggers.temporal import DateTimeTrigger
+from airflow.providers.standard.version_compat import AIRFLOW_V_2_10_PLUS
 from airflow.sensors.base import BaseSensorOperator
-from airflow.triggers.base import StartTriggerArgs
-from airflow.triggers.temporal import DateTimeTrigger
+
+try:
+    from airflow.triggers.base import StartTriggerArgs
+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.utils import timezone
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    try:
+        from airflow.sdk.definitions.context import Context
+    except ImportError:
+        # TODO: Remove once provider drops support for Airflow 2
+        from airflow.utils.context import Context
 
 
 class TimeSensor(BaseSensorOperator):
@@ -68,7 +89,7 @@ class TimeSensorAsync(BaseSensorOperator):
     """
 
     start_trigger_args = StartTriggerArgs(
-        trigger_cls="airflow.triggers.temporal.DateTimeTrigger",
+        trigger_cls="airflow.providers.standard.triggers.temporal.DateTimeTrigger",
         trigger_kwargs={"moment": "", "end_from_trigger": False},
         next_method="execute_complete",
         next_kwargs=None,
@@ -102,7 +123,9 @@ class TimeSensorAsync(BaseSensorOperator):
 
     def execute(self, context: Context) -> NoReturn:
         self.defer(
-            trigger=DateTimeTrigger(moment=self.target_datetime, end_from_trigger=self.end_from_trigger),
+            trigger=DateTimeTrigger(moment=self.target_datetime, end_from_trigger=self.end_from_trigger)
+            if AIRFLOW_V_2_10_PLUS
+            else DateTimeTrigger(moment=self.target_datetime),
             method_name="execute_complete",
         )
 

airflow/providers/standard/sensors/time_delta.py

@@ -17,41 +17,77 @@
 # under the License.
 from __future__ import annotations
 
-from datetime import timedelta
+from datetime import datetime, timedelta
 from time import sleep
 from typing import TYPE_CHECKING, Any, NoReturn
 
+from packaging.version import Version
+
 from airflow.configuration import conf
 from airflow.exceptions import AirflowSkipException
+from airflow.providers.standard.triggers.temporal import DateTimeTrigger, TimeDeltaTrigger
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.sensors.base import BaseSensorOperator
-from airflow.triggers.temporal import DateTimeTrigger, TimeDeltaTrigger
 from airflow.utils import timezone
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    try:
+        from airflow.sdk.definitions.context import Context
+    except ImportError:
+        # TODO: Remove once provider drops support for Airflow 2
+        from airflow.utils.context import Context
+
+
+def _get_airflow_version():
+    from airflow import __version__ as airflow_version
+
+    return Version(Version(airflow_version).base_version)
 
 
 class TimeDeltaSensor(BaseSensorOperator):
     """
-    Waits for a timedelta after the run's data interval.
+    Waits for a timedelta.
 
-    :param delta: time length to wait after the data interval before succeeding.
+    The delta will be evaluated against data_interval_end if present for the dag run,
+    otherwise run_after will be used.
+
+    :param delta: time to wait before succeeding.
 
     .. seealso::
         For more information on how to use this sensor, take a look at the guide:
         :ref:`howto/operator:TimeDeltaSensor`
 
-
     """
 
     def __init__(self, *, delta, **kwargs):
         super().__init__(**kwargs)
         self.delta = delta
 
-    def poke(self, context: Context):
-        target_dttm = context["data_interval_end"]
-        target_dttm += self.delta
-        self.log.info("Checking if the time (%s) has come", target_dttm)
+    def _derive_base_time(self, context: Context) -> datetime:
+        """
+        Get the "base time" against which the delta should be calculated.
+
+        If data_interval_end is populated, use it; else use run_after.
+        """
+        data_interval_end = context.get("data_interval_end")
+        if data_interval_end:
+            if not isinstance(data_interval_end, datetime):
+                raise ValueError("`data_interval_end` returned non-datetime object")
+
+            return data_interval_end
+
+        if not data_interval_end and not AIRFLOW_V_3_0_PLUS:
+            raise ValueError("`data_interval_end` not found in task context.")
+
+        dag_run = context.get("dag_run")
+        if not dag_run:
+            raise ValueError("`dag_run` not found in task context")
+        return dag_run.run_after
+
+    def poke(self, context: Context) -> bool:
+        base_time = self._derive_base_time(context=context)
+        target_dttm = base_time + self.delta
+        self.log.info("Checking if the delta has elapsed base_time=%s, delta=%s", base_time, self.delta)
         return timezone.utcnow() > target_dttm
 
 
@@ -75,19 +111,34 @@ class TimeDeltaSensorAsync(TimeDeltaSensor):
         self.end_from_trigger = end_from_trigger
 
     def execute(self, context: Context) -> bool | NoReturn:
-        target_dttm = context["data_interval_end"]
-        target_dttm += self.delta
+        base_time = self._derive_base_time(context=context)
+        target_dttm: datetime = base_time + self.delta
+
         if timezone.utcnow() > target_dttm:
             # If the target datetime is in the past, return immediately
             return True
         try:
-            trigger = DateTimeTrigger(moment=target_dttm, end_from_trigger=self.end_from_trigger)
+            if AIRFLOW_V_3_0_PLUS:
+                trigger = DateTimeTrigger(moment=target_dttm, end_from_trigger=self.end_from_trigger)
+            else:
+                trigger = DateTimeTrigger(moment=target_dttm)
         except (TypeError, ValueError) as e:
             if self.soft_fail:
                 raise AirflowSkipException("Skipping due to soft_fail is set to True.") from e
             raise
 
-        self.defer(trigger=trigger, method_name="execute_complete")
+        # todo: remove backcompat when min airflow version greater than 2.11
+        timeout: int | float | timedelta
+        if _get_airflow_version() >= Version("2.11.0"):
+            timeout = self.timeout
+        else:
+            timeout = timedelta(seconds=self.timeout)
+
+        self.defer(
+            trigger=trigger,
+            method_name="execute_complete",
+            timeout=timeout,
+        )
 
     def execute_complete(self, context: Context, event: Any = None) -> None:
         """Handle the event when the trigger fires and return immediately."""
@@ -121,7 +172,9 @@ class WaitSensor(BaseSensorOperator):
     def execute(self, context: Context) -> None:
         if self.deferrable:
             self.defer(
-                trigger=TimeDeltaTrigger(self.time_to_wait, end_from_trigger=True),
+                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/sensors/weekday.py

@@ -17,14 +17,19 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Iterable
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
 
+from airflow.providers.standard.utils.weekday import WeekDay
 from airflow.sensors.base import BaseSensorOperator
 from airflow.utils import timezone
-from airflow.utils.weekday import WeekDay
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    try:
+        from airflow.sdk.definitions.context import Context
+    except ImportError:
+        # TODO: Remove once provider drops support for Airflow 2
+        from airflow.utils.context import Context
 
 
 class DayOfWeekSensor(BaseSensorOperator):
@@ -49,7 +54,7 @@ class DayOfWeekSensor(BaseSensorOperator):
     **Example** (with :class:`~airflow.utils.weekday.WeekDay` enum): ::
 
         # import WeekDay Enum
-        from airflow.utils.weekday import WeekDay
+        from airflow.providers.standard.utils.weekday import WeekDay
 
         weekend_check = DayOfWeekSensor(
             task_id="weekend_check",
@@ -98,7 +103,20 @@ class DayOfWeekSensor(BaseSensorOperator):
             self.week_day,
             WeekDay(timezone.utcnow().isoweekday()).name,
         )
+
         if self.use_task_logical_date:
-            return context["logical_date"].isoweekday() in self._week_day_num
-        else:
-            return timezone.utcnow().isoweekday() in self._week_day_num
+            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 `run_after` should be provided in the task context when "
+                    "`use_task_logical_date` is True"
+                )
+
+            determined_weekday_num = (
+                logical_date.isoweekday() if logical_date else dag_run.run_after.isoweekday()  # type: ignore[union-attr]
+            )
+
+            return determined_weekday_num in self._week_day_num
+        return timezone.utcnow().isoweekday() in self._week_day_num

airflow/providers/standard/triggers/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.