diff --git a/statsig/http_worker.py b/statsig/http_worker.py index 557002e..bd27bec 100644 --- a/statsig/http_worker.py +++ b/statsig/http_worker.py @@ -6,7 +6,7 @@ from concurrent.futures.thread import ThreadPoolExecutor from decimal import Decimal from io import BytesIO -from urllib.parse import urlparse +from urllib.parse import urlparse, urljoin from typing import Callable, Tuple, Optional, Any, Dict, List import ijson @@ -160,14 +160,27 @@ def get_id_list( extra_tags = {} if id_list_file_id: extra_tags["id_list_file_id"] = id_list_file_id + proxy_preferred_url = self._rewrite_cdn_id_list_file_url_to_get_id_lists_proxy(url) resp = self._get_request( - url, + proxy_preferred_url, headers, log_on_exception, tag="get_id_list", get_text_value_only=True, extra_tags=extra_tags, ) + if ( + proxy_preferred_url != url + and (resp is None or not self._is_success_code(resp.status_code)) + ): + resp = self._get_request( + url, + headers, + log_on_exception, + tag="get_id_list", + get_text_value_only=True, + extra_tags=extra_tags, + ) if resp is not None and self._is_success_code(resp.status_code): return on_complete(resp) return on_complete(None) @@ -622,5 +635,18 @@ def __configure_endpoints(self, options: StatsigOptions) -> None: self.__api_for_get_id_lists = api_for_get_id_lists self.__api_for_log_event = api_for_log_event + def _rewrite_cdn_id_list_file_url_to_get_id_lists_proxy(self, url: str) -> str: + if self.__is_cdn_url(self.__api_for_get_id_lists): + return url + if not self.__is_cdn_url(url): + return url + parsed = urlparse(url) + if not parsed.path.startswith("/v1/download_id_list_file/"): + return url + suffix = parsed.path.removeprefix("/v1/") + if parsed.query: + suffix = f"{suffix}?{parsed.query}" + return urljoin(self.__api_for_get_id_lists, suffix) + def __is_cdn_url(self, url: str) -> bool: return url.startswith(STATSIG_CDN) diff --git a/tests/test_network.py b/tests/test_network.py index 773d230..f0ed1fe 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -182,6 +182,8 @@ def fake_request(_method, _url, *_args, **_kwargs): self.assertEqual(kwargs["request_path"], "/v1/get_id_lists") def test_id_list_network_latency_metric_includes_file_id_tag(self): + self.net._HttpWorker__api_for_get_id_lists = "https://api.statsigcdn.com/v1/" + def fake_request(_method, _url, *_args, **_kwargs): return RequestResult( data={}, @@ -209,6 +211,60 @@ def fake_request(_method, _url, *_args, **_kwargs): self.assertEqual(kwargs["request_path"], "/v1/download_id_list_file") self.assertEqual(kwargs["extra_tags"], {"id_list_file_id": "4PKKLINp6EZW3DNQ73sCxY"}) + def test_id_list_download_uses_get_id_lists_proxy_when_configured(self): + captured_urls = [] + self.net._HttpWorker__api_for_get_id_lists = "http://sfp-proxy.local/v1/" + + def fake_request(_method, url, *_args, **_kwargs): + captured_urls.append(url) + return RequestResult( + data={}, + status_code=200, + success=True, + error=None, + text="+1\r", + headers={"content-length": "3"}, + ) + + with patch.object(self.net, "_run_request_with_strict_timeout", side_effect=fake_request): + self.net.get_id_list( + lambda *_: None, + "https://api.statsigcdn.com/v1/download_id_list_file/foo%2Fbar?sv=2020-10-02&sig=abc&k=secret-test", + headers={}, + ) + + self.assertEqual( + captured_urls, + ["http://sfp-proxy.local/v1/download_id_list_file/foo%2Fbar?sv=2020-10-02&sig=abc&k=secret-test"], + ) + + def test_id_list_download_falls_back_to_cdn_when_proxy_fails(self): + captured_urls = [] + self.net._HttpWorker__api_for_get_id_lists = "http://sfp-proxy.local/v1/" + cdn_url = ( + "https://api.statsigcdn.com/v1/download_id_list_file/foo" + "?sv=2020-10-02&sig=abc&k=secret-test" + ) + proxy_url = "http://sfp-proxy.local/v1/download_id_list_file/foo?sv=2020-10-02&sig=abc&k=secret-test" + + def fake_request(_method, url, *_args, **_kwargs): + captured_urls.append(url) + if url == proxy_url: + return RequestResult(data=None, status_code=502, success=False, error=None) + return RequestResult( + data={}, + status_code=200, + success=True, + error=None, + text="+1\r", + headers={"content-length": "3"}, + ) + + with patch.object(self.net, "_run_request_with_strict_timeout", side_effect=fake_request): + self.net.get_id_list(lambda *_: None, cdn_url, headers={}) + + self.assertEqual(captured_urls, [proxy_url, cdn_url]) + def test_log_event_does_not_emit_network_latency(self): def fake_request(_method, _url, *_args, **_kwargs): return RequestResult(data={}, status_code=200, success=True, error=None)