HaimKortovich
2018f5a5f6
All checks were successful
Build and Publish / build-release (push) Successful in 1m8s
148 lines
4.3 KiB
Python
148 lines
4.3 KiB
Python
import boto3
|
|
import json
|
|
import os
|
|
import tempfile
|
|
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 read_bucket_info() -> dict:
|
|
if not settings.cosi_bucket_info_path:
|
|
raise ValueError("COSI_BUCKET_INFO_PATH not set")
|
|
with open(settings.cosi_bucket_info_path, "r") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def get_cosi_s3_config() -> dict:
|
|
bucket_info = read_bucket_info()
|
|
s3_conf = bucket_info["spec"]["secretS3"]
|
|
return {
|
|
"endpoint": s3_conf["endpoint"],
|
|
"access_key": s3_conf["accessKeyID"],
|
|
"secret_key": s3_conf["accessSecretKey"],
|
|
"bucket": bucket_info["spec"]["bucketName"]
|
|
}
|
|
|
|
|
|
def get_client():
|
|
if settings.cosi_bucket_info_path:
|
|
cosi_config = get_cosi_s3_config()
|
|
return boto3.client(
|
|
"s3",
|
|
endpoint_url=cosi_config["endpoint"],
|
|
aws_access_key_id=cosi_config["access_key"],
|
|
aws_secret_access_key=cosi_config["secret_key"],
|
|
config=Config(signature_version="s3v4"),
|
|
region_name="us-east-1"
|
|
)
|
|
return boto3.client(
|
|
"s3",
|
|
endpoint_url=settings.s3_endpoint,
|
|
aws_access_key_id=settings.s3_access_key_id,
|
|
aws_secret_access_key=settings.s3_secret_key,
|
|
config=Config(signature_version="s3v4"),
|
|
region_name=settings.s3_region
|
|
)
|
|
|
|
|
|
def get_bucket_name() -> str:
|
|
if settings.cosi_bucket_info_path:
|
|
return get_cosi_s3_config()["bucket"]
|
|
return settings.s3_bucket
|
|
|
|
|
|
def ensure_bucket_exists() -> None:
|
|
bucket_name = get_bucket_name()
|
|
client = get_client()
|
|
try:
|
|
client.head_bucket(Bucket=bucket_name)
|
|
logger.info(f"Bucket '{bucket_name}' already exists")
|
|
except client.exceptions.ClientError as e:
|
|
error_code = e.response['Error']['Code']
|
|
if error_code == '404':
|
|
try:
|
|
client.create_bucket(
|
|
Bucket=bucket_name,
|
|
CreateBucketConfiguration={
|
|
'LocationConstraint': 'us-east-1'
|
|
}
|
|
)
|
|
logger.info(f"Created bucket '{bucket_name}'")
|
|
except Exception as create_error:
|
|
logger.error(f"Failed to create bucket '{bucket_name}': {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:
|
|
bucket_name = get_bucket_name()
|
|
client = get_client()
|
|
|
|
file.file.seek(0, os.SEEK_END)
|
|
file_size = file.file.tell()
|
|
file.file.seek(0)
|
|
file_content = file.file.read()
|
|
file.file.seek(0)
|
|
|
|
extra_args = {"ContentType": content_type}
|
|
if metadata:
|
|
extra_args["Metadata"] = metadata
|
|
|
|
client.put_object(
|
|
Bucket=bucket_name,
|
|
Key=s3_key,
|
|
Body=file_content,
|
|
ContentLength=file_size,
|
|
ContentType=content_type,
|
|
Metadata=metadata
|
|
)
|
|
return s3_key
|
|
|
|
|
|
def delete_file(s3_key: str) -> None:
|
|
bucket_name = get_bucket_name()
|
|
client = get_client()
|
|
client.delete_object(Bucket=bucket_name, Key=s3_key)
|
|
|
|
|
|
def file_exists(s3_key: str) -> bool:
|
|
bucket_name = get_bucket_name()
|
|
client = get_client()
|
|
try:
|
|
client.head_object(Bucket=bucket_name, Key=s3_key)
|
|
return True
|
|
except client.exceptions.ClientError:
|
|
return False
|
|
|
|
|
|
def get_file_metadata(s3_key: str) -> dict:
|
|
bucket_name = get_bucket_name()
|
|
client = get_client()
|
|
response = client.head_object(Bucket=bucket_name, Key=s3_key)
|
|
return response.get("Metadata", {})
|
|
|
|
|
|
def download_to_temp(s3_key: str) -> str:
|
|
bucket_name = get_bucket_name()
|
|
client = get_client()
|
|
suffix = os.path.splitext(s3_key)[-1] or ".tmp"
|
|
tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
|
|
client.download_fileobj(bucket_name, s3_key, tmp)
|
|
tmp.close()
|
|
return tmp.name
|
|
|
|
|
|
def presigned_download_url(s3_key: str, expires_in: int = 3600) -> str:
|
|
bucket_name = get_bucket_name()
|
|
client = get_client()
|
|
return client.generate_presigned_url(
|
|
"get_object",
|
|
Params={"Bucket": bucket_name, "Key": s3_key},
|
|
ExpiresIn=expires_in
|
|
) |