Build client hooks (#4289)

Co-authored-by: Alexei Mochalov <[email protected]>

api/python/quilt3/__init__.py

@@ -9,7 +9,7 @@
 
 __version__ = Path(Path(__file__).parent, "VERSION").read_text().strip()
 
-from . import admin
+from . import admin, hooks
 from .api import (
     config,
     copy,

api/python/quilt3/data_transfer.py

@@ -41,7 +41,7 @@
 )
 from tqdm import tqdm
 
-from . import util
+from . import hooks, util
 from .session import get_boto3_session
 from .util import DISABLE_TQDM, PhysicalKey, QuiltException
 
@@ -153,15 +153,26 @@ def find_correct_client(self, api_type, bucket, param_dict):
     def get_boto_session(self):
         return get_boto3_session()
 
+    @staticmethod
+    def _build_client_base(session, client_kwargs):
+        return session.client('s3', **client_kwargs)
+
     def _build_client(self, is_unsigned):
         session = self.get_boto_session()
         conf_kwargs = {
             "max_pool_connections": MAX_CONCURRENCY,
         }
         if is_unsigned(session):
             conf_kwargs["signature_version"] = UNSIGNED
-
-        return session.client('s3', config=Config(**conf_kwargs))
+        client_kwargs = {
+            "config": Config(**conf_kwargs),
+        }
+        hook = hooks.get_build_s3_client_hook()
+        return (
+            self._build_client_base(session, client_kwargs)
+            if hook is None else
+            hook(self._build_client_base, session, client_kwargs)
+        )
 
     def _build_standard_client(self):
         s3_client = self._build_client(lambda session: session.get_credentials() is None)

api/python/quilt3/hooks.py

@@ -0,0 +1,62 @@
+import typing as T
+
+import boto3
+
+
+class BuildClientBase(T.Protocol):
+    def __call__(self, session: boto3.Session, client_kwargs: dict[str, T.Any], **kwargs): ...
+
+
+class BuildClientHook(T.Protocol):
+    def __call__(
+        self, build_client_base: BuildClientBase, session: boto3.Session, client_kwargs: dict[str, T.Any], **kwargs
+    ): ...
+
+
+_build_client_hook = None
+
+
+def get_build_s3_client_hook() -> T.Optional[BuildClientHook]:
+    """
+    Return build S3 client hook.
+    """
+
+    return _build_client_hook
+
+
+def set_build_s3_client_hook(hook: T.Optional[BuildClientHook]) -> T.Optional[BuildClientHook]:
+    """
+    Set build S3 client hook.
+
+    Example for overriding `ServerSideEncryption` parameter for certain S3 operations:
+
+    ```python
+    def event_handler(params, **kwargs):
+        # Be mindful with parameters you set here.
+        # Specifically it's not recommended to override/delete already set parameters
+        # because that can break quilt3 logic.
+        params.setdefault("ServerSideEncryption", "AES256")
+
+    def hook(build_client_base, session, client_kwargs, **kwargs):
+        client = build_client_base(session, client_kwargs, **kwargs)
+        # Docs for boto3 events system we use below:
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/events.html
+        for op in (
+            "CreateMultipartUpload",
+            "CopyObject",
+            "PutObject",
+        ):
+            client.meta.events.register(f"before-parameter-build.s3.{op}", event_handler)
+        return client
+    ```
+
+    Args:
+        hook: Build client hook.
+
+    Returns:
+        Old build client hook.
+    """
+    global _build_client_hook
+    old_hook = _build_client_hook
+    _build_client_hook = hook
+    return old_hook

api/python/tests/test_hooks.py

@@ -0,0 +1,45 @@
+from botocore.stub import Stubber
+
+from quilt3 import data_transfer, hooks, util
+
+
+def test_build_client_hooks():
+    try:
+        assert hooks.get_build_s3_client_hook() is None
+
+        stubber = None
+
+        def event_handler(params, **kwargs):
+            params.setdefault("ServerSideEncryption", "AES256")
+
+        def hook(build_client_base, session, client_kwargs):
+            client = build_client_base(session, client_kwargs)
+            # use register_first and * to ensure that our hook runs before the stubber's one
+            client.meta.events.register_first("before-parameter-build.*.*", event_handler)
+
+            nonlocal stubber
+            stubber = Stubber(client)
+            stubber.add_response(
+                "put_object",
+                {},
+                {
+                    "Bucket": "bucket",
+                    "Key": "key",
+                    "Body": b"data",
+                    "ServerSideEncryption": "AES256",
+                },
+            )
+            stubber.activate()
+
+            return client
+
+        assert hooks.set_build_s3_client_hook(hook) is None
+        assert hooks.get_build_s3_client_hook() is hook
+
+        data_transfer.put_bytes(b"data", util.PhysicalKey("bucket", "key", None))
+
+        assert stubber is not None
+        stubber.assert_no_pending_responses()
+    finally:
+        hooks.set_build_s3_client_hook(None)
+        assert hooks.get_build_s3_client_hook() is None

docs/CHANGELOG.md

@@ -13,6 +13,12 @@ Entries inside each section should be ordered by type:
 ## CLI
 !-->
 
+# unreleased - YYYY-MM-DD
+
+## Python API
+
+* [Added] `quilt3.hooks`: `set_build_s3_client_hook()` function for customizing S3 client ([#4289](https://github.com/quiltdata/quilt/pull/4289))
+
 # 6.2.0 - 2025-01-14
 
 ## CLI

docs/api-reference/Hooks.md

@@ -0,0 +1,40 @@
+
+# get\_build\_s3\_client\_hook() -> Optional[quilt3.hooks.BuildClientHook]  {#get\_build\_s3\_client\_hook}
+
+Return build S3 client hook.
+
+
+# set\_build\_s3\_client\_hook(hook: Optional[quilt3.hooks.BuildClientHook]) -> Optional[quilt3.hooks.BuildClientHook]  {#set\_build\_s3\_client\_hook}
+
+Set build S3 client hook.
+
+Example for overriding `ServerSideEncryption` parameter for certain S3 operations:
+
+```python
+def event_handler(params, **kwargs):
+    # Be mindful with parameters you set here.
+    # Specifically it's not recommended to override/delete already set parameters
+    # because that can break quilt3 logic.
+    params.setdefault("ServerSideEncryption", "AES256")
+
+def hook(build_client_base, session, client_kwargs, **kwargs):
+    client = build_client_base(session, client_kwargs, **kwargs)
+    # Docs for boto3 events system we use below:
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/events.html
+    for op in (
+        "CreateMultipartUpload",
+        "CopyObject",
+        "PutObject",
+    ):
+        client.meta.events.register(f"before-parameter-build.s3.{op}", event_handler)
+    return client
+```
+
+__Arguments__
+
+* __hook__:  Build client hook.
+
+__Returns__
+
+Old build client hook.
+

gendocs/pydocmd.yml

@@ -38,6 +38,12 @@ generate:
   - quilt3.admin.sso_config+
   - quilt3.admin.tabulator+
 
+- Hooks.md:
+  # don't do quilt3.hooks+ because pydocmd renders BuildClientBase and BuildClientHook
+  # in uninformative way
+  - quilt3.hooks.get_build_s3_client_hook
+  - quilt3.hooks.set_build_s3_client_hook
+
 # MkDocs pages configuration. The `<<` operator is sugar added by pydocmd
 # that allows you to use an external Markdown file (eg. your project's README)
 # in the documentation. The path must be relative to current working directory.