diff --git a/api/cohorts/constants.py b/api/cohorts/constants.py new file mode 100644 index 000000000000..79aa289ab9f5 --- /dev/null +++ b/api/cohorts/constants.py @@ -0,0 +1,10 @@ +COHORT_SYSTEM_TRAIT_KEY_PREFIX = "flagsmith_cohort_" +COHORT_MEMBERSHIP_APPLY_BATCH_SIZE = 100 +COHORT_MEMBERSHIP_APPLY_MAX_BATCHES_PER_RUN = 10 +DYNAMODB_THROTTLING_ERROR_CODES = frozenset( + { + "ProvisionedThroughputExceededException", + "RequestLimitExceeded", + "ThrottlingException", + } +) diff --git a/api/cohorts/metrics.py b/api/cohorts/metrics.py new file mode 100644 index 000000000000..cb9236f69064 --- /dev/null +++ b/api/cohorts/metrics.py @@ -0,0 +1,9 @@ +import prometheus_client + +flagsmith_cohorts_membership_deltas_applied_total = prometheus_client.Counter( + "flagsmith_cohorts_membership_deltas_applied_total", + "Total number of cohort membership ledger rows transitioned to their " + "applied state after the corresponding identity document write. " + "The `operation` label is either `add` or `remove`.", + ["operation"], +) diff --git a/api/cohorts/models.py b/api/cohorts/models.py index 661ac923a08e..cb2e9f910dd8 100644 --- a/api/cohorts/models.py +++ b/api/cohorts/models.py @@ -1,5 +1,6 @@ from django.db import models +from cohorts.constants import COHORT_SYSTEM_TRAIT_KEY_PREFIX from core.models import SoftDeleteExportableModel @@ -26,6 +27,10 @@ class Cohort(SoftDeleteExportableModel): version = models.PositiveIntegerField(default=0) created_at = models.DateTimeField(auto_now_add=True) + @property + def system_trait_key(self) -> str: + return f"{COHORT_SYSTEM_TRAIT_KEY_PREFIX}{self.uuid}" + class Meta: constraints = [ # Exactly one active cohort feeds a segment: two cohorts on one diff --git a/api/cohorts/services.py b/api/cohorts/services.py new file mode 100644 index 000000000000..91e9ba3bcc98 --- /dev/null +++ b/api/cohorts/services.py @@ -0,0 +1,68 @@ +import structlog +from django.db.models import QuerySet +from django.utils import timezone + +from cohorts.constants import COHORT_MEMBERSHIP_APPLY_BATCH_SIZE +from cohorts.metrics import flagsmith_cohorts_membership_deltas_applied_total +from cohorts.models import Cohort, CohortMembership, CohortMembershipState +from environments.dynamodb import DynamoIdentityWrapper + +logger = structlog.get_logger("cohorts") + +_PENDING_STATES = [ + CohortMembershipState.PENDING_ADD, + CohortMembershipState.PENDING_REMOVE, +] + + +def pending_memberships(cohort: Cohort) -> "QuerySet[CohortMembership]": + return CohortMembership.objects.filter(cohort=cohort, state__in=_PENDING_STATES) + + +def apply_pending_memberships(cohort: Cohort) -> bool: + identity_wrapper = DynamoIdentityWrapper() + environment_api_key: str = cohort.environment.api_key + trait_key = cohort.system_trait_key + batch = list( + pending_memberships(cohort).order_by("id")[:COHORT_MEMBERSHIP_APPLY_BATCH_SIZE] + ) + if not batch: + return False + added_ids: list[int] = [] + removed_ids: list[int] = [] + for row in batch: + if row.state == CohortMembershipState.PENDING_ADD: + identity_wrapper.set_system_trait( + environment_api_key=environment_api_key, + identifier=row.identifier, + trait_key=trait_key, + ) + added_ids.append(row.id) + else: + identity_wrapper.unset_system_trait( + environment_api_key=environment_api_key, + identifier=row.identifier, + trait_key=trait_key, + ) + removed_ids.append(row.id) + added_count = CohortMembership.objects.filter( + id__in=added_ids, state=CohortMembershipState.PENDING_ADD + ).update(state=CohortMembershipState.APPLIED, updated_at=timezone.now()) + removed_count, _ = CohortMembership.objects.filter( + id__in=removed_ids, state=CohortMembershipState.PENDING_REMOVE + ).delete() + flagsmith_cohorts_membership_deltas_applied_total.labels(operation="add").inc( + added_count + ) + flagsmith_cohorts_membership_deltas_applied_total.labels(operation="remove").inc( + removed_count + ) + if added_count or removed_count: + logger.info( + "membership.applied", + cohort__id=cohort.id, + environment__id=cohort.environment_id, + adds__count=added_count, + removes__count=removed_count, + ) + return pending_memberships(cohort).exists() diff --git a/api/cohorts/tasks.py b/api/cohorts/tasks.py new file mode 100644 index 000000000000..b992672619df --- /dev/null +++ b/api/cohorts/tasks.py @@ -0,0 +1,41 @@ +from datetime import timedelta + +import structlog +from botocore.exceptions import ClientError +from task_processor.decorators import register_task_handler +from task_processor.exceptions import TaskBackoffError + +from cohorts import services +from cohorts.constants import ( + COHORT_MEMBERSHIP_APPLY_MAX_BATCHES_PER_RUN, + DYNAMODB_THROTTLING_ERROR_CODES, +) +from cohorts.models import Cohort +from environments.dynamodb import DynamoIdentityWrapper + +logger = structlog.get_logger("cohorts") + + +@register_task_handler(timeout=timedelta(minutes=5)) +def apply_cohort_membership_deltas(cohort_id: int) -> None: + log = logger.bind(cohort__id=cohort_id) + if (cohort := Cohort.objects.filter(id=cohort_id).first()) is None: + log.info("membership.apply.skipped", reason="cohort_missing") + return + if not ( + cohort.environment.project.enable_dynamo_db + and DynamoIdentityWrapper().is_enabled + ): + log.info("membership.apply.skipped", reason="not_edge") + return + try: + for _ in range(COHORT_MEMBERSHIP_APPLY_MAX_BATCHES_PER_RUN): + if not services.apply_pending_memberships(cohort): + return + except ClientError as exc: + if exc.response["Error"]["Code"] in DYNAMODB_THROTTLING_ERROR_CODES: + log.warning("membership.apply.throttled") + raise TaskBackoffError() from exc + raise + # Still pending after this run's batch cap; continue in a fresh task run. + apply_cohort_membership_deltas.delay(kwargs={"cohort_id": cohort_id}) diff --git a/api/environments/dynamodb/constants.py b/api/environments/dynamodb/constants.py index 90097a39898e..fc46a8154293 100644 --- a/api/environments/dynamodb/constants.py +++ b/api/environments/dynamodb/constants.py @@ -9,6 +9,8 @@ DYNAMODB_MAX_BATCH_WRITE_ITEM_COUNT = 25 IDENTITIES_PAGINATION_LIMIT = 1000 +SYSTEM_TRAIT_WRITE_MAX_ATTEMPTS = 3 + # DynamoDB max item size is 400 KB (409,600 bytes). DOCUMENT_SIZE_HISTOGRAM_BUCKETS = ( 1_000, diff --git a/api/environments/dynamodb/wrappers/exceptions.py b/api/environments/dynamodb/wrappers/exceptions.py index e9e70f03c01b..907eea849e65 100644 --- a/api/environments/dynamodb/wrappers/exceptions.py +++ b/api/environments/dynamodb/wrappers/exceptions.py @@ -1,6 +1,16 @@ from decimal import Decimal +class SystemTraitWriteRaceError(Exception): + def __init__(self, composite_key: str) -> None: + super().__init__( + f"Gave up writing a system trait for identity {composite_key!r}: " + "concurrent writers kept changing the document between read and " + "conditional write." + ) + self.composite_key = composite_key + + class CapacityBudgetExceeded(Exception): def __init__( self, diff --git a/api/environments/dynamodb/wrappers/identity_wrapper.py b/api/environments/dynamodb/wrappers/identity_wrapper.py index dc5304168f22..8f29ffa54295 100644 --- a/api/environments/dynamodb/wrappers/identity_wrapper.py +++ b/api/environments/dynamodb/wrappers/identity_wrapper.py @@ -5,19 +5,29 @@ from typing import Iterable from boto3.dynamodb.conditions import Attr, Key +from botocore.exceptions import ClientError from django.conf import settings from django.core.exceptions import ObjectDoesNotExist from rest_framework.exceptions import NotFound from edge_api.identities.search import EdgeIdentitySearchData -from environments.dynamodb.constants import IDENTITIES_PAGINATION_LIMIT -from environments.dynamodb.wrappers.exceptions import CapacityBudgetExceeded +from environments.dynamodb.constants import ( + IDENTITIES_PAGINATION_LIMIT, + SYSTEM_TRAIT_WRITE_MAX_ATTEMPTS, +) +from environments.dynamodb.wrappers.exceptions import ( + CapacityBudgetExceeded, + SystemTraitWriteRaceError, +) from util.engine_models.context.mappers import ( is_context_in_segment, map_environment_identity_to_context, ) from util.engine_models.identities.models import IdentityModel -from util.mappers import map_identity_to_identity_document +from util.mappers import ( + map_engine_identity_to_identity_document, + map_identity_to_identity_document, +) from .base import BaseDynamoWrapper @@ -34,6 +44,17 @@ logger = logging.getLogger(__name__) +def _system_trait_value_matches( + stored_value: object, + document_value: bool | int | Decimal | str, +) -> bool: + # The bool check stops `Decimal(1) == True` false positives. + return ( + isinstance(stored_value, bool) == isinstance(document_value, bool) + and stored_value == document_value + ) + + class DynamoIdentityWrapper(BaseDynamoWrapper): def __init__(self) -> None: super().__init__() @@ -63,6 +84,117 @@ def write_identities(self, identities: Iterable["Identity"]): # type: ignore[no def get_item(self, composite_key: str) -> typing.Optional[dict]: # type: ignore[type-arg] return self.table.get_item(Key={"composite_key": composite_key}).get("Item") # type: ignore[union-attr] + def set_system_trait( + self, + *, + environment_api_key: str, + identifier: str, + trait_key: str, + trait_value: bool | int | float | str = True, + ) -> None: + """Idempotently set a system trait on an identity document. + + Writes only touch the `system_traits.` attribute, so + concurrent writes to other attributes are never overwritten; the + document is created if missing. Each write is conditional on the + document shape just read — a lost race re-reads and retries, and + `SystemTraitWriteRaceError` is raised once attempts are exhausted. + + Assumes stored documents never carry `system_traits` as NULL — the + document mapper omits the attribute when unset. + """ + composite_key = IdentityModel.generate_composite_key( + environment_api_key, identifier + ) + # DynamoDB rejects floats and returns all numbers as Decimal. + document_value: bool | int | Decimal | str = ( + Decimal(str(trait_value)) if isinstance(trait_value, float) else trait_value + ) + for _ in range(SYSTEM_TRAIT_WRITE_MAX_ATTEMPTS): + # Strongly consistent read: a replication-lagged hint would burn + # retry attempts on conditional writes that can never succeed. + document = self.table.get_item( # type: ignore[union-attr] + Key={"composite_key": composite_key}, ConsistentRead=True + ).get("Item") + system_traits = document.get("system_traits") if document else None + if isinstance(system_traits, dict) and _system_trait_value_matches( + system_traits.get(trait_key), document_value + ): + return + try: + if document is None: + self.table.put_item( # type: ignore[union-attr] + Item=map_engine_identity_to_identity_document( + IdentityModel( + identifier=identifier, + environment_api_key=environment_api_key, + system_traits={trait_key: trait_value}, + ) + ), + ConditionExpression="attribute_not_exists(composite_key)", + ) + elif isinstance(system_traits, dict): + self.table.update_item( # type: ignore[union-attr] + Key={"composite_key": composite_key}, + UpdateExpression="SET system_traits.#tk = :value", + ConditionExpression="attribute_exists(system_traits)", + ExpressionAttributeNames={"#tk": trait_key}, + ExpressionAttributeValues={":value": document_value}, + ) + else: + # If another writer created system_traits after we read + # the document, this write does nothing and their traits + # survive; the returned attributes tell us which happened. + response = self.table.update_item( # type: ignore[union-attr] + Key={"composite_key": composite_key}, + UpdateExpression=( + "SET system_traits = if_not_exists(system_traits, :init)" + ), + # Without this condition, update_item would re-create + # a just-deleted identity as an empty document + # containing nothing but this trait. + ConditionExpression="attribute_exists(composite_key)", + ExpressionAttributeValues={ + ":init": {trait_key: document_value} + }, + ReturnValues="ALL_NEW", + ) + written_traits = response["Attributes"].get("system_traits") + if isinstance(written_traits, dict) and _system_trait_value_matches( + written_traits.get(trait_key), document_value + ): + return + continue + return + except ClientError as exc: + if exc.response["Error"]["Code"] != "ConditionalCheckFailedException": + raise + raise SystemTraitWriteRaceError(composite_key) + + def unset_system_trait( + self, + *, + environment_api_key: str, + identifier: str, + trait_key: str, + ) -> None: + """Idempotently remove a system trait from an identity document.""" + composite_key = IdentityModel.generate_composite_key( + environment_api_key, identifier + ) + try: + self.table.update_item( # type: ignore[union-attr] + Key={"composite_key": composite_key}, + UpdateExpression="REMOVE system_traits.#tk", + # Failing this condition covers every no-op case at once: + # missing document, missing system_traits, or trait already absent. + ConditionExpression="attribute_exists(system_traits.#tk)", + ExpressionAttributeNames={"#tk": trait_key}, + ) + except ClientError as exc: + if exc.response["Error"]["Code"] != "ConditionalCheckFailedException": + raise + def delete_item(self, composite_key: str): # type: ignore[no-untyped-def] self.table.delete_item(Key={"composite_key": composite_key}) # type: ignore[union-attr] diff --git a/api/tests/unit/cohorts/conftest.py b/api/tests/unit/cohorts/conftest.py new file mode 100644 index 000000000000..adf99e3cbeb7 --- /dev/null +++ b/api/tests/unit/cohorts/conftest.py @@ -0,0 +1,26 @@ +import pytest + +from cohorts.models import Cohort +from environments.models import Environment +from projects.models import Project +from segments.models import Segment + + +@pytest.fixture() +def cohort(environment: Environment, segment: Segment) -> Cohort: + cohort: Cohort = Cohort.objects.create(environment=environment, segment=segment) + return cohort + + +@pytest.fixture() +def edge_cohort( + dynamo_enabled_project: Project, + dynamo_enabled_project_environment_one: Environment, +) -> Cohort: + segment = Segment.objects.create( + name="edge segment", project=dynamo_enabled_project + ) + cohort: Cohort = Cohort.objects.create( + environment=dynamo_enabled_project_environment_one, segment=segment + ) + return cohort diff --git a/api/tests/unit/cohorts/test_services.py b/api/tests/unit/cohorts/test_services.py new file mode 100644 index 000000000000..1737caea6465 --- /dev/null +++ b/api/tests/unit/cohorts/test_services.py @@ -0,0 +1,110 @@ +from pytest_mock import MockerFixture +from pytest_structlog import StructuredLogCapture + +from cohorts.models import Cohort, CohortMembership, CohortMembershipState +from cohorts.services import apply_pending_memberships +from environments.dynamodb import DynamoIdentityWrapper + + +def test_apply_pending_memberships__no_pending_rows__returns_false( + cohort: Cohort, +) -> None: + # Given + CohortMembership.objects.create( + cohort=cohort, identifier="user-1", state=CohortMembershipState.APPLIED + ) + + # When + result = apply_pending_memberships(cohort) + + # Then + assert result is False + + +def test_apply_pending_memberships__pending_rows__applies_and_flips( + cohort: Cohort, + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + api_key = cohort.environment.api_key + trait_key = cohort.system_trait_key + dynamodb_identity_wrapper.put_item( + { + "composite_key": f"{api_key}_member", + "identifier": "member", + "environment_api_key": api_key, + "system_traits": {trait_key: True}, + } + ) + CohortMembership.objects.create(cohort=cohort, identifier="joiner") + CohortMembership.objects.create( + cohort=cohort, + identifier="member", + state=CohortMembershipState.PENDING_REMOVE, + ) + + # When + result = apply_pending_memberships(cohort) + + # Then + assert result is False + joiner_document = dynamodb_identity_wrapper.get_item(f"{api_key}_joiner") + assert joiner_document is not None + assert joiner_document["system_traits"] == {trait_key: True} + member_document = dynamodb_identity_wrapper.get_item(f"{api_key}_member") + assert member_document is not None + assert member_document["system_traits"] == {} + assert list( + CohortMembership.objects.filter(cohort=cohort).values_list( + "identifier", "state" + ) + ) == [("joiner", CohortMembershipState.APPLIED)] + + +def test_apply_pending_memberships__more_rows_than_batch__returns_true( + cohort: Cohort, + dynamodb_identity_wrapper: DynamoIdentityWrapper, + mocker: MockerFixture, +) -> None: + # Given + mocker.patch("cohorts.services.COHORT_MEMBERSHIP_APPLY_BATCH_SIZE", 1) + CohortMembership.objects.create(cohort=cohort, identifier="user-1") + CohortMembership.objects.create(cohort=cohort, identifier="user-2") + + # When + result = apply_pending_memberships(cohort) + + # Then + assert result is True + assert ( + CohortMembership.objects.filter( + cohort=cohort, state=CohortMembershipState.APPLIED + ).count() + == 1 + ) + + +def test_apply_pending_memberships__row_transitioned_mid_write__not_flipped( + cohort: Cohort, + mocker: MockerFixture, + log: StructuredLogCapture, +) -> None: + # Given + membership = CohortMembership.objects.create(cohort=cohort, identifier="user-1") + wrapper_mock = mocker.patch("cohorts.services.DynamoIdentityWrapper").return_value + + def transition_row(**kwargs: str) -> None: + CohortMembership.objects.filter(id=membership.id).update( + state=CohortMembershipState.PENDING_REMOVE + ) + + wrapper_mock.set_system_trait.side_effect = transition_row + + # When + result = apply_pending_memberships(cohort) + + # Then + membership.refresh_from_db() + assert membership.state == CohortMembershipState.PENDING_REMOVE + assert result is True + assert not log.has("membership.applied") diff --git a/api/tests/unit/cohorts/test_tasks.py b/api/tests/unit/cohorts/test_tasks.py new file mode 100644 index 000000000000..9b00be6b7a81 --- /dev/null +++ b/api/tests/unit/cohorts/test_tasks.py @@ -0,0 +1,205 @@ +import pytest +from botocore.exceptions import ClientError +from prometheus_client import REGISTRY +from pytest_mock import MockerFixture +from pytest_structlog import StructuredLogCapture +from task_processor.exceptions import TaskBackoffError + +from cohorts import services +from cohorts.models import Cohort, CohortMembership, CohortMembershipState +from cohorts.tasks import apply_cohort_membership_deltas +from environments.dynamodb import DynamoIdentityWrapper + + +def test_apply_cohort_membership_deltas__pending_adds__applies_to_documents( + edge_cohort: Cohort, + dynamodb_identity_wrapper: DynamoIdentityWrapper, + log: StructuredLogCapture, +) -> None: + # Given + api_key = edge_cohort.environment.api_key + dynamodb_identity_wrapper.put_item( + { + "composite_key": f"{api_key}_seen", + "identifier": "seen", + "environment_api_key": api_key, + } + ) + CohortMembership.objects.create(cohort=edge_cohort, identifier="seen") + CohortMembership.objects.create(cohort=edge_cohort, identifier="never-seen") + + # When + apply_cohort_membership_deltas(cohort_id=edge_cohort.id) + + # Then + for identifier in ("seen", "never-seen"): + document = dynamodb_identity_wrapper.get_item(f"{api_key}_{identifier}") + assert document is not None + assert document["system_traits"] == {edge_cohort.system_trait_key: True} + assert ( + CohortMembership.objects.filter( + cohort=edge_cohort, state=CohortMembershipState.APPLIED + ).count() + == 2 + ) + assert log.has( + "membership.applied", cohort__id=edge_cohort.id, adds__count=2, removes__count=0 + ) + + +def test_apply_cohort_membership_deltas__pending_removes__drops_trait_and_rows( + edge_cohort: Cohort, + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + api_key = edge_cohort.environment.api_key + dynamodb_identity_wrapper.put_item( + { + "composite_key": f"{api_key}_member", + "identifier": "member", + "environment_api_key": api_key, + "system_traits": {edge_cohort.system_trait_key: True, "other": True}, + } + ) + CohortMembership.objects.create( + cohort=edge_cohort, + identifier="member", + state=CohortMembershipState.PENDING_REMOVE, + ) + + # When + apply_cohort_membership_deltas(cohort_id=edge_cohort.id) + + # Then + document = dynamodb_identity_wrapper.get_item(f"{api_key}_member") + assert document is not None + assert document["system_traits"] == {"other": True} + assert not CohortMembership.objects.filter(cohort=edge_cohort).exists() + + +def test_apply_cohort_membership_deltas__non_edge_project__skips( + cohort: Cohort, + mocker: MockerFixture, + log: StructuredLogCapture, +) -> None: + # Given + mocker.patch("cohorts.tasks.DynamoIdentityWrapper").return_value.is_enabled = True + membership = CohortMembership.objects.create(cohort=cohort, identifier="user-1") + + # When + apply_cohort_membership_deltas(cohort_id=cohort.id) + + # Then + membership.refresh_from_db() + assert membership.state == CohortMembershipState.PENDING_ADD + assert log.has("membership.apply.skipped", cohort__id=cohort.id, reason="not_edge") + + +def test_apply_cohort_membership_deltas__missing_cohort__skips( + db: None, + log: StructuredLogCapture, +) -> None: + # Given + missing_cohort_id = 404404 + + # When + apply_cohort_membership_deltas(cohort_id=missing_cohort_id) + + # Then + assert log.has( + "membership.apply.skipped", + cohort__id=missing_cohort_id, + reason="cohort_missing", + level="info", + ) + + +@pytest.mark.parametrize( + "error_code", + [ + "ProvisionedThroughputExceededException", + "RequestLimitExceeded", + "ThrottlingException", + ], +) +def test_apply_cohort_membership_deltas__dynamo_throttled__raises_backoff( + error_code: str, + edge_cohort: Cohort, + mocker: MockerFixture, + log: StructuredLogCapture, +) -> None: + # Given + mocker.patch("cohorts.tasks.DynamoIdentityWrapper").return_value.is_enabled = True + mocker.patch.object( + services, + "apply_pending_memberships", + side_effect=ClientError({"Error": {"Code": error_code}}, "UpdateItem"), + ) + + # When + with pytest.raises(TaskBackoffError): + apply_cohort_membership_deltas(cohort_id=edge_cohort.id) + + # Then + assert log.has("membership.apply.throttled", cohort__id=edge_cohort.id) + + +def test_apply_cohort_membership_deltas__other_client_error__reraises( + edge_cohort: Cohort, + mocker: MockerFixture, + log: StructuredLogCapture, +) -> None: + # Given + mocker.patch("cohorts.tasks.DynamoIdentityWrapper").return_value.is_enabled = True + mocker.patch.object( + services, + "apply_pending_memberships", + side_effect=ClientError({"Error": {"Code": "ValidationException"}}, "PutItem"), + ) + + # When + with pytest.raises(ClientError) as exc_info: + apply_cohort_membership_deltas(cohort_id=edge_cohort.id) + + # Then + assert exc_info.value.response["Error"]["Code"] == "ValidationException" + assert not log.has("membership.apply.throttled") + + +def test_apply_cohort_membership_deltas__more_rows_than_batch__drains_ledger( + edge_cohort: Cohort, + dynamodb_identity_wrapper: DynamoIdentityWrapper, + mocker: MockerFixture, +) -> None: + # Given + mocker.patch("cohorts.services.COHORT_MEMBERSHIP_APPLY_BATCH_SIZE", 1) + mocker.patch("cohorts.tasks.COHORT_MEMBERSHIP_APPLY_MAX_BATCHES_PER_RUN", 1) + CohortMembership.objects.create(cohort=edge_cohort, identifier="user-1") + CohortMembership.objects.create(cohort=edge_cohort, identifier="user-2") + + # When (synchronous task runner executes the re-enqueued task inline) + apply_cohort_membership_deltas(cohort_id=edge_cohort.id) + + # Then + assert ( + CohortMembership.objects.filter( + cohort=edge_cohort, state=CohortMembershipState.APPLIED + ).count() + == 2 + ) + + +def test_apply_cohort_membership_deltas__deltas_applied__increments_metric( + edge_cohort: Cohort, + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + CohortMembership.objects.create(cohort=edge_cohort, identifier="user-1") + metric = "flagsmith_cohorts_membership_deltas_applied_total" + before = REGISTRY.get_sample_value(metric, {"operation": "add"}) or 0.0 + + # When + apply_cohort_membership_deltas(cohort_id=edge_cohort.id) + + # Then + assert REGISTRY.get_sample_value(metric, {"operation": "add"}) == before + 1 diff --git a/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_identity_wrapper.py b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_identity_wrapper.py index f92da6c839b9..f5b2e2664c4d 100644 --- a/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_identity_wrapper.py +++ b/api/tests/unit/environments/dynamodb/wrappers/test_unit_dynamodb_identity_wrapper.py @@ -4,6 +4,7 @@ import pytest from boto3.dynamodb.conditions import Key from boto3.dynamodb.types import Binary +from botocore.exceptions import ClientError from django.core.exceptions import ObjectDoesNotExist from flag_engine.segments.constants import IN from mypy_boto3_dynamodb.service_resource import Table @@ -18,7 +19,10 @@ EdgeIdentitySearchType, ) from environments.dynamodb import DynamoIdentityWrapper -from environments.dynamodb.wrappers.exceptions import CapacityBudgetExceeded +from environments.dynamodb.wrappers.exceptions import ( + CapacityBudgetExceeded, + SystemTraitWriteRaceError, +) from environments.identities.models import Identity from environments.identities.traits.models import Trait from features.models import Feature, FeatureSegment, FeatureState @@ -693,3 +697,362 @@ def test_delete_all_identities__multiple_identities__deletes_only_matching_envir # Then assert flagsmith_identities_table.scan()["Count"] == 1 assert flagsmith_identities_table.scan()["Items"][0] == identity_three + + +def test_set_system_trait__document_with_system_traits__sets_only_given_key( + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + dynamodb_identity_wrapper.put_item( + { + "composite_key": "api-key_user-1", + "identifier": "user-1", + "environment_api_key": "api-key", + "identity_traits": [{"trait_key": "plan", "trait_value": "pro"}], + "system_traits": {"other": True}, + } + ) + + # When + dynamodb_identity_wrapper.set_system_trait( + environment_api_key="api-key", identifier="user-1", trait_key="cohort_x" + ) + + # Then + document = dynamodb_identity_wrapper.get_item("api-key_user-1") + assert document is not None + assert document["system_traits"] == {"other": True, "cohort_x": True} + assert document["identity_traits"] == [{"trait_key": "plan", "trait_value": "pro"}] + + +def test_set_system_trait__document_without_system_traits__creates_system_traits( + dynamodb_identity_wrapper: DynamoIdentityWrapper, + mocker: MockerFixture, +) -> None: + # Given + dynamodb_identity_wrapper.put_item( + { + "composite_key": "api-key_user-1", + "identifier": "user-1", + "environment_api_key": "api-key", + } + ) + read_spy = mocker.spy(dynamodb_identity_wrapper.table, "get_item") + + # When + dynamodb_identity_wrapper.set_system_trait( + environment_api_key="api-key", identifier="user-1", trait_key="cohort_x" + ) + + # Then + assert read_spy.call_count == 1 + document = dynamodb_identity_wrapper.get_item("api-key_user-1") + assert document is not None + assert document["system_traits"] == {"cohort_x": True} + + +def test_set_system_trait__missing_document__creates_document( + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + composite_key = "api-key_user-1" + + # When + dynamodb_identity_wrapper.set_system_trait( + environment_api_key="api-key", identifier="user-1", trait_key="cohort_x" + ) + + # Then + document = dynamodb_identity_wrapper.get_item(composite_key) + assert document is not None + assert document["identifier"] == "user-1" + assert document["system_traits"] == {"cohort_x": True} + + +def test_set_system_trait__system_traits_created_concurrently__merges_into_existing( + dynamodb_identity_wrapper: DynamoIdentityWrapper, + mocker: MockerFixture, +) -> None: + # Given + # The table holds a document whose system_traits appeared after the + # wrapper's first read (simulated by a stale first response without them). + dynamodb_identity_wrapper.put_item( + { + "composite_key": "api-key_user-1", + "identifier": "user-1", + "environment_api_key": "api-key", + "system_traits": {"other": True}, + } + ) + real_get_item = dynamodb_identity_wrapper.table.get_item # type: ignore[union-attr] + stale_responses: typing.Iterator[dict[str, typing.Any]] = iter( + [ + { + "Item": { + "composite_key": "api-key_user-1", + "identifier": "user-1", + "environment_api_key": "api-key", + } + } + ] + ) + mocker.patch.object( + dynamodb_identity_wrapper.table, + "get_item", + side_effect=lambda **kwargs: next(stale_responses, real_get_item(**kwargs)), + ) + + # When + dynamodb_identity_wrapper.set_system_trait( + environment_api_key="api-key", identifier="user-1", trait_key="cohort_x" + ) + + # Then + document = real_get_item(Key={"composite_key": "api-key_user-1"})["Item"] + assert document["system_traits"] == {"other": True, "cohort_x": True} + + +def test_set_system_trait__custom_trait_value__written_to_document( + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + dynamodb_identity_wrapper.put_item( + { + "composite_key": "api-key_user-1", + "identifier": "user-1", + "environment_api_key": "api-key", + "system_traits": {"cohort_x": True}, + } + ) + + # When + dynamodb_identity_wrapper.set_system_trait( + environment_api_key="api-key", + identifier="user-1", + trait_key="score", + trait_value=0.5, + ) + + # Then + document = dynamodb_identity_wrapper.get_item("api-key_user-1") + assert document is not None + assert document["system_traits"] == {"cohort_x": True, "score": Decimal("0.5")} + + +def test_set_system_trait__value_changed__overwrites( + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + dynamodb_identity_wrapper.put_item( + { + "composite_key": "api-key_user-1", + "identifier": "user-1", + "environment_api_key": "api-key", + "system_traits": {"tier": "silver"}, + } + ) + + # When + dynamodb_identity_wrapper.set_system_trait( + environment_api_key="api-key", + identifier="user-1", + trait_key="tier", + trait_value="gold", + ) + + # Then + document = dynamodb_identity_wrapper.get_item("api-key_user-1") + assert document is not None + assert document["system_traits"] == {"tier": "gold"} + + +def test_set_system_trait__already_set__skips_write( + dynamodb_identity_wrapper: DynamoIdentityWrapper, + mocker: MockerFixture, +) -> None: + # Given + dynamodb_identity_wrapper.put_item( + { + "composite_key": "api-key_user-1", + "identifier": "user-1", + "environment_api_key": "api-key", + "system_traits": {"cohort_x": True}, + } + ) + update_mock = mocker.patch.object(dynamodb_identity_wrapper.table, "update_item") + put_mock = mocker.patch.object(dynamodb_identity_wrapper.table, "put_item") + + # When + dynamodb_identity_wrapper.set_system_trait( + environment_api_key="api-key", identifier="user-1", trait_key="cohort_x" + ) + + # Then + update_mock.assert_not_called() + put_mock.assert_not_called() + + +def test_set_system_trait__stale_missing_document_read__retries_and_merges( + dynamodb_identity_wrapper: DynamoIdentityWrapper, + mocker: MockerFixture, +) -> None: + # Given + dynamodb_identity_wrapper.put_item( + { + "composite_key": "api-key_user-1", + "identifier": "user-1", + "environment_api_key": "api-key", + "system_traits": {"other": True}, + } + ) + real_get_item = dynamodb_identity_wrapper.table.get_item # type: ignore[union-attr] + stale_responses: typing.Iterator[dict[str, typing.Any]] = iter([{}]) + mocker.patch.object( + dynamodb_identity_wrapper.table, + "get_item", + side_effect=lambda **kwargs: next(stale_responses, real_get_item(**kwargs)), + ) + + # When + dynamodb_identity_wrapper.set_system_trait( + environment_api_key="api-key", identifier="user-1", trait_key="cohort_x" + ) + + # Then + document = real_get_item(Key={"composite_key": "api-key_user-1"})["Item"] + assert document["system_traits"] == {"other": True, "cohort_x": True} + + +def test_set_system_trait__conditional_writes_keep_losing__raises( + dynamodb_identity_wrapper: DynamoIdentityWrapper, + mocker: MockerFixture, +) -> None: + # Given + dynamodb_identity_wrapper.put_item( + { + "composite_key": "api-key_user-1", + "identifier": "user-1", + "environment_api_key": "api-key", + } + ) + real_get_item = dynamodb_identity_wrapper.table.get_item # type: ignore[union-attr] + mocker.patch.object(dynamodb_identity_wrapper.table, "get_item", return_value={}) + + # When + with pytest.raises(SystemTraitWriteRaceError): + dynamodb_identity_wrapper.set_system_trait( + environment_api_key="api-key", identifier="user-1", trait_key="cohort_x" + ) + + # Then + document = real_get_item(Key={"composite_key": "api-key_user-1"})["Item"] + assert "system_traits" not in document + + +def test_set_system_trait__unexpected_client_error__reraises( + dynamodb_identity_wrapper: DynamoIdentityWrapper, + mocker: MockerFixture, +) -> None: + # Given + mocker.patch.object( + dynamodb_identity_wrapper.table, + "put_item", + side_effect=ClientError({"Error": {"Code": "ValidationException"}}, "PutItem"), + ) + + # When + with pytest.raises(ClientError) as exc_info: + dynamodb_identity_wrapper.set_system_trait( + environment_api_key="api-key", identifier="user-1", trait_key="cohort_x" + ) + + # Then + assert exc_info.value.response["Error"]["Code"] == "ValidationException" + + +def test_unset_system_trait__member__removes_only_given_key( + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + dynamodb_identity_wrapper.put_item( + { + "composite_key": "api-key_user-1", + "identifier": "user-1", + "environment_api_key": "api-key", + "identity_traits": [{"trait_key": "plan", "trait_value": "pro"}], + "system_traits": {"cohort_x": True, "other": True}, + } + ) + + # When + dynamodb_identity_wrapper.unset_system_trait( + environment_api_key="api-key", identifier="user-1", trait_key="cohort_x" + ) + + # Then + document = dynamodb_identity_wrapper.get_item("api-key_user-1") + assert document is not None + assert document["system_traits"] == {"other": True} + assert document["identity_traits"] == [{"trait_key": "plan", "trait_value": "pro"}] + + +def test_unset_system_trait__missing_document__no_ghost_document( + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + composite_key = "api-key_never-seen" + + # When + dynamodb_identity_wrapper.unset_system_trait( + environment_api_key="api-key", identifier="never-seen", trait_key="cohort_x" + ) + + # Then + assert dynamodb_identity_wrapper.get_item(composite_key) is None + + +def test_unset_system_trait__unexpected_client_error__reraises( + dynamodb_identity_wrapper: DynamoIdentityWrapper, + mocker: MockerFixture, +) -> None: + # Given + mocker.patch.object( + dynamodb_identity_wrapper.table, + "update_item", + side_effect=ClientError( + {"Error": {"Code": "ValidationException"}}, "UpdateItem" + ), + ) + + # When + with pytest.raises(ClientError) as exc_info: + dynamodb_identity_wrapper.unset_system_trait( + environment_api_key="api-key", identifier="user-1", trait_key="cohort_x" + ) + + # Then + assert exc_info.value.response["Error"]["Code"] == "ValidationException" + + +def test_unset_system_trait__trait_absent__no_error( + dynamodb_identity_wrapper: DynamoIdentityWrapper, +) -> None: + # Given + dynamodb_identity_wrapper.put_item( + { + "composite_key": "api-key_user-1", + "identifier": "user-1", + "environment_api_key": "api-key", + "system_traits": {"other": True}, + } + ) + + # When + dynamodb_identity_wrapper.unset_system_trait( + environment_api_key="api-key", identifier="user-1", trait_key="cohort_x" + ) + + # Then + document = dynamodb_identity_wrapper.get_item("api-key_user-1") + assert document is not None + assert document["system_traits"] == {"other": True} diff --git a/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py b/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py index 32b3f0de2d75..ba2ba78db14c 100644 --- a/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py +++ b/api/tests/unit/util/mappers/test_unit_mappers_dynamodb.py @@ -9,6 +9,7 @@ from environments.dynamodb.constants import ( ENVIRONMENTS_V2_ENVIRONMENT_META_DOCUMENT_KEY, ) +from util.engine_models.identities.models import IdentityModel from util.mappers import dynamodb from util.mappers.engine import map_feature_state_to_engine @@ -154,6 +155,58 @@ def test_map_identity_to_identity_document__valid_identity__returns_expected_doc assert uuid.UUID(result["identity_uuid"]) # type: ignore[arg-type] +def test_map_engine_identity_to_identity_document__system_traits_set__included_in_document() -> ( + None +): + # Given + engine_identity = IdentityModel( + identifier="test_identity", + environment_api_key="api-key", + system_traits={"flagsmith_cohort_2b6d1f5f": True}, + ) + + # When + result = dynamodb.map_engine_identity_to_identity_document(engine_identity) + + # Then + assert result["system_traits"] == {"flagsmith_cohort_2b6d1f5f": True} + + +def test_map_engine_identity_to_identity_document__no_system_traits__key_absent() -> ( + None +): + # Given + engine_identity = IdentityModel( + identifier="test_identity", + environment_api_key="api-key", + ) + + # When + result = dynamodb.map_engine_identity_to_identity_document(engine_identity) + + # Then + assert "system_traits" not in result + + +def test_identity_document__system_traits_set__round_trip_preserves_system_traits() -> ( + None +): + # Given + document = dynamodb.map_engine_identity_to_identity_document( + IdentityModel( + identifier="test_identity", + environment_api_key="api-key", + system_traits={"flagsmith_cohort_2b6d1f5f": True}, + ) + ) + + # When + parsed = IdentityModel.model_validate(document) + + # Then + assert parsed.system_traits == {"flagsmith_cohort_2b6d1f5f": True} + + def test_map_environment_to_environment_v2_document__valid_environment__returns_expected_document( environment: "Environment", feature_state: "FeatureState", diff --git a/api/tests/unit/util/mappers/test_unit_mappers_sdk.py b/api/tests/unit/util/mappers/test_unit_mappers_sdk.py index 825fef2c22d9..aa2af8f3a0d6 100644 --- a/api/tests/unit/util/mappers/test_unit_mappers_sdk.py +++ b/api/tests/unit/util/mappers/test_unit_mappers_sdk.py @@ -3,6 +3,7 @@ import pytest from environments.identities.models import Identity +from util.mappers.engine import map_identity_to_engine from util.mappers.sdk import map_environment_to_sdk_document if TYPE_CHECKING: # pragma: no cover @@ -176,3 +177,26 @@ def test_map_environment_to_sdk_document__identity_overrides_disabled__returns_e "use_identity_composite_key_for_hashing": True, "use_identity_overrides_in_local_eval": False, } + + +def test_map_environment_to_sdk_document__system_traits_set__excluded_from_document( + mocker: "MockerFixture", + environment: "Environment", + identity: Identity, + identity_featurestate: "FeatureState", +) -> None: + # Given + engine_identity = map_identity_to_engine(identity, with_traits=False) + engine_identity.system_traits = {"flagsmith_cohort_2b6d1f5f": True} + mocker.patch( + "util.mappers.sdk.map_identity_to_engine", + return_value=engine_identity, + ) + + # When + result = map_environment_to_sdk_document(environment) + + # Then + assert result["identity_overrides"] == [ + engine_identity.model_dump(exclude={"system_traits"}) + ] diff --git a/api/util/engine_models/identities/models.py b/api/util/engine_models/identities/models.py index ce7173c08309..71518846788e 100644 --- a/api/util/engine_models/identities/models.py +++ b/api/util/engine_models/identities/models.py @@ -7,6 +7,7 @@ from util.engine_models.features.models import FeatureStateModel from util.engine_models.identities.traits.models import TraitModel +from util.engine_models.identities.traits.types import ContextValue from util.engine_models.utils.datetime import utcnow_with_tz from util.engine_models.utils.exceptions import DuplicateFeatureState @@ -42,6 +43,8 @@ class IdentityModel(BaseModel): default_factory=IdentityFeaturesList ) identity_traits: typing.List[TraitModel] = Field(default_factory=list) + # System-owned (e.g. cohort membership); unreachable by SDK and admin trait writes. + system_traits: typing.Optional[typing.Dict[str, ContextValue]] = None identity_uuid: UUID4 = Field(default_factory=uuid.uuid4) django_id: typing.Optional[int] = None diff --git a/api/util/mappers/dynamodb.py b/api/util/mappers/dynamodb.py index dda3ec27bd3e..f59458772629 100644 --- a/api/util/mappers/dynamodb.py +++ b/api/util/mappers/dynamodb.py @@ -53,7 +53,7 @@ TypeAdapter(EnvironmentV2MetaCompressed) ) -_NULLABLE_IDENTITY_KEY_ATTRIBUTES = {"dashboard_alias"} +_NULLABLE_IDENTITY_KEY_ATTRIBUTES = {"dashboard_alias", "system_traits"} def map_environment_to_environment_document( diff --git a/api/util/mappers/sdk.py b/api/util/mappers/sdk.py index 0d9e9865c603..26c23ec214f9 100644 --- a/api/util/mappers/sdk.py +++ b/api/util/mappers/sdk.py @@ -10,13 +10,17 @@ from environments.models import Environment -SDKDocumentValue: TypeAlias = dict[str, "SDKDocumentValue"] | str | bool | None | float +SDKDocumentValue: TypeAlias = ( + dict[str, "SDKDocumentValue"] | list["SDKDocumentValue"] | str | bool | None | float +) SDKDocument: TypeAlias = dict[str, SDKDocumentValue] -SDK_DOCUMENT_EXCLUDE = { - *IDENTITY_INTEGRATIONS_RELATION_NAMES, - "dynatrace_config", - "onboarding_pending", +SDK_DOCUMENT_EXCLUDE: dict[str, bool | dict[str, set[str]]] = { + **dict.fromkeys(IDENTITY_INTEGRATIONS_RELATION_NAMES, True), + "dynatrace_config": True, + "onboarding_pending": True, + # System-owned identity data must never reach local-eval SDKs. + "identity_overrides": {"__all__": {"system_traits"}}, } diff --git a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md index c012f1efc5e7..0c8a6b9066b4 100644 --- a/docs/docs/deployment-self-hosting/observability/_events-catalogue.md +++ b/docs/docs/deployment-self-hosting/observability/_events-catalogue.md @@ -71,6 +71,35 @@ Attributes: - `feature.count` - `organisation.id` +### `cohorts.membership.applied` + +Logged at `info` from: + - `api/cohorts/services.py:61` + +Attributes: + - `adds.count` + - `cohort.id` + - `environment.id` + - `removes.count` + +### `cohorts.membership.apply.skipped` + +Logged at `info` from: + - `api/cohorts/tasks.py:23` + - `api/cohorts/tasks.py:29` + +Attributes: + - `cohort.id` + - `reason` + +### `cohorts.membership.apply.throttled` + +Logged at `warning` from: + - `api/cohorts/tasks.py:37` + +Attributes: + - `cohort.id` + ### `core.encrypted_field.decrypt_failed` Logged at `warning` from: diff --git a/docs/docs/deployment-self-hosting/observability/_metrics-catalogue.md b/docs/docs/deployment-self-hosting/observability/_metrics-catalogue.md index 82c6e248c446..203a40345f6d 100644 --- a/docs/docs/deployment-self-hosting/observability/_metrics-catalogue.md +++ b/docs/docs/deployment-self-hosting/observability/_metrics-catalogue.md @@ -9,6 +9,15 @@ Labels: - `ci_commit_sha` - `version` +### `flagsmith_cohorts_membership_deltas_applied` + +Counter. + +Total number of cohort membership ledger rows transitioned to their applied state after the corresponding identity document write. The `operation` label is either `add` or `remove`. + +Labels: + - `operation` + ### `flagsmith_dynamo_environment_document_compression_ratio` Histogram.