Add ETag support for local evaluation polling

haacked · haacked · commit ef99809b765d · 2025-12-01T18:51:33.000-08:00
Add support for HTTP conditional requests using ETags to reduce bandwidth
when polling for feature flag definitions. When flag definitions haven't
changed, the server returns 304 Not Modified and the SDK skips processing.

- Add GetResponse dataclass to encapsulate response data, ETag, and status
- Update get() to send If-None-Match header and handle 304 responses
- Store ETag in client and pass it on subsequent polling requests
- Skip flag processing when 304 Not Modified is received
diff --git a/posthog/client.py b/posthog/client.py
@@ -232,6 +232,7 @@ def __init__(
         self.distinct_ids_feature_flags_reported = SizeLimitedDict(MAX_DICT_SIZE, set)
         self.flag_cache = self._initialize_flag_cache(flag_fallback_cache_url)
         self.flag_definition_version = 0
+        self._flags_etag: Optional[str] = None
         self.disabled = disabled
         self.disable_geoip = disable_geoip
         self.historical_migration = historical_migration
@@ -1183,11 +1184,24 @@ def _load_feature_flags(self):
                 f"/api/feature_flag/local_evaluation/?token={self.api_key}&send_cohorts",
                 self.host,
                 timeout=10,
+                etag=self._flags_etag,
             )
 
-            self.feature_flags = response["flags"] or []
-            self.group_type_mapping = response["group_type_mapping"] or {}
-            self.cohorts = response["cohorts"] or {}
+            # Update stored ETag
+            if response.etag:
+                self._flags_etag = response.etag
+
+            # If 304 Not Modified, flags haven't changed - skip processing
+            if response.not_modified:
+                self.log.debug(
+                    "[FEATURE FLAGS] Flags not modified (304), using cached data"
+                )
+                self._last_feature_flag_poll = datetime.now(tz=tzutc())
+                return
+
+            self.feature_flags = response.data["flags"] or []
+            self.group_type_mapping = response.data["group_type_mapping"] or {}
+            self.cohorts = response.data["cohorts"] or {}
 
             # Check if flag definitions changed and update version
             if self.flag_cache and old_flags_by_key != (
diff --git a/posthog/request.py b/posthog/request.py
@@ -1,5 +1,6 @@
 import json
 import logging
+from dataclasses import dataclass
 from datetime import date, datetime
 from gzip import GzipFile
 from io import BytesIO
@@ -12,6 +13,16 @@
 from posthog.utils import remove_trailing_slash
 from posthog.version import VERSION
 
+
+@dataclass
+class GetResponse:
+    """Response from a GET request with ETag support."""
+
+    data: Any
+    etag: Optional[str] = None
+    not_modified: bool = False
+
+
 # Retry on both connect and read errors
 # by default read errors will only retry idempotent HTTP methods (so not POST)
 adapter = requests.adapters.HTTPAdapter(
@@ -139,12 +150,13 @@ def remote_config(
     timeout: int = 15,
 ) -> Any:
     """Get remote config flag value from remote_config API endpoint"""
-    return get(
+    response = get(
         personal_api_key,
         f"/api/projects/@current/feature_flags/{key}/remote_config?token={project_api_key}",
         host,
         timeout,
     )
+    return response.data
 
 
 def batch_post(
@@ -162,15 +174,40 @@ def batch_post(
 
 
 def get(
-    api_key: str, url: str, host: Optional[str] = None, timeout: Optional[int] = None
-) -> requests.Response:
-    url = remove_trailing_slash(host or DEFAULT_HOST) + url
-    res = requests.get(
-        url,
-        headers={"Authorization": "Bearer %s" % api_key, "User-Agent": USER_AGENT},
-        timeout=timeout,
+    api_key: str,
+    url: str,
+    host: Optional[str] = None,
+    timeout: Optional[int] = None,
+    etag: Optional[str] = None,
+) -> GetResponse:
+    """
+    Make a GET request with optional ETag support.
+
+    If an etag is provided, sends If-None-Match header. Returns GetResponse with:
+    - not_modified=True and data=None if server returns 304
+    - not_modified=False and data=response if server returns 200
+    """
+    log = logging.getLogger("posthog")
+    full_url = remove_trailing_slash(host or DEFAULT_HOST) + url
+    headers = {"Authorization": "Bearer %s" % api_key, "User-Agent": USER_AGENT}
+
+    if etag:
+        headers["If-None-Match"] = etag
+
+    res = requests.get(full_url, headers=headers, timeout=timeout)
+
+    # Handle 304 Not Modified
+    if res.status_code == 304:
+        log.debug(f"GET {full_url} returned 304 Not Modified")
+        response_etag = res.headers.get("ETag")
+        return GetResponse(data=None, etag=response_etag or etag, not_modified=True)
+
+    # Handle normal response
+    data = _process_response(
+        res, success_message=f"GET {full_url} completed successfully"
     )
-    return _process_response(res, success_message=f"GET {url} completed successfully")
+    response_etag = res.headers.get("ETag")
+    return GetResponse(data=data, etag=response_etag, not_modified=False)
 
 
 class APIError(Exception):
diff --git a/posthog/test/test_client.py b/posthog/test/test_client.py
@@ -9,7 +9,7 @@
 
 from posthog.client import Client
 from posthog.contexts import get_context_session_id, new_context, set_context_session
-from posthog.request import APIError
+from posthog.request import APIError, GetResponse
 from posthog.test.test_utils import FAKE_TEST_API_KEY
 from posthog.types import FeatureFlag, LegacyFlagMetadata
 from posthog.version import VERSION
@@ -2095,13 +2095,21 @@ def test_enable_local_evaluation_false_disables_poller(
         self, patch_get, patch_poller
     ):
         """Test that when enable_local_evaluation=False, the poller is not started"""
-        patch_get.return_value = {
-            "flags": [
-                {"id": 1, "name": "Beta Feature", "key": "beta-feature", "active": True}
-            ],
-            "group_type_mapping": {},
-            "cohorts": {},
-        }
+        patch_get.return_value = GetResponse(
+            data={
+                "flags": [
+                    {
+                        "id": 1,
+                        "name": "Beta Feature",
+                        "key": "beta-feature",
+                        "active": True,
+                    }
+                ],
+                "group_type_mapping": {},
+                "cohorts": {},
+            },
+            etag='"test-etag"',
+        )
 
         client = Client(
             FAKE_TEST_API_KEY,
@@ -2123,13 +2131,21 @@ def test_enable_local_evaluation_false_disables_poller(
     @mock.patch("posthog.client.get")
     def test_enable_local_evaluation_true_starts_poller(self, patch_get, patch_poller):
         """Test that when enable_local_evaluation=True (default), the poller is started"""
-        patch_get.return_value = {
-            "flags": [
-                {"id": 1, "name": "Beta Feature", "key": "beta-feature", "active": True}
-            ],
-            "group_type_mapping": {},
-            "cohorts": {},
-        }
+        patch_get.return_value = GetResponse(
+            data={
+                "flags": [
+                    {
+                        "id": 1,
+                        "name": "Beta Feature",
+                        "key": "beta-feature",
+                        "active": True,
+                    }
+                ],
+                "group_type_mapping": {},
+                "cohorts": {},
+            },
+            etag='"test-etag"',
+        )
 
         client = Client(
             FAKE_TEST_API_KEY,
diff --git a/posthog/test/test_feature_flags.py b/posthog/test/test_feature_flags.py
@@ -11,7 +11,7 @@
     match_property,
     relative_date_parse_for_feature_flag_matching,
 )
-from posthog.request import APIError
+from posthog.request import APIError, GetResponse
 from posthog.test.test_utils import FAKE_TEST_API_KEY
 
 
@@ -2348,23 +2348,27 @@ def test_production_style_multivariate_dependency_chain(
     @mock.patch("posthog.client.Poller")
     @mock.patch("posthog.client.get")
     def test_load_feature_flags(self, patch_get, patch_poll):
-        patch_get.return_value = {
-            "flags": [
-                {
-                    "id": 1,
-                    "name": "Beta Feature",
-                    "key": "beta-feature",
-                    "active": True,
-                },
-                {
-                    "id": 2,
-                    "name": "Alpha Feature",
-                    "key": "alpha-feature",
-                    "active": False,
-                },
-            ],
-            "group_type_mapping": {"0": "company"},
-        }
+        patch_get.return_value = GetResponse(
+            data={
+                "flags": [
+                    {
+                        "id": 1,
+                        "name": "Beta Feature",
+                        "key": "beta-feature",
+                        "active": True,
+                    },
+                    {
+                        "id": 2,
+                        "name": "Alpha Feature",
+                        "key": "alpha-feature",
+                        "active": False,
+                    },
+                ],
+                "group_type_mapping": {"0": "company"},
+                "cohorts": {},
+            },
+            etag='"abc123"',
+        )
         client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
         with freeze_time("2020-01-01T12:01:00.0000Z"):
             client.load_feature_flags()
@@ -2375,6 +2379,106 @@ def test_load_feature_flags(self, patch_get, patch_poll):
             client._last_feature_flag_poll.isoformat(), "2020-01-01T12:01:00+00:00"
         )
         self.assertEqual(patch_poll.call_count, 1)
+        # Verify ETag is stored
+        self.assertEqual(client._flags_etag, '"abc123"')
+
+    @mock.patch("posthog.client.Poller")
+    @mock.patch("posthog.client.get")
+    def test_load_feature_flags_sends_etag_on_subsequent_requests(
+        self, patch_get, patch_poll
+    ):
+        """Test that the ETag is sent in If-None-Match header on subsequent requests"""
+        patch_get.return_value = GetResponse(
+            data={
+                "flags": [{"id": 1, "key": "beta-feature", "active": True}],
+                "group_type_mapping": {},
+                "cohorts": {},
+            },
+            etag='"initial-etag"',
+        )
+        client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
+        client.load_feature_flags()
+
+        # First call should have no etag
+        first_call_kwargs = patch_get.call_args_list[0][1]
+        self.assertIsNone(first_call_kwargs.get("etag"))
+
+        # Simulate second call
+        client._load_feature_flags()
+
+        # Second call should have the etag
+        second_call_kwargs = patch_get.call_args_list[1][1]
+        self.assertEqual(second_call_kwargs.get("etag"), '"initial-etag"')
+
+    @mock.patch("posthog.client.Poller")
+    @mock.patch("posthog.client.get")
+    def test_load_feature_flags_304_not_modified(self, patch_get, patch_poll):
+        """Test that 304 Not Modified responses skip flag processing"""
+        # First response with flags
+        initial_response = GetResponse(
+            data={
+                "flags": [{"id": 1, "key": "beta-feature", "active": True}],
+                "group_type_mapping": {"0": "company"},
+                "cohorts": {},
+            },
+            etag='"test-etag"',
+        )
+        # Second response is 304 Not Modified
+        not_modified_response = GetResponse(
+            data=None,
+            etag='"test-etag"',
+            not_modified=True,
+        )
+        patch_get.side_effect = [initial_response, not_modified_response]
+
+        client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
+        client.load_feature_flags()
+
+        # Verify initial flags are loaded
+        self.assertEqual(len(client.feature_flags), 1)
+        self.assertEqual(client.feature_flags[0]["key"], "beta-feature")
+        self.assertEqual(client.group_type_mapping, {"0": "company"})
+
+        # Second call with 304
+        client._load_feature_flags()
+
+        # Flags should still be the same (not cleared)
+        self.assertEqual(len(client.feature_flags), 1)
+        self.assertEqual(client.feature_flags[0]["key"], "beta-feature")
+        self.assertEqual(client.group_type_mapping, {"0": "company"})
+
+    @mock.patch("posthog.client.Poller")
+    @mock.patch("posthog.client.get")
+    def test_load_feature_flags_etag_updated_on_new_response(
+        self, patch_get, patch_poll
+    ):
+        """Test that ETag is updated when flags change"""
+        patch_get.side_effect = [
+            GetResponse(
+                data={
+                    "flags": [{"id": 1, "key": "flag-v1", "active": True}],
+                    "group_type_mapping": {},
+                    "cohorts": {},
+                },
+                etag='"etag-v1"',
+            ),
+            GetResponse(
+                data={
+                    "flags": [{"id": 1, "key": "flag-v2", "active": True}],
+                    "group_type_mapping": {},
+                    "cohorts": {},
+                },
+                etag='"etag-v2"',
+            ),
+        ]
+
+        client = Client(FAKE_TEST_API_KEY, personal_api_key="test")
+        client.load_feature_flags()
+        self.assertEqual(client._flags_etag, '"etag-v1"')
+
+        client._load_feature_flags()
+        self.assertEqual(client._flags_etag, '"etag-v2"')
+        self.assertEqual(client.feature_flags[0]["key"], "flag-v2")
 
     def test_load_feature_flags_wrong_key(self):
         client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
