Daniil
108 lines
3.5 KiB
Python
108 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from typing import BinaryIO
|
|
|
|
import boto3 # type: ignore[import-untyped]
|
|
import boto3.session # type: ignore[import-untyped]
|
|
from botocore.config import Config # type: ignore[import-untyped]
|
|
from botocore.exceptions import ClientError # type: ignore[import-untyped]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class S3Config:
|
|
access_key: str
|
|
secret_key: str
|
|
bucket_name: str
|
|
endpoint_url_internal: str | None
|
|
endpoint_url_public: str | None
|
|
presign_expires_seconds: int = 3600
|
|
|
|
|
|
class S3StorageBackend:
|
|
def __init__(self, cfg: S3Config) -> None:
|
|
self._cfg = cfg
|
|
self._bucket_ready = False
|
|
|
|
session = boto3.session.Session()
|
|
common = {
|
|
"aws_access_key_id": cfg.access_key,
|
|
"aws_secret_access_key": cfg.secret_key,
|
|
"region_name": "us-east-1",
|
|
"config": Config(signature_version="s3v4", s3={"addressing_style": "path"}),
|
|
}
|
|
|
|
self._client = session.client(
|
|
"s3", endpoint_url=cfg.endpoint_url_internal, **common
|
|
)
|
|
|
|
presign_endpoint = cfg.endpoint_url_public or cfg.endpoint_url_internal
|
|
self._presign_client = session.client(
|
|
"s3", endpoint_url=presign_endpoint, **common
|
|
)
|
|
|
|
def ensure_bucket(self) -> None:
|
|
if self._bucket_ready:
|
|
return
|
|
|
|
try:
|
|
self._client.head_bucket(Bucket=self._cfg.bucket_name)
|
|
except ClientError as e:
|
|
code = str(e.response.get("Error", {}).get("Code", ""))
|
|
if code in {"404", "NoSuchBucket"}:
|
|
self._client.create_bucket(Bucket=self._cfg.bucket_name)
|
|
else:
|
|
raise
|
|
|
|
self._bucket_ready = True
|
|
|
|
def upload_fileobj(
|
|
self, key: str, fileobj: BinaryIO, *, content_type: str | None
|
|
) -> None:
|
|
self.ensure_bucket()
|
|
extra_args = {"ContentType": content_type} if content_type else None
|
|
self._client.upload_fileobj(
|
|
Fileobj=fileobj,
|
|
Bucket=self._cfg.bucket_name,
|
|
Key=key,
|
|
ExtraArgs=extra_args,
|
|
)
|
|
|
|
def download_fileobj(self, key: str, fileobj: BinaryIO) -> None:
|
|
self.ensure_bucket()
|
|
self._client.download_fileobj(self._cfg.bucket_name, key, fileobj)
|
|
|
|
def exists(self, key: str) -> bool:
|
|
self.ensure_bucket()
|
|
try:
|
|
self._client.head_object(Bucket=self._cfg.bucket_name, Key=key)
|
|
return True
|
|
except ClientError as e:
|
|
code = str(e.response.get("Error", {}).get("Code", ""))
|
|
if code in {"404", "NoSuchKey"}:
|
|
return False
|
|
raise
|
|
|
|
def size(self, key: str) -> int:
|
|
self.ensure_bucket()
|
|
resp = self._client.head_object(Bucket=self._cfg.bucket_name, Key=key)
|
|
return int(resp.get("ContentLength", 0))
|
|
|
|
def delete(self, key: str) -> None:
|
|
self.ensure_bucket()
|
|
self._client.delete_object(Bucket=self._cfg.bucket_name, Key=key)
|
|
|
|
def read(self, key: str) -> bytes:
|
|
self.ensure_bucket()
|
|
resp = self._client.get_object(Bucket=self._cfg.bucket_name, Key=key)
|
|
body = resp["Body"].read()
|
|
return body
|
|
|
|
def generate_url(self, key: str) -> str:
|
|
self.ensure_bucket()
|
|
return self._presign_client.generate_presigned_url(
|
|
ClientMethod="get_object",
|
|
Params={"Bucket": self._cfg.bucket_name, "Key": key},
|
|
ExpiresIn=self._cfg.presign_expires_seconds,
|
|
)
|