fix(schemas): keep stored Decimals when validating Dynamo documents

src/flagsmith_schemas/pydantic_types.py

@@ -1,17 +1,39 @@
 from datetime import datetime
 from decimal import Decimal
+from functools import partial
+from typing import Any
 from uuid import UUID
 
-from pydantic import AfterValidator, BeforeValidator, ValidateAs
+from pydantic import (
+    AfterValidator,
+    BeforeValidator,
+    ValidateAs,
+    ValidatorFunctionWrapHandler,
+    WrapValidator,
+)
 
 from flagsmith_schemas.validators import (
+    is_integral_decimal,
     validate_dynamo_feature_state_value,
     validate_identity_feature_states,
     validate_multivariate_feature_state_values,
 )
 
+
+def _keep_stored_decimal(
+    value: Any, handler: ValidatorFunctionWrapHandler, *, integral: bool
+) -> Decimal:
+    # DynamoDB reads numbers back as `Decimal`s: keep them as they are.
+    if isinstance(value, Decimal) and (not integral or is_integral_decimal(value)):
+        return value
+    validated: Decimal = handler(value)
+    return validated
+
+
 ValidateDecimalAsFloat = ValidateAs(float, lambda v: Decimal(str(v)))
 ValidateDecimalAsInt = ValidateAs(int, lambda v: Decimal(v))
+KeepStoredFloat = WrapValidator(partial(_keep_stored_decimal, integral=False))
+KeepStoredInt = WrapValidator(partial(_keep_stored_decimal, integral=True))
 ValidateStrAsISODateTime = ValidateAs(datetime, lambda dt: dt.isoformat())
 ValidateStrAsUUID = ValidateAs(UUID, str)
 

src/flagsmith_schemas/types.py

@@ -22,6 +22,8 @@
     from pydantic_core import core_schema
 
     from flagsmith_schemas.pydantic_types import (
+        KeepStoredFloat,
+        KeepStoredInt,
         ValidateDecimalAsFloat,
         ValidateDecimalAsInt,
         ValidateDynamoFeatureStateValue,
@@ -36,6 +38,8 @@
     def WithJsonSchema(_: object) -> object:
         return ...
 
+    KeepStoredFloat = ...
+    KeepStoredInt = ...
     ValidateDecimalAsFloat = ...
     ValidateDecimalAsInt = ...
     ValidateDynamoFeatureStateValue = ...
@@ -76,14 +80,14 @@ def _validate_json_gzipped(data: Any) -> bytes:
             )
 
 
-DynamoInt: TypeAlias = Annotated[Decimal, ValidateDecimalAsInt]
+DynamoInt: TypeAlias = Annotated[Decimal, ValidateDecimalAsInt, KeepStoredInt]
 """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 = Annotated[Decimal, ValidateDecimalAsFloat]
+DynamoFloat: TypeAlias = Annotated[Decimal, ValidateDecimalAsFloat, KeepStoredFloat]
 """A float value stored in DynamoDB.
 
 DynamoDB represents all numbers as `Decimal`.

src/flagsmith_schemas/validators.py

@@ -8,6 +8,14 @@
     from flagsmith_schemas.types import DynamoFeatureValue
 
 
+def is_integral_decimal(value: typing.Any) -> typing.TypeGuard[Decimal]:
+    return (
+        isinstance(value, Decimal)
+        and value.is_finite()
+        and value == value.to_integral_value()
+    )
+
+
 def validate_dynamo_feature_state_value(
     value: typing.Any,
 ) -> "DynamoFeatureValue":
@@ -23,6 +31,8 @@ def validate_dynamo_feature_state_value(
         return value
     if isinstance(value, int):
         return Decimal(value)
+    if is_integral_decimal(value):
+        return value
     return str(value)
 
 

tests/integration/flagsmith_schemas/test_dynamodb.py

@@ -970,6 +970,36 @@ def test_document__validate_json__expected_result(
     assert document == expected_result
 
 
+@pytest.mark.parametrize(
+    ("document_type", "json_data_filename"),
+    [
+        (Environment, "flagsmith_environments.json"),
+        (EnvironmentAPIKey, "flagsmith_environment_api_key.json"),
+        (Identity, "flagsmith_identities.json"),
+        (EnvironmentV2Meta, "flagsmith_environments_v2:_META.json"),
+        (
+            EnvironmentV2IdentityOverride,
+            "flagsmith_environments_v2:identity_override.json",
+        ),
+    ],
+)
+def test_document__validate_python_stored_document__returns_unchanged(
+    request: pytest.FixtureRequest,
+    document_type: type[T],
+    json_data_filename: str,
+) -> None:
+    # Given
+    type_adapter = TypeAdapter(document_type)
+    json_data = request.path.parent.joinpath(f"data/{json_data_filename}").read_text()
+    stored_document = type_adapter.validate_json(json_data)
+
+    # When
+    document = type_adapter.validate_python(stored_document)
+
+    # Then
+    assert document == stored_document
+
+
 def test_type_adapter__identity_duplicate_features__raises_expected(
     mocker: MockerFixture,
 ) -> None:

tests/integration/flagsmith_schemas/test_types.py

@@ -1,9 +1,15 @@
 import gzip
+from decimal import Decimal
 
 import pytest
 from pydantic import TypeAdapter, ValidationError
 
-from flagsmith_schemas.types import DynamoFeatureValue, JsonGzipped
+from flagsmith_schemas.types import (
+    DynamoContextValue,
+    DynamoFeatureValue,
+    DynamoInt,
+    JsonGzipped,
+)
 
 
 def test_dynamo_feature_value__not_int__coerces_to_str() -> None:
@@ -17,6 +23,71 @@ def test_dynamo_feature_value__not_int__coerces_to_str() -> None:
     assert result == "12.34"
 
 
+@pytest.mark.parametrize(
+    ("value", "expected_result"),
+    [
+        pytest.param(Decimal("1234"), Decimal("1234"), id="integer"),
+        pytest.param(Decimal("12.34"), "12.34", id="non_integer"),
+        pytest.param(Decimal("Infinity"), "Infinity", id="infinity"),
+    ],
+)
+def test_dynamo_feature_value__stored_decimal__returns_expected(
+    value: Decimal,
+    expected_result: DynamoFeatureValue,
+) -> None:
+    # Given
+    type_adapter: TypeAdapter[DynamoFeatureValue] = TypeAdapter(DynamoFeatureValue)
+
+    # When
+    result = type_adapter.validate_python(value)
+
+    # Then
+    assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        pytest.param(Decimal("42"), id="integer"),
+        pytest.param(Decimal("1.5"), id="non_integer"),
+    ],
+)
+def test_dynamo_context_value__stored_decimal__returns_unchanged(
+    value: Decimal,
+) -> None:
+    # Given
+    type_adapter: TypeAdapter[DynamoContextValue] = TypeAdapter(DynamoContextValue)
+
+    # When
+    result = type_adapter.validate_python(value)
+
+    # Then
+    assert isinstance(result, Decimal)
+    assert str(result) == str(value)
+
+
+@pytest.mark.parametrize(
+    ("value", "expected_error_type"),
+    [
+        pytest.param(Decimal("1.5"), "int_from_float", id="non_integer"),
+        pytest.param(Decimal("Infinity"), "finite_number", id="infinity"),
+    ],
+)
+def test_dynamo_int__stored_non_integer_decimal__raises_expected(
+    value: Decimal,
+    expected_error_type: str,
+) -> None:
+    # Given
+    type_adapter: TypeAdapter[DynamoInt] = TypeAdapter(DynamoInt)
+
+    # When
+    with pytest.raises(ValidationError) as exc_info:
+        type_adapter.validate_python(value)
+
+    # Then
+    assert exc_info.value.errors()[0]["type"] == expected_error_type
+
+
 def test_dynamo_feature_value__long_string__raises_expected() -> None:
     # Given
     type_adapter: TypeAdapter[DynamoFeatureValue] = TypeAdapter(DynamoFeatureValue)