diff --git a/sdk/python/feast/infra/compute_engines/spark/utils.py b/sdk/python/feast/infra/compute_engines/spark/utils.py index 8c84c9f17a6..fb84b698dbc 100644 --- a/sdk/python/feast/infra/compute_engines/spark/utils.py +++ b/sdk/python/feast/infra/compute_engines/spark/utils.py @@ -9,7 +9,11 @@ from pyspark.sql import SparkSession from feast.infra.common.serde import SerializedArtifacts -from feast.utils import _convert_arrow_to_proto, _run_pyarrow_field_mapping +from feast.utils import ( + _convert_arrow_to_proto, + _run_pyarrow_field_mapping, + get_user_agent, +) try: import boto3 @@ -88,6 +92,7 @@ def _ensure_s3a_event_log_dir(spark_config: Dict[str, str]) -> None: config=BotoConfig( signature_version="s3v4", s3={"addressing_style": addressing_style}, + user_agent_extra=get_user_agent(), ), ) resp = s3.list_objects_v2(Bucket=bucket, Prefix=prefix, MaxKeys=1) diff --git a/sdk/python/feast/infra/offline_stores/ibis.py b/sdk/python/feast/infra/offline_stores/ibis.py index 3aebc4e903e..e3f6b1585a8 100644 --- a/sdk/python/feast/infra/offline_stores/ibis.py +++ b/sdk/python/feast/infra/offline_stores/ibis.py @@ -477,8 +477,15 @@ def point_in_time_join( def list_s3_files(path: str, endpoint_url: str) -> List[str]: import boto3 + from botocore.config import Config - s3 = boto3.client("s3", endpoint_url=endpoint_url) + from feast.utils import get_user_agent + + s3 = boto3.client( + "s3", + endpoint_url=endpoint_url, + config=Config(user_agent_extra=get_user_agent()), + ) if path.startswith("s3://"): path = path[len("s3://") :] bucket, prefix = path.split("/", 1) diff --git a/sdk/python/feast/infra/registry/s3.py b/sdk/python/feast/infra/registry/s3.py index 246636600a4..3c4b31719af 100644 --- a/sdk/python/feast/infra/registry/s3.py +++ b/sdk/python/feast/infra/registry/s3.py @@ -10,10 +10,11 @@ from feast.infra.registry.registry_store import RegistryStore from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto from feast.repo_config import RegistryConfig -from feast.utils import _utc_now +from feast.utils import _utc_now, get_user_agent try: import boto3 + from botocore.config import Config except ImportError as e: from feast.errors import FeastExtrasDependencyImportError @@ -28,8 +29,11 @@ def __init__(self, registry_config: RegistryConfig, repo_path: Path): self._key = self._uri.path.lstrip("/") self._boto_extra_args = registry_config.s3_additional_kwargs or {} + # FEAST_S3_ENDPOINT_URL may point at Amazon S3 or an S3-compatible endpoint. self.s3_client = boto3.resource( - "s3", endpoint_url=os.environ.get("FEAST_S3_ENDPOINT_URL") + "s3", + endpoint_url=os.environ.get("FEAST_S3_ENDPOINT_URL"), + config=Config(user_agent_extra=get_user_agent()), ) def get_registry_proto(self): diff --git a/sdk/python/feast/infra/utils/aws_utils.py b/sdk/python/feast/infra/utils/aws_utils.py index 39fa815f7e3..2d384874952 100644 --- a/sdk/python/feast/infra/utils/aws_utils.py +++ b/sdk/python/feast/infra/utils/aws_utils.py @@ -53,7 +53,10 @@ def get_s3_resource(aws_region: str): """ Get the S3 resource for the given AWS region. """ - return boto3.resource("s3", config=Config(region_name=aws_region)) + return boto3.resource( + "s3", + config=Config(region_name=aws_region, user_agent_extra=get_user_agent()), + ) def get_bucket_and_key(s3_path: str) -> Tuple[str, str]: @@ -716,7 +719,10 @@ def get_account_id() -> str: def list_s3_files(aws_region: str, path: str) -> List[str]: - s3 = boto3.client("s3", config=Config(region_name=aws_region)) + s3 = boto3.client( + "s3", + config=Config(region_name=aws_region, user_agent_extra=get_user_agent()), + ) if path.startswith("s3://"): path = path[len("s3://") :] bucket, prefix = path.split("/", 1)