perf(feature_flags): validate schema once per fetched document (#8431)

* perf(feature_flags): validate schema once per fetched document

FeatureFlags.get_configuration built a SchemaValidator and walked the
whole document on every call. evaluate and get_enabled_features both call
it, so a handler evaluating five flags validated the full document five
times per invocation, even when the store served it from cache.

Skip validation when the store returns the same dict object that was
last validated. The Parameters cache hands back the same object until
expiry, so identity is a reliable signal for a cache hit, and a fresh
document is always validated. A strong reference to the last validated
document is kept so its id() cannot be recycled.

AppConfigStore with an envelope produced a new object per call from the
JMESPath query, which would defeat the check. Memoise the extraction on
the raw document's identity so envelope users benefit as well.

Closes #8426

* docs(feature_flags): document configuration immutability

---------

Co-authored-by: Leandro <lcdama@amazon.pt>

Author: dreamorosi

aws_lambda_powertools/utilities/feature_flags/appconfig.py

@@ -86,6 +86,10 @@ def __init__(
             boto3_client=boto3_client,
             boto3_session=boto3_session,
         )
+        # Memoised envelope extraction: (raw document, extracted config). The Parameters cache hands back the
+        # same raw dict until expiry, so we can reuse the extracted result rather than re-running the JMESPath
+        # query and producing a new object on every call.
+        self._last_extracted: tuple[dict[str, Any], dict[str, Any]] | None = None
 
         # Override the user agent to use "feature_flags" instead of "parameters"
         self._register_feature_flags_user_agent()
@@ -140,11 +144,17 @@ def get_configuration(self) -> dict[str, Any]:
         config = self.get_raw_configuration
 
         if self.envelope:
+            if self._last_extracted is not None and self._last_extracted[0] is config:
+                self.logger.debug("Envelope enabled; reusing previously extracted config for cached document")
+                return self._last_extracted[1]
+
             self.logger.debug("Envelope enabled; extracting data from config", extra={"envelope": self.envelope})
-            config = jmespath_utils.query(
+            extracted = jmespath_utils.query(
                 data=config,
                 envelope=self.envelope,
                 jmespath_options=self.jmespath_options,
             )
+            self._last_extracted = (config, extracted)
+            return extracted
 
         return config

aws_lambda_powertools/utilities/feature_flags/feature_flags.py

@@ -81,6 +81,9 @@ def __init__(self, store: StoreProvider, logger: logging.Logger | Logger | None
         self.store = store
         self.logger = logger or logging.getLogger(__name__)
         self._exception_handlers: dict[Exception, Callable] = {}
+        # Last document that passed schema validation. We keep a strong reference so its id() can't be
+        # recycled by a different object, which lets us safely skip re-validation on store cache hits.
+        self._last_validated_config: dict | None = None
 
     def _match_by_action(self, action: str, condition_value: Any, context_value: Any) -> bool:
         try:
@@ -210,8 +213,17 @@ def get_configuration(self) -> dict:
         # parse result conf as JSON, keep in cache for max age defined in store
         self.logger.debug(f"Fetching schema from registered store, store={self.store}")
         config: dict = self.store.get_configuration()
+
+        # Stores that serve from cache (e.g. AppConfigStore via Parameters) return the same dict object until
+        # expiry, so identity is a reliable signal that we've already validated this exact document.
+        # A store that applies an envelope returns a fresh object each time and will still be validated.
+        if config is self._last_validated_config:
+            self.logger.debug("Schema already validated, skipping validation")
+            return config
+
         validator = schema.SchemaValidator(schema=config, logger=self.logger)
         validator.validate()
+        self._last_validated_config = config
 
         return config
 

docs/utilities/feature_flags.md

@@ -369,6 +369,9 @@ You can override `max_age` parameter when instantiating the store.
 
 You can access the configuration fetched from the store via `get_raw_configuration` property within the store instance.
 
+???+ warning
+    Treat the returned configuration as read-only. If you need to modify it, create a deep copy first.
+
 === "getting_stored_features.py"
 
     ```python hl_lines="9"
@@ -555,6 +558,10 @@ You can create your own custom FeatureFlags store provider by inheriting the `St
 * **`get_raw_configuration()`** – get the raw configuration from the store provider and return the parsed JSON dictionary
 * **`get_configuration()`** – get the configuration from the store provider, parsing it as a JSON dictionary. If an envelope is set, extract the envelope data
 
+Feature Flags can reuse a previously validated configuration when a store returns the same dictionary. Custom store
+providers must return a new dictionary when the configuration changes instead of modifying a previously returned
+dictionary in place.
+
 Here are an example of implementing a custom store provider using Amazon S3, a popular object storage.
 
 ???+ note

tests/functional/feature_flags/_boto3/test_feature_flags.py

@@ -1702,3 +1702,70 @@ def catch_exception(exc):
             context={"tenant_id": "not a list value"},
             default=False,
         )
+
+
+# Test schema validation is performed once per fetched document (#8426)
+def test_schema_validated_once_for_cached_document(mocker, config):
+    # GIVEN a store that serves the same document object on every call (e.g. a Parameters cache hit)
+    mocked_app_config_schema = {"my_feature": {"default": True}}
+    feature_flags = init_feature_flags(mocker, mocked_app_config_schema, config)
+    validate = mocker.spy(schema.SchemaValidator, "validate")
+
+    # WHEN evaluating several flags in the same invocation
+    for _ in range(5):
+        assert feature_flags.evaluate(name="my_feature", context={}, default=False) is True
+    feature_flags.get_enabled_features(context={})
+
+    # THEN the schema is validated only once
+    assert validate.call_count == 1
+
+
+def test_schema_revalidated_when_store_returns_new_document(mocker, config):
+    # GIVEN a store that returns a fresh document object on each call (e.g. after cache expiry)
+    first = {"my_feature": {"default": True}}
+    second = {"my_feature": {"default": False}}
+    store = init_fetcher_side_effect(mocker, config, side_effect=[first, second, second])
+    feature_flags = FeatureFlags(store=store)
+    validate = mocker.spy(schema.SchemaValidator, "validate")
+
+    # WHEN evaluating across a document change, then again on the same document
+    assert feature_flags.evaluate(name="my_feature", context={}, default=False) is True
+    assert feature_flags.evaluate(name="my_feature", context={}, default=False) is False
+    assert feature_flags.evaluate(name="my_feature", context={}, default=False) is False
+
+    # THEN each distinct document is validated exactly once
+    assert validate.call_count == 2
+
+
+def test_schema_invalid_document_is_never_cached_as_validated(mocker, config):
+    # GIVEN a store that first returns an invalid document, then a valid one
+    invalid = {"my_feature": {"default": "not a bool"}}
+    valid = {"my_feature": {"default": True}}
+    store = init_fetcher_side_effect(mocker, config, side_effect=[invalid, invalid, valid])
+    feature_flags = FeatureFlags(store=store)
+
+    # WHEN the invalid document is served twice
+    # THEN validation fails both times rather than being skipped after the first failure
+    with pytest.raises(schema.SchemaValidationError):
+        feature_flags.evaluate(name="my_feature", context={}, default=False)
+    with pytest.raises(schema.SchemaValidationError):
+        feature_flags.evaluate(name="my_feature", context={}, default=False)
+
+    # AND a subsequent valid document is validated and evaluated normally
+    assert feature_flags.evaluate(name="my_feature", context={}, default=False) is True
+
+
+def test_envelope_extraction_reused_for_cached_document(mocker, config):
+    # GIVEN a store with an envelope, served from cache
+    mocked_app_config_schema = {"app": {"features": {"my_feature": {"default": True}}}}
+    feature_flags = init_feature_flags(mocker, mocked_app_config_schema, config, envelope="app.features")
+    validate = mocker.spy(schema.SchemaValidator, "validate")
+
+    # WHEN evaluating several times
+    first = feature_flags.get_configuration()
+    for _ in range(3):
+        assert feature_flags.evaluate(name="my_feature", context={}, default=False) is True
+
+    # THEN the extracted document is the same object each time and validated only once
+    assert feature_flags.get_configuration() is first
+    assert validate.call_count == 1