add async connection pool

Author: faviansamatha

aws_advanced_python_wrapper/errors.py

@@ -45,3 +45,27 @@ class FailoverSuccessError(FailoverError):
 
 class ReadWriteSplittingError(AwsWrapperError):
     __module__ = "aws_advanced_python_wrapper"
+
+
+class AsyncConnectionPoolError(AwsWrapperError):
+    __module__ = "aws_advanced_python_wrapper"
+
+
+class PoolNotInitializedError(AsyncConnectionPoolError):
+    __module__ = "aws_advanced_python_wrapper"
+
+
+class PoolClosingError(AsyncConnectionPoolError):
+    __module__ = "aws_advanced_python_wrapper"
+
+
+class PoolExhaustedError(AsyncConnectionPoolError):
+    __module__ = "aws_advanced_python_wrapper"
+
+
+class ConnectionReleasedError(AsyncConnectionPoolError):
+    __module__ = "aws_advanced_python_wrapper"
+
+
+class PoolSizeLimitError(AsyncConnectionPoolError):
+    __module__ = "aws_advanced_python_wrapper"

aws_advanced_python_wrapper/iam_plugin.py

@@ -114,7 +114,7 @@ def _connect(self, host_info: HostInfo, props: Properties, connect_func: Callabl
 
             is_cached_token = (token_info is not None and not token_info.is_expired())
             if not self._plugin_service.is_login_exception(error=e) or not is_cached_token:
-                raise AwsWrapperError(Messages.get_formatted("IamAuthPlugin.ConnectException", e)) from e
+                raise
 
             # Login unsuccessful with cached token
             # Try to generate a new token and try to connect again

aws_advanced_python_wrapper/tortoise/async_support/async_wrapper.py

@@ -37,38 +37,6 @@ async def close_aws_wrapper(connection: AwsWrapperConnection) -> None:
         """Close an AWS wrapper connection asynchronously."""
         await asyncio.to_thread(connection.close)
 
-
-class AwsCursorAsyncWrapper:
-    """Wraps sync AwsCursor cursor with async support."""
-
-    def __init__(self, sync_cursor):
-        self._cursor = sync_cursor
-
-    async def execute(self, query, params=None):
-        """Execute a query asynchronously."""
-        return await asyncio.to_thread(self._cursor.execute, query, params)
-
-    async def executemany(self, query, params_list):
-        """Execute multiple queries asynchronously."""
-        return await asyncio.to_thread(self._cursor.executemany, query, params_list)
-
-    async def fetchall(self):
-        """Fetch all results asynchronously."""
-        return await asyncio.to_thread(self._cursor.fetchall)
-
-    async def fetchone(self):
-        """Fetch one result asynchronously."""
-        return await asyncio.to_thread(self._cursor.fetchone)
-
-    async def close(self):
-        """Close cursor asynchronously."""
-        return await asyncio.to_thread(self._cursor.close)
-
-    def __getattr__(self, name):
-        """Delegate non-async attributes to the wrapped cursor."""
-        return getattr(self._cursor, name)
-
-
 class AwsConnectionAsyncWrapper(AwsWrapperConnection):
     """Wraps sync AwsConnection with async cursor support."""
 
@@ -96,6 +64,10 @@ async def set_autocommit(self, value: bool):
         """Set autocommit mode."""
         return await asyncio.to_thread(setattr, self._wrapped_connection, 'autocommit', value)
 
+    async def close(self):
+        """Close the connection asynchronously."""
+        return await asyncio.to_thread(self._wrapped_connection.close)
+
     def __getattr__(self, name):
         """Delegate all other attributes/methods to the wrapped connection."""
         return getattr(self._wrapped_connection, name)
@@ -105,3 +77,33 @@ def __del__(self):
         if hasattr(self, '_wrapped_connection'):
             # Let the wrapped connection handle its own cleanup
             pass
+
+class AwsCursorAsyncWrapper:
+    """Wraps sync AwsCursor cursor with async support."""
+
+    def __init__(self, sync_cursor):
+        self._cursor = sync_cursor
+
+    async def execute(self, query, params=None):
+        """Execute a query asynchronously."""
+        return await asyncio.to_thread(self._cursor.execute, query, params)
+
+    async def executemany(self, query, params_list):
+        """Execute multiple queries asynchronously."""
+        return await asyncio.to_thread(self._cursor.executemany, query, params_list)
+
+    async def fetchall(self):
+        """Fetch all results asynchronously."""
+        return await asyncio.to_thread(self._cursor.fetchall)
+
+    async def fetchone(self):
+        """Fetch one result asynchronously."""
+        return await asyncio.to_thread(self._cursor.fetchone)
+
+    async def close(self):
+        """Close cursor asynchronously."""
+        return await asyncio.to_thread(self._cursor.close)
+
+    def __getattr__(self, name):
+        """Delegate non-async attributes to the wrapped cursor."""
+        return getattr(self._cursor, name)

aws_advanced_python_wrapper/tortoise/backends/base/client.py

@@ -16,6 +16,7 @@
 
 from typing import Any, Callable, Dict, Generic, cast
 
+import asyncio
 import mysql.connector
 from tortoise.backends.base.client import (BaseDBAsyncClient, T_conn,
                                            TransactionalDBClient,
@@ -110,3 +111,85 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         finally:
             await AwsWrapperAsyncConnector.close_aws_wrapper(self.client._connection)
             connections.reset(self.token)
+
+class TortoiseAwsClientPooledConnectionWrapper(Generic[T_conn]):
+    """Manages acquiring from and releasing connections to a pool."""
+
+    __slots__ = ("client", "connection", "_pool_init_lock", "with_db")
+
+    def __init__(
+        self, 
+        client: BaseDBAsyncClient, 
+        pool_init_lock: asyncio.Lock, 
+        with_db: bool = True
+    ) -> None:
+        self.client = client
+        self.connection: T_conn | None = None
+        self._pool_init_lock = pool_init_lock
+        self.with_db = with_db
+
+    async def ensure_connection(self) -> None:
+        """Ensure the connection pool is initialized."""
+        if not self.client._pool:
+            async with self._pool_init_lock:
+                if not self.client._pool:
+                    await self.client.create_connection(with_db=self.with_db)
+
+    async def __aenter__(self) -> T_conn:
+        """Acquire connection from pool."""
+        await self.ensure_connection()
+        self.connection = await self.client._pool.acquire()
+        return self.connection
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Close connection and release back to pool."""
+        if self.connection:
+            await self.connection.release()
+
+class TortoiseAwsClientPooledTransactionContext(TransactionContext):
+    """Transaction context that uses a pool to acquire connections."""
+
+    __slots__ = ("client", "connection_name", "token", "_pool_init_lock", "connection")
+
+    def __init__(self, client: TransactionalDBClient, pool_init_lock: asyncio.Lock) -> None:
+        self.client = client
+        self.connection_name = client.connection_name
+        self._pool_init_lock = pool_init_lock
+        self.connection = None
+
+    async def ensure_connection(self) -> None:
+        """Ensure the connection pool is initialized."""
+        if not self.client._parent._pool:
+            # a safeguard against multiple concurrent tasks trying to initialize the pool
+            async with self._pool_init_lock:
+                if not self.client._parent._pool:
+                    await self.client._parent.create_connection(with_db=True)
+
+    async def __aenter__(self) -> TransactionalDBClient:
+        """Enter transaction context."""
+        await self.ensure_connection()
+        
+        # Set the context variable so the current task sees a TransactionWrapper connection
+        self.token = connections.set(self.connection_name, self.client)
+        
+        # Create connection and begin transaction
+        self.connection = await self.client._parent._pool.acquire()
+        self.client._connection =  self.connection
+        await self.client.begin()
+        return self.client
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Exit transaction context with proper cleanup."""
+        try:
+            if not self.client._finalized:
+                if exc_type:
+                    # Can't rollback a transaction that already failed
+                    if exc_type is not TransactionManagementError:
+                        await self.client.rollback()
+                else:
+                    await self.client.commit()
+        finally:
+            if self.client._connection:
+                await self.client._connection.release()
+                # self.client._connection = None
+            connections.reset(self.token)

aws_advanced_python_wrapper/tortoise/backends/mysql/client.py

@@ -31,14 +31,16 @@
 
 from aws_advanced_python_wrapper.errors import AwsWrapperError, FailoverError
 from aws_advanced_python_wrapper.tortoise.async_support.async_wrapper import (
-    AwsConnectionAsyncWrapper)
+    AwsConnectionAsyncWrapper, AwsWrapperAsyncConnector)
 from aws_advanced_python_wrapper.tortoise.backends.base.client import (
     AwsBaseDBAsyncClient, AwsTransactionalDBClient,
-    TortoiseAwsClientConnectionWrapper, TortoiseAwsClientTransactionContext)
+    TortoiseAwsClientPooledConnectionWrapper, TortoiseAwsClientPooledTransactionContext)
 from aws_advanced_python_wrapper.tortoise.backends.mysql.executor import \
     AwsMySQLExecutor
 from aws_advanced_python_wrapper.tortoise.backends.mysql.schema_generator import \
     AwsMySQLSchemaGenerator
+from aws_advanced_python_wrapper.async_connection_pool import AsyncConnectionPool, PoolConfig
+from dataclasses import fields
 from aws_advanced_python_wrapper.utils.log import Logger
 
 logger = Logger(__name__)
@@ -125,6 +127,21 @@ def __init__(
         # Initialize state
         self._template: Dict[str, Any] = {}
         self._connection = None
+        self._pool_init_lock: asyncio.Lock = asyncio.Lock()
+        self._pool: Optional[AsyncConnectionPool] = None
+
+        # Pool configuration
+        default_pool_config = {field.name: field.default for field in fields(PoolConfig)}
+        self._pool_config = PoolConfig(
+            min_size = self.extra.pop("min_size", default_pool_config["min_size"]),
+            max_size = self.extra.pop("max_size", default_pool_config["max_size"]),
+            timeout = self.extra.pop("pool_timeout", default_pool_config["timeout"]),
+            max_lifetime = self.extra.pop("pool_lifetime", default_pool_config["max_lifetime"]),
+            max_idle_time = self.extra.pop("pool_max_idle_time", default_pool_config["max_idle_time"]),
+            health_check_interval = self.extra.pop("pool_health_check_interval", default_pool_config["health_check_interval"]),
+            pre_ping = self.extra.pop("pre_ping", default_pool_config["pre_ping"])
+        )
+        
 
     def _init_connection_templates(self) -> None:
         """Initialize connection templates for with/without database."""
@@ -140,6 +157,29 @@ def _init_connection_templates(self) -> None:
         self._template_with_db = {**base_template, "database": self.database}
         self._template_no_db = {**base_template, "database": None}
 
+    async def _init_pool(self) -> None:
+        """Initialize the connection pool."""
+        if self._pool is not None:
+            return
+
+        async def create_connection():
+            return await AwsWrapperAsyncConnector.connect_with_aws_wrapper(mysql.connector.Connect, **self._template)
+
+        async def health_check(conn):
+            is_closed = await asyncio.to_thread(lambda: conn._wrapped_connection.is_closed)
+            if is_closed:
+                raise Exception("Connection is closed")
+            else:
+                print("NOT CLOSED!")
+
+        self._pool: AsyncConnectionPool = AsyncConnectionPool(
+            creator=create_connection,
+            # closer=close_connection,
+            health_check=health_check,
+            config=self._pool_config
+        )
+        await self._pool.initialize()
+
     # Connection Management
     async def create_connection(self, with_db: bool) -> None:
         """Initialize connection pool and configure database settings."""
@@ -154,18 +194,24 @@ async def create_connection(self, with_db: bool) -> None:
         # Set template based on database requirement
         self._template = self._template_with_db if with_db else self._template_no_db
 
+        await self._init_pool()
+        print("Pool is initialized")
+        print(self._pool.get_stats())
+
     async def close(self) -> None:
         """Close connections - AWS wrapper handles cleanup internally."""
-        pass
+        if hasattr(self, '_pool') and self._pool:
+            await self._pool.close()
+            self._pool = None
 
     def acquire_connection(self):
         """Acquire a connection from the pool."""
         return self._acquire_connection(with_db=True)
 
-    def _acquire_connection(self, with_db: bool) -> TortoiseAwsClientConnectionWrapper:
+    def _acquire_connection(self, with_db: bool) -> TortoiseAwsClientPooledConnectionWrapper:
         """Create connection wrapper for specified database mode."""
-        return TortoiseAwsClientConnectionWrapper(
-            self, mysql.connector.Connect, with_db=with_db
+        return TortoiseAwsClientPooledConnectionWrapper(
+            self, pool_init_lock=self._pool_init_lock, with_db=with_db
         )
 
     # Database Operations
@@ -253,7 +299,7 @@ async def _execute_script(self, query: str, with_db: bool) -> None:
     # Transaction Support
     def _in_transaction(self) -> TransactionContext:
         """Create a new transaction context."""
-        return TortoiseAwsClientTransactionContext(TransactionWrapper(self))
+        return TortoiseAwsClientPooledTransactionContext(TransactionWrapper(self), self._pool_init_lock)
 
 
 class TransactionWrapper(AwsMySQLClient, AwsTransactionalDBClient):
@@ -262,6 +308,7 @@ class TransactionWrapper(AwsMySQLClient, AwsTransactionalDBClient):
     def __init__(self, connection: AwsMySQLClient) -> None:
         self.connection_name = connection.connection_name
         self._connection: AwsConnectionAsyncWrapper = connection._connection
+
         self._lock = asyncio.Lock()
         self._savepoint: Optional[str] = None
         self._finalized: bool = False

tests/integration/container/tortoise/test_tortoise_common.py

@@ -40,7 +40,6 @@ async def setup_tortoise(conn_utils, plugins="aurora_connection_tracker", **kwar
         plugins=plugins,
         **kwargs,
     )
-
     config = {
         "connections": {
             "default": db_url
@@ -53,17 +52,13 @@ async def setup_tortoise(conn_utils, plugins="aurora_connection_tracker", **kwar
         }
     }
 
-    from aws_advanced_python_wrapper.tortoise.sql_alchemy_tortoise_connection_provider import \
-        setup_tortoise_connection_provider
-    setup_tortoise_connection_provider()
     await Tortoise.init(config=config)
-    await Tortoise.generate_schemas()
 
+    await Tortoise.generate_schemas()
     await clear_test_models()
 
     yield
 
-    await clear_test_models()
     await reset_tortoise()
 
 

tests/integration/container/tortoise/test_tortoise_config.py

@@ -18,8 +18,6 @@
 
 # Import to register the aws-mysql backend
 import aws_advanced_python_wrapper.tortoise  # noqa: F401
-from aws_advanced_python_wrapper.tortoise.sql_alchemy_tortoise_connection_provider import \
-    setup_tortoise_connection_provider
 from tests.integration.container.tortoise.models.test_models import User
 from tests.integration.container.tortoise.test_tortoise_common import \
     reset_tortoise
@@ -67,7 +65,6 @@ async def setup_tortoise_dict_config(self, conn_utils):
             }
         }
 
-        setup_tortoise_connection_provider()
         await Tortoise.init(config=config)
         await Tortoise.generate_schemas()
         await self._clear_all_test_models()
@@ -121,7 +118,6 @@ async def setup_tortoise_multi_db(self, conn_utils):
             }
         }
 
-        setup_tortoise_connection_provider()
         await Tortoise.init(config=config)
 
         # Create second database
@@ -169,7 +165,6 @@ async def setup_tortoise_with_router(self, conn_utils):
             "routers": ["tests.integration.container.tortoise.router.test_router.TestRouter"]
         }
 
-        setup_tortoise_connection_provider()
         await Tortoise.init(config=config)
         await Tortoise.generate_schemas()
         await self._clear_all_test_models()