import os
import csv
import zipfile
import logging
import mysql.connector
from datetime import datetime
import time
import concurrent.futures
import glob

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('data_export.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

def export_table_to_csv(table_name, csv_path, db_config, batch_size=1000):
    """導(dǎo)出單個(gè)表的數(shù)據(jù)到CSV文件，分批處理"""
    conn = None
    cursor = None
    total_rows = 0
    try:
        conn = mysql.connector.connect(**db_config)
        cursor = conn.cursor()

        # 獲取數(shù)據(jù)并寫入CSV
        with open(csv_path, 'w', newline='', encoding='utf-8') as csvfile:
            writer = csv.writer(csvfile)
            
            # 執(zhí)行查詢并獲取列名
            cursor.execute(f"SELECT * FROM `{table_name}`")
            columns = [col[0] for col in cursor.description]
            writer.writerow(columns)
            
            # 分批獲取數(shù)據(jù)
            while True:
                rows = cursor.fetchmany(batch_size)
                if not rows:
                    break
                writer.writerows(rows)
                total_rows += len(rows)
                logger.debug(f"{table_name} 已導(dǎo)出 {total_rows} 行")

        logger.info(f"{table_name} CSV導(dǎo)出完成，總行數(shù)：{total_rows}")
        return total_rows

    except Exception as e:
        logger.error(f"導(dǎo)出表 {table_name} 失敗: {str(e)}", exc_info=True)
        raise
    finally:
        if cursor:
            cursor.close()
        if conn and conn.is_connected():
            conn.close()

def compress_to_zip(source_path, zip_path):
    """壓縮文件為ZIP格式"""
    try:
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            zipf.write(source_path, arcname=os.path.basename(source_path))
        logger.info(f"成功壓縮 {source_path} 到 {zip_path}")
    except Exception as e:
        logger.error(f"壓縮 {source_path} 失敗: {str(e)}", exc_info=True)
        raise

def process_table(table_name, db_config, csv_dir, zip_dir):
    """處理單個(gè)表的導(dǎo)出和壓縮"""
    start_time = time.time()
    logger.info(f"開始處理表: {table_name}")
    status = "成功"
    rows_exported = 0

    try:
        # 定義文件路徑
        csv_filename = f"{table_name}.csv"
        zip_filename = f"{table_name}.zip"
        csv_path = os.path.join(csv_dir, csv_filename)
        zip_path = os.path.join(zip_dir, zip_filename)

        # 導(dǎo)出CSV
        rows_exported = export_table_to_csv(table_name, csv_path, db_config)
        
        # 壓縮文件
        compress_to_zip(csv_path, zip_path)

    except Exception as e:
        status = f"失敗: {str(e)}"
        # 清理可能存在的中間文件
        for path in [csv_path, zip_path]:
            if path and os.path.exists(path):
                try:
                    os.remove(path)
                    logger.warning(f"已清理文件: {path}")
                except Exception as clean_error:
                    logger.error(f"清理文件失敗: {clean_error}")

    finally:
        duration = time.time() - start_time
        log_message = (
            f"表處理完成 - 表名: {table_name}, "
            f"狀態(tài): {status}, "
            f"導(dǎo)出行數(shù): {rows_exported}, "
            f"耗時(shí): {duration:.2f}秒"
        )
        logger.info(log_message)

def unzip_files(zip_dir, unzip_dir):
    """解壓指定目錄中的所有ZIP文件"""
    zip_files = glob.glob(os.path.join(zip_dir, '*.zip'))
    if not zip_files:
        logger.warning("未找到ZIP文件，跳過解壓")
        return

    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures = []
        for zip_path in zip_files:
            futures.append(executor.submit(
                lambda: extract_zip(zip_path, unzip_dir)
            ))
        for future in concurrent.futures.as_completed(futures):
            try:
                future.result()
            except Exception as e:
                logger.error(f"解壓過程中發(fā)生錯(cuò)誤: {str(e)}")

def extract_zip(zip_path, unzip_dir):
    """解壓?jiǎn)蝹€(gè)ZIP文件"""
    try:
        start_time = time.time()
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(unzip_dir)
        duration = time.time() - start_time
        logger.info(f"解壓完成: {zip_path} => {unzip_dir} (耗時(shí): {duration:.2f}秒)")
    except Exception as e:
        logger.error(f"解壓 {zip_path} 失敗: {str(e)}", exc_info=True)
        raise

def main():
    # 配置參數(shù)
    db_config = {
        'host': 'localhost',
        'user': 'your_username',
        'password': 'your_password',
        'database': 'your_database'
    }
    
    # 目錄配置
    base_dir = os.path.dirname(os.path.abspath(__file__))
    csv_dir = os.path.join(base_dir, 'csv_exports')
    zip_dir = os.path.join(base_dir, 'zip_archives')
    unzip_dir = os.path.join(base_dir, 'unzipped_files')

    # 創(chuàng)建目錄
    for dir_path in [csv_dir, zip_dir, unzip_dir]:
        os.makedirs(dir_path, exist_ok=True)
        logger.info(f"目錄已準(zhǔn)備: {dir_path}")

    # 獲取所有表名
    try:
        conn = mysql.connector.connect(**db_config)
        cursor = conn.cursor()
        cursor.execute("SHOW TABLES")
        tables = [table[0] for table in cursor.fetchall()]
        logger.info(f"發(fā)現(xiàn) {len(tables)} 個(gè)需要處理的表")
    except Exception as e:
        logger.error(f"獲取數(shù)據(jù)庫表失敗: {str(e)}", exc_info=True)
        return
    finally:
        if 'cursor' in locals():
            cursor.close()
        if 'conn' in locals() and conn.is_connected():
            conn.close()

    # 處理所有表（多線程導(dǎo)出和壓縮）
    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
        futures = []
        for table in tables:
            futures.append(executor.submit(
                process_table,
                table,
                db_config,
                csv_dir,
                zip_dir
            ))

        # 處理任務(wù)結(jié)果
        for future in concurrent.futures.as_completed(futures):
            try:
                future.result()
            except Exception as e:
                logger.error(f"表處理異常: {str(e)}")

    # 解壓所有ZIP文件（多線程解壓）
    logger.info("開始解壓所有ZIP文件")
    unzip_files(zip_dir, unzip_dir)
    logger.info("全部處理流程完成")

if __name__ == "__main__":
    main()

pip install mysql-connector-python

python script.py