Coverage for src/chainalysis/sql/transactional.py: 100%

55 statements  

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

1import pandas as pd 

2 

3from chainalysis.constants import BASE_URL, TRANSACTIONAL_ENDPOINTS 

4from chainalysis.util_functions.exceptions import ( 

5 DataSolutionsSDKException, 

6 ForbiddenException, 

7 NotFoundException, 

8 UnauthorizedException, 

9 UnhandledException, 

10) 

11from chainalysis.util_functions.requests import issue_request 

12 

13 

14class Transactional: 

15 """ 

16 This class provides methods to handle transactional queries. 

17 It supports fetching results as JSON or pandas DataFrame, 

18 and provides query execution statistics. 

19 """ 

20 

21 def __init__(self, api_key: str): 

22 """ 

23 Initialize the Transactional object with an API key. 

24 

25 :param api_key: API key for authenticating requests. 

26 :type api_key: str 

27 """ 

28 self.api_key = api_key 

29 self._status_code = 0 

30 self.results = {} 

31 self._stats = {} 

32 self.json_response = {} 

33 self.status = "error" 

34 self.dataframe_data = None 

35 

36 def __call__( 

37 self, query: str, parameters: dict[str, str] = {}, options: dict[str, str] = {} 

38 ) -> "Transactional": 

39 """ 

40 Execute a SQL query using the provided parameters and options. 

41 

42 :param query: SQL query template with placeholders for parameters. 

43 :type query: str 

44 :param parameters: Parameters to format the SQL query. 

45 :type parameters: dict[str, str] 

46 :param options: Additional options for query execution. 

47 :type options: dict[str, str] 

48 :return: Returns self for chaining or method chaining. 

49 :rtype: Transactional 

50 """ 

51 

52 query_execution_url = ( 

53 f"{BASE_URL['base_url']}/{TRANSACTIONAL_ENDPOINTS['query_execution']}" 

54 ) 

55 

56 body = { 

57 "sql": query, 

58 "parameters": parameters, 

59 "options": options, 

60 } 

61 

62 try: 

63 self.json_response = issue_request( 

64 api_key=self.api_key, 

65 url=query_execution_url, 

66 body=body, 

67 method="POST", 

68 ) 

69 self.status = self.json_response["status"] 

70 if self.status == "error": 

71 self.error_message = self.json_response["message"] 

72 self.error_details = self.json_response.get("details") 

73 self.exception = DataSolutionsSDKException(self.error_message) 

74 else: 

75 self._stats = self.json_response["stats"] 

76 self.results = self.json_response["results"] 

77 self._status_code = 200 

78 except (UnauthorizedException, ForbiddenException, NotFoundException) as e: 

79 raise e 

80 except DataSolutionsSDKException as e: 

81 self.exception = e.get_exception() 

82 self._status_code = e.status_code 

83 except Exception as e: 

84 self.exception = UnhandledException( 

85 details=e, 

86 ) 

87 return self 

88 

89 def json(self) -> dict: 

90 """ 

91 Return the JSON data of the results. 

92 

93 :raises Exception: Raises an exception if the query resulted in an error. 

94 :return: JSON results of the SQL query. 

95 :rtype: dict 

96 """ 

97 

98 if self.status != "error": 

99 return self.results 

100 else: 

101 raise self.exception 

102 

103 def df(self) -> pd.DataFrame: 

104 """ 

105 Convert query results into a pandas DataFrame. 

106 

107 :raises Exception: Raises an exception if the query resulted in an error. 

108 :return: DataFrame containing the results of the SQL query. 

109 :rtype: pd.DataFrame 

110 """ 

111 if self.status != "error": 

112 if self.dataframe_data is None: 

113 self.dataframe_data = pd.DataFrame(self.results) 

114 return self.dataframe_data 

115 elif self.exception: 

116 raise self.exception 

117 

118 def stats(self) -> dict: 

119 """ 

120 Get the statistics of the executed query. 

121 

122 :return: Statistics of the query execution. 

123 :rtype: dict 

124 """ 

125 if self.status != "error": 

126 return self._stats 

127 else: 

128 raise self.exception 

129 

130 def was_successful(self) -> bool: 

131 """ 

132 Determine if the query executed successfully. 

133 

134 :return: True if the query was successful, False otherwise. 

135 :rtype: bool 

136 """ 

137 if self.status != "error": 

138 return True 

139 return False 

140 

141 def status_code(self) -> int: 

142 """ 

143 Return the status code of the query. 

144 

145 :return: The status code of the query. 

146 :rtype: int 

147 """ 

148 return self._status_code