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

1import pandas as pd 

2from typing import List 

3from mindsdb.integrations.libs.api_handler import APITable 

4from mindsdb.integrations.utilities.handlers.query_utilities import SELECTQueryParser, SELECTQueryExecutor 

5from mindsdb.utilities import log 

6from mindsdb_sql_parser import ast 

7 

8logger = log.getLogger(__name__) 

9 

10 

11class ZipCodeBaseCodeLocationTable(APITable): 

12 """The ZipCodeBase Location Table implementation""" 

13 

14 def select(self, query: ast.Select) -> pd.DataFrame: 

15 """Pulls data from the https://app.zipcodebase.com/documentation#search API 

16 

17 Parameters 

18 ---------- 

19 query : ast.Select 

20 Given SQL SELECT query 

21 

22 Returns 

23 ------- 

24 pd.DataFrame 

25 Location of the codes matching the query 

26 

27 Raises 

28 ------ 

29 ValueError 

30 If the query contains an unsupported condition 

31 """ 

32 

33 select_statement_parser = SELECTQueryParser( 

34 query, 

35 'code_to_location', 

36 self.get_columns() 

37 ) 

38 

39 selected_columns, where_conditions, order_by_conditions, result_limit = select_statement_parser.parse_query() 

40 

41 search_params = {} 

42 subset_where_conditions = [] 

43 

44 for op, arg1, arg2 in where_conditions: 

45 if arg1 == "codes": 

46 if op == '=': 

47 search_params["codes"] = arg2 

48 else: 

49 raise NotImplementedError("Only '=' operator is supported for codes column.") 

50 

51 elif arg1 in self.get_columns(): 

52 subset_where_conditions.append([op, arg1, arg2]) 

53 

54 filter_flag = "codes" in search_params 

55 

56 if not filter_flag: 

57 raise NotImplementedError("`codes` column has to be present in where clause.") 

58 

59 code_to_location_df = pd.DataFrame(columns=self.get_columns()) 

60 

61 response = self.handler.client.code_to_location(search_params.get("codes")) 

62 

63 self.check_res(res=response) 

64 

65 content = response["content"] 

66 

67 code_to_location_df = pd.json_normalize(self.clean_resp(content["results"])) 

68 

69 select_statement_executor = SELECTQueryExecutor( 

70 code_to_location_df, 

71 selected_columns, 

72 subset_where_conditions, 

73 order_by_conditions, 

74 result_limit 

75 ) 

76 

77 code_to_location_df = select_statement_executor.execute_query() 

78 

79 return code_to_location_df 

80 

81 def clean_resp(self, data): 

82 clean_data = [] 

83 for k, v in data.items(): 

84 clean_data.extend(v) 

85 return clean_data 

86 

87 def check_res(self, res): 

88 if res["code"] != 200: 

89 raise Exception("Error fetching results - " + res["error"]) 

90 

91 def get_columns(self) -> List[str]: 

92 """Gets all columns to be returned in pandas DataFrame responses 

93 Returns 

94 ------- 

95 List[str] 

96 List of columns 

97 """ 

98 

99 return [ 

100 "postal_code", 

101 "country_code", 

102 "latitude", 

103 "longitude", 

104 "city", 

105 "state", 

106 "city_en", 

107 "state_en", 

108 "state_code", 

109 "province", 

110 "province_code" 

111 ] 

112 

113 

114class ZipCodeBaseCodeInRadiusTable(APITable): 

115 """The ZipCodeBase Codes within Radius Table implementation""" 

116 

117 def select(self, query: ast.Select) -> pd.DataFrame: 

118 """Pulls data from the https://app.zipcodebase.com/documentation#radius API 

119 

120 Parameters 

121 ---------- 

122 query : ast.Select 

123 Given SQL SELECT query 

124 

125 Returns 

126 ------- 

127 pd.DataFrame 

128 codes within the radius 

129 

130 Raises 

131 ------ 

132 ValueError 

133 If the query contains an unsupported condition 

134 """ 

135 

136 select_statement_parser = SELECTQueryParser( 

137 query, 

138 'codes_within_radius', 

139 self.get_columns() 

140 ) 

141 

142 selected_columns, where_conditions, order_by_conditions, result_limit = select_statement_parser.parse_query() 

143 

144 search_params = {} 

145 subset_where_conditions = [] 

146 

147 for op, arg1, arg2 in where_conditions: 

148 if arg1 == "code": 

149 if op == '=': 

150 search_params["code"] = arg2 

151 else: 

152 raise NotImplementedError("Only '=' operator is supported for code column.") 

153 

154 if arg1 == "radius": 

155 if op == '=': 

156 search_params["radius"] = arg2 

157 else: 

158 raise NotImplementedError("Only '=' operator is supported for radius column.") 

159 

160 if arg1 == "country": 

161 if op == '=': 

162 search_params["country"] = arg2 

163 else: 

164 raise NotImplementedError("Only '=' operator is supported for country column.") 

165 

166 if arg1 == "unit": 

167 if op == '=': 

168 search_params["unit"] = arg2 

169 else: 

170 raise NotImplementedError("Only '=' operator is supported for unit column.") 

171 

172 elif arg1 in self.get_columns(): 

173 subset_where_conditions.append([op, arg1, arg2]) 

174 

175 filter_flag = ("code" in search_params) and ("radius" in search_params) and ("country" in search_params) 

176 

177 if not filter_flag: 

178 raise NotImplementedError("`codes`, `radius` and `country` columns have to be present in where clause.") 

179 

180 code_to_location_df = pd.DataFrame(columns=self.get_columns()) 

181 

182 response = self.handler.client.codes_within_radius(search_params.get("code"), search_params.get("radius"), search_params.get("country"), search_params.get("unit", "km")) 

183 

184 self.check_res(res=response) 

185 

186 content = response["content"] 

187 

188 logger.info(f"response size - {len(content['results'])}") 

189 code_to_location_df = pd.json_normalize(content["results"]) 

190 

191 select_statement_executor = SELECTQueryExecutor( 

192 code_to_location_df, 

193 selected_columns, 

194 subset_where_conditions, 

195 order_by_conditions, 

196 result_limit 

197 ) 

198 

199 code_to_location_df = select_statement_executor.execute_query() 

200 

201 return code_to_location_df 

202 

203 def check_res(self, res): 

204 if res["code"] != 200: 

205 raise Exception("Error fetching results - " + res["error"]) 

206 

207 def get_columns(self) -> List[str]: 

208 """Gets all columns to be returned in pandas DataFrame responses 

209 Returns 

210 ------- 

211 List[str] 

212 List of columns 

213 """ 

214 

215 return [ 

216 "code", 

217 "city", 

218 "state", 

219 "city_en", 

220 "state_en", 

221 "distance" 

222 ] 

223 

224 

225class ZipCodeBaseCodeByCityTable(APITable): 

226 """The ZipCodeBase Codes within a City Table implementation""" 

227 

228 def select(self, query: ast.Select) -> pd.DataFrame: 

229 """Pulls data from the https://app.zipcodebase.com/documentation#city API 

230 

231 Parameters 

232 ---------- 

233 query : ast.Select 

234 Given SQL SELECT query 

235 

236 Returns 

237 ------- 

238 pd.DataFrame 

239 codes within the city 

240 

241 Raises 

242 ------ 

243 ValueError 

244 If the query contains an unsupported condition 

245 """ 

246 

247 select_statement_parser = SELECTQueryParser( 

248 query, 

249 'codes_by_city', 

250 self.get_columns() 

251 ) 

252 

253 selected_columns, where_conditions, order_by_conditions, result_limit = select_statement_parser.parse_query() 

254 

255 search_params = {} 

256 subset_where_conditions = [] 

257 

258 for op, arg1, arg2 in where_conditions: 

259 if arg1 == "city": 

260 if op == '=': 

261 search_params["city"] = arg2 

262 else: 

263 raise NotImplementedError("Only '=' operator is supported for city column.") 

264 

265 if arg1 == "country": 

266 if op == '=': 

267 search_params["country"] = arg2 

268 else: 

269 raise NotImplementedError("Only '=' operator is supported for country column.") 

270 

271 elif arg1 in self.get_columns(): 

272 subset_where_conditions.append([op, arg1, arg2]) 

273 

274 filter_flag = ("city" in search_params) and ("country" in search_params) 

275 

276 if not filter_flag: 

277 raise NotImplementedError("`city` and `country` columns have to be present in where clause.") 

278 

279 codes_by_city_df = pd.DataFrame(columns=self.get_columns()) 

280 

281 response = self.handler.client.codes_by_city(search_params.get("city"), search_params.get("country")) 

282 

283 self.check_res(res=response) 

284 

285 content = response["content"] 

286 

287 logger.info(f"response size - {len(content['results'])}") 

288 codes_by_city_df = pd.json_normalize({"codes": content["results"]}) 

289 

290 select_statement_executor = SELECTQueryExecutor( 

291 codes_by_city_df, 

292 selected_columns, 

293 subset_where_conditions, 

294 order_by_conditions, 

295 result_limit 

296 ) 

297 

298 codes_by_city_df = select_statement_executor.execute_query() 

299 

300 return codes_by_city_df 

301 

302 def check_res(self, res): 

303 if res["code"] != 200: 

304 raise Exception("Error fetching results - " + res["error"]) 

305 

306 def get_columns(self) -> List[str]: 

307 """Gets all columns to be returned in pandas DataFrame responses 

308 Returns 

309 ------- 

310 List[str] 

311 List of columns 

312 """ 

313 

314 return [ 

315 "codes" 

316 ] 

317 

318 

319class ZipCodeBaseCodeByStateTable(APITable): 

320 """The ZipCodeBase Codes within a State Table implementation""" 

321 

322 def select(self, query: ast.Select) -> pd.DataFrame: 

323 """Pulls data from the https://app.zipcodebase.com/documentation#state API 

324 

325 Parameters 

326 ---------- 

327 query : ast.Select 

328 Given SQL SELECT query 

329 

330 Returns 

331 ------- 

332 pd.DataFrame 

333 codes within the state 

334 

335 Raises 

336 ------ 

337 ValueError 

338 If the query contains an unsupported condition 

339 """ 

340 

341 select_statement_parser = SELECTQueryParser( 

342 query, 

343 'codes_by_state', 

344 self.get_columns() 

345 ) 

346 

347 selected_columns, where_conditions, order_by_conditions, result_limit = select_statement_parser.parse_query() 

348 

349 search_params = {} 

350 subset_where_conditions = [] 

351 

352 for op, arg1, arg2 in where_conditions: 

353 if arg1 == "state": 

354 if op == '=': 

355 search_params["state"] = arg2 

356 else: 

357 raise NotImplementedError("Only '=' operator is supported for state column.") 

358 

359 if arg1 == "country": 

360 if op == '=': 

361 search_params["country"] = arg2 

362 else: 

363 raise NotImplementedError("Only '=' operator is supported for country column.") 

364 

365 elif arg1 in self.get_columns(): 

366 subset_where_conditions.append([op, arg1, arg2]) 

367 

368 filter_flag = ("state" in search_params) and ("country" in search_params) 

369 

370 if not filter_flag: 

371 raise NotImplementedError("`state` and `country` columns have to be present in where clause.") 

372 

373 codes_by_state_df = pd.DataFrame(columns=self.get_columns()) 

374 

375 response = self.handler.client.codes_by_state(search_params.get("state"), search_params.get("country")) 

376 

377 self.check_res(res=response) 

378 

379 content = response["content"] 

380 

381 logger.info(f"response size - {len(content['results'])}") 

382 codes_by_state_df = pd.json_normalize({"codes": content["results"]}) 

383 

384 select_statement_executor = SELECTQueryExecutor( 

385 codes_by_state_df, 

386 selected_columns, 

387 subset_where_conditions, 

388 order_by_conditions, 

389 result_limit 

390 ) 

391 

392 codes_by_state_df = select_statement_executor.execute_query() 

393 

394 return codes_by_state_df 

395 

396 def check_res(self, res): 

397 if res["code"] != 200: 

398 raise Exception("Error fetching results - " + res["error"]) 

399 

400 def get_columns(self) -> List[str]: 

401 """Gets all columns to be returned in pandas DataFrame responses 

402 Returns 

403 ------- 

404 List[str] 

405 List of columns 

406 """ 

407 

408 return [ 

409 "codes" 

410 ] 

411 

412 

413class ZipCodeBaseStatesByCountryTable(APITable): 

414 """The ZipCodeBase Provinces/states within a country Table implementation""" 

415 

416 def select(self, query: ast.Select) -> pd.DataFrame: 

417 """Pulls data from the https://app.zipcodebase.com/documentation#provinces API 

418 

419 Parameters 

420 ---------- 

421 query : ast.Select 

422 Given SQL SELECT query 

423 

424 Returns 

425 ------- 

426 pd.DataFrame 

427 states within a country 

428 

429 Raises 

430 ------ 

431 ValueError 

432 If the query contains an unsupported condition 

433 """ 

434 

435 select_statement_parser = SELECTQueryParser( 

436 query, 

437 'states_by_country', 

438 self.get_columns() 

439 ) 

440 

441 selected_columns, where_conditions, order_by_conditions, result_limit = select_statement_parser.parse_query() 

442 

443 search_params = {} 

444 subset_where_conditions = [] 

445 

446 for op, arg1, arg2 in where_conditions: 

447 

448 if arg1 == "country": 

449 if op == '=': 

450 search_params["country"] = arg2 

451 else: 

452 raise NotImplementedError("Only '=' operator is supported for country column.") 

453 

454 elif arg1 in self.get_columns(): 

455 subset_where_conditions.append([op, arg1, arg2]) 

456 

457 filter_flag = ("country" in search_params) 

458 

459 if not filter_flag: 

460 raise NotImplementedError("`country` column has to be present in where clause.") 

461 

462 states_by_country_df = pd.DataFrame(columns=self.get_columns()) 

463 

464 response = self.handler.client.states_by_country(search_params.get("country")) 

465 

466 self.check_res(res=response) 

467 

468 content = response["content"] 

469 

470 logger.info(f"response size - {len(content['results'])}") 

471 states_by_country_df = pd.json_normalize({"states": content["results"]}) 

472 

473 select_statement_executor = SELECTQueryExecutor( 

474 states_by_country_df, 

475 selected_columns, 

476 subset_where_conditions, 

477 order_by_conditions, 

478 result_limit 

479 ) 

480 

481 states_by_country_df = select_statement_executor.execute_query() 

482 

483 return states_by_country_df 

484 

485 def check_res(self, res): 

486 if res["code"] != 200: 

487 raise Exception("Error fetching results - " + res["error"]) 

488 

489 def get_columns(self) -> List[str]: 

490 """Gets all columns to be returned in pandas DataFrame responses 

491 Returns 

492 ------- 

493 List[str] 

494 List of columns 

495 """ 

496 

497 return [ 

498 "states" 

499 ]