fix: better handling for when c extension isn't available for mysql connector/python instead of erroring out

Author: karenc-bq

aws_advanced_python_wrapper/mysql_driver_dialect.py

@@ -23,9 +23,20 @@
 from concurrent.futures import Executor, ThreadPoolExecutor, TimeoutError
 from inspect import signature
 
-from mysql.connector import CMySQLConnection, MySQLConnection
+CMYSQL_ENABLED = False
+
+from mysql.connector import MySQLConnection
 from mysql.connector.cursor import MySQLCursor
-from mysql.connector.cursor_cext import CMySQLCursor
+
+try:
+    from mysql.connector import CMySQLConnection
+    from mysql.connector.cursor_cext import CMySQLCursor
+
+    CMYSQL_ENABLED = True
+
+except ImportError as exc:
+    # Do nothing
+    pass
 
 from aws_advanced_python_wrapper.driver_dialect import DriverDialect
 from aws_advanced_python_wrapper.driver_dialect_codes import DriverDialectCodes
@@ -62,13 +73,21 @@ class MySQLDriverDialect(DriverDialect):
         "Cursor.fetchall"
     }
 
+    @staticmethod
+    def _is_mysql_connection(conn: Connection | object) -> bool:
+        return isinstance(conn, MySQLConnection) or (CMYSQL_ENABLED and isinstance(conn, CMySQLConnection))
+
+    @staticmethod
+    def _is_cmysql_cursor(obj: object) -> bool:
+        return CMYSQL_ENABLED and isinstance(obj, CMySQLCursor)
+
     def is_dialect(self, connect_func: Callable) -> bool:
         if MySQLDriverDialect.TARGET_DRIVER_CODE not in str(signature(connect_func)):
             return MySQLDriverDialect.TARGET_DRIVER_CODE.lower() in (connect_func.__module__ + connect_func.__qualname__).lower()
         return True
 
     def is_closed(self, conn: Connection) -> bool:
-        if isinstance(conn, CMySQLConnection) or isinstance(conn, MySQLConnection):
+        if MySQLDriverDialect._is_mysql_connection(conn):
 
             # is_connected validates the connection using a ping().
             # If there are any unread results from previous executions an error will be thrown.
@@ -86,14 +105,14 @@ def is_closed(self, conn: Connection) -> bool:
         raise UnsupportedOperationError(Messages.get_formatted("DriverDialect.UnsupportedOperationError", self._driver_name, "is_connected"))
 
     def get_autocommit(self, conn: Connection) -> bool:
-        if isinstance(conn, CMySQLConnection) or isinstance(conn, MySQLConnection):
+        if MySQLDriverDialect._is_mysql_connection(conn):
             return conn.autocommit
 
         raise UnsupportedOperationError(
             Messages.get_formatted("DriverDialect.UnsupportedOperationError", self._driver_name, "autocommit"))
 
     def set_autocommit(self, conn: Connection, autocommit: bool):
-        if isinstance(conn, CMySQLConnection) or isinstance(conn, MySQLConnection):
+        if MySQLDriverDialect._is_mysql_connection(conn):
             conn.autocommit = autocommit
             return
 
@@ -112,24 +131,24 @@ def abort_connection(self, conn: Connection):
                 "abort_connection"))
 
     def can_execute_query(self, conn: Connection) -> bool:
-        if isinstance(conn, CMySQLConnection) or isinstance(conn, MySQLConnection):
+        if MySQLDriverDialect._is_mysql_connection(conn):
             if conn.unread_result:
                 return conn.can_consume_results
         return True
 
     def is_in_transaction(self, conn: Connection) -> bool:
-        if isinstance(conn, CMySQLConnection) or isinstance(conn, MySQLConnection):
+        if MySQLDriverDialect._is_mysql_connection(conn):
             return bool(conn.in_transaction)
 
         raise UnsupportedOperationError(
             Messages.get_formatted("DriverDialect.UnsupportedOperationError", self._driver_name,
                                    "in_transaction"))
 
     def get_connection_from_obj(self, obj: object) -> Any:
-        if isinstance(obj, CMySQLConnection) or isinstance(obj, MySQLConnection):
+        if MySQLDriverDialect._is_mysql_connection(obj):
             return obj
 
-        if isinstance(obj, CMySQLCursor):
+        if MySQLDriverDialect._is_cmysql_cursor(obj):
             try:
                 conn = None
 
@@ -140,15 +159,15 @@ def get_connection_from_obj(self, obj: object) -> Any:
                 if conn is None:
                     return None
 
-                if isinstance(conn, CMySQLConnection) or isinstance(conn, MySQLConnection):
+                if MySQLDriverDialect._is_mysql_connection(conn):
                     return conn
 
             except ReferenceError:
                 return None
 
         if isinstance(obj, MySQLCursor):
             try:
-                if isinstance(obj._connection, CMySQLConnection) or isinstance(obj._connection, MySQLConnection):
+                if MySQLDriverDialect._is_mysql_connection(obj._connection):
                     return obj._connection
             except ReferenceError:
                 return None