From c4d0803035baf64895dde64da73c75ac5c191d38 Mon Sep 17 00:00:00 2001 From: Bhavani Ravi Date: Thu, 23 Oct 2025 05:07:25 +0530 Subject: [PATCH 1/2] Migrate yandex provider to common.compat --- providers/yandex/pyproject.toml | 3 +++ .../airflow/providers/yandex/hooks/yandex.py | 2 +- .../src/airflow/providers/yandex/links/yq.py | 4 ++-- .../providers/yandex/operators/dataproc.py | 4 ++-- .../airflow/providers/yandex/operators/yq.py | 4 ++-- .../providers/yandex/version_compat.py | 21 ------------------- 6 files changed, 10 insertions(+), 28 deletions(-) diff --git a/providers/yandex/pyproject.toml b/providers/yandex/pyproject.toml index 34bfbabe79e69..76d70aac341b8 100644 --- a/providers/yandex/pyproject.toml +++ b/providers/yandex/pyproject.toml @@ -59,6 +59,7 @@ dependencies = [ "apache-airflow>=2.10.0", "yandexcloud>=0.308.0; python_version < '3.13'", "yandex-query-client>=0.1.4; python_version < '3.13'", + "apache-airflow-providers-common-compat>=1.8.0", ] # The optional dependencies should be modified in place in the generated file @@ -73,8 +74,10 @@ dev = [ "apache-airflow", "apache-airflow-task-sdk", "apache-airflow-devel-common", + "apache-airflow-providers-common-compat", # Additional devel dependencies (do not remove this line and add extra development dependencies) "responses>=0.25.0", + ] # To build docs: diff --git a/providers/yandex/src/airflow/providers/yandex/hooks/yandex.py b/providers/yandex/src/airflow/providers/yandex/hooks/yandex.py index 78c89368adf56..86f20c8d5ea8c 100644 --- a/providers/yandex/src/airflow/providers/yandex/hooks/yandex.py +++ b/providers/yandex/src/airflow/providers/yandex/hooks/yandex.py @@ -20,6 +20,7 @@ import yandexcloud +from airflow.providers.common.compat.sdk import BaseHook from airflow.providers.yandex.utils.credentials import ( CredentialsType, get_credentials, @@ -28,7 +29,6 @@ from airflow.providers.yandex.utils.defaults import conn_name_attr, conn_type, default_conn_name, hook_name from airflow.providers.yandex.utils.fields import get_field_from_extras from airflow.providers.yandex.utils.user_agent import provider_user_agent -from airflow.providers.yandex.version_compat import BaseHook class YandexCloudBaseHook(BaseHook): diff --git a/providers/yandex/src/airflow/providers/yandex/links/yq.py b/providers/yandex/src/airflow/providers/yandex/links/yq.py index 7bd416e889831..c9f5195c1d34d 100644 --- a/providers/yandex/src/airflow/providers/yandex/links/yq.py +++ b/providers/yandex/src/airflow/providers/yandex/links/yq.py @@ -18,11 +18,11 @@ from typing import TYPE_CHECKING -from airflow.providers.yandex.version_compat import BaseOperatorLink, XCom +from airflow.providers.common.compat.sdk import BaseOperatorLink, XCom if TYPE_CHECKING: from airflow.models.taskinstancekey import TaskInstanceKey - from airflow.providers.yandex.version_compat import BaseOperator, Context + from airflow.providers.common.compat.sdk import BaseOperator, Context XCOM_WEBLINK_KEY = "web_link" diff --git a/providers/yandex/src/airflow/providers/yandex/operators/dataproc.py b/providers/yandex/src/airflow/providers/yandex/operators/dataproc.py index d55569c8a23bc..e9375b3baaa2e 100644 --- a/providers/yandex/src/airflow/providers/yandex/operators/dataproc.py +++ b/providers/yandex/src/airflow/providers/yandex/operators/dataproc.py @@ -23,11 +23,11 @@ import yandexcloud +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.yandex.hooks.dataproc import DataprocHook -from airflow.providers.yandex.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.providers.yandex.version_compat import Context + from airflow.providers.common.compat.sdk import Context @dataclass diff --git a/providers/yandex/src/airflow/providers/yandex/operators/yq.py b/providers/yandex/src/airflow/providers/yandex/operators/yq.py index 1c450174e1681..8c294133a3917 100644 --- a/providers/yandex/src/airflow/providers/yandex/operators/yq.py +++ b/providers/yandex/src/airflow/providers/yandex/operators/yq.py @@ -20,12 +20,12 @@ from functools import cached_property from typing import TYPE_CHECKING, Any +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.yandex.hooks.yq import YQHook from airflow.providers.yandex.links.yq import YQLink -from airflow.providers.yandex.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.providers.yandex.version_compat import Context + from airflow.providers.common.compat.sdk import Context class YQExecuteQueryOperator(BaseOperator): diff --git a/providers/yandex/src/airflow/providers/yandex/version_compat.py b/providers/yandex/src/airflow/providers/yandex/version_compat.py index d315b0d25897b..fe3ccfcfd3d7a 100644 --- a/providers/yandex/src/airflow/providers/yandex/version_compat.py +++ b/providers/yandex/src/airflow/providers/yandex/version_compat.py @@ -30,28 +30,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) - -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperator, BaseOperatorLink - from airflow.sdk.definitions.context import Context - from airflow.sdk.execution_time.xcom import XCom -else: - from airflow.models import BaseOperator, XCom - from airflow.models.baseoperatorlink import BaseOperatorLink # type: ignore[no-redef] - from airflow.utils.context import Context - - __all__ = [ "AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", - "BaseHook", - "BaseOperator", - "BaseOperatorLink", - "Context", - "XCom", ] From 03b1b9bca845708089910288433e74cb9b877e6f Mon Sep 17 00:00:00 2001 From: Bhavani Ravi Date: Thu, 23 Oct 2025 13:13:19 +0530 Subject: [PATCH 2/2] update testcases --- .../yandex/tests/unit/yandex/hooks/test_dataproc.py | 11 ++--------- .../yandex/tests/unit/yandex/hooks/test_yandex.py | 10 +--------- providers/yandex/tests/unit/yandex/hooks/test_yq.py | 9 +-------- providers/yandex/tests/unit/yandex/links/test_yq.py | 6 +----- .../tests/unit/yandex/operators/test_dataproc.py | 10 +--------- .../yandex/tests/unit/yandex/operators/test_yq.py | 10 +--------- 6 files changed, 7 insertions(+), 49 deletions(-) diff --git a/providers/yandex/tests/unit/yandex/hooks/test_dataproc.py b/providers/yandex/tests/unit/yandex/hooks/test_dataproc.py index adc81828ebb32..5b5a96bd8337a 100644 --- a/providers/yandex/tests/unit/yandex/hooks/test_dataproc.py +++ b/providers/yandex/tests/unit/yandex/hooks/test_dataproc.py @@ -23,19 +23,12 @@ pytest.importorskip("yandexcloud") -try: - import importlib.util - - if not importlib.util.find_spec("airflow.sdk.bases.hook"): - raise ImportError - - BASEHOOK_PATCH_PATH = "airflow.sdk.bases.hook.BaseHook" -except ImportError: - BASEHOOK_PATCH_PATH = "airflow.hooks.base.BaseHook" from airflow.models import Connection from airflow.providers.yandex.hooks.dataproc import DataprocHook +BASEHOOK_PATCH_PATH = "airflow.providers.common.compat.sdk.BaseHook" + # Airflow connection with type "yandexcloud" must be created CONNECTION_ID = "yandexcloud_default" diff --git a/providers/yandex/tests/unit/yandex/hooks/test_yandex.py b/providers/yandex/tests/unit/yandex/hooks/test_yandex.py index c0fe981669493..545a8bdcb07f1 100644 --- a/providers/yandex/tests/unit/yandex/hooks/test_yandex.py +++ b/providers/yandex/tests/unit/yandex/hooks/test_yandex.py @@ -27,15 +27,7 @@ from tests_common.test_utils.config import conf_vars -try: - import importlib.util - - if not importlib.util.find_spec("airflow.sdk.bases.hook"): - raise ImportError - - BASEHOOK_PATCH_PATH = "airflow.sdk.bases.hook.BaseHook" -except ImportError: - BASEHOOK_PATCH_PATH = "airflow.hooks.base.BaseHook" +BASEHOOK_PATCH_PATH = "airflow.providers.common.compat.sdk.BaseHook" class TestYandexHook: diff --git a/providers/yandex/tests/unit/yandex/hooks/test_yq.py b/providers/yandex/tests/unit/yandex/hooks/test_yq.py index e523c173f81e2..942511513bc48 100644 --- a/providers/yandex/tests/unit/yandex/hooks/test_yq.py +++ b/providers/yandex/tests/unit/yandex/hooks/test_yq.py @@ -26,15 +26,8 @@ from airflow.models import Connection -try: - import importlib.util +BASEHOOK_PATCH_PATH = "airflow.providers.common.compat.sdk.BaseHook" - if not importlib.util.find_spec("airflow.sdk.bases.hook"): - raise ImportError - - BASEHOOK_PATCH_PATH = "airflow.sdk.bases.hook.BaseHook" -except ImportError: - BASEHOOK_PATCH_PATH = "airflow.hooks.base.BaseHook" yandexcloud = pytest.importorskip("yandexcloud") from airflow.providers.yandex.hooks.yq import YQHook diff --git a/providers/yandex/tests/unit/yandex/links/test_yq.py b/providers/yandex/tests/unit/yandex/links/test_yq.py index 64baa18ad89b0..c742f28da7bb0 100644 --- a/providers/yandex/tests/unit/yandex/links/test_yq.py +++ b/providers/yandex/tests/unit/yandex/links/test_yq.py @@ -21,16 +21,12 @@ import pytest from airflow.models.taskinstance import TaskInstance +from airflow.providers.common.compat.sdk import XCom from airflow.providers.yandex.links.yq import YQLink from tests_common.test_utils.mock_operators import MockOperator from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk.execution_time.xcom import XCom -else: - from airflow.models import XCom - yandexcloud = pytest.importorskip("yandexcloud") diff --git a/providers/yandex/tests/unit/yandex/operators/test_dataproc.py b/providers/yandex/tests/unit/yandex/operators/test_dataproc.py index 66cb883909e44..7c4f1f39d0854 100644 --- a/providers/yandex/tests/unit/yandex/operators/test_dataproc.py +++ b/providers/yandex/tests/unit/yandex/operators/test_dataproc.py @@ -66,15 +66,7 @@ # https://cloud.yandex.com/docs/logging/concepts/log-group LOG_GROUP_ID = "my_log_group_id" -try: - import importlib.util - - if not importlib.util.find_spec("airflow.sdk.bases.hook"): - raise ImportError - - BASEHOOK_PATCH_PATH = "airflow.sdk.bases.hook.BaseHook" -except ImportError: - BASEHOOK_PATCH_PATH = "airflow.hooks.base.BaseHook" +BASEHOOK_PATCH_PATH = "airflow.providers.common.compat.sdk.BaseHook" class TestDataprocClusterCreateOperator: diff --git a/providers/yandex/tests/unit/yandex/operators/test_yq.py b/providers/yandex/tests/unit/yandex/operators/test_yq.py index 3fa707f9c20a3..18a931018cb00 100644 --- a/providers/yandex/tests/unit/yandex/operators/test_yq.py +++ b/providers/yandex/tests/unit/yandex/operators/test_yq.py @@ -33,15 +33,7 @@ from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS -try: - import importlib.util - - if not importlib.util.find_spec("airflow.sdk.bases.hook"): - raise ImportError - - BASEHOOK_PATCH_PATH = "airflow.sdk.bases.hook.BaseHook" -except ImportError: - BASEHOOK_PATCH_PATH = "airflow.hooks.base.BaseHook" +BASEHOOK_PATCH_PATH = "airflow.providers.common.compat.sdk.BaseHook" OAUTH_TOKEN = "my_oauth_token" FOLDER_ID = "my_folder_id"