flagsmith-engine/flag_engine/identities/models.py at e1caa277f1f46ef209f0bacf2be7fd3aa90adefd · Flagsmith/flagsmith-engine · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
import datetime
import typing
import uuid

from pydantic import UUID4, BaseModel, Field, computed_field, model_validator
from pydantic_collections import BaseCollectionModel

from flag_engine.features.models import FeatureStateModel
from flag_engine.identities.traits.models import TraitModel
from flag_engine.utils.datetime import utcnow_with_tz
from flag_engine.utils.exceptions import DuplicateFeatureState


class IdentityFeaturesList(BaseCollectionModel[FeatureStateModel]):  # type: ignore[misc,no-any-unimported]
    @staticmethod
    def _ensure_unique_feature_ids(
        value: typing.Sequence[FeatureStateModel],
    ) -> None:
        for i, feature_state in enumerate(value, start=1):
            if feature_state.feature.id in [
                feature_state.feature.id for feature_state in value[i:]
            ]:
                raise DuplicateFeatureState(
                    f"Feature state for feature id={feature_state.feature.id} already exists"
                )

    @model_validator(mode="after")
    def ensure_unique_feature_ids(self) -> "IdentityFeaturesList":
        self._ensure_unique_feature_ids(self.root)
        return self

    def append(self, feature_state: "FeatureStateModel") -> None:
        self._ensure_unique_feature_ids([*self, feature_state])
        super().append(feature_state)


class IdentityModel(BaseModel):
    identifier: str
    environment_api_key: str
    created_date: datetime.datetime = Field(default_factory=utcnow_with_tz)
    identity_features: IdentityFeaturesList = Field(
        default_factory=IdentityFeaturesList
    )
    identity_traits: typing.List[TraitModel] = Field(default_factory=list)
    identity_uuid: UUID4 = Field(default_factory=uuid.uuid4)
    django_id: typing.Optional[int] = None

    dashboard_alias: typing.Optional[str] = None

    @computed_field  # type: ignore[prop-decorator]
    @property
    def composite_key(self) -> str:
        return self.generate_composite_key(self.environment_api_key, self.identifier)

    @staticmethod
    def generate_composite_key(env_key: str, identifier: str) -> str:
        return f"{env_key}_{identifier}"

    def get_hash_key(self, use_identity_composite_key_for_hashing: bool) -> str:
        return (
            self.composite_key
            if use_identity_composite_key_for_hashing
            else self.identifier
        )

    def update_traits(
        self, traits: typing.List[TraitModel]
    ) -> typing.Tuple[typing.List[TraitModel], bool]:
        existing_traits = {trait.trait_key: trait for trait in self.identity_traits}
        traits_changed = False

        for trait in traits:
            existing_trait = existing_traits.get(trait.trait_key)

            if trait.trait_value is None and existing_trait:
                existing_traits.pop(trait.trait_key)
                traits_changed = True

            elif getattr(existing_trait, "trait_value", None) != trait.trait_value:
                existing_traits[trait.trait_key] = trait
                traits_changed = True

        self.identity_traits = list(existing_traits.values())
        return self.identity_traits, traits_changed

    def prune_features(self, valid_feature_names: typing.List[str]) -> None:
        self.identity_features = IdentityFeaturesList(
            [
                fs
                for fs in self.identity_features
                if fs.feature.name in valid_feature_names
            ]
        )