feat: Clone identity flag states (#3773)

api/conftest.py

@@ -704,6 +704,7 @@ def flagsmith_identities_table(dynamodb: DynamoDBServiceResource) -> Table:
             {"AttributeName": "composite_key", "AttributeType": "S"},
             {"AttributeName": "environment_api_key", "AttributeType": "S"},
             {"AttributeName": "identifier", "AttributeType": "S"},
+            {"AttributeName": "identity_uuid", "AttributeType": "S"},
         ],
         GlobalSecondaryIndexes=[
             {
@@ -713,7 +714,12 @@ def flagsmith_identities_table(dynamodb: DynamoDBServiceResource) -> Table:
                     {"AttributeName": "identifier", "KeyType": "RANGE"},
                 ],
                 "Projection": {"ProjectionType": "ALL"},
-            }
+            },
+            {
+                "IndexName": "identity_uuid-index",
+                "KeySchema": [{"AttributeName": "identity_uuid", "KeyType": "HASH"}],
+                "Projection": {"ProjectionType": "ALL"},
+            },
         ],
         BillingMode="PAY_PER_REQUEST",
     )

api/edge_api/identities/models.py

@@ -241,3 +241,20 @@ def _get_changes(self) -> IdentityChangeset:
 
     def _reset_initial_state(self):
         self._initial_state = copy.deepcopy(self)
+
+    def clone_flag_states_from(self, source_identity: "EdgeIdentity") -> None:
+        """
+        Clone the feature states from the source identity to the target identity.
+        """
+        # Delete identity_target's feature states
+        for feature_state in list(self.feature_overrides):
+            self.remove_feature_override(feature_state=feature_state)
+
+        # Clone identity_source's feature states to identity_target
+        for feature_in_source in source_identity.feature_overrides:
+            feature_state_target = FeatureStateModel(
+                feature=feature_in_source.feature,
+                feature_state_value=feature_in_source.feature_state_value,
+                enabled=feature_in_source.enabled,
+            )
+            self.add_feature_override(feature_state_target)

api/edge_api/identities/serializers.py

@@ -264,3 +264,10 @@ def to_representation(self, instance: IdentityOverrideV2):
 
 class GetEdgeIdentityOverridesSerializer(serializers.Serializer):
     results = GetEdgeIdentityOverridesResultSerializer(many=True)
+
+
+class EdgeIdentitySourceIdentityRequestSerializer(serializers.Serializer):
+    source_identity_uuid = serializers.CharField(
+        required=True,
+        help_text="UUID of the source identity to clone feature states from.",
+    )

api/edge_api/identities/views.py

@@ -7,7 +7,7 @@
 from django.shortcuts import get_object_or_404
 from django.utils.decorators import method_decorator
 from drf_yasg.utils import swagger_auto_schema
-from flag_engine.identities.models import IdentityModel
+from flag_engine.identities.models import IdentityFeaturesList, IdentityModel
 from flag_engine.identities.traits.models import TraitModel
 from pyngo import drf_error_details
 from rest_framework import status, viewsets
@@ -38,6 +38,7 @@
     EdgeIdentityFsQueryparamSerializer,
     EdgeIdentityIdentifierSerializer,
     EdgeIdentitySerializer,
+    EdgeIdentitySourceIdentityRequestSerializer,
     EdgeIdentityTraitsSerializer,
     EdgeIdentityWithIdentifierFeatureStateDeleteRequestBody,
     EdgeIdentityWithIdentifierFeatureStateRequestBody,
@@ -202,7 +203,6 @@ class EdgeIdentityFeatureStateViewSet(viewsets.ModelViewSet):
     lookup_field = "featurestate_uuid"
 
     serializer_class = EdgeIdentityFeatureStateSerializer
-
     # Patch is not supported
     http_method_names = [
         "get",
@@ -215,10 +215,9 @@ class EdgeIdentityFeatureStateViewSet(viewsets.ModelViewSet):
     ]
     pagination_class = None
 
-    def initial(self, request, *args, **kwargs):
-        super().initial(request, *args, **kwargs)
+    def get_identity(self, edge_identity_identity_uuid: str) -> EdgeIdentity:
         identity_document = EdgeIdentity.dynamo_wrapper.get_item_from_uuid_or_404(
-            self.kwargs["edge_identity_identity_uuid"]
+            edge_identity_identity_uuid
         )
 
         if (
@@ -233,8 +232,15 @@ def initial(self, request, *args, **kwargs):
                 environment__api_key=identity.environment_api_key
             ).values_list("feature__name", flat=True)
         )
-        identity.synchronise_features(valid_feature_names)
-        self.identity = identity
+        identity.synchronise_features(valid_feature_names=valid_feature_names)
+
+        return identity
+
+    def initial(self, request, *args, **kwargs):
+        super().initial(request, *args, **kwargs)
+        self.identity: EdgeIdentity = self.get_identity(
+            edge_identity_identity_uuid=self.kwargs["edge_identity_identity_uuid"]
+        )
 
     def get_object(self):
         feature_state = self.identity.get_feature_state_by_featurestate_uuid(
@@ -261,7 +267,7 @@ def list(self, request, *args, **kwargs):
         )
         q_params_serializer.is_valid(raise_exception=True)
 
-        identity_features = self.identity.feature_overrides
+        identity_features: IdentityFeaturesList = self.identity.feature_overrides
 
         feature = q_params_serializer.data.get("feature")
         if feature:
@@ -300,6 +306,35 @@ def all(self, request, *args, **kwargs):
 
         return Response(serializer.data)
 
+    @swagger_auto_schema(
+        request_body=EdgeIdentitySourceIdentityRequestSerializer(),
+        responses={200: IdentityAllFeatureStatesSerializer(many=True)},
+    )
+    @action(detail=False, methods=["POST"], url_path="clone-from-given-identity")
+    def clone_from_given_identity(self, request, *args, **kwargs) -> Response:
+        """
+        Clone feature states from a given source identity.
+        """
+        # Get and validate source identity
+        serializer = EdgeIdentitySourceIdentityRequestSerializer(
+            data=request.data, context={"request": request}
+        )
+        serializer.is_valid(raise_exception=True)
+
+        source_identity: EdgeIdentity = self.get_identity(
+            edge_identity_identity_uuid=serializer.validated_data[
+                "source_identity_uuid"
+            ]
+        )
+
+        self.identity.clone_flag_states_from(source_identity)
+        self.identity.save(
+            user=request.user.id,
+            master_api_key=getattr(request, "master_api_key", None),
+        )
+
+        return self.all(request, *args, **kwargs)
+
 
 class EdgeIdentityWithIdentifierFeatureStateView(APIView):
     permission_classes = [IsAuthenticated, EdgeIdentityWithIdentifierViewPermissions]

api/environments/dynamodb/wrappers/environment_wrapper.py

@@ -38,7 +38,8 @@ def write_environments(self, environments: Iterable["Environment"]) -> None:
 
 
 class DynamoEnvironmentWrapper(BaseDynamoEnvironmentWrapper):
-    table_name = settings.ENVIRONMENTS_TABLE_NAME_DYNAMO
+    def get_table_name(self) -> str | None:
+        return settings.ENVIRONMENTS_TABLE_NAME_DYNAMO
 
     def write_environments(self, environments: Iterable["Environment"]):
         with self.table.batch_writer() as writer:

api/environments/identities/models.py

@@ -141,6 +141,15 @@ def get_all_feature_states(
 
         return list(identity_flags.values())
 
+    def get_overridden_feature_states(self) -> dict[int, FeatureState]:
+        """
+        Get all overridden feature states for an identity.
+
+        :return: dict[int, FeatureState] - Key: feature ID. Value: Overridden feature_state.
+        """
+
+        return {fs.feature_id: fs for fs in self.identity_features.all()}
+
     def get_segments(
         self, traits: typing.List[Trait] = None, overrides_only: bool = False
     ) -> typing.List[Segment]:

api/environments/identities/serializers.py

@@ -140,3 +140,10 @@ def get_segment(self, instance) -> typing.Optional[typing.Dict[str, typing.Any]]
                 instance=instance.feature_segment.segment
             ).data
         return None
+
+
+class IdentitySourceIdentityRequestSerializer(serializers.Serializer):
+    source_identity_id = serializers.IntegerField(
+        required=True,
+        help_text="ID of the source identity to clone feature states from.",
+    )

api/features/models.py

@@ -947,6 +947,48 @@ def _is_more_recent_version(self, other: "FeatureState") -> bool:
             and self.version > other.version
         ) or (self.version is not None and other.version is None)
 
+    @staticmethod
+    def copy_identity_feature_states(
+        target_identity: "Identity", source_identity: "Identity"
+    ) -> None:
+        target_feature_states: dict[int, FeatureState] = (
+            target_identity.get_overridden_feature_states()
+        )
+        source_feature_states: dict[int, FeatureState] = (
+            source_identity.get_overridden_feature_states()
+        )
+
+        # Delete own feature states not in source_identity
+        feature_states_to_delete = list(
+            target_feature_states.keys() - source_feature_states.keys()
+        )
+        for feature_state_id in feature_states_to_delete:
+            target_feature_states[feature_state_id].delete()
+
+        # Clone source_identity's feature states to target_identity
+        for source_feature_id, source_feature_state in source_feature_states.items():
+            # Get target feature_state if exists in target identity or create new one
+            target_feature_state: FeatureState = target_feature_states.get(
+                source_feature_id
+            ) or FeatureState.objects.create(
+                environment=target_identity.environment,
+                identity=target_identity,
+                feature=source_feature_state.feature,
+            )
+
+            # Copy enabled value from source feature_state
+            target_feature_state.enabled = source_feature_states[
+                source_feature_id
+            ].enabled
+
+            # Copy feature state value from source feature_state
+            target_feature_state.feature_state_value.copy_from(
+                source_feature_state.feature_state_value
+            )
+
+            # Save changes to target feature_state
+            target_feature_state.save()
+
 
 class FeatureStateValue(
     AbstractBaseFeatureValueModel,
@@ -972,6 +1014,14 @@ def clone(self, feature_state: FeatureState) -> "FeatureStateValue":
         clone.save()
         return clone
 
+    def copy_from(self, source_feature_state_value: "FeatureStateValue"):
+        # Copy feature state type and values from given feature state value.
+        self.type = source_feature_state_value.type
+        self.boolean_value = source_feature_state_value.boolean_value
+        self.integer_value = source_feature_state_value.integer_value
+        self.string_value = source_feature_state_value.string_value
+        self.save()
+
     def get_update_log_message(self, history_instance) -> typing.Optional[str]:
         fs = self.feature_state
 

api/features/views.py

@@ -27,6 +27,7 @@
 from environments.identities.models import Identity
 from environments.identities.serializers import (
     IdentityAllFeatureStatesSerializer,
+    IdentitySourceIdentityRequestSerializer,
 )
 from environments.models import Environment
 from environments.permissions.permissions import (
@@ -661,6 +662,34 @@ def all(self, request, *args, **kwargs):
 
         return Response(serializer.data)
 
+    @swagger_auto_schema(
+        request_body=IdentitySourceIdentityRequestSerializer(),
+        responses={200: IdentityAllFeatureStatesSerializer(many=True)},
+    )
+    @action(methods=["POST"], detail=False, url_path="clone-from-given-identity")
+    def clone_from_given_identity(self, request, *args, **kwargs) -> Response:
+        """
+        Clone feature states from a given source identity.
+        """
+        serializer = IdentitySourceIdentityRequestSerializer(
+            data=request.data, context={"request": request}
+        )
+        serializer.is_valid(raise_exception=True)
+        # Get and validate source and target identities
+        target_identity = get_object_or_404(
+            queryset=Identity, pk=self.kwargs["identity_pk"]
+        )
+        source_identity = get_object_or_404(
+            queryset=Identity, pk=request.data.get("source_identity_id")
+        )
+
+        # Clone feature states
+        FeatureState.copy_identity_feature_states(
+            target_identity=target_identity, source_identity=source_identity
+        )
+
+        return self.all(request, *args, **kwargs)
+
 
 @method_decorator(
     name="list",

api/tests/conftest.py

@@ -90,7 +90,22 @@ def dynamo_enabled_project_environment_one_document(
 @pytest.fixture()
 def dynamo_environment_wrapper(
     flagsmith_environment_table: Table,
+    settings: SettingsWrapper,
 ) -> DynamoEnvironmentWrapper:
+    settings.ENVIRONMENTS_TABLE_NAME_DYNAMO = flagsmith_environment_table.name
     wrapper = DynamoEnvironmentWrapper()
     wrapper.table_name = flagsmith_environment_table.name
     return wrapper
+
+
+@pytest.fixture()
+def app_settings_for_dynamodb(
+    settings: SettingsWrapper,
+    flagsmith_environment_table: Table,
+    flagsmith_environments_v2_table: Table,
+    flagsmith_identities_table: Table,
+) -> None:
+    settings.ENVIRONMENTS_TABLE_NAME_DYNAMO = flagsmith_environment_table.name
+    settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO = flagsmith_environments_v2_table.name
+    settings.IDENTITIES_TABLE_NAME_DYNAMO = flagsmith_identities_table.name
+    return

api/tests/integration/conftest.py

@@ -100,7 +100,10 @@ def environment(
 
 @pytest.fixture()
 def dynamo_enabled_environment(
-    admin_client, dynamo_enabled_project, environment_api_key
+    admin_client: APIClient,
+    dynamo_enabled_project: int,
+    environment_api_key: str,
+    app_settings_for_dynamodb: None,
 ) -> int:
     environment_data = {
         "name": "Test Environment",