apache-airflow-providers-standard 1.0.0.dev1__py3-none-any.whl → 1.1.0__py3-none-any.whl

airflow/providers/standard/operators/bash.py

@@ -19,19 +19,29 @@ from __future__ import annotations
 
 import os
 import shutil
-import warnings
+import tempfile
+from collections.abc import Container, Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Callable, Container, Sequence, cast
+from typing import TYPE_CHECKING, Any, Callable, cast
 
 from airflow.exceptions import AirflowException, AirflowSkipException
-from airflow.hooks.subprocess import SubprocessHook
 from airflow.models.baseoperator import BaseOperator
-from airflow.utils.operator_helpers import context_to_airflow_vars
-from airflow.utils.types import ArgNotSet
+from airflow.providers.standard.hooks.subprocess import SubprocessHook, SubprocessResult, working_directory
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk.execution_time.context import context_to_airflow_vars
+else:
+    from airflow.utils.operator_helpers import context_to_airflow_vars  # type: ignore[no-redef, attr-defined]
 
 if TYPE_CHECKING:
-    from airflow.models.taskinstance import TaskInstance
-    from airflow.utils.context import Context
+    from airflow.utils.types import ArgNotSet
+
+    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 BashOperator(BaseOperator):
@@ -62,7 +72,10 @@ class BashOperator(BaseOperator):
     :param cwd: Working directory to execute the command in (templated).
         If None (default), the command is run in a temporary directory.
         To use current DAG folder as the working directory,
-        you might set template ``{{ dag_run.dag.folder }}``.
+        you might set template ``{{ task.dag.folder }}``.
+        When bash_command is a '.sh' or '.bash' file, Airflow must have write
+        access to the working directory. The script will be rendered (Jinja
+        template) into a new temporary file in this directory.
     :param output_processor: Function to further process the output of the bash script
         (default is lambda output: output).
 
@@ -93,14 +106,18 @@ class BashOperator(BaseOperator):
 
         .. code-block:: python
 
-            bash_command = "set -e; python3 script.py '{{ next_execution_date }}'"
+            bash_command = "set -e; python3 script.py '{{ data_interval_end }}'"
 
     .. note::
 
-        Add a space after the script name when directly calling a ``.sh`` script with the
-        ``bash_command`` argument -- for example ``bash_command="my_script.sh "``.  This
-        is because Airflow tries to apply load this file and process it as a Jinja template to
-        it ends with ``.sh``, which will likely not be what most users want.
+        To simply execute a ``.sh`` or ``.bash`` script (without any Jinja template), add a space after the
+        script name ``bash_command`` argument -- for example ``bash_command="my_script.sh "``. This
+        is because Airflow tries to load this file and process it as a Jinja template when
+        it ends with ``.sh`` or ``.bash``.
+
+        If you have Jinja template in your script, do not put any blank space. And add the script's directory
+        in the DAG's ``template_searchpath``. If you specify a ``cwd``, Airflow must have write access to
+        this directory. The script will be rendered (Jinja template) into a new temporary file in this directory.
 
     .. warning::
 
@@ -149,7 +166,6 @@ class BashOperator(BaseOperator):
         env: dict[str, str] | None = None,
         append_env: bool = False,
         output_encoding: str = "utf-8",
-        skip_exit_code: int | None = None,
         skip_on_exit_code: int | Container[int] | None = 99,
         cwd: str | None = None,
         output_processor: Callable[[str], Any] = lambda result: result,
@@ -159,11 +175,6 @@ class BashOperator(BaseOperator):
         self.bash_command = bash_command
         self.env = env
         self.output_encoding = output_encoding
-        if skip_exit_code is not None:
-            warnings.warn(
-                "skip_exit_code is deprecated. Please use skip_on_exit_code", DeprecationWarning, stacklevel=2
-            )
-            skip_on_exit_code = skip_exit_code
         self.skip_on_exit_code = (
             skip_on_exit_code
             if isinstance(skip_on_exit_code, Container)
@@ -174,33 +185,16 @@ class BashOperator(BaseOperator):
         self.cwd = cwd
         self.append_env = append_env
         self.output_processor = output_processor
-
-        # When using the @task.bash decorator, the Bash command is not known until the underlying Python
-        # callable is executed and therefore set to NOTSET initially. This flag is useful during execution to
-        # determine whether the bash_command value needs to re-rendered.
-        self._init_bash_command_not_set = isinstance(self.bash_command, ArgNotSet)
+        self._is_inline_cmd = None
+        if isinstance(bash_command, str):
+            self._is_inline_cmd = self._is_inline_command(bash_command=bash_command)
 
     @cached_property
     def subprocess_hook(self):
         """Returns hook for running the bash command."""
         return SubprocessHook()
 
-    @staticmethod
-    def refresh_bash_command(ti: TaskInstance) -> None:
-        """
-        Rewrite the underlying rendered bash_command value for a task instance in the metadatabase.
-
-        TaskInstance.get_rendered_template_fields() cannot be used because this will retrieve the
-        RenderedTaskInstanceFields from the metadatabase which doesn't have the runtime-evaluated bash_command
-        value.
-
-        :meta private:
-        """
-        from airflow.models.renderedtifields import RenderedTaskInstanceFields
-
-        RenderedTaskInstanceFields._update_runtime_evaluated_template_fields(ti)
-
-    def get_env(self, context):
+    def get_env(self, context) -> dict:
         """Build the set of environment variables to be exposed for the bash command."""
         system_env = os.environ.copy()
         env = self.env
@@ -220,7 +214,7 @@ class BashOperator(BaseOperator):
         return env
 
     def execute(self, context: Context):
-        bash_path = shutil.which("bash") or "bash"
+        bash_path: str = shutil.which("bash") or "bash"
         if self.cwd is not None:
             if not os.path.exists(self.cwd):
                 raise AirflowException(f"Can not find the cwd: {self.cwd}")
@@ -228,29 +222,52 @@ class BashOperator(BaseOperator):
                 raise AirflowException(f"The cwd {self.cwd} must be a directory")
         env = self.get_env(context)
 
-        # Because the bash_command value is evaluated at runtime using the @task.bash decorator, the
-        # RenderedTaskInstanceField data needs to be rewritten and the bash_command value re-rendered -- the
-        # latter because the returned command from the decorated callable could contain a Jinja expression.
-        # Both will ensure the correct Bash command is executed and that the Rendered Template view in the UI
-        # displays the executed command (otherwise it will display as an ArgNotSet type).
-        if self._init_bash_command_not_set:
-            ti = cast("TaskInstance", context["ti"])
-            self.refresh_bash_command(ti)
+        if self._is_inline_cmd:
+            result = self._run_inline_command(bash_path=bash_path, env=env)
+        else:
+            result = self._run_rendered_script_file(bash_path=bash_path, env=env)
 
-        result = self.subprocess_hook.run_command(
-            command=[bash_path, "-c", self.bash_command],
-            env=env,
-            output_encoding=self.output_encoding,
-            cwd=self.cwd,
-        )
         if result.exit_code in self.skip_on_exit_code:
             raise AirflowSkipException(f"Bash command returned exit code {result.exit_code}. Skipping.")
-        elif result.exit_code != 0:
+        if result.exit_code != 0:
             raise AirflowException(
                 f"Bash command failed. The command returned a non-zero exit code {result.exit_code}."
             )
 
         return self.output_processor(result.output)
 
+    def _run_inline_command(self, bash_path: str, env: dict) -> SubprocessResult:
+        """Pass the bash command as string directly in the subprocess."""
+        return self.subprocess_hook.run_command(
+            command=[bash_path, "-c", self.bash_command],
+            env=env,
+            output_encoding=self.output_encoding,
+            cwd=self.cwd,
+        )
+
+    def _run_rendered_script_file(self, bash_path: str, env: dict) -> SubprocessResult:
+        """
+        Save the bash command into a file and execute this file.
+
+        This allows for longer commands, and prevents "Argument list too long error".
+        """
+        with working_directory(cwd=self.cwd) as cwd:
+            with tempfile.NamedTemporaryFile(mode="w", dir=cwd, suffix=".sh") as file:
+                file.write(cast("str", self.bash_command))
+                file.flush()
+
+                bash_script = os.path.basename(file.name)
+                return self.subprocess_hook.run_command(
+                    command=[bash_path, bash_script],
+                    env=env,
+                    output_encoding=self.output_encoding,
+                    cwd=cwd,
+                )
+
+    @classmethod
+    def _is_inline_command(cls, bash_command: str) -> bool:
+        """Return True if the bash command is an inline string. False if it's a bash script file."""
+        return not bash_command.endswith(tuple(cls.template_ext))
+
     def on_kill(self) -> None:
         self.subprocess_hook.send_sigterm()

airflow/providers/standard/operators/branch.py

@@ -0,0 +1,105 @@
+#
+# 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.
+"""Branching operators."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
+
+from airflow.models.baseoperator import BaseOperator
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.providers.standard.utils.skipmixin import SkipMixin
+else:
+    from airflow.models.skipmixin import SkipMixin
+
+if TYPE_CHECKING:
+    from airflow.sdk.definitions.context import Context
+    from airflow.sdk.types import RuntimeTaskInstanceProtocol
+
+
+class BranchMixIn(SkipMixin):
+    """Utility helper which handles the branching as one-liner."""
+
+    def do_branch(self, context: Context, branches_to_execute: str | Iterable[str]) -> str | Iterable[str]:
+        """Implement the handling of branching including logging."""
+        self.log.info("Branch into %s", branches_to_execute)
+        branch_task_ids = self._expand_task_group_roots(context["ti"], branches_to_execute)
+        self.skip_all_except(context["ti"], branch_task_ids)
+        return branches_to_execute
+
+    def _expand_task_group_roots(
+        self, ti: RuntimeTaskInstanceProtocol, branches_to_execute: str | Iterable[str]
+    ) -> Iterable[str]:
+        """Expand any task group into its root task ids."""
+        if TYPE_CHECKING:
+            assert ti.task
+
+        task = ti.task
+        dag = task.dag
+        if TYPE_CHECKING:
+            assert dag
+
+        if branches_to_execute is None:
+            return
+        elif isinstance(branches_to_execute, str) or not isinstance(branches_to_execute, Iterable):
+            branches_to_execute = [branches_to_execute]
+
+        for branch in branches_to_execute:
+            if branch in dag.task_group_dict:
+                tg = dag.task_group_dict[branch]
+                root_ids = [root.task_id for root in tg.roots]
+                self.log.info("Expanding task group %s into %s", tg.group_id, root_ids)
+                yield from root_ids
+            else:
+                yield branch
+
+
+class BaseBranchOperator(BaseOperator, BranchMixIn):
+    """
+    A base class for creating operators with branching functionality, like to BranchPythonOperator.
+
+    Users should create a subclass from this operator and implement the function
+    `choose_branch(self, context)`. This should run whatever business logic
+    is needed to determine the branch, and return one of the following:
+    - A single task_id (as a str)
+    - A single task_group_id (as a str)
+    - A list containing a combination of task_ids and task_group_ids
+
+    The operator will continue with the returned task_id(s) and/or task_group_id(s), and all other
+    tasks directly downstream of this operator will be skipped.
+    """
+
+    inherits_from_skipmixin = True
+
+    def choose_branch(self, context: Context) -> str | Iterable[str]:
+        """
+        Abstract method to choose which branch to run.
+
+        Subclasses should implement this, running whatever logic is
+        necessary to choose a branch and returning a task_id or list of
+        task_ids.
+
+        :param context: Context dictionary as passed to execute()
+        """
+        raise NotImplementedError
+
+    def execute(self, context: Context):
+        return self.do_branch(context, self.choose_branch(context))

airflow/providers/standard/operators/datetime.py

@@ -17,14 +17,19 @@
 from __future__ import annotations
 
 import datetime
-from typing import TYPE_CHECKING, Iterable
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
-from airflow.operators.branch import BaseBranchOperator
+from airflow.providers.standard.operators.branch import BaseBranchOperator
 from airflow.utils import timezone
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    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 BranchDateTimeOperator(BaseBranchOperator):
@@ -44,7 +49,7 @@ class BranchDateTimeOperator(BaseBranchOperator):
     :param target_lower: target lower bound.
     :param target_upper: target upper bound.
     :param use_task_logical_date: If ``True``, uses task's logical date to compare with targets.
-        Execution date is useful for backfilling. If ``False``, uses system's date.
+        Logical date is useful for backfilling. If ``False``, uses system's date.
     """
 
     def __init__(
@@ -72,9 +77,14 @@ class BranchDateTimeOperator(BaseBranchOperator):
 
     def choose_branch(self, context: Context) -> str | Iterable[str]:
         if self.use_task_logical_date:
-            now = context["logical_date"]
+            now = context.get("logical_date")
+            if not now:
+                dag_run = context.get("dag_run")
+                now = dag_run.run_after  # type: ignore[union-attr, assignment]
         else:
             now = timezone.coerce_datetime(timezone.utcnow())
+        if TYPE_CHECKING:
+            assert isinstance(now, datetime.datetime)
         lower, upper = target_times_as_dates(now, self.target_lower, self.target_upper)
         lower = timezone.coerce_datetime(lower, self.dag.timezone)
         upper = timezone.coerce_datetime(upper, self.dag.timezone)

airflow/providers/standard/operators/empty.py

@@ -0,0 +1,39 @@
+# 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
+
+from airflow.models.baseoperator import BaseOperator
+
+if TYPE_CHECKING:
+    from airflow.sdk.definitions.context import Context
+
+
+class EmptyOperator(BaseOperator):
+    """
+    Operator that does literally nothing.
+
+    It can be used to group tasks in a DAG.
+    The task is evaluated by the scheduler but never processed by the executor.
+    """
+
+    ui_color = "#e8f7e4"
+    inherits_from_empty_operator = True
+
+    def execute(self, context: Context):
+        pass

airflow/providers/standard/operators/latest_only.py

@@ -0,0 +1,127 @@
+#
+# 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.
+"""Contains an operator to run downstream tasks only for the latest scheduled DagRun."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+import pendulum
+
+from airflow.providers.standard.operators.branch import BaseBranchOperator
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
+from airflow.utils.types import DagRunType
+
+if TYPE_CHECKING:
+    from pendulum.datetime import DateTime
+
+    from airflow.models import DagRun
+
+    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 LatestOnlyOperator(BaseBranchOperator):
+    """
+    Skip tasks that are not running during the most recent schedule interval.
+
+    If the task is run outside the latest schedule interval (i.e. run_type == DagRunType.MANUAL),
+    all directly downstream tasks will be skipped.
+
+    Note that downstream tasks are never skipped if the given DAG_Run is
+    marked as externally triggered.
+
+    Note that when used with timetables that produce zero-length or point-in-time data intervals
+    (e.g., ``DeltaTriggerTimetable``), this operator assumes each run is the latest
+    and does not skip downstream tasks.
+    """
+
+    ui_color = "#e9ffdb"  # nyanza
+
+    def choose_branch(self, context: Context) -> str | Iterable[str]:
+        # If the DAG Run is externally triggered, then return without
+        # skipping downstream tasks
+        dag_run: DagRun = context["dag_run"]  # type: ignore[assignment]
+        if dag_run.run_type == DagRunType.MANUAL:
+            self.log.info("Manually triggered DAG_Run: allowing execution to proceed.")
+            return list(self.get_direct_relative_ids(upstream=False))
+
+        dates = self._get_compare_dates(dag_run)
+
+        if dates is None:
+            self.log.info("Last scheduled execution: allowing execution to proceed.")
+            return list(self.get_direct_relative_ids(upstream=False))
+
+        now = pendulum.now("UTC")
+        left_window, right_window = dates
+        self.log.info(
+            "Checking latest only with left_window: %s right_window: %s now: %s",
+            left_window,
+            right_window,
+            now,
+        )
+
+        if not left_window < now <= right_window:
+            self.log.info("Not latest execution, skipping downstream.")
+            # we return an empty list, thus the parent BaseBranchOperator
+            # won't exclude any downstream tasks from skipping.
+            return []
+
+        self.log.info("Latest, allowing execution to proceed.")
+        return list(self.get_direct_relative_ids(upstream=False))
+
+    def _get_compare_dates(self, dag_run: DagRun) -> tuple[DateTime, DateTime] | None:
+        dagrun_date: DateTime
+        if AIRFLOW_V_3_0_PLUS:
+            dagrun_date = dag_run.logical_date or dag_run.run_after
+        else:
+            dagrun_date = dag_run.logical_date
+
+        from airflow.timetables.base import DataInterval, TimeRestriction
+
+        current_interval = DataInterval(
+            start=dag_run.data_interval_start or dagrun_date,
+            end=dag_run.data_interval_end or dagrun_date,
+        )
+
+        time_restriction = TimeRestriction(
+            earliest=None, latest=current_interval.end - timedelta(microseconds=1), catchup=True
+        )
+        if prev_info := self.dag.timetable.next_dagrun_info(
+            last_automated_data_interval=current_interval,
+            restriction=time_restriction,
+        ):
+            left = prev_info.data_interval.end
+        else:
+            left = current_interval.start
+
+        time_restriction = TimeRestriction(earliest=current_interval.end, latest=None, catchup=True)
+        next_info = self.dag.timetable.next_dagrun_info(
+            last_automated_data_interval=current_interval,
+            restriction=time_restriction,
+        )
+
+        if not next_info:
+            return None
+
+        return (left, next_info.data_interval.end)