feat: Add ability to give boto extra args for registry config (#3219)

* feature: add ability to give boto extra args in config

Signed-off-by: Max Z <[email protected]>

* rename: s3_additional_kwargs

Signed-off-by: Max Z <[email protected]>

Signed-off-by: Max Z <[email protected]>

Author: mzwiessele

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
+        )

sdk/python/feast/repo_config.py

@@ -113,6 +113,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`"""