feat: Annotate DynamoDB schemas with validation code

src/flagsmith_schemas/constants.py

@@ -0,0 +1,4 @@
+from importlib.util import find_spec
+
+PYDANTIC_INSTALLED = find_spec("pydantic") is not None
+MAX_STRING_FEATURE_STATE_VALUE_LENGTH = 20_000

src/flagsmith_schemas/dynamodb.py

@@ -1,24 +1,35 @@
 """
 The types in this module describe the Edge API's data model.
 They are used to type DynamoDB documents representing Flagsmith entities.
+
+These types can be used with Pydantic for validation and serialization
+when `pydantic` is installed.
+Otherwise, they serve as documentation for the structure of the data stored in DynamoDB.
 """
 
-from typing import Literal
+from typing import Annotated, Literal
 
 from typing_extensions import NotRequired, TypedDict
 
+from flagsmith_schemas.constants import PYDANTIC_INSTALLED
 from flagsmith_schemas.types import (
     ConditionOperator,
-    ContextValue,
     DateTimeStr,
+    DynamoContextValue,
+    DynamoFeatureValue,
     DynamoFloat,
     DynamoInt,
     FeatureType,
-    FeatureValue,
     RuleType,
     UUIDStr,
 )
 
+if PYDANTIC_INSTALLED:
+    from flagsmith_schemas.pydantic_types import (
+        ValidateIdentityFeatureStatesList,
+        ValidateMultivariateFeatureValuesList,
+    )
+
 
 class Feature(TypedDict):
     """Represents a Flagsmith feature, defined at project level."""
@@ -35,7 +46,7 @@ class MultivariateFeatureOption(TypedDict):
 
     id: NotRequired[DynamoInt | None]
     """Unique identifier for the multivariate feature option in Core. This is used by Core UI to display the selected option for an identity override for a multivariate feature."""
-    value: FeatureValue
+    value: DynamoFeatureValue
     """The feature state value that should be served when this option's parent multivariate feature state is selected by the engine."""
 
 
@@ -69,15 +80,15 @@ class FeatureState(TypedDict):
     """The feature that this feature state is for."""
     enabled: bool
     """Whether the feature is enabled or disabled."""
-    feature_state_value: FeatureValue
+    feature_state_value: DynamoFeatureValue
     """The value for this feature state."""
     django_id: NotRequired[DynamoInt | None]
     """Unique identifier for the feature state in Core. If feature state created via Core's `edge-identities` APIs in Core, this can be missing or `None`."""
     featurestate_uuid: NotRequired[UUIDStr]
     """The UUID for this feature state. Should be used if `django_id` is `None`. If not set, should be generated."""
     feature_segment: NotRequired[FeatureSegment | None]
     """Segment override data, if this feature state is for a segment override."""
-    multivariate_feature_state_values: NotRequired[list[MultivariateFeatureStateValue]]
+    multivariate_feature_state_values: "NotRequired[Annotated[list[MultivariateFeatureStateValue], ValidateMultivariateFeatureValuesList]]"
     """List of multivariate feature state values, if this feature state is for a multivariate feature.
 
     Total `percentage_allocation` sum of the child multivariate feature state values must be less or equal to 100.
@@ -89,7 +100,7 @@ class Trait(TypedDict):
 
     trait_key: str
     """Key of the trait."""
-    trait_value: ContextValue
+    trait_value: DynamoContextValue
     """Value of the trait."""
 
 
@@ -275,7 +286,9 @@ class Identity(TypedDict):
     """
     created_date: DateTimeStr
     """Creation timestamp."""
-    identity_features: NotRequired[list[FeatureState]]
+    identity_features: (
+        "NotRequired[Annotated[list[FeatureState], ValidateIdentityFeatureStatesList]]"
+    )
     """List of identity overrides for this identity."""
     identity_traits: list[Trait]
     """List of traits associated with this identity."""

src/flagsmith_schemas/pydantic_types.py

@@ -0,0 +1,22 @@
+from datetime import datetime
+from decimal import Decimal
+from uuid import UUID
+
+from pydantic import AfterValidator, BeforeValidator, ValidateAs
+
+from flagsmith_schemas.validators import (
+    validate_dynamo_feature_state_value,
+    validate_identity_feature_states,
+    validate_multivariate_feature_state_values,
+)
+
+ValidateDecimalAsFloat = ValidateAs(float, lambda v: Decimal(str(v)))
+ValidateDecimalAsInt = ValidateAs(int, lambda v: Decimal(v))
+ValidateStrAsISODateTime = ValidateAs(datetime, lambda dt: dt.isoformat())
+ValidateStrAsUUID = ValidateAs(UUID, str)
+
+ValidateDynamoFeatureStateValue = BeforeValidator(validate_dynamo_feature_state_value)
+ValidateIdentityFeatureStatesList = AfterValidator(validate_identity_feature_states)
+ValidateMultivariateFeatureValuesList = AfterValidator(
+    validate_multivariate_feature_state_values
+)

src/flagsmith_schemas/types.py

@@ -1,36 +1,60 @@
 from decimal import Decimal
-from typing import Literal, TypeAlias
+from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias
 
-DynamoInt: TypeAlias = Decimal
+from flagsmith_schemas.constants import PYDANTIC_INSTALLED
+
+if PYDANTIC_INSTALLED:
+    from flagsmith_schemas.pydantic_types import (
+        ValidateDecimalAsFloat,
+        ValidateDecimalAsInt,
+        ValidateDynamoFeatureStateValue,
+        ValidateStrAsISODateTime,
+        ValidateStrAsUUID,
+    )
+elif not TYPE_CHECKING:
+    # This code runs at runtime when Pydantic is not installed.
+    # We could use PEP 649 strings with `Annotated`, but Pydantic is inconsistent in how it parses them.
+    # Define dummy types instead.
+    ValidateDecimalAsFloat = ...
+    ValidateDecimalAsInt = ...
+    ValidateDynamoFeatureStateValue = ...
+    ValidateStrAsISODateTime = ...
+    ValidateStrAsUUID = ...
+
+
+DynamoInt: TypeAlias = Annotated[Decimal, ValidateDecimalAsInt]
 """An integer value stored in DynamoDB.
 
 DynamoDB represents all numbers as `Decimal`.
 `DynamoInt` indicates that the value should be treated as an integer.
 """
 
-DynamoFloat: TypeAlias = Decimal
+DynamoFloat: TypeAlias = Annotated[Decimal, ValidateDecimalAsFloat]
 """A float value stored in DynamoDB.
 
 DynamoDB represents all numbers as `Decimal`.
 `DynamoFloat` indicates that the value should be treated as a float.
 """
 
-UUIDStr: TypeAlias = str
+UUIDStr: TypeAlias = Annotated[str, ValidateStrAsUUID]
 """A string representing a UUID."""
 
-DateTimeStr: TypeAlias = str
+DateTimeStr: TypeAlias = Annotated[str, ValidateStrAsISODateTime]
 """A string representing a date and time in ISO 8601 format."""
 
 FeatureType = Literal["STANDARD", "MULTIVARIATE"]
 """Represents the type of a Flagsmith feature. Multivariate features include multiple weighted values."""
 
-FeatureValue: TypeAlias = object
+DynamoFeatureValue: TypeAlias = Annotated[
+    DynamoInt | bool | str | None,
+    ValidateDynamoFeatureStateValue,
+]
 """Represents the value of a Flagsmith feature. Can be stored a boolean, an integer, or a string.
 
 The default (SaaS) maximum length for strings is 20000 characters.
 """
 
-ContextValue: TypeAlias = DynamoInt | DynamoFloat | bool | str
+DynamoContextValue: TypeAlias = DynamoInt | DynamoFloat | bool | str
 """Represents a scalar value in the Flagsmith context, e.g., of an identity trait.
 Here's how we store different types:
 - Numeric string values (int, float) are stored as numbers.

src/flagsmith_schemas/validators.py

@@ -0,0 +1,55 @@
+import typing
+from decimal import Decimal
+
+from flagsmith_schemas.constants import MAX_STRING_FEATURE_STATE_VALUE_LENGTH
+
+if typing.TYPE_CHECKING:
+    from flagsmith_schemas.dynamodb import FeatureState, MultivariateFeatureStateValue
+    from flagsmith_schemas.types import DynamoFeatureValue
+
+
+def validate_dynamo_feature_state_value(
+    value: typing.Any,
+) -> "DynamoFeatureValue":
+    if isinstance(value, bool | None):
+        return value
+    if isinstance(value, str):
+        if len(value) > MAX_STRING_FEATURE_STATE_VALUE_LENGTH:
+            raise ValueError(
+                "Dynamo feature state value string length cannot exceed "
+                f"{MAX_STRING_FEATURE_STATE_VALUE_LENGTH} characters "
+                f"(got {len(value)} characters)."
+            )
+        return value
+    if isinstance(value, int):
+        return Decimal(value)
+    return str(value)
+
+
+def validate_multivariate_feature_state_values(
+    values: "list[MultivariateFeatureStateValue]",
+) -> "list[MultivariateFeatureStateValue]":
+    total_percentage = sum(value["percentage_allocation"] for value in values)
+    if total_percentage > 100:
+        raise ValueError(
+            "Total `percentage_allocation` of multivariate feature state values "
+            "cannot exceed 100."
+        )
+    return values
+
+
+def validate_identity_feature_states(
+    values: "list[FeatureState]",
+) -> "list[FeatureState]":
+    seen: set[Decimal] = set()
+
+    for feature_state in values:
+        feature_id = feature_state["feature"]["id"]
+        if feature_id in seen:
+            raise ValueError(
+                f"Feature id={feature_id} cannot have multiple "
+                "feature states for a single identity."
+            )
+        seen.add(feature_id)
+
+    return values