Coverage for polar/checkout/ip_geolocation.py: 44%
30 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-12-05 16:17 +0000
1import argparse 1a
2import sys 1a
3from typing import Annotated, cast 1a
5import ipinfo_db 1a
6import ipinfo_db.reader 1a
7from fastapi import Depends, Request 1a
9from polar.config import settings 1a
11DATABASE_PATH = ( 1a
12 settings.IP_GEOLOCATION_DATABASE_DIRECTORY_PATH
13 / settings.IP_GEOLOCATION_DATABASE_NAME
14)
17async def _get_client_dependency(request: Request) -> "IPGeolocationClient | None": 1a
18 """
19 Retrieve the IPInfo database client from the FastAPI request state.
20 """
21 return request.state.ip_geolocation_client
24IPGeolocationClient = Annotated[ipinfo_db.Client, Depends(_get_client_dependency)] 1a
27def _download_database(access_token: str) -> None: 1a
28 """
29 Download the IP to Country ASN database.
31 This should not be called when starting the server or during a request but
32 at build time.
34 Args:
35 access_token: IPInfo access token.
36 """
37 client = ipinfo_db.Client(access_token, DATABASE_PATH, replace=True)
38 client.close()
41def get_client() -> IPGeolocationClient: 1a
42 """
43 Open the IP to Country ASN database.
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)
56def get_ip_country(client: IPGeolocationClient, ip: str) -> str | None: 1a
57 """
58 Get the country alpha-2 code for the given IP address, if available.
60 Args:
61 client: IPInfo database client.
62 ip: IP address.
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
74if __name__ == "__main__": 74 ↛ 75line 74 didn't jump to line 75 because the condition on line 74 was never true1a
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()
81 _download_database(args.access_token)
82 sys.stdout.write(f"Database downloaded to {DATABASE_PATH}\n")
84__all__ = ["get_client", "get_ip_country", "IPGeolocationClient"] 1a