diff --git a/aws_advanced_python_wrapper/aio/driver_dialect/psycopg.py b/aws_advanced_python_wrapper/aio/driver_dialect/psycopg.py index 61425098..223aa79d 100644 --- a/aws_advanced_python_wrapper/aio/driver_dialect/psycopg.py +++ b/aws_advanced_python_wrapper/aio/driver_dialect/psycopg.py @@ -162,8 +162,8 @@ async def abort_connection(self, conn: Any) -> None: # EFM exists for. # # Shutting the underlying socket down (SHUT_RDWR) -- the async mirror of - # sync PgDriverDialect.abort_connection and of JDBC Connection.abort() -- - # makes this event loop's selector see the fd become readable/errored, so + # sync PgDriverDialect.abort_connection -- makes this event loop's + # selector see the fd become readable/errored, so # the suspended read wakes immediately (even on a dead host) with an # OSError/OperationalError the failover plugin classifies as a connection # loss. We detach (not close) the fd so the connection still owns it and diff --git a/aws_advanced_python_wrapper/blue_green_plugin.py b/aws_advanced_python_wrapper/blue_green_plugin.py index e8618a05..3e65d3a1 100644 --- a/aws_advanced_python_wrapper/blue_green_plugin.py +++ b/aws_advanced_python_wrapper/blue_green_plugin.py @@ -1529,7 +1529,7 @@ def _get_status_of_created(self) -> BlueGreenStatus: """ New connect requests: go to blue or green hosts; default behaviour; no routing. Existing connections: default behaviour; no action. - Execute JDBC calls: default behaviour; no action. + Method execution: default behaviour; no action. """ return BlueGreenStatus( self._bg_id, @@ -1546,7 +1546,7 @@ def _get_status_of_preparation(self): New connect requests to green: route to corresponding IP address. New connect requests with IP address: default behaviour; no routing. Existing connections: default behaviour; no action. - Execute JDBC calls: default behaviour; no action. + Method execution: default behaviour; no action. """ if self._is_switchover_timer_expired(): @@ -1603,7 +1603,7 @@ def _get_status_of_in_progress(self) -> BlueGreenStatus: New connect requests to green: suspend. New connect requests with IP address: suspend. Existing connections: default behaviour; no action. - Execute JDBC calls: suspend. + Method execution: suspend. """ if self._is_switchover_timer_expired(): diff --git a/aws_advanced_python_wrapper/custom_endpoint_plugin.py b/aws_advanced_python_wrapper/custom_endpoint_plugin.py index 9ec3aa68..08361ae0 100644 --- a/aws_advanced_python_wrapper/custom_endpoint_plugin.py +++ b/aws_advanced_python_wrapper/custom_endpoint_plugin.py @@ -266,7 +266,7 @@ def __init__(self, plugin_service: PluginService, props: Properties): self._monitors.register_monitor_type( CustomEndpointMonitor, expiration_timeout_ns=self._idle_monitor_expiration_ms * 1_000_000, - inactive_timeout_ns=1 * 60 * 1_000_000_000) # 1 minute, matches JDBC + inactive_timeout_ns=1 * 60 * 1_000_000_000) # 1 minute CustomEndpointPlugin._SUBSCRIBED_METHODS.update(self._plugin_service.network_bound_methods) diff --git a/aws_advanced_python_wrapper/mysql_driver_dialect.py b/aws_advanced_python_wrapper/mysql_driver_dialect.py index 1a503c1d..4b6f2821 100644 --- a/aws_advanced_python_wrapper/mysql_driver_dialect.py +++ b/aws_advanced_python_wrapper/mysql_driver_dialect.py @@ -164,8 +164,8 @@ def abort_connection(self, conn: Connection): # operation so the owning thread's blocked recv returns promptly, WITHOUT # freeing the connection (the owning thread closes it -- freeing it here # would race a cross-thread use-after-free in the driver, the env-4 SIGSEGV). - # Thread-safe equivalent of JDBC's Connection.abort(). Only the pure-Python - # connector exposes the raw socket; best-effort no-op for the C extension. + # Only the pure-Python connector exposes the raw socket; best-effort no-op + # for the C extension. if not MySQLDriverDialect._is_mysql_connection(conn): raise UnsupportedOperationError( Messages.get_formatted( diff --git a/aws_advanced_python_wrapper/pg_driver_dialect.py b/aws_advanced_python_wrapper/pg_driver_dialect.py index 269e3ff2..41501134 100644 --- a/aws_advanced_python_wrapper/pg_driver_dialect.py +++ b/aws_advanced_python_wrapper/pg_driver_dialect.py @@ -92,8 +92,7 @@ def abort_connection(self, conn: Connection): # which defeats the EFM's purpose on exactly the network-failure case # it exists for. # - # Shutting the underlying socket down is the thread-safe equivalent of - # JDBC's Connection.abort(): it unblocks the owning thread's recv + # Shutting the underlying socket down unblocks the owning thread's recv # immediately (even on a dead host) WITHOUT freeing any struct, so # there is no SSL_free to race. The owning thread observes the broken # connection and closes it on its own thread (the only safe place for diff --git a/aws_advanced_python_wrapper/plugin_service.py b/aws_advanced_python_wrapper/plugin_service.py index 76eae5d2..5d505644 100644 --- a/aws_advanced_python_wrapper/plugin_service.py +++ b/aws_advanced_python_wrapper/plugin_service.py @@ -985,9 +985,25 @@ def get_factory_weights(factory_types: List[Type[PluginFactory]]) -> Dict[Type[P return weights - def must_use_pipeline(self, method: DbApiMethod): + def must_use_pipeline(self, method: DbApiMethod) -> bool: + """Whether this method has to run through the plugin pipeline. + + The pipeline is required when the method always uses it, when the chain has not been + built yet (nothing to decide on), when a real plugin is subscribed, or when telemetry + is on (so per-plugin NESTED spans are still emitted). + + The trailing ``is_network_bound_method`` term keeps network-bound methods on the + pipeline even when nothing else requires it: DefaultPlugin.execute also applies + DriverDialect.execute's socket timeout and its interrupt-and-wait cleanup. Skipping + that for a network-bound method lets a later close/reuse race a still-running operation + (env-4 SIGSEGV), so those methods stay on the pipeline regardless of subscriptions. + """ plugin_chain_info: Optional[PluginChainCallableInfo] = self._function_cache[method.id] - return method.always_use_pipeline or plugin_chain_info is None or plugin_chain_info.is_subscribed or self._telemetry_in_use + return (method.always_use_pipeline + or plugin_chain_info is None + or plugin_chain_info.is_subscribed + or self._telemetry_in_use + or self._container.plugin_service.is_network_bound_method(method.method_name)) def execute(self, target: object, method: DbApiMethod, target_driver_func: Callable, *args, **kwargs) -> Any: plugin_service = self._container.plugin_service @@ -1044,36 +1060,44 @@ def _execute_with_subscribed_plugins( pipeline_func_info = self._make_pipeline(method.method_name) self._function_cache[method.id] = pipeline_func_info - # Execute only if method needs to use pipeline, or a plugin is subscribed to this method - if method.always_use_pipeline or pipeline_func_info.is_subscribed: + # Execute only if the method needs to use the pipeline, or a plugin is subscribed to it. + if self.must_use_pipeline(method): return pipeline_func_info.func(plugin_func, target_driver_func, method.method_name, plugin_to_skip) - else: - return target_driver_func() + + result = target_driver_func() + + # DefaultPlugin.execute refreshes the cached in-transaction state after every method except + # close; failover and read_write_splitting read it to decide whether a transaction is open. + plugin_service = self._container.plugin_service + if method != DbApiMethod.CONNECTION_CLOSE and plugin_service.current_connection is not None: + plugin_service.update_in_transaction() + + return result + + def _subscribed_plugins(self, method_name: str) -> List[Plugin]: + all_methods_marker = DbApiMethod.ALL.method_name + return [ + plugin for plugin in self._plugins + if all_methods_marker in plugin.subscribed_methods or method_name in plugin.subscribed_methods + ] # Builds the plugin pipeline function chain. The pipeline is built in a way that allows plugins to perform logic # both before and after the target driver function call. def _make_pipeline(self, method_name: str) -> PluginChainCallableInfo: - pipeline_func: Optional[Callable] = None - num_plugins: int = len(self._plugins) - is_subscribed: bool = False - - # Build the pipeline starting at the end and working backwards - for i in range(num_plugins - 1, -1, -1): - plugin: Plugin = self._plugins[i] + subscribed = self._subscribed_plugins(method_name) + if not subscribed: + raise AwsWrapperError(Messages.get("PluginManager.PipelineNone")) - subscribed_methods: Set[str] = plugin.subscribed_methods - is_plugin_subscribed = DbApiMethod.ALL.method_name in subscribed_methods or method_name in subscribed_methods - is_subscribed |= is_plugin_subscribed + # DefaultPlugin subscribes to "*" and is appended to every plugin list, so counting it here would + # pin is_subscribed to True for every method and make the bypass in _execute_with_subscribed_plugins + # unreachable. + is_subscribed = any(not isinstance(plugin, DefaultPlugin) for plugin in subscribed) - if is_plugin_subscribed: - if pipeline_func is None: - # Defines the call to DefaultPlugin, which is the last plugin in the pipeline - pipeline_func = self._create_base_pipeline_func(plugin) - continue - pipeline_func = self._extend_pipeline_func(plugin, pipeline_func) + # Build the pipeline starting at the end and working backwards + pipeline_func = self._create_base_pipeline_func(subscribed[-1]) + for plugin in reversed(subscribed[:-1]): + pipeline_func = self._extend_pipeline_func(plugin, pipeline_func) - if pipeline_func is None: - raise AwsWrapperError(Messages.get("PluginManager.PipelineNone")) return PluginChainCallableInfo(pipeline_func, is_subscribed) def _create_base_pipeline_func(self, plugin: Plugin): 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 ee8d6e3d..1a0ba75e 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 @@ -272,7 +272,7 @@ IamAuthUtils.GeneratedNewAuthToken=Generated new authentication token = {} LimitlessPlugin.FailedToConnectToHost=[LimitlessPlugin] Failed to connect to host {}. LimitlessPlugin.UnsupportedDialectOrDatabase=[LimitlessPlugin] Unsupported dialect '{}' encountered. Please ensure the connection parameters are correct, and refer to the documentation to ensure that the connecting database is compatible with the Limitless Connection Plugin. -LimitlessQueryHelper.UnsupportedDialectOrDatabase=[LimitlessQueryHelper] Unsupported dialect '{}' encountered. Please ensure JDBC connection parameters are correct, and refer to the documentation to ensure that the connecting database is compatible with the Limitless Connection Plugin. +LimitlessQueryHelper.UnsupportedDialectOrDatabase=[LimitlessQueryHelper] Unsupported dialect '{}' encountered. Please ensure connection parameters are correct, and refer to the documentation to ensure that the connecting database is compatible with the Limitless Connection Plugin. LimitlessRouterMonitor.errorDuringMonitoringStop=[LimitlessRouterMonitor] Stopping monitoring after unhandled error was thrown in Limitless Router Monitoring thread for host {}. Error: {} LimitlessRouterMonitor.InterruptedErrorDuringMonitoring=[LimitlessRouterMonitor] Limitless Router Monitoring thread for host {} was interrupted. diff --git a/tests/integration/container/test_blue_green_deployment.py b/tests/integration/container/test_blue_green_deployment.py index 94acffa6..aab38443 100644 --- a/tests/integration/container/test_blue_green_deployment.py +++ b/tests/integration/container/test_blue_green_deployment.py @@ -1027,7 +1027,7 @@ def green_iam_connectivity_monitor( else: self.logger.debug(f"[DirectGreenIamIp{thread_prefix} @ {host_id}] Thread exception: {e}") result_queue.append(TimeHolder(start_ns, perf_counter_ns(), error=str(e))) - # TODO: is 'Access Denied' the error message in Python as well as JDBC? + # TODO: confirm 'Access Denied' is the error message surfaced in Python. if notify_on_first_error and "access denied" in str(e).lower(): results.green_node_changed_name_time_ns.compare_and_set(0, perf_counter_ns()) self.logger.debug( diff --git a/tests/unit/test_aio_host_list_provider.py b/tests/unit/test_aio_host_list_provider.py index ae089d6c..8cf91b0b 100644 --- a/tests/unit/test_aio_host_list_provider.py +++ b/tests/unit/test_aio_host_list_provider.py @@ -542,7 +542,7 @@ def test_topology_monitor_enters_high_freq_mode_on_writer_change(): async def _run_briefly(): monitor.start() # Allow enough ticks for writer-change detection - await asyncio.sleep(0.08) + await asyncio.sleep(1) await monitor.stop() asyncio.run(_run_briefly()) @@ -628,7 +628,7 @@ def test_topology_monitor_ignores_requests_after_writer_confirmed(): async def _run(): monitor.start() - await asyncio.sleep(0.08) # let at least one tick complete + await asyncio.sleep(1) # let at least one tick complete ignore_during = monitor.should_ignore_refresh_request() await monitor.stop() return ignore_during diff --git a/tests/unit/test_plugin_manager.py b/tests/unit/test_plugin_manager.py index fd3b758d..97c6e88b 100644 --- a/tests/unit/test_plugin_manager.py +++ b/tests/unit/test_plugin_manager.py @@ -610,3 +610,105 @@ def notify_host_list_changed(self, changes: Dict[str, Set[HostEvent]]): def notify_connection_changed(self, changes: Set[ConnectionEvent]) -> OldConnectionSuggestedAction: self._calls.append(type(self).__name__ + ":notify_connection_changed") raise AwsWrapperError() + + +def test_default_plugin_excluded_from_is_subscribed(mocker, mock_telemetry_factory): + # DefaultPlugin subscribes to "*", but it must not mark a method as subscribed on its own -- + # otherwise the direct-call bypass in _execute_with_subscribed_plugins is unreachable. + mocker.patch.object(PluginManager, "__init__", lambda w, x, y, z: None) + manager = PluginManager(mocker.MagicMock(), mocker.MagicMock(), mocker.MagicMock()) + manager._plugins = [DefaultPlugin(mocker.MagicMock(), mocker.MagicMock())] + manager._telemetry_factory = mock_telemetry_factory + + assert not manager._make_pipeline(DbApiMethod.CURSOR_EXECUTE.method_name).is_subscribed + assert not manager._make_pipeline(DbApiMethod.CONNECT.method_name).is_subscribed + + # A real subscribing plugin still sets the flag, and only for the methods it subscribes to. + manager._plugins = [TestPluginTwo([]), DefaultPlugin(mocker.MagicMock(), mocker.MagicMock())] + assert manager._make_pipeline(DbApiMethodTest.TEST_CALL_A.method_name).is_subscribed + assert not manager._make_pipeline(DbApiMethod.CURSOR_FETCHALL.method_name).is_subscribed + + +def test_unsubscribed_method_bypasses_pipeline(mocker, container, mock_telemetry_factory): + # With only DefaultPlugin in the chain, a non-network-bound method skips the pipeline entirely + # but must still refresh the cached in-transaction state. + calls = [] + container.plugin_service.is_network_bound_method.side_effect = \ + lambda name: name == DbApiMethod.CURSOR_EXECUTE.method_name + container.plugin_service.update_in_transaction.side_effect = \ + lambda *args: calls.append("update_in_transaction") + container.plugin_service.driver_dialect.execute.side_effect = \ + lambda method_name, func, *args, **kwargs: (calls.append("dialect.execute"), func())[1] + + mocker.patch.object(PluginManager, "__init__", lambda w, x, y, z: None) + manager = PluginManager(mocker.MagicMock(), mocker.MagicMock(), mocker.MagicMock()) + manager._plugins = [DefaultPlugin(container.plugin_service, mocker.MagicMock())] + manager._container = container + manager._telemetry_factory = mock_telemetry_factory + manager._telemetry_factory.open_telemetry_context.return_value = None + manager._telemetry_in_use = False + manager._function_cache = [None] * (DbApiMethod.ALL.id + 1) + + def _execute(method): + return manager._execute_with_subscribed_plugins( + method, + lambda plugin, next_func: plugin.execute(mocker.MagicMock(), method.method_name, next_func), + lambda: (calls.append("target"), "result_value")[1]) + + # Not network bound -> bypass, no DriverDialect.execute, transaction state still updated. + assert _execute(DbApiMethod.CURSOR_LASTROWID) == "result_value" + assert calls == ["target", "update_in_transaction"] + + # Network bound -> stays on the pipeline so the socket timeout guard is preserved. + calls.clear() + assert _execute(DbApiMethod.CURSOR_EXECUTE) == "result_value" + assert calls == ["dialect.execute", "target", "update_in_transaction"] + + # Telemetry on -> back on the pipeline even for the otherwise-bypassable method, so the + # per-plugin NESTED spans are still emitted. + calls.clear() + manager._telemetry_in_use = True + manager._function_cache = [None] * (DbApiMethod.ALL.id + 1) + assert _execute(DbApiMethod.CURSOR_LASTROWID) == "result_value" + assert calls == ["dialect.execute", "target", "update_in_transaction"] + + +def test_must_use_pipeline(mocker, container, mock_telemetry_factory): + # must_use_pipeline is the single authority for the bypass decision, so each term matters. + container.plugin_service.is_network_bound_method.side_effect = \ + lambda name: name == DbApiMethod.CURSOR_EXECUTE.method_name + + mocker.patch.object(PluginManager, "__init__", lambda w, x, y, z: None) + manager = PluginManager(mocker.MagicMock(), mocker.MagicMock(), mocker.MagicMock()) + manager._plugins = [DefaultPlugin(container.plugin_service, mocker.MagicMock())] + manager._container = container + manager._telemetry_factory = mock_telemetry_factory + manager._telemetry_in_use = False + manager._function_cache = [None] * (DbApiMethod.ALL.id + 1) + + # Chain not built yet -> nothing to decide on, so the pipeline is required. + assert manager.must_use_pipeline(DbApiMethod.CURSOR_LASTROWID) + + # Built, unsubscribed, not network bound, telemetry off -> bypass allowed. + manager._function_cache[DbApiMethod.CURSOR_LASTROWID.id] = \ + manager._make_pipeline(DbApiMethod.CURSOR_LASTROWID.method_name) + assert not manager.must_use_pipeline(DbApiMethod.CURSOR_LASTROWID) + + # always_use_pipeline and network-bound methods are always required. + assert manager.must_use_pipeline(DbApiMethod.CONNECT) + manager._function_cache[DbApiMethod.CURSOR_EXECUTE.id] = \ + manager._make_pipeline(DbApiMethod.CURSOR_EXECUTE.method_name) + assert manager.must_use_pipeline(DbApiMethod.CURSOR_EXECUTE) + + # Telemetry re-enables the pipeline for the otherwise-bypassable method. + manager._telemetry_in_use = True + assert manager.must_use_pipeline(DbApiMethod.CURSOR_LASTROWID) + + # A real subscribed plugin also forces the pipeline. + subscriber = mocker.MagicMock() + subscriber.subscribed_methods = {DbApiMethod.CURSOR_LASTROWID.method_name} + manager._plugins = [subscriber, DefaultPlugin(container.plugin_service, mocker.MagicMock())] + manager._telemetry_in_use = False + manager._function_cache[DbApiMethod.CURSOR_LASTROWID.id] = \ + manager._make_pipeline(DbApiMethod.CURSOR_LASTROWID.method_name) + assert manager.must_use_pipeline(DbApiMethod.CURSOR_LASTROWID)