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

airflow/providers/standard/triggers/external_task.py

@@ -0,0 +1,288 @@
+# 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
+
+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.
+    :param logical_dates: A list of logical dates for the external dag.
+    """
+
+    def __init__(
+        self,
+        external_dag_id: str,
+        run_ids: list[str] | None = None,
+        execution_dates: list[datetime] | None = None,
+        logical_dates: list[datetime] | None = None,
+        external_task_ids: typing.Collection[str] | None = None,
+        external_task_group_id: str | None = None,
+        failed_states: typing.Iterable[str] | None = None,
+        skipped_states: typing.Iterable[str] | None = None,
+        allowed_states: typing.Iterable[str] | None = None,
+        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
+        self.logical_dates = logical_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
+            data["logical_dates"] = self.logical_dates
+        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."""
+        if AIRFLOW_V_3_0_PLUS:
+            get_count_func = self._get_count_af_3
+            run_id_or_dates = (self.run_ids or self.logical_dates) or []
+        else:
+            get_count_func = self._get_count
+            run_id_or_dates = self.execution_dates or []
+
+        while True:
+            if self.failed_states:
+                failed_count = await get_count_func(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 get_count_func(self.skipped_states)
+                if skipped_count > 0:
+                    yield TriggerEvent({"status": "skipped"})
+                    return
+            allowed_count = await get_count_func(self.allowed_states)
+
+            if allowed_count == len(run_id_or_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)
+
+    async def _get_count_af_3(self, states):
+        from airflow.providers.standard.utils.sensor_helper import _get_count_by_matched_states
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        params = {
+            "dag_id": self.external_dag_id,
+            "logical_dates": self.logical_dates,
+            "run_ids": self.run_ids,
+        }
+        if self.external_task_ids:
+            count = await sync_to_async(RuntimeTaskInstance.get_ti_count)(
+                task_ids=self.external_task_ids,  # type: ignore[arg-type]
+                states=states,
+                **params,
+            )
+        elif self.external_task_group_id:
+            run_id_task_state_map = await sync_to_async(RuntimeTaskInstance.get_task_states)(
+                task_group_id=self.external_task_group_id,
+                **params,
+            )
+            count = await sync_to_async(_get_count_by_matched_states)(
+                run_id_task_state_map=run_id_task_state_map,
+                states=states,
+            )
+        else:
+            count = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+                dag_id=self.external_dag_id,
+                logical_dates=self.logical_dates,
+                run_ids=self.run_ids,
+                states=states,
+            )
+        if self.external_task_ids:
+            return count / len(self.external_task_ids)
+        return count
+
+    @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,
+            "run_ids": self.run_ids,
+            "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."""
+        runs_ids_or_dates = 0
+        if self.run_ids:
+            runs_ids_or_dates = len(self.run_ids)
+        elif self.execution_dates:
+            runs_ids_or_dates = len(self.execution_dates)
+
+        if AIRFLOW_V_3_0_PLUS:
+            event = await self.validate_count_dags_af_3(runs_ids_or_dates_len=runs_ids_or_dates)
+            yield TriggerEvent(event)
+            return
+        else:
+            while True:
+                num_dags = await self.count_dags()  # type: ignore[call-arg]
+                if num_dags == runs_ids_or_dates:
+                    yield TriggerEvent(self.serialize())
+                    return
+                await asyncio.sleep(self.poll_interval)
+
+    async def validate_count_dags_af_3(self, runs_ids_or_dates_len: int = 0) -> tuple[str, dict[str, Any]]:
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        cls_path, data = self.serialize()
+
+        while True:
+            num_dags = await sync_to_async(RuntimeTaskInstance.get_dr_count)(
+                dag_id=self.dag_id,
+                run_ids=self.run_ids,
+                states=self.states,  # type: ignore[arg-type]
+                logical_dates=self.execution_dates,
+            )
+            if num_dags == runs_ids_or_dates_len:
+                if isinstance(self.run_ids, list):
+                    for run_id in self.run_ids:
+                        state = await sync_to_async(RuntimeTaskInstance.get_dagrun_state)(
+                            dag_id=self.dag_id,
+                            run_id=run_id,
+                        )
+                        data[run_id] = state
+                        return cls_path, data
+            await asyncio.sleep(self.poll_interval)
+
+    if not AIRFLOW_V_3_0_PLUS:
+        from airflow.utils.session import NEW_SESSION, provide_session  # type: ignore[misc]
+
+        @sync_to_async
+        @provide_session
+        def count_dags(self, *, session: Session = NEW_SESSION) -> int:
+            """Count how many dag runs in the database match our criteria."""
+            _dag_run_date_condition = (
+                DagRun.run_id.in_(self.run_ids)
+                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

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

@@ -0,0 +1,113 @@
+# 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
+        if moment.tzinfo is None:
+            raise ValueError("You cannot pass naive datetimes")
+        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

@@ -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.