dc89e45675
- Updated type hints from List/Optional to list/None for better readability and consistency across the codebase. - Refactored import statements for better organization and clarity. - Enhanced the ScheduledTaskBase schema to use modern type hints. - Cleaned up unnecessary comments and whitespace in various files. - Improved error handling and logging in task execution handlers. - Updated test cases to reflect changes in type hints and ensure compatibility with the new structure.
66 lines
1.9 KiB
Python
66 lines
1.9 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
|
|
from app.core.seeds import seed_all_data
|
|
|
|
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 and create tables if they do not exist."""
|
|
logger = get_logger(__name__)
|
|
try:
|
|
logger.info("Initializing database tables")
|
|
async with engine.begin() as conn:
|
|
await conn.run_sync(SQLModel.metadata.create_all)
|
|
logger.info("Database tables created successfully")
|
|
|
|
# Seed initial data
|
|
await seed_initial_data()
|
|
|
|
except Exception:
|
|
logger.exception("Failed to initialize database")
|
|
raise
|
|
|
|
|
|
async def seed_initial_data() -> None:
|
|
"""Seed initial data into the database."""
|
|
logger = get_logger(__name__)
|
|
logger.info("Starting initial data seeding")
|
|
|
|
async with AsyncSession(engine) as session:
|
|
await seed_all_data(session)
|