algorithm_platform/data_processing/data_operation/test.py

#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :test.py
# @Time      :2025/3/13 14:19
# @Author    :David
# @Company: shenyang JY
from enum import Enum
import paramiko
from datetime import datetime, timedelta
from typing import Optional
import os
import zipfile
import shutil
import tempfile

# 配置信息
SFTP_HOST = '192.168.1.33'
SFTP_PORT = 2022
SFTP_USERNAME = 'liudawei'
SFTP_PASSWORD = 'liudawei@123'
# 在原配置部分添加以下配置
DEST_SFTP_HOST = 'dest_sftp.example.com'
DEST_SFTP_PORT = 22
DEST_SFTP_USERNAME = 'dest_username'
DEST_SFTP_PASSWORD = 'dest_password'
DEFAULT_TARGET_DIR = 'cdq'  # 默认上传目录

# 更新后的三级映射
MAPPINGS = {
    'koi': {('Zone', '1.0'): {'J00645'}},
    'lucky': {}, 'seer': {('lgb', '1.0'): {'J00578'}}
}

def get_next_target_time(current_time=None):
    """获取下一个目标时刻"""
    if current_time is None:
        current_time = datetime.now()

    target_hours = [0, 6, 12, 18]
    current_hour = current_time.hour

    for hour in sorted(target_hours):
        if current_hour < hour:
            return current_time.replace(hour=hour, minute=0, second=0, microsecond=0)

    # 如果当前时间超过所有目标小时，使用次日0点
    return (current_time + timedelta(days=1)).replace(hour=0, minute=0, second=0, microsecond=0)


def download_files_via_sftp(mappings, datetime_str, local_temp_dir, model_type):
    """
    封装SFTP连接和文件下载的完整流程
    :param mappings: 文件映射配置
    :param datetime_str: 日期时间字符串，用于文件名
    :param local_temp_dir: 本地临时目录路径
    """
    transport = None
    sftp = None
    try:
        # 创建SSH传输通道
        transport = paramiko.Transport((SFTP_HOST, SFTP_PORT))
        transport.connect(username=SFTP_USERNAME, password=SFTP_PASSWORD)

        # 创建SFTP客户端
        sftp = paramiko.SFTPClient.from_transport(transport)

        # 执行文件下载
        for engineer in mappings:
            datetime_str = datetime_str if engineer == 'koi' else 2025012000
            remote_base = f"/{engineer}/"
            try:
                sftp.chdir(remote_base)
            except FileNotFoundError:
                print(f"工程师目录不存在: {remote_base}")
                continue

            for model_version in mappings[engineer]:
                target_file = f"jy_{engineer}.{'.'.join(model_version)}_{datetime_str}.zip"
                remote_path = os.path.join(remote_base, target_file).replace("\\", "/")
                local_path = os.path.join(local_temp_dir, target_file).replace("\\", "/")

                try:
                    sftp.get(remote_path, local_path)
                    print(f"下载成功: {remote_path} -> {local_path}")
                except Exception as e:
                    print(f"文件下载失败 {remote_path}: {str(e)}")

    except paramiko.AuthenticationException:
        print("认证失败，请检查用户名和密码")
    except paramiko.SSHException as e:
        print(f"SSH连接异常: {str(e)}")
    except Exception as e:
        print(f"未知错误: {str(e)}")
    finally:
        # 遍历到最后一个中短期，确保连接关闭
        if model_type == 'zcq':
            if sftp:
                sftp.close()
            if transport and transport.is_active():
                transport.close()


def upload_to_sftp(local_path: str, target_dir: str) -> bool:
    """上传文件到SFTP服务器

    Args:
        local_path: 本地文件路径
        target_dir: 远程目标目录

    Returns:
        上传是否成功 (True/False)
    """
    transport: Optional[paramiko.Transport] = None
    sftp: Optional[paramiko.SFTPClient] = None

    try:
        # 验证本地文件存在
        if not os.path.isfile(local_path):
            raise FileNotFoundError(f"本地文件不存在: {local_path}")

        # 创建SFTP连接
        transport = paramiko.Transport((DEST_SFTP_HOST, DEST_SFTP_PORT))
        transport.connect(username=DEST_SFTP_USERNAME, password=DEST_SFTP_PASSWORD)
        sftp = paramiko.SFTPClient.from_transport(transport)

        # 执行上传
        remote_filename = os.path.basename(local_path)
        remote_path = f"{target_dir}/{remote_filename}"
        sftp.put(local_path, remote_path)
        print(f"[SUCCESS] 上传完成: {remote_path}")
        return True

    except Exception as e:
        print(f"[ERROR] 上传失败: {str(e)}")
        return False
    finally:
        # 确保资源释放
        if sftp:
            sftp.close()
        if transport and transport.is_active():
            transport.close()

def process_zips(mappings, local_temp_dir, datetime_str, final_collect_dir):
    """处理所有下载的ZIP文件并收集场站目录"""
    for engineer in mappings:
        datetime_str = datetime_str if engineer == 'koi' else 2025012000
        for model_version in mappings[engineer]:
            target_file = f"jy_{engineer}.{'.'.join(model_version)}_{datetime_str}_dq.zip"
            zip_path = os.path.join(local_temp_dir, target_file).replace("\\", "/")
            station_codes = mappings[engineer][model_version]

            if not os.path.exists(zip_path):
                continue

            # 创建临时解压目录
            with tempfile.TemporaryDirectory() as temp_extract:
                # 解压ZIP文件
                try:
                    with zipfile.ZipFile(zip_path, 'r') as zf:
                        zf.extractall(temp_extract)
                except zipfile.BadZipFile:
                    print(f"无效的ZIP文件: {zip_path}")
                    continue

                # 收集场站目录
                for root, dirs, files in os.walk(temp_extract):
                    for dir_name in dirs:
                        if dir_name in station_codes:
                            src = os.path.join(root, dir_name)
                            dest = os.path.join(final_collect_dir, dir_name)

                            if not os.path.exists(dest):
                                shutil.copytree(src, dest)
                                print(f"已收集场站: {dir_name}")


def create_final_zip(final_collect_dir: str, datetime_str: str, model_type: str) -> str:
    """创建ZIP压缩包并返回完整路径

    Args:
        final_collect_dir: 需要打包的源目录
        datetime_str: 时间戳字符串
        model_type: 模型类型标识

    Returns:
        生成的ZIP文件完整路径
    """
    # 确保缓存目录存在
    os.makedirs('../cache/ftp', exist_ok=True)

    # 构造标准化文件名
    zip_filename = f"jy_algo_{datetime_str}_{model_type}.zip"
    output_path = os.path.join('../cache/ftp', zip_filename)

    try:
        with zipfile.ZipFile(output_path, 'w', zipfile.ZIP_DEFLATED) as zf:
            for root, _, files in os.walk(final_collect_dir):
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, final_collect_dir)
                    zf.write(file_path, arcname)
        print(f"[SUCCESS] ZIP创建成功: {output_path}")
        return output_path
    except Exception as e:
        print(f"[ERROR] 创建ZIP失败: {str(e)}")
        raise


def clean_up_file(file_path: str) -> None:
    """安全删除本地文件"""
    try:
        if os.path.exists(file_path):
            os.remove(file_path)
            print(f"[CLEANUP] 已删除本地文件: {file_path}")
    except Exception as e:
        print(f"[WARNING] 文件删除失败: {str(e)}")

def prod_data_handler(mappings):
    # 创建临时工作目录
    for model_type in ['cdq', 'dq', 'zcq']:
        with tempfile.TemporaryDirectory() as local_temp_dir:
            final_collect_dir = os.path.join(local_temp_dir, 'collected_stations')
            os.makedirs(final_collect_dir, exist_ok=True)

            # 计算目标时间
            target_time = get_next_target_time()
            datetime_str = target_time.strftime("%Y%m%d%H")
            print(f"目标时间: {datetime_str}")
            datetime_str = '2025012412'

            # 下载文件
            download_files_via_sftp(mappings, datetime_str, local_temp_dir, model_type)

            # 处理下载的文件
            process_zips(mappings, local_temp_dir, datetime_str, final_collect_dir)

            # 创建最终ZIP
            zip_path  = create_final_zip(final_collect_dir, datetime_str, model_type)

            # 上传打包ZIP文件
            if upload_to_sftp(zip_path, f"/{model_type}"):
                # 步骤3: 上传成功后清理
                clean_up_file(zip_path)
            else:
                print("[WARNING] 上传未成功，保留本地文件")

if __name__ == "__main__":
    prod_data_handler()