You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
69 lines
2.5 KiB
69 lines
2.5 KiB
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import anyio
|
|
import boto3
|
|
from botocore.exceptions import ClientError
|
|
|
|
|
|
class S3StorageAdapter:
|
|
def __init__(self, config: dict[str, Any]) -> None:
|
|
self.config = config
|
|
|
|
async def check(self) -> dict[str, str]:
|
|
client = self._create_client()
|
|
|
|
def _check() -> dict[str, str]:
|
|
client.head_bucket(Bucket=self.config["bucket"])
|
|
return {"status": "healthy", "detail": self.config["bucket"]}
|
|
|
|
return await anyio.to_thread.run_sync(_check)
|
|
|
|
async def put_file(self, source_path: str, storage_key: str) -> dict[str, str | None]:
|
|
client = self._create_client()
|
|
|
|
def _put() -> dict[str, str | None]:
|
|
extra_args: dict[str, Any] = {}
|
|
client.upload_file(source_path, self.config["bucket"], storage_key, ExtraArgs=extra_args)
|
|
metadata = client.head_object(Bucket=self.config["bucket"], Key=storage_key)
|
|
return {"etag": metadata.get("ETag"), "storage_key": storage_key}
|
|
|
|
return await anyio.to_thread.run_sync(_put)
|
|
|
|
async def download_file(self, storage_key: str, destination_path: Path) -> Path:
|
|
client = self._create_client()
|
|
|
|
def _download() -> Path:
|
|
destination_path.parent.mkdir(parents=True, exist_ok=True)
|
|
client.download_file(self.config["bucket"], storage_key, str(destination_path))
|
|
return destination_path
|
|
|
|
return await anyio.to_thread.run_sync(_download)
|
|
|
|
async def object_exists(self, storage_key: str) -> bool:
|
|
client = self._create_client()
|
|
|
|
def _exists() -> bool:
|
|
try:
|
|
client.head_object(Bucket=self.config["bucket"], Key=storage_key)
|
|
return True
|
|
except ClientError as exc:
|
|
error_code = exc.response.get("Error", {}).get("Code")
|
|
if error_code in {"404", "NoSuchKey", "NotFound"}:
|
|
return False
|
|
raise
|
|
|
|
return await anyio.to_thread.run_sync(_exists)
|
|
|
|
def _create_client(self):
|
|
client_kwargs = {
|
|
"service_name": "s3",
|
|
"region_name": self.config.get("region"),
|
|
"endpoint_url": self.config.get("endpoint_url"),
|
|
"aws_access_key_id": self.config.get("access_key_id"),
|
|
"aws_secret_access_key": self.config.get("secret_access_key"),
|
|
}
|
|
cleaned_kwargs = {key: value for key, value in client_kwargs.items() if value is not None}
|
|
return boto3.client(**cleaned_kwargs)
|
|
|