Coverage for src/kwai/core/db/uow.py: 100%

1"""Module that implements a unit of work pattern.""" 

2 

3from kwai.core.db.database import Database 

4 

5 

6class UnitOfWork: 

7 """A unit of work implementation.""" 

8 

9 def __init__(self, database: Database): 

10 self._database = database 

11 

12 async def __aenter__(self): 

13 """Enter the unit of work.""" 

14 await self._database.begin() 

15 return self 

16 

17 async def __aexit__(self, exc_type, exc_val, exc_tb): 

18 """Exit the unit of work. 

19 

20 When an exception occurred, the transaction will be rollbacked. 

21 """ 

22 if exc_type: 

23 await self._database.rollback() 

24 else: 

25 await self._database.commit()