Coverage for bc/kwai-bc-club/src/kwai_bc_club/repositories/country_db_repository.py: 100%

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

2 

3from kwai_core.db.database import Database 

4from sql_smith.query import SelectQuery 

5 

6from kwai_bc_club.domain.country import CountryEntity, CountryIdentifier 

7from kwai_bc_club.repositories._tables import CountryRow 

8from kwai_bc_club.repositories.country_repository import ( 

9 CountryNotFoundException, 

10 CountryRepository, 

11) 

12 

13 

14class CountryDbRepository(CountryRepository): 

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

16 

17 def __init__(self, database: Database): 

18 self._database = database 

19 

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

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

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

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

24 ) 

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

26 if row: 

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

28 

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

30 

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

32 new_id = await self._database.insert( 

33 CountryRow.__table_name__, CountryRow.persist(country) 

34 ) 

35 return country.set_id(CountryIdentifier(new_id)) 

36 

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

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