update-server/minio_tool.py

from minio import Minio
from datetime import timedelta
import os
from dotenv import load_dotenv
import io
import socket
import shutil
from pathlib import Path
import urllib3
from urllib.parse import urlparse
from urllib3.util import Retry, Timeout

load_dotenv()
BASE_DIR = Path(__file__).resolve().parent


def env_bool(name: str, default: bool = False) -> bool:
    value = os.getenv(name)
    return default if value is None else value.strip().lower() in {"1", "true", "yes", "on"}

MINIO_ENDPOINT = os.getenv("MINIO_ENDPOINT")
MINIO_ACCESS_KEY = os.getenv("MINIO_ACCESS_KEY")
MINIO_SECRET_KEY = os.getenv("MINIO_SECRET_KEY")
MINIO_BUCKET = os.getenv("MINIO_BUCKET")
SIGN_EXPIRE_MIN = int(os.getenv("SIGN_EXPIRE_MIN") or 60)
LOCAL_UPLOAD_ROOT = Path(os.getenv("LOCAL_UPLOAD_ROOT", "local_uploads")).expanduser()
if not LOCAL_UPLOAD_ROOT.is_absolute():
    LOCAL_UPLOAD_ROOT = BASE_DIR / LOCAL_UPLOAD_ROOT
LOCAL_FILE_URL_BASE = os.getenv("LOCAL_FILE_URL_BASE")
MINIO_SECURE = env_bool("MINIO_SECURE")
MINIO_CONNECT_TIMEOUT_SEC = float(os.getenv("MINIO_CONNECT_TIMEOUT_SEC", "2"))
MINIO_READ_TIMEOUT_SEC = float(os.getenv("MINIO_READ_TIMEOUT_SEC", "5"))
MINIO_RETRY_TOTAL = int(os.getenv("MINIO_RETRY_TOTAL", "1"))
LOCAL_UPLOAD_ROOT.mkdir(parents=True, exist_ok=True)

http_client = urllib3.PoolManager(
    timeout=Timeout(connect=MINIO_CONNECT_TIMEOUT_SEC, read=MINIO_READ_TIMEOUT_SEC),
    retries=Retry(total=MINIO_RETRY_TOTAL, connect=MINIO_RETRY_TOTAL,
                  read=MINIO_RETRY_TOTAL, status=MINIO_RETRY_TOTAL, backoff_factor=0.2),
)

mc = Minio(
    MINIO_ENDPOINT,
    access_key=MINIO_ACCESS_KEY,
    secret_key=MINIO_SECRET_KEY,
    secure=MINIO_SECURE,
    http_client=http_client,
)

public_endpoint_raw = os.getenv("MINIO_PUBLIC_ENDPOINT", "").strip()
if public_endpoint_raw:
    parsed_public_endpoint = urlparse(
        public_endpoint_raw if "://" in public_endpoint_raw else f"http://{public_endpoint_raw}"
    )
    public_mc = Minio(
        parsed_public_endpoint.netloc,
        access_key=MINIO_ACCESS_KEY,
        secret_key=MINIO_SECRET_KEY,
        secure=parsed_public_endpoint.scheme == "https",
    )
else:
    public_mc = mc


def get_url(file_path: str) -> str:
    bucket = MINIO_BUCKET
    full_object_path = file_path
    try:
        mc.stat_object(bucket, full_object_path)
        url = public_mc.presigned_get_object(
            bucket, full_object_path, expires=timedelta(minutes=SIGN_EXPIRE_MIN)
        )
        return url
    except Exception as err:
        print(f"generate url error: {err}")
        local_path = LOCAL_UPLOAD_ROOT / file_path
        if local_path.exists():
            if LOCAL_FILE_URL_BASE:
                return LOCAL_FILE_URL_BASE.rstrip('/') + '/' + file_path
            host = os.getenv('SERVER_HOST') or '127.0.0.1'
            if host == '0.0.0.0':
                try:
                    host = socket.gethostbyname(socket.gethostname())
                except Exception:
                    host = '127.0.0.1'
            port = os.getenv('SERVER_PORT') or '8000'
            return f"http://{host}:{port}/static/{file_path}"
        scheme = "https" if MINIO_SECURE else "http"
        return f"{scheme}://{MINIO_ENDPOINT}/{bucket}/{full_object_path}"


def put_file(object_path: str, data: bytes, size: int):
    bucket = MINIO_BUCKET
    try:
        bio = io.BytesIO(data)
        mc.put_object(bucket, object_path, bio, length=size)
        return {'storage': 'minio', 'path': object_path}
    except Exception as e:
        print(f"minio put_file error: {e}")
        local_path = LOCAL_UPLOAD_ROOT / object_path
        local_path.parent.mkdir(parents=True, exist_ok=True)
        with open(local_path, 'wb') as f:
            f.write(data)
        print(f"local fallback stored {local_path}")
        return {'storage': 'local', 'path': str(local_path)}



def put_file_stream(object_path: str, source, size: int):
    bucket = MINIO_BUCKET
    try:
        source.seek(0)
        mc.put_object(bucket, object_path, source, length=size)
        return {'storage': 'minio', 'path': object_path}
    except Exception as e:
        print(f"minio put_file_stream error: {e}")
        local_path = LOCAL_UPLOAD_ROOT / object_path
        local_path.parent.mkdir(parents=True, exist_ok=True)
        source.seek(0)
        with open(local_path, 'wb') as target:
            shutil.copyfileobj(source, target, length=1024 * 1024)
        print(f"local fallback stored {local_path}")
        return {'storage': 'local', 'path': str(local_path)}


def remove_prefix(prefix: str):
    bucket = MINIO_BUCKET
    try:
        objects = mc.list_objects(bucket, prefix=prefix, recursive=True)
        for obj in objects:
            try:
                mc.remove_object(bucket, obj.object_name)
            except Exception as e:
                print(f"failed remove object {obj.object_name}: {e}")
    except Exception as e:
        print(f"minio remove_prefix error: {e}")
    local_prefix = LOCAL_UPLOAD_ROOT / prefix
    if local_prefix.exists():
        try:
            shutil.rmtree(local_prefix)
        except Exception as e:
            print(f"local remove_prefix error: {e}")

def iter_file(object_path: str, chunk_size: int = 1024 * 1024):
    """Stream an object without creating a second full local copy."""
    response = None
    emitted = 0
    try:
        response = mc.get_object(MINIO_BUCKET, object_path)
        while True:
            chunk = response.read(chunk_size)
            if not chunk: break
            emitted += len(chunk)
            yield chunk
        return
    except Exception as err:
        if emitted:
            raise RuntimeError(f"MinIO stream interrupted after {emitted} bytes: {object_path}") from err
        print(f"MinIO stream fallback for {object_path}: {err}")
    finally:
        if response is not None:
            response.close(); response.release_conn()
    local_path = LOCAL_UPLOAD_ROOT / object_path
    with open(local_path, "rb") as source:
        while True:
            chunk = source.read(chunk_size)
            if not chunk: break
            yield chunk