test: fix failing integration tests

.github/workflows/integration_tests.yml

@@ -11,14 +11,67 @@ permissions:
   contents: read    # This is required for actions/checkout
 
 jobs:
-  build-integration-tests:
-    name: Run Integration Tests
+  lts-integration-tests:
+    name: Run LTS Integration Tests
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
         python-version: [ "3.8", "3.11" ]
-        engine-version: [ "lts", "latest"]
+        environment: [ "mysql", "pg" ]
+
+    steps:
+      - name: 'Clone repository'
+        uses: actions/checkout@v4
+
+      - name: 'Set up JDK 8'
+        uses: actions/setup-java@v4
+        with:
+          distribution: 'corretto'
+          java-version: 8
+
+      - name: Install poetry
+        shell: bash
+        run: |
+          pipx install poetry==1.8.2
+          poetry config virtualenvs.prefer-active-python true
+
+      - name: Install dependencies
+        run: poetry install
+
+      - name: 'Configure AWS Credentials'
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: arn:aws:iam::${{ secrets.AWS_ACCOUNT_ID }}:role/${{ secrets.AWS_DEPLOY_ROLE }}
+          role-session-name: python_integration_tests
+          role-duration-seconds: 21600
+          aws-region: ${{ secrets.AWS_DEFAULT_REGION }}
+
+      - name: 'Run LTS Integration Tests'
+        run: |
+          ./gradlew --no-parallel --no-daemon test-python-${{ matrix.python-version }}-${{ matrix.environment }} --info
+        env:
+          RDS_CLUSTER_DOMAIN: ${{ secrets.DB_CONN_SUFFIX }}
+          RDS_DB_REGION: ${{ secrets.AWS_DEFAULT_REGION }}
+          AURORA_MYSQL_DB_ENGINE_VERSION: lts
+          AURORA_PG_ENGINE_VERSION: lts
+
+      - name: 'Archive LTS results'
+        if: always()
+        uses: actions/upload-artifact@v4
+        with:
+          name: pytest-integration-report-${{ matrix.python-version }}-${{ matrix.environment }}-lts
+          path: ./tests/integration/container/reports
+          retention-days: 5
+
+  latest-integration-tests:
+    name: Run Latest Integration Tests
+    runs-on: ubuntu-latest
+    needs: lts-integration-tests
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [ "3.8", "3.11" ]
         environment: ["mysql", "pg"]
 
     steps:
@@ -48,19 +101,19 @@ jobs:
           role-duration-seconds: 21600
           aws-region: ${{ secrets.AWS_DEFAULT_REGION }}
 
-      - name: 'Run Integration Tests'
+      - name: 'Run Latest Integration Tests'
         run: |
           ./gradlew --no-parallel --no-daemon test-python-${{ matrix.python-version }}-${{ matrix.environment }} --info
         env:
           RDS_CLUSTER_DOMAIN: ${{ secrets.DB_CONN_SUFFIX }}
           RDS_DB_REGION: ${{ secrets.AWS_DEFAULT_REGION }}
-          AURORA_MYSQL_DB_ENGINE_VERSION: ${{ matrix.engine-version }}
-          AURORA_PG_ENGINE_VERSION: ${{ matrix.engine-version }}
+          AURORA_MYSQL_DB_ENGINE_VERSION: latest
+          AURORA_PG_ENGINE_VERSION: latest
 
-      - name: 'Archive results'
+      - name: 'Archive Latest results'
         if: always()
         uses: actions/upload-artifact@v4
         with:
-          name: pytest-integration-report-${{ matrix.python-version }}-${{ matrix.environment }}-${{ matrix.engine-version }}
+          name: pytest-integration-report-${{ matrix.python-version }}-${{ matrix.environment }}-latest
           path: ./tests/integration/container/reports
           retention-days: 5

aws_advanced_python_wrapper/driver_dialect.py

@@ -38,10 +38,11 @@ class DriverDialect(ABC):
     Driver dialects help the driver-agnostic AWS Python Driver interface with the driver-specific functionality of the underlying Python Driver.
     """
     _QUERY = "SELECT 1"
+    _ALL_METHODS = "*"
 
-    _executor: ClassVar[Executor] = ThreadPoolExecutor()
+    _executor: ClassVar[Executor] = ThreadPoolExecutor(thread_name_prefix="DriverDialectExecutor")
     _dialect_code: str = DriverDialectCodes.GENERIC
-    _network_bound_methods: Set[str] = {"*"}
+    _network_bound_methods: Set[str] = {_ALL_METHODS}
     _read_only: bool = False
     _autocommit: bool = False
     _driver_name: str = "Generic"
@@ -127,7 +128,7 @@ def execute(
             *args: Any,
             exec_timeout: Optional[float] = None,
             **kwargs: Any) -> Cursor:
-        if method_name not in self._network_bound_methods:
+        if DriverDialect._ALL_METHODS not in self.network_bound_methods and method_name not in self.network_bound_methods:
             return exec_func()
 
         if exec_timeout is None:

aws_advanced_python_wrapper/failover_plugin.py

@@ -351,22 +351,26 @@ def _invalidate_current_connection(self):
         """
         conn = self._plugin_service.current_connection
         if conn is None:
-            return
+            return None
+
+        driver_dialect = self._plugin_service.driver_dialect
 
         if self._plugin_service.is_in_transaction:
             self._plugin_service.update_in_transaction(True)
             try:
+                driver_dialect.execute("Connection.rollback", lambda: conn.rollback())
                 conn.rollback()
             except Exception:
                 pass
 
-        driver_dialect = self._plugin_service.driver_dialect
-        if driver_dialect is not None and not driver_dialect.is_closed(conn):
+        if not driver_dialect.is_closed(conn):
             try:
-                conn.close()
+                return driver_dialect.execute("Connection.close", lambda: conn.close())
             except Exception:
                 pass
 
+        return None
+
     def _invalid_invocation_on_closed_connection(self):
         if not self._closed_explicitly:
             self._is_closed = False

aws_advanced_python_wrapper/pg_driver_dialect.py

@@ -48,6 +48,7 @@ class PgDriverDialect(DriverDialect):
         "Connection.is_read_only",
         "Connection.set_read_only",
         "Connection.rollback",
+        "Connection.close",
         "Connection.cursor",
         "Cursor.close",
         "Cursor.callproc",

aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties

@@ -382,8 +382,6 @@ RoundRobinHostSelector.RoundRobinInvalidHostWeightPairs= [RoundRobinHostSelector
 WeightedRandomHostSelector.WeightedRandomInvalidHostWeightPairs= [WeightedRandomHostSelector] The provided host weight pairs have not been configured correctly. Please ensure the provided host weight pairs is a comma separated list of pairs, each pair in the format of <host>:<weight>. Weight values must be an integer greater than or equal to the default weight value of 1. Weight pair: '{}'
 WeightedRandomHostSelector.WeightedRandomInvalidDefaultWeight=[WeightedRandomHostSelector] The provided default weight value is not valid. Weight values must be an integer greater than or equal to 1.
 
-SlidingExpirationCache.CleaningUp=[SlidingExpirationCache] Cleaning up...
-
 SqlAlchemyPooledConnectionProvider.PoolNone=[SqlAlchemyPooledConnectionProvider] Attempted to find or create a pool for '{}' but the result of the attempt evaluated to None.
 SqlAlchemyPooledConnectionProvider.UnableToCreateDefaultKey=[SqlAlchemyPooledConnectionProvider] Unable to create a default key for internal connection pools. By default, the user parameter is used, but the given user evaluated to None or the empty string (""). Please ensure you have passed a valid user in the connection properties.
 

aws_advanced_python_wrapper/sqlalchemy_driver_dialect.py

@@ -67,7 +67,7 @@ def abort_connection(self, conn: Connection):
         if isinstance(conn, PoolProxiedConnection):
             conn = conn.driver_connection
             if conn is None:
-                return
+                return None
 
         return self._underlying_driver.abort_connection(conn)
 
@@ -122,6 +122,6 @@ def transfer_session_state(self, from_conn: Connection, to_conn: Connection):
             to_driver_conn = to_conn.driver_connection
 
         if from_driver_conn is None or to_driver_conn is None:
-            return
+            return None
 
         return self._underlying_driver.transfer_session_state(from_driver_conn, to_driver_conn)

aws_advanced_python_wrapper/states/session_state.py

@@ -92,3 +92,6 @@ def copy(self):
         new_session_state.readonly = self.readonly.copy()
 
         return new_session_state
+
+    def __str__(self):
+        return f"autocommit: {self.auto_commit}, readonly: {self.readonly}"

aws_advanced_python_wrapper/utils/pg_exception_handler.py

@@ -27,9 +27,19 @@ class PgExceptionHandler(ExceptionHandler):
     _PAM_AUTHENTICATION_FAILED_MSG = "PAM authentication failed"
     _CONNECTION_FAILED = "connection failed"
     _CONSUMING_INPUT_FAILED = "consuming input failed"
+    _CONNECTION_SOCKET_CLOSED = "connection socket closed"
 
-    _NETWORK_ERRORS: List[str]
-    _ACCESS_ERRORS: List[str]
+    _NETWORK_ERROR_MESSAGES: List[str] = [
+        _CONNECTION_FAILED,
+        _CONSUMING_INPUT_FAILED,
+        _CONNECTION_SOCKET_CLOSED
+    ]
+    _ACCESS_ERROR_MESSAGES: List[str] = [
+        _PASSWORD_AUTHENTICATION_FAILED_MSG,
+        _PAM_AUTHENTICATION_FAILED_MSG
+    ]
+    _NETWORK_ERROR_CODES: List[str]
+    _ACCESS_ERROR_CODES: List[str]
 
     def is_network_exception(self, error: Optional[Exception] = None, sql_state: Optional[str] = None) -> bool:
         if isinstance(error, QueryTimeoutError) or isinstance(error, ConnectionTimeout):
@@ -43,15 +53,15 @@ def is_network_exception(self, error: Optional[Exception] = None, sql_state: Opt
                 # getattr may throw an AttributeError if the error does not have a `sqlstate` attribute
                 pass
 
-        if sql_state is not None and sql_state in self._NETWORK_ERRORS:
+        if sql_state is not None and sql_state in self._NETWORK_ERROR_CODES:
             return True
 
         if isinstance(error, OperationalError):
             if len(error.args) == 0:
                 return False
             # Check the error message if this is a generic error
             error_msg: str = error.args[0]
-            return self._CONNECTION_FAILED in error_msg or self._CONSUMING_INPUT_FAILED in error_msg
+            return any(msg in error_msg for msg in self._NETWORK_ERROR_MESSAGES)
 
         return False
 
@@ -63,7 +73,7 @@ def is_login_exception(self, error: Optional[Exception] = None, sql_state: Optio
             if sql_state is None and hasattr(error, "sqlstate") and error.sqlstate is not None:
                 sql_state = error.sqlstate
 
-            if sql_state is not None and sql_state in self._ACCESS_ERRORS:
+            if sql_state is not None and sql_state in self._ACCESS_ERROR_CODES:
                 return True
 
             if isinstance(error, OperationalError):
@@ -72,15 +82,14 @@ def is_login_exception(self, error: Optional[Exception] = None, sql_state: Optio
 
                 # Check the error message if this is a generic error
                 error_msg: str = error.args[0]
-                if self._PASSWORD_AUTHENTICATION_FAILED_MSG in error_msg \
-                        or self._PAM_AUTHENTICATION_FAILED_MSG in error_msg:
+                if any(msg in error_msg for msg in self._ACCESS_ERROR_MESSAGES):
                     return True
 
         return False
 
 
 class SingleAzPgExceptionHandler(PgExceptionHandler):
-    _NETWORK_ERRORS: List[str] = [
+    _NETWORK_ERROR_CODES: List[str] = [
         "53",  # insufficient resources
         "57P01",  # admin shutdown
         "57P02",  # crash shutdown
@@ -92,14 +101,14 @@ class SingleAzPgExceptionHandler(PgExceptionHandler):
         "XX"  # internal error(backend)
     ]
 
-    _ACCESS_ERRORS: List[str] = [
+    _ACCESS_ERROR_CODES: List[str] = [
         "28000",  # PAM authentication errors
         "28P01"
     ]
 
 
 class MultiAzPgExceptionHandler(PgExceptionHandler):
-    _NETWORK_ERRORS: List[str] = [
+    _NETWORK_ERROR_CODES: List[str] = [
         "28000",  # access denied during reboot, this should be considered a temporary failure
         "53",  # insufficient resources
         "57P01",  # admin shutdown
@@ -112,4 +121,4 @@ class MultiAzPgExceptionHandler(PgExceptionHandler):
         "XX"  # internal error(backend)
     ]
 
-    _ACCESS_ERRORS: List[str] = ["28P01"]
+    _ACCESS_ERROR_CODES: List[str] = ["28P01"]

aws_advanced_python_wrapper/utils/sliding_expiration_cache.py

@@ -126,7 +126,6 @@ def _cleanup_thread_internal(self):
         while True:
             try:
                 sleep(self._cleanup_interval_ns / 1_000_000_000)
-                logger.debug("SlidingExpirationCache.CleaningUp")
                 self._cleanup_time_ns.set(perf_counter_ns() + self._cleanup_interval_ns)
                 keys = [key for key, _ in self._cdict.items()]
                 for key in keys:

aws_advanced_python_wrapper/utils/telemetry/xray_telemetry.py

@@ -78,7 +78,7 @@ def post_copy(context: XRayTelemetryContext, trace_level: TelemetryTraceLevel):
         return
 
     if trace_level in [TelemetryTraceLevel.FORCE_TOP_LEVEL, TelemetryTraceLevel.TOP_LEVEL]:
-        with ThreadPoolExecutor() as executor:
+        with ThreadPoolExecutor(thread_name_prefix=context.get_name()) as executor:
             future = executor.submit(_clone_and_close_context, context, trace_level)
             future.result()
     else:

aws_advanced_python_wrapper/writer_failover_handler.py

@@ -120,9 +120,9 @@ def get_result_from_future(self, current_topology: Tuple[HostInfo, ...]) -> Writ
 
             executor = ThreadPoolExecutor(thread_name_prefix="WriterFailoverHandlerExecutor")
             try:
+                futures = [executor.submit(self.reconnect_to_writer, writer_host),
+                           executor.submit(self.wait_for_new_writer, current_topology, writer_host)]
                 try:
-                    futures = [executor.submit(self.reconnect_to_writer, writer_host),
-                               executor.submit(self.wait_for_new_writer, current_topology, writer_host)]
                     for future in as_completed(futures, timeout=self._max_failover_timeout_sec):
                         result = future.result()
                         if result.is_connected:
@@ -132,9 +132,10 @@ def get_result_from_future(self, current_topology: Tuple[HostInfo, ...]) -> Writ
                             return result
                 except TimeoutError:
                     self._timeout_event.set()
-                finally:
-                    self._timeout_event.set()
+                    for future in futures:
+                        future.cancel()
             finally:
+                self._timeout_event.set()
                 executor.shutdown(wait=False)
 
         return WriterFailoverHandlerImpl.failed_writer_failover_result

tests/integration/container/conftest.py

@@ -28,6 +28,8 @@
     DriverDialectManager
 from aws_advanced_python_wrapper.exception_handling import ExceptionManager
 from aws_advanced_python_wrapper.host_list_provider import RdsHostListProvider
+from aws_advanced_python_wrapper.host_monitoring_plugin import \
+    MonitoringThreadContainer
 from aws_advanced_python_wrapper.plugin_service import PluginServiceImpl
 from aws_advanced_python_wrapper.utils.log import Logger
 from aws_advanced_python_wrapper.utils.rdsutils import RdsUtils
@@ -67,6 +69,8 @@ def pytest_runtest_setup(item):
     else:
         TestEnvironment.get_current().set_current_driver(None)
 
+    logger.info("Starting test preparation for: " + test_name)
+
     segment: Optional[Segment] = None
     if TestEnvironmentFeatures.TELEMETRY_TRACES_ENABLED in TestEnvironment.get_current().get_features():
         segment = xray_recorder.begin_segment("test: setup")
@@ -92,7 +96,7 @@ def pytest_runtest_setup(item):
         # Wait up to 5min
         instances: List[str] = list()
         start_time = timeit.default_timer()
-        while (len(instances) != request.get_num_of_instances()
+        while (len(instances) < request.get_num_of_instances()
                or len(instances) == 0
                or not rds_utility.is_db_instance_writer(instances[0])) and (
                 timeit.default_timer() - start_time) < 300:  # 5 min
@@ -135,6 +139,7 @@ def pytest_runtest_setup(item):
         DatabaseDialectManager._known_endpoint_dialects.clear()
         CustomEndpointPlugin._monitors.clear()
         CustomEndpointMonitor._custom_endpoint_info_cache.clear()
+        MonitoringThreadContainer.clean_up()
 
         ConnectionProviderManager.reset_provider()
         DatabaseDialectManager.reset_custom_dialect()