diff --git a/posthog/geoip.py b/posthog/geoip.py index 7a749c0b294c2..1038e297ba8fa 100644 --- a/posthog/geoip.py +++ b/posthog/geoip.py @@ -6,7 +6,6 @@ logger = structlog.get_logger(__name__) - try: geoip: Optional[GeoIP2] = GeoIP2(cache=8) # Cache setting corresponds to MODE_MEMORY: Load database into memory. Pure Python. @@ -26,6 +25,9 @@ "time_zone", ] +# GeoIP2 returns the city name as 'city', but we want to map it to 'city_name' +GEOIP_KEY_MAPPING = {"city": "city_name"} + def get_geoip_properties(ip_address: Optional[str]) -> dict[str, str]: """ @@ -52,6 +54,8 @@ def get_geoip_properties(ip_address: Optional[str]) -> dict[str, str]: properties = {} for key, value in geoip_properties.items(): - if value and key in VALID_GEOIP_PROPERTIES: - properties[f"$geoip_{key}"] = value + if value: + mapped_key = GEOIP_KEY_MAPPING.get(key, key) + if mapped_key in VALID_GEOIP_PROPERTIES: + properties[f"$geoip_{mapped_key}"] = value return properties diff --git a/posthog/helpers/tests/test_geoip.py b/posthog/helpers/tests/test_geoip.py index 31fc0e7038dac..64b5d173a069f 100644 --- a/posthog/helpers/tests/test_geoip.py +++ b/posthog/helpers/tests/test_geoip.py @@ -11,20 +11,25 @@ uk_ip = "31.28.64.3" us_ip_v6 = "2600:6c52:7a00:11c:1b6:b7b0:ea19:6365" local_ip = "127.0.0.1" +mexico_ip = "187.188.10.252" +australia_ip_2 = "13.106.122.3" @pytest.mark.parametrize( - "test_input,expected", + "test_input,expected_country,expected_city", [ - (australia_ip, "Australia"), - (uk_ip, "United Kingdom"), - (us_ip_v6, "United States"), + (australia_ip, "Australia", "Sydney"), + (uk_ip, "United Kingdom", "Baldock"), + (us_ip_v6, "United States", "San Luis Obispo"), + (mexico_ip, "Mexico", "Coyoacán"), + (australia_ip, "Australia", "Sydney"), ], ) -def test_geoip_results(test_input, expected): +def test_geoip_results(test_input, expected_country, expected_city): properties = get_geoip_properties(test_input) - assert properties["$geoip_country_name"] == expected - assert len(properties) == 6 + assert properties["$geoip_country_name"] == expected_country + assert properties["$geoip_city_name"] == expected_city + assert len(properties) == 7 class TestGeoIPDBError(TestCase): diff --git a/rust/feature-flags/src/geoip.rs b/rust/feature-flags/src/geoip.rs index 8ffbf0bfa34e4..966b370ccb5e4 100644 --- a/rust/feature-flags/src/geoip.rs +++ b/rust/feature-flags/src/geoip.rs @@ -33,7 +33,7 @@ impl GeoIpClient { /// Checks if the given IP address is valid. fn is_valid_ip(&self, ip: &str) -> bool { - ip != "127.0.0.1" || ip != "::1" + ip != "127.0.0.1" && ip != "::1" } /// Looks up the city data for the given IP address. @@ -56,25 +56,23 @@ impl GeoIpClient { /// Returns a dictionary of geoip properties for the given ip address. pub fn get_geoip_properties(&self, ip_address: Option<&str>) -> HashMap { - match ip_address { - None => { - info!("No IP address provided; returning empty properties"); - HashMap::new() + let ip = match ip_address { + Some(ip) if self.is_valid_ip(ip) => ip, + _ => { + info!("No valid IP address provided; returning empty properties"); + return HashMap::new(); } - Some(ip) if !self.is_valid_ip(ip) => { - info!("Returning empty properties for IP: {}", ip); + }; + + match IpAddr::from_str(ip) { + Ok(addr) => self + .lookup_city(ip, addr) + .map(|city| extract_properties(&city)) + .unwrap_or_default(), + Err(_) => { + error!("Invalid IP address: {}", ip); HashMap::new() } - Some(ip) => match IpAddr::from_str(ip) { - Ok(addr) => self - .lookup_city(ip, addr) - .map(|city| extract_properties(&city)) - .unwrap_or_default(), - Err(_) => { - error!("Invalid IP address: {}", ip); - HashMap::new() - } - }, } } } @@ -173,12 +171,17 @@ mod tests { initialize(); let service = create_test_service(); let test_cases = vec![ - ("13.106.122.3", "Australia"), - ("31.28.64.3", "United Kingdom"), - ("2600:6c52:7a00:11c:1b6:b7b0:ea19:6365", "United States"), + ("13.106.122.3", "Australia", "Sydney"), + ("31.28.64.3", "United Kingdom", "Baldock"), + ( + "2600:6c52:7a00:11c:1b6:b7b0:ea19:6365", + "United States", + "San Luis Obispo", + ), + ("187.188.10.252", "Mexico", "Coyoacán"), ]; - for (ip, expected_country) in test_cases { + for (ip, expected_country, expected_city) in test_cases { let result = service.get_geoip_properties(Some(ip)); info!("GeoIP lookup result for IP {}: {:?}", ip, result); info!( @@ -190,6 +193,10 @@ mod tests { result.get("$geoip_country_name"), Some(&expected_country.to_string()) ); + assert_eq!( + result.get("$geoip_city_name"), + Some(&expected_city.to_string()) + ); assert_eq!(result.len(), 7); } }