Experimental plugin autoloading (boto#3514)

Author: SamRemis

botocore/credentials.py

@@ -50,6 +50,7 @@
     InstanceMetadataFetcher,
     JSONFileCache,
     SSOTokenLoader,
+    create_nested_client,
     parse_key_val_file,
     resolve_imds_endpoint_mode,
 )
@@ -265,7 +266,9 @@ def _get_client_creator(session, region_name):
     def client_creator(service_name, **kwargs):
         create_client_kwargs = {'region_name': region_name}
         create_client_kwargs.update(**kwargs)
-        return session.create_client(service_name, **create_client_kwargs)
+        return create_nested_client(
+            session, service_name, **create_client_kwargs
+        )
 
     return client_creator
 

botocore/plugin.py

@@ -0,0 +1,85 @@
+# Copyright 2025 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. A copy of
+# the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file 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.
+"""
+NOTE: This module is considered private and is subject to abrupt breaking
+changes without prior announcement. Please do not use it directly.
+"""
+
+import importlib
+import logging
+import os
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Optional
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class PluginContext:
+    """
+    Encapsulation of plugins tracked within the `_plugin_context` context variable.
+    """
+
+    plugins: Optional[str] = None
+
+
+_plugin_context = ContextVar("_plugin_context")
+
+
+def get_plugin_context():
+    """Get the current `_plugin_context` context variable if set, else None."""
+    return _plugin_context.get(None)
+
+
+def set_plugin_context(ctx):
+    """Set the current `_plugin_context` context variable."""
+    token = _plugin_context.set(ctx)
+    return token
+
+
+def reset_plugin_context(token):
+    """Reset the current `_plugin_context` context variable."""
+    _plugin_context.reset(token)
+
+
+def get_botocore_plugins():
+    context = get_plugin_context()
+    if context is not None:
+        plugins = context.plugins
+        if plugins is None:
+            context.plugins = os.environ.get('BOTOCORE_EXPERIMENTAL__PLUGINS')
+        else:
+            return plugins
+    return os.environ.get('BOTOCORE_EXPERIMENTAL__PLUGINS')
+
+
+def load_client_plugins(client, plugins):
+    for plugin_name, module_name in plugins.items():
+        log.debug(
+            "Importing client plugin %s from module %s",
+            plugin_name,
+            module_name,
+        )
+        try:
+            module = importlib.import_module(module_name)
+            module.initialize_client_plugin(client)
+        except ModuleNotFoundError:
+            log.debug(
+                "Failed to locate the following plugin module: %s.",
+                plugin_name,
+            )
+        except Exception as e:
+            log.debug(
+                "Error raised during the loading of %s: %s", plugin_name, e
+            )

botocore/session.py

@@ -68,6 +68,7 @@
 from botocore.loaders import create_loader
 from botocore.model import ServiceModel
 from botocore.parsers import ResponseParserFactory
+from botocore.plugin import get_botocore_plugins, load_client_plugins
 from botocore.regions import EndpointResolver
 from botocore.useragent import UserAgentString
 from botocore.utils import (
@@ -1039,6 +1040,7 @@ def create_client(
         monitor = self._get_internal_component('monitor')
         if monitor is not None:
             monitor.register(client.meta.events)
+        self._register_client_plugins(client)
         return client
 
     def _resolve_region_name(self, region_name, config):
@@ -1164,6 +1166,24 @@ def _get_ignored_credentials(self, aws_session_token, aws_account_id):
             credential_inputs.append('aws_account_id')
         return ', '.join(credential_inputs) if credential_inputs else None
 
+    def _register_client_plugins(self, client):
+        plugins_list = get_botocore_plugins()
+        if plugins_list == "DISABLED" or not plugins_list:
+            return
+
+        client_plugins = {}
+        for plugin in plugins_list.split(','):
+            try:
+                name, module = [part.strip() for part in plugin.split('=')]
+                client_plugins[name] = module
+            except ValueError:
+                logger.warning(
+                    f"Invalid plugin format: {plugin}. Expected 'name=module'"
+                )
+
+        if client_plugins:
+            load_client_plugins(client, client_plugins)
+
 
 class ComponentLocator:
     """Service locator for session components."""

botocore/tokens.py

@@ -32,6 +32,7 @@
     CachedProperty,
     JSONFileCache,
     SSOTokenLoader,
+    create_nested_client,
     get_token_from_environment,
 )
 
@@ -256,7 +257,7 @@ def _client(self):
             region_name=self._sso_config["sso_region"],
             signature_version=UNSIGNED,
         )
-        return self._session.create_client("sso-oidc", config=config)
+        return create_nested_client(self._session, "sso-oidc", config=config)
 
     def _attempt_create_token(self, token):
         response = self._client.create_token(

botocore/utils.py

@@ -87,6 +87,11 @@
     UnsupportedS3ControlArnError,
     UnsupportedS3ControlConfigurationError,
 )
+from botocore.plugin import (
+    PluginContext,
+    reset_plugin_context,
+    set_plugin_context,
+)
 
 logger = logging.getLogger(__name__)
 DEFAULT_METADATA_SERVICE_TIMEOUT = 1
@@ -364,6 +369,18 @@ def is_global_accesspoint(context):
     return is_global
 
 
+def create_nested_client(session, service_name, **kwargs):
+    # If a client is created from within a plugin based on the environment variable,
+    # an infinite loop could arise.  Any clients created from within another client
+    # must use this method to prevent infinite loops.
+    ctx = PluginContext(plugins="DISABLED")
+    token = set_plugin_context(ctx)
+    try:
+        return session.create_client(service_name, **kwargs)
+    finally:
+        reset_plugin_context(token)
+
+
 class _RetriesExceededError(Exception):
     """Internal exception used when the number of retries are exceeded."""
 

tests/functional/plugins/disabled_test_plugin.py

@@ -0,0 +1,29 @@
+# Copyright 2025 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. A copy of
+# the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file 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.
+class DisabledPluginModule:
+    def __init__(self):
+        self.invocations = 0
+
+    def register_event(self, client):
+        client.meta.events.register('before-call', self.increment_invocations)
+
+    def increment_invocations(self, **kwargs):
+        self.invocations += 1
+
+
+plugin_instance = DisabledPluginModule()
+
+
+def initialize_client_plugin(client):
+    plugin_instance.register_event(client)
+    return plugin_instance

tests/functional/plugins/first_of_two_plugins.py

@@ -0,0 +1,29 @@
+# Copyright 2025 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. A copy of
+# the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file 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.
+class FirstPluginModule:
+    def __init__(self):
+        self.invocations = 0
+
+    def register_event(self, client):
+        client.meta.events.register('before-call', self.increment_invocations)
+
+    def increment_invocations(self, **kwargs):
+        self.invocations += 1
+
+
+plugin_instance = FirstPluginModule()
+
+
+def initialize_client_plugin(client):
+    plugin_instance.register_event(client)
+    return plugin_instance

tests/functional/plugins/recursive_plugin_test_module.py

@@ -0,0 +1,46 @@
+# Copyright 2025 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. A copy of
+# the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file 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.
+from botocore import utils
+from botocore.session import get_session
+from tests import ClientHTTPStubber
+
+
+class RecursivePluginModule:
+    """A mock recursive plugin for testing nested client creation."""
+
+    def __init__(self):
+        self.called = False
+
+    def register_event(self, client):
+        client.meta.events.register(
+            'before-call.dynamodb.*', self.create_client
+        )
+
+    def create_client(self, **kwargs):
+        self.called = True
+        session = get_session()
+        session.set_credentials('key', 'secret')
+        client = utils.create_nested_client(
+            session, "dynamodb", region_name="us-west-2"
+        )
+        with ClientHTTPStubber(client) as http_stubber:
+            http_stubber.add_response(status=200, body=b'')
+            client.list_tables()
+
+
+plugin_instance = RecursivePluginModule()
+
+
+def initialize_client_plugin(client):
+    plugin_instance.register_event(client)
+    return plugin_instance

tests/functional/plugins/second_of_two_plugins.py

@@ -0,0 +1,29 @@
+# Copyright 2025 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. A copy of
+# the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file 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.
+class SecondPluginModule:
+    def __init__(self):
+        self.invocations = 0
+
+    def register_event(self, client):
+        client.meta.events.register('before-call', self.increment_invocations)
+
+    def increment_invocations(self, **kwargs):
+        self.invocations += 1
+
+
+plugin_instance = SecondPluginModule()
+
+
+def initialize_client_plugin(client):
+    plugin_instance.register_event(client)
+    return plugin_instance

tests/functional/plugins/simple_plugin_test_module.py

@@ -0,0 +1,31 @@
+# Copyright 2025 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. A copy of
+# the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file 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.
+class TestPluginModule:
+    """A mock plugin module for testing client plugin loading."""
+
+    def __init__(self):
+        self.events_seen = []
+
+    def register_event(self, client):
+        client.meta.events.register('before-call', self.increment_calls)
+
+    def increment_calls(self, **kwargs):
+        self.events_seen.append(kwargs)
+
+
+plugin_instance = TestPluginModule()
+
+
+def initialize_client_plugin(client):
+    plugin_instance.register_event(client)
+    return plugin_instance