Sergey Antropoff
90 lines
3.2 KiB
Python
90 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
from contextlib import asynccontextmanager
|
|
from typing import Any, AsyncIterator
|
|
|
|
from aiobotocore.session import get_session
|
|
|
|
from app.config import get_settings
|
|
|
|
|
|
class ObjectStorage:
|
|
def __init__(self) -> None:
|
|
self.settings = get_settings()
|
|
|
|
def _client_kwargs(self) -> dict[str, Any]:
|
|
return {
|
|
"endpoint_url": self.settings.s3_endpoint_url,
|
|
"aws_access_key_id": self.settings.s3_access_key,
|
|
"aws_secret_access_key": self.settings.s3_secret_key,
|
|
"region_name": self.settings.s3_region,
|
|
"use_ssl": self.settings.s3_use_ssl,
|
|
}
|
|
|
|
@asynccontextmanager
|
|
async def client(self) -> AsyncIterator[Any]:
|
|
session = get_session()
|
|
async with session.create_client("s3", **self._client_kwargs()) as client:
|
|
yield client
|
|
|
|
async def ensure_bucket(self) -> None:
|
|
if not self.settings.s3_create_bucket:
|
|
return
|
|
async with self.client() as client:
|
|
try:
|
|
await client.head_bucket(Bucket=self.settings.s3_bucket)
|
|
except Exception:
|
|
await client.create_bucket(Bucket=self.settings.s3_bucket)
|
|
|
|
async def put_bytes(self, key: str, data: bytes, content_type: str = "application/octet-stream") -> None:
|
|
async with self.client() as client:
|
|
await client.put_object(
|
|
Bucket=self.settings.s3_bucket,
|
|
Key=key,
|
|
Body=data,
|
|
ContentType=content_type,
|
|
)
|
|
|
|
async def get_bytes(self, key: str) -> bytes:
|
|
async with self.client() as client:
|
|
resp = await client.get_object(Bucket=self.settings.s3_bucket, Key=key)
|
|
async with resp["Body"] as stream:
|
|
return await stream.read()
|
|
|
|
async def delete(self, key: str) -> None:
|
|
async with self.client() as client:
|
|
await client.delete_object(Bucket=self.settings.s3_bucket, Key=key)
|
|
|
|
async def delete_prefix(self, prefix: str = "wraps/") -> int:
|
|
"""Delete all objects under prefix. Returns number of deleted keys."""
|
|
deleted = 0
|
|
async with self.client() as client:
|
|
token: str | None = None
|
|
while True:
|
|
kwargs: dict[str, Any] = {
|
|
"Bucket": self.settings.s3_bucket,
|
|
"Prefix": prefix,
|
|
"MaxKeys": 1000,
|
|
}
|
|
if token:
|
|
kwargs["ContinuationToken"] = token
|
|
resp = await client.list_objects_v2(**kwargs)
|
|
contents = resp.get("Contents") or []
|
|
if contents:
|
|
# delete_objects accepts up to 1000 keys
|
|
await client.delete_objects(
|
|
Bucket=self.settings.s3_bucket,
|
|
Delete={
|
|
"Objects": [{"Key": obj["Key"]} for obj in contents],
|
|
"Quiet": True,
|
|
},
|
|
)
|
|
deleted += len(contents)
|
|
if not resp.get("IsTruncated"):
|
|
break
|
|
token = resp.get("NextContinuationToken")
|
|
return deleted
|
|
|
|
|
|
storage = ObjectStorage()
|