Agent/knowhub/knowhub_db/migrate_resources.py

#!/usr/bin/env python3
"""
迁移 SQLite resources 表中的 text 数据到 PostgreSQL
删除 code 类型数据
"""

import os
import json
import sqlite3
import psycopg2
from psycopg2.extras import RealDictCursor
from dotenv import load_dotenv
from datetime import datetime

load_dotenv()

CREATE_RESOURCES_TABLE = """
CREATE TABLE IF NOT EXISTS resources (
    id              TEXT PRIMARY KEY,
    title           TEXT,
    body            TEXT,
    secure_body     TEXT,
    content_type    TEXT,
    metadata        JSONB,
    sort_order      INTEGER DEFAULT 0,
    submitted_by    TEXT,
    created_at      BIGINT,
    updated_at      BIGINT
) WITH (appendoptimized=false);
"""

def migrate_resources():
    # 连接 SQLite
    sqlite_conn = sqlite3.connect('/root/Agent/knowhub/knowhub.db')
    sqlite_conn.row_factory = sqlite3.Row
    sqlite_cur = sqlite_conn.cursor()

    # 连接 PostgreSQL
    pg_conn = psycopg2.connect(
        host=os.getenv('KNOWHUB_DB'),
        port=int(os.getenv('KNOWHUB_PORT', 5432)),
        user=os.getenv('KNOWHUB_USER'),
        password=os.getenv('KNOWHUB_PASSWORD'),
        database=os.getenv('KNOWHUB_DB_NAME')
    )
    pg_conn.autocommit = True
    pg_cur = pg_conn.cursor(cursor_factory=RealDictCursor)

    # 创建表
    pg_cur.execute(CREATE_RESOURCES_TABLE)
    print("✅ resources 表已创建")

    # 读取 text 类型数据
    sqlite_cur.execute("SELECT * FROM resources WHERE content_type='text'")
    texts = sqlite_cur.fetchall()
    print(f"\n找到 {len(texts)} 条 text 数据")

    success = 0
    for text in texts:
        try:
            # 转换时间戳
            created_at = updated_at = None
            if text['created_at']:
                dt = datetime.fromisoformat(text['created_at'].replace('+00:00', ''))
                created_at = int(dt.timestamp())
            if text['updated_at']:
                dt = datetime.fromisoformat(text['updated_at'].replace('+00:00', ''))
                updated_at = int(dt.timestamp())

            pg_cur.execute("""
                INSERT INTO resources (id, title, body, secure_body, content_type,
                                       metadata, sort_order, submitted_by, created_at, updated_at)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                ON CONFLICT (id) DO UPDATE SET
                    title = EXCLUDED.title,
                    body = EXCLUDED.body,
                    secure_body = EXCLUDED.secure_body,
                    metadata = EXCLUDED.metadata,
                    sort_order = EXCLUDED.sort_order,
                    updated_at = EXCLUDED.updated_at
            """, (
                text['id'],
                text['title'],
                text['body'],
                text['secure_body'],
                text['content_type'],
                text['metadata'],
                text['sort_order'] or 0,
                text['submitted_by'],
                created_at,
                updated_at
            ))
            success += 1
        except Exception as e:
            print(f"❌ 迁移失败 {text['id']}: {e}")

    print(f"✅ text 数据迁移完成: {success}/{len(texts)}")

    # 删除 code 类型数据
    sqlite_cur.execute("SELECT COUNT(*) FROM resources WHERE content_type='code'")
    code_count = sqlite_cur.fetchone()[0]
    sqlite_cur.execute("DELETE FROM resources WHERE content_type='code'")
    sqlite_conn.commit()
    print(f"✅ 已删除 {code_count} 条 code 数据")

    sqlite_conn.close()
    pg_conn.close()

if __name__ == "__main__":
    migrate_resources()