Added custom pool

Author: faviansamatha

aws_advanced_python_wrapper/async_connection_pool.py

@@ -169,7 +169,8 @@ def _run(self):
                             len(endpoints),
                             endpoint_hostnames)
 
-                        sleep(self._refresh_rate_ns / 1_000_000_000)
+                        if self._stop_event.wait(self._refresh_rate_ns / 1_000_000_000):
+                            break
                         continue
 
                     endpoint_info = CustomEndpointInfo.from_db_cluster_endpoint(endpoints[0])
@@ -178,7 +179,8 @@ def _run(self):
                     if cached_info is not None and cached_info == endpoint_info:
                         elapsed_time = perf_counter_ns() - start_ns
                         sleep_duration = max(0, self._refresh_rate_ns - elapsed_time)
-                        sleep(sleep_duration / 1_000_000_000)
+                        if self._stop_event.wait(sleep_duration / 1_000_000_000):
+                            break
                         continue
 
                     logger.debug(
@@ -196,7 +198,8 @@ def _run(self):
 
                     elapsed_time = perf_counter_ns() - start_ns
                     sleep_duration = max(0, self._refresh_rate_ns - elapsed_time)
-                    sleep(sleep_duration / 1_000_000_000)
+                    if self._stop_event.wait(sleep_duration / 1_000_000_000):
+                        break
                     continue
                 except InterruptedError as e:
                     raise e
@@ -219,7 +222,6 @@ def close(self):
         CustomEndpointMonitor._custom_endpoint_info_cache.remove(self._custom_endpoint_host_info.host)
         self._stop_event.set()
 
-
 class CustomEndpointPlugin(Plugin):
     """
     A plugin that analyzes custom endpoints for custom endpoint information and custom endpoint changes, such as adding

aws_advanced_python_wrapper/custom_endpoint_plugin.py

@@ -69,3 +69,7 @@ class ConnectionReleasedError(AsyncConnectionPoolError):
 
 class PoolSizeLimitError(AsyncConnectionPoolError):
     __module__ = "aws_advanced_python_wrapper"
+
+
+class PoolHealthCheckError(AsyncConnectionPoolError):
+    __module__ = "aws_advanced_python_wrapper"

aws_advanced_python_wrapper/errors.py

@@ -14,6 +14,16 @@
 
 from tortoise.backends.base.config_generator import DB_LOOKUP
 
+
+def cast_to_bool(value):
+    """Generic function to cast various types to boolean."""
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        return value.lower() in ('true', '1', 'yes', 'on')
+    return bool(value)
+
+
 # Register AWS MySQL backend
 DB_LOOKUP["aws-mysql"] = {
     "engine": "aws_advanced_python_wrapper.tortoise.backends.mysql",
@@ -29,9 +39,9 @@
         "minsize": int,
         "maxsize": int,
         "connect_timeout": int,
-        "echo": bool,
-        "use_unicode": bool,
-        "ssl": bool,
-        "use_pure": bool,
+        "echo": cast_to_bool,
+        "use_unicode": cast_to_bool,
+        "ssl": cast_to_bool,
+        "use_pure": cast_to_bool
     },
 }

aws_advanced_python_wrapper/tortoise/__init__.py

@@ -37,103 +37,26 @@ class AwsTransactionalDBClient(TransactionalDBClient):
     _parent: AwsBaseDBAsyncClient
     pass
 
-
-class TortoiseAwsClientConnectionWrapper(Generic[T_conn]):
-    """Manages acquiring from and releasing connections to a pool."""
-
-    __slots__ = ("client", "connection", "connect_func", "with_db")
-
-    def __init__(
-        self,
-        client: AwsBaseDBAsyncClient,
-        connect_func: Callable,
-        with_db: bool = True
-    ) -> None:
-        self.connect_func = connect_func
-        self.client = client
-        self.connection: AwsConnectionAsyncWrapper | None = None
-        self.with_db = with_db
-
-    async def ensure_connection(self) -> None:
-        """Ensure the connection pool is initialized."""
-        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 AwsWrapperAsyncConnector.connect_with_aws_wrapper(self.connect_func, **self.client._template)
-        return cast("T_conn", 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 AwsWrapperAsyncConnector.close_aws_wrapper(self.connection)
-
-
-class TortoiseAwsClientTransactionContext(TransactionContext):
-    """Transaction context that uses a pool to acquire connections."""
-
-    __slots__ = ("client", "connection_name", "token")
-
-    def __init__(self, client: AwsTransactionalDBClient) -> None:
-        self.client: AwsTransactionalDBClient = client
-        self.connection_name = client.connection_name
-
-    async def ensure_connection(self) -> None:
-        """Ensure the connection pool is initialized."""
-        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.client._connection = await AwsWrapperAsyncConnector.connect_with_aws_wrapper(
-            mysql.connector.Connect,
-            **self.client._parent._template
-        )
-        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:
-            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")
+    __slots__ = ("client", "connection", "_pool_init_lock",)
 
     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)
+                    await self.client.create_connection(with_db=True)
 
     async def __aenter__(self) -> T_conn:
         """Acquire connection from pool."""

aws_advanced_python_wrapper/tortoise/backends/base/client.py

@@ -135,10 +135,10 @@ def __init__(
         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"]),
+            acquire_conn_timeout = self.extra.pop("acquire_conn_timeout", default_pool_config["acquire_conn_timeout"]),
+            max_conn_lifetime = self.extra.pop("max_conn_lifetime", default_pool_config["max_conn_lifetime"]),
+            max_conn_idle_time = self.extra.pop("max_conn_idle_time", default_pool_config["max_conn_idle_time"]),
+            health_check_interval = self.extra.pop("health_check_interval", default_pool_config["health_check_interval"]),
             pre_ping = self.extra.pop("pre_ping", default_pool_config["pre_ping"])
         )
 
@@ -165,17 +165,8 @@ async def _init_pool(self) -> None:
         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()
@@ -195,8 +186,13 @@ async def create_connection(self, with_db: bool) -> None:
         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())
+
+    def _disable_pool_for_testing(self) -> None:
+        """Disable pool initialization for unit testing."""
+        self._pool = None
+        async def _no_op():
+            pass
+        self._init_pool = _no_op
 
     async def close(self) -> None:
         """Close connections - AWS wrapper handles cleanup internally."""
@@ -206,25 +202,21 @@ async def close(self) -> 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) -> TortoiseAwsClientPooledConnectionWrapper:
-        """Create connection wrapper for specified database mode."""
         return TortoiseAwsClientPooledConnectionWrapper(
-            self, pool_init_lock=self._pool_init_lock, with_db=with_db
+            self, pool_init_lock=self._pool_init_lock
         )
 
     # Database Operations
     async def db_create(self) -> None:
         """Create the database."""
         await self.create_connection(with_db=False)
-        await self._execute_script(f"CREATE DATABASE {self.database};", False)
+        await self.execute_script(f"CREATE DATABASE {self.database};")
         await self.close()
 
     async def db_delete(self) -> None:
         """Delete the database."""
         await self.create_connection(with_db=False)
-        await self._execute_script(f"DROP DATABASE {self.database};", False)
+        await self.execute_script(f"DROP DATABASE {self.database};")
         await self.close()
 
     # Query Execution Methods
@@ -279,14 +271,10 @@ async def execute_query_dict(self, query: str, values: Optional[List[Any]] = Non
         """Execute a query and return only the results as dictionaries."""
         return (await self.execute_query(query, values))[1]
 
+    @translate_exceptions
     async def execute_script(self, query: str) -> None:
         """Execute a script query."""
-        await self._execute_script(query, True)
-
-    @translate_exceptions
-    async def _execute_script(self, query: str, with_db: bool) -> None:
-        """Execute a multi-statement query by parsing and running statements sequentially."""
-        async with self._acquire_connection(with_db) as connection:
+        async with self.acquire_connection() as connection:
             logger.debug(f"Executing script: {query}")
             async with connection.cursor() as cursor:
                 # Parse multi-statement queries since MySQL Connector doesn't handle them well

aws_advanced_python_wrapper/tortoise/backends/mysql/client.py

@@ -62,16 +62,17 @@ def conn_utils():
 
 def pytest_runtest_setup(item):
     test_name: Optional[str] = None
+    full_test_name = item.nodeid  # Full test path including class and method
+    
     if hasattr(item, "callspec"):
         current_driver = item.callspec.params.get("test_driver")
         TestEnvironment.get_current().set_current_driver(current_driver)
-        test_name = item.callspec.id
+        test_name = f"{item.name}[{item.callspec.id}]"
     else:
         TestEnvironment.get_current().set_current_driver(None)
-        # Fallback to item.name if no callspec (for non-parameterized tests)
-        test_name = getattr(item, 'name', None) or str(item)
+        test_name = item.name
 
-    logger.info(f"Starting test preparation for: {test_name}")
+    logger.info(f"Starting test preparation for: {test_name} (full: {full_test_name})")
 
     segment: Optional[Segment] = None
     if TestEnvironmentFeatures.TELEMETRY_TRACES_ENABLED in TestEnvironment.get_current().get_features():
@@ -82,6 +83,7 @@ def pytest_runtest_setup(item):
                                .get_info().get_request().get_target_python_version().name)
         if test_name is not None:
             segment.put_annotation("test_name", test_name)
+            segment.put_annotation("full_test_name", full_test_name)
 
     info = TestEnvironment.get_current().get_info()
     request = info.get_request()

tests/integration/container/conftest.py

@@ -72,6 +72,7 @@ def create_custom_endpoint(self, rds_utils):
         finally:
             try:
                 rds_client.delete_db_cluster_endpoint(DBClusterEndpointIdentifier=self.endpoint_id)
+                self._wait_until_endpoint_deleted(rds_client)
             except ClientError as e:
                 if e.response['Error']['Code'] != 'DBClusterEndpointNotFoundFault':
                     pass  # Ignore if endpoint doesn't exist
@@ -108,18 +109,46 @@ def _wait_until_endpoint_available(self, rds_client):
 
         if not available:
             pytest.fail(f"Timed out waiting for custom endpoint to become available: {self.endpoint_id}")
+    
+    def _wait_until_endpoint_deleted(self, rds_client):
+        """Wait for the custom endpoint to be deleted."""
+        end_ns = perf_counter_ns() + 5 * 60 * 1_000_000_000  # 5 minutes
+        
+        while perf_counter_ns() < end_ns:
+            try:
+                rds_client.describe_db_cluster_endpoints(DBClusterEndpointIdentifier=self.endpoint_id)
+                sleep(5)  # Still exists, keep waiting
+            except ClientError as e:
+                if e.response['Error']['Code'] == 'DBClusterEndpointNotFoundFault':
+                    return  # Successfully deleted
+                raise  # Other error, re-raise
 
     @pytest_asyncio.fixture
-    async def setup_tortoise_custom_endpoint(self, conn_utils, create_custom_endpoint):
+    async def setup_tortoise_custom_endpoint(self, conn_utils, create_custom_endpoint, request):
         """Setup Tortoise with custom endpoint plugin."""
-        async for result in setup_tortoise(conn_utils, plugins="custom_endpoint,aurora_connection_tracker", host=create_custom_endpoint):
+        plugins, user = request.param
+        user_value = getattr(conn_utils, user) if user != "default" else None
+        
+        kwargs = {}
+        if "fastest_response_strategy" in plugins:
+            kwargs["reader_host_selector_strategy"] = "fastest_response"
+            
+        async for result in setup_tortoise(conn_utils, plugins=plugins, host=create_custom_endpoint, user=user_value, **kwargs):
             yield result
 
+    @pytest.mark.parametrize("setup_tortoise_custom_endpoint", [
+        ("custom_endpoint,aurora_connection_tracker", "default"),
+        ("failover,iam,aurora_connection_tracker,custom_endpoint,fastest_response_strategy", "iam_user")
+    ], indirect=True)
     @pytest.mark.asyncio
     async def test_basic_read_operations(self, setup_tortoise_custom_endpoint):
         """Test basic read operations with custom endpoint plugin."""
         await run_basic_read_operations("Custom Test", "custom")
 
+    @pytest.mark.parametrize("setup_tortoise_custom_endpoint", [
+        ("custom_endpoint,aurora_connection_tracker", "default"),
+        ("failover,iam,aurora_connection_tracker,custom_endpoint,fastest_response_strategy", "iam_user")
+    ], indirect=True)
     @pytest.mark.asyncio
     async def test_basic_write_operations(self, setup_tortoise_custom_endpoint):
         """Test basic write operations with custom endpoint plugin."""