83239cb4fa
- Created alembic.ini configuration file for Alembic migrations. - Added README file for Alembic with a brief description. - Implemented env.py for Alembic to manage database migrations. - Created script.py.mako template for migration scripts. - Added initial migration script to create database tables. - Created a migration script to add initial plan and playlist data. - Updated database initialization to run Alembic migrations. - Enhanced credit service to automatically recharge user credits based on their plan. - Implemented delete_task method in scheduler service to remove scheduled tasks. - Updated scheduler API to reflect task deletion instead of cancellation. - Added CLI tool for managing database migrations. - Updated tests to cover new functionality for task deletion and credit recharge. - Updated pyproject.toml and lock files to include Alembic as a dependency.
62 lines
1.7 KiB
Python
62 lines
1.7 KiB
Python
from collections.abc import AsyncGenerator, Callable
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
|
|
from sqlmodel import SQLModel
|
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
|
|
|
# Import all models to ensure SQLModel metadata discovery
|
|
import app.models # noqa: F401
|
|
from app.core.config import settings
|
|
from app.core.logging import get_logger
|
|
|
|
engine: AsyncEngine = create_async_engine(
|
|
settings.DATABASE_URL,
|
|
echo=settings.DATABASE_ECHO,
|
|
)
|
|
|
|
|
|
async def get_db() -> AsyncGenerator[AsyncSession]:
|
|
"""Get a database session for dependency injection."""
|
|
logger = get_logger(__name__)
|
|
async with AsyncSession(engine) as session:
|
|
try:
|
|
yield session
|
|
except Exception:
|
|
logger.exception("Database session error")
|
|
await session.rollback()
|
|
raise
|
|
finally:
|
|
await session.close()
|
|
|
|
|
|
def get_session_factory() -> Callable[[], AsyncSession]:
|
|
"""Get a session factory function for services."""
|
|
|
|
def session_factory() -> AsyncSession:
|
|
return AsyncSession(engine)
|
|
|
|
return session_factory
|
|
|
|
|
|
async def init_db() -> None:
|
|
"""Initialize the database using Alembic migrations."""
|
|
logger = get_logger(__name__)
|
|
try:
|
|
logger.info("Running database migrations")
|
|
# Run Alembic migrations programmatically
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
|
|
# Get the alembic config
|
|
alembic_cfg = Config("alembic.ini")
|
|
|
|
# Run migrations to the latest revision
|
|
command.upgrade(alembic_cfg, "head")
|
|
logger.info("Database migrations completed successfully")
|
|
|
|
except Exception:
|
|
logger.exception("Failed to initialize database")
|
|
raise
|
|
|
|
|