feat(s3): tag S3 client user-agent for attribution

Signed-off-by: Gonzalo Peña-Castellanos <goanpeca@gmail.com>

Author: goanpeca

python/whylogs/api/_s3_client.py

@@ -0,0 +1,39 @@
+"""Shared construction for the boto3 S3 clients that whylogs builds itself.
+
+Works with Amazon S3 and any S3-compatible object store (for example
+Cloudflare R2, MinIO, or Backblaze B2) by passing an ``endpoint_url`` and
+matching credentials.
+"""
+
+from typing import Any, Optional
+
+import boto3
+from botocore.client import BaseClient
+from botocore.config import Config
+
+
+def _whylogs_version() -> str:
+    try:
+        from importlib import metadata
+    except ImportError:  # Python < 3.8
+        import importlib_metadata as metadata  # type: ignore
+
+    try:
+        return metadata.version("whylogs")
+    except metadata.PackageNotFoundError:  # type: ignore
+        return "dev"
+
+
+def build_s3_client(config: Optional[Config] = None, **kwargs: Any) -> BaseClient:
+    """Build an S3 client for clients whylogs creates itself.
+
+    The whylogs identifier is appended to ``user_agent_extra`` so it is added to,
+    not substituted for, any value already present on a caller-supplied ``config``.
+    A custom ``endpoint_url`` for an S3-compatible store may be passed through
+    ``kwargs``.
+    """
+    suffix = f"whylogs/{_whylogs_version()}"
+    existing = getattr(config, "user_agent_extra", None)
+    user_agent_extra = f"{existing} {suffix}" if existing else suffix
+    config = (config or Config()).merge(Config(user_agent_extra=user_agent_extra))
+    return boto3.client("s3", config=config, **kwargs)

python/whylogs/api/reader/s3.py

@@ -1,10 +1,11 @@
 from tempfile import NamedTemporaryFile
 from typing import Optional
 
-import boto3
 from botocore.client import BaseClient
+from botocore.config import Config
 
 from whylogs import ResultSet
+from whylogs.api._s3_client import build_s3_client
 from whylogs.api.reader.reader import Reader
 
 
@@ -15,6 +16,10 @@ class S3Reader(Reader):
     >**IMPORTANT**: In order to correctly connect to your Amazon S3 bucket, make sure you have
     the following environment variables set: `[AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY]`.
 
+    The same reader targets any S3-compatible object store by passing an ``endpoint_url``
+    (and matching credentials). To reuse a fully configured client instead, pass it as
+    ``s3_client`` and it is used as provided.
+
     Parameters
     ----------
     bucket_name: str, optional
@@ -23,6 +28,13 @@ class S3Reader(Reader):
     object_name: str, optional
         The s3's object name. It basically states the location where the file goes to.
         Also made optional, so it can be defined through the `option` method
+    s3_client: BaseClient, optional
+        A boto3 S3 client. When omitted, a client is built internally.
+    endpoint_url: str, optional
+        Endpoint URL of an S3-compatible object store. Leave unset for AWS S3.
+        Only applies to the internally built client.
+    config: botocore.config.Config, optional
+        Extra botocore configuration merged into the internally built client.
 
     Examples
     --------
@@ -40,8 +52,10 @@ def __init__(
         object_name: Optional[str] = None,
         bucket_name: Optional[str] = None,
         s3_client: Optional[BaseClient] = None,
+        endpoint_url: Optional[str] = None,
+        config: Optional[Config] = None,
     ):
-        self.s3_client = s3_client or boto3.client("s3")
+        self.s3_client = s3_client or build_s3_client(config=config, endpoint_url=endpoint_url)
         self.object_name = object_name or None
         self.bucket_name = bucket_name or ""
 

python/whylogs/api/writer/s3.py

@@ -2,10 +2,11 @@
 import os
 from typing import Any, List, Optional, Tuple, Union
 
-import boto3
 from botocore.client import BaseClient
+from botocore.config import Config
 from botocore.exceptions import ClientError
 
+from whylogs.api._s3_client import build_s3_client
 from whylogs.api.writer import Writer
 from whylogs.api.writer.writer import _Writable
 from whylogs.core.utils import deprecated_alias
@@ -20,11 +21,16 @@ class S3Writer(Writer):
     >**IMPORTANT**: In order to correctly connect to your Amazon S3 bucket, make sure you have
     the following environment variables set: `[AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY]`
 
+    The same writer targets any S3-compatible object store by passing an ``endpoint_url``
+    (and matching credentials). To reuse a fully configured client instead, pass it as
+    ``s3_client`` and it is used as provided.
+
     Parameters
     ----------
     s3_client: BaseClient, optional
         The s3 client used to authenticate and perform operations on the s3 bucket.
-        Should be a BaseClient from the boto3 library
+        Should be a BaseClient from the boto3 library. When omitted, a client is built
+        internally.
     base_prefix: str, optional
         The base file prefix for s3, in order to organize. A placeholder 'profile' will take place if None is provided.
     bucket_name: str, optional
@@ -33,6 +39,11 @@ class S3Writer(Writer):
     object_name: str, optional
         The s3's object name. It basically states the location where the file goes to.
         Also made optional, so it can be defined through the `option` method
+    endpoint_url: str, optional
+        Endpoint URL of an S3-compatible object store. Leave unset for AWS S3.
+        Only applies to the internally built client.
+    config: botocore.config.Config, optional
+        Extra botocore configuration merged into the internally built client.
     Returns
     -------
         None
@@ -49,6 +60,15 @@ class S3Writer(Writer):
     profile.writer("s3").option(bucket_name="my_bucket").write()
     ```
 
+    Writing to an S3-compatible endpoint:
+
+    ```python
+    profile.writer("s3").option(
+        bucket_name="my_bucket",
+        endpoint_url="https://<s3-compatible-endpoint>",
+    ).write()
+    ```
+
     """
 
     def __init__(
@@ -57,8 +77,10 @@ def __init__(
         base_prefix: Optional[str] = None,
         bucket_name: Optional[str] = None,
         object_name: Optional[str] = None,
+        endpoint_url: Optional[str] = None,
+        config: Optional[Config] = None,
     ):
-        self.s3_client = s3_client or boto3.client("s3")
+        self.s3_client = s3_client or build_s3_client(config=config, endpoint_url=endpoint_url)
         self.base_prefix = base_prefix or "profile"
         self.bucket_name = bucket_name or ""
         self.object_name = object_name or None