Migrate MYSQL provider to common.compat

Author: kirubankamaraj

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",

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

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):

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):

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):

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):

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",
 ]

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(

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",