Coverage for polar/checkout/ip_geolocation.py: 47%

1import argparse 

2import sys 

3from typing import Annotated, cast 

4 

5import ipinfo_db 

6import ipinfo_db.reader 

7from fastapi import Depends, Request 

8 

9from polar.config import settings 

10 

11DATABASE_PATH = ( 

12 settings.IP_GEOLOCATION_DATABASE_DIRECTORY_PATH 

13 / settings.IP_GEOLOCATION_DATABASE_NAME 

14) 

15 

16 

17async def _get_client_dependency(request: Request) -> "IPGeolocationClient | None": 

18 """ 

19 Retrieve the IPInfo database client from the FastAPI request state. 

20 """ 

21 return request.state.ip_geolocation_client 

22 

23 

24IPGeolocationClient = Annotated[ipinfo_db.Client, Depends(_get_client_dependency)] 

25 

26 

27def _download_database(access_token: str) -> None: 

28 """ 

29 Download the IP to Country ASN database. 

30 

31 This should not be called when starting the server or during a request but 

32 at build time. 

33 

34 Args: 

35 access_token: IPInfo access token. 

36 """ 

37 client = ipinfo_db.Client(access_token, DATABASE_PATH, replace=True) 

38 client.close() 

39 

40 

41def get_client() -> IPGeolocationClient: 

42 """ 

43 Open the IP to Country ASN database. 

44 

45 Returns: 

46 IPInfo database client. 

47 """ 

48 if not DATABASE_PATH.exists(): 

49 raise FileNotFoundError( 

50 f"Database not found at {DATABASE_PATH}. " 

51 "Please run `python -m polar.checkout.ip_geolocation ACCESS_TOKEN`." 

52 ) 

53 return ipinfo_db.Client(path=DATABASE_PATH) 

54 

55 

56def get_ip_country(client: IPGeolocationClient, ip: str) -> str | None: 

57 """ 

58 Get the country alpha-2 code for the given IP address, if available. 

59 

60 Args: 

61 client: IPInfo database client. 

62 ip: IP address. 

63 

64 Returns: 

65 Country alpha-2 code. 

66 """ 

67 ret = cast(str | None, client.getCountry(ip)) 

68 # Convert empty str into None response to ease empty case checking 

69 if ret == "": 

70 return None 

71 return ret 

72 

73 

74if __name__ == "__main__": 74 ↛ 75line 74 didn't jump to line 75 because the condition on line 74 was never true

75 parser = argparse.ArgumentParser( 

76 description="Download the IP to Country ASN database." 

77 ) 

78 parser.add_argument("access_token", type=str, help="IPInfo access token") 

79 args = parser.parse_args() 

80 

81 _download_database(args.access_token) 

82 sys.stdout.write(f"Database downloaded to {DATABASE_PATH}\n") 

83 

84__all__ = ["get_client", "get_ip_country", "IPGeolocationClient"]