Coverage for mindsdb / integrations / handlers / serpstack_handler / serpstack_handler.py: 0%

1import requests 

2 

3from mindsdb.utilities import log 

4from mindsdb_sql_parser import parse_sql 

5 

6from mindsdb.integrations.libs.api_handler import APIHandler 

7 

8from mindsdb.integrations.libs.response import ( 

9 HandlerStatusResponse as StatusResponse, 

10 HandlerResponse as Response 

11) 

12 

13from .serpstack_tables import ( 

14 OrganicResultsTable, ImageResultsTable, 

15 VideoResultsTable, NewsResultsTable, ShoppingResultsTable 

16) 

17 

18 

19logger = log.getLogger(__name__) 

20 

21 

22class SerpstackHandler(APIHandler): 

23 """A class for handling connections and interactions with the Serpstack API. 

24 Attributes: 

25 api_key (str): API access key for the Serpstack API. 

26 is_connected (bool): Whether or not the API client is connected to Serpstack. 

27 """ 

28 

29 name = 'serpstack' 

30 

31 def __init__(self, name: str = None, **kwargs): 

32 super().__init__(name) 

33 

34 connection_data = kwargs.get('connection_data', {}) 

35 self.connection_data = connection_data 

36 

37 self.access_key = None 

38 self.base_url = None 

39 self.is_connected = False 

40 

41 if 'access_key' in self.connection_data: 

42 self.access_key = self.connection_data['access_key'] 

43 

44 # register tables 

45 organic_results = OrganicResultsTable(self) 

46 self._register_table('organic_results', organic_results) 

47 

48 image_results = ImageResultsTable(self) 

49 self._register_table('image_results', image_results) 

50 

51 video_results = VideoResultsTable(self) 

52 self._register_table('video_results', video_results) 

53 

54 news_results = NewsResultsTable(self) 

55 self._register_table('news_results', news_results) 

56 

57 shopping_results = ShoppingResultsTable(self) 

58 self._register_table('shopping_results', shopping_results) 

59 

60 def connect(self): 

61 """ Sets up connection and returns the base URL to be used""" 

62 

63 if self.is_connected: 

64 return self.base_url 

65 

66 if not self.access_key: 

67 logger.error("No access key provided for Serpstack API") 

68 return None 

69 

70 try: 

71 url = f"https://api.serpstack.com/search?access_key={self.access_key}" 

72 api_request = requests.get(url) 

73 api_response = api_request.json() 

74 # error 105 means that user is on a free plan 

75 if api_response['error']['code'] == 105: 

76 self.base_url = "http://api.serpstack.com/search" 

77 # error 310 means that missing search query, which means that user can use https 

78 elif api_response['error']['code'] == 310: 

79 self.base_url = "https://api.serpstack.com/search" 

80 # any other error suggests issues with the account 

81 else: 

82 logger.error(f"Failed to connect to Serpstack API: {api_response['error']['info']}") 

83 return None 

84 

85 self.is_connected = True 

86 return self.base_url 

87 

88 except Exception as e: 

89 logger.error(f"Failed to connect to Serpstack API: {str(e)}") 

90 return None 

91 

92 def check_connection(self) -> StatusResponse: 

93 """ Checks connection to Serpstack API""" 

94 response = StatusResponse(False) 

95 

96 try: 

97 self.connect() 

98 response.success = True 

99 response.copy_storage = True 

100 except Exception as e: 

101 response.error_message = ( 

102 f"Failed to connect to Serpstack API: {str(e)}" 

103 ) 

104 logger.error(response.error_message) 

105 response.success = False 

106 

107 self.is_connected = response.success 

108 

109 return response 

110 

111 def native_query(self, query: str = None) -> Response: 

112 """Receive and process a raw query. 

113 Parameters 

114 ---------- 

115 query : str 

116 query in a native format 

117 Returns 

118 ------- 

119 StatusResponse 

120 Request status 

121 """ 

122 ast = parse_sql(query) 

123 return self.query(ast)