Spaces:
Sleeping
Sleeping
File size: 1,451 Bytes
047d92b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | import os
import threading
from contextlib import asynccontextmanager
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
class DatabaseConfig:
_thread_local = threading.local()
@classmethod
def get_engine(cls):
if not hasattr(cls._thread_local, "engine"):
cls._thread_local.engine = create_async_engine(
url=os.getenv("SQLALCHEMY_DATABASE_URI"),
pool_recycle=1800,
pool_size=5,
max_overflow=45,
pool_pre_ping=True,
pool_timeout=60,
)
return cls._thread_local.engine
@classmethod
def _get_session_factory(cls):
if not hasattr(cls._thread_local, "session_factory"):
cls._thread_local.session_factory = async_sessionmaker(
bind=cls.get_engine(),
autoflush=False,
autocommit=False,
expire_on_commit=False,
)
return cls._thread_local.session_factory
@classmethod
@asynccontextmanager
async def async_session(cls):
session_factory = cls._get_session_factory()
async with session_factory() as session:
try:
yield session
await session.commit()
except Exception:
await session.rollback()
raise
finally:
await session.close()
|