shhhoaib/pakshield-backend

1from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker2from sqlalchemy.orm import DeclarativeBase3from app.config import get_settings4 5settings = get_settings()6 7_engine = None8 9 10def get_engine():11    global _engine12    if _engine is None:13        connect_args = {}14        if settings.database_url.startswith("sqlite"):15            connect_args = {"check_same_thread": False}16        _engine = create_async_engine(17            settings.database_url,18            echo=False,19            connect_args=connect_args,20        )21    return _engine22 23 24async_session = async_sessionmaker(get_engine(), class_=AsyncSession, expire_on_commit=False)25 26 27class Base(DeclarativeBase):28    pass29 30 31async def get_db():32    async with async_session() as session:33        try:34            yield session35        finally:36            await session.close()37 38 39async def init_db():40    async with get_engine().begin() as conn:41        await conn.run_sync(Base.metadata.create_all)42 43    async with get_engine().begin() as conn:44        def _migrate(conn):45            import sqlalchemy as sa46            inspector = sa.inspect(conn)47            columns = [c["name"] for c in inspector.get_columns("scan_records")]48            for col in ["feature_scores_json", "security_json", "reasons_json"]:49                if col not in columns:50                    conn.execute(sa.text(f"ALTER TABLE scan_records ADD COLUMN {col} TEXT"))51        await conn.run_sync(_migrate)52