Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions pymfdata/rdb/connection.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
from asyncio import current_task
from contextlib import asynccontextmanager, contextmanager
from typing import AsyncIterable, Callable, Union, Optional
from typing import Callable, Union, Optional

from sqlalchemy.engine import Engine, create_engine
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_scoped_session, create_async_engine
Expand Down Expand Up @@ -34,26 +34,26 @@ def init_session_factory(self, autocommit: bool = False, autoflush: bool = False
async def session(self) -> Callable[..., AsyncSession]:
assert self._session_factory is not None

session: AsyncSession = self._session_factory()
session: Union[AsyncSession, async_scoped_session] = self._session_factory()
try:
yield session
except Exception:
await session.rollback()
raise
finally:
await session.close()
await session.remove()

async def get_db_session(self) -> Callable[..., AsyncSession]:
assert self._session_factory is not None

session: AsyncSession = self._session_factory()
session: Union[AsyncSession, async_scoped_session] = self._session_factory()
try:
yield session
except Exception:
await session.rollback()
raise
finally:
await session.close()
await session.remove()

@property
def engine(self):
Expand Down