cognee/cognitive_architecture/infrastructure/databases/relational/utils/with_rollback.py
2024-03-12 13:42:51 +01:00

18 lines
595 B
Python

import logging
from contextlib import asynccontextmanager
from sqlalchemy.ext.asyncio import async_scoped_session
logger = logging.getLogger(__name__)
@asynccontextmanager
async def with_rollback(session: async_scoped_session):
"""Provide a transactional scope around a series of operations."""
try:
# async with session.begin():
yield session
await session.commit()
await session.remove()
except Exception as exception:
await session.rollback()
logger.error("Session rolled back due to: %s", str(exception))
raise exception