Agent/knowhub/scripts/backup_before_strategy_refactor.py

#!/usr/bin/env python3
"""
备份 strategy 抽象化重构前的 DB 状态。
在 DB 内建 bk_20260422_<table> 表，快照以下 6 张表：
  strategy / requirement_strategy / strategy_capability
  strategy_resource / strategy_knowledge / knowledge

回滚时用：
  DELETE FROM <table>; INSERT INTO <table> SELECT * FROM bk_20260422_<table>;
  （或按 version 过滤回滚 knowledge 表）
"""
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

TABLES = [
    'strategy',
    'requirement_strategy',
    'strategy_capability',
    'strategy_resource',
    'strategy_knowledge',
    'knowledge',
]
BK_PREFIX = 'bk_20260422_'


def main():
    s = PostgreSQLCapabilityStore()
    cur = s._get_cursor()
    try:
        print(f'{"="*60}', flush=True)
        print(f'Backup strategy refactor prep — 2026-04-22', flush=True)
        print(f'{"="*60}\n', flush=True)

        results = []
        for t in TABLES:
            bk = f'{BK_PREFIX}{t}'
            # 检查 bk 表是否已存在
            cur.execute("""SELECT EXISTS (
              SELECT 1 FROM information_schema.tables WHERE table_name=%s)""", (bk,))
            exists = cur.fetchone()['exists']
            if exists:
                cur.execute(f'SELECT COUNT(*) c FROM {bk}')
                bk_count = cur.fetchone()['c']
                print(f'⚠️  {bk} 已存在（{bk_count} 行），跳过 CREATE', flush=True)
                # still count source
                cur.execute(f'SELECT COUNT(*) c FROM {t}')
                src_count = cur.fetchone()['c']
                results.append((t, src_count, bk_count, 'existed'))
                continue

            # 计数
            cur.execute(f'SELECT COUNT(*) c FROM {t}')
            src_count = cur.fetchone()['c']
            print(f'backing up {t} ({src_count} rows) → {bk}', flush=True)

            # CREATE TABLE AS (AnalyticDB 的 greenplum 分支支持)
            cur.execute(f'CREATE TABLE {bk} AS SELECT * FROM {t}')

            # 验证
            cur.execute(f'SELECT COUNT(*) c FROM {bk}')
            bk_count = cur.fetchone()['c']
            ok = '✓' if bk_count == src_count else '❌'
            print(f'  {ok} {bk}: {bk_count} rows', flush=True)
            results.append((t, src_count, bk_count, 'created'))

        print(f'\n{"="*60}', flush=True)
        print(f'{"Table":32s} {"src":>8s} {"bk":>8s} {"status":>10s}', flush=True)
        print(f'{"-"*60}', flush=True)
        all_ok = True
        for t, src, bk_n, status in results:
            ok = '✓' if src == bk_n else '❌'
            if src != bk_n: all_ok = False
            print(f'{t:32s} {src:>8d} {bk_n:>8d} {status:>10s} {ok}', flush=True)
        print(f'{"="*60}', flush=True)
        if all_ok:
            print('\n✅ Backup complete. All counts match.', flush=True)
        else:
            print('\n❌ Counts mismatch. INVESTIGATE before proceeding.', flush=True)

        # 显示回滚指令
        print('\n回滚指令（当需要时）：', flush=True)
        for t in TABLES:
            bk = f'{BK_PREFIX}{t}'
            if t == 'knowledge':
                print(f"  DELETE FROM {t} WHERE version = 'howard_strategy_instance';", flush=True)
                print(f"  -- 不删 v0 老数据；新 version 删除后相当于恢复到 1046 行状态", flush=True)
            else:
                print(f'  DELETE FROM {t}; INSERT INTO {t} SELECT * FROM {bk};', flush=True)
    finally:
        cur.close()
        s.close()


if __name__ == '__main__':
    main()