apache-airflow-providers-standard 1.9.1rc1__py3-none-any.whl → 1.10.3__py3-none-any.whl

airflow/providers/standard/triggers/hitl.py

@@ -16,7 +16,7 @@
 # under the License.
 from __future__ import annotations
 
-from airflow.exceptions import AirflowOptionalProviderFeatureException
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
 from airflow.providers.standard.version_compat import AIRFLOW_V_3_1_PLUS
 
 if not AIRFLOW_V_3_1_PLUS:
@@ -30,6 +30,9 @@ from uuid import UUID
 
 from asgiref.sync import sync_to_async
 
+from airflow.providers.common.compat.sdk import ParamValidationError
+from airflow.sdk import Param
+from airflow.sdk.definitions.param import ParamsDict
 from airflow.sdk.execution_time.hitl import (
     HITLUser,
     get_hitl_detail_content_detail,
@@ -43,7 +46,7 @@ class HITLTriggerEventSuccessPayload(TypedDict, total=False):
     """Minimum required keys for a success Human-in-the-loop TriggerEvent."""
 
     chosen_options: list[str]
-    params_input: dict[str, Any]
+    params_input: dict[str, dict[str, Any]]
     responded_by_user: HITLUser | None
     responded_at: datetime
     timedout: bool
@@ -53,7 +56,7 @@ class HITLTriggerEventFailurePayload(TypedDict):
     """Minimum required keys for a failed Human-in-the-loop TriggerEvent."""
 
     error: str
-    error_type: Literal["timeout", "unknown"]
+    error_type: Literal["timeout", "unknown", "validation"]
 
 
 class HITLTrigger(BaseTrigger):
@@ -64,7 +67,7 @@ class HITLTrigger(BaseTrigger):
         *,
         ti_id: UUID,
         options: list[str],
-        params: dict[str, Any],
+        params: dict[str, dict[str, Any]],
         defaults: list[str] | None = None,
         multiple: bool = False,
         timeout_datetime: datetime | None,
@@ -80,7 +83,21 @@ class HITLTrigger(BaseTrigger):
         self.defaults = defaults
         self.timeout_datetime = timeout_datetime
 
-        self.params = params
+        self.params = ParamsDict(
+            {
+                k: Param(
+                    v.pop("value"),
+                    **v,
+                )
+                if HITLTrigger._is_param(v)
+                else Param(v)
+                for k, v in params.items()
+            },
+        )
+
+    @staticmethod
+    def _is_param(value: Any) -> bool:
+        return isinstance(value, dict) and all(key in value for key in ("description", "schema", "value"))
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         """Serialize HITLTrigger arguments and classpath."""
@@ -90,99 +107,131 @@ class HITLTrigger(BaseTrigger):
                 "ti_id": self.ti_id,
                 "options": self.options,
                 "defaults": self.defaults,
-                "params": self.params,
+                "params": {k: self.params.get_param(k).serialize() for k in self.params},
                 "multiple": self.multiple,
                 "timeout_datetime": self.timeout_datetime,
                 "poke_interval": self.poke_interval,
             },
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:
-        """Loop until the Human-in-the-loop response received or timeout reached."""
-        while True:
-            if self.timeout_datetime and self.timeout_datetime < utcnow():
-                # Fetch latest HITL detail before fallback
-                resp = await sync_to_async(get_hitl_detail_content_detail)(ti_id=self.ti_id)
-                # Response already received, yield success and exit
-                if resp.response_received and resp.chosen_options:
-                    if TYPE_CHECKING:
-                        assert resp.responded_by_user is not None
-                        assert resp.responded_at is not None
-
-                    self.log.info(
-                        "[HITL] responded_by=%s (id=%s) options=%s at %s (timeout fallback skipped)",
-                        resp.responded_by_user.name,
-                        resp.responded_by_user.id,
-                        resp.chosen_options,
-                        resp.responded_at,
-                    )
-                    yield TriggerEvent(
-                        HITLTriggerEventSuccessPayload(
-                            chosen_options=resp.chosen_options,
-                            params_input=resp.params_input or {},
-                            responded_at=resp.responded_at,
-                            responded_by_user=HITLUser(
-                                id=resp.responded_by_user.id,
-                                name=resp.responded_by_user.name,
-                            ),
-                            timedout=False,
-                        )
-                    )
-                    return
-
-                if self.defaults is None:
-                    yield TriggerEvent(
-                        HITLTriggerEventFailurePayload(
-                            error="The timeout has passed, and the response has not yet been received.",
-                            error_type="timeout",
-                        )
-                    )
-                    return
-
-                resp = await sync_to_async(update_hitl_detail_response)(
-                    ti_id=self.ti_id,
-                    chosen_options=self.defaults,
-                    params_input=self.params,
+    async def _handle_timeout(self) -> TriggerEvent:
+        """Handle HITL timeout logic and yield appropriate event."""
+        resp = await sync_to_async(get_hitl_detail_content_detail)(ti_id=self.ti_id)
+
+        # Case 1: Response arrived just before timeout
+        if resp.response_received and resp.chosen_options:
+            if TYPE_CHECKING:
+                assert resp.responded_by_user is not None
+                assert resp.responded_at is not None
+
+            chosen_options_list = list(resp.chosen_options or [])
+            self.log.info(
+                "[HITL] responded_by=%s (id=%s) options=%s at %s (timeout fallback skipped)",
+                resp.responded_by_user.name,
+                resp.responded_by_user.id,
+                chosen_options_list,
+                resp.responded_at,
+            )
+            return TriggerEvent(
+                HITLTriggerEventSuccessPayload(
+                    chosen_options=chosen_options_list,
+                    params_input=resp.params_input or {},
+                    responded_at=resp.responded_at,
+                    responded_by_user=HITLUser(
+                        id=resp.responded_by_user.id,
+                        name=resp.responded_by_user.name,
+                    ),
+                    timedout=False,
                 )
-                if TYPE_CHECKING:
-                    assert resp.responded_at is not None
-                self.log.info(
-                    "[HITL] timeout reached before receiving response, fallback to default %s", self.defaults
+            )
+
+        # Case 2: No defaults defined → failure
+        if self.defaults is None:
+            return TriggerEvent(
+                HITLTriggerEventFailurePayload(
+                    error="The timeout has passed, and the response has not yet been received.",
+                    error_type="timeout",
                 )
-                yield TriggerEvent(
-                    HITLTriggerEventSuccessPayload(
-                        chosen_options=self.defaults,
-                        params_input=self.params,
-                        responded_by_user=None,
-                        responded_at=resp.responded_at,
-                        timedout=True,
+            )
+
+        # Case 3: Timeout fallback to default
+        resp = await sync_to_async(update_hitl_detail_response)(
+            ti_id=self.ti_id,
+            chosen_options=self.defaults,
+            params_input=self.params.dump(),
+        )
+        if TYPE_CHECKING:
+            assert resp.responded_at is not None
+
+        self.log.info(
+            "[HITL] timeout reached before receiving response, fallback to default %s",
+            self.defaults,
+        )
+        return TriggerEvent(
+            HITLTriggerEventSuccessPayload(
+                chosen_options=self.defaults,
+                params_input=self.params.dump(),
+                responded_by_user=None,
+                responded_at=resp.responded_at,
+                timedout=True,
+            )
+        )
+
+    async def _handle_response(self):
+        """Check if HITL response is ready and yield success if so."""
+        resp = await sync_to_async(get_hitl_detail_content_detail)(ti_id=self.ti_id)
+        if TYPE_CHECKING:
+            assert resp.responded_by_user is not None
+            assert resp.responded_at is not None
+
+        if not (resp.response_received and resp.chosen_options):
+            return None
+
+        # validate input
+        if params_input := resp.params_input:
+            try:
+                for key, value in params_input.items():
+                    self.params[key] = value
+            except ParamValidationError as err:
+                return TriggerEvent(
+                    HITLTriggerEventFailurePayload(
+                        error=str(err),
+                        error_type="validation",
                     )
                 )
+
+        chosen_options_list = list(resp.chosen_options or [])
+        self.log.info(
+            "[HITL] responded_by=%s (id=%s) options=%s at %s",
+            resp.responded_by_user.name,
+            resp.responded_by_user.id,
+            chosen_options_list,
+            resp.responded_at,
+        )
+        return TriggerEvent(
+            HITLTriggerEventSuccessPayload(
+                chosen_options=chosen_options_list,
+                params_input=params_input or {},
+                responded_at=resp.responded_at,
+                responded_by_user=HITLUser(
+                    id=resp.responded_by_user.id,
+                    name=resp.responded_by_user.name,
+                ),
+                timedout=False,
+            )
+        )
+
+    async def run(self) -> AsyncIterator[TriggerEvent]:
+        """Loop until the Human-in-the-loop response received or timeout reached."""
+        while True:
+            if self.timeout_datetime and self.timeout_datetime < utcnow():
+                event = await self._handle_timeout()
+                yield event
                 return
 
-            resp = await sync_to_async(get_hitl_detail_content_detail)(ti_id=self.ti_id)
-            if resp.response_received and resp.chosen_options:
-                if TYPE_CHECKING:
-                    assert resp.responded_by_user is not None
-                    assert resp.responded_at is not None
-                self.log.info(
-                    "[HITL] responded_by=%s (id=%s) options=%s at %s",
-                    resp.responded_by_user.name,
-                    resp.responded_by_user.id,
-                    resp.chosen_options,
-                    resp.responded_at,
-                )
-                yield TriggerEvent(
-                    HITLTriggerEventSuccessPayload(
-                        chosen_options=resp.chosen_options,
-                        params_input=resp.params_input or {},
-                        responded_at=resp.responded_at,
-                        responded_by_user=HITLUser(
-                            id=resp.responded_by_user.id,
-                            name=resp.responded_by_user.name,
-                        ),
-                        timedout=False,
-                    )
-                )
+            event = await self._handle_response()
+            if event:
+                yield event
                 return
+
             await asyncio.sleep(self.poke_interval)

airflow/providers/standard/utils/openlineage.py

@@ -0,0 +1,185 @@
+# 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 logging
+from typing import TYPE_CHECKING
+
+from airflow.providers.common.compat.openlineage.check import require_openlineage_version
+from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException
+
+if TYPE_CHECKING:
+    from airflow.models import TaskInstance
+    from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI
+
+log = logging.getLogger(__name__)
+
+OPENLINEAGE_PROVIDER_MIN_VERSION = "2.8.0"
+
+
+def _is_openlineage_provider_accessible() -> bool:
+    """
+    Check if the OpenLineage provider is accessible.
+
+    This function attempts to import the necessary OpenLineage modules and checks if the provider
+    is enabled and the listener is available.
+
+    Returns:
+        bool: True if the OpenLineage provider is accessible, False otherwise.
+    """
+    try:
+        from airflow.providers.openlineage.conf import is_disabled
+        from airflow.providers.openlineage.plugins.listener import get_openlineage_listener
+    except (ImportError, AttributeError):
+        log.debug("OpenLineage provider could not be imported.")
+        return False
+
+    if is_disabled():
+        log.debug("OpenLineage provider is disabled.")
+        return False
+
+    if not get_openlineage_listener():
+        log.debug("OpenLineage listener could not be found.")
+        return False
+
+    return True
+
+
+@require_openlineage_version(provider_min_version=OPENLINEAGE_PROVIDER_MIN_VERSION)
+def _get_openlineage_parent_info(ti: TaskInstance | RuntimeTI) -> dict[str, str]:
+    """Get OpenLineage metadata about the parent task."""
+    from airflow.providers.openlineage.plugins.macros import (
+        lineage_job_name,
+        lineage_job_namespace,
+        lineage_root_job_name,
+        lineage_root_job_namespace,
+        lineage_root_run_id,
+        lineage_run_id,
+    )
+
+    return {
+        "parentRunId": lineage_run_id(ti),
+        "parentJobName": lineage_job_name(ti),
+        "parentJobNamespace": lineage_job_namespace(),
+        "rootParentRunId": lineage_root_run_id(ti),
+        "rootParentJobName": lineage_root_job_name(ti),
+        "rootParentJobNamespace": lineage_root_job_namespace(ti),
+    }
+
+
+def _inject_openlineage_parent_info_to_dagrun_conf(
+    dr_conf: dict | None, ol_parent_info: dict[str, str]
+) -> dict:
+    """
+    Safely inject OpenLineage parent and root run metadata into a DAG run configuration.
+
+    This function adds parent and root job/run identifiers derived from the given TaskInstance into the
+    `openlineage` section of the DAG run configuration. If an `openlineage` key already exists, it is
+    preserved and extended, but no existing parent or root identifiers are overwritten.
+
+    The function performs several safety checks:
+    - If conf is not a dictionary or contains a non-dict `openlineage` section, conf is returned unmodified.
+    - If `openlineage` section contains any parent/root lineage identifiers, conf is returned unmodified.
+
+    Args:
+        dr_conf: The original DAG run configuration dictionary or None.
+        ol_parent_info: OpenLineage metadata about the parent task
+
+    Returns:
+        A modified DAG run conf with injected OpenLineage parent and root metadata,
+        or the original conf if injection is not possible.
+    """
+    current_ol_dr_conf = {}
+    if isinstance(dr_conf, dict) and dr_conf.get("openlineage"):
+        current_ol_dr_conf = dr_conf["openlineage"]
+        if not isinstance(current_ol_dr_conf, dict):
+            log.warning(
+                "Existing 'openlineage' section of DagRun conf is not a dictionary; "
+                "skipping injection of parent metadata."
+            )
+            return dr_conf
+        forbidden_keys = (
+            "parentRunId",
+            "parentJobName",
+            "parentJobNamespace",
+            "rootParentRunId",
+            "rootJobName",
+            "rootJobNamespace",
+        )
+
+        if existing := [k for k in forbidden_keys if k in current_ol_dr_conf]:
+            log.warning(
+                "'openlineage' section of DagRun conf already contains parent or root "
+                "identifiers: `%s`; skipping injection to avoid overwriting existing values.",
+                ", ".join(existing),
+            )
+            return dr_conf
+
+    return {**(dr_conf or {}), **{"openlineage": {**ol_parent_info, **current_ol_dr_conf}}}
+
+
+def safe_inject_openlineage_properties_into_dagrun_conf(
+    dr_conf: dict | None, ti: TaskInstance | RuntimeTI | None
+) -> dict | None:
+    """
+    Safely inject OpenLineage parent task metadata into a DAG run conf.
+
+    This function checks whether the OpenLineage provider is accessible and supports parent information
+    injection. If so, it enriches the DAG run conf with OpenLineage metadata about the parent task
+    to improve lineage tracking. The function does not modify other conf fields, will not overwrite
+    any existing content, and safely returns the original configuration if OpenLineage is unavailable,
+    unsupported, or an error occurs during injection.
+
+    :param dr_conf: The original DAG run configuration dictionary.
+    :param ti: The TaskInstance whose metadata may be injected.
+
+    :return: A potentially enriched DAG run conf with OpenLineage parent information,
+        or the original conf if injection was skipped or failed.
+    """
+    try:
+        if ti is None:
+            log.debug("Task instance not provided - dagrun conf not modified.")
+            return dr_conf
+
+        if not _is_openlineage_provider_accessible():
+            log.debug("OpenLineage provider not accessible - dagrun conf not modified.")
+            return dr_conf
+
+        ol_parent_info = _get_openlineage_parent_info(ti=ti)
+
+        log.info("Injecting openlineage parent task information into dagrun conf.")
+        new_conf = _inject_openlineage_parent_info_to_dagrun_conf(
+            dr_conf=dr_conf.copy() if isinstance(dr_conf, dict) else dr_conf,
+            ol_parent_info=ol_parent_info,
+        )
+        return new_conf
+    except AirflowOptionalProviderFeatureException:
+        log.info(
+            "Current OpenLineage provider version doesn't support parent information in "
+            "the DagRun conf. Upgrade `apache-airflow-providers-openlineage>=%s` to use this feature. "
+            "DagRun conf has not been modified by OpenLineage.",
+            OPENLINEAGE_PROVIDER_MIN_VERSION,
+        )
+        return dr_conf
+    except Exception as e:
+        log.warning(
+            "An error occurred while trying to inject OpenLineage information into dagrun conf. "
+            "DagRun conf has not been modified by OpenLineage. Error: %s",
+            str(e),
+        )
+        log.debug("Error details: ", exc_info=e)
+        return dr_conf

airflow/providers/standard/utils/python_virtualenv.py

@@ -19,16 +19,18 @@
 
 from __future__ import annotations
 
+import logging
 import os
+import shlex
 import shutil
+import subprocess
 import warnings
 from pathlib import Path
 
 import jinja2
 from jinja2 import select_autoescape
 
-from airflow.configuration import conf
-from airflow.utils.process_utils import execute_in_subprocess
+from airflow.providers.common.compat.sdk import conf
 
 
 def _is_uv_installed() -> bool:
@@ -132,6 +134,37 @@ def _index_urls_to_uv_env_vars(index_urls: list[str] | None = None) -> dict[str,
     return uv_index_env_vars
 
 
+def _execute_in_subprocess(cmd: list[str], cwd: str | None = None, env: dict[str, str] | None = None) -> None:
+    """
+    Execute a process and stream output to logger.
+
+    :param cmd: command and arguments to run
+    :param cwd: Current working directory passed to the Popen constructor
+    :param env: Additional environment variables to set for the subprocess.
+    """
+    log = logging.getLogger(__name__)
+
+    log.info("Executing cmd: %s", " ".join(shlex.quote(c) for c in cmd))
+    with subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        bufsize=0,
+        close_fds=False,
+        cwd=cwd,
+        env=env,
+    ) as proc:
+        log.info("Output:")
+        if proc.stdout:
+            with proc.stdout:
+                for line in iter(proc.stdout.readline, b""):
+                    log.info("%s", line.decode().rstrip())
+
+        exit_code = proc.wait()
+    if exit_code != 0:
+        raise subprocess.CalledProcessError(exit_code, cmd)
+
+
 def prepare_virtualenv(
     venv_directory: str,
     python_bin: str,
@@ -167,9 +200,10 @@ def prepare_virtualenv(
 
     if _use_uv():
         venv_cmd = _generate_uv_cmd(venv_directory, python_bin, system_site_packages)
+        _execute_in_subprocess(venv_cmd, env={**os.environ, **_index_urls_to_uv_env_vars(index_urls)})
     else:
         venv_cmd = _generate_venv_cmd(venv_directory, python_bin, system_site_packages)
-    execute_in_subprocess(venv_cmd)
+        _execute_in_subprocess(venv_cmd)
 
     pip_cmd = None
     if requirements is not None and len(requirements) != 0:
@@ -188,7 +222,7 @@ def prepare_virtualenv(
             )
 
     if pip_cmd:
-        execute_in_subprocess(pip_cmd, env={**os.environ, **_index_urls_to_uv_env_vars(index_urls)})
+        _execute_in_subprocess(pip_cmd, env={**os.environ, **_index_urls_to_uv_env_vars(index_urls)})
 
     return f"{venv_directory}/bin/python"
 

airflow/providers/standard/utils/python_virtualenv_script.jinja2

@@ -40,6 +40,23 @@ if sys.version_info >= (3,6):
         pass
 {% endif %}
 
+try:
+    from airflow.sdk.execution_time import task_runner
+except ModuleNotFoundError:
+    pass
+else:
+  {#-
+    We are in an Airflow 3.x environment, try and set up supervisor comms so
+    virtualenv can access Vars/Conn/XCom/etc that normal tasks can
+
+    We don't use the walrus operator (`:=`) below as it is possible people can
+    be using this on pre-3.8 versions of python, and while Airflow doesn't
+    support them, it's easy to not break it not using that operator here.
+  #}
+  reinit_supervisor_comms = getattr(task_runner, "reinit_supervisor_comms", None)
+  if reinit_supervisor_comms:
+      reinit_supervisor_comms()
+
 # Script
 {{ python_callable_source }}
 
@@ -49,12 +66,10 @@ if sys.version_info >= (3,6):
 import types
 
 {{ modified_dag_module_name }}  = types.ModuleType("{{ modified_dag_module_name }}")
-
 {{ modified_dag_module_name }}.{{ python_callable }} = {{ python_callable }}
-
 sys.modules["{{modified_dag_module_name}}"] = {{modified_dag_module_name}}
 
-{% endif%}
+{%- endif -%}
 
 {% if op_args or op_kwargs %}
 with open(sys.argv[1], "rb") as file:

airflow/providers/standard/utils/sensor_helper.py

@@ -16,6 +16,7 @@
 # under the License.
 from __future__ import annotations
 
+from collections.abc import Collection
 from typing import TYPE_CHECKING, Any, cast
 
 from sqlalchemy import func, select, tuple_
@@ -27,7 +28,7 @@ from airflow.utils.session import NEW_SESSION, provide_session
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
-    from sqlalchemy.sql import Executable
+    from sqlalchemy.sql import Select
 
 
 @provide_session
@@ -59,6 +60,7 @@ def _get_count(
             session.scalar(
                 _count_stmt(TI, states, dttm_filter, external_dag_id).where(TI.task_id.in_(external_task_ids))
             )
+            or 0
         ) / len(external_task_ids)
     elif external_task_group_id:
         external_task_group_task_ids = _get_external_task_group_task_ids(
@@ -68,20 +70,25 @@ def _get_count(
             count = 0
         else:
             count = (
-                session.scalar(
-                    _count_stmt(TI, states, dttm_filter, external_dag_id).where(
-                        tuple_(TI.task_id, TI.map_index).in_(external_task_group_task_ids)
+                (
+                    session.scalar(
+                        _count_stmt(TI, states, dttm_filter, external_dag_id).where(
+                            tuple_(TI.task_id, TI.map_index).in_(external_task_group_task_ids)
+                        )
                     )
+                    or 0
                 )
                 / len(external_task_group_task_ids)
                 * len(dttm_filter)
             )
     else:
-        count = session.scalar(_count_stmt(DR, states, dttm_filter, external_dag_id))
+        count = session.scalar(_count_stmt(DR, states, dttm_filter, external_dag_id)) or 0
     return cast("int", count)
 
 
-def _count_stmt(model, states, dttm_filter, external_dag_id) -> Executable:
+def _count_stmt(
+    model: type[DagRun] | type[TaskInstance], states: list[str], dttm_filter: list[Any], external_dag_id: str
+) -> Select[tuple[int]]:
     """
     Get the count of records against dttm filter and states.
 
@@ -97,7 +104,9 @@ def _count_stmt(model, states, dttm_filter, external_dag_id) -> Executable:
     )
 
 
-def _get_external_task_group_task_ids(dttm_filter, external_task_group_id, external_dag_id, session):
+def _get_external_task_group_task_ids(
+    dttm_filter: list[Any], external_task_group_id: str, external_dag_id: str, session: Session
+) -> list[tuple[str, int]]:
     """
     Get the count of records against dttm filter and states.
 
@@ -107,6 +116,8 @@ def _get_external_task_group_task_ids(dttm_filter, external_task_group_id, exter
     :param session: airflow session object
     """
     refreshed_dag_info = SerializedDagModel.get_dag(external_dag_id, session=session)
+    if not refreshed_dag_info:
+        return [(external_task_group_id, -1)]
     task_group = refreshed_dag_info.task_group_dict.get(external_task_group_id)
 
     if task_group:
@@ -129,7 +140,7 @@ def _get_external_task_group_task_ids(dttm_filter, external_task_group_id, exter
 
 def _get_count_by_matched_states(
     run_id_task_state_map: dict[str, dict[str, Any]],
-    states: list[str],
+    states: Collection[str],
 ):
     count = 0
     for _, task_states in run_id_task_state_map.items():

airflow/providers/standard/utils/skipmixin.py

@@ -21,7 +21,7 @@ from collections.abc import Iterable, Sequence
 from types import GeneratorType
 from typing import TYPE_CHECKING
 
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
@@ -63,7 +63,7 @@ class SkipMixin(LoggingMixin):
         """
         # Import is internal for backward compatibility when importing PythonOperator
         # from airflow.providers.common.compat.standard.operators
-        from airflow.exceptions import DownstreamTasksSkipped
+        from airflow.providers.common.compat.sdk import DownstreamTasksSkipped
 
         #  The following could be applied only for non-mapped tasks,
         #  as future mapped tasks have not been expanded yet. Such tasks