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

airflow/providers/standard/sensors/time_delta.py

@@ -0,0 +1,134 @@
+#
+# 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 datetime import timedelta
+from time import sleep
+from typing import TYPE_CHECKING, Any, NoReturn
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowSkipException
+from airflow.providers.standard.utils.version_references 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
+
+
+class TimeDeltaSensor(BaseSensorOperator):
+    """
+    Waits for a timedelta after the run's data interval.
+
+    :param delta: time length to wait after the data interval 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)
+        return timezone.utcnow() > target_dttm
+
+
+class TimeDeltaSensorAsync(TimeDeltaSensor):
+    """
+    A deferrable drop-in replacement for TimeDeltaSensor.
+
+    Will defers itself to avoid taking up a worker slot while it is waiting.
+
+    :param delta: time length to wait after the data interval before succeeding.
+    :param end_from_trigger: End the task directly from the triggerer without going into the worker.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/operator:TimeDeltaSensorAsync`
+
+    """
+
+    def __init__(self, *, end_from_trigger: bool = False, delta, **kwargs) -> None:
+        super().__init__(delta=delta, **kwargs)
+        self.end_from_trigger = end_from_trigger
+
+    def execute(self, context: Context) -> bool | NoReturn:
+        target_dttm = context["data_interval_end"]
+        target_dttm += self.delta
+        if timezone.utcnow() > target_dttm:
+            # If the target datetime is in the past, return immediately
+            return True
+        try:
+            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")
+
+    def execute_complete(self, context: Context, event: Any = None) -> None:
+        """Handle the event when the trigger fires and return immediately."""
+        return None
+
+
+class WaitSensor(BaseSensorOperator):
+    """
+    A sensor that waits a specified period of time before completing.
+
+    This differs from TimeDeltaSensor because the time to wait is measured from the start of the task, not
+    the data_interval_end of the DAG run.
+
+    :param time_to_wait: time length to wait after the task starts before succeeding.
+    :param deferrable: Run sensor in deferrable mode
+    """
+
+    def __init__(
+        self,
+        time_to_wait: timedelta | int,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.deferrable = deferrable
+        if isinstance(time_to_wait, int):
+            self.time_to_wait = timedelta(minutes=time_to_wait)
+        else:
+            self.time_to_wait = time_to_wait
+
+    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),
+                method_name="execute_complete",
+            )
+        else:
+            sleep(int(self.time_to_wait.total_seconds()))

airflow/providers/standard/sensors/weekday.py

@@ -0,0 +1,104 @@
+#
+# 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.sensors.base import BaseSensorOperator
+from airflow.utils import timezone
+from airflow.utils.weekday import WeekDay
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class DayOfWeekSensor(BaseSensorOperator):
+    """
+    Waits until the first specified day of the week.
+
+    For example, if the execution day of the task is '2018-12-22' (Saturday)
+    and you pass 'FRIDAY', the task will wait until next Friday.
+
+    **Example** (with single day): ::
+
+        weekend_check = DayOfWeekSensor(
+            task_id="weekend_check", week_day="Saturday", use_task_logical_date=True, dag=dag
+        )
+
+    **Example** (with multiple day using set): ::
+
+        weekend_check = DayOfWeekSensor(
+            task_id="weekend_check", week_day={"Saturday", "Sunday"}, use_task_logical_date=True, dag=dag
+        )
+
+    **Example** (with :class:`~airflow.utils.weekday.WeekDay` enum): ::
+
+        # import WeekDay Enum
+        from airflow.utils.weekday import WeekDay
+
+        weekend_check = DayOfWeekSensor(
+            task_id="weekend_check",
+            week_day={WeekDay.SATURDAY, WeekDay.SUNDAY},
+            use_task_logical_date=True,
+            dag=dag,
+        )
+
+    :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 week_day. Execution Date is Useful for backfilling.
+        If ``False``, uses system's day of the week. Useful when you
+        don't want to run anything on weekdays on the system.
+
+    .. seealso::
+        For more information on how to use this sensor, take a look at the guide:
+        :ref:`howto/operator:DayOfWeekSensor`
+
+    """
+
+    def __init__(
+        self,
+        *,
+        week_day: str | Iterable[str] | WeekDay | Iterable[WeekDay],
+        use_task_logical_date: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        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 poke(self, context: Context) -> bool:
+        self.log.info(
+            "Poking until weekday is in %s, Today is %s",
+            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

airflow/providers/standard/utils/__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/utils/python_virtualenv.py

@@ -0,0 +1,209 @@
+#
+# 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.
+"""Utilities for creating a virtual environment."""
+
+from __future__ import annotations
+
+import os
+import shutil
+import sys
+from pathlib import Path
+
+import jinja2
+from jinja2 import select_autoescape
+
+from airflow.configuration import conf
+from airflow.utils.process_utils import execute_in_subprocess
+
+
+def _is_uv_installed() -> bool:
+    """
+    Verify whether the uv tool is installed by checking if it's included in the system PATH or installed as a package.
+
+    :return: True if it is. Whichever way of checking it works, is fine.
+    """
+    return bool(shutil.which("uv"))
+
+
+def _use_uv() -> bool:
+    """
+    Check if the uv tool should be used.
+
+    :return: True if uv should be used.
+    """
+    venv_install_method = conf.get("standard", "venv_install_method", fallback="auto").lower()
+    if venv_install_method == "auto":
+        return _is_uv_installed()
+    elif venv_install_method == "uv":
+        return True
+    return False
+
+
+def _generate_uv_cmd(tmp_dir: str, python_bin: str, system_site_packages: bool) -> list[str]:
+    """Build the command to install the venv via UV."""
+    cmd = ["uv", "venv", "--allow-existing", "--seed"]
+    if python_bin is not None:
+        cmd += ["--python", python_bin]
+    if system_site_packages:
+        cmd.append("--system-site-packages")
+    cmd.append(tmp_dir)
+    return cmd
+
+
+def _generate_venv_cmd(tmp_dir: str, python_bin: str, system_site_packages: bool) -> list[str]:
+    """We are using venv command instead of venv module to allow creation of venv for different python versions."""
+    if python_bin is None:
+        python_bin = sys.executable
+    cmd = [python_bin, "-m", "venv", tmp_dir]
+    if system_site_packages:
+        cmd.append("--system-site-packages")
+    return cmd
+
+
+def _generate_uv_install_cmd_from_file(
+    tmp_dir: str, requirements_file_path: str, pip_install_options: list[str]
+) -> list[str]:
+    return [
+        "uv",
+        "pip",
+        "install",
+        "--python",
+        f"{tmp_dir}/bin/python",
+        *pip_install_options,
+        "-r",
+        requirements_file_path,
+    ]
+
+
+def _generate_pip_install_cmd_from_file(
+    tmp_dir: str, requirements_file_path: str, pip_install_options: list[str]
+) -> list[str]:
+    return [f"{tmp_dir}/bin/pip", "install", *pip_install_options, "-r", requirements_file_path]
+
+
+def _generate_uv_install_cmd_from_list(
+    tmp_dir: str, requirements: list[str], pip_install_options: list[str]
+) -> list[str]:
+    return ["uv", "pip", "install", "--python", f"{tmp_dir}/bin/python", *pip_install_options, *requirements]
+
+
+def _generate_pip_install_cmd_from_list(
+    tmp_dir: str, requirements: list[str], pip_install_options: list[str]
+) -> list[str]:
+    return [f"{tmp_dir}/bin/pip", "install", *pip_install_options, *requirements]
+
+
+def _generate_pip_conf(conf_file: Path, index_urls: list[str]) -> None:
+    if index_urls:
+        pip_conf_options = f"index-url = {index_urls[0]}"
+        if len(index_urls) > 1:
+            pip_conf_options += f"\nextra-index-url = {' '.join(x for x in index_urls[1:])}"
+    else:
+        pip_conf_options = "no-index = true"
+    conf_file.write_text(f"[global]\n{pip_conf_options}")
+
+
+def prepare_virtualenv(
+    venv_directory: str,
+    python_bin: str,
+    system_site_packages: bool,
+    requirements: list[str] | None = None,
+    requirements_file_path: str | None = None,
+    pip_install_options: list[str] | None = None,
+    index_urls: list[str] | None = None,
+) -> str:
+    """
+    Create a virtual environment and install the additional python packages.
+
+    :param venv_directory: The path for directory where the environment will be created.
+    :param python_bin: Path to the Python executable.
+    :param system_site_packages: Whether to include system_site_packages in your virtualenv.
+        See virtualenv documentation for more information.
+    :param requirements: List of additional python packages.
+    :param requirements_file_path: Path to the ``requirements.txt`` file.
+    :param pip_install_options: a list of pip install options when installing requirements
+        See 'pip install -h' for available options
+    :param index_urls: an optional list of index urls to load Python packages from.
+        If not provided the system pip conf will be used to source packages from.
+    :return: Path to a binary file with Python in a virtual environment.
+    """
+    if pip_install_options is None:
+        pip_install_options = []
+
+    if requirements is not None and requirements_file_path is not None:
+        raise ValueError("Either requirements OR requirements_file_path has to be passed, but not both")
+
+    if index_urls is not None:
+        _generate_pip_conf(Path(venv_directory) / "pip.conf", index_urls)
+
+    if _use_uv():
+        venv_cmd = _generate_uv_cmd(venv_directory, python_bin, system_site_packages)
+    else:
+        venv_cmd = _generate_venv_cmd(venv_directory, python_bin, system_site_packages)
+    execute_in_subprocess(venv_cmd)
+
+    pip_cmd = None
+    if requirements is not None and len(requirements) != 0:
+        if _use_uv():
+            pip_cmd = _generate_uv_install_cmd_from_list(venv_directory, requirements, pip_install_options)
+        else:
+            pip_cmd = _generate_pip_install_cmd_from_list(venv_directory, requirements, pip_install_options)
+    if requirements_file_path is not None and requirements_file_path:
+        if _use_uv():
+            pip_cmd = _generate_uv_install_cmd_from_file(
+                venv_directory, requirements_file_path, pip_install_options
+            )
+        else:
+            pip_cmd = _generate_pip_install_cmd_from_file(
+                venv_directory, requirements_file_path, pip_install_options
+            )
+
+    if pip_cmd:
+        execute_in_subprocess(pip_cmd)
+
+    return f"{venv_directory}/bin/python"
+
+
+def write_python_script(
+    jinja_context: dict,
+    filename: str,
+    render_template_as_native_obj: bool = False,
+):
+    """
+    Render the python script to a file to execute in the virtual environment.
+
+    :param jinja_context: The jinja context variables to unpack and replace with its placeholders in the
+        template file.
+    :param filename: The name of the file to dump the rendered script to.
+    :param render_template_as_native_obj: If ``True``, rendered Jinja template would be converted
+        to a native Python object
+    """
+    template_loader = jinja2.FileSystemLoader(searchpath=os.path.dirname(__file__))
+    template_env: jinja2.Environment
+    if render_template_as_native_obj:
+        template_env = jinja2.nativetypes.NativeEnvironment(
+            loader=template_loader, undefined=jinja2.StrictUndefined
+        )
+    else:
+        template_env = jinja2.Environment(
+            loader=template_loader,
+            undefined=jinja2.StrictUndefined,
+            autoescape=select_autoescape(["html", "xml"]),
+        )
+    template = template_env.get_template("python_virtualenv_script.jinja2")
+    template.stream(**jinja_context).dump(filename)

airflow/providers/standard/utils/python_virtualenv_script.jinja2

@@ -0,0 +1,101 @@
+{#
+ 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 {{ pickling_library }}
+import sys
+
+{% if expect_airflow %}
+ {# Check whether Airflow is available in the environment.
+ # If it is, we'll want to ensure that we integrate any macros that are being provided
+ # by plugins prior to unpickling the task context. #}
+if sys.version_info >= (3,6):
+    try:
+        from airflow.plugins_manager import integrate_macros_plugins
+        integrate_macros_plugins()
+    except ImportError:
+        {# Airflow is not available in this environment, therefore we won't
+         # be able to integrate any plugin macros. #}
+        pass
+{% endif %}
+
+# Script
+{{ python_callable_source }}
+
+# monkey patching for the cases when python_callable is part of the dag module.
+{% if modified_dag_module_name is defined %}
+
+import types
+
+{{ modified_dag_module_name }}  = types.ModuleType("{{ modified_dag_module_name }}")
+
+{{ modified_dag_module_name }}.{{ python_callable }} = {{ python_callable }}
+
+sys.modules["{{modified_dag_module_name}}"] = {{modified_dag_module_name}}
+
+{% endif%}
+
+{% if op_args or op_kwargs %}
+with open(sys.argv[1], "rb") as file:
+    arg_dict = {{ pickling_library }}.load(file)
+{% else %}
+arg_dict = {"args": [], "kwargs": {}}
+{% endif %}
+
+{% if string_args_global | default(true) -%}
+# Read string args
+with open(sys.argv[3], "r") as file:
+    virtualenv_string_args = list(map(lambda x: x.strip(), list(file)))
+{% endif %}
+
+{% if use_airflow_context | default(false) -%}
+if len(sys.argv) > 5:
+    import json
+    from types import ModuleType
+
+    from airflow.providers.standard.operators import python as airflow_python
+    from airflow.serialization.serialized_objects import BaseSerialization
+
+
+    class _MockPython(ModuleType):
+        @staticmethod
+        def get_current_context():
+            with open(sys.argv[5]) as file:
+                context = json.load(file)
+                return BaseSerialization.deserialize(context, use_pydantic_models=True)
+
+        def __getattr__(self, name: str):
+            return getattr(airflow_python, name)
+
+
+    MockPython = _MockPython("MockPython")
+    sys.modules["airflow.providers.standard.operators.python"] = MockPython
+{% endif %}
+
+try:
+    res = {{ python_callable }}(*arg_dict["args"], **arg_dict["kwargs"])
+except Exception as e:
+    with open(sys.argv[4], "w") as file:
+        file.write(str(e))
+    raise
+
+# Write output
+with open(sys.argv[2], "wb") as file:
+    if res is not None:
+        {{ pickling_library }}.dump(res, file)

airflow/providers/standard/utils/version_references.py

@@ -0,0 +1,26 @@
+#
+# 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 packaging.version import Version
+
+from airflow import __version__ as airflow_version
+
+AIRFLOW_VERSION = Version(airflow_version)
+AIRFLOW_V_2_10_PLUS = Version(AIRFLOW_VERSION.base_version) >= Version("2.10.0")
+AIRFLOW_V_3_0_PLUS = Version(AIRFLOW_VERSION.base_version) >= Version("3.0.0")