chore: improve chain creation readability and fix must_use_pipeline dead code

Author: karenc-bq

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):

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)