ai-contract-platform/scripts/backup_law_library.py

#!/usr/bin/env python3
"""
法律库备份脚本

功能：
1. 全量备份数据库
2. 增量备份法律条文
3. 备份验证和恢复测试

用法：
    python scripts/backup_law_library.py --type full
    python scripts/backup_law_library.py --type incremental
"""

import argparse
import os
import sys
import subprocess
from datetime import datetime
from pathlib import Path

# 配置
BACKUP_DIR = "backups/legal-library"
DATABASE_URL = os.environ.get("DATABASE_URL", "")

def create_backup_dir():
    """创建备份目录"""
    Path(BACKUP_DIR).mkdir(parents=True, exist_ok=True)
    return BACKUP_DIR

def full_backup():
    """全量备份"""
    print("Starting full backup...")
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_file = os.path.join(BACKUP_DIR, f"full_backup_{timestamp}.sql")

    # 使用pg_dump进行备份
    cmd = [
        "pg_dump",
        DATABASE_URL.replace("postgresql://", ""),
        "--format=plain",
        "--file", backup_file
    ]

    # 替代方案：直接导出关键表
    if not DATABASE_URL:
        print("ERROR: DATABASE_URL not set")
        return None

    import psycopg
    import json

    conn = psycopg.connect(DATABASE_URL)
    cur = conn.cursor()

    # 导出法律条文
    print("Exporting legal_articles...")
    cur.execute("SELECT * FROM legal_articles ORDER BY id")
    columns = [desc[0] for desc in cur.description]

    data_file = os.path.join(BACKUP_DIR, f"legal_articles_{timestamp}.json")
    with open(data_file, 'w', encoding='utf-8') as f:
        for row in cur.fetchall():
            record = dict(zip(columns, row))
            # 处理特殊类型
            for key, value in record.items():
                if hasattr(value, 'isoformat'):
                    record[key] = value.isoformat()
            f.write(json.dumps(record, ensure_ascii=False) + '\n')

    # 统计
    cur.execute("SELECT COUNT(*) FROM legal_articles")
    total = cur.fetchone()[0]

    conn.close()

    print(f"Backup completed: {total} articles")
    print(f"Backup file: {data_file}")

    # 生成备份报告
    report = {
        "timestamp": timestamp,
        "type": "full",
        "total_articles": total,
        "backup_file": data_file,
        "status": "success"
    }

    report_file = os.path.join(BACKUP_DIR, f"backup_report_{timestamp}.json")
    with open(report_file, 'w', encoding='utf-8') as f:
        json.dump(report, f, ensure_ascii=False, indent=2)

    return data_file

def incremental_backup():
    """增量备份"""
    print("Starting incremental backup...")
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    if not DATABASE_URL:
        print("ERROR: DATABASE_URL not set")
        return None

    import psycopg
    import json

    conn = psycopg.connect(DATABASE_URL)
    cur = conn.cursor()

    # 获取上次备份时间（假设每24小时备份一次）
    cur.execute('''
        SELECT MAX(updated_at) FROM legal_articles
        WHERE updated_at IS NOT NULL
    ''')
    last_update = cur.fetchone()[0]

    if last_update:
        print(f"Last update: {last_update}")

        # 导出更新的记录
        cur.execute('''
            SELECT * FROM legal_articles
            WHERE updated_at >= NOW() - INTERVAL '1 day'
            ORDER BY id
        ''')

        columns = [desc[0] for desc in cur.description]
        data_file = os.path.join(BACKUP_DIR, f"incremental_{timestamp}.json")

        count = 0
        with open(data_file, 'w', encoding='utf-8') as f:
            for row in cur.fetchall():
                record = dict(zip(columns, row))
                for key, value in record.items():
                    if hasattr(value, 'isoformat'):
                        record[key] = value.isoformat()
                f.write(json.dumps(record, ensure_ascii=False) + '\n')
                count += 1

        print(f"Incremental backup: {count} articles")
        print(f"Backup file: {data_file}")

        conn.close()
        return data_file

    conn.close()
    return None

def verify_backup(backup_file):
    """验证备份"""
    print(f"Verifying backup: {backup_file}")

    if not os.path.exists(backup_file):
        print("ERROR: Backup file not found")
        return False

    # 检查文件大小
    size = os.path.getsize(backup_file)
    print(f"File size: {size / 1024 / 1024:.2f} MB")

    # 检查记录数
    with open(backup_file, 'r', encoding='utf-8') as f:
        count = sum(1 for _ in f)

    print(f"Records: {count}")

    return count > 0

def main():
    parser = argparse.ArgumentParser(description='Backup law library')
    parser.add_argument('--type', choices=['full', 'incremental', 'verify'],
                       default='full', help='Backup type')
    parser.add_argument('--file', help='Backup file to verify')
    args = parser.parse_args()

    create_backup_dir()

    if args.type == 'full':
        backup_file = full_backup()
        if backup_file:
            verify_backup(backup_file)
    elif args.type == 'incremental':
        backup_file = incremental_backup()
        if backup_file:
            verify_backup(backup_file)
    elif args.type == 'verify':
        if args.file:
            verify_backup(args.file)
        else:
            print("ERROR: --file required for verify")

    return 0

if __name__ == "__main__":
    sys.exit(main())