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

airflow/providers/standard/operators/trigger_dagrun.py

@@ -1,292 +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 datetime
-import json
-import time
-from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any
-
-from sqlalchemy import select
-from sqlalchemy.orm.exc import NoResultFound
-
-from airflow.api.common.trigger_dag import trigger_dag
-from airflow.configuration import conf
-from airflow.exceptions import (
-    AirflowException,
-    AirflowSkipException,
-    DagNotFound,
-    DagRunAlreadyExists,
-)
-from airflow.models import BaseOperator, BaseOperatorLink
-from airflow.models.dag import DagModel
-from airflow.models.dagbag import DagBag
-from airflow.models.dagrun import DagRun
-from airflow.models.xcom import XCom
-from airflow.providers.standard.triggers.external_task import DagStateTrigger
-from airflow.utils import timezone
-from airflow.utils.helpers import build_airflow_url_with_query
-from airflow.utils.session import provide_session
-from airflow.utils.state import DagRunState
-from airflow.utils.types import DagRunTriggeredByType, DagRunType
-
-XCOM_LOGICAL_DATE_ISO = "trigger_logical_date_iso"
-XCOM_RUN_ID = "trigger_run_id"
-
-
-if TYPE_CHECKING:
-    from sqlalchemy.orm.session import Session
-
-    from airflow.models.taskinstancekey import TaskInstanceKey
-
-    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 TriggerDagRunLink(BaseOperatorLink):
-    """
-    Operator link for TriggerDagRunOperator.
-
-    It allows users to access DAG triggered by task using TriggerDagRunOperator.
-    """
-
-    name = "Triggered DAG"
-
-    def get_link(self, operator: BaseOperator, *, ti_key: TaskInstanceKey) -> str:
-        from airflow.models.renderedtifields import RenderedTaskInstanceFields
-
-        if TYPE_CHECKING:
-            assert isinstance(operator, TriggerDagRunOperator)
-
-        if template_fields := RenderedTaskInstanceFields.get_templated_fields(ti_key):
-            trigger_dag_id: str = template_fields.get("trigger_dag_id", operator.trigger_dag_id)
-        else:
-            trigger_dag_id = operator.trigger_dag_id
-
-        # Fetch the correct dag_run_id for the triggerED dag which is
-        # stored in xcom during execution of the triggerING task.
-        triggered_dag_run_id = XCom.get_value(ti_key=ti_key, key=XCOM_RUN_ID)
-
-        query = {"dag_id": trigger_dag_id, "dag_run_id": triggered_dag_run_id}
-        return build_airflow_url_with_query(query)
-
-
-class TriggerDagRunOperator(BaseOperator):
-    """
-    Triggers a DAG run for a specified DAG ID.
-
-    Note that if database isolation mode is enabled, not all features are supported.
-
-    :param trigger_dag_id: The ``dag_id`` of the DAG to trigger (templated).
-    :param trigger_run_id: The run ID to use for the triggered DAG run (templated).
-        If not provided, a run ID will be automatically generated.
-    :param conf: Configuration for the DAG run (templated).
-    :param logical_date: Logical date for the triggered DAG (templated).
-    :param reset_dag_run: Whether clear existing DAG run if already exists.
-        This is useful when backfill or rerun an existing DAG run.
-        This only resets (not recreates) the DAG run.
-        DAG run conf is immutable and will not be reset on rerun of an existing DAG run.
-        When reset_dag_run=False and dag run exists, DagRunAlreadyExists will be raised.
-        When reset_dag_run=True and dag run exists, existing DAG run will be cleared to rerun.
-    :param wait_for_completion: Whether or not wait for DAG run completion. (default: False)
-    :param poke_interval: Poke interval to check DAG run status when wait_for_completion=True.
-        (default: 60)
-    :param allowed_states: Optional list of allowed DAG run states of the triggered DAG. This is useful when
-        setting ``wait_for_completion`` to True. Must be a valid DagRunState.
-        Default is ``[DagRunState.SUCCESS]``.
-    :param failed_states: Optional list of failed or disallowed DAG run states of the triggered DAG. This is
-        useful when setting ``wait_for_completion`` to True. Must be a valid DagRunState.
-        Default is ``[DagRunState.FAILED]``.
-    :param skip_when_already_exists: Set to true to mark the task as SKIPPED if a DAG run of the triggered
-        DAG for the same logical date already exists.
-    :param deferrable: If waiting for completion, whether or not to defer the task until done,
-        default is ``False``.
-    """
-
-    template_fields: Sequence[str] = (
-        "trigger_dag_id",
-        "trigger_run_id",
-        "logical_date",
-        "conf",
-        "wait_for_completion",
-        "skip_when_already_exists",
-    )
-    template_fields_renderers = {"conf": "py"}
-    ui_color = "#ffefeb"
-    operator_extra_links = [TriggerDagRunLink()]
-
-    def __init__(
-        self,
-        *,
-        trigger_dag_id: str,
-        trigger_run_id: str | None = None,
-        conf: dict | None = None,
-        logical_date: str | datetime.datetime | None = None,
-        reset_dag_run: bool = False,
-        wait_for_completion: bool = False,
-        poke_interval: int = 60,
-        allowed_states: list[str | DagRunState] | None = None,
-        failed_states: list[str | DagRunState] | None = None,
-        skip_when_already_exists: bool = False,
-        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
-        **kwargs,
-    ) -> None:
-        super().__init__(**kwargs)
-        self.trigger_dag_id = trigger_dag_id
-        self.trigger_run_id = trigger_run_id
-        self.conf = conf
-        self.reset_dag_run = reset_dag_run
-        self.wait_for_completion = wait_for_completion
-        self.poke_interval = poke_interval
-        if allowed_states:
-            self.allowed_states = [DagRunState(s) for s in allowed_states]
-        else:
-            self.allowed_states = [DagRunState.SUCCESS]
-        if failed_states or failed_states == []:
-            self.failed_states = [DagRunState(s) for s in failed_states]
-        else:
-            self.failed_states = [DagRunState.FAILED]
-        self.skip_when_already_exists = skip_when_already_exists
-        self._defer = deferrable
-
-        if logical_date is not None and not isinstance(logical_date, (str, datetime.datetime)):
-            type_name = type(logical_date).__name__
-            raise TypeError(
-                f"Expected str or datetime.datetime type for parameter 'logical_date'. Got {type_name}"
-            )
-
-        self.logical_date = logical_date
-
-    def execute(self, context: Context):
-        if self.logical_date is None or isinstance(self.logical_date, datetime.datetime):
-            parsed_logical_date = self.logical_date
-        else:
-            parsed_logical_date = timezone.parse(self.logical_date)
-
-        try:
-            json.dumps(self.conf)
-        except TypeError:
-            raise AirflowException("conf parameter should be JSON Serializable")
-
-        if self.trigger_run_id:
-            run_id = str(self.trigger_run_id)
-        else:
-            run_id = DagRun.generate_run_id(
-                run_type=DagRunType.MANUAL,
-                logical_date=parsed_logical_date,
-                run_after=parsed_logical_date or timezone.utcnow(),
-            )
-
-        try:
-            dag_run = trigger_dag(
-                dag_id=self.trigger_dag_id,
-                run_id=run_id,
-                conf=self.conf,
-                logical_date=parsed_logical_date,
-                replace_microseconds=False,
-                triggered_by=DagRunTriggeredByType.OPERATOR,
-            )
-
-        except DagRunAlreadyExists as e:
-            if self.reset_dag_run:
-                dag_run = e.dag_run
-                self.log.info("Clearing %s on %s", self.trigger_dag_id, dag_run.run_id)
-
-                # Get target dag object and call clear()
-                dag_model = DagModel.get_current(self.trigger_dag_id)
-                if dag_model is None:
-                    raise DagNotFound(f"Dag id {self.trigger_dag_id} not found in DagModel")
-
-                # Note: here execution fails on database isolation mode. Needs structural changes for AIP-72
-                dag_bag = DagBag(dag_folder=dag_model.fileloc, read_dags_from_db=True)
-                dag = dag_bag.get_dag(self.trigger_dag_id)
-                dag.clear(run_id=dag_run.run_id)
-            else:
-                if self.skip_when_already_exists:
-                    raise AirflowSkipException(
-                        "Skipping due to skip_when_already_exists is set to True and DagRunAlreadyExists"
-                    )
-                raise e
-        if dag_run is None:
-            raise RuntimeError("The dag_run should be set here!")
-        # Store the run id from the dag run (either created or found above) to
-        # be used when creating the extra link on the webserver.
-        ti = context["task_instance"]
-        ti.xcom_push(key=XCOM_RUN_ID, value=dag_run.run_id)
-
-        if self.wait_for_completion:
-            # Kick off the deferral process
-            if self._defer:
-                self.defer(
-                    trigger=DagStateTrigger(
-                        dag_id=self.trigger_dag_id,
-                        states=self.allowed_states + self.failed_states,
-                        run_ids=[run_id],
-                        poll_interval=self.poke_interval,
-                    ),
-                    method_name="execute_complete",
-                )
-            # wait for dag to complete
-            while True:
-                self.log.info(
-                    "Waiting for %s on %s to become allowed state %s ...",
-                    self.trigger_dag_id,
-                    run_id,
-                    self.allowed_states,
-                )
-                time.sleep(self.poke_interval)
-
-                # Note: here execution fails on database isolation mode. Needs structural changes for AIP-72
-                dag_run.refresh_from_db()
-                state = dag_run.state
-                if state in self.failed_states:
-                    raise AirflowException(f"{self.trigger_dag_id} failed with failed states {state}")
-                if state in self.allowed_states:
-                    self.log.info("%s finished with allowed state %s", self.trigger_dag_id, state)
-                    return
-
-    @provide_session
-    def execute_complete(self, context: Context, session: Session, event: tuple[str, dict[str, Any]]):
-        # This run_ids is parsed from the return trigger event
-        provided_run_id = event[1]["run_ids"][0]
-        try:
-            # Note: here execution fails on database isolation mode. Needs structural changes for AIP-72
-            dag_run = session.execute(
-                select(DagRun).where(DagRun.dag_id == self.trigger_dag_id, DagRun.run_id == provided_run_id)
-            ).scalar_one()
-        except NoResultFound:
-            raise AirflowException(
-                f"No DAG run found for DAG {self.trigger_dag_id} and run ID {provided_run_id}"
-            )
-
-        state = dag_run.state
-
-        if state in self.failed_states:
-            raise AirflowException(f"{self.trigger_dag_id} failed with failed state {state}")
-        if state in self.allowed_states:
-            self.log.info("%s finished with allowed state %s", self.trigger_dag_id, state)
-            return
-
-        raise AirflowException(
-            f"{self.trigger_dag_id} return {state} which is not in {self.failed_states}"
-            f" or {self.allowed_states}"
-        )

airflow/providers/standard/sensors/bash.py

@@ -1,120 +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 os
-from collections.abc import Sequence
-from subprocess import PIPE, STDOUT, Popen
-from tempfile import NamedTemporaryFile, TemporaryDirectory, gettempdir
-from typing import TYPE_CHECKING
-
-from airflow.exceptions import AirflowFailException
-from airflow.sensors.base import BaseSensorOperator
-
-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 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