diff --git a/src/edge_proxy/environments.py b/src/edge_proxy/environments.py index 8127a55..d3b6c7e 100644 --- a/src/edge_proxy/environments.py +++ b/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) diff --git a/tests/fixtures/response_data.py b/tests/fixtures/response_data.py index 927a923..b177227 100644 --- a/tests/fixtures/response_data.py +++ b/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, diff --git a/tests/test_environments.py b/tests/test_environments.py index e9bbc88..edc55c2 100644 --- a/tests/test_environments.py +++ b/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,