From 6d95d5eeee8cede22e29136ccf5dc7208a4e1e39 Mon Sep 17 00:00:00 2001 From: karezche <64801825+karenc-bq@users.noreply.github.com> Date: Tue, 4 Aug 2026 07:56:15 -0700 Subject: [PATCH 1/2] chore: improve chain creation readability and fix must_use_pipeline dead code --- aws_advanced_python_wrapper/plugin_service.py | 73 ++++++++----- tests/unit/test_plugin_manager.py | 102 ++++++++++++++++++ 2 files changed, 151 insertions(+), 24 deletions(-) diff --git a/aws_advanced_python_wrapper/plugin_service.py b/aws_advanced_python_wrapper/plugin_service.py index 76eae5d28..51cbcf6ec 100644 --- a/aws_advanced_python_wrapper/plugin_service.py +++ b/aws_advanced_python_wrapper/plugin_service.py @@ -985,9 +985,26 @@ 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. + + Mirrors JDBC ``ConnectionPluginManager.mustUsePipeline``: 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). + + Intentional deviation from JDBC: the trailing ``is_network_bound_method`` term. JDBC's + DefaultConnectionPlugin is a thin passthrough, but Python's 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 +1061,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/tests/unit/test_plugin_manager.py b/tests/unit/test_plugin_manager.py index fd3b758de..97c6e88bf 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) From 49cd6a1865f1667dbc90f3500bd18410e6cbac30 Mon Sep 17 00:00:00 2001 From: karezche <64801825+karenc-bq@users.noreply.github.com> Date: Tue, 4 Aug 2026 16:32:46 -0700 Subject: [PATCH 2/2] test: fix flakey unit tests --- tests/unit/test_aio_host_list_provider.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_aio_host_list_provider.py b/tests/unit/test_aio_host_list_provider.py index ae089d6ca..8cf91b0b2 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