102 lines
3.2 KiB
Python
102 lines
3.2 KiB
Python
import boto3
|
|
import tempfile
|
|
import os
|
|
from botocore.client import Config
|
|
from fastapi import UploadFile
|
|
from app.config import settings
|
|
from app.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
def get_client():
|
|
return boto3.client(
|
|
"s3",
|
|
endpoint_url=settings.s3_endpoint,
|
|
aws_access_key_id=settings.s3_access_key,
|
|
aws_secret_access_key=settings.s3_secret_key,
|
|
config=Config(signature_version="s3v4"),
|
|
region_name=settings.s3_region
|
|
)
|
|
|
|
def ensure_bucket_exists() -> None:
|
|
"""Ensure the S3 bucket exists, create it if it doesn't exist.
|
|
|
|
Raises:
|
|
Exception: If bucket creation fails (service will fail to start)
|
|
"""
|
|
client = get_client()
|
|
try:
|
|
client.head_bucket(Bucket=settings.s3_bucket)
|
|
logger.info(f"Bucket '{settings.s3_bucket}' already exists")
|
|
except client.exceptions.ClientError as e:
|
|
error_code = e.response['Error']['Code']
|
|
if error_code == '404':
|
|
try:
|
|
client.create_bucket(
|
|
Bucket=settings.s3_bucket,
|
|
CreateBucketConfiguration={
|
|
'LocationConstraint': settings.s3_region
|
|
}
|
|
)
|
|
logger.info(f"Created bucket '{settings.s3_bucket}'")
|
|
except Exception as create_error:
|
|
logger.error(f"Failed to create bucket '{settings.s3_bucket}': {create_error}")
|
|
raise
|
|
else:
|
|
logger.error(f"Error checking bucket: {e}")
|
|
raise
|
|
|
|
def upload_file(file: UploadFile, s3_key: str, content_type: str, metadata: dict = None) -> str:
|
|
"""Upload file to S3 with metadata"""
|
|
client = get_client()
|
|
|
|
extra_args = {"ContentType": content_type}
|
|
if metadata:
|
|
extra_args["Metadata"] = metadata
|
|
|
|
client.upload_fileobj(
|
|
file.file,
|
|
settings.s3_bucket,
|
|
s3_key,
|
|
ExtraArgs=extra_args
|
|
)
|
|
return s3_key
|
|
|
|
def delete_file(s3_key: str) -> None:
|
|
"""Delete file from S3"""
|
|
client = get_client()
|
|
client.delete_object(Bucket=settings.s3_bucket, Key=s3_key)
|
|
|
|
def file_exists(s3_key: str) -> bool:
|
|
"""Check if file exists in S3"""
|
|
client = get_client()
|
|
try:
|
|
client.head_object(Bucket=settings.s3_bucket, Key=s3_key)
|
|
return True
|
|
except client.exceptions.ClientError:
|
|
return False
|
|
|
|
def get_file_metadata(s3_key: str) -> dict:
|
|
"""Get file metadata from S3"""
|
|
client = get_client()
|
|
response = client.head_object(Bucket=settings.s3_bucket, Key=s3_key)
|
|
return response.get("Metadata", {})
|
|
|
|
def download_to_temp(s3_key: str) -> str:
|
|
"""Download file from S3 to temp file"""
|
|
client = get_client()
|
|
suffix = os.path.splitext(s3_key)[-1] or ".tmp"
|
|
tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
|
|
client.download_fileobj(settings.s3_bucket, s3_key, tmp)
|
|
tmp.close()
|
|
return tmp.name
|
|
|
|
def presigned_download_url(s3_key: str, expires_in: int = 3600) -> str:
|
|
"""Generate presigned download URL"""
|
|
client = get_client()
|
|
return client.generate_presigned_url(
|
|
"get_object",
|
|
Params={"Bucket": settings.s3_bucket, "Key": s3_key},
|
|
ExpiresIn=expires_in
|
|
)
|