apache-airflow-providers-standard 0.0.1__py3-none-any.whl

airflow/providers/standard/operators/weekday.py

@@ -0,0 +1,120 @@
+#
+# 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 typing import TYPE_CHECKING, Iterable
+
+from airflow.operators.branch import BaseBranchOperator
+from airflow.utils import timezone
+from airflow.utils.weekday import WeekDay
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class BranchDayOfWeekOperator(BaseBranchOperator):
+    """
+    Branches into one of two lists of tasks depending on the current day.
+
+    For more information on how to use this operator, take a look at the guide:
+    :ref:`howto/operator:BranchDayOfWeekOperator`
+
+    **Example** (with single day):
+
+    .. code-block:: python
+
+        from airflow.operators.empty import EmptyOperator
+        from airflow.operators.weekday import BranchDayOfWeekOperator
+
+        monday = EmptyOperator(task_id="monday")
+        other_day = EmptyOperator(task_id="other_day")
+
+        monday_check = BranchDayOfWeekOperator(
+            task_id="monday_check",
+            week_day="Monday",
+            use_task_logical_date=True,
+            follow_task_ids_if_true="monday",
+            follow_task_ids_if_false="other_day",
+        )
+        monday_check >> [monday, other_day]
+
+    **Example** (with :class:`~airflow.utils.weekday.WeekDay` enum):
+
+    .. code-block:: python
+
+        # import WeekDay Enum
+        from airflow.utils.weekday import WeekDay
+        from airflow.operators.empty import EmptyOperator
+        from airflow.operators.weekday import BranchDayOfWeekOperator
+
+        workday = EmptyOperator(task_id="workday")
+        weekend = EmptyOperator(task_id="weekend")
+        weekend_check = BranchDayOfWeekOperator(
+            task_id="weekend_check",
+            week_day={WeekDay.SATURDAY, WeekDay.SUNDAY},
+            use_task_logical_date=True,
+            follow_task_ids_if_true="weekend",
+            follow_task_ids_if_false="workday",
+        )
+        # add downstream dependencies as you would do with any branch operator
+        weekend_check >> [workday, weekend]
+
+    :param follow_task_ids_if_true: task_id, task_group_id, or a list of task_ids and/or task_group_ids
+        to follow if criteria met.
+    :param follow_task_ids_if_false: task_id, task_group_id, or a list of task_ids and/or task_group_ids
+        to follow if criteria not met.
+    :param week_day: Day of the week to check (full name). Optionally, a set
+        of days can also be provided using a set. Example values:
+
+        * ``"MONDAY"``,
+        * ``{"Saturday", "Sunday"}``
+        * ``{WeekDay.TUESDAY}``
+        * ``{WeekDay.SATURDAY, WeekDay.SUNDAY}``
+
+        To use `WeekDay` enum, import it from `airflow.utils.weekday`
+
+    :param use_task_logical_date: If ``True``, uses task's logical date to compare
+        with is_today. Execution Date is Useful for backfilling.
+        If ``False``, uses system's day of the week.
+    """
+
+    def __init__(
+        self,
+        *,
+        follow_task_ids_if_true: str | Iterable[str],
+        follow_task_ids_if_false: str | Iterable[str],
+        week_day: str | Iterable[str] | WeekDay | Iterable[WeekDay],
+        use_task_logical_date: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.follow_task_ids_if_true = follow_task_ids_if_true
+        self.follow_task_ids_if_false = follow_task_ids_if_false
+        self.week_day = week_day
+        self.use_task_logical_date = use_task_logical_date
+        self._week_day_num = WeekDay.validate_week_day(week_day)
+
+    def choose_branch(self, context: Context) -> str | Iterable[str]:
+        if self.use_task_logical_date:
+            now = context["logical_date"]
+        else:
+            now = timezone.make_naive(timezone.utcnow(), self.dag.timezone)
+
+        if now.isoweekday() in self._week_day_num:
+            return self.follow_task_ids_if_true
+        return self.follow_task_ids_if_false

airflow/providers/standard/sensors/__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.

airflow/providers/standard/sensors/bash.py

@@ -0,0 +1,114 @@
+#
+# 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 os
+from subprocess import PIPE, STDOUT, Popen
+from tempfile import NamedTemporaryFile, TemporaryDirectory, gettempdir
+from typing import TYPE_CHECKING, Sequence
+
+from airflow.exceptions import AirflowFailException
+from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class BashSensor(BaseSensorOperator):
+    """
+    Executes a bash command/script.
+
+    Return True if and only if the return code is 0.
+
+    :param bash_command: The command, set of commands or reference to a
+        bash script (must be '.sh') to be executed.
+
+    :param env: If env is not None, it must be a mapping that defines the
+        environment variables for the new process; these are used instead
+        of inheriting the current process environment, which is the default
+        behavior. (templated)
+    :param output_encoding: output encoding of bash command.
+    :param retry_exit_code: If task exits with this code, treat the sensor
+        as not-yet-complete and retry the check later according to the
+        usual retry/timeout settings. Any other non-zero return code will
+        be treated as an error, and cause the sensor to fail. If set to
+        ``None`` (the default), any non-zero exit code will cause a retry
+        and the task will never raise an error except on time-out.
+
+    .. seealso::
+        For more information on how to use this sensor,take a look at the guide:
+        :ref:`howto/operator:BashSensor`
+    """
+
+    template_fields: Sequence[str] = ("bash_command", "env")
+
+    def __init__(
+        self, *, bash_command, env=None, output_encoding="utf-8", retry_exit_code: int | None = None, **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.bash_command = bash_command
+        self.env = env
+        self.output_encoding = output_encoding
+        self.retry_exit_code = retry_exit_code
+
+    def poke(self, context: Context):
+        """Execute the bash command in a temporary directory."""
+        bash_command = self.bash_command
+        self.log.info("Tmp dir root location: %s", gettempdir())
+        with TemporaryDirectory(prefix="airflowtmp") as tmp_dir, NamedTemporaryFile(
+            dir=tmp_dir, prefix=self.task_id
+        ) as f:
+            f.write(bytes(bash_command, "utf_8"))
+            f.flush()
+            fname = f.name
+            script_location = tmp_dir + "/" + fname
+            self.log.info("Temporary script location: %s", script_location)
+            self.log.info("Running command: %s", bash_command)
+
+            with Popen(
+                ["bash", fname],
+                stdout=PIPE,
+                stderr=STDOUT,
+                close_fds=True,
+                cwd=tmp_dir,
+                env=self.env,
+                preexec_fn=os.setsid,
+            ) as resp:
+                if resp.stdout:
+                    self.log.info("Output:")
+                    for line in iter(resp.stdout.readline, b""):
+                        self.log.info(line.decode(self.output_encoding).strip())
+                resp.wait()
+                self.log.info("Command exited with return code %s", resp.returncode)
+
+                # zero code means success, the sensor can go green
+                if resp.returncode == 0:
+                    return True
+
+                # we have a retry exit code, sensor retries if return code matches, otherwise error
+                elif self.retry_exit_code is not None:
+                    if resp.returncode == self.retry_exit_code:
+                        self.log.info("Return code matches retry code, will retry later")
+                        return False
+                    else:
+                        raise AirflowFailException(f"Command exited with return code {resp.returncode}")
+
+                # backwards compatibility: sensor retries no matter the error code
+                else:
+                    self.log.info("Non-zero return code and no retry code set, will retry later")
+                    return False

airflow/providers/standard/sensors/date_time.py

@@ -0,0 +1,152 @@
+#
+# 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
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, NoReturn, Sequence
+
+from airflow.providers.standard.utils.version_references import AIRFLOW_V_3_0_PLUS
+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
+
+
+from airflow.triggers.temporal import DateTimeTrigger
+from airflow.utils import timezone
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class DateTimeSensor(BaseSensorOperator):
+    """
+    Waits until the specified datetime.
+
+    A major advantage of this sensor is idempotence for the ``target_time``.
+    It handles some cases for which ``TimeSensor`` and ``TimeDeltaSensor`` are not suited.
+
+    **Example** 1 :
+        If a task needs to wait for 11am on each ``execution_date``. Using
+        ``TimeSensor`` or ``TimeDeltaSensor``, all backfill tasks started at
+        1am have to wait for 10 hours. This is unnecessary, e.g. a backfill
+        task with ``{{ ds }} = '1970-01-01'`` does not need to wait because
+        ``1970-01-01T11:00:00`` has already passed.
+
+    **Example** 2 :
+        If a DAG is scheduled to run at 23:00 daily, but one of the tasks is
+        required to run at 01:00 next day, using ``TimeSensor`` will return
+        ``True`` immediately because 23:00 > 01:00. Instead, we can do this:
+
+        .. code-block:: python
+
+            DateTimeSensor(
+                task_id="wait_for_0100",
+                target_time="{{ next_execution_date.tomorrow().replace(hour=1) }}",
+            )
+
+    :param target_time: datetime after which the job succeeds. (templated)
+    """
+
+    template_fields: Sequence[str] = ("target_time",)
+
+    def __init__(self, *, target_time: str | datetime.datetime, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+        # self.target_time can't be a datetime object as it is a template_field
+        if isinstance(target_time, datetime.datetime):
+            self.target_time = target_time.isoformat()
+        elif isinstance(target_time, str):
+            self.target_time = target_time
+        else:
+            raise TypeError(
+                f"Expected str or datetime.datetime type for target_time. Got {type(target_time)}"
+            )
+
+    def poke(self, context: Context) -> bool:
+        self.log.info("Checking if the time (%s) has come", self.target_time)
+        return timezone.utcnow() > timezone.parse(self.target_time)
+
+
+class DateTimeSensorAsync(DateTimeSensor):
+    """
+    Wait until the specified datetime occurs.
+
+    Deferring itself to avoid taking up a worker slot while it is waiting.
+    It is a drop-in replacement for DateTimeSensor.
+
+    :param target_time: datetime after which the job succeeds. (templated)
+    :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.
+    :param end_from_trigger: End the task directly from the triggerer without going into the worker.
+    """
+
+    start_trigger_args = StartTriggerArgs(
+        trigger_cls="airflow.triggers.temporal.DateTimeTrigger",
+        trigger_kwargs={"moment": "", "end_from_trigger": False},
+        next_method="execute_complete",
+        next_kwargs=None,
+        timeout=None,
+    )
+    start_from_trigger = False
+
+    def __init__(
+        self,
+        *,
+        start_from_trigger: bool = False,
+        end_from_trigger: bool = False,
+        trigger_kwargs: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.end_from_trigger = end_from_trigger
+
+        self.start_from_trigger = start_from_trigger
+        if self.start_from_trigger:
+            self.start_trigger_args.trigger_kwargs = dict(
+                moment=timezone.parse(self.target_time),
+                end_from_trigger=self.end_from_trigger,
+            )
+
+    def execute(self, context: Context) -> NoReturn:
+        self.defer(
+            method_name="execute_complete",
+            trigger=DateTimeTrigger(
+                moment=timezone.parse(self.target_time),
+                end_from_trigger=self.end_from_trigger,
+            )
+            if AIRFLOW_V_3_0_PLUS
+            else DateTimeTrigger(moment=timezone.parse(self.target_time)),
+        )
+
+    def execute_complete(self, context: Context, event: Any = None) -> None:
+        """Handle the event when the trigger fires and return immediately."""
+        return None

airflow/providers/standard/sensors/python.py

@@ -0,0 +1,80 @@
+#
+# 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 typing import TYPE_CHECKING, Any, Callable, Mapping, Sequence
+
+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:
+    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
+        else:
+            return PokeReturnValue(bool(return_value))

airflow/providers/standard/sensors/time.py

@@ -0,0 +1,130 @@
+#
+# 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
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, NoReturn
+
+from airflow.providers.standard.utils.version_references import AIRFLOW_V_2_10_PLUS
+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
+
+
+from airflow.triggers.temporal import DateTimeTrigger
+from airflow.utils import timezone
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class TimeSensor(BaseSensorOperator):
+    """
+    Waits until the specified time of the day.
+
+    :param target_time: time after which the job succeeds
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/operator:TimeSensor`
+
+    """
+
+    def __init__(self, *, target_time: datetime.time, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.target_time = target_time
+
+    def poke(self, context: Context) -> bool:
+        self.log.info("Checking if the time (%s) has come", self.target_time)
+        return timezone.make_naive(timezone.utcnow(), self.dag.timezone).time() > self.target_time
+
+
+class TimeSensorAsync(BaseSensorOperator):
+    """
+    Waits until the specified time of the day.
+
+    This frees up a worker slot while it is waiting.
+
+    :param target_time: time after which the job succeeds
+    :param start_from_trigger: Start the task directly from the triggerer without going into the worker.
+    :param end_from_trigger: End 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:TimeSensorAsync`
+    """
+
+    start_trigger_args = StartTriggerArgs(
+        trigger_cls="airflow.triggers.temporal.DateTimeTrigger",
+        trigger_kwargs={"moment": "", "end_from_trigger": False},
+        next_method="execute_complete",
+        next_kwargs=None,
+        timeout=None,
+    )
+    start_from_trigger = False
+
+    def __init__(
+        self,
+        *,
+        target_time: datetime.time,
+        start_from_trigger: bool = False,
+        trigger_kwargs: dict[str, Any] | None = None,
+        end_from_trigger: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.start_from_trigger = start_from_trigger
+        self.end_from_trigger = end_from_trigger
+        self.target_time = target_time
+
+        aware_time = timezone.coerce_datetime(
+            datetime.datetime.combine(datetime.datetime.today(), self.target_time, self.dag.timezone)
+        )
+
+        self.target_datetime = timezone.convert_to_utc(aware_time)
+        if self.start_from_trigger:
+            self.start_trigger_args.trigger_kwargs = dict(
+                moment=self.target_datetime, end_from_trigger=self.end_from_trigger
+            )
+
+    def execute(self, context: Context) -> NoReturn:
+        self.defer(
+            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",
+        )
+
+    def execute_complete(self, context: Context, event: Any = None) -> None:
+        """Handle the event when the trigger fires and return immediately."""
+        return None