chore: add logs to tests

Author: karenc-bq

.github/workflows/integration_tests.yml

@@ -19,7 +19,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version: [ "3.11", "3.12", "3.13"  ]
-        environment: [ "mysql", "pg" ]
+        environment: [ "pg" ]
 
     steps:
       - name: 'Clone repository'
@@ -73,7 +73,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version: [ "3.11", "3.12", "3.13" ]
-        environment: [ "mysql", "pg" ]
+        environment: [ "pg" ]
 
     steps:
       - name: 'Clone repository'

aws_advanced_python_wrapper/read_write_splitting_plugin.py

@@ -17,6 +17,8 @@
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Callable, Optional, Protocol, Set, Tuple
 
+from aws_advanced_python_wrapper.sqlalchemy_driver_dialect import SqlAlchemyDriverDialect
+
 if TYPE_CHECKING:
     from aws_advanced_python_wrapper.driver_dialect import DriverDialect
     from aws_advanced_python_wrapper.host_list_provider import HostListProviderService
@@ -402,13 +404,20 @@ def _close_connection_if_idle(self, internal_conn: Optional[Connection]):
         current_conn = self._plugin_service.current_connection
         driver_dialect = self._plugin_service.driver_dialect
 
+        logger.debug(f"current connection during close if idle: {current_conn}: {self._plugin_service.current_host_info}")
         if internal_conn == current_conn:
             # Connection is in use, do not close
+            logger.debug(f"current connection is in use. do not close")
             return
 
         try:
             if self._is_connection_usable(internal_conn, driver_dialect):
-                driver_dialect.execute(DbApiMethod.CONNECTION_CLOSE.method_name, lambda: internal_conn.close())
+                if isinstance(driver_dialect, SqlAlchemyDriverDialect):
+                    logger.debug(f"Invalidating pooled connection {internal_conn}")
+                    driver_dialect.execute(DbApiMethod.CONNECTION_CLOSE.method_name, lambda: internal_conn.invalidate(soft=True))
+                else:
+                    logger.debug(f"ReadWriteSplittingPlugin.ClosingIdleConnection: {internal_conn}")
+                    driver_dialect.execute(DbApiMethod.CONNECTION_CLOSE.method_name, lambda: internal_conn.close())
         except Exception:
             # Ignore exceptions during cleanup - connection might already be dead
             pass

tests/integration/container/test_custom_endpoint.py

@@ -61,7 +61,7 @@ def rds_utils(self):
         return RdsTestUtility(region)
 
     @pytest.fixture(scope='class')
-    def props(self):
+    def props_with_failover(self):
         p: Properties = Properties(
             {"plugins": "custom_endpoint,read_write_splitting,failover", "connect_timeout": 10_000, "autocommit": True, "cluster_id": "cluster1"})
 
@@ -77,6 +77,24 @@ def props(self):
 
         return p
 
+    @pytest.fixture(scope='class')
+    def props(self):
+        p: Properties = Properties(
+            {"plugins": "custom_endpoint,read_write_splitting", "connect_timeout": 10_000, "autocommit": True, "cluster_id": "cluster1"})
+
+        features = TestEnvironment.get_current().get_features()
+        if TestEnvironmentFeatures.TELEMETRY_TRACES_ENABLED in features \
+                or TestEnvironmentFeatures.TELEMETRY_METRICS_ENABLED in features:
+            WrapperProperties.ENABLE_TELEMETRY.set(p, True)
+            WrapperProperties.TELEMETRY_SUBMIT_TOPLEVEL.set(p, True)
+        if TestEnvironmentFeatures.TELEMETRY_TRACES_ENABLED in features:
+            WrapperProperties.TELEMETRY_TRACES_BACKEND.set(p, "XRAY")
+        if TestEnvironmentFeatures.TELEMETRY_METRICS_ENABLED in features:
+            WrapperProperties.TELEMETRY_METRICS_BACKEND.set(p, "OTLP")
+
+        return p
+
+
     @pytest.fixture(scope='class', autouse=True)
     def setup_and_teardown(self):
         env_info = TestEnvironment.get_current().get_info()
@@ -221,13 +239,13 @@ def wait_until_endpoint_has_members(self, rds_client, expected_members: Set[str]
         duration_sec = (perf_counter_ns() - start_ns) / 1_000_000_000
         self.logger.debug(f"wait_until_endpoint_has_specified_members took {duration_sec} seconds.")
 
-    def test_custom_endpoint_failover(self, test_driver: TestDriver, conn_utils, props, rds_utils):
-        props["failover_mode"] = "reader_or_writer"
+    def test_custom_endpoint_failover(self, test_driver: TestDriver, conn_utils, props_with_failover, rds_utils):
+        props_with_failover["failover_mode"] = "reader_or_writer"
 
         target_driver_connect = DriverHelper.get_connect_func(test_driver)
         kwargs = conn_utils.get_connect_params()
         kwargs["host"] = self.endpoint_info["Endpoint"]
-        conn = AwsWrapperConnection.connect(target_driver_connect, **kwargs, **props)
+        conn = AwsWrapperConnection.connect(target_driver_connect, **kwargs, **props_with_failover)
 
         endpoint_members = self.endpoint_info["StaticMembers"]
         instance_id = rds_utils.query_instance_id(conn)
@@ -281,7 +299,7 @@ def _setup_custom_endpoint_role(self, target_driver_connect, conn_kwargs, rds_ut
         self.logger.debug("Custom endpoint instance successfully set to role: " + host_role.name)
 
     def test_custom_endpoint_read_write_splitting__with_custom_endpoint_changes__with_reader_as_init_conn(
-            self, test_driver: TestDriver, conn_utils, props, rds_utils):
+            self, test_driver: TestDriver, conn_utils, props_with_failover, rds_utils):
         '''
         Will test for the following scenario:
         1. Initially connect to a reader instance via the custom endpoint.
@@ -297,13 +315,13 @@ def test_custom_endpoint_read_write_splitting__with_custom_endpoint_changes__wit
         kwargs["host"] = self.endpoint_info["Endpoint"]
         # This setting is not required for the test, but it allows us to also test re-creation of expired monitors since
         # it takes more than 30 seconds to modify the cluster endpoint (usually around 140s).
-        props["custom_endpoint_idle_monitor_expiration_ms"] = 30_000
-        props["wait_for_custom_endpoint_info_timeout_ms"] = 30_000
+        props_with_failover["custom_endpoint_idle_monitor_expiration_ms"] = 30_000
+        props_with_failover["wait_for_custom_endpoint_info_timeout_ms"] = 30_000
 
         # Ensure that we are starting with a reader connection
         self._setup_custom_endpoint_role(target_driver_connect, kwargs, rds_utils, HostRole.READER)
 
-        conn = AwsWrapperConnection.connect(target_driver_connect, **kwargs, **props)
+        conn = AwsWrapperConnection.connect(target_driver_connect, **kwargs, **props_with_failover)
         endpoint_members = self.endpoint_info["StaticMembers"]
         original_reader_id = rds_utils.query_instance_id(conn)
         assert original_reader_id in endpoint_members
@@ -350,28 +368,28 @@ def test_custom_endpoint_read_write_splitting__with_custom_endpoint_changes__wit
 
     def test_custom_endpoint_read_write_splitting__with_custom_endpoint_changes__with_writer_as_init_conn(
             self, test_driver: TestDriver, conn_utils, props, rds_utils):
-        '''
+        """
         Will test for the following scenario:
-        1. Iniitially connect to the writer instance via the custom endpoint.
+        1. Initially connect to the writer instance via the custom endpoint.
         2. Attempt to switch to reader instance - should succeed, but will still use writer instance as reader.
         3. Modify the custom endpoint to add a reader instance as a static member.
         4. Switch to reader instance - should succeed.
         5. Switch back to writer instance - should succeed.
         6. Modify the custom endpoint to remove the reader instance as a static member.
         7. Attempt to switch to reader instance - should fail since the custom endpoint no longer has the reader instance.
-        '''
+        """
 
         target_driver_connect = DriverHelper.get_connect_func(test_driver)
         kwargs = conn_utils.get_connect_params()
         kwargs["host"] = self.endpoint_info["Endpoint"]
         # This setting is not required for the test, but it allows us to also test re-creation of expired monitors since
         # it takes more than 30 seconds to modify the cluster endpoint (usually around 140s).
-        props["custom_endpoint_idle_monitor_expiration_ms"] = 30_000
-        props["wait_for_custom_endpoint_info_timeout_ms"] = 30_000
+        props_with_failover["custom_endpoint_idle_monitor_expiration_ms"] = 30_000
+        props_with_failover["wait_for_custom_endpoint_info_timeout_ms"] = 30_000
 
         # Ensure that we are starting with a writer connection
         self._setup_custom_endpoint_role(target_driver_connect, kwargs, rds_utils, HostRole.WRITER)
-        conn = AwsWrapperConnection.connect(target_driver_connect, **kwargs, **props)
+        conn = AwsWrapperConnection.connect(target_driver_connect, **kwargs, **props_with_failover)
 
         endpoint_members = self.endpoint_info["StaticMembers"]
         original_writer_id = str(rds_utils.query_instance_id(conn))

tests/integration/container/test_read_write_splitting.py

@@ -832,6 +832,7 @@ def test_pooled_connection__cluster_url_failover(
                          TestEnvironmentFeatures.NETWORK_OUTAGES_ENABLED,
                          TestEnvironmentFeatures.ABORT_CONNECTION_SUPPORTED])
     @disable_on_engines([DatabaseEngine.MYSQL])
+    @pytest.mark.repeat(10)
     def test_pooled_connection__failover_failed(
         self,
         test_environment: TestEnvironment,