Coverage for src/kwai/modules/club/repositories/country_db_repository.py: 100%

1"""Module that implements a country repository with a database.""" 

2 

3from sql_smith.query import SelectQuery 

4 

5from kwai.core.db.database import Database 

6from kwai.core.domain.entity import Entity 

7from kwai.modules.club.domain.country import CountryEntity, CountryIdentifier 

8from kwai.modules.club.repositories._tables import CountryRow 

9from kwai.modules.club.repositories.country_repository import ( 

10 CountryNotFoundException, 

11 CountryRepository, 

12) 

13 

14 

15class CountryDbRepository(CountryRepository): 

16 """A repository for countries in a database.""" 

17 

18 def __init__(self, database: Database): 

19 self._database = database 

20 

21 async def get_by_iso_2(self, iso_2: str) -> CountryEntity: 

22 query: SelectQuery = Database.create_query_factory().select() 

23 query.from_(CountryRow.__table_name__).columns(*CountryRow.get_aliases()).where( 

24 CountryRow.field("iso_2").eq(iso_2) 

25 ) 

26 row = await self._database.fetch_one(query) 

27 if row: 

28 return CountryRow.map(row).create_country() 

29 

30 raise CountryNotFoundException(f"Country with iso 2 {iso_2} does not exist.") 

31 

32 async def create(self, country: CountryEntity) -> CountryEntity: 

33 new_id = await self._database.insert( 

34 CountryRow.__table_name__, CountryRow.persist(country) 

35 ) 

36 return Entity.replace(country, id_=CountryIdentifier(new_id)) 

37 

38 async def delete(self, country: CountryEntity): 

39 await self._database.delete(country.id.value, CountryRow.__table_name__)