53 lines
1.5 KiB
Python

from dataclasses import dataclass
from typing import BinaryIO
import boto3
from botocore.config import Config
from django.conf import settings
@dataclass(frozen=True)
class StoredObject:
bucket: str
object_key: str
content_type: str
size_bytes: int
class TosStorage:
def __init__(self) -> None:
tos = settings.TOS
self.bucket = tos["bucket"]
self.client = boto3.client(
"s3",
endpoint_url=tos["endpoint"],
aws_access_key_id=tos["access_key_id"],
aws_secret_access_key=tos["secret_access_key"],
region_name="cn-shanghai",
config=Config(s3={"addressing_style": "virtual"}),
)
def upload_fileobj(self, *, fileobj: BinaryIO, object_key: str, content_type: str) -> StoredObject:
fileobj.seek(0, 2)
size = fileobj.tell()
fileobj.seek(0)
self.client.upload_fileobj(
fileobj,
self.bucket,
object_key,
ExtraArgs={"ContentType": content_type},
)
return StoredObject(
bucket=self.bucket,
object_key=object_key,
content_type=content_type,
size_bytes=size,
)
def presigned_get_url(self, *, object_key: str, expires_in: int = 3600) -> str:
return self.client.generate_presigned_url(
"get_object",
Params={"Bucket": self.bucket, "Key": object_key},
ExpiresIn=expires_in,
)