diff --git a/sdk/python/feast/infra/registry/s3.py b/sdk/python/feast/infra/registry/s3.py index d3772910f56..0a94c942e18 100644 --- a/sdk/python/feast/infra/registry/s3.py +++ b/sdk/python/feast/infra/registry/s3.py @@ -25,6 +25,7 @@ def __init__(self, registry_config: RegistryConfig, repo_path: Path): self._uri = urlparse(uri) self._bucket = self._uri.hostname self._key = self._uri.path.lstrip("/") + self._boto_extra_args = registry_config.s3_additional_kwargs or {} self.s3_client = boto3.resource( "s3", endpoint_url=os.environ.get("FEAST_S3_ENDPOINT_URL") @@ -77,4 +78,6 @@ def _write_registry(self, registry_proto: RegistryProto): file_obj = TemporaryFile() file_obj.write(registry_proto.SerializeToString()) file_obj.seek(0) - self.s3_client.Bucket(self._bucket).put_object(Body=file_obj, Key=self._key) + self.s3_client.Bucket(self._bucket).put_object( + Body=file_obj, Key=self._key, **self._boto_extra_args + ) diff --git a/sdk/python/feast/repo_config.py b/sdk/python/feast/repo_config.py index 47a5ae321d9..9088b2d9962 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -112,6 +112,9 @@ class RegistryConfig(FeastBaseModel): set to infinity by setting TTL to 0 seconds, which means the cache will only be loaded once and will never expire. Users can manually refresh the cache by calling feature_store.refresh_registry() """ + s3_additional_kwargs: Optional[Dict[str, str]] + """ Dict[str, str]: Extra arguments to pass to boto3 when writing the registry file to S3. """ + class RepoConfig(FeastBaseModel): """Repo config. Typically loaded from `feature_store.yaml`"""