58 lines
1.5 KiB
Python
58 lines
1.5 KiB
Python
import time
|
|
from contextlib import asynccontextmanager
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession
|
|
|
|
from app.core._config import DATABASE_URL
|
|
from app.core.logger import make_log
|
|
|
|
|
|
def _to_async_dsn(url: str) -> str:
|
|
# Convert psycopg2 DSN to asyncpg DSN
|
|
# postgresql+psycopg2://user:pass@host:5432/db -> postgresql+asyncpg://user:pass@host:5432/db
|
|
return url.replace("+psycopg2", "+asyncpg")
|
|
|
|
|
|
# Async engine for PostgreSQL
|
|
engine = create_async_engine(
|
|
_to_async_dsn(DATABASE_URL),
|
|
pool_size=10,
|
|
max_overflow=20,
|
|
pool_timeout=30,
|
|
pool_recycle=1800,
|
|
pool_pre_ping=True,
|
|
)
|
|
|
|
AsyncSessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
|
|
|
|
|
|
async def wait_db_ready():
|
|
ready = False
|
|
while not ready:
|
|
try:
|
|
async with engine.connect() as conn:
|
|
await conn.execute(text("SELECT 1"))
|
|
ready = True
|
|
except Exception as e:
|
|
make_log("SQL", 'PostgreSQL is not ready yet: ' + str(e), level='debug')
|
|
time.sleep(1)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def db_session(auto_commit: bool = False):
|
|
session: AsyncSession = AsyncSessionLocal()
|
|
try:
|
|
yield session
|
|
if auto_commit:
|
|
await session.commit()
|
|
except BaseException as e:
|
|
await session.rollback()
|
|
raise e
|
|
finally:
|
|
await session.close()
|
|
|
|
|
|
def new_session() -> AsyncSession:
|
|
return AsyncSessionLocal()
|