30 lines
929 B
Python
30 lines
929 B
Python
"""Keep synchronous compatibility reads off the Gateway's async connection pool."""
|
|
|
|
import asyncio
|
|
from contextvars import ContextVar
|
|
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
from sqlalchemy.pool import NullPool
|
|
|
|
_isolated_factory = ContextVar("isolated_db_factory", default=None)
|
|
|
|
|
|
def isolated_session_factory():
|
|
return _isolated_factory.get()
|
|
|
|
|
|
async def run_isolated_db_read(coro):
|
|
from deerflow.persistence.engine import get_engine
|
|
|
|
source = get_engine()
|
|
if source is None:
|
|
return await coro
|
|
# Never return a connection created on this temporary loop to the main pool.
|
|
engine = create_async_engine(source.url, poolclass=NullPool)
|
|
token = _isolated_factory.set(async_sessionmaker(engine, expire_on_commit=False))
|
|
try:
|
|
return await asyncio.wait_for(coro, 30)
|
|
finally:
|
|
_isolated_factory.reset(token)
|
|
await engine.dispose()
|