database.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
  2. from sqlalchemy.orm import sessionmaker
  3. from sqlalchemy import text
  4. from app.models import Base
  5. from app.config import get_settings
  6. settings = get_settings()
  7. # Convert sqlite:/// to sqlite+aiosqlite:/// for async support
  8. db_url = settings.database_url.replace("sqlite://", "sqlite+aiosqlite://")
  9. # Create async engine
  10. engine = create_async_engine(db_url, echo=True)
  11. # Create async session factory
  12. async_session = async_sessionmaker(
  13. engine, class_=AsyncSession, expire_on_commit=False
  14. )
  15. async def init_db():
  16. """Initialize database tables."""
  17. async with engine.begin() as conn:
  18. # Check if migration is needed (users table doesn't exist)
  19. result = await conn.execute(text(
  20. "SELECT name FROM sqlite_master WHERE type='table' AND name='users'"
  21. ))
  22. needs_migration = result.fetchone() is None
  23. # Check if we have old schema (listening_sessions exists but no users table)
  24. result = await conn.execute(text(
  25. "SELECT name FROM sqlite_master WHERE type='table' AND name='listening_sessions'"
  26. ))
  27. has_old_schema = result.fetchone() is not None
  28. if needs_migration and has_old_schema:
  29. # Need to run migration
  30. print("Existing database detected without users table - migration required")
  31. print("Please run: python -m app.migrations.add_multi_user")
  32. print("Or delete absrecommend.db to start fresh")
  33. raise RuntimeError("Database migration required")
  34. # Create all tables (will skip existing ones)
  35. await conn.run_sync(Base.metadata.create_all)
  36. async def get_db():
  37. """Dependency for getting database session."""
  38. async with async_session() as session:
  39. try:
  40. yield session
  41. finally:
  42. await session.close()