Coverage for src/chainalysis/util_functions/requests.py: 25%

28 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-08-20 16:53 -0400

1from requests import JSONDecodeError, request 

2from tenacity import retry, retry_if_exception, stop_after_attempt, wait_exponential 

3 

4from chainalysis.util_functions.exceptions import ( 

5 BadRequest, 

6 DataSolutionsAPIException, 

7 ForbiddenException, 

8 UnauthorizedException, 

9) 

10 

11 

12def get_headers(api_key: str) -> dict: 

13 """ 

14 Generate headers for an HTTP request. 

15 

16 This function creates a dictionary of headers required for an HTTP request, 

17 including the API key for authorization and JSON content-type headers which 

18 indicate what sort of information will be sent with the request. 

19 

20 :param api_key: The API key used for authenticating the request. 

21 :return: A dictionary containing the headers for the HTTP request. 

22 """ 

23 

24 return { 

25 "X-API-KEY": api_key, 

26 "Content-Type": "application/json", 

27 "Accept": "application/json", 

28 } 

29 

30 

31def retry_condition(exception): 

32 if isinstance(exception, (BadRequest, UnauthorizedException, ForbiddenException)): 

33 return False 

34 return True 

35 

36 

37@retry( 

38 wait=wait_exponential(multiplier=1, min=4, max=10), 

39 retry=retry_if_exception(retry_condition), 

40 stop=stop_after_attempt(10), 

41) 

42def issue_request( 

43 url: str, 

44 api_key: str, 

45 params: dict = {}, 

46 body: dict = {}, 

47 method: str = "GET", 

48) -> dict: 

49 """ 

50 Helper method to issue a request to the Data Solutions API. 

51 This method will automatically retry the request, and handle 

52 basic error checking. 

53 """ 

54 

55 # Encode params and body into a single variable data 

56 data = {**(params or {}), **(body or {})} 

57 

58 headers = get_headers(api_key) 

59 

60 response = request( 

61 method, 

62 url, 

63 headers=headers, 

64 json=data, 

65 ) 

66 

67 if response.status_code >= 300: 

68 try: # Try to decode the json response of the error 

69 json_output = response.json() 

70 error_message = json_output.get("message") 

71 

72 except JSONDecodeError: 

73 raise DataSolutionsAPIException( 

74 message="Unexpected response from the API - response was not JSON", 

75 ) 

76 if response.status_code == 400: 

77 raise BadRequest( 

78 message=error_message, 

79 ) 

80 raise DataSolutionsAPIException( 

81 message=error_message, 

82 ) 

83 

84 try: # Try to decode the json response of the successful request 

85 response_json = response.json() 

86 except JSONDecodeError: 

87 raise DataSolutionsAPIException( 

88 message="Unexpected response from the API - response was not JSON", 

89 ) 

90 

91 return response_json