From b65dec234b51410738c293028ae0c5cccf568b8b Mon Sep 17 00:00:00 2001 From: karezche <64801825+karenc-bq@users.noreply.github.com> Date: Wed, 11 Nov 2026 03:24:39 -0700 Subject: [PATCH] feat: connection lifetime plugin --- CHANGELOG.md | 1 + README.md | 2 + .../aio/connection_lifetime_plugin.py | 302 +++++++++++++++ .../aio/plugin_factory.py | 10 + .../connection_lifetime_plugin.py | 360 ++++++++++++++++++ aws_advanced_python_wrapper/errors.py | 25 ++ aws_advanced_python_wrapper/pep249.py | 4 +- aws_advanced_python_wrapper/plugin_service.py | 4 + ...dvanced_python_wrapper_messages.properties | 13 + .../sqlalchemy_dialects/mysql.py | 8 +- .../sqlalchemy_dialects/mysql_async.py | 8 +- .../sqlalchemy_dialects/pg.py | 8 +- .../sqlalchemy_dialects/pg_async.py | 8 +- .../utils/properties.py | 13 + docs/README.md | 1 + .../PluginChainCompatibility.md | 6 + .../SqlAlchemySupport.md | 1 + .../UsingThePythonWrapper.md | 1 + tests/unit/test_aio_plugin_factory.py | 1 + tests/unit/test_aio_sqlalchemy_dialect.py | 24 ++ tests/unit/test_sqlalchemy_dialects.py | 16 + 21 files changed, 806 insertions(+), 10 deletions(-) create mode 100644 aws_advanced_python_wrapper/aio/connection_lifetime_plugin.py create mode 100644 aws_advanced_python_wrapper/connection_lifetime_plugin.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 03eec8a1..f1442886 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ## [Unreleased] ### :magic_wand: Added +* The `connection_lifetime` plugin, bounding how long a physical connection is used. Once `connection_lifetime_ms` has elapsed, the next network-bound call either closes the connection or reconnects through the host the application originally supplied, depending on `connection_lifetime_action`. Expiration is deferred while a transaction is open, and reports through `ConnectionLifetimeSuccessError`, `ConnectionLifetimeExpiredError` or `ConnectionLifetimeFailedError` so a connection pool can tell a usable connection from a dead one. Available in both the synchronous and asynchronous wrapper. See [Connection Lifetime Plugin](./docs/using-the-python-wrapper/using-plugins/UsingTheConnectionLifetimePlugin.md). * The `failover` and `failover_v2` plugins currently rewraps all errors to an AwsWrapperError even when the errors do not trigger failover. Set the `preserve_driver_exception_type` configuration parameter to `True` to re-raise the target driver errors as-is for the failover plugins. Failover errors are unaffected and are still raised as `FailoverSuccessError`, `FailoverFailedError` or `TransactionResolutionUnknownError`. This is useful for ORMs to properly classify the errors. ([Issue #1275](https://github.com/aws/aws-advanced-python-wrapper/issues/1275)) ## [3.1.0] - 2026-08-24 diff --git a/README.md b/README.md index 2587072e..0307a904 100644 --- a/README.md +++ b/README.md @@ -107,6 +107,8 @@ The following table lists the connection properties used with the AWS Advanced P | `failover_reader_connect_timeout_sec` | [Failover Plugin](docs/using-the-python-wrapper/using-plugins/UsingTheFailoverPlugin.md) | | `failover_timeout_sec` | [Failover Plugin](docs/using-the-python-wrapper/using-plugins/UsingTheFailoverPlugin.md) | | `failover_writer_reconnect_interval_sec` | [Failover Plugin](docs/using-the-python-wrapper/using-plugins/UsingTheFailoverPlugin.md) | +| `connection_lifetime_ms` | [Connection Lifetime Plugin](docs/using-the-python-wrapper/using-plugins/UsingTheConnectionLifetimePlugin.md) | +| `connection_lifetime_action` | [Connection Lifetime Plugin](docs/using-the-python-wrapper/using-plugins/UsingTheConnectionLifetimePlugin.md) | | `failure_detection_count` | [Host Monitoring Plugin](docs/using-the-python-wrapper/using-plugins/UsingTheHostMonitoringPlugin.md) | | `failure_detection_enabled` | [Host Monitoring Plugin](docs/using-the-python-wrapper/using-plugins/UsingTheHostMonitoringPlugin.md) | | `failure_detection_interval_ms` | [Host Monitoring Plugin](docs/using-the-python-wrapper/using-plugins/UsingTheHostMonitoringPlugin.md) | diff --git a/aws_advanced_python_wrapper/aio/connection_lifetime_plugin.py b/aws_advanced_python_wrapper/aio/connection_lifetime_plugin.py new file mode 100644 index 00000000..c3090c6d --- /dev/null +++ b/aws_advanced_python_wrapper/aio/connection_lifetime_plugin.py @@ -0,0 +1,302 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Async counterpart of :class:`ConnectionLifetimePlugin`. + +Same contract as the sync plugin -- stamp each physical connection, and on +the first network-bound call after ``connection_lifetime_ms`` has elapsed +either close the connection or replace it with one on the host the +application originally asked for. Property parsing and the action enum are +shared with the sync module so the two cannot disagree on what a +configuration value means. +""" + +from __future__ import annotations + +import asyncio +from time import monotonic_ns +from typing import TYPE_CHECKING, Any, Callable, Optional, Set +from weakref import WeakKeyDictionary + +from aws_advanced_python_wrapper.aio.plugin import AsyncPlugin +from aws_advanced_python_wrapper.connection_lifetime_plugin import ( + LIFECYCLE_METHODS, ConnectionLifetimeAction, + parse_connection_lifetime_action, parse_connection_lifetime_ns) +from aws_advanced_python_wrapper.errors import (AwsWrapperError, + ConnectionLifetimeExpiredError, + ConnectionLifetimeFailedError, + ConnectionLifetimeSuccessError) +from aws_advanced_python_wrapper.pep249_methods import DbApiMethod +from aws_advanced_python_wrapper.utils.log import Logger +from aws_advanced_python_wrapper.utils.messages import Messages +from aws_advanced_python_wrapper.utils.notifications import ( + ConnectionEvent, OldConnectionSuggestedAction) + +if TYPE_CHECKING: + from aws_advanced_python_wrapper.aio.driver_dialect.base import \ + AsyncDriverDialect + from aws_advanced_python_wrapper.aio.plugin_service import \ + AsyncPluginService + from aws_advanced_python_wrapper.hostinfo import HostInfo + from aws_advanced_python_wrapper.utils.properties import Properties + +logger = Logger(__name__) + + +class AsyncConnectionLifetimePlugin(AsyncPlugin): + def __init__( + self, + plugin_service: AsyncPluginService, + props: Properties, + clock_ns: Callable[[], int] = monotonic_ns): + self._plugin_service = plugin_service + self._props = props + # Injected only by the tests, which need to move time without sleeping. + self._clock_ns = clock_ns + self._connection_lifetime_ns = parse_connection_lifetime_ns(props) + self._action = parse_connection_lifetime_action(props) + + subscribed = set(plugin_service.network_bound_methods) + subscribed.update({ + DbApiMethod.CONNECT.method_name, + DbApiMethod.FORCE_CONNECT.method_name, + DbApiMethod.NOTIFY_CONNECTION_CHANGED.method_name, + }) + self._subscribed_methods = subscribed + + # See the sync plugin for why the keys are weak. + self._connection_creation_times: WeakKeyDictionary = WeakKeyDictionary() + self._original_host_info: Optional[HostInfo] = None + self._timed_connection: Optional[Any] = None + self._connection_creation_time_ns: Optional[int] = None + self._terminal = False + + @property + def subscribed_methods(self) -> Set[str]: + return self._subscribed_methods + + async def connect( + self, + target_driver_func: Callable, + driver_dialect: AsyncDriverDialect, + host_info: HostInfo, + props: Properties, + is_initial_connection: bool, + connect_func: Callable) -> Any: + return await self._connect(host_info, is_initial_connection, connect_func) + + async def force_connect( + self, + target_driver_func: Callable, + driver_dialect: AsyncDriverDialect, + host_info: HostInfo, + props: Properties, + is_initial_connection: bool, + force_connect_func: Callable) -> Any: + return await self._connect(host_info, is_initial_connection, force_connect_func) + + async def _connect( + self, + host_info: HostInfo, + is_initial_connection: bool, + connect_func: Callable) -> Any: + # Sync parity (connection_lifetime_plugin.py): the reconnect host is + # the application's own, captured from the first initial connect + # because the plugin service does not yet carry it during + # init_host_provider. + if is_initial_connection and self._original_host_info is None: + self._original_host_info = host_info + + connection = await connect_func() + self._record_connection_creation_time(connection) + return connection + + def _record_connection_creation_time(self, connection: Optional[Any]) -> None: + if connection is None: + return + try: + self._connection_creation_times[connection] = self._clock_ns() + except TypeError: + # Not weak-referenceable; stamped on first use instead. + pass + + def notify_connection_changed( + self, changes: Set[ConnectionEvent]) -> Optional[OldConnectionSuggestedAction]: + if (ConnectionEvent.INITIAL_CONNECTION in changes + or ConnectionEvent.CONNECTION_OBJECT_CHANGED in changes): + self._select_current_connection_timestamp() + return OldConnectionSuggestedAction.NO_OPINION + + async def execute( + self, + target: object, + method_name: str, + execute_func: Callable, + *args: Any, + **kwargs: Any) -> Any: + if (method_name in LIFECYCLE_METHODS + or not self._plugin_service.is_network_bound_method(method_name)): + return await execute_func() + + if self._terminal: + raise ConnectionLifetimeExpiredError( + Messages.get("ConnectionLifetimePlugin.ConnectionClosed")) + + if not self._is_current_connection_expired(): + return await execute_func() + + # An open transaction finishes on the connection it started on, so + # commit and rollback are not cut off mid-way. Deferring does not + # extend the lifetime. + if self._plugin_service.is_in_transaction: + logger.debug("ConnectionLifetimePlugin.ActionDeferred", method_name) + return await execute_func() + + if self._action == ConnectionLifetimeAction.CLOSE: + await self._close_current_connection() + else: + await self._reconnect() + + def _select_current_connection_timestamp(self) -> None: + current_connection = self._plugin_service.current_connection + creation_time = None + if current_connection is not None: + creation_time = self._connection_creation_times.get(current_connection) + if creation_time is None: + creation_time = self._clock_ns() + self._record_connection_creation_time(current_connection) + + self._timed_connection = current_connection + self._connection_creation_time_ns = creation_time + + def _is_current_connection_expired(self) -> bool: + current_connection = self._plugin_service.current_connection + if current_connection is not self._timed_connection or self._connection_creation_time_ns is None: + self._select_current_connection_timestamp() + + creation_time = self._connection_creation_time_ns + # Inclusive boundary: a 1000 ms lifetime expires at exactly 1000 ms. + return (creation_time is not None + and self._clock_ns() - creation_time >= self._connection_lifetime_ns) + + async def _close_current_connection(self) -> None: + self._terminal = True + current_connection = self._plugin_service.current_connection + close_failure: Optional[Exception] = None + + if current_connection is not None: + self._connection_creation_times.pop(current_connection, None) + try: + await self._close(current_connection) + except Exception as e: + close_failure = e + logger.warning("ConnectionLifetimePlugin.CloseConnectionFailed", e) + + error = ConnectionLifetimeExpiredError( + Messages.get("ConnectionLifetimePlugin.ConnectionClosed")) + if close_failure is not None: + raise error from close_failure + raise error + + async def _reconnect(self) -> None: + previous_connection = self._plugin_service.current_connection + replacement_connection = None + try: + reconnect_host_info = self._original_host_info + if reconnect_host_info is None: + reconnect_host_info = self._plugin_service.initial_connection_host_info + if reconnect_host_info is None: + raise AwsWrapperError( + Messages.get("ConnectionLifetimePlugin.OriginalHostUnavailable")) + + logger.debug("ConnectionLifetimePlugin.Reconnecting", reconnect_host_info.url) + replacement_connection = await self._plugin_service.connect( + reconnect_host_info, self._props, self) + if replacement_connection is None: + raise AwsWrapperError( + Messages.get("ConnectionLifetimePlugin.ReconnectReturnedNone")) + replacement_creation_time = self._clock_ns() + + await self._plugin_service.set_current_connection( + replacement_connection, reconnect_host_info) + await self._forget_closed_connection(previous_connection) + self._connection_creation_times[replacement_connection] = replacement_creation_time + self._timed_connection = replacement_connection + self._connection_creation_time_ns = replacement_creation_time + except Exception as e: + await self._cleanup_failed_replacement(replacement_connection, previous_connection) + raise ConnectionLifetimeFailedError( + Messages.get("ConnectionLifetimePlugin.UnableToReconnect")) from e + + raise ConnectionLifetimeSuccessError( + Messages.get("ConnectionLifetimePlugin.ConnectionChanged")) + + async def _forget_closed_connection(self, connection: Optional[Any]) -> None: + if connection is None: + return + + try: + closed = await self._plugin_service.driver_dialect.is_closed(connection) + except Exception: + # A connection that cannot report its state is unusable; do not retain it. + closed = True + + if closed: + self._connection_creation_times.pop(connection, None) + + async def _cleanup_failed_replacement( + self, + replacement_connection: Optional[Any], + previous_connection: Optional[Any]) -> None: + if (replacement_connection is not None + and replacement_connection is self._plugin_service.current_connection + and replacement_connection is not previous_connection): + # The replacement was installed and the failure came after, so + # neither connection can be relied on from here. + self._terminal = True + await self._close_after_failed_installation(replacement_connection) + await self._close_after_failed_installation(previous_connection) + return + + await self._close_unattached_replacement(replacement_connection) + + async def _close_after_failed_installation(self, connection: Optional[Any]) -> None: + if connection is None: + return + + self._connection_creation_times.pop(connection, None) + try: + await self._close(connection) + except Exception as e: + logger.warning("ConnectionLifetimePlugin.FailedInstallationConnectionCloseFailed", e) + + async def _close_unattached_replacement(self, replacement_connection: Optional[Any]) -> None: + if (replacement_connection is None + or replacement_connection is self._plugin_service.current_connection): + return + + self._connection_creation_times.pop(replacement_connection, None) + try: + await self._close(replacement_connection) + except Exception as e: + logger.warning("ConnectionLifetimePlugin.ReplacementCloseFailed", e) + + @staticmethod + async def _close(connection: Any) -> None: + # aiomysql's close() is a coroutine, psycopg's async close() is too, + # but a mock or a sync driver connection returns None -- mirrors + # AsyncPluginService's own old-connection disposal. + result = connection.close() + if asyncio.iscoroutine(result): + await result diff --git a/aws_advanced_python_wrapper/aio/plugin_factory.py b/aws_advanced_python_wrapper/aio/plugin_factory.py index 41c16df9..ddcc448a 100644 --- a/aws_advanced_python_wrapper/aio/plugin_factory.py +++ b/aws_advanced_python_wrapper/aio/plugin_factory.py @@ -178,6 +178,14 @@ def get_instance( return AsyncAuroraConnectionTrackerPlugin(plugin_service) +class _ConnectionLifetimeFactory: + def get_instance( + self, plugin_service, props, host_list_provider=None): + from aws_advanced_python_wrapper.aio.connection_lifetime_plugin import \ + AsyncConnectionLifetimePlugin + return AsyncConnectionLifetimePlugin(plugin_service, props) + + class _ConnectTimeFactory: def get_instance( self, plugin_service, props, host_list_provider=None): @@ -275,6 +283,7 @@ def get_instance( "federated_auth": _FederatedAuthFactory(), "okta": _OktaAuthFactory(), "aurora_connection_tracker": _AuroraConnectionTrackerFactory(), + "connection_lifetime": _ConnectionLifetimeFactory(), "connect_time": _ConnectTimeFactory(), "execute_time": _ExecuteTimeFactory(), "dev": _DeveloperFactory(), @@ -297,6 +306,7 @@ def get_instance( # (the v1/v2 sync codes collapse onto one async factory, which takes the # v1 slot: failover 400, host_monitoring 500). PLUGIN_FACTORY_WEIGHTS: Dict[Type[Any], int] = { + _ConnectionLifetimeFactory: 30, _CustomEndpointFactory: 40, _InitialConnectionFactory: 50, _AuroraConnectionTrackerFactory: 100, diff --git a/aws_advanced_python_wrapper/connection_lifetime_plugin.py b/aws_advanced_python_wrapper/connection_lifetime_plugin.py new file mode 100644 index 00000000..9b2be801 --- /dev/null +++ b/aws_advanced_python_wrapper/connection_lifetime_plugin.py @@ -0,0 +1,360 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). +# You may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Bounds how long the wrapper keeps a physical connection. + +The plugin stamps every physical connection it sees created, and before each +network-bound call checks whether the current one has outlived +``connection_lifetime_ms``. An expired connection is either closed or replaced +by a fresh one on the host the application originally asked for. There is no +background timer: an expired connection stays open until the next +network-bound call. +""" + +from __future__ import annotations + +from enum import Enum +from time import monotonic_ns +from typing import TYPE_CHECKING, Any, Callable, Optional, Set +from weakref import WeakKeyDictionary + +from aws_advanced_python_wrapper.errors import (AwsWrapperError, + ConnectionLifetimeExpiredError, + ConnectionLifetimeFailedError, + ConnectionLifetimeSuccessError) +from aws_advanced_python_wrapper.pep249_methods import DbApiMethod +from aws_advanced_python_wrapper.plugin import Plugin, PluginFactory +from aws_advanced_python_wrapper.utils.log import Logger +from aws_advanced_python_wrapper.utils.messages import Messages +from aws_advanced_python_wrapper.utils.notifications import ( + ConnectionEvent, OldConnectionSuggestedAction) +from aws_advanced_python_wrapper.utils.properties import (Properties, + WrapperProperties) + +if TYPE_CHECKING: + from aws_advanced_python_wrapper.driver_dialect import DriverDialect + from aws_advanced_python_wrapper.hostinfo import HostInfo + from aws_advanced_python_wrapper.pep249 import Connection + from aws_advanced_python_wrapper.plugin_service import PluginService + +logger = Logger(__name__) + +# Teardown calls never trigger the lifetime action, even though the driver +# dialects count them as network-bound. Closing an expired connection has to +# keep working -- raising here would mean a connection could not be disposed +# of, and would surface out of a ``with`` block's exit in place of whatever +# the body raised. +LIFECYCLE_METHODS = frozenset({ + DbApiMethod.CONNECTION_CLOSE.method_name, + DbApiMethod.CONNECTION_IS_CLOSED.method_name, + DbApiMethod.CURSOR_CLOSE.method_name, +}) + + +class ConnectionLifetimeAction(Enum): + CLOSE = "close" + RECONNECT = "reconnect" + + +def parse_connection_lifetime_ns(props: Properties) -> int: + """Read ``connection_lifetime_ms`` as nanoseconds, or raise. + + The property is required and has no default, so an absent, non-numeric or + non-positive value is a configuration error worth failing the connect on + rather than silently disabling the plugin. + """ + raw = WrapperProperties.CONNECTION_LIFETIME_MS.get(props) + try: + lifetime_ms = int(raw) # type: ignore[arg-type] + except (TypeError, ValueError): + raise AwsWrapperError( + Messages.get_formatted("ConnectionLifetimePlugin.InvalidConnectionLifetime", raw)) + + if lifetime_ms < 1: + raise AwsWrapperError( + Messages.get_formatted("ConnectionLifetimePlugin.InvalidConnectionLifetime", raw)) + + return lifetime_ms * 1_000_000 + + +def parse_connection_lifetime_action(props: Properties) -> ConnectionLifetimeAction: + """Read ``connection_lifetime_action``, ignoring surrounding space and case.""" + configured = WrapperProperties.CONNECTION_LIFETIME_ACTION.get(props) + try: + return ConnectionLifetimeAction(str(configured).strip().lower()) + except ValueError: + raise AwsWrapperError( + Messages.get_formatted( + "ConnectionLifetimePlugin.InvalidConnectionLifetimeAction", configured)) + + +class ConnectionLifetimePlugin(Plugin): + def __init__( + self, + plugin_service: PluginService, + props: Properties, + clock_ns: Callable[[], int] = monotonic_ns): + self._plugin_service = plugin_service + self._props = props + # Injected only by the tests, which need to move time without sleeping. + self._clock_ns = clock_ns + self._connection_lifetime_ns = parse_connection_lifetime_ns(props) + self._action = parse_connection_lifetime_action(props) + + subscribed = set(plugin_service.network_bound_methods) + subscribed.update({ + DbApiMethod.CONNECT.method_name, + DbApiMethod.FORCE_CONNECT.method_name, + DbApiMethod.NOTIFY_CONNECTION_CHANGED.method_name, + }) + self._subscribed_methods = subscribed + + # Keyed on the connection object, so two physical connections that + # compare equal are still timed separately. Weak keys keep a + # connection the pipeline opened and then discarded from being held + # alive by this map; a connection that cannot be weak-referenced is + # simply not pre-stamped and gets its timestamp on first use instead. + self._connection_creation_times: WeakKeyDictionary = WeakKeyDictionary() + self._original_host_info: Optional[HostInfo] = None + self._timed_connection: Optional[Connection] = None + self._connection_creation_time_ns: Optional[int] = None + self._terminal = False + + @property + def subscribed_methods(self) -> Set[str]: + return self._subscribed_methods + + def connect( + self, + target_driver_func: Callable, + driver_dialect: DriverDialect, + host_info: HostInfo, + props: Properties, + is_initial_connection: bool, + connect_func: Callable) -> Connection: + return self._connect(host_info, is_initial_connection, connect_func) + + def force_connect( + self, + target_driver_func: Callable, + driver_dialect: DriverDialect, + host_info: HostInfo, + props: Properties, + is_initial_connection: bool, + force_connect_func: Callable) -> Connection: + return self._connect(host_info, is_initial_connection, force_connect_func) + + def _connect( + self, + host_info: HostInfo, + is_initial_connection: bool, + connect_func: Callable) -> Connection: + # The host to reconnect through is the one the application asked for, + # not one a later routing, read/write-splitting or failover decision + # landed on. It cannot be read from the plugin service in + # init_host_provider: AwsWrapperConnection.__init__ only fills that + # field in the refresh_host_list() that runs after the + # init_host_provider pipeline, so it is still None inside the hook. + # The first initial connect carries the same host, and this plugin is + # the outermost one in the chain, so it sees that host before any + # nested connect a routing plugin makes. + # + # Note this names the host the reconnect *asks* for. Plugins beneath + # this one may route the replacement elsewhere -- initial_connection + # substitutes a writer instance for a cluster endpoint on a + # non-initial connect. + if is_initial_connection and self._original_host_info is None: + self._original_host_info = host_info + + connection = connect_func() + self._record_connection_creation_time(connection) + return connection + + def _record_connection_creation_time(self, connection: Optional[Connection]) -> None: + if connection is None: + return + try: + self._connection_creation_times[connection] = self._clock_ns() + except TypeError: + # Not weak-referenceable; _select_current_connection_timestamp + # stamps it on first use. + pass + + def notify_connection_changed(self, changes: Set[ConnectionEvent]) -> OldConnectionSuggestedAction: + if (ConnectionEvent.INITIAL_CONNECTION in changes + or ConnectionEvent.CONNECTION_OBJECT_CHANGED in changes): + self._select_current_connection_timestamp() + return OldConnectionSuggestedAction.NO_OPINION + + def execute(self, target: type, method_name: str, execute_func: Callable, *args: Any, **kwargs: Any) -> Any: + if (method_name in LIFECYCLE_METHODS + or not self._plugin_service.is_network_bound_method(method_name)): + return execute_func() + + if self._terminal: + raise ConnectionLifetimeExpiredError( + Messages.get("ConnectionLifetimePlugin.ConnectionClosed")) + + if not self._is_current_connection_expired(): + return execute_func() + + # Let an open transaction finish on the connection it started on, so + # commit and rollback are not cut off mid-way. Deferring does not + # extend the lifetime: the action runs on the first eligible call once + # the transaction ends. + if self._plugin_service.is_in_transaction: + logger.debug("ConnectionLifetimePlugin.ActionDeferred", method_name) + return execute_func() + + if self._action == ConnectionLifetimeAction.CLOSE: + self._close_current_connection() + else: + self._reconnect() + + def _select_current_connection_timestamp(self) -> None: + current_connection = self._plugin_service.current_connection + creation_time = None + if current_connection is not None: + creation_time = self._connection_creation_times.get(current_connection) + if creation_time is None: + creation_time = self._clock_ns() + self._record_connection_creation_time(current_connection) + + self._timed_connection = current_connection + self._connection_creation_time_ns = creation_time + + def _is_current_connection_expired(self) -> bool: + current_connection = self._plugin_service.current_connection + if current_connection is not self._timed_connection or self._connection_creation_time_ns is None: + self._select_current_connection_timestamp() + + creation_time = self._connection_creation_time_ns + # The boundary is inclusive: a 1000 ms lifetime expires at exactly 1000 ms. + return (creation_time is not None + and self._clock_ns() - creation_time >= self._connection_lifetime_ns) + + def _close_current_connection(self) -> None: + # Set before closing: once the physical connection is on its way out + # the wrapper connection is unusable whether or not close() succeeds. + self._terminal = True + current_connection = self._plugin_service.current_connection + close_failure: Optional[Exception] = None + + if current_connection is not None: + self._connection_creation_times.pop(current_connection, None) + try: + current_connection.close() + except Exception as e: + close_failure = e + logger.warning("ConnectionLifetimePlugin.CloseConnectionFailed", e) + + error = ConnectionLifetimeExpiredError( + Messages.get("ConnectionLifetimePlugin.ConnectionClosed")) + if close_failure is not None: + raise error from close_failure + raise error + + def _reconnect(self) -> None: + previous_connection = self._plugin_service.current_connection + replacement_connection = None + try: + reconnect_host_info = self._original_host_info + if reconnect_host_info is None: + # Fall back to whatever the plugin service still reports; a + # pipeline that never ran an initial connect through this + # plugin (a pooled connection handed over, say) has no + # application host recorded. + reconnect_host_info = self._plugin_service.initial_connection_host_info + if reconnect_host_info is None: + raise AwsWrapperError( + Messages.get("ConnectionLifetimePlugin.OriginalHostUnavailable")) + + logger.debug("ConnectionLifetimePlugin.Reconnecting", reconnect_host_info.url) + replacement_connection = self._plugin_service.connect( + reconnect_host_info, self._props, self) + if replacement_connection is None: + raise AwsWrapperError( + Messages.get("ConnectionLifetimePlugin.ReconnectReturnedNone")) + replacement_creation_time = self._clock_ns() + + # Installs the replacement and disposes the old connection. The + # timer is only reset once this has succeeded. + self._plugin_service.set_current_connection(replacement_connection, reconnect_host_info) + self._forget_closed_connection(previous_connection) + self._connection_creation_times[replacement_connection] = replacement_creation_time + self._timed_connection = replacement_connection + self._connection_creation_time_ns = replacement_creation_time + except Exception as e: + self._cleanup_failed_replacement(replacement_connection, previous_connection) + raise ConnectionLifetimeFailedError( + Messages.get("ConnectionLifetimePlugin.UnableToReconnect")) from e + + raise ConnectionLifetimeSuccessError( + Messages.get("ConnectionLifetimePlugin.ConnectionChanged")) + + def _forget_closed_connection(self, connection: Optional[Connection]) -> None: + if connection is None: + return + + try: + closed = self._plugin_service.driver_dialect.is_closed(connection) + except Exception: + # A connection that cannot report its state is unusable; do not retain it. + closed = True + + if closed: + self._connection_creation_times.pop(connection, None) + + def _cleanup_failed_replacement( + self, + replacement_connection: Optional[Connection], + previous_connection: Optional[Connection]) -> None: + if (replacement_connection is not None + and replacement_connection is self._plugin_service.current_connection + and replacement_connection is not previous_connection): + # The replacement was installed and the failure came after, so + # neither connection can be relied on from here. + self._terminal = True + self._close_after_failed_installation(replacement_connection) + self._close_after_failed_installation(previous_connection) + return + + self._close_unattached_replacement(replacement_connection) + + def _close_after_failed_installation(self, connection: Optional[Connection]) -> None: + if connection is None: + return + + self._connection_creation_times.pop(connection, None) + try: + connection.close() + except Exception as e: + logger.warning("ConnectionLifetimePlugin.FailedInstallationConnectionCloseFailed", e) + + def _close_unattached_replacement(self, replacement_connection: Optional[Connection]) -> None: + if (replacement_connection is None + or replacement_connection is self._plugin_service.current_connection): + return + + self._connection_creation_times.pop(replacement_connection, None) + try: + replacement_connection.close() + except Exception as e: + logger.warning("ConnectionLifetimePlugin.ReplacementCloseFailed", e) + + +class ConnectionLifetimePluginFactory(PluginFactory): + @staticmethod + def get_instance(plugin_service: PluginService, props: Properties) -> Plugin: + return ConnectionLifetimePlugin(plugin_service, props) diff --git a/aws_advanced_python_wrapper/errors.py b/aws_advanced_python_wrapper/errors.py index 7af9a752..7c371a7f 100644 --- a/aws_advanced_python_wrapper/errors.py +++ b/aws_advanced_python_wrapper/errors.py @@ -67,6 +67,31 @@ class FailoverSuccessError(FailoverError): __module__ = "aws_advanced_python_wrapper" +class ConnectionLifetimeError(OperationalError): + __module__ = "aws_advanced_python_wrapper" + + +class ConnectionLifetimeExpiredError(ConnectionLifetimeError): + # connection_lifetime_action=close: the physical connection is gone and + # the wrapper connection is permanently unusable. + __module__ = "aws_advanced_python_wrapper" + sqlstate = "08003" + + +class ConnectionLifetimeSuccessError(ConnectionLifetimeError): + # connection_lifetime_action=reconnect succeeded: the wrapper connection + # now holds a replacement physical connection and may be retried once the + # caller re-creates its cursors. + __module__ = "aws_advanced_python_wrapper" + sqlstate = "08S02" + + +class ConnectionLifetimeFailedError(ConnectionLifetimeError): + # connection_lifetime_action=reconnect failed: no usable connection. + __module__ = "aws_advanced_python_wrapper" + sqlstate = "08001" + + class ReadWriteSplittingError(AwsWrapperError, InterfaceError): __module__ = "aws_advanced_python_wrapper" diff --git a/aws_advanced_python_wrapper/pep249.py b/aws_advanced_python_wrapper/pep249.py index 3e486b17..6f82e655 100644 --- a/aws_advanced_python_wrapper/pep249.py +++ b/aws_advanced_python_wrapper/pep249.py @@ -47,8 +47,8 @@ class Warning(Exception): class Error(Exception): __module__ = "aws_advanced_python_wrapper" - errno = None - sqlstate = None + errno: Optional[int] = None + sqlstate: Optional[str] = None class InterfaceError(Error): diff --git a/aws_advanced_python_wrapper/plugin_service.py b/aws_advanced_python_wrapper/plugin_service.py index b87b44c0..859d4803 100644 --- a/aws_advanced_python_wrapper/plugin_service.py +++ b/aws_advanced_python_wrapper/plugin_service.py @@ -55,6 +55,8 @@ AwsSecretsManagerPluginFactory from aws_advanced_python_wrapper.connect_time_plugin import \ ConnectTimePluginFactory +from aws_advanced_python_wrapper.connection_lifetime_plugin import \ + ConnectionLifetimePluginFactory from aws_advanced_python_wrapper.connection_provider import ( ConnectionProvider, ConnectionProviderManager) from aws_advanced_python_wrapper.database_dialect import ( @@ -884,6 +886,7 @@ class PluginManager(CanReleaseResources): "fastest_response_strategy": FastestResponseStrategyPluginFactory, "stale_dns": StaleDnsPluginFactory, "custom_endpoint": CustomEndpointPluginFactory, + "connection_lifetime": ConnectionLifetimePluginFactory, "connect_time": ConnectTimePluginFactory, "execute_time": ExecuteTimePluginFactory, "dev": DeveloperPluginFactory, @@ -900,6 +903,7 @@ class PluginManager(CanReleaseResources): # the highest values. The first plugin of the list will have the lowest weight, and the # last one will have the highest weight. PLUGIN_FACTORY_WEIGHTS: Dict[Type[PluginFactory], int] = { + ConnectionLifetimePluginFactory: 30, CustomEndpointPluginFactory: 40, AuroraInitialConnectionStrategyPluginFactory: 50, AuroraConnectionTrackerPluginFactory: 100, diff --git a/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties b/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties index a87c4d28..6023978b 100644 --- a/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties +++ b/aws_advanced_python_wrapper/resources/aws_advanced_python_wrapper_messages.properties @@ -116,6 +116,19 @@ conftest.ExceptionWhileObtainingInstanceIDs=[conftest] An exception was thrown w ConnectTimePlugin.ConnectTime=[ConnectTimePlugin] Connected in {} nanos. +ConnectionLifetimePlugin.InvalidConnectionLifetime=[ConnectionLifetimePlugin] The configuration property 'connection_lifetime_ms' is required and must be a positive whole number of milliseconds. Received: '{}'. +ConnectionLifetimePlugin.InvalidConnectionLifetimeAction=[ConnectionLifetimePlugin] The configuration property 'connection_lifetime_action' must be either 'close' or 'reconnect'. Received: '{}'. +ConnectionLifetimePlugin.ConnectionClosed=[ConnectionLifetimePlugin] The connection has reached its maximum lifetime and is closed. +ConnectionLifetimePlugin.ActionDeferred=[ConnectionLifetimePlugin] The connection has reached its maximum lifetime, but the lifetime action is deferred while a transaction is active. Executing '{}' on the current connection. +ConnectionLifetimePlugin.CloseConnectionFailed=[ConnectionLifetimePlugin] Unable to close the expired connection: {} +ConnectionLifetimePlugin.OriginalHostUnavailable=[ConnectionLifetimePlugin] The original application host is unavailable for connection-lifetime reconnection. +ConnectionLifetimePlugin.Reconnecting=[ConnectionLifetimePlugin] The connection has reached its maximum lifetime. Reconnecting to the original application host '{}'. +ConnectionLifetimePlugin.ReconnectReturnedNone=[ConnectionLifetimePlugin] The connection-lifetime reconnect attempt did not return a connection. +ConnectionLifetimePlugin.ReplacementCloseFailed=[ConnectionLifetimePlugin] Unable to close the unused replacement connection after a connection-lifetime reconnect failure: {} +ConnectionLifetimePlugin.FailedInstallationConnectionCloseFailed=[ConnectionLifetimePlugin] Unable to close a connection after a replacement connection was only partially installed: {} +ConnectionLifetimePlugin.ConnectionChanged=[ConnectionLifetimePlugin] The active connection changed because its configured maximum lifetime was reached. Re-create any cursors before retrying the operation. +ConnectionLifetimePlugin.UnableToReconnect=[ConnectionLifetimePlugin] Unable to establish a replacement connection after the configured maximum connection lifetime was reached. + ConnectionProvider.UnsupportedHostSelectorStrategy=[ConnectionProvider] Unsupported host selection strategy '{}' specified for this connection provider '{}'. Please visit the documentation for all supported strategies. CustomEndpointMonitor.DetectedChangeInCustomEndpointInfo=[CustomEndpointMonitor] Detected change in custom endpoint info for '{}':\n{} diff --git a/aws_advanced_python_wrapper/sqlalchemy_dialects/mysql.py b/aws_advanced_python_wrapper/sqlalchemy_dialects/mysql.py index e4706293..b836d069 100644 --- a/aws_advanced_python_wrapper/sqlalchemy_dialects/mysql.py +++ b/aws_advanced_python_wrapper/sqlalchemy_dialects/mysql.py @@ -102,8 +102,12 @@ def is_disconnect(self, e, connection, cursor): # through to upstream's e.errno probe and raised AttributeError out of # SA's unguarded is_disconnect call. See # sqlalchemy_dialects/_exception_handling.py. - from aws_advanced_python_wrapper.errors import (FailoverError, - FailoverFailedError) + from aws_advanced_python_wrapper.errors import ( + ConnectionLifetimeError, ConnectionLifetimeSuccessError, + FailoverError, FailoverFailedError) + + if isinstance(e, ConnectionLifetimeError): + return not isinstance(e, ConnectionLifetimeSuccessError) # Catch the whole FailoverError family (FailoverSuccessError, # FailoverFailedError, AND TransactionResolutionUnknownError) before the diff --git a/aws_advanced_python_wrapper/sqlalchemy_dialects/mysql_async.py b/aws_advanced_python_wrapper/sqlalchemy_dialects/mysql_async.py index 58d056ca..c90062bf 100644 --- a/aws_advanced_python_wrapper/sqlalchemy_dialects/mysql_async.py +++ b/aws_advanced_python_wrapper/sqlalchemy_dialects/mysql_async.py @@ -231,8 +231,12 @@ def is_disconnect(self, e, connection, cursor): # return True so SA invalidates and the creator retries. # _AsyncDriverErrorNormalizeMixin handles the do_execute path; # this handles the cursor-creation path that runs earlier. - from aws_advanced_python_wrapper.errors import (FailoverError, - FailoverFailedError) + from aws_advanced_python_wrapper.errors import ( + ConnectionLifetimeError, ConnectionLifetimeSuccessError, + FailoverError, FailoverFailedError) + + if isinstance(e, ConnectionLifetimeError): + return not isinstance(e, ConnectionLifetimeSuccessError) # Catch the whole FailoverError family -- including # TransactionResolutionUnknownError -- before upstream probes diff --git a/aws_advanced_python_wrapper/sqlalchemy_dialects/pg.py b/aws_advanced_python_wrapper/sqlalchemy_dialects/pg.py index a5a9728c..80541b48 100644 --- a/aws_advanced_python_wrapper/sqlalchemy_dialects/pg.py +++ b/aws_advanced_python_wrapper/sqlalchemy_dialects/pg.py @@ -104,8 +104,12 @@ def is_disconnect(self, e, connection, cursor): # bound to the new writer via plugin_service.current_connection; # SA should reuse it, not invalidate). # - FailoverFailedError → True (no usable connection; invalidate). - from aws_advanced_python_wrapper.errors import (FailoverError, - FailoverFailedError) + from aws_advanced_python_wrapper.errors import ( + ConnectionLifetimeError, ConnectionLifetimeSuccessError, + FailoverError, FailoverFailedError) + + if isinstance(e, ConnectionLifetimeError): + return not isinstance(e, ConnectionLifetimeSuccessError) # Catch the whole FailoverError family -- including # TransactionResolutionUnknownError -- before the upstream is_disconnect diff --git a/aws_advanced_python_wrapper/sqlalchemy_dialects/pg_async.py b/aws_advanced_python_wrapper/sqlalchemy_dialects/pg_async.py index 132640d7..992309a1 100644 --- a/aws_advanced_python_wrapper/sqlalchemy_dialects/pg_async.py +++ b/aws_advanced_python_wrapper/sqlalchemy_dialects/pg_async.py @@ -154,8 +154,12 @@ def is_disconnect(self, e, connection, cursor): # - FailoverFailedError → True (no usable connection). # Complements _AsyncDriverErrorNormalizeMixin for the # cursor-creation path that runs before do_execute. - from aws_advanced_python_wrapper.errors import (FailoverError, - FailoverFailedError) + from aws_advanced_python_wrapper.errors import ( + ConnectionLifetimeError, ConnectionLifetimeSuccessError, + FailoverError, FailoverFailedError) + + if isinstance(e, ConnectionLifetimeError): + return not isinstance(e, ConnectionLifetimeSuccessError) # Catch the whole FailoverError family -- including # TransactionResolutionUnknownError -- before upstream probes diff --git a/aws_advanced_python_wrapper/utils/properties.py b/aws_advanced_python_wrapper/utils/properties.py index 9623c83f..f9911360 100644 --- a/aws_advanced_python_wrapper/utils/properties.py +++ b/aws_advanced_python_wrapper/utils/properties.py @@ -355,6 +355,19 @@ class WrapperProperties: 5_000, ) + # ConnectionLifetimePlugin + CONNECTION_LIFETIME_MS = WrapperProperty( + "connection_lifetime_ms", + """Maximum lifetime of a physical connection, in milliseconds. Required by the connection_lifetime + plugin, and must be at least 1.""", + ) + CONNECTION_LIFETIME_ACTION = WrapperProperty( + "connection_lifetime_action", + """Action to take once a physical connection reaches its configured lifetime. Supported values are + 'close' and 'reconnect'.""", + "reconnect", + ) + # Host Availability Strategy DEFAULT_HOST_AVAILABILITY_STRATEGY = WrapperProperty( "default_host_availability_strategy", diff --git a/docs/README.md b/docs/README.md index a1ceef71..63ed2eda 100644 --- a/docs/README.md +++ b/docs/README.md @@ -27,6 +27,7 @@ - [Reader Selection Strategies](using-the-python-wrapper/ReaderSelectionStrategies.md) - [Simple Read Write Splitting Plugin](using-the-python-wrapper/using-plugins/UsingTheSimpleReadWriteSplittingPlugin.md) - [Fastest Response Strategy Plugin](using-the-python-wrapper/using-plugins/UsingTheFastestResponseStrategyPlugin.md) + - [Connection Lifetime Plugin](using-the-python-wrapper/using-plugins/UsingTheConnectionLifetimePlugin.md) - [Host Availability Strategy](using-the-python-wrapper/HostAvailabilityStrategy.md) - [Development Guide](./development-guide/DevelopmentGuide.md) - [Setup](./development-guide/DevelopmentGuide.md#setup) diff --git a/docs/using-the-python-wrapper/PluginChainCompatibility.md b/docs/using-the-python-wrapper/PluginChainCompatibility.md index 02b54384..4f7e1828 100644 --- a/docs/using-the-python-wrapper/PluginChainCompatibility.md +++ b/docs/using-the-python-wrapper/PluginChainCompatibility.md @@ -16,6 +16,12 @@ a chain you haven't used before. |---------------------|-----------------------------------|-----| | `read_write_splitting` | `failover_v2` (not `failover`) | `failover_v2` starts the cluster topology monitor on the **initial** connect. With plain `failover`, the host list stays at `{writer-only}` until a later failover signal fires, and `conn.read_only = True` flips silently fall back to the writer rather than splitting to a reader. | +## Host-selection interactions + +| Plugin | Interacts with | Why | +|-----------------------|-------------------------------------------------|-----| +| `connection_lifetime` | `read_write_splitting`, `srw`, `gdb_rw`, `failover`, `failover_v2` | With `connection_lifetime_action=reconnect` the replacement connection always goes to the host the application originally supplied, never to the reader a splitting plugin switched to or the writer a failover landed on. An expired reader connection therefore returns to the original endpoint. Use `connection_lifetime_action=close` if the application should pick its own host again. | + ## Driver-specific incompatibilities | Plugin | Incompatible with | Why | diff --git a/docs/using-the-python-wrapper/SqlAlchemySupport.md b/docs/using-the-python-wrapper/SqlAlchemySupport.md index b4d5afd7..1a30432d 100644 --- a/docs/using-the-python-wrapper/SqlAlchemySupport.md +++ b/docs/using-the-python-wrapper/SqlAlchemySupport.md @@ -182,6 +182,7 @@ Plugins are configured identically to non-SA usage — via the `plugins` connect | [Fastest Response Strategy Plugin](using-plugins/UsingTheFastestResponseStrategyPlugin.md) | `fastest_response_strategy` | ✓ | | [Blue/Green Deployment Plugin](using-plugins/UsingTheBlueGreenPlugin.md) | `bg` | ✓ | | [Limitless Plugin](using-plugins/UsingTheLimitlessPlugin.md) | `limitless` | ✓ | +| [Connection Lifetime Plugin](using-plugins/UsingTheConnectionLifetimePlugin.md) | `connection_lifetime` | ✓ | | [Read Write Splitting Plugin](using-plugins/UsingTheReadWriteSplittingPlugin.md) | `read_write_splitting` | ✗ | | [Simple Read Write Splitting Plugin](using-plugins/UsingTheSimpleReadWriteSplittingPlugin.md) | `srw` | ✗ | diff --git a/docs/using-the-python-wrapper/UsingThePythonWrapper.md b/docs/using-the-python-wrapper/UsingThePythonWrapper.md index 384dc1fb..1ef5676a 100644 --- a/docs/using-the-python-wrapper/UsingThePythonWrapper.md +++ b/docs/using-the-python-wrapper/UsingThePythonWrapper.md @@ -107,6 +107,7 @@ The AWS Advanced Python Wrapper has several built-in plugins that are available | [Fastest Response Strategy Plugin](./using-plugins/UsingTheFastestResponseStrategyPlugin.md) | `fastest_response_strategy` | Aurora | A host selection strategy plugin that uses a host monitoring service to monitor each reader host's response time and choose the host with the fastest response. | None | | [Blue/Green Deployment Plugin](./using-plugins/UsingTheBlueGreenPlugin.md) | `bg` | Aurora,
RDS Instance | Enables client-side Blue/Green Deployment support. | None | | [Limitless Plugin](./using-plugins/UsingTheLimitlessPlugin.md) | `limitless` | Aurora | Enables client-side load-balancing of Transaction Routers on Amazon Aurora Limitless Databases. | None | +| [Connection Lifetime Plugin](./using-plugins/UsingTheConnectionLifetimePlugin.md) | `connection_lifetime` | Any database | Limits the lifetime of a physical connection and either closes it or reconnects through the host the application originally supplied, before the next network-bound call. | None | In addition to the built-in plugins, you can also create custom plugins more suitable for your needs. For more information, see [Custom Plugins](../development-guide/LoadablePlugins.md#using-custom-plugins). diff --git a/tests/unit/test_aio_plugin_factory.py b/tests/unit/test_aio_plugin_factory.py index c2cbd6af..b2dfcaf7 100644 --- a/tests/unit/test_aio_plugin_factory.py +++ b/tests/unit/test_aio_plugin_factory.py @@ -258,6 +258,7 @@ async def _body() -> None: # SRW endpoints are required at construction time. "srw_read_endpoint": "reader.example.com", "srw_write_endpoint": "writer.example.com", + "connection_lifetime_ms": "60000", }) # Use Aurora-capable hlp so failover / rws / custom_endpoint # don't error. diff --git a/tests/unit/test_aio_sqlalchemy_dialect.py b/tests/unit/test_aio_sqlalchemy_dialect.py index 28a1c5a4..cde72fc8 100644 --- a/tests/unit/test_aio_sqlalchemy_dialect.py +++ b/tests/unit/test_aio_sqlalchemy_dialect.py @@ -338,3 +338,27 @@ class _NativeLike: called_arg = fetch_mock.call_args.args[0] assert called_arg is native + + +# ---- is_disconnect classification --------------------------------------- + + +def test_async_dialects_classify_the_connection_lifetime_error_family(): + """Only a successful reconnect leaves the pool slot's wrapper usable. + + The close action and a failed reconnect leave nothing to reuse, and none + of the three carry the driver attributes upstream's ``is_disconnect`` + probes, so they have to be classified before falling through. + """ + from aws_advanced_python_wrapper.errors import ( + ConnectionLifetimeExpiredError, ConnectionLifetimeFailedError, + ConnectionLifetimeSuccessError) + from aws_advanced_python_wrapper.sqlalchemy_dialects.mysql_async import \ + AwsWrapperMySQLAiomysqlAsyncDialect + + for dialect_cls in (AwsWrapperPGPsycopgAsyncDialect, + AwsWrapperMySQLAiomysqlAsyncDialect): + d = dialect_cls.__new__(dialect_cls) + assert d.is_disconnect(ConnectionLifetimeSuccessError("x"), None, None) is False + assert d.is_disconnect(ConnectionLifetimeExpiredError("x"), None, None) is True + assert d.is_disconnect(ConnectionLifetimeFailedError("x"), None, None) is True diff --git a/tests/unit/test_sqlalchemy_dialects.py b/tests/unit/test_sqlalchemy_dialects.py index 073535d7..10104f0c 100644 --- a/tests/unit/test_sqlalchemy_dialects.py +++ b/tests/unit/test_sqlalchemy_dialects.py @@ -74,6 +74,22 @@ def test_is_disconnect_handles_full_failover_error_family(dialect_cls): assert d.is_disconnect(FailoverFailedError("x"), None, None) is True +@pytest.mark.parametrize( + "dialect_cls", [AwsWrapperMySQLConnectorDialect, AwsWrapperPGPsycopgDialect]) +def test_is_disconnect_handles_the_connection_lifetime_error_family(dialect_cls): + # Only a successful reconnect leaves the pool slot's wrapper usable; the + # close action and a failed reconnect leave nothing to reuse. None of + # these carry ``errno``, so they have to be classified before the + # upstream probe, same as the FailoverError family above. + from aws_advanced_python_wrapper.errors import ( + ConnectionLifetimeExpiredError, ConnectionLifetimeFailedError, + ConnectionLifetimeSuccessError) + d = dialect_cls.__new__(dialect_cls) # avoid full dialect __init__/dbapi + assert d.is_disconnect(ConnectionLifetimeSuccessError("x"), None, None) is False + assert d.is_disconnect(ConnectionLifetimeExpiredError("x"), None, None) is True + assert d.is_disconnect(ConnectionLifetimeFailedError("x"), None, None) is True + + # Registration uses the SA . convention: the wrapper plugs in # as a driver under the stock ``postgresql`` / ``mysql`` dialects, so the URL is # ``postgresql+aws_wrapper_psycopg`` / ``mysql+aws_wrapper_mysqlconnector``.