Agent/migrate_to_resource.py

#!/usr/bin/env python3
"""
安全迁移脚本：contents → resources + knowledge.resource_ids
在服务器上执行此脚本，然后再拉取新代码

变更内容：
1. contents表 → resources表
2. 为resources表添加新字段（secure_body, content_type, metadata, updated_at）
3. 为knowledge表添加resource_ids字段
"""

import sqlite3
import sys
from pathlib import Path

def migrate_resources_table(conn, cursor):
    """迁移resources表（原contents表）"""
    print("\n" + "="*60)
    print("步骤1: 迁移 resources 表")
    print("="*60)

    # 检查contents表是否存在
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='contents'")
    if not cursor.fetchone():
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='resources'")
        if cursor.fetchone():
            print("✅ resources表已存在，跳过此步骤")
            return True
        else:
            print("⚠️  既没有contents表也没有resources表，将在首次启动时创建")
            return False

    # 检查contents表中的数据
    cursor.execute("SELECT COUNT(*) FROM contents")
    count = cursor.fetchone()[0]
    print(f"contents表中有 {count} 条记录")

    # 检查是否已有新字段
    cursor.execute("PRAGMA table_info(contents)")
    columns = {row[1] for row in cursor.fetchall()}
    print(f"现有字段: {columns}")

    # 添加新字段（如果不存在）
    migrations = []
    if "secure_body" not in columns:
        migrations.append("ALTER TABLE contents ADD COLUMN secure_body TEXT DEFAULT ''")
    if "content_type" not in columns:
        migrations.append("ALTER TABLE contents ADD COLUMN content_type TEXT DEFAULT 'text'")
    if "metadata" not in columns:
        migrations.append("ALTER TABLE contents ADD COLUMN metadata TEXT DEFAULT '{}'")
    if "updated_at" not in columns:
        migrations.append("ALTER TABLE contents ADD COLUMN updated_at TEXT DEFAULT ''")

    if migrations:
        print(f"\n添加 {len(migrations)} 个新字段...")
        for sql in migrations:
            print(f"  {sql}")
            cursor.execute(sql)
        conn.commit()
        print("✅ 字段添加完成")

    # 重命名表
    print("\n重命名表: contents → resources")
    try:
        cursor.execute("ALTER TABLE contents RENAME TO resources")
        conn.commit()
        print("✅ 表重命名完成")
    except Exception as e:
        print(f"❌ 重命名失败: {e}")
        return False

    # 验证
    cursor.execute("SELECT COUNT(*) FROM resources")
    new_count = cursor.fetchone()[0]
    print(f"resources表中有 {new_count} 条记录")

    if new_count == count:
        print(f"✅ 数据完整，{count} 条记录全部保留")
        return True
    else:
        print(f"⚠️  数据不一致: 原 {count} 条 → 现 {new_count} 条")
        return False


def migrate_knowledge_table(conn, cursor):
    """为knowledge表添加resource_ids字段"""
    print("\n" + "="*60)
    print("步骤2: 更新 knowledge 表")
    print("="*60)

    # 检查knowledge表是否存在
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='knowledge'")
    if not cursor.fetchone():
        print("⚠️  knowledge表不存在，跳过此步骤")
        return True

    # 检查是否已有resource_ids字段
    cursor.execute("PRAGMA table_info(knowledge)")
    columns = {row[1] for row in cursor.fetchall()}
    print(f"现有字段: {columns}")

    if "resource_ids" in columns:
        print("✅ resource_ids字段已存在，跳过此步骤")
        return True

    # 添加resource_ids字段
    print("\n添加 resource_ids 字段...")
    try:
        cursor.execute("ALTER TABLE knowledge ADD COLUMN resource_ids TEXT DEFAULT '[]'")
        conn.commit()
        print("✅ resource_ids字段添加完成")
        return True
    except Exception as e:
        print(f"❌ 添加字段失败: {e}")
        return False


def migrate():
    """主迁移函数"""
def migrate():
    """主迁移函数"""
    # 查找数据库文件
    db_path = Path("knowhub.db")
    if not db_path.exists():
        print("❌ 找不到 knowhub.db")
        print("请在包含数据库的目录中运行此脚本")
        sys.exit(1)

    print(f"数据库路径: {db_path.absolute()}")

    conn = sqlite3.connect(str(db_path))
    cursor = conn.cursor()

    # 显示当前表状态
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
    tables = {row[0] for row in cursor.fetchall()}
    print(f"\n当前表: {tables}")

    # 执行迁移
    success = True

    # 步骤1: 迁移resources表
    if not migrate_resources_table(conn, cursor):
        success = False

    # 步骤2: 更新knowledge表
    if not migrate_knowledge_table(conn, cursor):
        success = False

    # 最终验证
    print("\n" + "="*60)
    print("迁移总结")
    print("="*60)

    cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
    tables = {row[0] for row in cursor.fetchall()}
    print(f"迁移后的表: {tables}")

    # 检查resources表
    if "resources" in tables:
        cursor.execute("SELECT COUNT(*) FROM resources")
        count = cursor.fetchone()[0]
        print(f"✅ resources表: {count} 条记录")
    else:
        print("⚠️  resources表不存在")

    # 检查knowledge表
    if "knowledge" in tables:
        cursor.execute("SELECT COUNT(*) FROM knowledge")
        count = cursor.fetchone()[0]
        cursor.execute("PRAGMA table_info(knowledge)")
        columns = {row[1] for row in cursor.fetchall()}
        has_resource_ids = "resource_ids" in columns
        print(f"✅ knowledge表: {count} 条记录, resource_ids字段: {'存在' if has_resource_ids else '不存在'}")
    else:
        print("⚠️  knowledge表不存在")

    conn.close()

    if success:
        print("\n✅ 迁移完成！现在可以拉取新代码并重启服务")
    else:
        print("\n⚠️  迁移过程中有警告，请检查上述输出")

    return success


if __name__ == "__main__":
    print("=" * 60)
    print("KnowHub 数据库迁移")
    print("变更内容：")
    print("  1. contents表 → resources表")
    print("  2. resources表添加新字段")
    print("  3. knowledge表添加resource_ids字段")
    print("=" * 60)

    try:
        success = migrate()
        sys.exit(0 if success else 1)
    except Exception as e:
        print(f"\n❌ 迁移失败: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)