Spaces:
Paused
Paused
File size: 3,144 Bytes
ec17af3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | import os
import shutil
from pathlib import Path
from typing import BinaryIO
import boto3
from app.config import settings
class StorageClient:
def __init__(self) -> None:
self.use_s3 = settings.storage_backend == "s3" or all(
[
settings.s3_endpoint,
settings.s3_access_key,
settings.s3_secret_key,
settings.s3_bucket,
]
)
Path(settings.storage_dir).mkdir(parents=True, exist_ok=True)
if self.use_s3:
missing = [
name
for name, value in {
"S3_ENDPOINT": settings.s3_endpoint,
"S3_ACCESS_KEY": settings.s3_access_key,
"S3_SECRET_KEY": settings.s3_secret_key,
"S3_BUCKET": settings.s3_bucket,
}.items()
if not value
]
if missing:
raise RuntimeError(f"S3 storage is enabled but missing: {', '.join(missing)}")
self.s3 = boto3.client(
"s3",
endpoint_url=settings.s3_endpoint,
aws_access_key_id=settings.s3_access_key,
aws_secret_access_key=settings.s3_secret_key,
region_name=settings.s3_region or "auto",
)
self.bucket = settings.s3_bucket
else:
self.s3 = None
self.bucket = None
def save_file(self, key: str, file_obj: BinaryIO) -> str:
if self.use_s3 and self.s3 and self.bucket:
self.s3.upload_fileobj(file_obj, self.bucket, key)
return f"s3://{self.bucket}/{key}"
path = Path(settings.storage_dir) / key
path.parent.mkdir(parents=True, exist_ok=True)
with open(path, "wb") as out:
shutil.copyfileobj(file_obj, out)
return str(path)
def delete_file(self, key_or_path: str) -> None:
if key_or_path.startswith("s3://") and self.use_s3 and self.s3 and self.bucket:
key = key_or_path.replace(f"s3://{self.bucket}/", "", 1)
self.s3.delete_object(Bucket=self.bucket, Key=key)
cache_path = Path(settings.storage_dir) / "_s3_cache" / key
if cache_path.exists():
cache_path.unlink()
return
if os.path.exists(key_or_path):
os.remove(key_or_path)
def get_local_path(self, key_or_path: str) -> str:
if key_or_path.startswith("s3://"):
if not (self.use_s3 and self.s3 and self.bucket):
raise RuntimeError("S3 file requires configured S3 storage")
prefix = f"s3://{self.bucket}/"
key = key_or_path[len(prefix):] if key_or_path.startswith(prefix) else key_or_path.split("/", 3)[-1]
cache_path = Path(settings.storage_dir) / "_s3_cache" / key
if not cache_path.exists():
cache_path.parent.mkdir(parents=True, exist_ok=True)
self.s3.download_file(self.bucket, key, str(cache_path))
return str(cache_path)
return key_or_path
storage_client = StorageClient()
|