Add SQLAlchemy ORM plugin tests

aws_advanced_python_wrapper/sqlalchemy/mysql_orm_dialect.py

@@ -12,12 +12,27 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-# aws_advanced_python_wrapper/sqlalchemy/sqlalchemy_mysqlconnector_dialect.py
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import mysql.connector
+from mysql.connector import CMySQLConnection
+from mysql.connector.errors import Error
 from sqlalchemy.dialects.mysql.mysqlconnector import \
     MySQLDialect_mysqlconnector
+from sqlalchemy.engine import default
+
+from aws_advanced_python_wrapper import AwsWrapperConnection
+from aws_advanced_python_wrapper.errors import AwsWrapperError
+
+if TYPE_CHECKING:
+    from sqlalchemy import Connection
 
 
 class SqlAlchemyOrmMysqlDialect(MySQLDialect_mysqlconnector):
+    supports_statement_cache = True
+
     """
     SQLAlchemy dialect for AWS Advanced Python Wrapper with mysqlconnector. Extends the SQLAlchemy MySQL mysqlconnector dialect.
     This dialect is not related to the DriverDialect or DatabaseDialect classes used by our driver. Instead, it is used
@@ -27,3 +42,143 @@ class SqlAlchemyOrmMysqlDialect(MySQLDialect_mysqlconnector):
 
     name = 'mysql'
     driver = 'aws_wrapper_mysqlconnector'
+
+    @classmethod
+    def import_dbapi(cls):
+        """
+        Return the DB-API 2.0 module.
+        SQLAlchemy calls this to get the driver module.
+        """
+        import aws_advanced_python_wrapper
+        return aws_advanced_python_wrapper
+
+    def create_connect_args(self, url):
+        """
+        Transform SQLAlchemy URL into connection arguments.
+        Must include the 'target' parameter for our wrapper driver.
+        """
+        # Extract standard connection parameters
+        opts = url.translate_connect_args(username='user')
+
+        # Add query string parameters
+        opts.update(url.query)
+
+        # Add the required 'target' parameter for our wrapper
+        if 'target' not in opts:
+            opts['target'] = mysql.connector.Connect
+        if 'wrapper_plugins' not in opts:
+            opts['plugins'] = "aurora_connection_tracker,failover_v2"
+        else:
+            opts['plugins'] = opts['wrapper_plugins']
+            opts.pop('wrapper_plugins', None)
+        if 'connect_timeout' in opts:
+            opts['connect_timeout'] = int(opts['connect_timeout'])
+
+        # Return empty args list and kwargs dict
+        return [], opts
+
+    def _detect_charset(self, connection: Connection) -> str:
+        if isinstance(connection, CMySQLConnection):
+            return connection.charset
+        else:
+            raise Exception("Could not detect charset because connection was not a CMySQLConnection.")
+
+    def _extract_error_code(self, exception: BaseException) -> int:
+        if isinstance(exception, AwsWrapperError):
+            err = exception.driver_error
+            if err and isinstance(err, Error):
+                return err.errno
+            else:
+                raise Exception("Could not extract error code because driver_error was not a BaseException.")
+        else:
+            raise Exception("Could not extract error code because exception was not an AwsWrapperError.")
+
+    def initialize(self, connection):
+        """
+        Override initialization to handle type introspection.
+        The parent class tries to use TypeInfo.fetch() which requires
+        a native SQLAlchemy connection, not AwsWrapperConnection.
+        """
+
+        # Unwrap SQLAlchemy's connection object
+        wrapper_conn, wrapper_parent = self._get_wrapper_connection_and_parent(connection)
+
+        # this is driver-based, does not need server version info
+        # and is fairly critical for even basic SQL operations
+        self._connection_charset: Optional[str] = self._detect_charset(
+            wrapper_conn.target_connection
+        )
+
+        # call super().initialize() because we need to have
+        # server_version_info set up.  in 1.4 under python 2 only this does the
+        # "check unicode returns" thing, which is the one area that some
+        # SQL gets compiled within initialize() currently
+        default.DefaultDialect.initialize(self, connection)
+
+        self._detect_sql_mode(connection)
+        self._detect_ansiquotes(connection)  # depends on sql mode
+        self._detect_casing(connection)
+        if self._server_ansiquotes:
+            # if ansiquotes == True, build a new IdentifierPreparer
+            # with the new setting
+            self.identifier_preparer = self.preparer(
+                self, server_ansiquotes=self._server_ansiquotes
+            )
+
+        self.supports_sequences = (
+                self.is_mariadb and self.server_version_info >= (10, 3)
+        )
+
+        self.supports_for_update_of = (
+                self._is_mysql and self.server_version_info >= (8,)
+        )
+
+        self.use_mysql_for_share = (
+                self._is_mysql and self.server_version_info >= (8, 0, 1)
+        )
+
+        self._needs_correct_for_88718_96365 = (
+                not self.is_mariadb and self.server_version_info >= (8,)
+        )
+
+        self.delete_returning = (
+                self.is_mariadb and self.server_version_info >= (10, 0, 5)
+        )
+
+        self.insert_returning = (
+                self.is_mariadb and self.server_version_info >= (10, 5)
+        )
+
+        self._requires_alias_for_on_duplicate_key = (
+                self._is_mysql and self.server_version_info >= (8, 0, 20)
+        )
+
+        self._warn_for_known_db_issues()
+
+    def _get_wrapper_connection_and_parent(self, connection):
+        """
+        Traverse the connection chain to find AwsWrapperConnection and its parent connection.
+
+        Args:
+            connection: SQLAlchemy Connection object
+
+        Returns:
+            AwsWrapperConnection instance or None, parent connection of AwsWrapperConnection or None
+        """
+        # Start with the DBAPI connection
+        parent = connection
+        child = connection.connection
+
+        # Traverse up to 5 levels deep (reasonable limit)
+        for _ in range(5):
+            if isinstance(child, AwsWrapperConnection):
+                return child, parent
+
+            # Try to go deeper if there's a .connection attribute
+            if hasattr(child, 'connection'):
+                parent = child
+                child = child.connection
+            else:
+                break
+
+        return None

tests/integration/container/sqlalchemy/test_sqlalchemy_basic.py

@@ -29,7 +29,6 @@
                             subqueryload)
 from sqlalchemy.sql import func
 
-from tests.integration.container.utils.rds_test_utility import RdsTestUtility
 from ..utils.conditions import (disable_on_features, enable_on_deployments,
                                 enable_on_engines)
 from ..utils.database_engine import DatabaseEngine
@@ -114,41 +113,29 @@ class Book(Base):
                       TestEnvironmentFeatures.BLUE_GREEN_DEPLOYMENT,
                       TestEnvironmentFeatures.PERFORMANCE])
 class TestSqlAlchemy:
-    @pytest.fixture(scope='class')
-    def rds_utils(self):
-        region: str = TestEnvironment.get_current().get_info().get_region()
-        return RdsTestUtility(region)
-
-
-    @pytest.fixture(scope="class")
+    @pytest.fixture(scope="function")
     def engine(self, conn_utils):
         conn_str = f'mysql+aws_wrapper_mysqlconnector://{conn_utils.user}:{conn_utils.password}@{conn_utils.writer_cluster_host}:{conn_utils.port}/{conn_utils.dbname}'
         engine = create_engine(conn_str)
         Base.metadata.create_all(engine)
         yield engine
         Base.metadata.drop_all(engine)
+        engine.dispose()
 
-    @pytest.fixture(scope="class")
-    def Session(self, engine):
-        Session = sessionmaker(bind=engine)
-        yield Session
-
-    @pytest.fixture(scope="class")
-    def session(self, Session):
-        session = Session()
+    @pytest.fixture(scope="function")
+    def session(self, engine):
+        session = sessionmaker(bind=engine)()
         yield session
         session.rollback()
         session.close()
 
-    def test_sqlalchemy_backend_configuration(self, test_environment: TestEnvironment, engine):
+    def test_sqlalchemy_backend_configuration(self, test_environment: TestEnvironment, session):
         """Test SQLAlchemy backend configuration with empty plugins"""
         # Verify that the connection is using the AWS wrapper
-        with engine.connect() as connection:
-            assert connection.connection is not None
+        assert session.connection().connection is not None
 
         # Test basic connection functionality
-        with Session(engine) as session:
-            assert session.query(TestModel).count() == 0
+        assert session.query(TestModel).count() == 0
 
     def test_sqlalchemy_basic_model_operations(self, session, test_environment: TestEnvironment):
         """Test basic SQLAlchemy ORM operations (CRUD)"""