Merge pull request #159 from Flagsmith/feat/environment-document-304

feat: Set If-Modified-Since header for environment document requests

Author: rolodato

src/edge_proxy/environments.py

@@ -1,8 +1,10 @@
 import typing
 from datetime import datetime
+from email.utils import formatdate
 from functools import lru_cache
 
 import httpx
+import starlette.status
 import structlog
 from flag_engine.engine import (
     get_environment_feature_state,
@@ -22,7 +24,7 @@
     map_traits_to_response_data,
 )
 from edge_proxy.models import IdentityWithTraits
-from edge_proxy.settings import AppSettings
+from edge_proxy.settings import AppSettings, EnvironmentKeyPair
 
 logger = structlog.get_logger(__name__)
 
@@ -58,9 +60,7 @@ async def refresh_environment_caches(self):
         received_error = False
         for key_pair in self.settings.environment_key_pairs:
             try:
-                environment_document = await self._fetch_document(
-                    key_pair.server_side_key
-                )
+                environment_document = await self._fetch_document(key_pair)
                 if self.cache.put_environment(
                     environment_api_key=key_pair.client_side_key,
                     environment_document=environment_document,
@@ -142,11 +142,41 @@ def get_environment(self, client_side_key: str) -> dict[str, typing.Any]:
             return environment_document
         raise FlagsmithUnknownKeyError(client_side_key)
 
-    async def _fetch_document(self, server_side_key: str) -> dict[str, typing.Any]:
+    async def _fetch_document(
+        self, key_pair: EnvironmentKeyPair
+    ) -> dict[str, typing.Any]:
+        headers = {
+            "X-Environment-Key": key_pair.server_side_key,
+        }
+        environment_document = self.cache.get_environment(
+            environment_api_key=key_pair.client_side_key
+        )
+        if environment_document:
+            updated_at: str = environment_document.get("updated_at")
+            if updated_at:
+                try:
+                    epoch_seconds = datetime.fromisoformat(updated_at).timestamp()
+                    # Same implementation as https://docs.djangoproject.com/en/4.2/ref/utils/#django.utils.http.http_date
+                    headers["If-Modified-Since"] = formatdate(
+                        epoch_seconds, usegmt=True
+                    )
+                except ValueError:
+                    logger.warning(
+                        f"failed to parse updated_at, environment={key_pair.client_side_key} updated_at={updated_at}"
+                    )
+            else:
+                logger.warning(
+                    f"received environment with no updated_at: {key_pair.client_side_key}"
+                )
         response = await self._client.get(
             url=f"{self.settings.api_url}/environment-document/",
-            headers={"X-Environment-Key": server_side_key},
+            headers=headers,
         )
+        if response.status_code == starlette.status.HTTP_304_NOT_MODIFIED:
+            assert environment_document, (
+                f"GET /environment-document returned 304 without a cached document. environment={key_pair.client_side_key}"
+            )
+            return environment_document
         response.raise_for_status()
         return orjson.loads(response.text)
 

tests/fixtures/response_data.py

@@ -91,6 +91,7 @@
 
 
 environment_1 = {
+    "updated_at": "1969-07-20T20:17:40Z",
     "feature_states": [
         _environment_feature_state_1,
         _environment_feature_state_2,

tests/test_environments.py

@@ -196,6 +196,41 @@ async def test_refresh_environment_caches_clears_endpoint_caches_if_environment_
     assert environment_service.get_flags_response_data.cache_info().currsize == 0
 
 
+@pytest.mark.asyncio
+async def test_refresh_environment_caches_sets_last_modified_if_environment_was_cached(
+    mocker: MockerFixture,
+):
+    # Given
+    # An EnvironmentService that saves the last If-Modified-Since request header its client used
+    if_modified_since = ""
+
+    def save_if_modified_since(**kwargs):
+        nonlocal if_modified_since
+        if headers := kwargs.get("headers"):
+            if_modified_since = headers.get("If-Modified-Since")
+        return mocker.MagicMock(
+            text=orjson.dumps(environment_1),
+        )
+
+    mocked_client = mocker.AsyncMock()
+    mocked_client.get.side_effect = save_if_modified_since
+    environment_service = EnvironmentService(settings=settings, client=mocked_client)
+
+    # When
+    # We refresh its environment caches
+    await environment_service.refresh_environment_caches()
+    # Then
+    # No If-Modified-Since request header is sent initially
+    assert not if_modified_since
+
+    # When
+    # We refresh the caches while having an existing valid cache
+    await environment_service.refresh_environment_caches()
+    # Then
+    # If-Modified-Since is set to the environment's updated_at in HTTP date format
+    assert if_modified_since == "Sun, 20 Jul 1969 20:17:40 GMT"
+
+
 @pytest.mark.asyncio
 async def test_get_identity_flags_response_skips_cache_for_different_identity(
     mocker: MockerFixture,