sage/migrate_llm_catelog.py

69 lines
2.1 KiB
Python

import asyncio
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. Create table if not exists
print("Creating llm_catalog_rel table...")
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,
INDEX idx_llm (llmid),
INDEX idx_catelog (llmcatelogid)
)
"""
try:
await sor.sqlExe(create_sql, {})
print("Table llm_catalog_rel created or exists.")
except Exception as e:
print(f"Create table warning: {e}")
# 2. Migrate data
print("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 = 0
for r in rows:
new_id = getID()
data = {
'id': new_id,
'llmid': r['id'],
'llmcatelogid': r['llmcatelogid']
}
try:
await sor.C('llm_catalog_rel', data)
migrated += 1
except Exception as e:
print(f"Insert error for {r['id']}: {e}")
print(f"Migration complete. Migrated {migrated} records.")
# 3. Drop column
print("Dropping column llmcatelogid from llm...")
try:
await sor.sqlExe("alter table llm drop column llmcatelogid", {})
print("Column dropped.")
except Exception as e:
print(f"Drop column error: {e}")
return True
if __name__ == '__main__':
success = asyncio.get_event_loop().run_until_complete(migrate())
if success:
print('Migration complete.')
else:
print('Migration failed.')