[download id list] support download id list file from sfp

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)

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)