Spaces:
Sleeping
Sleeping
File size: 4,756 Bytes
db7c1e8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | #!/usr/bin/env python3
"""
Database Migration Script for AI Backend with RAG + Authentication
This script creates all required tables in the Neon Postgres database
"""
import asyncio
import sys
import logging
import os
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy import text
# Add the src directory to the path so we can import our modules
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
src_dir = os.path.join(project_root, "src")
sys.path.insert(0, src_dir)
from db.base import Base
from config.settings import settings
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
handlers=[
logging.StreamHandler(sys.stdout)
]
)
logger = logging.getLogger(__name__)
async def create_database_tables():
"""Create all required database tables"""
logger.info("Starting database migration...")
# Create async engine
engine = create_async_engine(
settings.neon_db_url,
echo=True, # Set to True to log SQL queries during migration
pool_pre_ping=True,
pool_size=5,
max_overflow=10,
)
try:
logger.info("Connecting to database...")
# Create all tables
async with engine.begin() as conn:
logger.info("Creating tables...")
await conn.run_sync(Base.metadata.create_all)
logger.info("Tables created successfully")
# Verify tables exist by checking if we can access them
tables = Base.metadata.tables.keys()
logger.info(f"Created {len(tables)} tables: {list(tables)}")
# Test connection by running a simple query
result = await conn.execute(text("SELECT 1"))
logger.info("Database connection test successful")
logger.info("Database migration completed successfully!")
return True
except Exception as e:
logger.error(f"Database migration failed: {e}")
return False
finally:
# Dispose of the engine to close all connections
await engine.dispose()
async def check_existing_tables():
"""Check if tables already exist in the database"""
engine = create_async_engine(
settings.neon_db_url,
pool_pre_ping=True,
)
try:
async with engine.begin() as conn:
# Get list of existing tables
result = await conn.execute(text("""
SELECT table_name
FROM information_schema.tables
WHERE table_schema = 'public'
"""))
existing_tables = [row[0] for row in result]
logger.info(f"Existing tables in database: {existing_tables}")
# Check which of our models' tables already exist
our_tables = [table.name for table in Base.metadata.tables.values()]
existing_our_tables = [table for table in our_tables if table in existing_tables]
if existing_our_tables:
logger.warning(f"The following tables already exist: {existing_our_tables}")
return existing_our_tables
else:
logger.info("No existing tables from our models found.")
return []
except Exception as e:
logger.error(f"Error checking existing tables: {e}")
return []
finally:
await engine.dispose()
async def run_migration():
"""Main function to run the migration"""
logger.info("Migration script started")
# Check for existing tables
existing_tables = await check_existing_tables()
if existing_tables:
logger.info("Some tables already exist. Migration will continue and handle existing tables gracefully.")
else:
logger.info("No existing tables found. Proceeding with fresh migration.")
# Create tables
success = await create_database_tables()
if success:
logger.info("Migration completed successfully!")
return True
else:
logger.error("Migration failed!")
return False
if __name__ == "__main__":
# Check if we're running in a virtual environment or have the required dependencies
try:
import sqlalchemy
import asyncpg
except ImportError as e:
logger.error(f"Required dependencies not found: {e}")
logger.error("Please make sure you've installed the requirements: pip install -r requirements.txt")
sys.exit(1)
# Run the migration
success = asyncio.run(run_migration())
if success:
logger.info("Migration script completed successfully!")
sys.exit(0)
else:
logger.error("Migration script failed!")
sys.exit(1) |