Simple PG workflow working

Author: aaron-congo

aws_advanced_python_wrapper/__init__.py

@@ -15,8 +15,20 @@
 from logging import DEBUG, getLogger
 
 from .cleanup import release_resources
+from .driver_info import DriverInfo
 from .utils.utils import LogUtils
 from .wrapper import AwsWrapperConnection
+from aws_advanced_python_wrapper.pep249 import (
+    Error,
+    InterfaceError,
+    DatabaseError,
+    DataError,
+    OperationalError,
+    IntegrityError,
+    InternalError,
+    ProgrammingError,
+    NotSupportedError
+)
 
 # PEP249 compliance
 connect = AwsWrapperConnection.connect
@@ -32,9 +44,19 @@
     'set_logger',
     'apilevel',
     'threadsafety',
-    'paramstyle'
+    'paramstyle',
+    'Error',
+    'InterfaceError',
+    'DatabaseError',
+    'DataError',
+    'OperationalError',
+    'IntegrityError',
+    'InternalError',
+    'ProgrammingError',
+    'NotSupportedError'
 ]
 
+__version__ = DriverInfo.DRIVER_VERSION
 
 def set_logger(name='aws_advanced_python_wrapper', level=DEBUG, format_string=None):
     LogUtils.setup_logger(getLogger(name), level, format_string)

aws_advanced_python_wrapper/sqlalchemy/orm_dialect.py

@@ -0,0 +1,139 @@
+# aws_advanced_python_wrapper/sqlalchemy/sqlalchemy_psycopg_dialect.py
+from psycopg import Connection
+from sqlalchemy.dialects.postgresql.psycopg import PGDialect_psycopg
+import re
+
+class SqlAlchemyOrmPgDialect(PGDialect_psycopg):
+    """
+    SQLAlchemy dialect for AWS Advanced Python Wrapper.
+    Extends PostgreSQL psycopg dialect with Aurora-aware connection handling.
+    """
+
+    name = 'postgresql'
+    driver = 'aws_wrapper'
+
+    def __init__(self, **kwargs):
+        # Skip parent's version check since we're a wrapper, not psycopg itself
+        super(PGDialect_psycopg, self).__init__(**kwargs)
+
+        # Dynamically detect the actual psycopg version we're wrapping to ensure
+        # SQLAlchemy uses the correct feature set and SQL generation
+        try:
+            import psycopg
+            m = re.match(r"(\d+)\.(\d+)(?:\.(\d+))?", psycopg.__version__)
+            if m:
+                self.psycopg_version = tuple(
+                    int(x) for x in m.group(1, 2, 3) if x is not None
+                )
+            else:
+                self.psycopg_version = (3, 0, 2)  # Minimum supported
+        except (ImportError, AttributeError):
+            self.psycopg_version = (3, 0, 2)
+
+    @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 'target' parameter for the wrapper.
+        """
+        # 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 your wrapper
+        if 'target' not in opts:
+            opts['target'] = Connection.connect
+
+        # Return empty args list and kwargs dict
+        return ([], opts)
+
+    def on_connect(self):
+        """
+        Return a callable that will be executed on new connections. This can be used if we need to set any session-level
+        parameters.
+        """
+
+        def set_session_params(conn):
+            # Set any Aurora-specific session parameters
+            cursor = conn.cursor()
+            try:
+                # Example: Set statement timeout
+                cursor.execute("SET statement_timeout = '60s'")
+            finally:
+                cursor.close()
+
+        return set_session_params
+
+    def get_isolation_level(self, dbapi_connection):
+        """Get the current isolation level"""
+        cursor = dbapi_connection.cursor()
+        try:
+            cursor.execute("SHOW transaction_isolation")
+            val = cursor.fetchone()
+            if val:
+                # Extract first element from tuple and format
+                return val.upper().replace(' ', '_')
+            return 'READ_COMMITTED'  # PostgreSQL's default
+        finally:
+            cursor.close()
+
+    def initialize(self, connection):
+        """
+        Override initialization to handle type introspection.
+        The parent class tries to use TypeInfo.fetch() which requires
+        a native psycopg connection, not our wrapper.
+        """
+        # Find the AwsWrapperConnection at whatever nesting level
+        wrapper_conn = self._get_wrapper_connection(connection)
+
+        if wrapper_conn and hasattr(wrapper_conn, 'connection'):
+            # Get the underlying psycopg connection
+            underlying_conn = wrapper_conn.connection
+
+            # Temporarily swap the entire connection chain
+            original_dbapi_conn = connection.connection
+            connection.connection = underlying_conn
+
+            try:
+                # Call parent initialization with native psycopg connection
+                super().initialize(connection)
+            finally:
+                # Restore original connection chain
+                connection.connection = original_dbapi_conn
+        else:
+            # If we can't find wrapper or it doesn't expose underlying connection,
+            # skip type introspection (custom types won't be auto-configured)
+            pass
+
+    def _get_wrapper_connection(self, connection):
+        """
+        Traverse the connection chain to find AwsWrapperConnection.
+        Handles variable nesting depths depending on pool configuration.
+        """
+        from aws_advanced_python_wrapper import AwsWrapperConnection
+
+        # Start with the DBAPI connection
+        current = connection.connection
+
+        # Traverse up to 5 levels deep (reasonable limit)
+        for _ in range(5):
+            if isinstance(current, AwsWrapperConnection):
+                return current
+
+            # Try to go deeper if there's a .connection attribute
+            if hasattr(current, 'connection'):
+                current = current.connection
+            else:
+                break
+
+        return None

pyproject.toml

@@ -84,3 +84,5 @@ filterwarnings = [
     'ignore:Exception during reset or similar:pytest.PytestUnhandledThreadExceptionWarning'
 ]
 
+[tool.poetry.plugins."sqlalchemy.dialects"]
+"postgresql.aws_wrapper" = "aws_advanced_python_wrapper.sqlalchemy.orm_dialect:SqlAlchemyOrmPgDialect"

tests/unit/test_sqlalchemy_orm.py

@@ -0,0 +1,66 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License").
+#  You may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from sqlalchemy import create_engine, Column, Integer, String
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+class TestSqlAlchemyORM:
+    def test_basic_workflow(self):
+        # Step 1: Create engine (connection to database)
+        engine = create_engine('postgresql+aws_wrapper://pguser:pgpassword@mydb.cluster-XYZ.us-west-1.rds.amazonaws.com:5432/somedb')
+
+        # Step 2: Define base class for declarative models
+        Base = declarative_base()
+
+        # Step 3: Define model class (separate from database operations)
+        class User(Base):
+            __tablename__ = 'users'
+
+            id = Column(Integer, primary_key=True)
+            name = Column(String(50))
+            email = Column(String(100))
+
+        # Step 4: Create tables
+        Base.metadata.create_all(engine)
+
+        # Step 5: Create session factory
+        Session = sessionmaker(bind=engine)
+
+        # Step 6: Use session for database operations
+        session = Session()
+
+        # INSERT - Create new object and add to session
+        new_user = User(name='John Doe', email='john@example.com')
+        session.add(new_user)
+        session.commit()  # Explicit commit required
+
+        # SELECT - Query using session
+        users = session.query(User).filter(User.name == 'John Doe').all()
+        for user in users:
+            print(f"{user.name}: {user.email}")
+
+
+        # UPDATE - Modify object and commit
+        user = session.query(User).filter(User.name == "John Doe").first()
+        user.email = 'newemail@example.com'
+        session.commit()  # Changes tracked by session
+
+        # DELETE - Remove object from session
+        user_to_delete = session.query(User).filter(User.name == "John Doe").first()
+        session.delete(user_to_delete)
+        session.commit()
+
+        # Always close session when done
+        session.close()