Add identity overrides support

 - Bump flag-engine, pydantic
 - Store identity overrides by identifier in the cache
 - Better names for `api_poll_frequency`, `api_poll_timeout`
 - Typing improvements

pyproject.toml

@@ -2,16 +2,17 @@
 dynamic = ["version"]
 name = "edge-proxy"
 dependencies = [
-  "fastapi",
-  "flagsmith-flag-engine>4,<5",
-  "httpx",
-  "marshmallow",
-  "orjson",
-  "pydantic<2",
-  "python-decouple",
-  "python-dotenv",
-  "structlog",
-  "uvicorn",
+    "fastapi",
+    "flagsmith-flag-engine>5",
+    "httpx",
+    "marshmallow",
+    "orjson",
+    "pydantic",
+    "python-decouple",
+    "python-dotenv",
+    "structlog",
+    "uvicorn",
+    "pydantic-settings>=2.2.1",
 ]
 requires-python = ">= 3.12"
 

requirements-dev.lock

@@ -8,6 +8,8 @@
 #   with-sources: false
 
 -e file:.
+annotated-types==0.6.0
+    # via pydantic
 anyio==4.3.0
     # via httpx
     # via starlette
@@ -26,7 +28,7 @@ fastapi==0.110.1
     # via edge-proxy
 filelock==3.13.3
     # via virtualenv
-flagsmith-flag-engine==4.1.0
+flagsmith-flag-engine==5.1.1
     # via edge-proxy
 freezegun==1.4.0
     # via pytest-freezegun
@@ -58,13 +60,18 @@ platformdirs==4.2.0
 pluggy==1.4.0
     # via pytest
 pre-commit==3.7.0
-pydantic==1.10.15
+pydantic==2.7.1
     # via edge-proxy
     # via fastapi
     # via flagsmith-flag-engine
     # via pydantic-collections
+    # via pydantic-settings
 pydantic-collections==0.5.4
     # via flagsmith-flag-engine
+pydantic-core==2.18.2
+    # via pydantic
+pydantic-settings==2.2.1
+    # via edge-proxy
 pytest==8.1.1
     # via pytest-asyncio
     # via pytest-freezegun
@@ -78,10 +85,11 @@ python-decouple==3.8
     # via edge-proxy
 python-dotenv==1.0.1
     # via edge-proxy
+    # via pydantic-settings
 pyyaml==6.0.1
     # via pre-commit
 reorder-python-imports==3.12.0
-semver==2.13.0
+semver==3.0.2
     # via flagsmith-flag-engine
 setuptools==69.2.0
     # via nodeenv
@@ -98,6 +106,7 @@ typing-extensions==4.11.0
     # via fastapi
     # via pydantic
     # via pydantic-collections
+    # via pydantic-core
 uvicorn==0.29.0
     # via edge-proxy
 virtualenv==20.25.1

requirements.lock

@@ -8,6 +8,8 @@
 #   with-sources: false
 
 -e file:.
+annotated-types==0.6.0
+    # via pydantic
 anyio==4.3.0
     # via httpx
     # via starlette
@@ -18,7 +20,7 @@ click==8.1.7
     # via uvicorn
 fastapi==0.110.1
     # via edge-proxy
-flagsmith-flag-engine==4.1.0
+flagsmith-flag-engine==5.1.1
     # via edge-proxy
 h11==0.14.0
     # via httpcore
@@ -36,18 +38,24 @@ orjson==3.10.0
     # via edge-proxy
 packaging==24.0
     # via marshmallow
-pydantic==1.10.15
+pydantic==2.7.1
     # via edge-proxy
     # via fastapi
     # via flagsmith-flag-engine
     # via pydantic-collections
+    # via pydantic-settings
 pydantic-collections==0.5.4
     # via flagsmith-flag-engine
+pydantic-core==2.18.2
+    # via pydantic
+pydantic-settings==2.2.1
+    # via edge-proxy
 python-decouple==3.8
     # via edge-proxy
 python-dotenv==1.0.1
     # via edge-proxy
-semver==2.13.0
+    # via pydantic-settings
+semver==3.0.2
     # via flagsmith-flag-engine
 sniffio==1.3.1
     # via anyio
@@ -60,5 +68,6 @@ typing-extensions==4.11.0
     # via fastapi
     # via pydantic
     # via pydantic-collections
+    # via pydantic-core
 uvicorn==0.29.0
     # via edge-proxy

src/edge_proxy/cache.py

@@ -1,4 +1,5 @@
-import typing
+from collections import defaultdict
+from typing import Any
 from abc import ABC
 
 
@@ -9,7 +10,7 @@ def __init__(self, *args, **kwargs):
     def put_environment(
         self,
         environment_api_key: str,
-        environment_document: typing.Dict[str, typing.Any],
+        environment_document: dict[str, Any],
     ) -> bool:
         """
         Update the environment cache for the given key with the given environment document.
@@ -26,29 +27,54 @@ def put_environment(
     def _put_environment(
         self,
         environment_api_key: str,
-        environment_document: typing.Dict[str, typing.Any],
+        environment_document: dict[str, Any],
     ) -> None:
         raise NotImplementedError()
 
-    def get_environment(
-        self, environment_api_key: str
-    ) -> typing.Dict[str, typing.Any] | None:
+    def get_environment(self, environment_api_key: str) -> dict[str, Any] | None:
+        raise NotImplementedError()
+
+    def get_identity(
+        self,
+        environment_api_key: str,
+        identifier: str,
+    ) -> dict[str, Any]:
         raise NotImplementedError()
 
 
+_LocalCacheDict = dict[str, dict[str, Any]]
+
+
 class LocalMemEnvironmentsCache(BaseEnvironmentsCache):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._cache = {}
+        self._environment_cache: _LocalCacheDict = {}
+        self._identity_override_cache = defaultdict[str, _LocalCacheDict](dict)
 
     def _put_environment(
         self,
         environment_api_key: str,
-        environment_document: typing.Dict[str, typing.Any],
+        environment_document: dict[str, Any],
     ) -> None:
-        self._cache[environment_api_key] = environment_document
+        self._environment_cache[environment_api_key] = environment_document
+        for identity_document in environment_document.get("identity_overrides") or []:
+            if identifier := identity_document.get("identifier"):
+                self._identity_override_cache[environment_api_key][identifier] = (
+                    identity_document
+                )
 
     def get_environment(
-        self, environment_api_key
-    ) -> typing.Dict[str, typing.Any] | None:
-        return self._cache.get(environment_api_key)
+        self,
+        environment_api_key: str,
+    ) -> dict[str, Any] | None:
+        return self._environment_cache.get(environment_api_key)
+
+    def get_identity(
+        self,
+        environment_api_key: str,
+        identifier: str,
+    ) -> dict[str, Any]:
+        return self._identity_override_cache[environment_api_key].get(identifier) or {
+            "environment_api_key": environment_api_key,
+            "identifier": identifier,
+        }

src/edge_proxy/environments.py

@@ -9,7 +9,7 @@
     get_environment_feature_states,
     get_identity_feature_states,
 )
-from flag_engine.environments.builders import build_environment_model
+from flag_engine.environments.models import EnvironmentModel
 from flag_engine.identities.models import IdentityModel
 from orjson import orjson
 
@@ -76,7 +76,7 @@ def get_flags_response_data(
         self, environment_key: str, feature: str = None
     ) -> dict[str, typing.Any]:
         environment_document = self.get_environment(environment_key)
-        environment = build_environment_model(environment_document)
+        environment = EnvironmentModel.model_validate(environment_document)
 
         if feature:
             feature_state = get_environment_feature_state(environment, feature)
@@ -102,9 +102,12 @@ def get_identity_response_data(
         self, input_data: IdentityWithTraits, environment_key: str
     ) -> dict[str, typing.Any]:
         environment_document = self.get_environment(environment_key)
-        environment = build_environment_model(environment_document)
-        identity = IdentityModel(
-            identifier=input_data.identifier, environment_api_key=environment_key
+        environment = EnvironmentModel.model_validate(environment_document)
+        identity = IdentityModel.model_validate(
+            self.cache.get_identity(
+                environment_api_key=environment_key,
+                identifier=input_data.identifier,
+            )
         )
         trait_models = input_data.traits
         flags = filter_out_server_key_only_feature_states(

src/edge_proxy/mappers.py

@@ -30,4 +30,4 @@ def map_feature_states_to_response_data(
 def map_traits_to_response_data(
     traits: list[TraitModel],
 ) -> list[dict[str, Any]]:
-    return [trait.dict() for trait in traits]
+    return [trait.model_dump() for trait in traits]

src/edge_proxy/settings.py

@@ -4,12 +4,13 @@
 import sys
 from enum import Enum
 from pathlib import Path
-from typing import Any, Dict, List, Tuple
+from typing import Any
 
 import structlog
 
-from pydantic import BaseModel, BaseSettings, HttpUrl, IPvAnyAddress, Field
-from pydantic.env_settings import SettingsSourceCallable
+from pydantic import AliasChoices, BaseModel, HttpUrl, IPvAnyAddress, Field
+
+from pydantic_settings import BaseSettings, PydanticBaseSettingsSource
 
 
 CONFIG_PATH = os.environ.get(
@@ -41,7 +42,7 @@ def to_logging_log_level(self) -> int:
 def ensure_defaults() -> None:
     if not os.path.exists(CONFIG_PATH):
         defaults = AppSettings()
-        defaults_json = defaults.json(indent=4, exclude_none=True)
+        defaults_json = defaults.model_dump_json(indent=4, exclude_none=True)
         print(defaults_json, file=sys.stdout)
         try:
             with open(CONFIG_PATH, "w") as fp:
@@ -54,7 +55,7 @@ def ensure_defaults() -> None:
             )
 
 
-def json_config_settings_source(settings: BaseSettings) -> Dict[str, Any]:
+def json_config_settings_source() -> dict[str, Any]:
     """
     A simple settings source that loads variables from a JSON file
     at the project's root.
@@ -92,7 +93,7 @@ class ServerSettings(BaseModel):
 
 
 class AppSettings(BaseModel):
-    environment_key_pairs: List[EnvironmentKeyPair] = [
+    environment_key_pairs: list[EnvironmentKeyPair] = [
         EnvironmentKeyPair(
             server_side_key="ser.environment_key",
             client_side_key="environment_key",
@@ -101,28 +102,35 @@ class AppSettings(BaseModel):
     api_url: HttpUrl = "https://edge.api.flagsmith.com/api/v1"
     api_poll_frequency_seconds: int = Field(
         default=10,
-        aliases=["api_poll_frequency_seconds", "api_poll_frequency"],
+        validation_alias=AliasChoices(
+            "api_poll_frequency_seconds",
+            "api_poll_frequency",
+        ),
     )
     api_poll_timeout_seconds: int = Field(
         default=5,
-        aliases=["api_poll_timeout_seconds", "api_poll_timeout"],
+        validation_alias=AliasChoices(
+            "api_poll_timeout_seconds",
+            "api_poll_timeout",
+        ),
     )
     endpoint_caches: EndpointCachesSettings | None = None
-    allow_origins: List[str] = ["*"]
+    allow_origins: list[str] = Field(default_factory=lambda: ["*"])
     logging: LoggingSettings = LoggingSettings()
     server: ServerSettings = ServerSettings()
 
 
-class AppConfig(AppSettings, BaseSettings):
-    class Config:
-        @classmethod
-        def customise_sources(
-            cls,
-            init_settings: SettingsSourceCallable,
-            env_settings: SettingsSourceCallable,
-            file_secret_settings: SettingsSourceCallable,
-        ) -> Tuple[SettingsSourceCallable, ...]:
-            return init_settings, env_settings, json_config_settings_source
+class AppConfig(AppSettings, BaseSettings, config=""):
+    @classmethod
+    def settings_customise_sources(
+        cls,
+        settings_cls: type[BaseSettings],
+        init_settings: PydanticBaseSettingsSource,
+        env_settings: PydanticBaseSettingsSource,
+        dotenv_settings: PydanticBaseSettingsSource,
+        file_secret_settings: PydanticBaseSettingsSource,
+    ) -> tuple[PydanticBaseSettingsSource, ...]:
+        return init_settings, env_settings, json_config_settings_source
 
 
 def get_settings() -> AppConfig:

tests/conftest.py

@@ -5,6 +5,10 @@
 from fastapi.testclient import TestClient
 
 
+if typing.TYPE_CHECKING:
+    from edge_proxy.environments import EnvironmentService
+
+
 @pytest.fixture
 def environment_1_feature_states_response_list() -> typing.List[dict]:
     return [
@@ -35,9 +39,27 @@ def environment_1_feature_states_response_list_response_with_segment_override(
     return environment_1_feature_states_response_list
 
 
+@pytest.fixture
+def environment_1_feature_states_response_list_response_with_identity_override(
+    environment_1_feature_states_response_list: list[dict[str, typing.Any]],
+) -> list[dict[str, typing.Any]]:
+    environment_1_feature_states_response_list[0]["feature_state_value"] = (
+        "identity_override"
+    )
+    environment_1_feature_states_response_list[0]["enabled"] = True
+    return environment_1_feature_states_response_list
+
+
 @pytest.fixture(autouse=True)
 def skip_json_config_settings_source(mocker: MockerFixture) -> None:
-    mocker.patch("edge_proxy.settings.json_config_settings_source", lambda _: {})
+    mocker.patch("edge_proxy.settings.json_config_settings_source", dict)
+
+
+@pytest.fixture
+def environment_service() -> "EnvironmentService":
+    from edge_proxy.server import environment_service
+
+    return environment_service
 
 
 @pytest.fixture