Coverage for src/integry/resources/base.py: 75%

1from typing import AsyncIterator, Awaitable, Literal, Type, TypeVar 

2import httpx 

3from pydantic import BaseModel 

4 

5from integry.exceptions import NotFound 

6from integry.utils.common import get_hash 

7 

8 

9class BaseResource: 

10 http_client: httpx.AsyncClient 

11 app_key: str 

12 app_secret: str 

13 name: Literal["functions"] | Literal["apps"] 

14 

15 def __init__( 

16 self, 

17 client: httpx.AsyncClient, 

18 app_key: str, 

19 app_secret: str, 

20 resource: Literal["functions"] | Literal["apps"], 

21 ): 

22 self.http_client = client 

23 self.app_key = app_key 

24 self.app_secret = app_secret 

25 self.name = resource 

26 

27 def _get_response_data_or_raise(self, response: httpx.Response): 

28 if response.status_code == 404: 28 ↛ 29line 28 didn't jump to line 29 because the condition on line 28 was never true

29 data = response.json() 

30 raise NotFound(data.get("detail")) 

31 

32 if response.status_code != 200: 32 ↛ 33line 32 didn't jump to line 33 because the condition on line 32 was never true

33 raise Exception(response.content) 

34 

35 return response.json() 

36 

37 def _get_signed_request_headers(self, user_id: str): 

38 hash = get_hash(self.app_secret, user_id) 

39 return {"App-Key": self.app_key, "User-ID": user_id, "Hash": hash} 

40 

41 

42ResourceModel = TypeVar("ResourceModel", bound=BaseModel) 

43ResourcePage = TypeVar("ResourcePage", bound=BaseModel) 

44 

45 

46class AsyncPaginator(AsyncIterator[ResourceModel], Awaitable[ResourcePage]): 

47 _cursor_key = "_cursor" 

48 

49 def __init__( 

50 self, 

51 resource: BaseResource, 

52 user_id: str, 

53 query_string: str, 

54 model: Type[ResourceModel], 

55 paginated_response_model: Type[ResourcePage], 

56 explicit_cursor: str = "", 

57 ): 

58 self._resource = resource 

59 self._user_id = user_id 

60 

61 self._query_string = query_string 

62 self._data = [] 

63 self._cursor = "" 

64 self._explicit_cursor = explicit_cursor 

65 

66 self._model = model 

67 self._paginated_response_model = paginated_response_model 

68 

69 def __aiter__(self): 

70 return self 

71 

72 async def __anext__(self): 

73 if not self._data and self._cursor is not None: 

74 await self._populate_next_page() 

75 

76 if not self._data: 

77 raise StopAsyncIteration 

78 

79 return self._data.pop(0) 

80 

81 def __await__(self): 

82 return self._get_data().__await__() 

83 

84 async def _populate_next_page(self): 

85 self._data, self._cursor = await self._get_next_page(self._cursor) 

86 

87 async def _get_data(self) -> ResourcePage: 

88 data, cursor = await self._get_next_page(self._explicit_cursor) 

89 

90 if cursor is None: 90 ↛ 91line 90 didn't jump to line 91 because the condition on line 90 was never true

91 cursor = "" 

92 

93 return self._paginated_response_model( 

94 cursor=cursor, **{self._resource.name: data} 

95 ) 

96 

97 async def _get_next_page(self, cursor) -> tuple[list[ResourceModel], str | None]: 

98 response = await self._resource.http_client.post( 

99 f"{self._resource.name}/list/{self._query_string}", 

100 headers=self._resource._get_signed_request_headers(self._user_id), 

101 # json={self._cursor_key: cursor}, 

102 ) 

103 response_data = self._resource._get_response_data_or_raise(response) 

104 

105 data = [ 

106 self._model(**record, _resource=self._resource) 

107 for record in response_data.get(self._resource.name, []) 

108 ] 

109 next_cursor = response_data.get(self._cursor_key) 

110 return data, next_cursor