Skip to content

Commit 2ce3181

Browse files
fix: address CodeRabbit review feedback for S3 store
- Close streaming body properly to prevent connection leaks - Add regional bucket creation support for non-us-east-1 regions Co-authored-by: William Easton <strawgate@users.noreply.github.com>
1 parent 05c62a4 commit 2ce3181

File tree

1 file changed

+15
-3
lines changed
  • key-value/key-value-aio/src/key_value/aio/stores/s3

1 file changed

+15
-3
lines changed

key-value/key-value-aio/src/key_value/aio/stores/s3/store.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -206,7 +206,18 @@ async def _setup(self) -> None:
206206
import contextlib
207207

208208
with contextlib.suppress(self._connected_client.exceptions.BucketAlreadyOwnedByYou): # pyright: ignore[reportUnknownMemberType]
209-
await self._connected_client.create_bucket(Bucket=self._bucket_name) # pyright: ignore[reportUnknownMemberType]
209+
# Build create_bucket parameters
210+
create_params: dict[str, Any] = {"Bucket": self._bucket_name}
211+
212+
# Get region from client metadata
213+
region_name = getattr(self._connected_client.meta, "region_name", None) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
214+
215+
# For regions other than us-east-1, we need to specify LocationConstraint
216+
# Skip this for custom endpoints (LocalStack, MinIO) which may not support it
217+
if region_name and region_name != "us-east-1" and not self._endpoint_url:
218+
create_params["CreateBucketConfiguration"] = {"LocationConstraint": region_name}
219+
220+
await self._connected_client.create_bucket(**create_params) # pyright: ignore[reportUnknownMemberType]
210221
else:
211222
# Re-raise authentication, permission, or other errors
212223
raise
@@ -253,8 +264,9 @@ async def _get_managed_entry(self, *, key: str, collection: str) -> ManagedEntry
253264
Key=s3_key,
254265
)
255266

256-
# Read the object body
257-
body_bytes = await response["Body"].read() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
267+
# Read the object body and ensure the streaming body is closed
268+
async with response["Body"] as stream: # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
269+
body_bytes = await stream.read() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
258270
json_value = body_bytes.decode("utf-8") # pyright: ignore[reportUnknownMemberType]
259271

260272
# Deserialize to ManagedEntry

0 commit comments

Comments
 (0)