from app.dao.chat import ChatDAO
from app.database import async_session_maker
from app.dao.user import UserDAO


class UnitOfWork:
	def __init__(self):
		self.session_factory = async_session_maker

	async def __aenter__(self):
		self.session = self.session_factory()

		self.user = UserDAO(self.session)
		self.chat = ChatDAO(self.session)

	async def __aexit__(self, *args):
		await self.rollback()
		await self.session.close()

	async def commit(self):
		await self.session.commit()

	async def rollback(self):
		await self.session.rollback()