From f72c764e59d44f3c50bafd0cd2aef2dcf51af07b Mon Sep 17 00:00:00 2001 From: Gagan Date: Wed, 30 Oct 2024 16:39:37 +0530 Subject: [PATCH] feat(export): Add support for edge identities data (#4654) Co-authored-by: Matthew Elwell --- api/edge_api/identities/export.py | 209 ++++++++++++++ api/import_export/export.py | 53 +++- api/poetry.lock | 6 +- api/pyproject.toml | 2 +- .../test_unit_import_export_export.py | 272 ++++++++++++++++++ 5 files changed, 531 insertions(+), 11 deletions(-) create mode 100644 api/edge_api/identities/export.py diff --git a/api/edge_api/identities/export.py b/api/edge_api/identities/export.py new file mode 100644 index 000000000000..ebfebb85f498 --- /dev/null +++ b/api/edge_api/identities/export.py @@ -0,0 +1,209 @@ +import logging +import typing +import uuid +from decimal import Decimal + +from django.utils import timezone +from flag_engine.identities.traits.types import map_any_value_to_trait_value + +from edge_api.identities.models import EdgeIdentity +from environments.identities.traits.models import Trait +from features.models import Feature, FeatureState +from features.multivariate.models import MultivariateFeatureOption + +EXPORT_EDGE_IDENTITY_PAGINATION_LIMIT = 20000 + +logger = logging.getLogger() + + +def export_edge_identity_and_overrides( # noqa: C901 + environment_api_key: str, +) -> tuple[list, list, list]: + kwargs = { + "environment_api_key": environment_api_key, + "limit": EXPORT_EDGE_IDENTITY_PAGINATION_LIMIT, + } + identity_export = [] + traits_export = [] + identity_override_export = [] + + feature_id_to_uuid: dict[int, str] = get_feature_uuid_cache(environment_api_key) + mv_feature_option_id_to_uuid: dict[int, str] = get_mv_feature_option_uuid_cache( + environment_api_key + ) + while True: + response = EdgeIdentity.dynamo_wrapper.get_all_items(**kwargs) + for item in response["Items"]: + identifier = item["identifier"] + # export identity + identity_export.append( + export_edge_identity( + identifier, environment_api_key, item["created_date"] + ) + ) + # export traits + for trait in item["identity_traits"]: + traits_export.append( + export_edge_trait(trait, identifier, environment_api_key) + ) + for override in item["identity_features"]: + featurestate_uuid = override["featurestate_uuid"] + feature_id = override["feature"]["id"] + if feature_id not in feature_id_to_uuid: + logging.warning("Feature with id %s does not exist", feature_id) + continue + + feature_uuid = feature_id_to_uuid[feature_id] + + # export feature state + identity_override_export.append( + export_edge_feature_state( + identifier, + environment_api_key, + featurestate_uuid, + feature_uuid, + override["enabled"], + ) + ) + featurestate_value = override["feature_state_value"] + if featurestate_value is not None: + # export feature state value + identity_override_export.append( + export_featurestate_value(featurestate_value, featurestate_uuid) + ) + if mvfsv_overrides := override["multivariate_feature_state_values"]: + for mvfsv_override in mvfsv_overrides: + mv_feature_option_id = mvfsv_override[ + "multivariate_feature_option" + ]["id"] + if mv_feature_option_id not in mv_feature_option_id_to_uuid: + logging.warning( + "MultivariateFeatureOption with id %s does not exist", + mv_feature_option_id, + ) + continue + + mv_feature_option_uuid = mv_feature_option_id_to_uuid[ + mv_feature_option_id + ] + percentage_allocation = float( + mvfsv_override["percentage_allocation"] + ) + # export mv feature state value + identity_override_export.append( + export_mv_featurestate_value( + featurestate_uuid, + mv_feature_option_uuid, + percentage_allocation, + ) + ) + if "LastEvaluatedKey" not in response: + break + kwargs["start_key"] = response["LastEvaluatedKey"] + return identity_export, traits_export, identity_override_export + + +def get_feature_uuid_cache(environment_api_key: str) -> dict[int, str]: + qs = Feature.objects.filter( + project__environments__api_key=environment_api_key + ).values_list("id", "uuid") + return {feature_id: feature_uuid for feature_id, feature_uuid in qs} + + +def get_mv_feature_option_uuid_cache(environment_api_key: str) -> dict[int, str]: + qs = MultivariateFeatureOption.objects.filter( + feature__project__environments__api_key=environment_api_key + ).values_list("id", "uuid") + return {mvfso_id: mvfso_uuid for mvfso_id, mvfso_uuid in qs} + + +def export_edge_trait(trait: dict, identifier: str, environment_api_key: str) -> dict: + trait_value = map_any_value_to_trait_value(trait["trait_value"]) + trait_value_data = Trait.generate_trait_value_data(trait_value) + return { + "model": "traits.trait", + "fields": { + "identity": [identifier, environment_api_key], + "created_date": timezone.now().isoformat(), + "trait_key": trait["trait_key"], + **trait_value_data, + }, + } + + +def export_edge_identity( + identifier: str, environment_api_key: str, created_date: str +) -> dict: + return { + "model": "identities.identity", + "fields": { + "identifier": identifier, + "created_date": created_date, + "environment": [environment_api_key], + }, + } + + +def export_edge_feature_state( + identifier: str, + environment_api_key: str, + featurestate_uuid: str, + feature_uuid: str, + enabled: bool, +) -> dict: + # NOTE: All of the datetime columns are not part of + # dynamo but are part of the django model + # hence we are setting them to current time + return { + "model": "features.featurestate", + "fields": { + "uuid": featurestate_uuid, + "created_at": timezone.now().isoformat(), + "updated_at": timezone.now().isoformat(), + "live_from": timezone.now().isoformat(), + "feature": [feature_uuid], + "environment": [environment_api_key], + "identity": [ + identifier, + environment_api_key, + ], + "feature_segment": None, + "enabled": enabled, + "version": 1, + }, + } + + +def export_featurestate_value( + featurestate_value: typing.Any, featurestate_uuid: str +) -> dict: + if isinstance(featurestate_value, Decimal): + if featurestate_value.as_tuple().exponent == 0: + featurestate_value = int(featurestate_value) + + fsv_data = FeatureState().generate_feature_state_value_data(featurestate_value) + fsv_data.pop("feature_state") + + return { + "model": "features.featurestatevalue", + "fields": { + "uuid": uuid.uuid4(), + "feature_state": [featurestate_uuid], + **fsv_data, + }, + } + + +def export_mv_featurestate_value( + featurestate_uuid: str, mv_feature_option_uuid: int, percentage_allocation: float +) -> dict: + + return { + "model": "multivariate.multivariatefeaturestatevalue", + "fields": { + "uuid": uuid.uuid4(), + "feature_state": [featurestate_uuid], + "multivariate_feature_option": [mv_feature_option_uuid], + "percentage_allocation": percentage_allocation, + }, + } diff --git a/api/import_export/export.py b/api/import_export/export.py index 85aec04a0222..fd373bbc7255 100644 --- a/api/import_export/export.py +++ b/api/import_export/export.py @@ -8,8 +8,9 @@ import boto3 from django.core import serializers from django.core.serializers.json import DjangoJSONEncoder -from django.db.models import Model, Q +from django.db.models import F, Model, Q +from edge_api.identities.export import export_edge_identity_and_overrides from environments.identities.models import Identity from environments.identities.traits.models import Trait from environments.models import Environment, EnvironmentAPIKey, Webhook @@ -76,6 +77,7 @@ def full_export(organisation_id: int) -> typing.List[dict]: *export_identities(organisation_id), *export_features(organisation_id), *export_metadata(organisation_id), + *export_edge_identities(organisation_id), ] @@ -115,13 +117,25 @@ def export_projects(organisation_id: int) -> typing.List[dict]: _EntityExportConfig(Segment, default_filter), _EntityExportConfig( SegmentRule, - Q(segment__project__organisation__id=organisation_id) - | Q(rule__segment__project__organisation__id=organisation_id), + Q( + segment__project__organisation__id=organisation_id, + segment_id=F("segment__version_of"), + ) + | Q( + rule__segment__project__organisation__id=organisation_id, + rule__segment_id=F("rule__segment__version_of"), + ), ), _EntityExportConfig( Condition, - Q(rule__segment__project__organisation__id=organisation_id) - | Q(rule__rule__segment__project__organisation__id=organisation_id), + Q( + rule__segment__project__organisation__id=organisation_id, + rule__segment_id=F("rule__segment__version_of"), + ) + | Q( + rule__rule__segment__project__organisation__id=organisation_id, + rule__rule__segment_id=F("rule__rule__segment__version_of"), + ), ), _EntityExportConfig(Tag, default_filter), _EntityExportConfig(DataDogConfiguration, default_filter), @@ -150,12 +164,20 @@ def export_identities(organisation_id: int) -> typing.List[dict]: traits = _export_entities( _EntityExportConfig( Trait, - Q(identity__environment__project__organisation__id=organisation_id), + Q( + identity__environment__project__organisation__id=organisation_id, + identity__environment__project__enable_dynamo_db=False, + ), ), ) + identities = _export_entities( _EntityExportConfig( - Identity, Q(environment__project__organisation__id=organisation_id) + Identity, + Q( + environment__project__organisation__id=organisation_id, + environment__project__enable_dynamo_db=False, + ), ), ) @@ -166,6 +188,23 @@ def export_identities(organisation_id: int) -> typing.List[dict]: return [*identities, *traits] +def export_edge_identities(organisation_id: int) -> typing.List[dict]: + identities = [] + traits = [] + identity_overrides = [] + for environment in Environment.objects.filter( + project__organisation__id=organisation_id, project__enable_dynamo_db=True + ): + exported_identities, exported_traits, exported_overrides = ( + export_edge_identity_and_overrides(environment.api_key) + ) + identities.extend(exported_identities) + traits.extend(exported_traits) + identity_overrides.extend(exported_overrides) + + return [*identities, *traits, *identity_overrides] + + def export_features(organisation_id: int) -> typing.List[dict]: """ Export all features and related entities, except ChangeRequests. diff --git a/api/poetry.lock b/api/poetry.lock index 3db2269f0a73..a9160aeebb48 100644 --- a/api/poetry.lock +++ b/api/poetry.lock @@ -1365,12 +1365,12 @@ resolved_reference = "f3809f6d592b2c6cfdfa88e0b345ce722ac47727" [[package]] name = "flagsmith-flag-engine" -version = "5.2.0" +version = "5.3.0" description = "Flag engine for the Flagsmith API." optional = false python-versions = "*" files = [ - {file = "flagsmith-flag-engine-5.2.0.tar.gz", hash = "sha256:247da79a9fee55a55de1440fcd112de9b56cef9d5feb8c783319a9e0620e3a72"}, + {file = "flagsmith-flag-engine-5.3.0.tar.gz", hash = "sha256:87007f6a312cf11b2c201acd54b30f17de8aa039c3c56af431f1ed3c743fa84c"}, ] [package.dependencies] @@ -4182,4 +4182,4 @@ files = [ [metadata] lock-version = "2.0" python-versions = ">=3.11, <3.13" -content-hash = "508c20bd1cbdf2c8d1173d13b27369a7a9a0be82350f429467af94f53e80ebb9" +content-hash = "27333f5bbd3bb607cdb7d728dae6d0a6a11658cba09b45d69a6ee5a744111ad5" diff --git a/api/pyproject.toml b/api/pyproject.toml index 4825c7577c6c..aaac71259708 100644 --- a/api/pyproject.toml +++ b/api/pyproject.toml @@ -142,7 +142,7 @@ environs = "~9.2.0" django-lifecycle = "~1.0.0" drf-writable-nested = "~0.6.2" django-filter = "~2.4.0" -flagsmith-flag-engine = "^5.2.0" +flagsmith-flag-engine = "^5.3.0" boto3 = "~1.28.78" slack-sdk = "~3.9.0" asgiref = "~3.8.1" diff --git a/api/tests/unit/import_export/test_unit_import_export_export.py b/api/tests/unit/import_export/test_unit_import_export_export.py index 19b48a50f6db..8cb374d28666 100644 --- a/api/tests/unit/import_export/test_unit_import_export_export.py +++ b/api/tests/unit/import_export/test_unit_import_export_export.py @@ -1,6 +1,8 @@ import json import re +import typing import uuid +from decimal import Decimal import boto3 from core.constants import STRING @@ -9,7 +11,10 @@ from django.core.serializers.json import DjangoJSONEncoder from flag_engine.segments.constants import ALL_RULE, EQUAL from moto import mock_s3 +from mypy_boto3_dynamodb.service_resource import Table +from pytest_mock import MockerFixture +from environments.identities.models import Identity from environments.models import Environment, EnvironmentAPIKey, Webhook from features.feature_types import MULTIVARIATE from features.models import Feature, FeatureSegment, FeatureState @@ -17,6 +22,7 @@ from features.workflows.core.models import ChangeRequest from import_export.export import ( S3OrganisationExporter, + export_edge_identities, export_environments, export_features, export_metadata, @@ -96,6 +102,36 @@ def test_export_project(organisation): # TODO: test whether the export is importable +def test_export_project__only_live_segments_are_exported( + organisation: Organisation, project: Project +): + # Given + # a segment + segment = Segment.objects.create(project=project, name="test segment") + segment_rule = SegmentRule.objects.create(segment=segment, type=ALL_RULE) + segment_condition = Condition.objects.create( + rule=segment_rule, operator=EQUAL, property="foo", value="bar" + ) + + # A segment version that is not live + segmet2 = Segment.objects.create( + project=project, name="test segment 2", version_of=segment + ) + segment_rule2 = SegmentRule.objects.create(segment=segmet2, type=ALL_RULE) + Condition.objects.create( + rule=segment_rule2, operator=EQUAL, property="foo", value="bar" + ) + # When + export = export_projects(organisation.id) + + # Then + # only the project and the live segment should be exported + assert len(export) == 4 + assert export[1]["fields"]["uuid"] == str(segment.uuid) + assert export[2]["fields"]["uuid"] == str(segment_rule.uuid) + assert export[3]["fields"]["uuid"] == str(segment_condition.uuid) + + def test_export_environments(project): # Given # an environment @@ -278,6 +314,242 @@ def test_export_features_with_environment_feature_version( # TODO: test whether the export is importable +def test_export_edge_identities( + flagsmith_identities_table: Table, + project: Project, + environment: Environment, + multivariate_feature: Feature, + multivariate_options: typing.List[MultivariateFeatureOption], + mocker: MockerFixture, +) -> None: + # Given + project.enable_dynamo_db = True + project.save() + + # First, let's create some features(to override) + int_feature = Feature.objects.create( + project=project, name="int_feature", initial_value=11 + ) + float_feature = Feature.objects.create( + project=project, name="float_feature", initial_value=11.1 + ) + bool_feature = Feature.objects.create( + project=project, name="bool_feature", initial_value=True + ) + + # Let's create another feature that we are not going to override + Feature.objects.create(project=project, name="string_feature", initial_value="foo") + + # another mv feature that we will override + # using the option id that does not exists + second_mv_feature = Feature.objects.create( + project=project, name="mv_feature_with_deleted_option", type=MULTIVARIATE + ) + + mv_option = multivariate_options[0] + + identity_identifier = "Development_user_123456" + mv_override_fs_uuid = "b7c3d9e9-0bcc-4e60-8264-43e84b00fcbd" + int_override_fs_uuid = "c6f9cec7-f27b-4e4f-80ff-5a2dfa3d4d20" + float_override_fs_uuid = "b90eafdc-56f3-45ba-965f-e245007f3050" + bool_override_fs_uuid = "2dab9fe3-49df-41ec-adc1-30f5dfe0b855" + + identity_document = { + "composite_key": f"{environment.api_key}_{identity_identifier}", + "created_date": "2024-09-22T07:27:27.770956+00:00", + "django_id": None, + "environment_api_key": environment.api_key, + "identifier": identity_identifier, + "identity_features": [ + { + "django_id": None, + "enabled": False, + "feature": { + "id": multivariate_feature.id, + "name": multivariate_feature.name, + "type": "MULTIVARIATE", + }, + "featurestate_uuid": mv_override_fs_uuid, + "feature_segment": None, + "feature_state_value": "control", + "multivariate_feature_state_values": [ + { + "id": None, + "multivariate_feature_option": { + "id": mv_option.id, + "value": mv_option.string_value, + }, + "mv_fs_value_uuid": "1897c9df-b8fa-4870-a077-f48eadbf3aac", + "percentage_allocation": 100, + } + ], + }, + { + "django_id": None, + "enabled": True, + "feature": { + "id": int_feature.id, + "name": int_feature.name, + "type": "STANDARD", + }, + "featurestate_uuid": int_override_fs_uuid, + "feature_segment": None, + "feature_state_value": 123, + "multivariate_feature_state_values": [], + }, + { + "django_id": None, + "enabled": True, + "feature": { + "id": float_feature.id, + "name": int_feature.name, + "type": "STANDARD", + }, + "featurestate_uuid": float_override_fs_uuid, + "feature_segment": None, + "feature_state_value": Decimal("123.123"), + "multivariate_feature_state_values": [], + }, + { + "django_id": None, + "enabled": True, + "feature": { + "id": bool_feature.id, + "name": int_feature.name, + "type": "STANDARD", + }, + "featurestate_uuid": bool_override_fs_uuid, + "feature_segment": None, + "feature_state_value": False, + "multivariate_feature_state_values": [], + }, + # A feature that does not exists anymore(to make sure we skip it during export) + { + "django_id": None, + "enabled": True, + "featurestate_uuid": "53bd193f-da11-40c8-b694-3261f28c720c", + "feature": { + "id": 9999, + "name": "feature_that_does_not_exists", + "type": "STANDARD", + }, + }, + # An mv override with an option that does not exists anymore(to make sure we skip it during export) + { + "django_id": None, + "enabled": False, + "feature": { + "id": second_mv_feature.id, + "name": second_mv_feature.name, + "type": "MULTIVARIATE", + }, + "featurestate_uuid": "53bd193f-da11-40c8-b694-3261f28c720d", + "feature_segment": None, + "feature_state_value": "control", + "multivariate_feature_state_values": [ + { + "id": None, + "multivariate_feature_option": { + "id": 999999, # does not exists + "value": mv_option.string_value, + }, + "mv_fs_value_uuid": "1897c9df-b8fa-4870-a077-f48eadbf3aac", + "percentage_allocation": 100, + } + ], + }, + ], + "identity_traits": [ + {"trait_key": "int_trait", "trait_value": 123}, + {"trait_key": "float_trait", "trait_value": Decimal("123.123")}, + {"trait_key": "str_trait", "trait_value": "some-string"}, + {"trait_key": "bool_trait", "trait_value": True}, + ], + "identity_uuid": "37ecaac3-70dd-4135-b2ee-9b2e3ffdc028", + } + flagsmith_identities_table.put_item(Item=identity_document) + + # another identity to test pagination + flagsmith_identities_table.put_item( + Item={ + "composite_key": f"{environment.api_key}_identity_one", + "environment_api_key": environment.api_key, + "identifier": "identity_two", + "created_date": "2024-09-22T07:27:27.770956+00:00", + "identity_traits": [], + "identity_features": [], + } + ) + + # When + mocker.patch("edge_api.identities.export.EXPORT_EDGE_IDENTITY_PAGINATION_LIMIT", 1) + export_json = export_edge_identities(project.organisation_id) + + # Let's load the data + file_path = f"/tmp/{uuid.uuid4()}.json" + with open(file_path, "a+") as f: + f.write(json.dumps(export_json, cls=DjangoJSONEncoder)) + f.seek(0) + + call_command("loaddata", f.name, format="json") + # Then + # the identity was created + assert Identity.objects.count() == 2 + identity = Identity.objects.get(identifier=identity_identifier) + + traits = identity.get_all_user_traits() + + assert len(traits) == 4 + int_trait = traits[0] + assert int_trait.trait_key == "int_trait" + assert int_trait.trait_value == 123 + + float_trait = traits[1] + assert float_trait.trait_key == "float_trait" + assert float_trait.trait_value == 123.123 + + str_trait = traits[2] + assert str_trait.trait_key == "str_trait" + assert str_trait.trait_value == "some-string" + + bool_trait = traits[3] + assert bool_trait.trait_key == "bool_trait" + assert bool_trait.trait_value is True + + all_feature_states = identity.get_all_feature_states() + assert len(all_feature_states) == 6 + + actual_mv_override = all_feature_states[0] + assert str(actual_mv_override.uuid) == mv_override_fs_uuid + assert ( + actual_mv_override.get_feature_state_value(identity=identity) + == mv_option.string_value + ) + + actual_int_override = all_feature_states[1] + assert str(actual_int_override.uuid) == int_override_fs_uuid + assert actual_int_override.get_feature_state_value(identity=identity) == 123 + + actual_float_override = all_feature_states[2] + assert str(actual_float_override.uuid) == float_override_fs_uuid + assert actual_float_override.get_feature_state_value(identity=identity) == "123.123" + + actual_bool_override = all_feature_states[3] + assert str(actual_bool_override.uuid) == bool_override_fs_uuid + assert actual_bool_override.get_feature_state_value(identity=identity) is False + + actual_string_fs = all_feature_states[4] + assert actual_string_fs.get_feature_state_value(identity=identity) == "foo" + assert actual_string_fs.identity is None + + override_without_mv_option = all_feature_states[5] + assert ( + override_without_mv_option.get_feature_state_value(identity=identity) + == "control" + ) + assert override_without_mv_option.identity == identity + + @mock_s3 def test_organisation_exporter_export_to_s3(organisation): # Given