fix: /api/v1/projects/{project_pk}/segments/ is broken for compressed environments (#6913)

Author: khvn26

api/environments/dynamodb/migrator.py

@@ -108,8 +108,6 @@ def migrate(self):  # type: ignore[no-untyped-def]
         api_keys = EnvironmentAPIKey.objects.filter(environment__project_id=project_id)
         api_key_wrapper.write_api_keys(api_keys)
 
-        identity_wrapper = DynamoIdentityWrapper(
-            environment_wrapper=environment_wrapper
-        )
+        identity_wrapper = DynamoIdentityWrapper()
         identity_wrapper.write_identities(self.iter_identities_in_chunks(project_id))
         self.project_metadata.finish_identity_migration()  # type: ignore[no-untyped-call]

api/environments/dynamodb/wrappers/identity_wrapper.py

@@ -13,15 +13,13 @@
 from environments.dynamodb.constants import IDENTITIES_PAGINATION_LIMIT
 from environments.dynamodb.wrappers.exceptions import CapacityBudgetExceeded
 from util.engine_models.context.mappers import (
-    get_context_segments,
+    is_context_in_segment,
     map_environment_identity_to_context,
 )
-from util.engine_models.environments.models import EnvironmentModel
 from util.engine_models.identities.models import IdentityModel
 from util.mappers import map_identity_to_identity_document
 
 from .base import BaseDynamoWrapper
-from .environment_wrapper import DynamoEnvironmentWrapper
 
 if typing.TYPE_CHECKING:
     from boto3.dynamodb.conditions import ConditionBase
@@ -37,12 +35,8 @@
 
 
 class DynamoIdentityWrapper(BaseDynamoWrapper):
-    def __init__(
-        self,
-        environment_wrapper: DynamoEnvironmentWrapper | None = None,
-    ) -> None:
+    def __init__(self) -> None:
         super().__init__()
-        self._environment_wrapper = environment_wrapper or DynamoEnvironmentWrapper()
 
     def get_table_name(self) -> str | None:  # type: ignore[override]
         return settings.IDENTITIES_TABLE_NAME_DYNAMO
@@ -188,22 +182,29 @@ def get_segment_ids(
         identity_pk: str = None,  # type: ignore[assignment]
         identity_model: IdentityModel = None,  # type: ignore[assignment]
     ) -> list:  # type: ignore[type-arg]
+        from environments.models import Environment
+        from util.mappers.engine import map_segment_to_engine
+
         if not (identity_pk or identity_model):
             raise ValueError("Must provide one of identity_pk or identity_model.")
 
         with suppress(ObjectDoesNotExist):
             identity = identity_model or IdentityModel.model_validate(
                 self.get_item_from_uuid(identity_pk)
             )
-            environment = EnvironmentModel.model_validate(
-                self._environment_wrapper.get_item(identity.environment_api_key)
+            environment = Environment.objects.select_related("project").get(
+                api_key=identity.environment_api_key,
             )
+            segments = environment.project.get_segments_from_cache()
             context = map_environment_identity_to_context(
                 environment=environment,
                 identity=identity,
                 override_traits=None,
             )
-            segments = get_context_segments(context, environment.project.segments)
-            return [segment.id for segment in segments]
+            return [
+                segment.id
+                for segment in segments
+                if is_context_in_segment(context, map_segment_to_engine(segment))
+            ]
 
         return []

api/environments/tasks.py

@@ -50,7 +50,7 @@ def delete_environment_from_dynamo(api_key: str, environment_id: str):  # type:
     environment_wrapper.delete_environment(api_key)
 
     # Delete identities
-    identity_wrapper = DynamoIdentityWrapper(environment_wrapper=environment_wrapper)
+    identity_wrapper = DynamoIdentityWrapper()
     identity_wrapper.delete_all_identities(api_key)
 
     # Delete environment_v2 documents

api/tests/integration/edge_api/identities/test_edge_identity_featurestates_viewset.py

@@ -22,7 +22,6 @@
 )
 from environments.dynamodb import (
     DynamoEnvironmentV2Wrapper,
-    DynamoEnvironmentWrapper,
     DynamoIdentityWrapper,
 )
 from environments.models import Environment
@@ -1128,7 +1127,6 @@ def test_edge_identity_clone_flag_states_from(
     flagsmith_identities_table: Table,
     dynamodb_identity_wrapper: DynamoIdentityWrapper,
     dynamodb_wrapper_v2: DynamoEnvironmentV2Wrapper,
-    dynamo_environment_wrapper: DynamoEnvironmentWrapper,
 ) -> None:
     mocker.patch(
         "environments.dynamodb.services.DynamoIdentityWrapper",
@@ -1140,11 +1138,6 @@ def test_edge_identity_clone_flag_states_from(
         autospec=True,
         return_value=dynamodb_wrapper_v2,
     )
-    mocker.patch(
-        "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper",
-        autospec=True,
-        return_value=dynamo_environment_wrapper,
-    )
 
     def create_identity(identifier: str) -> EdgeIdentity:
         identity_model = IdentityModel(

api/tests/unit/environments/dynamodb/test_unit_migrator.py

@@ -43,9 +43,7 @@ def test_migrate_calls_internal_methods_with_correct_arguments(  # type: ignore[
     identity_migrator.migrate()  # type: ignore[no-untyped-call]
 
     # Then
-    mocked_identity_wrapper.assert_called_with(
-        environment_wrapper=mocked_environment_wrapper.return_value
-    )
+    mocked_identity_wrapper.assert_called_once_with()
 
     args, kwargs = mocked_identity_wrapper.return_value.write_identities.call_args
     assert kwargs == {}

api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_identity_wrapper.py

@@ -3,9 +3,11 @@
 
 import pytest
 from boto3.dynamodb.conditions import Key
+from boto3.dynamodb.types import Binary
 from django.core.exceptions import ObjectDoesNotExist
 from flag_engine.segments.constants import IN
 from mypy_boto3_dynamodb.service_resource import Table
+from pytest_django.fixtures import SettingsWrapper
 from pytest_mock import MockerFixture
 from rest_framework.exceptions import NotFound
 
@@ -27,7 +29,7 @@
 from segments.models import Condition, Segment, SegmentRule
 from util.engine_models.identities.models import IdentityModel
 from util.mappers import (
-    map_environment_to_environment_document,
+    map_environment_to_compressed_environment_document,
     map_identity_to_identity_document,
 )
 
@@ -313,12 +315,6 @@ def test_get_segment_ids_returns_correct_segment_ids(  # type: ignore[no-untyped
 
     identity_document = map_identity_to_identity_document(identity)
     identity_uuid = identity_document["identity_uuid"]
-    environment_document = map_environment_to_environment_document(environment)
-
-    mocked_environment_wrapper = mocker.patch(
-        "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper"
-    )
-    mocked_environment_wrapper.return_value.get_item.return_value = environment_document
 
     dynamo_identity_wrapper = DynamoIdentityWrapper()
     mocked_get_item_from_uuid = mocker.patch.object(
@@ -331,9 +327,6 @@ def test_get_segment_ids_returns_correct_segment_ids(  # type: ignore[no-untyped
     # Then
     assert segment_ids == [identity_matching_segment.id]
     mocked_get_item_from_uuid.assert_called_with(identity_uuid)
-    mocked_environment_wrapper.return_value.get_item.assert_called_with(
-        environment.api_key
-    )
 
 
 def test_get_segment_ids_with_segment_feature_overrides(
@@ -388,12 +381,6 @@ def test_get_segment_ids_with_segment_feature_overrides(
 
     identity_document = map_identity_to_identity_document(identity)
     identity_uuid = identity_document["identity_uuid"]
-    environment_document = map_environment_to_environment_document(environment)
-
-    mocked_environment_wrapper = mocker.patch(
-        "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper"
-    )
-    mocked_environment_wrapper.return_value.get_item.return_value = environment_document
 
     dynamo_identity_wrapper = DynamoIdentityWrapper()
     mocker.patch.object(
@@ -430,12 +417,6 @@ def test_get_segment_ids_returns_segment_using_in_operator_for_integer_traits(
 
     identity_document = map_identity_to_identity_document(identity)
     identity_uuid = identity_document["identity_uuid"]
-    environment_document = map_environment_to_environment_document(environment)
-
-    mocked_environment_wrapper = mocker.patch(
-        "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper"
-    )
-    mocked_environment_wrapper.return_value.get_item.return_value = environment_document
 
     dynamo_identity_wrapper = DynamoIdentityWrapper()
     mocker.patch.object(
@@ -495,12 +476,6 @@ def test_get_segment_ids_with_identity_model(identity, environment, mocker):  #
     # Given
     identity_document = map_identity_to_identity_document(identity)
     identity_model = IdentityModel.parse_obj(identity_document)
-    environment_document = map_environment_to_environment_document(environment)
-
-    mocked_environment_wrapper = mocker.patch(
-        "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper"
-    )
-    mocked_environment_wrapper.return_value.get_item.return_value = environment_document
 
     dynamo_identity_wrapper = DynamoIdentityWrapper()
     mocker.patch.object(
@@ -514,6 +489,47 @@ def test_get_segment_ids_with_identity_model(identity, environment, mocker):  #
     assert segment_ids == []
 
 
+def test_get_segment_ids__compressed_environment_in_dynamo__returns_correct_segment_ids(
+    identity: "Identity",
+    identity_matching_segment: "Segment",
+    dynamodb_identity_wrapper: DynamoIdentityWrapper,
+    flagsmith_identities_table: Table,
+    flagsmith_environment_table: Table,
+    settings: "SettingsWrapper",
+) -> None:
+    """Regression test for https://github.com/Flagsmith/flagsmith/issues/6912
+
+    Previously, get_segment_ids read the environment document from DynamoDB
+    and failed with a ValidationError when the document contained compressed
+    (gzipped Binary) `project` and `feature_states` fields.
+    """
+    # Given - identity written to DynamoDB
+    identity_document = map_identity_to_identity_document(identity)
+    flagsmith_identities_table.put_item(Item=identity_document)
+    identity_uuid = str(identity_document["identity_uuid"])
+
+    # And - a compressed environment document in DynamoDB
+    settings.ENVIRONMENTS_TABLE_NAME_DYNAMO = flagsmith_environment_table.name
+    compressed_result = map_environment_to_compressed_environment_document(
+        identity.environment,
+    )
+    flagsmith_environment_table.put_item(Item=compressed_result.document)
+
+    # Verify the document actually has compressed Binary fields
+    stored = flagsmith_environment_table.get_item(
+        Key={"api_key": identity.environment.api_key},
+    )["Item"]
+    assert stored.get("compressed") is True
+    assert isinstance(stored["project"], Binary)
+    assert isinstance(stored["feature_states"], Binary)
+
+    # When
+    segment_ids = dynamodb_identity_wrapper.get_segment_ids(identity_uuid)
+
+    # Then
+    assert segment_ids == [identity_matching_segment.id]
+
+
 def test_identity_wrapper__iter_all_items_paginated__returns_expected(
     identity: "Identity",
     mocker: "MockerFixture",
@@ -643,41 +659,6 @@ def test_identity_wrapper__iter_all_items_paginated__capacity_budget_set__raises
     )
 
 
-def test_get_segment_ids__called_multiple_times__reuses_environment_wrapper(
-    project: "Project",
-    environment: "Environment",
-    identity: "Identity",
-    mocker: "MockerFixture",
-) -> None:
-    # Given
-    identity_document = map_identity_to_identity_document(identity)
-    environment_document = map_environment_to_environment_document(environment)
-
-    mocked_environment_wrapper_class = mocker.patch(
-        "environments.dynamodb.wrappers.identity_wrapper.DynamoEnvironmentWrapper"
-    )
-    mocked_environment_wrapper_class.return_value.get_item.return_value = (
-        environment_document
-    )
-
-    dynamo_identity_wrapper = DynamoIdentityWrapper()
-    mocker.patch.object(
-        dynamo_identity_wrapper,
-        "get_item_from_uuid",
-        return_value=identity_document,
-    )
-    identity_uuid = str(identity_document["identity_uuid"])
-
-    # When
-    dynamo_identity_wrapper.get_segment_ids(identity_uuid)
-    dynamo_identity_wrapper.get_segment_ids(identity_uuid)
-    dynamo_identity_wrapper.get_segment_ids(identity_uuid)
-
-    # Then - DynamoEnvironmentWrapper should be instantiated once
-    # (during DynamoIdentityWrapper.__init__), not once per get_segment_ids call.
-    assert mocked_environment_wrapper_class.call_count == 1
-
-
 def test_delete_all_identities__deletes_all_identities_documents_from_dynamodb(
     flagsmith_identities_table: Table,
     dynamodb_identity_wrapper: DynamoIdentityWrapper,

api/util/engine_models/context/mappers.py

@@ -6,7 +6,6 @@
 """
 
 import typing
-from typing import Protocol
 
 from flag_engine.context.types import (
     EvaluationContext,
@@ -23,14 +22,12 @@
 from util.engine_models.identities.traits.models import TraitModel
 from util.engine_models.segments.models import SegmentModel, SegmentRuleModel
 
-
-class EnvironmentProtocol(Protocol):
-    api_key: str
-    name: str | None
+if typing.TYPE_CHECKING:
+    from environments.models import Environment
 
 
 def map_environment_identity_to_context(
-    environment: EnvironmentProtocol,
+    environment: "Environment",
     identity: IdentityModel,
     override_traits: typing.Optional[typing.List[TraitModel]],
 ) -> EvaluationContext:
@@ -40,8 +37,7 @@ def map_environment_identity_to_context(
     Vendored from flagsmith-flag-engine's fix/missing-export branch and adapted
     to return v10's EvaluationContext TypedDict.
 
-    :param environment: Any object with `api_key` and `name` attributes
-        (e.g. Django Environment or Pydantic EnvironmentModel).
+    :param environment: An Environment object.
     :param identity: The identity model object (Pydantic IdentityModel).
     :param override_traits: A list of TraitModel objects, to be used in place of
         `identity.identity_traits` if provided.
@@ -184,19 +180,3 @@ def is_context_in_segment(
 
     segment_context = map_segment_to_segment_context(segment)
     return v10_is_context_in_segment(context, segment_context)
-
-
-def get_context_segments(
-    context: EvaluationContext,
-    segments: typing.List[SegmentModel],
-) -> typing.List[SegmentModel]:
-    """
-    Get the list of segments that match a given evaluation context.
-
-    This is a compatibility function for code that expects the old API.
-
-    :param context: The EvaluationContext (TypedDict).
-    :param segments: List of SegmentModel objects to check.
-    :return: List of matching SegmentModel objects.
-    """
-    return [segment for segment in segments if is_context_in_segment(context, segment)]