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

1from mindsdb.integrations.handlers.zipcodebase_handler.zipcodebase_tables import ( 

2 ZipCodeBaseCodeLocationTable, 

3 ZipCodeBaseCodeInRadiusTable, 

4 ZipCodeBaseCodeByCityTable, 

5 ZipCodeBaseCodeByStateTable, 

6 ZipCodeBaseStatesByCountryTable 

7) 

8from mindsdb.integrations.handlers.zipcodebase_handler.zipcodebase import ZipCodeBaseClient 

9from mindsdb.integrations.libs.api_handler import APIHandler 

10from mindsdb.integrations.libs.response import ( 

11 HandlerStatusResponse as StatusResponse, 

12) 

13 

14from mindsdb.utilities import log 

15from mindsdb_sql_parser import parse_sql 

16 

17 

18logger = log.getLogger(__name__) 

19 

20 

21class ZipCodeBaseHandler(APIHandler): 

22 """The ZipCodeBase handler implementation""" 

23 

24 def __init__(self, name: str, **kwargs): 

25 """Initialize the ZipCodeBase handler. 

26 Parameters 

27 ---------- 

28 name : str 

29 name of a handler instance 

30 """ 

31 super().__init__(name) 

32 

33 connection_data = kwargs.get("connection_data", {}) 

34 self.connection_data = connection_data 

35 self.kwargs = kwargs 

36 self.client = ZipCodeBaseClient(self.connection_data["api_key"]) 

37 self.is_connected = False 

38 

39 code_to_location_data = ZipCodeBaseCodeLocationTable(self) 

40 self._register_table("code_to_location", code_to_location_data) 

41 

42 codes_within_radius_data = ZipCodeBaseCodeInRadiusTable(self) 

43 self._register_table("codes_within_radius", codes_within_radius_data) 

44 

45 codes_by_city_data = ZipCodeBaseCodeByCityTable(self) 

46 self._register_table("codes_by_city", codes_by_city_data) 

47 

48 codes_by_state_data = ZipCodeBaseCodeByStateTable(self) 

49 self._register_table("codes_by_state", codes_by_state_data) 

50 

51 states_by_country_data = ZipCodeBaseStatesByCountryTable(self) 

52 self._register_table("states_by_country", states_by_country_data) 

53 

54 def connect(self) -> StatusResponse: 

55 """Set up the connection required by the handler. 

56 Returns 

57 ------- 

58 StatusResponse 

59 connection object 

60 """ 

61 resp = StatusResponse(False) 

62 status = self.client.remaining_requests() 

63 if status["code"] != 200: 

64 resp.success = False 

65 resp.error_message = status["error"] 

66 return resp 

67 self.is_connected = True 

68 return resp 

69 

70 def check_connection(self) -> StatusResponse: 

71 """Check connection to the handler. 

72 Returns 

73 ------- 

74 StatusResponse 

75 Status confirmation 

76 """ 

77 response = StatusResponse(False) 

78 

79 try: 

80 status = self.client.remaining_requests() 

81 if status["code"] == 200: 

82 logger.info("Authentication successful") 

83 response.success = True 

84 else: 

85 response.success = False 

86 logger.info("Error connecting to ZipCodeBase. " + status["error"]) 

87 response.error_message = status["error"] 

88 except Exception as e: 

89 logger.error(f"Error connecting to ZipCodeBase: {e}!") 

90 response.error_message = e 

91 

92 self.is_connected = response.success 

93 return response 

94 

95 def native_query(self, query: str) -> StatusResponse: 

96 """Receive and process a raw query. 

97 Parameters 

98 ---------- 

99 query : str 

100 query in a native format 

101 Returns 

102 ------- 

103 StatusResponse 

104 Request status 

105 """ 

106 ast = parse_sql(query) 

107 return self.query(ast)