65 lines
2.0 KiB
Python
65 lines
2.0 KiB
Python
|
"""
|
||
|
|
Migration script to add kb_type, priority, is_enabled, doc_count to knowledge_bases.
|
||
|
|
Run: python scripts/migrations/run_migration.py
|
||
|
|
"""
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
|
||
|
|
|
||
|
|
from sqlalchemy import text
|
||
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
||
|
|
from sqlalchemy.orm import sessionmaker
|
||
|
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
||
|
|
|
||
|
|
from app.core.config import get_settings
|
||
|
|
|
||
|
|
|
||
|
|
async def run_migration():
|
||
|
|
"""Run the migration to add new columns to knowledge_bases table."""
|
||
|
|
settings = get_settings()
|
||
|
|
engine = create_async_engine(settings.database_url, echo=True)
|
||
|
|
async_session_maker = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
||
|
|
|
||
|
|
migration_sql = """
|
||
|
|
-- Add kb_type column
|
||
|
|
ALTER TABLE knowledge_bases
|
||
|
|
ADD COLUMN IF NOT EXISTS kb_type VARCHAR DEFAULT 'general';
|
||
|
|
|
||
|
|
-- Add priority column
|
||
|
|
ALTER TABLE knowledge_bases
|
||
|
|
ADD COLUMN IF NOT EXISTS priority INTEGER DEFAULT 0;
|
||
|
|
|
||
|
|
-- Add is_enabled column
|
||
|
|
ALTER TABLE knowledge_bases
|
||
|
|
ADD COLUMN IF NOT EXISTS is_enabled BOOLEAN DEFAULT TRUE;
|
||
|
|
|
||
|
|
-- Add doc_count column
|
||
|
|
ALTER TABLE knowledge_bases
|
||
|
|
ADD COLUMN IF NOT EXISTS doc_count INTEGER DEFAULT 0;
|
||
|
|
"""
|
||
|
|
|
||
|
|
async with async_session_maker() as session:
|
||
|
|
for statement in migration_sql.strip().split(';'):
|
||
|
|
statement = statement.strip()
|
||
|
|
if statement and not statement.startswith('--'):
|
||
|
|
try:
|
||
|
|
await session.execute(text(statement))
|
||
|
|
print(f"Executed: {statement[:50]}...")
|
||
|
|
except Exception as e:
|
||
|
|
if "already exists" in str(e).lower() or "duplicate" in str(e).lower():
|
||
|
|
print(f"Skipped (already exists): {statement[:50]}...")
|
||
|
|
else:
|
||
|
|
raise
|
||
|
|
|
||
|
|
await session.commit()
|
||
|
|
print("\nMigration completed successfully!")
|
||
|
|
|
||
|
|
await engine.dispose()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
asyncio.run(run_migration())
|