ai-flow-python/app/util/google_storage_util.py

import io
import json
import logging
import os

from app.config import settings
from app.util.file_util import generate_unique_filename

logger = logging.getLogger(__name__)


def upload_file(
    local_path: str,
    gcs_key: str,
    project_id: str | None = None,
    bucket_name: str | None = None,
    credentials_json: str | None = None,
) -> str | None:
    if not local_path or not gcs_key:
        logger.error("Local file path or GCS key is blank.")
        return None

    if not os.path.isfile(local_path):
        logger.error("Local file does not exist: %s", local_path)
        return None

    pid = project_id or settings.gcs_project_id
    bn = bucket_name or settings.gcs_bucket_name
    cred = credentials_json or settings.gcs_credentials_json
    if not cred and settings.gcs_credentials_file:
        try:
            with open(settings.gcs_credentials_file, "r", encoding="utf-8") as f:
                cred = f.read()
        except Exception as e:
            logger.error("Failed to read GCS credentials file: %s", e)

    try:
        from google.cloud import storage as gcs_storage
        from google.oauth2 import service_account

        if cred:
            info = json.loads(cred)
            credentials = service_account.Credentials.from_service_account_info(info)
            client = gcs_storage.Client(project=pid, credentials=credentials)
        else:
            client = gcs_storage.Client(project=pid)

        bucket = client.bucket(bn)
        blob = bucket.blob(gcs_key)
        blob.upload_from_filename(local_path)

        gcs_url = f"https://storage.googleapis.com/{bn}/{gcs_key}"
        logger.debug("Successfully uploaded: %s", gcs_url)
        return gcs_url
    except Exception as e:
        logger.error("Failed to upload file to GCS: %s -> %s, %s", local_path, gcs_key, e)
        return None


def upload_files(
    local_paths: list[str],
    gcs_prefix: str,
    project_id: str | None = None,
    bucket_name: str | None = None,
    credentials_json: str | None = None,
) -> list[str]:
    uploaded: list[str] = []
    if not local_paths:
        return uploaded

    if not gcs_prefix:
        gcs_prefix = "uploads/"
    if not gcs_prefix.endswith("/"):
        gcs_prefix += "/"

    for path in local_paths:
        try:
            name = generate_unique_filename(os.path.basename(path))
            key = gcs_prefix + name
            url = upload_file(path, key, project_id, bucket_name, credentials_json)
            if url:
                uploaded.append(url)
        except Exception as e:
            logger.error("Failed to upload file: %s, %s", path, e)

    return uploaded