apache-airflow-providers-standard 0.1.0rc1__py3-none-any.whl → 1.0.0.dev0__py3-none-any.whl

airflow/providers/standard/sensors/python.py

@@ -1,85 +0,0 @@
-#
-# 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
-        else:
-            return PokeReturnValue(bool(return_value))

airflow/providers/standard/triggers/__init__.py

@@ -1,16 +0,0 @@
-# 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/triggers/external_task.py

@@ -1,211 +0,0 @@
-# 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 asyncio
-import typing
-from typing import Any
-
-from asgiref.sync import sync_to_async
-from sqlalchemy import func
-
-from airflow.models import DagRun
-from airflow.providers.standard.utils.sensor_helper import _get_count
-from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
-from airflow.triggers.base import BaseTrigger, TriggerEvent
-from airflow.utils.session import NEW_SESSION, provide_session
-
-if typing.TYPE_CHECKING:
-    from datetime import datetime
-
-    from sqlalchemy.orm import Session
-
-    from airflow.utils.state import DagRunState
-
-
-class WorkflowTrigger(BaseTrigger):
-    """
-    A trigger to monitor tasks, task group and dag execution in Apache Airflow.
-
-    :param external_dag_id: The ID of the external dag.
-    :param run_ids: A list of run ids for the external dag.
-    :param external_task_ids: A collection of external task IDs to wait for.
-    :param external_task_group_id: The ID of the external task group to wait for.
-    :param failed_states: States considered as failed for external tasks.
-    :param skipped_states: States considered as skipped for external tasks.
-    :param allowed_states: States considered as successful for external tasks.
-    :param poke_interval: The interval (in seconds) for poking the external tasks.
-    :param soft_fail: If True, the trigger will not fail the entire dag on external task failure.
-    """
-
-    def __init__(
-        self,
-        external_dag_id: str,
-        run_ids: list[str] | None = None,
-        execution_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,
-        poke_interval: float = 2.0,
-        soft_fail: bool = False,
-        **kwargs,
-    ):
-        self.external_dag_id = external_dag_id
-        self.external_task_ids = external_task_ids
-        self.external_task_group_id = external_task_group_id
-        self.failed_states = failed_states
-        self.skipped_states = skipped_states
-        self.allowed_states = allowed_states
-        self.run_ids = run_ids
-        self.poke_interval = poke_interval
-        self.soft_fail = soft_fail
-        self.execution_dates = execution_dates
-        super().__init__(**kwargs)
-
-    def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serialize the trigger param and module path."""
-        data: dict[str, typing.Any] = {
-            "external_dag_id": self.external_dag_id,
-            "external_task_ids": self.external_task_ids,
-            "external_task_group_id": self.external_task_group_id,
-            "failed_states": self.failed_states,
-            "skipped_states": self.skipped_states,
-            "allowed_states": self.allowed_states,
-            "poke_interval": self.poke_interval,
-            "soft_fail": self.soft_fail,
-        }
-        if AIRFLOW_V_3_0_PLUS:
-            data["run_ids"] = self.run_ids
-        else:
-            data["execution_dates"] = self.execution_dates
-
-        return "airflow.providers.standard.triggers.external_task.WorkflowTrigger", data
-
-    async def run(self) -> typing.AsyncIterator[TriggerEvent]:
-        """Check periodically tasks, task group or dag status."""
-        while True:
-            if self.failed_states:
-                failed_count = await self._get_count(self.failed_states)
-                if failed_count > 0:
-                    yield TriggerEvent({"status": "failed"})
-                    return
-                else:
-                    yield TriggerEvent({"status": "success"})
-                    return
-            if self.skipped_states:
-                skipped_count = await self._get_count(self.skipped_states)
-                if skipped_count > 0:
-                    yield TriggerEvent({"status": "skipped"})
-                    return
-            allowed_count = await self._get_count(self.allowed_states)
-            _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates
-            if allowed_count == len(_dates):  # type: ignore[arg-type]
-                yield TriggerEvent({"status": "success"})
-                return
-            self.log.info("Sleeping for %s seconds", self.poke_interval)
-            await asyncio.sleep(self.poke_interval)
-
-    @sync_to_async
-    def _get_count(self, states: typing.Iterable[str] | None) -> int:
-        """
-        Get the count of records against dttm filter and states. Async wrapper for _get_count.
-
-        :param states: task or dag states
-        :return The count of records.
-        """
-        return _get_count(
-            dttm_filter=self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates,
-            external_task_ids=self.external_task_ids,
-            external_task_group_id=self.external_task_group_id,
-            external_dag_id=self.external_dag_id,
-            states=states,
-        )
-
-
-class DagStateTrigger(BaseTrigger):
-    """
-    Waits asynchronously for a dag to complete for a specific run_id.
-
-    :param dag_id: The dag_id that contains the task you want to wait for
-    :param states: allowed states, default is ``['success']``
-    :param run_ids: The run_id of dag run.
-    :param poll_interval: The time interval in seconds to check the state.
-        The default value is 5.0 sec.
-    """
-
-    def __init__(
-        self,
-        dag_id: str,
-        states: list[DagRunState],
-        run_ids: list[str] | None = None,
-        execution_dates: list[datetime] | None = None,
-        poll_interval: float = 5.0,
-    ):
-        super().__init__()
-        self.dag_id = dag_id
-        self.states = states
-        self.run_ids = run_ids
-        self.execution_dates = execution_dates
-        self.poll_interval = poll_interval
-
-    def serialize(self) -> tuple[str, dict[str, typing.Any]]:
-        """Serialize DagStateTrigger arguments and classpath."""
-        data = {
-            "dag_id": self.dag_id,
-            "states": self.states,
-            "poll_interval": self.poll_interval,
-        }
-
-        if AIRFLOW_V_3_0_PLUS:
-            data["run_ids"] = self.run_ids
-        else:
-            data["execution_dates"] = self.execution_dates
-
-        return "airflow.providers.standard.triggers.external_task.DagStateTrigger", data
-
-    async def run(self) -> typing.AsyncIterator[TriggerEvent]:
-        """Check periodically if the dag run exists, and has hit one of the states yet, or not."""
-        while True:
-            # mypy confuses typing here
-            num_dags = await self.count_dags()  # type: ignore[call-arg]
-            _dates = self.run_ids if AIRFLOW_V_3_0_PLUS else self.execution_dates
-            if num_dags == len(_dates):  # type: ignore[arg-type]
-                yield TriggerEvent(self.serialize())
-                return
-            await asyncio.sleep(self.poll_interval)
-
-    @sync_to_async
-    @provide_session
-    def count_dags(self, *, session: Session = NEW_SESSION) -> int | None:
-        """Count how many dag runs in the database match our criteria."""
-        _dag_run_date_condition = (
-            DagRun.run_id.in_(self.run_ids)
-            if AIRFLOW_V_3_0_PLUS
-            else DagRun.execution_date.in_(self.execution_dates)
-        )
-        count = (
-            session.query(func.count("*"))  # .count() is inefficient
-            .filter(
-                DagRun.dag_id == self.dag_id,
-                DagRun.state.in_(self.states),
-                _dag_run_date_condition,
-            )
-            .scalar()
-        )
-        return typing.cast(int, count)

airflow/providers/standard/triggers/file.py

@@ -1,131 +0,0 @@
-# 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 asyncio
-import datetime
-import os
-from collections.abc import AsyncIterator
-from glob import glob
-from typing import Any
-
-from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
-
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.triggers.base import BaseEventTrigger, BaseTrigger, TriggerEvent
-else:
-    from airflow.triggers.base import (  # type: ignore
-        BaseTrigger,
-        BaseTrigger as BaseEventTrigger,
-        TriggerEvent,
-    )
-
-
-class FileTrigger(BaseTrigger):
-    """
-    A trigger that fires exactly once after it finds the requested file or folder.
-
-    :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 poke_interval: Time that the job should wait in between each try
-    """
-
-    def __init__(
-        self,
-        filepath: str,
-        recursive: bool = False,
-        poke_interval: float = 5.0,
-        **kwargs,
-    ):
-        super().__init__()
-        self.filepath = filepath
-        self.recursive = recursive
-        self.poke_interval = poke_interval
-
-    def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serialize FileTrigger arguments and classpath."""
-        return (
-            "airflow.providers.standard.triggers.file.FileTrigger",
-            {
-                "filepath": self.filepath,
-                "recursive": self.recursive,
-                "poke_interval": self.poke_interval,
-            },
-        )
-
-    async def run(self) -> AsyncIterator[TriggerEvent]:
-        """Loop until the relevant files are found."""
-        while True:
-            for path in glob(self.filepath, recursive=self.recursive):
-                if os.path.isfile(path):
-                    mod_time_f = os.path.getmtime(path)
-                    mod_time = datetime.datetime.fromtimestamp(mod_time_f).strftime("%Y%m%d%H%M%S")
-                    self.log.info("Found File %s last modified: %s", path, mod_time)
-                    yield TriggerEvent(True)
-                    return
-                for _, _, files in os.walk(self.filepath):
-                    if files:
-                        yield TriggerEvent(True)
-                        return
-            await asyncio.sleep(self.poke_interval)
-
-
-class FileDeleteTrigger(BaseEventTrigger):
-    """
-    A trigger that fires exactly once after it finds the requested file and then delete the file.
-
-    The difference between ``FileTrigger`` and ``FileDeleteTrigger`` is ``FileDeleteTrigger`` can only find a
-    specific file.
-
-    :param filepath: File (relative to the base path set within the connection).
-    :param poke_interval: Time that the job should wait in between each try
-    """
-
-    def __init__(
-        self,
-        filepath: str,
-        poke_interval: float = 5.0,
-        **kwargs,
-    ):
-        super().__init__()
-        self.filepath = filepath
-        self.poke_interval = poke_interval
-
-    def serialize(self) -> tuple[str, dict[str, Any]]:
-        """Serialize FileDeleteTrigger arguments and classpath."""
-        return (
-            "airflow.providers.standard.triggers.file.FileDeleteTrigger",
-            {
-                "filepath": self.filepath,
-                "poke_interval": self.poke_interval,
-            },
-        )
-
-    async def run(self) -> AsyncIterator[TriggerEvent]:
-        """Loop until the relevant file is found."""
-        while True:
-            if os.path.isfile(self.filepath):
-                mod_time_f = os.path.getmtime(self.filepath)
-                mod_time = datetime.datetime.fromtimestamp(mod_time_f).strftime("%Y%m%d%H%M%S")
-                self.log.info("Found file %s last modified: %s", self.filepath, mod_time)
-                os.remove(self.filepath)
-                self.log.info("File %s has been deleted", self.filepath)
-                yield TriggerEvent(True)
-                return
-            await asyncio.sleep(self.poke_interval)

airflow/providers/standard/triggers/temporal.py

@@ -1,114 +0,0 @@
-# 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 asyncio
-import datetime
-from collections.abc import AsyncIterator
-from typing import Any
-
-import pendulum
-
-from airflow.exceptions import AirflowException
-from airflow.providers.standard.version_compat import AIRFLOW_V_2_10_PLUS
-from airflow.triggers.base import BaseTrigger, TriggerEvent
-from airflow.utils import timezone
-
-if AIRFLOW_V_2_10_PLUS:
-    from airflow.triggers.base import TaskSuccessEvent
-
-
-class DateTimeTrigger(BaseTrigger):
-    """
-    Trigger based on a datetime.
-
-    A trigger that fires exactly once, at the given datetime, give or take
-    a few seconds.
-
-    The provided datetime MUST be in UTC.
-
-    :param moment: when to yield event
-    :param end_from_trigger: whether the trigger should mark the task successful after time condition
-        reached or resume the task after time condition reached.
-    """
-
-    def __init__(self, moment: datetime.datetime, *, end_from_trigger: bool = False) -> None:
-        super().__init__()
-        if not isinstance(moment, datetime.datetime):
-            raise TypeError(f"Expected datetime.datetime type for moment. Got {type(moment)}")
-        # Make sure it's in UTC
-        elif moment.tzinfo is None:
-            raise ValueError("You cannot pass naive datetimes")
-        else:
-            self.moment: pendulum.DateTime = timezone.convert_to_utc(moment)
-        if not AIRFLOW_V_2_10_PLUS and end_from_trigger:
-            raise AirflowException("end_from_trigger is only supported in Airflow 2.10 and later. ")
-
-        self.end_from_trigger = end_from_trigger
-
-    def serialize(self) -> tuple[str, dict[str, Any]]:
-        return (
-            "airflow.providers.standard.triggers.temporal.DateTimeTrigger",
-            {"moment": self.moment, "end_from_trigger": self.end_from_trigger},
-        )
-
-    async def run(self) -> AsyncIterator[TriggerEvent]:
-        """
-        Loop until the relevant time is met.
-
-        We do have a two-phase delay to save some cycles, but sleeping is so
-        cheap anyway that it's pretty loose. We also don't just sleep for
-        "the number of seconds until the time" in case the system clock changes
-        unexpectedly, or handles a DST change poorly.
-        """
-        # Sleep in successively smaller increments starting from 1 hour down to 10 seconds at a time
-        self.log.info("trigger starting")
-        for step in 3600, 60, 10:
-            seconds_remaining = (self.moment - pendulum.instance(timezone.utcnow())).total_seconds()
-            while seconds_remaining > 2 * step:
-                self.log.info("%d seconds remaining; sleeping %s seconds", seconds_remaining, step)
-                await asyncio.sleep(step)
-                seconds_remaining = (self.moment - pendulum.instance(timezone.utcnow())).total_seconds()
-        # Sleep a second at a time otherwise
-        while self.moment > pendulum.instance(timezone.utcnow()):
-            self.log.info("sleeping 1 second...")
-            await asyncio.sleep(1)
-        if self.end_from_trigger:
-            self.log.info("Sensor time condition reached; marking task successful and exiting")
-            yield TaskSuccessEvent()
-        else:
-            self.log.info("yielding event with payload %r", self.moment)
-            yield TriggerEvent(self.moment)
-
-
-class TimeDeltaTrigger(DateTimeTrigger):
-    """
-    Create DateTimeTriggers based on delays.
-
-    Subclass to create DateTimeTriggers based on time delays rather
-    than exact moments.
-
-    While this is its own distinct class here, it will serialise to a
-    DateTimeTrigger class, since they're operationally the same.
-
-    :param delta: how long to wait
-    :param end_from_trigger: whether the trigger should mark the task successful after time condition
-        reached or resume the task after time condition reached.
-    """
-
-    def __init__(self, delta: datetime.timedelta, *, end_from_trigger: bool = False) -> None:
-        super().__init__(moment=timezone.utcnow() + delta, end_from_trigger=end_from_trigger)

airflow/providers/standard/utils/__init__.py

@@ -1,16 +0,0 @@
-# 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.