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

airflow/providers/standard/operators/smooth.py

@@ -0,0 +1,38 @@
+#
+# 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 SmoothOperator(BaseOperator):
+    """Operator that logs a YouTube link to Sade song "Smooth Operator"."""
+
+    ui_color = "#e8f7e4"
+    yt_link: str = "https://www.youtube.com/watch?v=4TYv2PhG89A"
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+
+    def execute(self, context: Context):
+        self.log.info("Enjoy Sade - Smooth Operator: %s", self.yt_link)

airflow/providers/standard/operators/trigger_dagrun.py

@@ -0,0 +1,370 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import datetime
+import json
+import time
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from sqlalchemy import select
+from sqlalchemy.orm.exc import NoResultFound
+
+from airflow.api.common.trigger_dag import trigger_dag
+from airflow.configuration import conf
+from airflow.exceptions import (
+    AirflowException,
+    AirflowSkipException,
+    DagNotFound,
+    DagRunAlreadyExists,
+)
+from airflow.models import BaseOperator
+from airflow.models.dag import DagModel
+from airflow.models.dagbag import DagBag
+from airflow.models.dagrun import DagRun
+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.state import DagRunState
+from airflow.utils.types import NOTSET, ArgNotSet, DagRunType
+
+XCOM_LOGICAL_DATE_ISO = "trigger_logical_date_iso"
+XCOM_RUN_ID = "trigger_run_id"
+
+
+if TYPE_CHECKING:
+    from sqlalchemy.orm.session import Session
+
+    from airflow.models.taskinstancekey import TaskInstanceKey
+
+    try:
+        from airflow.sdk.definitions.context import Context
+    except ImportError:
+        # TODO: Remove once provider drops support for Airflow 2
+        from airflow.utils.context import Context
+
+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]
+
+
+class TriggerDagRunLink(BaseOperatorLink):
+    """
+    Operator link for TriggerDagRunOperator.
+
+    It allows users to access DAG triggered by task using TriggerDagRunOperator.
+    """
+
+    name = "Triggered DAG"
+
+    def get_link(self, operator: BaseOperator, *, ti_key: TaskInstanceKey) -> str:
+        if TYPE_CHECKING:
+            assert isinstance(operator, TriggerDagRunOperator)
+
+        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.
+        triggered_dag_run_id = XCom.get_value(ti_key=ti_key, key=XCOM_RUN_ID)
+
+        if AIRFLOW_V_3_0_PLUS:
+            from airflow.utils.helpers import build_airflow_dagrun_url
+
+            return build_airflow_dagrun_url(dag_id=trigger_dag_id, run_id=triggered_dag_run_id)
+        from airflow.utils.helpers import build_airflow_url_with_query  # type:ignore[attr-defined]
+
+        query = {"dag_id": trigger_dag_id, "dag_run_id": triggered_dag_run_id}
+        return build_airflow_url_with_query(query)
+
+
+class TriggerDagRunOperator(BaseOperator):
+    """
+    Triggers a DAG run for a specified DAG ID.
+
+    Note that if database isolation mode is enabled, not all features are supported.
+
+    :param trigger_dag_id: The ``dag_id`` of the DAG to trigger (templated).
+    :param trigger_run_id: The run ID to use for the triggered DAG run (templated).
+        If not provided, a run ID will be automatically generated.
+    :param conf: Configuration for the DAG run (templated).
+    :param logical_date: Logical date for the triggered DAG (templated).
+    :param reset_dag_run: Whether clear existing DAG run if already exists.
+        This is useful when backfill or rerun an existing DAG run.
+        This only resets (not recreates) the DAG run.
+        DAG run conf is immutable and will not be reset on rerun of an existing DAG run.
+        When reset_dag_run=False and dag run exists, DagRunAlreadyExists will be raised.
+        When reset_dag_run=True and dag run exists, existing DAG run will be cleared to rerun.
+    :param wait_for_completion: Whether or not wait for DAG run completion. (default: False)
+    :param poke_interval: Poke interval to check DAG run status when wait_for_completion=True.
+        (default: 60)
+    :param allowed_states: Optional list of allowed DAG run states of the triggered DAG. This is useful when
+        setting ``wait_for_completion`` to True. Must be a valid DagRunState.
+        Default is ``[DagRunState.SUCCESS]``.
+    :param failed_states: Optional list of failed or disallowed DAG run states of the triggered DAG. This is
+        useful when setting ``wait_for_completion`` to True. Must be a valid DagRunState.
+        Default is ``[DagRunState.FAILED]``.
+    :param skip_when_already_exists: Set to true to mark the task as SKIPPED if a DAG run of the triggered
+        DAG for the same logical date already exists.
+    :param deferrable: If waiting for completion, whether or not to defer the task until done,
+        default is ``False``.
+    """
+
+    template_fields: Sequence[str] = (
+        "trigger_dag_id",
+        "trigger_run_id",
+        "logical_date",
+        "conf",
+        "wait_for_completion",
+        "skip_when_already_exists",
+    )
+    template_fields_renderers = {"conf": "py"}
+    ui_color = "#ffefeb"
+    operator_extra_links = [TriggerDagRunLink()]
+
+    def __init__(
+        self,
+        *,
+        trigger_dag_id: str,
+        trigger_run_id: str | None = None,
+        conf: dict | None = None,
+        logical_date: str | datetime.datetime | None | ArgNotSet = NOTSET,
+        reset_dag_run: bool = False,
+        wait_for_completion: bool = False,
+        poke_interval: int = 60,
+        allowed_states: list[str | DagRunState] | None = None,
+        failed_states: list[str | DagRunState] | None = None,
+        skip_when_already_exists: bool = False,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.trigger_dag_id = trigger_dag_id
+        self.trigger_run_id = trigger_run_id
+        self.conf = conf
+        self.reset_dag_run = reset_dag_run
+        self.wait_for_completion = wait_for_completion
+        self.poke_interval = poke_interval
+        if allowed_states:
+            self.allowed_states = [DagRunState(s) for s in allowed_states]
+        else:
+            self.allowed_states = [DagRunState.SUCCESS]
+        if failed_states is not None:
+            self.failed_states = [DagRunState(s) for s in failed_states]
+        else:
+            self.failed_states = [DagRunState.FAILED]
+        self.skip_when_already_exists = skip_when_already_exists
+        self._defer = deferrable
+        self.logical_date = logical_date
+        if logical_date is NOTSET:
+            self.logical_date = NOTSET
+        elif logical_date is None or isinstance(logical_date, (str, datetime.datetime)):
+            self.logical_date = logical_date
+        else:
+            raise TypeError(
+                f"Expected str, datetime.datetime, or None for parameter 'logical_date'. Got {type(logical_date).__name__}"
+            )
+
+    def execute(self, context: Context):
+        if self.logical_date is NOTSET:
+            # If no logical_date is provided we will set utcnow()
+            parsed_logical_date = timezone.utcnow()
+        elif self.logical_date is None or isinstance(self.logical_date, datetime.datetime):
+            parsed_logical_date = self.logical_date  # type: ignore
+        elif isinstance(self.logical_date, str):
+            parsed_logical_date = timezone.parse(self.logical_date)
+
+        try:
+            json.dumps(self.conf)
+        except TypeError:
+            raise ValueError("conf parameter should be JSON Serializable")
+
+        if self.trigger_run_id:
+            run_id = str(self.trigger_run_id)
+        else:
+            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,
+            deferrable=self._defer,
+        )
+
+    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,
+                execution_date=parsed_logical_date,
+                replace_microseconds=False,
+            )
+
+        except DagRunAlreadyExists as e:
+            if self.reset_dag_run:
+                dag_run = e.dag_run
+                self.log.info("Clearing %s on %s", self.trigger_dag_id, dag_run.run_id)
+
+                # Get target dag object and call clear()
+                dag_model = DagModel.get_current(self.trigger_dag_id)
+                if dag_model is None:
+                    raise DagNotFound(f"Dag id {self.trigger_dag_id} not found in DagModel")
+
+                # Note: here execution fails on database isolation mode. Needs structural changes for AIP-72
+                dag_bag = DagBag(dag_folder=dag_model.fileloc, read_dags_from_db=True)
+                dag = dag_bag.get_dag(self.trigger_dag_id)
+                dag.clear(start_date=dag_run.logical_date, end_date=dag_run.logical_date)
+            else:
+                if self.skip_when_already_exists:
+                    raise AirflowSkipException(
+                        "Skipping due to skip_when_already_exists is set to True and DagRunAlreadyExists"
+                    )
+                raise e
+        if dag_run is None:
+            raise RuntimeError("The dag_run should be set here!")
+        # Store the run id from the dag run (either created or found above) to
+        # be used when creating the extra link on the webserver.
+        ti = context["task_instance"]
+        ti.xcom_push(key=XCOM_RUN_ID, value=dag_run.run_id)
+
+        if self.wait_for_completion:
+            # Kick off the deferral process
+            if self._defer:
+                self.defer(
+                    trigger=DagStateTrigger(
+                        dag_id=self.trigger_dag_id,
+                        states=self.allowed_states + self.failed_states,
+                        execution_dates=[dag_run.logical_date],
+                        run_ids=[run_id],
+                        poll_interval=self.poke_interval,
+                    ),
+                    method_name="execute_complete",
+                )
+            # wait for dag to complete
+            while True:
+                self.log.info(
+                    "Waiting for %s on %s to become allowed state %s ...",
+                    self.trigger_dag_id,
+                    run_id,
+                    self.allowed_states,
+                )
+                time.sleep(self.poke_interval)
+
+                # Note: here execution fails on database isolation mode. Needs structural changes for AIP-72
+                dag_run.refresh_from_db()
+                state = dag_run.state
+                if state in self.failed_states:
+                    raise AirflowException(f"{self.trigger_dag_id} failed with failed states {state}")
+                if state in self.allowed_states:
+                    self.log.info("%s finished with allowed state %s", self.trigger_dag_id, state)
+                    return
+
+    def execute_complete(self, context: Context, event: tuple[str, dict[str, Any]]):
+        if AIRFLOW_V_3_0_PLUS:
+            self._trigger_dag_run_af_3_execute_complete(event=event)
+        else:
+            self._trigger_dag_run_af_2_execute_complete(event=event)
+
+    def _trigger_dag_run_af_3_execute_complete(self, event: tuple[str, dict[str, Any]]):
+        run_ids = event[1]["run_ids"]
+        event_data = event[1]
+        failed_run_id_conditions = []
+
+        for run_id in run_ids:
+            state = event_data.get(run_id)
+            if state in self.failed_states:
+                failed_run_id_conditions.append(run_id)
+                continue
+            if state in self.allowed_states:
+                self.log.info(
+                    "%s finished with allowed state %s for run_id %s",
+                    self.trigger_dag_id,
+                    state,
+                    run_id,
+                )
+
+        if failed_run_id_conditions:
+            raise AirflowException(
+                f"{self.trigger_dag_id} failed with failed states {self.failed_states} for run_ids"
+                f" {failed_run_id_conditions}"
+            )
+
+    if not AIRFLOW_V_3_0_PLUS:
+        from airflow.utils.session import NEW_SESSION, provide_session  # type: ignore[misc]
+
+        @provide_session
+        def _trigger_dag_run_af_2_execute_complete(
+            self, event: tuple[str, dict[str, Any]], session: Session = NEW_SESSION
+        ):
+            # 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.execution_date == provided_logical_date
+                    )
+                ).scalar_one()
+            except NoResultFound:
+                raise AirflowException(
+                    f"No DAG run found for DAG {self.trigger_dag_id} and logical date {self.logical_date}"
+                )
+
+            state = dag_run.state
+
+            if state in self.failed_states:
+                raise AirflowException(f"{self.trigger_dag_id} failed with failed state {state}")
+            if state in self.allowed_states:
+                self.log.info("%s finished with allowed state %s", self.trigger_dag_id, state)
+                return
+
+            raise AirflowException(
+                f"{self.trigger_dag_id} return {state} which is not in {self.failed_states}"
+                f" or {self.allowed_states}"
+            )

airflow/providers/standard/operators/weekday.py

@@ -17,14 +17,19 @@
 # under the License.
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Iterable
+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:
-    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):
@@ -38,7 +43,8 @@ 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")
         other_day = EmptyOperator(task_id="other_day")
@@ -57,8 +63,9 @@ class BranchDayOfWeekOperator(BaseBranchOperator):
     .. code-block:: python
 
         # import WeekDay Enum
-        from airflow.utils.weekday import WeekDay
-        from airflow.operators.empty import EmptyOperator
+        from airflow.providers.standard.utils.weekday import WeekDay
+        from airflow.providers.standard.operators.empty import EmptyOperator
+        from airflow.operators.weekday import BranchDayOfWeekOperator
 
         workday = EmptyOperator(task_id="workday")
         weekend = EmptyOperator(task_id="weekend")
@@ -109,10 +116,13 @@ class BranchDayOfWeekOperator(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.make_naive(timezone.utcnow(), self.dag.timezone)
 
-        if now.isoweekday() in self._week_day_num:
+        if now.isoweekday() in self._week_day_num:  # type: ignore[union-attr]
             return self.follow_task_ids_if_true
         return self.follow_task_ids_if_false

airflow/providers/standard/sensors/bash.py

@@ -0,0 +1,118 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import os
+from collections.abc import Sequence
+from subprocess import PIPE, STDOUT, Popen
+from tempfile import NamedTemporaryFile, TemporaryDirectory, gettempdir
+from typing import TYPE_CHECKING
+
+from airflow.exceptions import AirflowFailException
+from airflow.sensors.base import BaseSensorOperator
+
+if TYPE_CHECKING:
+    try:
+        from airflow.sdk.definitions.context import Context
+    except ImportError:
+        # TODO: Remove once provider drops support for Airflow 2
+        from airflow.utils.context import Context
+
+
+class BashSensor(BaseSensorOperator):
+    """
+    Executes a bash command/script.
+
+    Return True if and only if the return code is 0.
+
+    :param bash_command: The command, set of commands or reference to a
+        bash script (must be '.sh') to be executed.
+
+    :param env: If env is not None, it must be a mapping that defines the
+        environment variables for the new process; these are used instead
+        of inheriting the current process environment, which is the default
+        behavior. (templated)
+    :param output_encoding: output encoding of bash command.
+    :param retry_exit_code: If task exits with this code, treat the sensor
+        as not-yet-complete and retry the check later according to the
+        usual retry/timeout settings. Any other non-zero return code will
+        be treated as an error, and cause the sensor to fail. If set to
+        ``None`` (the default), any non-zero exit code will cause a retry
+        and the task will never raise an error except on time-out.
+
+    .. seealso::
+        For more information on how to use this sensor,take a look at the guide:
+        :ref:`howto/operator:BashSensor`
+    """
+
+    template_fields: Sequence[str] = ("bash_command", "env")
+
+    def __init__(
+        self, *, bash_command, env=None, output_encoding="utf-8", retry_exit_code: int | None = None, **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.bash_command = bash_command
+        self.env = env
+        self.output_encoding = output_encoding
+        self.retry_exit_code = retry_exit_code
+
+    def poke(self, context: Context):
+        """Execute the bash command in a temporary directory."""
+        bash_command = self.bash_command
+        self.log.info("Tmp dir root location: %s", gettempdir())
+        with (
+            TemporaryDirectory(prefix="airflowtmp") as tmp_dir,
+            NamedTemporaryFile(dir=tmp_dir, prefix=self.task_id) as f,
+        ):
+            f.write(bytes(bash_command, "utf_8"))
+            f.flush()
+            fname = f.name
+            script_location = tmp_dir + "/" + fname
+            self.log.info("Temporary script location: %s", script_location)
+            self.log.info("Running command: %s", bash_command)
+
+            with Popen(
+                ["bash", fname],
+                stdout=PIPE,
+                stderr=STDOUT,
+                close_fds=True,
+                cwd=tmp_dir,
+                env=self.env,
+                preexec_fn=os.setsid,
+            ) as resp:
+                if resp.stdout:
+                    self.log.info("Output:")
+                    for line in iter(resp.stdout.readline, b""):
+                        self.log.info(line.decode(self.output_encoding).strip())
+                resp.wait()
+                self.log.info("Command exited with return code %s", resp.returncode)
+
+                # zero code means success, the sensor can go green
+                if resp.returncode == 0:
+                    return True
+
+                # we have a retry exit code, sensor retries if return code matches, otherwise error
+                if self.retry_exit_code is not None:
+                    if resp.returncode == self.retry_exit_code:
+                        self.log.info("Return code matches retry code, will retry later")
+                        return False
+                    raise AirflowFailException(f"Command exited with return code {resp.returncode}")
+
+                # backwards compatibility: sensor retries no matter the error code
+                self.log.info("Non-zero return code and no retry code set, will retry later")
+                return False

airflow/providers/standard/sensors/date_time.py

@@ -18,15 +18,37 @@
 from __future__ import annotations
 
 import datetime
-from typing import TYPE_CHECKING, Any, NoReturn, Sequence
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, NoReturn
 
+from airflow.providers.standard.triggers.temporal import DateTimeTrigger
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.sensors.base import BaseSensorOperator
-from airflow.triggers.base import StartTriggerArgs
-from airflow.triggers.temporal import DateTimeTrigger
+
+try:
+    from airflow.triggers.base import StartTriggerArgs
+except ImportError:
+    # TODO: Remove this when min airflow version is 2.10.0 for standard provider
+    @dataclass
+    class StartTriggerArgs:  # type: ignore[no-redef]
+        """Arguments required for start task execution from triggerer."""
+
+        trigger_cls: str
+        next_method: str
+        trigger_kwargs: dict[str, Any] | None = None
+        next_kwargs: dict[str, Any] | None = None
+        timeout: datetime.timedelta | None = None
+
+
 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 DateTimeSensor(BaseSensorOperator):
@@ -37,7 +59,7 @@ class DateTimeSensor(BaseSensorOperator):
     It handles some cases for which ``TimeSensor`` and ``TimeDeltaSensor`` are not suited.
 
     **Example** 1 :
-        If a task needs to wait for 11am on each ``execution_date``. Using
+        If a task needs to wait for 11am on each ``logical_date``. Using
         ``TimeSensor`` or ``TimeDeltaSensor``, all backfill tasks started at
         1am have to wait for 10 hours. This is unnecessary, e.g. a backfill
         task with ``{{ ds }} = '1970-01-01'`` does not need to wait because
@@ -52,7 +74,7 @@ class DateTimeSensor(BaseSensorOperator):
 
             DateTimeSensor(
                 task_id="wait_for_0100",
-                target_time="{{ next_execution_date.tomorrow().replace(hour=1) }}",
+                target_time="{{ data_interval_end.tomorrow().replace(hour=1) }}",
             )
 
     :param target_time: datetime after which the job succeeds. (templated)
@@ -93,7 +115,7 @@ class DateTimeSensorAsync(DateTimeSensor):
     """
 
     start_trigger_args = StartTriggerArgs(
-        trigger_cls="airflow.triggers.temporal.DateTimeTrigger",
+        trigger_cls="airflow.providers.standard.triggers.temporal.DateTimeTrigger",
         trigger_kwargs={"moment": "", "end_from_trigger": False},
         next_method="execute_complete",
         next_kwargs=None,
@@ -125,7 +147,9 @@ class DateTimeSensorAsync(DateTimeSensor):
             trigger=DateTimeTrigger(
                 moment=timezone.parse(self.target_time),
                 end_from_trigger=self.end_from_trigger,
-            ),
+            )
+            if AIRFLOW_V_3_0_PLUS
+            else DateTimeTrigger(moment=timezone.parse(self.target_time)),
         )
 
     def execute_complete(self, context: Context, event: Any = None) -> None: