feat: add production migration script for llm_catalog_rel

2026-05-16 21:31:50 +08:00 · 2026-05-16 21:31:50 +08:00 · 2e86b7a008
commit 2e86b7a008
parent 051d15f1af
1 changed files with 124 additions and 0 deletions
--- a/migrate_llm_catelog_prod.py
+++ b/migrate_llm_catelog_prod.py
@ -0,0 +1,124 @@
+#!/usr/bin/env python3
+"""
+Production Migration Script for llm -> llm_catalog_rel.
+This script handles:
+1. Creating the new relationship table if it doesn't exist.
+2. Migrating existing single-type associations to the new table.
+3. Dropping the old column (optional but recommended).
+
+Usage:
+    cd /home/hermesai/repos/sage
+    ./py3/bin/python migrate_llm_catelog_prod.py
+"""
+import asyncio
+import sys
+sys.path.insert(0, '.')
+
+from sqlor.dbpools import DBPools
+from appPublic.jsonConfig import getConfig
+from appPublic.uniqueID import getID
+
+async def migrate():
+    config = getConfig('.')
+    db = DBPools(config.databases)
+    
+    # Sage 数据库
+    dbname = 'sage'
+    
+    async with db.sqlorContext(dbname) as sor:
+        # 1. 检查列是否存在
+        check_col_sql = "select count(*) as cnt from information_schema.columns where table_schema = database() and table_name = 'llm' and column_name = 'llmcatelogid'"
+        res = await sor.sqlExe(check_col_sql, {})
+        has_column = res[0].cnt > 0 if res else False
+
+        if not has_column:
+            print("[SKIP] Column 'llmcatelogid' does not exist in 'llm' table. Migration not needed.")
+            return True
+
+        # 2. 创建表 (如果不存在)
+        print("[STEP 1] Ensuring table 'llm_catalog_rel' exists...")
+        create_sql = """
+        CREATE TABLE IF NOT EXISTS llm_catalog_rel (
+            id VARCHAR(32) NOT NULL PRIMARY KEY,
+            llmid VARCHAR(32) NOT NULL,
+            llmcatelogid VARCHAR(32) NOT NULL,
+            UNIQUE KEY uq_llm_catelog (llmid, llmcatelogid),
+            INDEX idx_llm (llmid),
+            INDEX idx_catelog (llmcatelogid)
+        )
+        """
+        try:
+            await sor.sqlExe(create_sql, {})
+            print("  -> Table created or exists.")
+        except Exception as e:
+            print(f"  -> Table creation warning: {e}")
+
+        # 2.1 确保唯一索引存在
+        print("[STEP 1.1] Ensuring unique index 'uq_llm_catelog' exists...")
+        try:
+            await sor.sqlExe("ALTER TABLE llm_catalog_rel ADD UNIQUE INDEX uq_llm_catelog (llmid, llmcatelogid)", {})
+            print("  -> Unique index added.")
+        except Exception as e:
+            if "Duplicate key name" in str(e) or "Multiple primary key defined" in str(e):
+                print("  -> Unique index already exists.")
+            else:
+                print(f"  -> Index creation warning (may already exist): {e}")
+
+        # 3. 迁移数据
+        print("[STEP 2] Migrating data...")
+        sql = "select id, llmcatelogid from llm where llmcatelogid is not null and llmcatelogid != ''"
+        rows = await sor.sqlExe(sql, {})
+        print(f"  -> Found {len(rows)} records to migrate.")
+        
+        migrated_count = 0
+        error_count = 0
+        
+        for r in rows:
+            # Check if already migrated to avoid duplicates if re-running
+            check_rel_sql = "select count(*) as cnt from llm_catalog_rel where llmid = ${llmid}$ and llmcatelogid = ${catelogid}$"
+            rel_res = await sor.sqlExe(check_rel_sql, {'llmid': r['id'], 'catelogid': r['llmcatelogid']})
+            if rel_res and rel_res[0].cnt > 0:
+                continue
+
+            new_id = getID()
+            data = {
+                'id': new_id,
+                'llmid': r['id'],
+                'llmcatelogid': r['llmcatelogid']
+            }
+            try:
+                await sor.C('llm_catalog_rel', data)
+                migrated_count += 1
+            except Exception as e:
+                print(f"  -> Error migrating {r['id']}: {e}")
+                error_count += 1
+
+        print(f"  -> Migration complete. Inserted: {migrated_count}, Errors: {error_count}")
+
+        # 4. 删除旧列 (生产环境谨慎执行，建议确认迁移成功后手动执行)
+        # print("[STEP 3] Dropping old column 'llmcatelogid'...")
+        # try:
+        #     await sor.sqlExe("alter table llm drop column llmcatelogid", {})
+        #     print("  -> Column dropped.")
+        # except Exception as e:
+        #     print(f"  -> Drop column error: {e}")
+
+        print("\n[MIGRATION DONE]")
+        print("Please verify data in 'llm_catalog_rel'.")
+        print("If correct, you may manually run: alter table llm drop column llmcatelogid;")
+
+    return True
+
+if __name__ == '__main__':
+    print("Starting production migration...")
+    try:
+        success = asyncio.get_event_loop().run_until_complete(migrate())
+        if success:
+            sys.exit(0)
+        else:
+            sys.exit(1)
+    except Exception as e:
+        print(f"FATAL ERROR: {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)