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

airflow/providers/standard/LICENSE

@@ -199,55 +199,3 @@ 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.
-
-============================================================================
-   APACHE AIRFLOW SUBCOMPONENTS:
-
-   The Apache Airflow project contains subcomponents with separate copyright
-   notices and license terms. Your use of the source code for the these
-   subcomponents is subject to the terms and conditions of the following
-   licenses.
-
-
-========================================================================
-Third party Apache 2.0 licenses
-========================================================================
-
-The following components are provided under the Apache 2.0 License.
-See project link for details. The text of each license is also included
-at 3rd-party-licenses/LICENSE-[project].txt.
-
-    (ALv2 License) hue v4.3.0 (https://github.com/cloudera/hue/)
-    (ALv2 License) jqclock v2.3.0 (https://github.com/JohnRDOrazio/jQuery-Clock-Plugin)
-    (ALv2 License) bootstrap3-typeahead v4.0.2 (https://github.com/bassjobsen/Bootstrap-3-Typeahead)
-    (ALv2 License) connexion v2.7.0 (https://github.com/zalando/connexion)
-
-========================================================================
-MIT licenses
-========================================================================
-
-The following components are provided under the MIT License. See project link for details.
-The text of each license is also included at 3rd-party-licenses/LICENSE-[project].txt.
-
-    (MIT License) jquery v3.5.1 (https://jquery.org/license/)
-    (MIT License) dagre-d3 v0.6.4 (https://github.com/cpettitt/dagre-d3)
-    (MIT License) bootstrap v3.4.1 (https://github.com/twbs/bootstrap/)
-    (MIT License) d3-tip v0.9.1 (https://github.com/Caged/d3-tip)
-    (MIT License) dataTables v1.10.25 (https://datatables.net)
-    (MIT License) normalize.css v3.0.2 (http://necolas.github.io/normalize.css/)
-    (MIT License) ElasticMock v1.3.2 (https://github.com/vrcmarcos/elasticmock)
-    (MIT License) MomentJS v2.24.0 (http://momentjs.com/)
-    (MIT License) eonasdan-bootstrap-datetimepicker v4.17.49 (https://github.com/eonasdan/bootstrap-datetimepicker/)
-
-========================================================================
-BSD 3-Clause licenses
-========================================================================
-The following components are provided under the BSD 3-Clause license. See project links for details.
-The text of each license is also included at 3rd-party-licenses/LICENSE-[project].txt.
-
-    (BSD 3 License) d3 v5.16.0 (https://d3js.org)
-    (BSD 3 License) d3-shape v2.1.0 (https://github.com/d3/d3-shape)
-    (BSD 3 License) cgroupspy 0.2.1 (https://github.com/cloudsigma/cgroupspy)
-
-========================================================================
-See 3rd-party-licenses/LICENSES-ui.txt for packages used in `/airflow/www`

airflow/providers/standard/__init__.py

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

airflow/providers/standard/get_provider_info.py

@@ -15,8 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE
-# OVERWRITTEN WHEN PREPARING PACKAGES.
+# NOTE! THIS FILE IS AUTOMATICALLY GENERATED AND WILL BE OVERWRITTEN!
 #
 # IF YOU WANT TO MODIFY THIS FILE, YOU SHOULD MODIFY THE TEMPLATE
 # `get_provider_info_TEMPLATE.py.jinja2` IN the `dev/breeze/src/airflow_breeze/templates` DIRECTORY
@@ -28,9 +27,8 @@ def get_provider_info():
         "name": "Standard",
         "description": "Airflow Standard Provider\n",
         "state": "ready",
-        "source-date-epoch": 1734536895,
-        "versions": ["0.0.3", "0.0.2", "0.0.1"],
-        "dependencies": ["apache-airflow>=2.9.0", "apache-airflow-providers-common-sql>=1.20.0"],
+        "source-date-epoch": 1739964539,
+        "versions": ["0.1.0", "0.0.3", "0.0.2", "0.0.1"],
         "integrations": [
             {
                 "integration-name": "Standard",
@@ -51,6 +49,7 @@ def get_provider_info():
                     "airflow.providers.standard.operators.weekday",
                     "airflow.providers.standard.operators.bash",
                     "airflow.providers.standard.operators.python",
+                    "airflow.providers.standard.operators.empty",
                     "airflow.providers.standard.operators.generic_transfer",
                     "airflow.providers.standard.operators.trigger_dagrun",
                     "airflow.providers.standard.operators.latest_only",
@@ -106,4 +105,5 @@ def get_provider_info():
                 },
             }
         },
+        "dependencies": ["apache-airflow>=2.9.0", "apache-airflow-providers-common-sql>=1.20.0"],
     }

airflow/providers/standard/operators/bash.py

@@ -34,8 +34,11 @@ from airflow.utils.types import ArgNotSet
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session as SASession
 
-    from airflow.models.taskinstance import TaskInstance
-    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 BashOperator(BaseOperator):
@@ -66,7 +69,7 @@ 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.
@@ -198,7 +201,7 @@ class BashOperator(BaseOperator):
     # TODO: This should be replaced with Task SDK API call
     @staticmethod
     @provide_session
-    def refresh_bash_command(ti: TaskInstance, session: SASession = NEW_SESSION) -> None:
+    def refresh_bash_command(ti, session: SASession = NEW_SESSION) -> None:
         """
         Rewrite the underlying rendered bash_command value for a task instance in the metadatabase.
 
@@ -211,11 +214,6 @@ class BashOperator(BaseOperator):
         from airflow.models.renderedtifields import RenderedTaskInstanceFields
 
         """Update rendered task instance fields for cases where runtime evaluated, not templated."""
-        # Note: Need lazy import to break the partly loaded class loop
-        from airflow.models.taskinstance import TaskInstance
-
-        # If called via remote API the DAG needs to be re-loaded
-        TaskInstance.ensure_dag(ti, session=session)
 
         rtif = RenderedTaskInstanceFields(ti)
         RenderedTaskInstanceFields.write(rtif, session=session)

airflow/providers/standard/operators/datetime.py

@@ -25,7 +25,11 @@ from airflow.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):

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/generic_transfer.py

@@ -24,7 +24,11 @@ from airflow.hooks.base import BaseHook
 from airflow.models import BaseOperator
 
 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 GenericTransfer(BaseOperator):

airflow/providers/standard/operators/latest_only.py

@@ -28,7 +28,12 @@ from airflow.operators.branch import BaseBranchOperator
 
 if TYPE_CHECKING:
     from airflow.models import DAG, DagRun
-    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 LatestOnlyOperator(BaseBranchOperator):
@@ -52,7 +57,7 @@ class LatestOnlyOperator(BaseBranchOperator):
             self.log.info("Externally triggered DAG_Run: allowing execution to proceed.")
             return list(context["task"].get_direct_relative_ids(upstream=False))
 
-        dag: DAG = context["dag"]
+        dag: DAG = context["dag"]  # type: ignore[assignment]
         next_info = dag.next_dagrun_info(dag.get_run_data_interval(dag_run), restricted=False)
         now = pendulum.now("UTC")
 

airflow/providers/standard/operators/python.py

@@ -43,15 +43,10 @@ from airflow.exceptions import (
 )
 from airflow.models.baseoperator import BaseOperator
 from airflow.models.skipmixin import SkipMixin
-from airflow.models.taskinstance import _CURRENT_CONTEXT
 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.typing_compat import Literal
+from airflow.providers.standard.version_compat import AIRFLOW_V_2_10_PLUS, AIRFLOW_V_3_0_PLUS
 from airflow.utils import hashlib_wrapper
 from airflow.utils.context import context_copy_partial, context_merge
 from airflow.utils.file import get_unique_dag_module_name
@@ -61,9 +56,17 @@ from airflow.utils.process_utils import execute_in_subprocess, execute_in_subpro
 log = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
+    from typing import Literal
+
     from pendulum.datetime import DateTime
 
-    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
+
+    _SerializerTypeDef = Literal["pickle", "cloudpickle", "dill"]
 
 
 @cache
@@ -299,7 +302,8 @@ class ShortCircuitOperator(PythonOperator, SkipMixin):
         self.log.info("Skipping downstream tasks")
         if AIRFLOW_V_3_0_PLUS:
             self.skip(
-                dag_run=dag_run,
+                dag_id=dag_run.dag_id,
+                run_id=dag_run.run_id,
                 tasks=to_skip,
                 map_index=context["ti"].map_index,
             )
@@ -343,7 +347,6 @@ def _load_cloudpickle():
     return cloudpickle
 
 
-_SerializerTypeDef = Literal["pickle", "cloudpickle", "dill"]
 _SERIALIZERS: dict[_SerializerTypeDef, Any] = {
     "pickle": lazy_object_proxy.Proxy(_load_pickle),
     "dill": lazy_object_proxy.Proxy(_load_dill),
@@ -374,6 +377,9 @@ class _BasePythonVirtualenvOperator(PythonOperator, metaclass=ABCMeta):
         "yesterday_ds",
         "yesterday_ds_nodash",
     }
+    if AIRFLOW_V_3_0_PLUS:
+        BASE_SERIALIZABLE_CONTEXT_KEYS.add("task_reschedule_count")
+
     PENDULUM_SERIALIZABLE_CONTEXT_KEYS = {
         "data_interval_end",
         "data_interval_start",
@@ -388,6 +394,8 @@ class _BasePythonVirtualenvOperator(PythonOperator, metaclass=ABCMeta):
         "prev_execution_date",
         "prev_execution_date_success",
     }
+    if AIRFLOW_V_3_0_PLUS:
+        PENDULUM_SERIALIZABLE_CONTEXT_KEYS.add("start_date")
 
     AIRFLOW_SERIALIZABLE_CONTEXT_KEYS = {
         "macros",
@@ -415,7 +423,6 @@ class _BasePythonVirtualenvOperator(PythonOperator, metaclass=ABCMeta):
         skip_on_exit_code: int | Container[int] | None = None,
         env_vars: dict[str, str] | None = None,
         inherit_env: bool = True,
-        use_airflow_context: bool = False,
         **kwargs,
     ):
         if (
@@ -457,7 +464,6 @@ class _BasePythonVirtualenvOperator(PythonOperator, metaclass=ABCMeta):
         )
         self.env_vars = env_vars
         self.inherit_env = inherit_env
-        self.use_airflow_context = use_airflow_context
 
     @abstractmethod
     def _iter_serializable_context_keys(self):
@@ -516,7 +522,6 @@ class _BasePythonVirtualenvOperator(PythonOperator, metaclass=ABCMeta):
                 "pickling_library": self.serializer,
                 "python_callable": self.python_callable.__name__,
                 "python_callable_source": self.get_python_source(),
-                "use_airflow_context": self.use_airflow_context,
             }
 
             if inspect.getfile(self.python_callable) == self.dag.fileloc:
@@ -527,20 +532,6 @@ class _BasePythonVirtualenvOperator(PythonOperator, metaclass=ABCMeta):
                 filename=os.fspath(script_path),
                 render_template_as_native_obj=self.dag.render_template_as_native_obj,
             )
-            if self.use_airflow_context:
-                # TODO: replace with commented code when context serialization is implemented in AIP-72
-                raise AirflowException(
-                    "The `use_airflow_context=True` is not yet implemented. "
-                    "It will work in Airflow 3 after AIP-72 context "
-                    "serialization is ready."
-                )
-                # context = get_current_context()
-                # with create_session() as session:
-                #     dag_run, task_instance = context["dag_run"], context["task_instance"]
-                #     session.add_all([dag_run, task_instance])
-                #     serializable_context: dict[Encoding, Any] = # Get serializable context here
-                # with airflow_context_path.open("w+") as file:
-                #     json.dump(serializable_context, file)
 
             env_vars = dict(os.environ) if self.inherit_env else {}
             if self.env_vars:
@@ -583,7 +574,11 @@ class _BasePythonVirtualenvOperator(PythonOperator, metaclass=ABCMeta):
             return self._read_result(output_path)
 
     def determine_kwargs(self, context: Mapping[str, Any]) -> Mapping[str, Any]:
-        return KeywordParameters.determine(self.python_callable, self.op_args, context).serializing()
+        keyword_params = KeywordParameters.determine(self.python_callable, self.op_args, context)
+        if AIRFLOW_V_3_0_PLUS:
+            return keyword_params.unpacking()
+        else:
+            return keyword_params.serializing()  # type: ignore[attr-defined]
 
 
 class PythonVirtualenvOperator(_BasePythonVirtualenvOperator):
@@ -651,8 +646,6 @@ class PythonVirtualenvOperator(_BasePythonVirtualenvOperator):
         environment. If set to ``True``, the virtual environment will inherit the environment variables
         of the parent process (``os.environ``). If set to ``False``, the virtual environment will be
         executed with a clean environment.
-    :param use_airflow_context: Whether to provide ``get_current_context()`` to the python_callable.
-        NOT YET IMPLEMENTED - waits for AIP-72 context serialization.
     """
 
     template_fields: Sequence[str] = tuple(
@@ -680,7 +673,6 @@ class PythonVirtualenvOperator(_BasePythonVirtualenvOperator):
         venv_cache_path: None | os.PathLike[str] = None,
         env_vars: dict[str, str] | None = None,
         inherit_env: bool = True,
-        use_airflow_context: bool = False,
         **kwargs,
     ):
         if (
@@ -697,18 +689,6 @@ class PythonVirtualenvOperator(_BasePythonVirtualenvOperator):
             raise AirflowException(
                 "Passing non-string types (e.g. int or float) as python_version not supported"
             )
-        if use_airflow_context and (not expect_airflow and not system_site_packages):
-            raise AirflowException(
-                "The `use_airflow_context` parameter is set to True, but "
-                "expect_airflow and system_site_packages are set to False."
-            )
-        # TODO: remove when context serialization is implemented in AIP-72
-        if use_airflow_context and not AIRFLOW_V_3_0_PLUS:
-            raise AirflowException(
-                "The `use_airflow_context=True` is not yet implemented. "
-                "It will work in Airflow 3 after AIP-72 context "
-                "serialization is ready."
-            )
         if not requirements:
             self.requirements: list[str] = []
         elif isinstance(requirements, str):
@@ -737,7 +717,6 @@ class PythonVirtualenvOperator(_BasePythonVirtualenvOperator):
             skip_on_exit_code=skip_on_exit_code,
             env_vars=env_vars,
             inherit_env=inherit_env,
-            use_airflow_context=use_airflow_context,
             **kwargs,
         )
 
@@ -953,8 +932,6 @@ class ExternalPythonOperator(_BasePythonVirtualenvOperator):
         environment. If set to ``True``, the virtual environment will inherit the environment variables
         of the parent process (``os.environ``). If set to ``False``, the virtual environment will be
         executed with a clean environment.
-    :param use_airflow_context: Whether to provide ``get_current_context()`` to the python_callable.
-        NOT YET IMPLEMENTED - waits for AIP-72 context serialization.
     """
 
     template_fields: Sequence[str] = tuple({"python"}.union(PythonOperator.template_fields))
@@ -975,22 +952,10 @@ class ExternalPythonOperator(_BasePythonVirtualenvOperator):
         skip_on_exit_code: int | Container[int] | None = None,
         env_vars: dict[str, str] | None = None,
         inherit_env: bool = True,
-        use_airflow_context: bool = False,
         **kwargs,
     ):
         if not python:
             raise ValueError("Python Path must be defined in ExternalPythonOperator")
-        if use_airflow_context and not expect_airflow:
-            raise AirflowException(
-                "The `use_airflow_context` parameter is set to True, but expect_airflow is set to False."
-            )
-        # TODO: remove when context serialization is implemented in AIP-72
-        if use_airflow_context:
-            raise AirflowException(
-                "The `use_airflow_context=True` is not yet implemented. "
-                "It will work in Airflow 3 after AIP-72 context "
-                "serialization is ready."
-            )
         self.python = python
         self.expect_pendulum = expect_pendulum
         super().__init__(
@@ -1005,7 +970,6 @@ class ExternalPythonOperator(_BasePythonVirtualenvOperator):
             skip_on_exit_code=skip_on_exit_code,
             env_vars=env_vars,
             inherit_env=inherit_env,
-            use_airflow_context=use_airflow_context,
             **kwargs,
         )
 
@@ -1120,7 +1084,7 @@ class BranchExternalPythonOperator(ExternalPythonOperator, BranchMixIn):
         return self.do_branch(context, super().execute(context))
 
 
-def get_current_context() -> Context:
+def get_current_context() -> Mapping[str, Any]:
     """
     Retrieve the execution context dictionary without altering user method's signature.
 
@@ -1147,9 +1111,22 @@ def get_current_context() -> Context:
     Current context will only have value if this method was called after an operator
     was starting to execute.
     """
+    if AIRFLOW_V_3_0_PLUS:
+        from airflow.sdk import get_current_context
+
+        return get_current_context()
+    else:
+        return _get_current_context()
+
+
+def _get_current_context() -> Mapping[str, Any]:
+    # Airflow 2.x
+    # TODO: To be removed when Airflow 2 support is dropped
+    from airflow.models.taskinstance import _CURRENT_CONTEXT  # type: ignore[attr-defined]
+
     if not _CURRENT_CONTEXT:
-        raise AirflowException(
+        raise RuntimeError(
             "Current context was requested but no context was found! "
-            "Are you running within an airflow task?"
+            "Are you running within an Airflow task?"
         )
     return _CURRENT_CONTEXT[-1]

airflow/providers/standard/operators/trigger_dagrun.py

@@ -21,7 +21,7 @@ import datetime
 import json
 import time
 from collections.abc import Sequence
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any
 
 from sqlalchemy import select
 from sqlalchemy.orm.exc import NoResultFound
@@ -54,7 +54,12 @@ if TYPE_CHECKING:
     from sqlalchemy.orm.session import Session
 
     from airflow.models.taskinstancekey import TaskInstanceKey
-    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 TriggerDagRunLink(BaseOperatorLink):
@@ -68,29 +73,20 @@ class TriggerDagRunLink(BaseOperatorLink):
 
     def get_link(self, operator: BaseOperator, *, ti_key: TaskInstanceKey) -> str:
         from airflow.models.renderedtifields import RenderedTaskInstanceFields
-        from airflow.models.taskinstance import TaskInstance
 
-        ti = TaskInstance.get_task_instance(
-            dag_id=ti_key.dag_id, run_id=ti_key.run_id, task_id=ti_key.task_id, map_index=ti_key.map_index
-        )
         if TYPE_CHECKING:
-            assert ti is not None
+            assert isinstance(operator, TriggerDagRunOperator)
 
-        template_fields = RenderedTaskInstanceFields.get_templated_fields(ti)
-        untemplated_trigger_dag_id = cast(TriggerDagRunOperator, operator).trigger_dag_id
-        if template_fields:
-            trigger_dag_id = template_fields.get("trigger_dag_id", untemplated_trigger_dag_id)
+        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 = untemplated_trigger_dag_id
+            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,
-        }
+        query = {"dag_id": trigger_dag_id, "dag_run_id": triggered_dag_run_id}
         return build_airflow_url_with_query(query)
 
 
@@ -181,12 +177,10 @@ class TriggerDagRunOperator(BaseOperator):
         self.logical_date = logical_date
 
     def execute(self, context: Context):
-        if isinstance(self.logical_date, datetime.datetime):
+        if self.logical_date is None or isinstance(self.logical_date, datetime.datetime):
             parsed_logical_date = self.logical_date
-        elif isinstance(self.logical_date, str):
-            parsed_logical_date = timezone.parse(self.logical_date)
         else:
-            parsed_logical_date = timezone.utcnow()
+            parsed_logical_date = timezone.parse(self.logical_date)
 
         try:
             json.dumps(self.conf)
@@ -196,7 +190,11 @@ class TriggerDagRunOperator(BaseOperator):
         if self.trigger_run_id:
             run_id = str(self.trigger_run_id)
         else:
-            run_id = DagRun.generate_run_id(DagRunType.MANUAL, parsed_logical_date)
+            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(
@@ -211,7 +209,7 @@ class TriggerDagRunOperator(BaseOperator):
         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.logical_date)
+                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)
@@ -221,7 +219,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(start_date=dag_run.logical_date, end_date=dag_run.logical_date)
+                dag.clear(run_id=dag_run.run_id)
             else:
                 if self.skip_when_already_exists:
                     raise AirflowSkipException(
@@ -242,7 +240,7 @@ class TriggerDagRunOperator(BaseOperator):
                     trigger=DagStateTrigger(
                         dag_id=self.trigger_dag_id,
                         states=self.allowed_states + self.failed_states,
-                        logical_dates=[dag_run.logical_date],
+                        run_ids=[run_id],
                         poll_interval=self.poke_interval,
                     ),
                     method_name="execute_complete",
@@ -252,7 +250,7 @@ class TriggerDagRunOperator(BaseOperator):
                 self.log.info(
                     "Waiting for %s on %s to become allowed state %s ...",
                     self.trigger_dag_id,
-                    dag_run.logical_date,
+                    run_id,
                     self.allowed_states,
                 )
                 time.sleep(self.poke_interval)
@@ -268,18 +266,16 @@ class TriggerDagRunOperator(BaseOperator):
 
     @provide_session
     def execute_complete(self, context: Context, session: Session, event: tuple[str, dict[str, Any]]):
-        # This logical_date is parsed from the return trigger event
-        provided_logical_date = event[1]["logical_dates"][0]
+        # 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.logical_date == provided_logical_date
-                )
+                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 logical date {self.logical_date}"
+                f"No DAG run found for DAG {self.trigger_dag_id} and run ID {provided_run_id}"
             )
 
         state = dag_run.state

airflow/providers/standard/operators/weekday.py

@@ -25,7 +25,11 @@ from airflow.utils import timezone
 from airflow.utils.weekday import WeekDay
 
 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 BranchDayOfWeekOperator(BaseBranchOperator):
@@ -39,7 +43,7 @@ class BranchDayOfWeekOperator(BaseBranchOperator):
 
     .. code-block:: python
 
-        from airflow.operators.empty import EmptyOperator
+        from airflow.providers.standard.operators.empty import EmptyOperator
         from airflow.operators.weekday import BranchDayOfWeekOperator
 
         monday = EmptyOperator(task_id="monday")
@@ -60,7 +64,7 @@ class BranchDayOfWeekOperator(BaseBranchOperator):
 
         # import WeekDay Enum
         from airflow.utils.weekday import WeekDay
-        from airflow.operators.empty import EmptyOperator
+        from airflow.providers.standard.operators.empty import EmptyOperator
         from airflow.operators.weekday import BranchDayOfWeekOperator
 
         workday = EmptyOperator(task_id="workday")

airflow/providers/standard/sensors/bash.py

@@ -27,7 +27,11 @@ from airflow.exceptions import AirflowFailException
 from airflow.sensors.base import BaseSensorOperator
 
 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 BashSensor(BaseSensorOperator):