apache-airflow-providers-standard 0.1.1rc1__tar.gz → 0.2.0b1__tar.gz

{apache_airflow_providers_standard-0.1.1rc1 → apache_airflow_providers_standard-0.2.0b1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-standard
-Version: 0.1.1rc1
+Version: 0.2.0b1
 Summary: Provider package apache-airflow-providers-standard for Apache Airflow
 Keywords: airflow-provider,standard,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -20,13 +20,13 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: System :: Monitoring
-Requires-Dist: apache-airflow>=2.9.0rc0
+Requires-Dist: apache-airflow>=2.9.0
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-standard/0.1.1/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-standard/0.1.1
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-standard/0.2.0b1/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-standard/0.2.0b1
+Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
-Project-URL: Twitter, https://x.com/ApacheAirflow
 Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
 
 
@@ -54,7 +54,7 @@ Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
 
 Package ``apache-airflow-providers-standard``
 
-Release: ``0.1.1``
+Release: ``0.2.0b1``
 
 
 Airflow Standard Provider
@@ -67,7 +67,7 @@ This is a provider package for ``standard`` provider. All classes for this provi
 are in ``airflow.providers.standard`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-standard/0.1.1/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-standard/0.2.0b1/>`_.
 
 Installation
 ------------
@@ -88,5 +88,5 @@ PIP package         Version required
 ==================  ==================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-standard/0.1.1/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-standard/0.2.0b1/changelog.html>`_.
 

{apache_airflow_providers_standard-0.1.1rc1 → apache_airflow_providers_standard-0.2.0b1}/README.rst

@@ -23,7 +23,7 @@
 
 Package ``apache-airflow-providers-standard``
 
-Release: ``0.1.1``
+Release: ``0.2.0b1``
 
 
 Airflow Standard Provider
@@ -36,7 +36,7 @@ This is a provider package for ``standard`` provider. All classes for this provi
 are in ``airflow.providers.standard`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-standard/0.1.1/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-standard/0.2.0b1/>`_.
 
 Installation
 ------------
@@ -57,4 +57,4 @@ PIP package         Version required
 ==================  ==================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-standard/0.1.1/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-standard/0.2.0b1/changelog.html>`_.

{apache_airflow_providers_standard-0.1.1rc1 → apache_airflow_providers_standard-0.2.0b1}/pyproject.toml

@@ -25,7 +25,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-standard"
-version = "0.1.1.rc1"
+version = "0.2.0b1"
 description = "Provider package apache-airflow-providers-standard for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -57,7 +57,7 @@ requires-python = "~=3.9"
 # Make sure to run ``breeze static-checks --type update-providers-dependencies --all-files``
 # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build``
 dependencies = [
-    "apache-airflow>=2.9.0rc0",
+    "apache-airflow>=2.9.0",
 ]
 
 [dependency-groups]
@@ -79,12 +79,12 @@ apache-airflow-providers-fab = {workspace = true}
 apache-airflow-providers-standard = {workspace = true}
 
 [project.urls]
-"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-standard/0.1.1"
-"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-standard/0.1.1/changelog.html"
+"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-standard/0.2.0b1"
+"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-standard/0.2.0b1/changelog.html"
 "Bug Tracker" = "https://github.com/apache/airflow/issues"
 "Source Code" = "https://github.com/apache/airflow"
 "Slack Chat" = "https://s.apache.org/airflow-slack"
-"Twitter" = "https://x.com/ApacheAirflow"
+"Mastodon" = "https://fosstodon.org/@airflow"
 "YouTube" = "https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/"
 
 [project.entry-points."apache_airflow_provider"]

{apache_airflow_providers_standard-0.1.1rc1 → apache_airflow_providers_standard-0.2.0b1}/src/airflow/providers/standard/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "0.1.1"
+__version__ = "0.2.0b1"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.9.0"

{apache_airflow_providers_standard-0.1.1rc1 → apache_airflow_providers_standard-0.2.0b1}/src/airflow/providers/standard/get_provider_info.py

@@ -27,8 +27,8 @@ def get_provider_info():
         "name": "Standard",
         "description": "Airflow Standard Provider\n",
         "state": "ready",
-        "source-date-epoch": 1741509906,
-        "versions": ["0.1.1", "0.1.0", "0.0.3", "0.0.2", "0.0.1"],
+        "source-date-epoch": 1742480519,
+        "versions": ["0.2.0b1", "0.1.1", "0.1.0", "0.0.3", "0.0.2", "0.0.1"],
         "integrations": [
             {
                 "integration-name": "Standard",
@@ -53,6 +53,7 @@ def get_provider_info():
                     "airflow.providers.standard.operators.trigger_dagrun",
                     "airflow.providers.standard.operators.latest_only",
                     "airflow.providers.standard.operators.smooth",
+                    "airflow.providers.standard.operators.branch",
                 ],
             }
         ],

apache_airflow_providers_standard-0.2.0b1/src/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))

{apache_airflow_providers_standard-0.1.1rc1 → apache_airflow_providers_standard-0.2.0b1}/src/airflow/providers/standard/operators/datetime.py

@@ -21,7 +21,7 @@ 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:

{apache_airflow_providers_standard-0.1.1rc1 → apache_airflow_providers_standard-0.2.0b1}/src/airflow/providers/standard/operators/latest_only.py

@@ -24,7 +24,7 @@ from typing import TYPE_CHECKING
 
 import pendulum
 
-from airflow.operators.branch import BaseBranchOperator
+from airflow.providers.standard.operators.branch import BaseBranchOperator
 from airflow.utils.types import DagRunType
 
 if TYPE_CHECKING:

{apache_airflow_providers_standard-0.1.1rc1 → apache_airflow_providers_standard-0.2.0b1}/src/airflow/providers/standard/operators/python.py

@@ -42,9 +42,7 @@ from airflow.exceptions import (
     DeserializingResultError,
 )
 from airflow.models.baseoperator import BaseOperator
-from airflow.models.skipmixin import SkipMixin
 from airflow.models.variable import Variable
-from airflow.operators.branch import BranchMixIn
 from airflow.providers.standard.utils.python_virtualenv import prepare_virtualenv, write_python_script
 from airflow.providers.standard.version_compat import AIRFLOW_V_2_10_PLUS, AIRFLOW_V_3_0_PLUS
 from airflow.utils import hashlib_wrapper
@@ -53,6 +51,14 @@ from airflow.utils.file import get_unique_dag_module_name
 from airflow.utils.operator_helpers import KeywordParameters
 from airflow.utils.process_utils import execute_in_subprocess, execute_in_subprocess_with_kwargs
 
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.providers.standard.operators.branch import BranchMixIn
+    from airflow.providers.standard.utils.skipmixin import SkipMixin
+else:
+    from airflow.models.skipmixin import SkipMixin
+    from airflow.operators.branch import BranchMixIn
+
+
 log = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
@@ -212,13 +218,11 @@ class PythonOperator(BaseOperator):
         """
         try:
             from airflow.utils.operator_helpers import ExecutionCallableRunner
-
-            asset_events = self._asset_events if AIRFLOW_V_3_0_PLUS else self._dataset_events
-
-            runner = ExecutionCallableRunner(self.python_callable, asset_events, logger=self.log)
         except ImportError:
-            # Handle Pre Airflow 3.10 case where ExecutionCallableRunner was not available
+            # Handle Pre Airflow 2.10 case where ExecutionCallableRunner was not available
             return self.python_callable(*self.op_args, **self.op_kwargs)
+        asset_events = self._asset_events if AIRFLOW_V_3_0_PLUS else self._dataset_events
+        runner = ExecutionCallableRunner(self.python_callable, asset_events, logger=self.log)
         return runner.run(*self.op_args, **self.op_kwargs)
 
 
@@ -236,6 +240,8 @@ class BranchPythonOperator(PythonOperator, BranchMixIn):
     the DAG run's state to be inferred.
     """
 
+    inherits_from_skipmixin = True
+
     def execute(self, context: Context) -> Any:
         return self.do_branch(context, super().execute(context))
 
@@ -266,6 +272,8 @@ class ShortCircuitOperator(PythonOperator, SkipMixin):
         skipped but the ``trigger_rule`` defined for all other downstream tasks will be respected.
     """
 
+    inherits_from_skipmixin = True
+
     def __init__(self, *, ignore_downstream_trigger_rules: bool = True, **kwargs) -> None:
         super().__init__(**kwargs)
         self.ignore_downstream_trigger_rules = ignore_downstream_trigger_rules
@@ -295,25 +303,24 @@ class ShortCircuitOperator(PythonOperator, SkipMixin):
 
         to_skip = get_tasks_to_skip()
 
-        # this let's us avoid an intermediate list unless debug logging
+        # this lets us avoid an intermediate list unless debug logging
         if self.log.getEffectiveLevel() <= logging.DEBUG:
             self.log.debug("Downstream task IDs %s", to_skip := list(get_tasks_to_skip()))
 
         self.log.info("Skipping downstream tasks")
         if AIRFLOW_V_3_0_PLUS:
             self.skip(
-                dag_id=dag_run.dag_id,
-                run_id=dag_run.run_id,
+                ti=context["ti"],
                 tasks=to_skip,
-                map_index=context["ti"].map_index,
             )
         else:
-            self.skip(
-                dag_run=dag_run,
-                tasks=to_skip,
-                execution_date=cast("DateTime", dag_run.logical_date),  # type: ignore[call-arg, union-attr]
-                map_index=context["ti"].map_index,
-            )
+            if to_skip:
+                self.skip(
+                    dag_run=context["dag_run"],
+                    tasks=to_skip,
+                    execution_date=cast("DateTime", dag_run.logical_date),  # type: ignore[call-arg, union-attr]
+                    map_index=context["ti"].map_index,
+                )
 
         self.log.info("Done.")
         # returns the result of the super execute method as it is instead of returning None
@@ -868,6 +875,8 @@ class BranchPythonVirtualenvOperator(PythonVirtualenvOperator, BranchMixIn):
         :ref:`howto/operator:BranchPythonVirtualenvOperator`
     """
 
+    inherits_from_skipmixin = True
+
     def execute(self, context: Context) -> Any:
         return self.do_branch(context, super().execute(context))
 

{apache_airflow_providers_standard-0.1.1rc1 → apache_airflow_providers_standard-0.2.0b1}/src/airflow/providers/standard/operators/trigger_dagrun.py

@@ -38,13 +38,12 @@ from airflow.models import BaseOperator
 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.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.utils import timezone
 from airflow.utils.session import provide_session
 from airflow.utils.state import DagRunState
-from airflow.utils.types import DagRunTriggeredByType, DagRunType
+from airflow.utils.types import DagRunType
 
 XCOM_LOGICAL_DATE_ISO = "trigger_logical_date_iso"
 XCOM_RUN_ID = "trigger_run_id"
@@ -63,7 +62,9 @@ if TYPE_CHECKING:
 
 if AIRFLOW_V_3_0_PLUS:
     from airflow.sdk import BaseOperatorLink
+    from airflow.sdk.execution_time.xcom import XCom
 else:
+    from airflow.models import XCom  # type: ignore[no-redef]
     from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
 
 
@@ -77,15 +78,15 @@ class TriggerDagRunLink(BaseOperatorLink):
     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
+        trigger_dag_id = operator.trigger_dag_id
+        if not AIRFLOW_V_3_0_PLUS:
+            from airflow.models.renderedtifields import RenderedTaskInstanceFields
+
+            if template_fields := RenderedTaskInstanceFields.get_templated_fields(ti_key):
+                trigger_dag_id: str = template_fields.get("trigger_dag_id", operator.trigger_dag_id)  # type: ignore[no-redef]
 
         # Fetch the correct dag_run_id for the triggerED dag which is
         # stored in xcom during execution of the triggerING task.
@@ -173,7 +174,7 @@ class TriggerDagRunOperator(BaseOperator):
             self.allowed_states = [DagRunState(s) for s in allowed_states]
         else:
             self.allowed_states = [DagRunState.SUCCESS]
-        if failed_states or failed_states == []:
+        if failed_states is not None:
             self.failed_states = [DagRunState(s) for s in failed_states]
         else:
             self.failed_states = [DagRunState.FAILED]
@@ -197,25 +198,51 @@ class TriggerDagRunOperator(BaseOperator):
         try:
             json.dumps(self.conf)
         except TypeError:
-            raise AirflowException("conf parameter should be JSON Serializable")
+            raise ValueError("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(),
-            )
+            if AIRFLOW_V_3_0_PLUS:
+                run_id = DagRun.generate_run_id(
+                    run_type=DagRunType.MANUAL,
+                    logical_date=parsed_logical_date,
+                    run_after=parsed_logical_date or timezone.utcnow(),
+                )
+            else:
+                run_id = DagRun.generate_run_id(DagRunType.MANUAL, parsed_logical_date or timezone.utcnow())  # type: ignore[misc,call-arg]
 
+        if AIRFLOW_V_3_0_PLUS:
+            self._trigger_dag_af_3(context=context, run_id=run_id, parsed_logical_date=parsed_logical_date)
+        else:
+            self._trigger_dag_af_2(context=context, run_id=run_id, parsed_logical_date=parsed_logical_date)
+
+    def _trigger_dag_af_3(self, context, run_id, parsed_logical_date):
+        from airflow.exceptions import DagRunTriggerException
+
+        raise DagRunTriggerException(
+            trigger_dag_id=self.trigger_dag_id,
+            dag_run_id=run_id,
+            conf=self.conf,
+            logical_date=parsed_logical_date,
+            reset_dag_run=self.reset_dag_run,
+            skip_when_already_exists=self.skip_when_already_exists,
+            wait_for_completion=self.wait_for_completion,
+            allowed_states=self.allowed_states,
+            failed_states=self.failed_states,
+            poke_interval=self.poke_interval,
+        )
+
+        # TODO: Support deferral
+
+    def _trigger_dag_af_2(self, context, run_id, parsed_logical_date):
         try:
             dag_run = trigger_dag(
                 dag_id=self.trigger_dag_id,
                 run_id=run_id,
                 conf=self.conf,
-                logical_date=parsed_logical_date,
+                execution_date=parsed_logical_date,
                 replace_microseconds=False,
-                triggered_by=DagRunTriggeredByType.OPERATOR,
             )
 
         except DagRunAlreadyExists as e:
@@ -231,7 +258,7 @@ class TriggerDagRunOperator(BaseOperator):
                 # 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)
+                dag.clear(start_date=dag_run.logical_date, end_date=dag_run.logical_date)
             else:
                 if self.skip_when_already_exists:
                     raise AirflowSkipException(
@@ -252,6 +279,7 @@ class TriggerDagRunOperator(BaseOperator):
                     trigger=DagStateTrigger(
                         dag_id=self.trigger_dag_id,
                         states=self.allowed_states + self.failed_states,
+                        execution_dates=[dag_run.logical_date],
                         run_ids=[run_id],
                         poll_interval=self.poke_interval,
                     ),
@@ -278,16 +306,18 @@ class TriggerDagRunOperator(BaseOperator):
 
     @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]
+        # This logical_date is parsed from the return trigger event
+        provided_logical_date = event[1]["execution_dates"][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)
+                select(DagRun).where(
+                    DagRun.dag_id == self.trigger_dag_id, DagRun.execution_date == provided_logical_date
+                )
             ).scalar_one()
         except NoResultFound:
             raise AirflowException(
-                f"No DAG run found for DAG {self.trigger_dag_id} and run ID {provided_run_id}"
+                f"No DAG run found for DAG {self.trigger_dag_id} and logical date {self.logical_date}"
             )
 
         state = dag_run.state

{apache_airflow_providers_standard-0.1.1rc1 → apache_airflow_providers_standard-0.2.0b1}/src/airflow/providers/standard/operators/weekday.py

@@ -20,9 +20,9 @@ from __future__ import annotations
 from collections.abc import Iterable
 from typing import TYPE_CHECKING
 
-from airflow.operators.branch import BaseBranchOperator
+from airflow.providers.standard.operators.branch import BaseBranchOperator
+from airflow.providers.standard.utils.weekday import WeekDay
 from airflow.utils import timezone
-from airflow.utils.weekday import WeekDay
 
 if TYPE_CHECKING:
     try:
@@ -63,7 +63,7 @@ class BranchDayOfWeekOperator(BaseBranchOperator):
     .. code-block:: python
 
         # import WeekDay Enum
-        from airflow.utils.weekday import WeekDay
+        from airflow.providers.standard.utils.weekday import WeekDay
         from airflow.providers.standard.operators.empty import EmptyOperator
         from airflow.operators.weekday import BranchDayOfWeekOperator
 

{apache_airflow_providers_standard-0.1.1rc1 → apache_airflow_providers_standard-0.2.0b1}/src/airflow/providers/standard/sensors/time_delta.py

@@ -46,29 +46,48 @@ def _get_airflow_version():
 
 class TimeDeltaSensor(BaseSensorOperator):
     """
-    Waits for a timedelta after the run's data interval.
+    Waits for a timedelta.
 
-    :param delta: time length to wait after the data interval before succeeding.
+    The delta will be evaluated against data_interval_end if present for the dag run,
+    otherwise run_after will be used.
+
+    :param delta: time to wait 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):
-        data_interval_end = context["data_interval_end"]
+    def _derive_base_time(self, context: Context) -> datetime:
+        """
+        Get the "base time" against which the delta should be calculated.
+
+        If data_interval_end is populated, use it; else use run_after.
+        """
+        data_interval_end = context.get("data_interval_end")
+        if data_interval_end:
+            if not isinstance(data_interval_end, datetime):
+                raise ValueError("`data_interval_end` returned non-datetime object")
 
-        if not isinstance(data_interval_end, datetime):
-            raise ValueError("`data_interval_end` returned non-datetime object")
+            return data_interval_end
 
-        target_dttm: datetime = data_interval_end + self.delta
-        self.log.info("Checking if the time (%s) has come", target_dttm)
+        if not data_interval_end and not AIRFLOW_V_3_0_PLUS:
+            raise ValueError("`data_interval_end` not found in task context.")
+
+        dag_run = context.get("dag_run")
+        if not dag_run:
+            raise ValueError("`dag_run` not found in task context")
+        return dag_run.run_after
+
+    def poke(self, context: Context) -> bool:
+        base_time = self._derive_base_time(context=context)
+        target_dttm = base_time + self.delta
+        self.log.info("Checking if the delta has elapsed base_time=%s, delta=%s", base_time, self.delta)
         return timezone.utcnow() > target_dttm
 
 
@@ -92,12 +111,8 @@ class TimeDeltaSensorAsync(TimeDeltaSensor):
         self.end_from_trigger = end_from_trigger
 
     def execute(self, context: Context) -> bool | NoReturn:
-        data_interval_end = context["data_interval_end"]
-
-        if not isinstance(data_interval_end, datetime):
-            raise ValueError("`data_interval_end` returned non-datetime object")
-
-        target_dttm: datetime = data_interval_end + self.delta
+        base_time = self._derive_base_time(context=context)
+        target_dttm: datetime = base_time + self.delta
 
         if timezone.utcnow() > target_dttm:
             # If the target datetime is in the past, return immediately

{apache_airflow_providers_standard-0.1.1rc1 → apache_airflow_providers_standard-0.2.0b1}/src/airflow/providers/standard/sensors/weekday.py

@@ -20,9 +20,9 @@ from __future__ import annotations
 from collections.abc import Iterable
 from typing import TYPE_CHECKING
 
+from airflow.providers.standard.utils.weekday import WeekDay
 from airflow.sensors.base import BaseSensorOperator
 from airflow.utils import timezone
-from airflow.utils.weekday import WeekDay
 
 if TYPE_CHECKING:
     try:
@@ -54,7 +54,7 @@ class DayOfWeekSensor(BaseSensorOperator):
     **Example** (with :class:`~airflow.utils.weekday.WeekDay` enum): ::
 
         # import WeekDay Enum
-        from airflow.utils.weekday import WeekDay
+        from airflow.providers.standard.utils.weekday import WeekDay
 
         weekend_check = DayOfWeekSensor(
             task_id="weekend_check",
@@ -103,7 +103,21 @@ class DayOfWeekSensor(BaseSensorOperator):
             self.week_day,
             WeekDay(timezone.utcnow().isoweekday()).name,
         )
+
         if self.use_task_logical_date:
-            return context["logical_date"].isoweekday() in self._week_day_num
+            logical_date = context.get("logical_date")
+            dag_run = context.get("dag_run")
+
+            if not (logical_date or (dag_run and dag_run.run_after)):
+                raise ValueError(
+                    "Either `logical_date` or `run_after` should be provided in the task context when "
+                    "`use_task_logical_date` is True"
+                )
+
+            determined_weekday_num = (
+                logical_date.isoweekday() if logical_date else dag_run.run_after.isoweekday()  # type: ignore[union-attr]
+            )
+
+            return determined_weekday_num in self._week_day_num
         else:
             return timezone.utcnow().isoweekday() in self._week_day_num

apache_airflow_providers_standard-0.2.0b1/src/airflow/providers/standard/utils/skipmixin.py

@@ -0,0 +1,192 @@
+#
+# 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 Iterable, Sequence
+from types import GeneratorType
+from typing import TYPE_CHECKING
+
+from airflow.exceptions import AirflowException
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+if TYPE_CHECKING:
+    from airflow.models.operator import Operator
+    from airflow.sdk.definitions._internal.node import DAGNode
+    from airflow.sdk.types import RuntimeTaskInstanceProtocol
+
+# The key used by SkipMixin to store XCom data.
+XCOM_SKIPMIXIN_KEY = "skipmixin_key"
+
+# The dictionary key used to denote task IDs that are skipped
+XCOM_SKIPMIXIN_SKIPPED = "skipped"
+
+# The dictionary key used to denote task IDs that are followed
+XCOM_SKIPMIXIN_FOLLOWED = "followed"
+
+
+def _ensure_tasks(nodes: Iterable[DAGNode]) -> Sequence[Operator]:
+    from airflow.models.baseoperator import BaseOperator
+    from airflow.models.mappedoperator import MappedOperator
+
+    return [n for n in nodes if isinstance(n, (BaseOperator, MappedOperator))]
+
+
+# This class should only be used in Airflow 3.0 and later.
+class SkipMixin(LoggingMixin):
+    """A Mixin to skip Tasks Instances."""
+
+    @staticmethod
+    def _set_state_to_skipped(
+        tasks: Sequence[str | tuple[str, int]],
+        map_index: int | None,
+    ) -> None:
+        """
+        Set state of task instances to skipped from the same dag run.
+
+        Raises
+        ------
+        SkipDownstreamTaskInstances
+            If the task instances are not in the same dag run.
+        """
+        # Import is internal for backward compatibility when importing PythonOperator
+        # from airflow.providers.common.compat.standard.operators
+        from airflow.exceptions import DownstreamTasksSkipped
+
+        #  The following could be applied only for non-mapped tasks,
+        #  as future mapped tasks have not been expanded yet. Such tasks
+        #  have to be handled by NotPreviouslySkippedDep.
+        if tasks and map_index == -1:
+            raise DownstreamTasksSkipped(tasks=tasks)
+
+    def skip(
+        self,
+        ti: RuntimeTaskInstanceProtocol,
+        tasks: Iterable[DAGNode],
+    ):
+        """
+        Set tasks instances to skipped from the same dag run.
+
+        If this instance has a `task_id` attribute, store the list of skipped task IDs to XCom
+        so that NotPreviouslySkippedDep knows these tasks should be skipped when they
+        are cleared.
+
+        :param ti: the task instance for which to set the tasks to skipped
+        :param tasks: tasks to skip (not task_ids)
+        """
+        # SkipMixin may not necessarily have a task_id attribute. Only store to XCom if one is available.
+        task_id: str | None = getattr(self, "task_id", None)
+        task_list = _ensure_tasks(tasks)
+        if not task_list:
+            return
+
+        task_ids_list = [d.task_id for d in task_list]
+
+        if task_id is not None:
+            ti.xcom_push(
+                key=XCOM_SKIPMIXIN_KEY,
+                value={XCOM_SKIPMIXIN_SKIPPED: task_ids_list},
+            )
+
+        self._set_state_to_skipped(task_ids_list, ti.map_index)
+
+    def skip_all_except(
+        self,
+        ti: RuntimeTaskInstanceProtocol,
+        branch_task_ids: None | str | Iterable[str],
+    ):
+        """
+        Implement the logic for a branching operator.
+
+        Given a single task ID or list of task IDs to follow, this skips all other tasks
+        immediately downstream of this operator.
+
+        branch_task_ids is stored to XCom so that NotPreviouslySkippedDep knows skipped tasks or
+        newly added tasks should be skipped when they are cleared.
+        """
+        # Ensure we don't serialize a generator object
+        if branch_task_ids and isinstance(branch_task_ids, GeneratorType):
+            branch_task_ids = list(branch_task_ids)
+        log = self.log  # Note: need to catch logger form instance, static logger breaks pytest
+        if isinstance(branch_task_ids, str):
+            branch_task_id_set = {branch_task_ids}
+        elif isinstance(branch_task_ids, Iterable):
+            branch_task_id_set = set(branch_task_ids)
+            invalid_task_ids_type = {
+                (bti, type(bti).__name__) for bti in branch_task_id_set if not isinstance(bti, str)
+            }
+            if invalid_task_ids_type:
+                raise AirflowException(
+                    f"'branch_task_ids' expected all task IDs are strings. "
+                    f"Invalid tasks found: {invalid_task_ids_type}."
+                )
+        elif branch_task_ids is None:
+            branch_task_id_set = set()
+        else:
+            raise AirflowException(
+                "'branch_task_ids' must be either None, a task ID, or an Iterable of IDs, "
+                f"but got {type(branch_task_ids).__name__!r}."
+            )
+
+        log.info("Following branch %s", branch_task_id_set)
+
+        if TYPE_CHECKING:
+            assert ti.task
+
+        task = ti.task
+        dag = ti.task.dag
+
+        valid_task_ids = set(dag.task_ids)
+        invalid_task_ids = branch_task_id_set - valid_task_ids
+        if invalid_task_ids:
+            raise AirflowException(
+                "'branch_task_ids' must contain only valid task_ids. "
+                f"Invalid tasks found: {invalid_task_ids}."
+            )
+
+        downstream_tasks = _ensure_tasks(task.downstream_list)
+
+        if downstream_tasks:
+            # For a branching workflow that looks like this, when "branch" does skip_all_except("task1"),
+            # we intuitively expect both "task1" and "join" to execute even though strictly speaking,
+            # "join" is also immediately downstream of "branch" and should have been skipped. Therefore,
+            # we need a special case here for such empty branches: Check downstream tasks of branch_task_ids.
+            # In case the task to skip is also downstream of branch_task_ids, we add it to branch_task_ids and
+            # exclude it from skipping.
+            #
+            # branch  ----->  join
+            #   \            ^
+            #     v        /
+            #       task1
+            #
+            for branch_task_id in list(branch_task_id_set):
+                branch_task_id_set.update(dag.get_task(branch_task_id).get_flat_relative_ids(upstream=False))
+
+            skip_tasks = [
+                (t.task_id, ti.map_index) for t in downstream_tasks if t.task_id not in branch_task_id_set
+            ]
+
+            follow_task_ids = [t.task_id for t in downstream_tasks if t.task_id in branch_task_id_set]
+            log.info("Skipping tasks %s", skip_tasks)
+            ti.xcom_push(
+                key=XCOM_SKIPMIXIN_KEY,
+                value={XCOM_SKIPMIXIN_FOLLOWED: follow_task_ids},
+            )
+            #  The following could be applied only for non-mapped tasks,
+            #  as future mapped tasks have not been expanded yet. Such tasks
+            #  have to be handled by NotPreviouslySkippedDep.
+            self._set_state_to_skipped(skip_tasks, ti.map_index)  # type: ignore[arg-type]

apache_airflow_providers_standard-0.2.0b1/src/airflow/providers/standard/utils/weekday.py

@@ -0,0 +1,77 @@
+# 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.
+"""Get the ISO standard day number of the week from a given day string."""
+
+from __future__ import annotations
+
+import enum
+from collections.abc import Iterable
+
+
+@enum.unique
+class WeekDay(enum.IntEnum):
+    """Python Enum containing Days of the Week."""
+
+    MONDAY = 1
+    TUESDAY = 2
+    WEDNESDAY = 3
+    THURSDAY = 4
+    FRIDAY = 5
+    SATURDAY = 6
+    SUNDAY = 7
+
+    @classmethod
+    def get_weekday_number(cls, week_day_str: str):
+        """
+        Return the ISO Week Day Number for a Week Day.
+
+        :param week_day_str: Full Name of the Week Day. Example: "Sunday"
+        :return: ISO Week Day Number corresponding to the provided Weekday
+        """
+        sanitized_week_day_str = week_day_str.upper()
+
+        if sanitized_week_day_str not in cls.__members__:
+            raise AttributeError(f'Invalid Week Day passed: "{week_day_str}"')
+
+        return cls[sanitized_week_day_str]
+
+    @classmethod
+    def convert(cls, day: str | WeekDay) -> int:
+        """Return the day number in the week."""
+        if isinstance(day, WeekDay):
+            return day
+        return cls.get_weekday_number(week_day_str=day)
+
+    @classmethod
+    def validate_week_day(
+        cls,
+        week_day: str | WeekDay | Iterable[str] | Iterable[WeekDay],
+    ) -> set[int]:
+        """Validate each item of iterable and create a set to ease compare of values."""
+        if not isinstance(week_day, Iterable):
+            if isinstance(week_day, WeekDay):
+                week_day = {week_day}
+            else:
+                raise TypeError(
+                    f"Unsupported Type for week_day parameter: {type(week_day)}."
+                    "Input should be iterable type:"
+                    "str, set, list, dict or Weekday enum type"
+                )
+        if isinstance(week_day, str):
+            week_day = {week_day}
+
+        return {cls.convert(item) for item in week_day}