feat(aws): add S3 endpoint and user-agent support

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

Author: goanpeca

great_expectations/compatibility/aws.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+from importlib import metadata
+from typing import Any, Dict
+
 from great_expectations.compatibility.not_imported import NotImported
 
 BOTO_NOT_IMPORTED = NotImported(
@@ -32,6 +35,40 @@
 except ImportError:
     exceptions = BOTO_NOT_IMPORTED
 
+
+def _get_distribution_version() -> str:
+    try:
+        return metadata.version("great_expectations")
+    except metadata.PackageNotFoundError:
+        return "dev"
+
+
+def get_s3_boto3_options(boto3_options: Dict[str, Any]) -> Dict[str, Any]:
+    """Return boto3 client options with a ``great-expectations`` agent suffix.
+
+    Works with Amazon S3 and any S3-compatible object store (for example
+    Backblaze B2, Cloudflare R2, or MinIO). A caller-supplied ``config`` and
+    ``endpoint_url`` are preserved; the suffix is appended to an existing agent
+    string rather than replacing it.
+    """
+    suffix = f"great-expectations/{_get_distribution_version()}"
+    options = dict(boto3_options)
+
+    if not Config:
+        return options
+
+    config = options.get("config")
+    existing = getattr(config, "user_agent_extra", None) if config else None
+    user_agent_extra = f"{existing} {suffix}" if existing else suffix
+
+    if config:
+        options["config"] = config.merge(Config(user_agent_extra=user_agent_extra))
+    else:
+        options["config"] = Config(user_agent_extra=user_agent_extra)
+
+    return options
+
+
 try:
     import sqlalchemy_redshift
 except ImportError:

great_expectations/datasource/fluent/pandas_s3_datasource.py

@@ -68,7 +68,7 @@ def _get_s3_client(self) -> BaseClient:
                     "boto3_options", {}
                 )
                 try:
-                    s3_client = aws.boto3.client("s3", **boto3_options)
+                    s3_client = aws.boto3.client("s3", **aws.get_s3_boto3_options(boto3_options))
                 except Exception as e:
                     # Failure to create "s3_client" is most likely due invalid "boto3_options" dictionary.  # noqa: E501 # FIXME CoP
                     raise PandasS3DatasourceError(  # noqa: TRY003 # FIXME CoP

great_expectations/datasource/fluent/spark_s3_datasource.py

@@ -70,7 +70,7 @@ def _get_s3_client(self) -> BaseClient:
                     "boto3_options", {}
                 )
                 try:
-                    s3_client = aws.boto3.client("s3", **boto3_options)
+                    s3_client = aws.boto3.client("s3", **aws.get_s3_boto3_options(boto3_options))
                 except Exception as e:
                     # Failure to create "s3_client" is most likely due invalid "boto3_options" dictionary.  # noqa: E501 # FIXME CoP
                     raise SparkS3DatasourceError(  # noqa: TRY003 # FIXME CoP

great_expectations/execution_engine/pandas_execution_engine.py

@@ -167,7 +167,7 @@ def _instantiate_azure_client(self) -> None:
     def _instantiate_s3_client(self) -> None:
         # If s3_client was passed in (from data source) use it, otherwise create our own
         self._s3 = self._config.get("s3_client") or aws.boto3.client(
-            "s3", **self.config.get("boto3_options", {})
+            "s3", **aws.get_s3_boto3_options(self.config.get("boto3_options", {}))
         )
 
     def _instantiate_gcs_client(self) -> None: