chore: use cancel_safe from psycopg's v3.3.2

Author: karenc-bq

.github/workflows/integration_tests.yml

@@ -1,40 +1,105 @@
 name: Integration Tests
 
 on:
-  workflow_dispatch:
   push:
     branches:
       - main
+      - chore/run-integration-tests-on-codebuild
+    paths-ignore:
+      - '**/*.md'
+      - '**/*.jpg'
+      - '**/README.txt'
+      - '**/LICENSE.txt'
+      - 'docs/**'
+      - 'ISSUE_TEMPLATE/**'
+      - '**/remove-old-artifacts.yml'
+  pull_request_target:
+    branches:
+      - main
+    paths-ignore:
+      - '**/*.md'
+      - '**/*.jpg'
+      - '**/README.txt'
+      - '**/LICENSE.txt'
+      - 'docs/**'
+      - 'ISSUE_TEMPLATE/**'
+      - '**/remove-old-artifacts.yml'
+  workflow_dispatch:
+    inputs:
+      ref:
+        description: 'Git ref (branch, tag, or SHA) to test'
+        required: false
+        type: string
+      repository:
+        description: 'Repository to test in the form owner/repo (use for fork branches)'
+        required: false
+        type: string
 
 permissions:
   id-token: write   # This is required for requesting the JWT
   contents: read    # This is required for actions/checkout
 
 jobs:
+  approve:
+    # Auto-approve for non-fork scenarios (push, manual dispatch, or PR from the same repo)
+    if: >
+      github.event_name == 'push' ||
+      github.event_name == 'workflow_dispatch' ||
+      github.event.pull_request.head.repo.full_name == github.repository
+    runs-on: ubuntu-latest
+    steps:
+      - run: echo "Approved - not a fork PR"
+
+  approve-fork:
+    # Require manual maintainer approval (via the integration-tests environment) for fork PRs
+    if: >
+      github.event_name == 'pull_request_target' &&
+      github.event.pull_request.head.repo.full_name != github.repository
+    runs-on: ubuntu-latest
+    environment: integration-tests
+    steps:
+      - run: echo "Fork PR approved by maintainer"
+
   lts-integration-tests:
     name: Run LTS Integration Tests
+    needs: [ approve, approve-fork ]
+    if: |
+      always() &&
+      (needs.approve.result == 'success' || needs.approve-fork.result == 'success') &&
+      !(needs.approve.result == 'failure' || needs.approve-fork.result == 'failure')
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
-        python-version: [ "3.11", "3.12", "3.13"  ]
-        environment: [ "mysql", "pg" ]
+        python-version: [ "3.11" ]
+        environment: [ "pg" ]
 
     steps:
       - name: 'Clone repository'
         uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6
+        with:
+          ref: ${{ github.event.pull_request.head.sha || github.sha }}
 
       - name: 'Set up JDK 8'
         uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654  # v5
         with:
           distribution: 'corretto'
           java-version: 8
 
+      - name: 'Set up Python ${{ matrix.python-version }}'
+        uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405  # v6.2.0
+        with:
+          python-version: ${{ matrix.python-version }}
+
       - name: Install poetry
         shell: bash
         run: |
-          pipx install poetry==1.8.2
-          poetry config virtualenvs.prefer-active-python true
+          python -m pip install --upgrade pip
+          python -m pip install pipx
+          python -m pipx install poetry==1.8.2
+          PIPX_BIN_DIR=$(python -m pipx environment --value PIPX_BIN_DIR)
+          echo "$PIPX_BIN_DIR" >> "$GITHUB_PATH"
+          "$PIPX_BIN_DIR/poetry" config virtualenvs.prefer-active-python true
 
       - name: Install dependencies
         run: poetry install
@@ -71,24 +136,35 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [ "3.11", "3.12", "3.13" ]
-        environment: [ "mysql", "pg" ]
+        python-version: [ "3.11" ]
+        environment: [ "pg" ]
 
     steps:
       - name: 'Clone repository'
         uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6
+        with:
+          ref: ${{ github.event.pull_request.head.sha || github.sha }}
 
       - name: 'Set up JDK 8'
         uses: actions/setup-java@be666c2fcd27ec809703dec50e508c2fdc7f6654  # v5
         with:
           distribution: 'corretto'
           java-version: 8
 
+      - name: 'Set up Python ${{ matrix.python-version }}'
+        uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405  # v6.2.0
+        with:
+          python-version: ${{ matrix.python-version }}
+
       - name: Install poetry
         shell: bash
         run: |
-          pipx install poetry==1.8.2
-          poetry config virtualenvs.prefer-active-python true
+          python -m pip install --upgrade pip
+          python -m pip install pipx
+          python -m pipx install poetry==1.8.2
+          PIPX_BIN_DIR=$(python -m pipx environment --value PIPX_BIN_DIR)
+          echo "$PIPX_BIN_DIR" >> "$GITHUB_PATH"
+          "$PIPX_BIN_DIR/poetry" config virtualenvs.prefer-active-python true
 
       - name: Install dependencies
         run: poetry install

aws_advanced_python_wrapper/aurora_connection_tracker_plugin.py

@@ -21,6 +21,7 @@
                     Optional, Set)
 
 from aws_advanced_python_wrapper.utils.notifications import HostEvent
+from aws_advanced_python_wrapper.utils import services_container
 from aws_advanced_python_wrapper.utils.utils import Utils
 
 if TYPE_CHECKING:
@@ -52,6 +53,7 @@ class OpenedConnectionTracker:
     _shutdown_event: ClassVar[threading.Event] = threading.Event()
     _safe_to_check_closed_classes: ClassVar[Set[str]] = {"psycopg"}
     _default_sleep_time: ClassVar[int] = 30
+    _invalidate_executor_name: ClassVar[str] = "OpenedConnectionTrackerInvalidate"
 
     @classmethod
     def _start_prune_thread(cls):
@@ -166,9 +168,12 @@ def invalidate_all_connections(self, host_info: Optional[HostInfo] = None, host:
         with self._lock:
             connection_set: Optional[WeakSet] = self._opened_connections.get(instance_endpoint)
             connections_list = list(connection_set) if connection_set is not None else None
+            if connection_set is not None:
+                connection_set.clear()
+                self._opened_connections.pop(instance_endpoint, None)
 
-        if connections_list is not None:
-            self._log_connection_set(instance_endpoint, connection_set)
+        if connections_list:
+            self._log_connection_set(instance_endpoint, connections_list)
             self._invalidate_connections(connections_list)
 
     def remove_connection_tracking(self, host_info: HostInfo, connection: Connection | None):
@@ -211,9 +216,8 @@ def _task(connections_list: list):
                 pass
 
     def _invalidate_connections(self, connections_list: list):
-        invalidate_connection_thread: Thread = Thread(daemon=True, target=self._task,
-                                                      args=[connections_list])  # type: ignore
-        invalidate_connection_thread.start()
+        executor = services_container.get_thread_pool(self._invalidate_executor_name, drain_first=True)
+        executor.submit(self._task, connections_list)
 
     def log_opened_connections(self):
         with self._lock:

aws_advanced_python_wrapper/host_list_provider.py

@@ -464,7 +464,12 @@ def __init__(self, dialect: db_dialect.TopologyAwareDatabaseDialect, props: Prop
 
         self.instance_template: HostInfo = instance_template
         self._max_timeout_sec = WrapperProperties.AUXILIARY_QUERY_TIMEOUT_SEC.get_int(props)
-        self._thread_pool = services_container.get_thread_pool(self._executor_name)
+        # Need to drain topology query pools before monitor services shut down.
+        self._thread_pool = services_container.get_thread_pool(self._executor_name, drain_first=True)
+
+    @classmethod
+    def release_resources(cls, wait: bool = True) -> bool:
+        return services_container.release_thread_pool(cls._executor_name, wait=wait)
 
     def _validate_host_pattern(self, host: str):
         if not self._rds_utils.is_dns_pattern_valid(host):
@@ -495,8 +500,9 @@ def query_for_topology(
         :return: a tuple of :py:class:`HostInfo` objects representing the database topology. If the query results did not include a writer instance,
         an empty tuple will be returned.
         """
+        thread_pool = services_container.get_thread_pool(self._executor_name, drain_first=True)
         query_for_topology_func_with_timeout = preserve_transaction_status_with_timeout(
-                    self._thread_pool, self._max_timeout_sec, driver_dialect, conn)(self._query_for_topology)
+                    thread_pool, self._max_timeout_sec, driver_dialect, conn)(self._query_for_topology)
         x = query_for_topology_func_with_timeout(conn)
         return x
 

aws_advanced_python_wrapper/pg_driver_dialect.py

@@ -73,6 +73,11 @@ def is_closed(self, conn: Connection) -> bool:
 
     def abort_connection(self, conn: Connection):
         if isinstance(conn, psycopg.Connection):
+            try:
+                conn.cancel_safe()
+            except Exception:
+                # Best effort cancel
+                pass
             conn.close()
             return
         raise UnsupportedOperationError(

aws_advanced_python_wrapper/utils/services_container.py

@@ -17,7 +17,7 @@
 import threading
 from concurrent.futures import ThreadPoolExecutor
 from datetime import timedelta
-from typing import Dict, Optional
+from typing import Dict, Optional, Set
 
 from aws_advanced_python_wrapper.allowed_and_blocked_hosts import \
     AllowedAndBlockedHosts
@@ -37,6 +37,10 @@ def __init__(self) -> None:
         self._storage_service: Optional[StorageService] = None
         self._monitor_service: Optional[MonitorService] = None
         self._thread_pools: Dict[str, ThreadPoolExecutor] = {}
+        # Some service pools must be drained BEFORE the monitor service is released and connections are closed.
+        # This prevents worker threads like the topology util threads from continuing to using connections
+        # after they are closed, causing segfaults.
+        self._drain_first_pools: Set[str] = set()
         self._lock = threading.Lock()
 
     def _ensure_initialized(self) -> None:
@@ -63,19 +67,24 @@ def monitor_service(self) -> MonitorService:
         self._ensure_initialized()
         return self._monitor_service  # type: ignore
 
-    def get_thread_pool(self, name: str, max_workers: Optional[int] = None) -> ThreadPoolExecutor:
+    def get_thread_pool(self, name: str, max_workers: Optional[int] = None, drain_first: bool = False) -> ThreadPoolExecutor:
         pool = self._thread_pools.get(name)
         if pool is not None:
+            if drain_first:
+                self._drain_first_pools.add(name)
             return pool
         with self._lock:
             if name not in self._thread_pools:
                 self._thread_pools[name] = ThreadPoolExecutor(
                     max_workers=max_workers, thread_name_prefix=name)
+            if drain_first:
+                self._drain_first_pools.add(name)
             return self._thread_pools[name]
 
     def release_thread_pool(self, name: str, wait: bool = True) -> bool:
         with self._lock:
             pool = self._thread_pools.pop(name, None)
+            self._drain_first_pools.discard(name)
             if pool is not None:
                 try:
                     pool.shutdown(wait=wait)
@@ -85,6 +94,18 @@ def release_thread_pool(self, name: str, wait: bool = True) -> bool:
             return False
 
     def release_resources(self) -> None:
+        # Some thread pools need to be drained first before shutting down the monitor services.
+        # This prevents segfaults when monitor services shut down and close all the active monitoring connections.
+        with self._lock:
+            drain_names = list(self._drain_first_pools)
+        for name in drain_names:
+            pool = self._thread_pools.get(name)
+            if pool is not None:
+                try:
+                    pool.shutdown(wait=True)
+                except Exception as e:
+                    logger.debug("CoreServices.ErrorShuttingDownPool", name, e)
+
         if self._monitor_service is not None:
             try:
                 self._monitor_service.release_resources()
@@ -114,6 +135,7 @@ def release_resources(self) -> None:
                 except Exception as e:
                     logger.debug("CoreServices.ErrorShuttingDownPool", name, e)
             self._thread_pools.clear()
+            self._drain_first_pools.clear()
 
 
 _instance = _ServicesContainer()
@@ -132,8 +154,8 @@ def get_monitor_service() -> MonitorService:
     return _instance.monitor_service
 
 
-def get_thread_pool(name: str, max_workers: Optional[int] = None) -> ThreadPoolExecutor:
-    return _instance.get_thread_pool(name, max_workers)
+def get_thread_pool(name: str, max_workers: Optional[int] = None, drain_first: bool = False) -> ThreadPoolExecutor:
+    return _instance.get_thread_pool(name, max_workers, drain_first)
 
 
 def release_thread_pool(name: str, wait: bool = True) -> bool:

poetry.lock

@@ -41,8 +41,8 @@ flake8-type-checking = ">=2.9,<4.0"
 isort = "^5.13.2"
 pep8-naming = ">=0.14.1,<0.16.0"
 SQLAlchemy = "^2.0.30"
-psycopg = "^3.3.1"
-psycopg-binary = "^3.3.1"
+psycopg = "^3.3.2"
+psycopg-binary = "^3.3.2"
 mysql-connector-python = "^9.5.0"
 django = "^5.0"
 django-stubs = "^5.2.8"
@@ -60,8 +60,8 @@ pytest-html-merger = ">=0.0.10,<0.1.1"
 toxiproxy-python = "^0.1.1"
 parameterized = "^0.9.0"
 tabulate = ">=0.9,<0.11"
-psycopg = "^3.3.1"
-psycopg-binary = "^3.3.1"
+psycopg = "^3.3.2"
+psycopg-binary = "^3.3.2"
 mysql-connector-python = "^9.5.0"
 opentelemetry-exporter-otlp = "^1.22.0"
 opentelemetry-exporter-otlp-proto-grpc = "^1.22.0"

pyproject.toml

@@ -75,6 +75,7 @@ def props(self):
             "connect_timeout": 10,
             "monitoring-connect_timeout": 5,
             "monitoring-socket_timeout": 5,
+            "auxiliary_query_timeout_sec": 3,
             "autocommit": True,
             "cluster_id": "cluster1"
         })

tests/integration/container/test_aurora_failover.py

@@ -51,7 +51,7 @@ def rds_utils(self):
 
     @pytest.fixture
     def props(self):
-        p: Properties = Properties({"plugins": "read_write_splitting", "connect_timeout": 10, "autocommit": True})
+        p: Properties = Properties({"plugins": "read_write_splitting", "connect_timeout": 10, "auxiliary_query_timeout_sec": 3, "autocommit": True})
 
         if TestEnvironmentFeatures.TELEMETRY_TRACES_ENABLED in TestEnvironment.get_current().get_features() \
                 or TestEnvironmentFeatures.TELEMETRY_METRICS_ENABLED in TestEnvironment.get_current().get_features():

tests/integration/container/test_autoscaling.py

@@ -221,6 +221,7 @@ def test_failover_with_secrets_manager(
             "connect_timeout": 10,
             "monitoring-connect_timeout": 5,
             "monitoring-socket_timeout": 5,
+            "auxiliary_query_timeout_sec": 3,
             "topology_refresh_ms": 10,
         })