Agent/knowhub/scripts/phase2_schema_migration.py

#!/usr/bin/env python3
"""
Phase 2 schema migration：
  1. ALTER requirement_strategy 加 3 列（is_selected / coverage_score / coverage_explanation）
  2. CREATE TABLE requirement_knowledge（带 metadata）
  3. CREATE TABLE knowledge_resource（纯 junction）

幂等：IF NOT EXISTS / 检查列存在。autocommit=True 安全。
"""
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from knowhub.knowhub_db.pg_capability_store import PostgreSQLCapabilityStore


def column_exists(cur, table, col):
    cur.execute("""SELECT EXISTS (SELECT 1 FROM information_schema.columns
                   WHERE table_name=%s AND column_name=%s)""", (table, col))
    return cur.fetchone()['exists']


def table_exists(cur, table):
    cur.execute("""SELECT EXISTS (SELECT 1 FROM information_schema.tables
                   WHERE table_name=%s)""", (table,))
    return cur.fetchone()['exists']


def main():
    s = PostgreSQLCapabilityStore()
    cur = s._get_cursor()
    try:
        # ─────────────────────────────────────────────────
        # Step 1: ALTER requirement_strategy
        print('=== Step 1: ALTER requirement_strategy ===', flush=True)
        for col, typ in [('is_selected', 'BOOLEAN'),
                         ('coverage_score', 'FLOAT'),
                         ('coverage_explanation', 'TEXT')]:
            if column_exists(cur, 'requirement_strategy', col):
                print(f'  {col} 已存在，跳过', flush=True)
            else:
                cur.execute(f'ALTER TABLE requirement_strategy ADD COLUMN {col} {typ}')
                print(f'  ✓ 添加 {col} {typ}', flush=True)

        # ─────────────────────────────────────────────────
        # Step 2: CREATE requirement_knowledge
        print('\n=== Step 2: CREATE requirement_knowledge ===', flush=True)
        if table_exists(cur, 'requirement_knowledge'):
            print('  已存在，跳过', flush=True)
        else:
            cur.execute("""
                CREATE TABLE requirement_knowledge (
                    requirement_id        VARCHAR NOT NULL,
                    knowledge_id          VARCHAR NOT NULL,
                    is_selected           BOOLEAN,
                    coverage_score        FLOAT,
                    coverage_explanation  TEXT,
                    PRIMARY KEY (requirement_id, knowledge_id)
                ) DISTRIBUTED BY (requirement_id)
            """)
            print('  ✓ 已建表', flush=True)

        # ─────────────────────────────────────────────────
        # Step 3: CREATE knowledge_resource
        print('\n=== Step 3: CREATE knowledge_resource ===', flush=True)
        if table_exists(cur, 'knowledge_resource'):
            print('  已存在，跳过', flush=True)
        else:
            cur.execute("""
                CREATE TABLE knowledge_resource (
                    knowledge_id  VARCHAR NOT NULL,
                    resource_id   VARCHAR NOT NULL,
                    PRIMARY KEY (knowledge_id, resource_id)
                ) DISTRIBUTED BY (knowledge_id)
            """)
            print('  ✓ 已建表', flush=True)

        # ─────────────────────────────────────────────────
        # 验证
        print('\n=== 验证 ===', flush=True)
        cur.execute("""SELECT column_name, data_type FROM information_schema.columns
                       WHERE table_name='requirement_strategy' ORDER BY ordinal_position""")
        print('requirement_strategy 列:')
        for r in cur.fetchall():
            print(f'  {r["column_name"]:28s} {r["data_type"]}')
        for t in ['requirement_knowledge', 'knowledge_resource']:
            cur.execute("""SELECT column_name, data_type FROM information_schema.columns
                           WHERE table_name=%s ORDER BY ordinal_position""", (t,))
            print(f'\n{t} 列:')
            for r in cur.fetchall():
                print(f'  {r["column_name"]:28s} {r["data_type"]}')
    finally:
        cur.close()
        s.close()


if __name__ == '__main__':
    main()