TorqueWrench/backend/migrate_add_device_info.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
数据库迁移脚本：为work_results表添加device_sn和device_name字段
"""

import sqlite3
import os
from datetime import datetime


def migrate_add_device_info():
    """为work_results表添加设备信息字段"""
    db_path = "wrench.db"
    
    if not os.path.exists(db_path):
        print(f"数据库文件不存在: {db_path}")
        return
    
    print("="*60)
    print("数据库迁移：添加设备信息字段")
    print("="*60)
    
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        # 检查表结构
        cursor.execute("PRAGMA table_info(work_results)")
        columns = cursor.fetchall()
        column_names = [col[1] for col in columns]
        
        print(f"当前字段: {column_names}")
        
        # 检查是否需要添加字段
        needs_migration = False
        if 'device_sn' not in column_names:
            print("需要添加 device_sn 字段")
            needs_migration = True
        if 'device_name' not in column_names:
            print("需要添加 device_name 字段")
            needs_migration = True
        
        if not needs_migration:
            print("[OK] 字段已存在，无需迁移")
            conn.close()
            return
        
        print("\n开始迁移...")
        
        # SQLite不支持直接添加列，需要重建表
        # 1. 创建新表（包含新字段）
        print("1. 创建新表结构...")
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS work_results_new (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                trace_id TEXT NOT NULL,
                process_id TEXT NOT NULL,
                process_name TEXT,
                bolts TEXT NOT NULL,
                device_sn TEXT,
                device_name TEXT,
                submitted_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        
        # 2. 复制数据（旧字段映射到新表）
        print("2. 复制数据...")
        cursor.execute('''
            INSERT INTO work_results_new 
            (id, trace_id, process_id, process_name, bolts, device_sn, device_name, submitted_at)
            SELECT 
                id, trace_id, process_id, process_name, bolts, NULL, NULL, submitted_at
            FROM work_results
        ''')
        
        # 3. 删除旧表
        print("3. 删除旧表...")
        cursor.execute('DROP TABLE work_results')
        
        # 4. 重命名新表
        print("4. 重命名新表...")
        cursor.execute('ALTER TABLE work_results_new RENAME TO work_results')
        
        conn.commit()
        
        # 验证
        cursor.execute("PRAGMA table_info(work_results)")
        new_columns = cursor.fetchall()
        new_column_names = [col[1] for col in new_columns]
        
        print(f"\n[OK] 迁移完成！")
        print(f"新表字段: {new_column_names}")
        
        # 统计数据
        cursor.execute("SELECT COUNT(*) FROM work_results")
        count = cursor.fetchone()[0]
        print(f"[OK] 结果记录数量: {count}")
        
    except Exception as e:
        conn.rollback()
        print(f"[FAIL] 迁移失败: {e}")
        import traceback
        traceback.print_exc()
    finally:
        conn.close()
    
    print("="*60)


if __name__ == "__main__":
    migrate_add_device_info()