Coverage for src/kwai/modules/portal/repositories/author_db_query.py: 100%

1"""Module that implements an AuthorQuery for a database.""" 

2 

3from dataclasses import dataclass 

4from typing import Self 

5 

6from sql_smith.functions import on 

7 

8from kwai.core.db.database import Database 

9from kwai.core.db.database_query import DatabaseQuery 

10from kwai.core.db.table_row import JoinedTableRow 

11from kwai.core.domain.value_objects.unique_id import UniqueId 

12from kwai.modules.portal.domain.author import AuthorEntity 

13from kwai.modules.portal.repositories._tables import AuthorRow, UserRow 

14from kwai.modules.portal.repositories.author_query import AuthorQuery 

15 

16 

17@dataclass(kw_only=True, frozen=True, slots=True) 

18class AuthorQueryRow(JoinedTableRow): 

19 """A data transfer object for the author query.""" 

20 

21 author: AuthorRow 

22 user: UserRow 

23 

24 def create_entity(self) -> AuthorEntity: 

25 """Create an Author entity from a row.""" 

26 return self.author.create_entity( 

27 uuid=UniqueId.create_from_string(self.user.uuid) 

28 ) 

29 

30 

31class AuthorDbQuery(AuthorQuery, DatabaseQuery): 

32 """An author query using a database.""" 

33 

34 def __init__(self, database: Database): 

35 super().__init__(database) 

36 

37 def init(self): 

38 self._query.from_(AuthorRow.__table_name__).columns( 

39 *(AuthorRow.get_aliases()) + UserRow.get_aliases() 

40 ).join( 

41 UserRow.__table_name__, 

42 on(UserRow.column("id"), AuthorRow.column("user_id")), 

43 ) 

44 

45 @property 

46 def columns(self): 

47 return AuthorQueryRow.get_aliases() 

48 

49 def filter_by_id(self, id_) -> Self: 

50 self._query.and_where(AuthorRow.field("user_id").eq(id_.value)) 

51 return self 

52 

53 def filter_by_uuid(self, uuid: UniqueId): 

54 self._query.and_where(UserRow.field("uuid").eq(uuid)) 

55 return self