diff --git a/providers/mysql/pyproject.toml b/providers/mysql/pyproject.toml index 05ca23bf29b8c..97ac56046d43e 100644 --- a/providers/mysql/pyproject.toml +++ b/providers/mysql/pyproject.toml @@ -58,6 +58,7 @@ requires-python = ">=3.10" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.10.0", + "apache-airflow-providers-common-compat>=1.8.0", "apache-airflow-providers-common-sql>=1.20.0", # The mysqlclient package creates friction when installing on MacOS as it needs pkg-config to # Install and compile, and it's really only used by MySQL provider, so we can skip it on MacOS @@ -95,6 +96,7 @@ dev = [ "apache-airflow-task-sdk", "apache-airflow-devel-common", "apache-airflow-providers-amazon", + "apache-airflow-providers-common-compat", "apache-airflow-providers-common-sql", "apache-airflow-providers-openlineage", "apache-airflow-providers-presto", diff --git a/providers/mysql/src/airflow/providers/mysql/hooks/mysql.py b/providers/mysql/src/airflow/providers/mysql/hooks/mysql.py index affa6edd7471e..b7500bb8b8cef 100644 --- a/providers/mysql/src/airflow/providers/mysql/hooks/mysql.py +++ b/providers/mysql/src/airflow/providers/mysql/hooks/mysql.py @@ -30,12 +30,7 @@ logger = logging.getLogger(__name__) if TYPE_CHECKING: - from airflow.providers.mysql.version_compat import AIRFLOW_V_3_0_PLUS - - if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import Connection - else: - from airflow.models.connection import Connection # type: ignore[assignment] + from airflow.providers.common.compat.sdk import Connection try: from mysql.connector.abstracts import MySQLConnectionAbstract diff --git a/providers/mysql/src/airflow/providers/mysql/transfers/presto_to_mysql.py b/providers/mysql/src/airflow/providers/mysql/transfers/presto_to_mysql.py index 98d25c3c0bf89..868c1c11b6d13 100644 --- a/providers/mysql/src/airflow/providers/mysql/transfers/presto_to_mysql.py +++ b/providers/mysql/src/airflow/providers/mysql/transfers/presto_to_mysql.py @@ -20,16 +20,12 @@ from collections.abc import Sequence from typing import TYPE_CHECKING +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.mysql.hooks.mysql import MySqlHook -from airflow.providers.mysql.version_compat import BaseOperator from airflow.providers.presto.hooks.presto import PrestoHook 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.common.compat.sdk import Context class PrestoToMySqlOperator(BaseOperator): diff --git a/providers/mysql/src/airflow/providers/mysql/transfers/s3_to_mysql.py b/providers/mysql/src/airflow/providers/mysql/transfers/s3_to_mysql.py index 474433c2894d6..09633c3ea3ca3 100644 --- a/providers/mysql/src/airflow/providers/mysql/transfers/s3_to_mysql.py +++ b/providers/mysql/src/airflow/providers/mysql/transfers/s3_to_mysql.py @@ -21,15 +21,11 @@ from typing import TYPE_CHECKING from airflow.providers.amazon.aws.hooks.s3 import S3Hook +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.mysql.hooks.mysql import MySqlHook -from airflow.providers.mysql.version_compat import BaseOperator 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.common.compat.sdk import Context class S3ToMySqlOperator(BaseOperator): diff --git a/providers/mysql/src/airflow/providers/mysql/transfers/trino_to_mysql.py b/providers/mysql/src/airflow/providers/mysql/transfers/trino_to_mysql.py index 439b2d1de30e7..6308c3e7f66ab 100644 --- a/providers/mysql/src/airflow/providers/mysql/transfers/trino_to_mysql.py +++ b/providers/mysql/src/airflow/providers/mysql/transfers/trino_to_mysql.py @@ -20,16 +20,12 @@ from collections.abc import Sequence from typing import TYPE_CHECKING +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.mysql.hooks.mysql import MySqlHook -from airflow.providers.mysql.version_compat import BaseOperator from airflow.providers.trino.hooks.trino import TrinoHook 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.common.compat.sdk import Context class TrinoToMySqlOperator(BaseOperator): diff --git a/providers/mysql/src/airflow/providers/mysql/transfers/vertica_to_mysql.py b/providers/mysql/src/airflow/providers/mysql/transfers/vertica_to_mysql.py index 0dc911465983b..4250a3344271a 100644 --- a/providers/mysql/src/airflow/providers/mysql/transfers/vertica_to_mysql.py +++ b/providers/mysql/src/airflow/providers/mysql/transfers/vertica_to_mysql.py @@ -33,16 +33,12 @@ "installed in case you see compilation error during installation." ) +from airflow.providers.common.compat.sdk import BaseOperator from airflow.providers.mysql.hooks.mysql import MySqlHook -from airflow.providers.mysql.version_compat import BaseOperator from airflow.providers.vertica.hooks.vertica import VerticaHook 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.common.compat.sdk import Context class VerticaToMySqlOperator(BaseOperator): diff --git a/providers/mysql/src/airflow/providers/mysql/version_compat.py b/providers/mysql/src/airflow/providers/mysql/version_compat.py index 78cd58dd016ab..ff0f446dc5efc 100644 --- a/providers/mysql/src/airflow/providers/mysql/version_compat.py +++ b/providers/mysql/src/airflow/providers/mysql/version_compat.py @@ -29,12 +29,6 @@ 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 BaseOperator -else: - from airflow.models import BaseOperator - __all__ = [ "AIRFLOW_V_3_0_PLUS", - "BaseOperator", ] diff --git a/providers/mysql/tests/unit/mysql/hooks/test_mysql.py b/providers/mysql/tests/unit/mysql/hooks/test_mysql.py index ea011f3decdca..d72f560103731 100644 --- a/providers/mysql/tests/unit/mysql/hooks/test_mysql.py +++ b/providers/mysql/tests/unit/mysql/hooks/test_mysql.py @@ -25,8 +25,8 @@ import pytest import sqlalchemy -from airflow.models import Connection from airflow.models.dag import DAG +from airflow.providers.common.compat.sdk import Connection try: import MySQLdb.cursors @@ -52,6 +52,7 @@ class TestMySqlHookConn: def setup_method(self): self.connection = Connection( + conn_id="test_conn_id", conn_type="mysql", login="login", password="password", @@ -194,7 +195,14 @@ def test_get_uri(self, mock_connect, connection_params, expected_uri): @mock.patch("MySQLdb.connect") def test_get_conn_from_connection(self, mock_connect): - conn = Connection(login="login-conn", password="password-conn", host="host", schema="schema") + conn = Connection( + conn_id="test_conn_id", + conn_type="mysql", + login="login-conn", + password="password-conn", + host="host", + schema="schema", + ) hook = MySqlHook(connection=conn) hook.get_conn() mock_connect.assert_called_once_with( @@ -203,7 +211,14 @@ def test_get_conn_from_connection(self, mock_connect): @mock.patch("MySQLdb.connect") def test_get_conn_from_connection_with_schema(self, mock_connect): - conn = Connection(login="login-conn", password="password-conn", host="host", schema="schema") + conn = Connection( + conn_id="test_conn_id", + conn_type="mysql", + login="login-conn", + password="password-conn", + host="host", + schema="schema", + ) hook = MySqlHook(connection=conn, schema="schema-override") hook.get_conn() mock_connect.assert_called_once_with( diff --git a/providers/mysql/tests/unit/mysql/hooks/test_mysql_connector_python.py b/providers/mysql/tests/unit/mysql/hooks/test_mysql_connector_python.py index 5d2f0d1613b35..bbee9a69b1650 100644 --- a/providers/mysql/tests/unit/mysql/hooks/test_mysql_connector_python.py +++ b/providers/mysql/tests/unit/mysql/hooks/test_mysql_connector_python.py @@ -20,13 +20,15 @@ import json from unittest import mock -from airflow.models import Connection +from airflow.providers.common.compat.sdk import Connection from airflow.providers.mysql.hooks.mysql import MySqlHook class TestMySqlHookConnMySqlConnectorPython: def setup_method(self): self.connection = Connection( + conn_id="test_conn_id", + conn_type="mysql", login="login", password="password", host="host",