Coverage for graphqler / utils / request_utils.py: 74%

1from urllib3.exceptions import InsecureRequestWarning 

2from urllib3 import disable_warnings 

3from typing import Callable 

4from graphqler import config 

5 

6import time 

7import requests 

8import json 

9 

10 

11# The last time a request was made so that we can wait between requests 

12last_request_time = time.time() 

13session = None 

14 

15 

16def get_headers() -> dict: 

17 """Get the headers for the request. 

18 Authorization will be used from the AUTHORIZATION variable first, then from the CUSTOM_HEADERS variable. 

19 

20 Returns: 

21 dict: The headers for the request 

22 """ 

23 headers = {"Content-Type": "application/json"} 

24 if config.CUSTOM_HEADERS: 

25 headers.update(config.CUSTOM_HEADERS) 

26 

27 if config.AUTHORIZATION: 

28 headers["Authorization"] = f"{config.AUTHORIZATION}" 

29 

30 return headers 

31 

32 

33def get_proxies() -> dict: 

34 """Get the proxies for the request 

35 

36 Returns: 

37 dict: The proxies for the request 

38 """ 

39 if config.PROXY and "http:" in config.PROXY: 

40 return {"http": config.PROXY} 

41 elif config.PROXY and "https:" in config.PROXY: 

42 return {"https": config.PROXY} 

43 elif config.PROXY: 

44 return {"http": config.PROXY, "https": config.PROXY} 

45 else: 

46 return {} 

47 

48 

49def send_graphql_request(url: str, payload: str | dict | list, next: Callable[[dict], dict] | None = None) -> tuple[dict, requests.Response]: 

50 """Send GraphQL request to the specified endpoint 

51 

52 Args: 

53 url (str): URL of the graphql server 

54 payload (str | dict | list): The payload to send to the GraphQL API. If dict or string, must provide the query and variables keys 

55 next (Callable[[dict], dict], optional): Callback function in case there is action to be done after. Defaults to None. 

56 

57 Returns: 

58 tuple[dict, requests.Response]: Dictionary of the graphql response, and the request's response 

59 """ 

60 global last_request_time 

61 

62 # Make the body (if it's a string, add the key, if it's dict or list, assume the creator of the request knows what they are doing 

63 # (ie. added the query / variable keys themselves)) 

64 if isinstance(payload, str): 

65 body = {"query": payload} 

66 else: 

67 body = payload 

68 

69 # If the last request was made recently, wait for a bit 

70 time_since_last_request = time.time() - last_request_time 

71 if time_since_last_request < config.TIME_BETWEEN_REQUESTS: 

72 time.sleep(config.TIME_BETWEEN_REQUESTS - time_since_last_request) 

73 

74 # Make the request and set the last request time 

75 session = get_or_create_session() 

76 response = session.post( 

77 url=url, 

78 json=body, 

79 timeout=config.REQUEST_TIMEOUT, 

80 ) 

81 last_request_time = time.time() 

82 

83 if response.status_code != 200: 

84 return parse_response(response.text), response 

85 

86 # if next: 

87 # return next(json.loads(response.text)) 

88 

89 return parse_response(response.text), response 

90 

91 

92def parse_response(response_text: str) -> dict: 

93 """Parse the response and try to jsonify it 

94 

95 Args: 

96 response_text (str): The response text 

97 

98 Returns: 

99 dict: A dictionary of the response 

100 """ 

101 json_text = "" 

102 try: 

103 json_text = json.loads(response_text) 

104 return json_text 

105 except Exception: 

106 return {"errors": [response_text]} 

107 

108 

109def get_or_create_session() -> requests.Session: 

110 """Gets an existing session or creates a new one 

111 

112 Returns: 

113 requests.Session: The session 

114 """ 

115 global session 

116 

117 if session and isinstance(session, requests.Session): 

118 return session 

119 else: 

120 session = create_new_session() 

121 return session 

122 

123 

124def create_new_session() -> requests.Session: 

125 """Create a new session 

126 

127 Returns: 

128 requests.Session: The session 

129 """ 

130 session = requests.Session() 

131 session.headers.update(get_headers()) 

132 

133 # Set proxy if available 

134 if config.PROXY: 

135 session.proxies.update(get_proxies()) 

136 disable_warnings(InsecureRequestWarning) 

137 session.verify = False 

138 return session