diff --git a/providers/smtp/src/airflow/providers/smtp/hooks/smtp.py b/providers/smtp/src/airflow/providers/smtp/hooks/smtp.py index d0bd9fb17bded..e50f4b97e56d1 100644 --- a/providers/smtp/src/airflow/providers/smtp/hooks/smtp.py +++ b/providers/smtp/src/airflow/providers/smtp/hooks/smtp.py @@ -41,13 +41,10 @@ import aiosmtplib from airflow.exceptions import AirflowException, AirflowNotFoundException -from airflow.providers.smtp.version_compat import BaseHook +from airflow.providers.common.compat.sdk import BaseHook if TYPE_CHECKING: - try: - from airflow.sdk import Connection - except (ImportError, ModuleNotFoundError): - from airflow.models.connection import Connection # type: ignore[assignment] + from airflow.providers.common.compat.sdk import Connection def build_xoauth2_string(username: str, token: str) -> str: diff --git a/providers/smtp/src/airflow/providers/smtp/notifications/smtp.py b/providers/smtp/src/airflow/providers/smtp/notifications/smtp.py index 731743f0c5ff2..7933f275408cf 100644 --- a/providers/smtp/src/airflow/providers/smtp/notifications/smtp.py +++ b/providers/smtp/src/airflow/providers/smtp/notifications/smtp.py @@ -27,7 +27,7 @@ from airflow.providers.smtp.version_compat import AIRFLOW_V_3_1_PLUS if TYPE_CHECKING: - from airflow.sdk import Context + from airflow.providers.common.compat.sdk import Context class SmtpNotifier(BaseNotifier): diff --git a/providers/smtp/src/airflow/providers/smtp/operators/smtp.py b/providers/smtp/src/airflow/providers/smtp/operators/smtp.py index c35da707a5a69..9621def333f38 100644 --- a/providers/smtp/src/airflow/providers/smtp/operators/smtp.py +++ b/providers/smtp/src/airflow/providers/smtp/operators/smtp.py @@ -22,11 +22,11 @@ from typing import TYPE_CHECKING, Any from airflow.exceptions import AirflowException +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.smtp.hooks.smtp import SmtpHook -from airflow.providers.smtp.version_compat import BaseOperator if TYPE_CHECKING: - from airflow.providers.smtp.version_compat import Context + from airflow.providers.common.compat.sdk import Context class EmailOperator(BaseOperator): diff --git a/providers/smtp/src/airflow/providers/smtp/version_compat.py b/providers/smtp/src/airflow/providers/smtp/version_compat.py index 5c0789f3a9977..a4fae50317e81 100644 --- a/providers/smtp/src/airflow/providers/smtp/version_compat.py +++ b/providers/smtp/src/airflow/providers/smtp/version_compat.py @@ -35,22 +35,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_0_PLUS: - from airflow.sdk import BaseOperator - from airflow.sdk.definitions.context import Context -else: - from airflow.models import BaseOperator - from airflow.utils.context import Context - -if AIRFLOW_V_3_1_PLUS: - from airflow.sdk import BaseHook -else: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] - __all__ = [ "AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", - "BaseHook", - "BaseOperator", - "Context", ]