Bryan Van Deusen
61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
"""Database models and initialization."""
|
|
|
|
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
|
from sqlalchemy.orm import DeclarativeBase
|
|
|
|
from app.models.subscription import Subscription
|
|
from app.models.source import Source
|
|
from app.models.download import Download
|
|
from app.models.credential import Credential
|
|
from app.models.content import ContentItem
|
|
from app.models.setting import Setting
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
"""Base class for all models."""
|
|
pass
|
|
|
|
|
|
# Re-export models
|
|
__all__ = [
|
|
"Base",
|
|
"Subscription",
|
|
"Source",
|
|
"Download",
|
|
"Credential",
|
|
"ContentItem",
|
|
"Setting",
|
|
"init_db",
|
|
"get_session",
|
|
]
|
|
|
|
# Global engine and session factory
|
|
_engine = None
|
|
_session_factory = None
|
|
|
|
|
|
async def init_db(database_url: str):
|
|
"""Initialize database engine and create tables."""
|
|
global _engine, _session_factory
|
|
|
|
_engine = create_async_engine(database_url, echo=False)
|
|
_session_factory = async_sessionmaker(_engine, class_=AsyncSession, expire_on_commit=False)
|
|
|
|
# Import all models to register them with Base
|
|
from app.models.subscription import Subscription
|
|
from app.models.source import Source
|
|
from app.models.download import Download
|
|
from app.models.credential import Credential
|
|
from app.models.content import ContentItem
|
|
from app.models.setting import Setting
|
|
|
|
return _engine
|
|
|
|
|
|
async def get_session() -> AsyncSession:
|
|
"""Get a database session."""
|
|
if _session_factory is None:
|
|
raise RuntimeError("Database not initialized. Call init_db first.")
|
|
async with _session_factory() as session:
|
|
yield session
|