| from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
| from sqlalchemy.orm import DeclarativeBase
|
| from sqlalchemy import create_engine
|
| import os
|
|
|
| DATABASE_URL = os.getenv("DATABASE_URL", "sqlite+aiosqlite:///./data/1proxy.db")
|
|
|
|
|
| if "sqlite" in DATABASE_URL and "./data/" in DATABASE_URL:
|
| os.makedirs("./data", exist_ok=True)
|
|
|
|
|
| engine_kwargs = {
|
| "echo": False,
|
| "future": True,
|
| }
|
|
|
|
|
| if DATABASE_URL.startswith("postgresql"):
|
| engine_kwargs.update(
|
| {
|
| "pool_size": 20,
|
| "max_overflow": 30,
|
| "pool_pre_ping": True,
|
| "pool_recycle": 3600,
|
| "pool_timeout": 30,
|
| }
|
| )
|
|
|
| engine = create_async_engine(DATABASE_URL, **engine_kwargs)
|
|
|
| AsyncSessionLocal = async_sessionmaker(
|
| engine,
|
| class_=AsyncSession,
|
| expire_on_commit=False,
|
| )
|
|
|
|
|
| class Base(DeclarativeBase):
|
| pass
|
|
|
|
|
| async def get_db():
|
| async with AsyncSessionLocal() as session:
|
| try:
|
| yield session
|
| await session.commit()
|
| except Exception:
|
| await session.rollback()
|
| raise
|
| finally:
|
| await session.close()
|
|
|
|
|
| async def init_db():
|
| async with engine.begin() as conn:
|
| await conn.run_sync(Base.metadata.create_all)
|
|
|