From 64ce50201808bf589d83341404cdb9c99620013c Mon Sep 17 00:00:00 2001 From: kevinhongzl Date: Tue, 1 Jul 2025 00:36:25 +0800 Subject: [PATCH 1/2] Provider Migration: Replace BaseSensorOperator to Task SDK for datadog --- .../airflow/providers/datadog/sensors/datadog.py | 13 ++----------- .../src/airflow/providers/datadog/version_compat.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/providers/datadog/src/airflow/providers/datadog/sensors/datadog.py b/providers/datadog/src/airflow/providers/datadog/sensors/datadog.py index 251020febb871..d39f20352e665 100644 --- a/providers/datadog/src/airflow/providers/datadog/sensors/datadog.py +++ b/providers/datadog/src/airflow/providers/datadog/sensors/datadog.py @@ -24,19 +24,10 @@ from airflow.exceptions import AirflowException from airflow.providers.datadog.hooks.datadog import DatadogHook -from airflow.providers.datadog.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseSensorOperator -else: - from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] +from airflow.providers.datadog.version_compat 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 + from airflow.providers.datadog.version_compat import Context class DatadogSensor(BaseSensorOperator): diff --git a/providers/datadog/src/airflow/providers/datadog/version_compat.py b/providers/datadog/src/airflow/providers/datadog/version_compat.py index 48d122b669696..1fa00e51ec6cd 100644 --- a/providers/datadog/src/airflow/providers/datadog/version_compat.py +++ b/providers/datadog/src/airflow/providers/datadog/version_compat.py @@ -33,3 +33,16 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) + +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import BaseSensorOperator + from airflow.sdk.definitions.context import Context +else: + from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] + from airflow.utils.context import Context + +__all__ = [ + "AIRFLOW_V_3_0_PLUS", + "BaseSensorOperator", + "Context", +] From 87d2e0f21c2961a1360d862edd596ab2b0e05d1b Mon Sep 17 00:00:00 2001 From: Kaxil Naik Date: Tue, 1 Jul 2025 00:12:31 +0530 Subject: [PATCH 2/2] Apply suggestion from @kaxil --- .../datadog/src/airflow/providers/datadog/version_compat.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/providers/datadog/src/airflow/providers/datadog/version_compat.py b/providers/datadog/src/airflow/providers/datadog/version_compat.py index 1fa00e51ec6cd..52dee9e5be6f4 100644 --- a/providers/datadog/src/airflow/providers/datadog/version_compat.py +++ b/providers/datadog/src/airflow/providers/datadog/version_compat.py @@ -39,7 +39,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: from airflow.sdk.definitions.context import Context else: from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] - from airflow.utils.context import Context + from airflow.utils.context import Context # type: ignore[no-redef] __all__ = [ "AIRFLOW_V_3_0_PLUS",