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, *, always_commit: bool = False): 

10 """Initialize the unit of work. 

11 

12 When always_commit is True, the unit of work will always be committed, even 

13 when an exception occurred. 

14 """ 

15 self._database = database 

16 self._always_commit = always_commit 

17 

18 async def __aenter__(self): 

19 """Enter the unit of work.""" 

20 await self._database.begin() 

21 return self 

22 

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

24 """Exit the unit of work. 

25 

26 When an exception occurred and always_commit is False, the transaction will 

27 be rollbacked. 

28 """ 

29 if self._always_commit: 

30 await self._database.commit() 

31 return 

32 

33 if exc_type: 

34 await self._database.rollback() 

35 else: 

36 await self._database.commit()