Add SQLAlchemy ORM plugin tests

aws_advanced_python_wrapper/sqlalchemy/mysql_orm_dialect.py

@@ -12,9 +12,26 @@
 #  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
+from aws_advanced_python_wrapper.utils.properties import (Properties,
+                                                          PropertiesUtils)
+
+if TYPE_CHECKING:
+    from sqlalchemy import Connection
+
+    from aws_advanced_python_wrapper.hostinfo import HostInfo
 
 
 class SqlAlchemyOrmMysqlDialect(MySQLDialect_mysqlconnector):
@@ -27,3 +44,154 @@ 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"
+        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
+
+    def prepare_connect_info(self, host_info: HostInfo, props: Properties) -> Properties:
+        prop_copy: Properties = Properties(props.copy())
+
+        prop_copy["host"] = host_info.host
+
+        if host_info.is_port_specified():
+            prop_copy["port"] = str(host_info.port)
+
+        PropertiesUtils.remove_wrapper_props(prop_copy)
+        return prop_copy

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