fix: Too much data for declared Content-Length error when endpoint …

…caches enabled (#92)

* Fix exception with cached responses

* Lint / format

src/environments.py

@@ -4,7 +4,6 @@
 from functools import lru_cache
 
 import httpx
-from fastapi.responses import ORJSONResponse
 from flag_engine.engine import (
     get_environment_feature_state,
     get_environment_feature_states,
@@ -15,7 +14,7 @@
 from orjson import orjson
 
 from src.cache import BaseEnvironmentsCache, LocalMemEnvironmentsCache
-from src.exceptions import FlagsmithUnknownKeyError
+from src.exceptions import FeatureNotFoundError, FlagsmithUnknownKeyError
 from src.feature_utils import filter_out_server_key_only_feature_states
 from src.mappers import (
     map_feature_state_to_response_data,
@@ -73,7 +72,7 @@ async def refresh_environment_caches(self):
 
     def get_flags_response_data(
         self, environment_key: str, feature: str = None
-    ) -> ORJSONResponse:
+    ) -> dict[str, typing.Any]:
         environment_document = self.get_environment(environment_key)
         environment = build_environment_model(environment_document)
 
@@ -84,13 +83,7 @@ def get_flags_response_data(
                 feature_states=[feature_state],
                 environment=environment,
             ):
-                return ORJSONResponse(
-                    status_code=404,
-                    content={
-                        "status": "not_found",
-                        "message": f"feature '{feature}' not found",
-                    },
-                )
+                raise FeatureNotFoundError()
 
             data = map_feature_state_to_response_data(feature_state)
 
@@ -101,11 +94,11 @@ def get_flags_response_data(
             )
             data = map_feature_states_to_response_data(feature_states)
 
-        return ORJSONResponse(data)
+        return data
 
     def get_identity_response_data(
         self, input_data: IdentityWithTraits, environment_key: str
-    ) -> ORJSONResponse:
+    ) -> dict[str, typing.Any]:
         environment_document = self.get_environment(environment_key)
         environment = build_environment_model(environment_document)
         identity = IdentityModel(
@@ -127,7 +120,7 @@ def get_identity_response_data(
                 identity_hash_key=identity.composite_key,
             ),
         }
-        return ORJSONResponse(data)
+        return data
 
     def get_environment(self, client_side_key: str) -> dict[str, typing.Any]:
         if environment_document := self.cache.get_environment(client_side_key):

src/exceptions.py

@@ -1,2 +1,6 @@
 class FlagsmithUnknownKeyError(KeyError):
     pass
+
+
+class FeatureNotFoundError(Exception):
+    pass

src/main.py

@@ -12,7 +12,7 @@
 
 from .cache import LocalMemEnvironmentsCache
 from .environments import EnvironmentService
-from .exceptions import FlagsmithUnknownKeyError
+from .exceptions import FeatureNotFoundError, FlagsmithUnknownKeyError
 from .models import IdentityWithTraits
 from .settings import Settings
 from .sse import router as sse_router
@@ -51,15 +51,27 @@ async def health_check():
 
 @app.get("/api/v1/flags/", response_class=ORJSONResponse)
 async def flags(feature: str = None, x_environment_key: str = Header(None)):
-    return environment_service.get_flags_response_data(x_environment_key, feature)
+    try:
+        data = environment_service.get_flags_response_data(x_environment_key, feature)
+    except FeatureNotFoundError:
+        return ORJSONResponse(
+            status_code=404,
+            content={
+                "status": "not_found",
+                "message": f"feature '{feature}' not found",
+            },
+        )
+
+    return ORJSONResponse(data)
 
 
 @app.post("/api/v1/identities/", response_class=ORJSONResponse)
 async def identity(
     input_data: IdentityWithTraits,
     x_environment_key: str = Header(None),
 ):
-    return environment_service.get_identity_response_data(input_data, x_environment_key)
+    data = environment_service.get_identity_response_data(input_data, x_environment_key)
+    return ORJSONResponse(data)
 
 
 @app.on_event("startup")