Init

app/services/storage.py

@@ -0,0 +1,89 @@
+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()