BrandCultivation/utils/file_stream.py

import time
from core import get_logger, settings
from io import BytesIO
import os
import pandas as pd
import requests

logger = get_logger("utils.file_stream")


class FileStreamUtils:
    upload_url = settings.file_upload_url
    download_url = settings.file_download_url
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
        "Accept": "*/*",
    }
    if settings.file_service_cookie:
        headers["Cookie"] = settings.file_service_cookie

    @staticmethod
    def upload_files(reports_dir, files):
        files_id = {}
        for filename in files:
            file_path = os.path.join(reports_dir, f"{filename}.xlsx")
            start_time = time.time()
            try:
                with open(file_path, "rb") as f:
                    upload_files = {"file": (os.path.basename(file_path), f)}
                    response = requests.post(
                        FileStreamUtils.upload_url,
                        headers=FileStreamUtils.headers,
                        files=upload_files,
                        verify=True,
                    )
                    duration_ms = (time.time() - start_time) * 1000
                    if response.json().get("success"):
                        file_id = response.json()["data"]["file_info"]["fileid"]
                        files_id[filename] = file_id
                        logger.info(f"File uploaded: {filename} -> {file_id} ({duration_ms:.0f}ms)")
                    else:
                        logger.error(f"Upload failed for {filename}: {response.text}")
                        return None
            except requests.exceptions.RequestException as e:
                logger.error(f"Upload request error for {filename}: {e}", exc_info=True)
                return None
            except Exception as e:
                logger.error(f"Upload error for {filename}: {e}", exc_info=True)
                return None
        return files_id

    @staticmethod
    def download_file(file_id, file_type="xlsx"):
        """通过file_id从阿里云文件数据库下载文件"""
        start_time = time.time()
        try:
            response = requests.get(
                f"{FileStreamUtils.download_url}/{file_id}",
                headers=FileStreamUtils.headers,
                verify=True,
            )
            duration_ms = (time.time() - start_time) * 1000

            if response.status_code == 200:
                file_content = BytesIO(response.content)
                if file_type == "xlsx":
                    data = pd.read_excel(file_content, engine="openpyxl")
                elif file_type == "csv":
                    data = pd.read_csv(file_content)
                else:
                    raise ValueError(f"不支持的文件类型：{file_type}")
                logger.info(f"File downloaded: {file_id} ({duration_ms:.0f}ms, {len(response.content)} bytes)")
                return data
            else:
                logger.error(f"Download failed: file_id={file_id}, status={response.status_code}")
                return None
        except requests.exceptions.RequestException as e:
            logger.error(f"Download request error: file_id={file_id}, error={e}", exc_info=True)
            return None
        except Exception as e:
            logger.error(f"Download error: file_id={file_id}, error={e}", exc_info=True)
            return None