From 4b461df96f62351b3c684836ef0fbd9d1c5d7960 Mon Sep 17 00:00:00 2001 From: Max Z Date: Thu, 15 Sep 2022 15:40:07 +0100 Subject: [PATCH 1/2] feature: add ability to give boto extra args in config Signed-off-by: Max Z --- sdk/python/feast/infra/registry/s3.py | 3 ++- sdk/python/feast/repo_config.py | 3 +++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/sdk/python/feast/infra/registry/s3.py b/sdk/python/feast/infra/registry/s3.py index d3772910f5..b7bf8b14a3 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.boto_extra_args or {} self.s3_client = boto3.resource( "s3", endpoint_url=os.environ.get("FEAST_S3_ENDPOINT_URL") @@ -77,4 +78,4 @@ 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 47a5ae321d..98d67a2ace 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() """ + boto_extra_args: Optional[Dict[str, str]] + """ Dict[str, str]: Extra arguments to pass when writing the registry file to S3. """ + class RepoConfig(FeastBaseModel): """Repo config. Typically loaded from `feature_store.yaml`""" From 74a7422c8449dbd411f2775e9951806ff3bc3e02 Mon Sep 17 00:00:00 2001 From: Max Zwiessele Date: Thu, 15 Sep 2022 17:26:18 +0100 Subject: [PATCH 2/2] rename: s3_additional_kwargs Signed-off-by: Max Z --- sdk/python/feast/infra/registry/s3.py | 6 ++++-- sdk/python/feast/repo_config.py | 4 ++-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/sdk/python/feast/infra/registry/s3.py b/sdk/python/feast/infra/registry/s3.py index b7bf8b14a3..0a94c942e1 100644 --- a/sdk/python/feast/infra/registry/s3.py +++ b/sdk/python/feast/infra/registry/s3.py @@ -25,7 +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.boto_extra_args or {} + 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") @@ -78,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._boto_extra_args) + 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 98d67a2ace..9088b2d996 100644 --- a/sdk/python/feast/repo_config.py +++ b/sdk/python/feast/repo_config.py @@ -112,8 +112,8 @@ 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() """ - boto_extra_args: Optional[Dict[str, str]] - """ Dict[str, str]: Extra arguments to pass when writing the registry file to S3. """ + 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):