Skip to content

Commit

Permalink
misc
Browse files Browse the repository at this point in the history
  • Loading branch information
@gagantrivedi
gagantrivedi committed Sep 26, 2024
1 parent 4781b09 commit 9cb31b0
Show file tree
Hide file tree
Showing 3 changed files with 106 additions and 30 deletions.
65 changes: 42 additions & 23 deletions api/edge_api/identities/export.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
import logging
import typing
import uuid
from decimal import Decimal
from typing import Dict, List, Tuple
from typing import Dict, Tuple

from django.utils import timezone
from flag_engine.identities.traits.types import map_any_value_to_trait_value
Expand All @@ -11,23 +12,25 @@
from features.models import Feature, FeatureState
from features.multivariate.models import MultivariateFeatureOption

feature_id_to_uuid: Dict[int, str] = {}
EXPORT_EDGE_IDENTITY_PAGINATION_LIMIT = 20000

mv_feature_option_id_to_uuid: Dict[int, str] = {}
logger = logging.getLogger()

EXPORT_EDGE_IDENTITY_PAGINATION_LIMIT = 1000


def export_edge_identity_and_overrides(
def export_edge_identity_and_overrides( # noqa: C901
environment_api_key: str,
) -> Tuple[List, List, List]:
) -> 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] = {}

mv_feature_option_id_to_uuid: Dict[int, str] = {}
while True:
response = EdgeIdentity.dynamo_wrapper.get_all_items(**kwargs)
for item in response["Items"]:
Expand All @@ -46,13 +49,22 @@ def export_edge_identity_and_overrides(
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:
try:
feature_id_to_uuid[feature_id] = _get_feature_uuid(feature_id)
except Feature.DoesNotExist:
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_id,
feature_uuid,
override["enabled"],
)
)
Expand All @@ -65,10 +77,26 @@ def export_edge_identity_and_overrides(
if mvfsv := override["multivariate_feature_state_values"]:
# only one mvfs override can exists
mv_feature_option_id = mvfsv[0]["multivariate_feature_option"]["id"]
if mv_feature_option_id not in mv_feature_option_id_to_uuid:
try:
mv_feature_option_id_to_uuid[mv_feature_option_id] = (
_get_mv_feature_option_uuid(mv_feature_option_id)
)
except MultivariateFeatureOption.DoesNotExist:
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
]

# export mv feature state value
identity_override_export.append(
export_mv_featurestate_value(
featurestate_uuid, mv_feature_option_id
featurestate_uuid, mv_feature_option_uuid
)
)
if "LastEvaluatedKey" not in response:
Expand Down Expand Up @@ -116,14 +144,12 @@ def export_edge_feature_state(
identifier: str,
environment_api_key: str,
featurestate_uuid: str,
feature_id: int,
feature_uuid: str,
enabled: bool,
) -> dict:
global feature_id_to_uuid
if feature_id not in feature_id_to_uuid:
feature_id_to_uuid[feature_id] = _get_feature_uuid(feature_id)
feature_uuid = feature_id_to_uuid[feature_id]

# 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": {
Expand Down Expand Up @@ -165,16 +191,9 @@ def export_featurestate_value(


def export_mv_featurestate_value(
featurestate_uuid: str, mv_feature_option_id: int
featurestate_uuid: str, mv_feature_option_uuid: int
) -> dict:

global mv_feature_option_id_to_uuid
if mv_feature_option_id not in mv_feature_option_id_to_uuid:
mv_feature_option_id_to_uuid[mv_feature_option_id] = (
_get_mv_feature_option_uuid(mv_feature_option_id)
)
mv_feature_option_uuid = mv_feature_option_id_to_uuid[mv_feature_option_id]

return {
"model": "multivariate.multivariatefeaturestatevalue",
"fields": {
Expand Down
11 changes: 9 additions & 2 deletions api/import_export/export.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -152,13 +152,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,
),
),
)

Expand Down
60 changes: 55 additions & 5 deletions api/tests/unit/import_export/test_unit_import_export_export.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -291,7 +291,7 @@ def test_export_edge_identities(
multivariate_feature: Feature,
multivariate_options: typing.List[MultivariateFeatureOption],
mocker: MockerFixture,
):
) -> None:
# Given
project.enable_dynamo_db = True
project.save()
Expand All @@ -310,6 +310,12 @@ def test_export_edge_identities(
# 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"
Expand Down Expand Up @@ -387,6 +393,41 @@ def test_export_edge_identities(
"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},
Expand All @@ -397,12 +438,16 @@ def test_export_edge_identities(
"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": [],
}
)

Expand All @@ -422,9 +467,7 @@ def test_export_edge_identities(
assert Identity.objects.count() == 2
identity = Identity.objects.get(identifier=identity_identifier)

# With the traits that were part of the document
traits = identity.get_all_user_traits()
all_feature_states = identity.get_all_feature_states()

assert len(traits) == 4
int_trait = traits[0]
Expand All @@ -443,8 +486,8 @@ def test_export_edge_identities(
assert bool_trait.trait_key == "bool_trait"
assert bool_trait.trait_value is True

# And the feature states that were part of the document
assert len(all_feature_states) == 5
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
Expand All @@ -469,6 +512,13 @@ def test_export_edge_identities(
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):
Expand Down

0 comments on commit 9cb31b0

Please sign in to comment.