ai-flow-python/app/service/task/impl/file_upload_task.py

import logging
import os
import time
from datetime import datetime
from typing import Any

from app.service.task.base_task import BaseTask
from app.service.task.task_param import TaskParam
from app.service.task.task_result import TaskResult
from app.util import aliyun_util, google_storage_util
from app.util.download_util import download_file
from app.util.file_util import generate_unique_filename

logger = logging.getLogger(__name__)


class FileUploadTask(BaseTask):
    def get_type(self) -> str:
        return "file_upload"

    def execute(self, param: TaskParam) -> TaskResult:
        files_obj = param.inputs.get("files")
        files: list[str] = files_obj if isinstance(files_obj, list) else []

        try:
            provider = str(param.inputs.get("provider", ""))
            prefix = str(param.inputs.get("prefix", ""))
            date_folder = param.inputs.get("date_folder")

            logger.info("[file_upload] Starting: %d files, provider=%s, prefix=%s, date_folder=%s",
                         len(files), provider, prefix, date_folder)

            if not files:
                return TaskResult.failure("No files to upload")
            if provider not in ("aliyun", "google"):
                return TaskResult.failure(f"Unsupported target provider: {provider}")
            if not prefix or not prefix.strip():
                return TaskResult.failure("Key prefix cannot be empty")
            if not prefix.endswith("/"):
                prefix += "/"

            final_prefix = self._build_final_prefix(prefix, str(date_folder) if date_folder else None)
            logger.info("[file_upload] Final prefix: %s", final_prefix)

            t0 = time.time()
            local_files = self._process_files(files, param.workspace, param.step_id)
            logger.info("[file_upload] Resolved %d local files in %.2fs", len(local_files), time.time() - t0)

            uploaded_urls: list[str] = []
            upload_t0 = time.time()
            if provider == "aliyun":
                uploaded_urls = self._upload_to_aliyun(local_files, final_prefix)
            elif provider == "google":
                uploaded_urls = self._upload_to_google(local_files, final_prefix)

            logger.info("[file_upload] Completed: %d/%d files uploaded to %s in %.2fs",
                         len(uploaded_urls), len(local_files), provider, time.time() - upload_t0)
            return TaskResult.success({"files": uploaded_urls})
        except Exception as e:
            logger.error("[file_upload] Execution failed: %s", e, exc_info=True)
            return TaskResult.failure(f"File copy failed: {e}")

    @staticmethod
    def _build_final_prefix(prefix: str, date_folder: str | None) -> str:
        if not date_folder:
            return prefix
        try:
            date_path = datetime.now().strftime(date_folder)
            if not date_path.endswith("/"):
                date_path += "/"
            return prefix + date_path
        except Exception:
            return prefix

    @staticmethod
    def _process_files(files: list[str], workspace: str, step_id: str) -> list[str]:
        local_files: list[str] = []
        for f in files:
            try:
                if _is_url(f):
                    logger.info("[file_upload] Downloading URL: %s", f)
                    local_path = download_file(f, workspace + "/" + step_id)
                    if local_path:
                        local_files.append(local_path)
                else:
                    local_files.append(workspace + "/" + f)
            except Exception as e:
                logger.error("[file_upload] Error processing file: %s, %s", f, e)
        return local_files

    @staticmethod
    def _upload_to_aliyun(files: list[str], prefix: str) -> list[str]:
        uploaded: list[str] = []
        for path in files:
            if not os.path.isfile(path):
                logger.warning("[file_upload] File not found, skipping: %s", path)
                continue
            try:
                unique_name = generate_unique_filename(os.path.basename(path))
                key = prefix + unique_name
                logger.info("[file_upload] Uploading to Aliyun OSS: %s -> %s", os.path.basename(path), key)
                url = aliyun_util.upload_file(path, key)
                if url:
                    uploaded.append(url)
            except Exception as e:
                logger.error("[file_upload] Exception uploading file: %s, %s", path, e)
        return uploaded

    @staticmethod
    def _upload_to_google(files: list[str], prefix: str) -> list[str]:
        uploaded: list[str] = []
        for path in files:
            if not os.path.isfile(path):
                logger.warning("[file_upload] File not found, skipping: %s", path)
                continue
            try:
                unique_name = generate_unique_filename(os.path.basename(path))
                key = prefix + unique_name
                logger.info("[file_upload] Uploading to Google GCS: %s -> %s", os.path.basename(path), key)
                url = google_storage_util.upload_file(path, key)
                if url:
                    uploaded.append(url)
            except Exception as e:
                logger.error("[file_upload] Exception uploading file: %s, %s", path, e)
        return uploaded


def _is_url(s: str) -> bool:
    if not s:
        return False
    lower = s.strip().lower()
    return any(lower.startswith(p) for p in ("http://", "https://", "ftp://", "s3://", "gs://"))