feat(schemas): Support partially gzip-compressed environment documents (#157)

Co-authored-by: Evandro Myller <[email protected]>

Author: khvn26

pyproject.toml

@@ -26,6 +26,7 @@ optional-dependencies = { test-tools = [
     "django-health-check",
     "prometheus-client (>=0.0.16)",
 ], flagsmith-schemas = [
+    "simplejson",
     "typing_extensions",
     "flagsmith-flag-engine>6",
 ] }

src/flagsmith_schemas/dynamodb.py

@@ -20,6 +20,7 @@
     DynamoFloat,
     DynamoInt,
     FeatureType,
+    JsonGzipped,
     UUIDStr,
 )
 
@@ -200,19 +201,14 @@ class Webhook(TypedDict):
     """Secret used to sign webhook payloads."""
 
 
-class _EnvironmentFields(TypedDict):
+class _EnvironmentBaseFields(TypedDict):
     """Common fields for Environment documents."""
 
     name: NotRequired[str]
     """Environment name. Defaults to an empty string if not set."""
     updated_at: NotRequired[DateTimeStr | None]
     """Last updated timestamp. If not set, current timestamp should be assumed."""
 
-    project: Project
-    """Project-specific data for this environment."""
-    feature_states: list[FeatureState]
-    """List of feature states representing the environment defaults."""
-
     allow_client_traits: NotRequired[bool]
     """Whether the SDK API should allow clients to set traits for this environment. Identical to project-level's `persist_trait_data` setting. Defaults to `True`."""
     hide_sensitive_data: NotRequired[bool]
@@ -240,7 +236,52 @@ class _EnvironmentFields(TypedDict):
     """Webhook configuration."""
 
 
-### Root document schemas below. Indexed fields are marked as **INDEXED** in the docstrings. ###
+class _EnvironmentV1Fields(TypedDict):
+    """Common fields for environment documents in `flagsmith_environments`."""
+
+    api_key: str
+    """Public client-side API key for the environment. **INDEXED**."""
+    id: DynamoInt
+    """Unique identifier for the environment in Core."""
+
+
+class _EnvironmentV2MetaFields(TypedDict):
+    """Common fields for environment documents in `flagsmith_environments_v2`."""
+
+    environment_id: str
+    """Unique identifier for the environment in Core. Same as `Environment.id`, but string-typed to reduce coupling with Core's type definitions **INDEXED**."""
+    environment_api_key: str
+    """Public client-side API key for the environment. **INDEXED**."""
+    document_key: Literal["_META"]
+    """The fixed document key for the environment v2 document. Always `"_META"`. **INDEXED**."""
+
+    id: DynamoInt
+    """Unique identifier for the environment in Core. Exists for compatibility with the API environment document schema."""
+
+
+class _EnvironmentBaseFieldsUncompressed(TypedDict):
+    """Common fields for uncompressed environment documents."""
+
+    project: Project
+    """Project-specific data for this environment."""
+    feature_states: list[FeatureState]
+    """List of feature states representing the environment defaults."""
+    compressed: NotRequired[Literal[False]]
+    """Either `False` or absent to indicate the data is uncompressed."""
+
+
+class _EnvironmentBaseFieldsCompressed(TypedDict):
+    """Common fields for compressed environment documents."""
+
+    project: JsonGzipped[Project]
+    """Project-specific data for this environment. **COMPRESSED**."""
+    feature_states: JsonGzipped[list[FeatureState]]
+    """List of feature states representing the environment defaults. **COMPRESSED**."""
+    compressed: Literal[True]
+    """Always `True` to indicate the data is compressed."""
+
+
+### Root document schemas below. Indexed fields are marked as **INDEXED** in the docstrings. Compressed fields are marked as **COMPRESSED**. ###
 
 
 class EnvironmentAPIKey(TypedDict):
@@ -295,33 +336,50 @@ class Identity(TypedDict):
     """Unique identifier for the identity in Core. If identity created via Core's `edge-identities` API, this can be missing or `None`."""
 
 
-class Environment(_EnvironmentFields):
+class Environment(
+    _EnvironmentBaseFieldsUncompressed,
+    _EnvironmentV1Fields,
+    _EnvironmentBaseFields,
+):
     """Represents a Flagsmith environment. Carries all necessary data for flag evaluation within the environment.
 
     **DynamoDB table**: `flagsmith_environments`
     """
 
-    api_key: str
-    """Public client-side API key for the environment. **INDEXED**."""
-    id: DynamoInt
-    """Unique identifier for the environment in Core."""
 
+class EnvironmentCompressed(
+    _EnvironmentBaseFieldsCompressed,
+    _EnvironmentV1Fields,
+    _EnvironmentBaseFields,
+):
+    """Represents a Flagsmith environment. Carries all necessary data for flag evaluation within the environment.
+    Has compressed fields.
+
+    **DynamoDB table**: `flagsmith_environments`
+    """
 
-class EnvironmentV2Meta(_EnvironmentFields):
+
+class EnvironmentV2Meta(
+    _EnvironmentBaseFieldsUncompressed,
+    _EnvironmentV2MetaFields,
+    _EnvironmentBaseFields,
+):
     """Represents a Flagsmith environment. Carries all necessary data for flag evaluation within the environment.
 
     **DynamoDB table**: `flagsmith_environments_v2`
     """
 
-    environment_id: str
-    """Unique identifier for the environment in Core. Same as `Environment.id`, but string-typed to reduce coupling with Core's type definitions **INDEXED**."""
-    environment_api_key: str
-    """Public client-side API key for the environment. **INDEXED**."""
-    document_key: Literal["_META"]
-    """The fixed document key for the environment v2 document. Always `"_META"`. **INDEXED**."""
 
-    id: DynamoInt
-    """Unique identifier for the environment in Core. Exists for compatibility with the API environment document schema."""
+class EnvironmentV2MetaCompressed(
+    _EnvironmentBaseFieldsCompressed,
+    _EnvironmentV2MetaFields,
+    _EnvironmentBaseFields,
+):
+    """Represents a Flagsmith environment. Carries all necessary data for flag evaluation within the environment.
+    Has compressed fields.
+
+    **DynamoDB table**: `flagsmith_environments_v2`
+    """
 
 
 class EnvironmentV2IdentityOverride(TypedDict):

src/flagsmith_schemas/types.py

@@ -1,10 +1,25 @@
 from decimal import Decimal
-from typing import TYPE_CHECKING, Annotated, Literal, TypeAlias
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    Generic,
+    Literal,
+    SupportsBytes,
+    TypeAlias,
+    TypeVar,
+    get_args,
+)
 
 from flagsmith_schemas.constants import PYDANTIC_INSTALLED
 
 if PYDANTIC_INSTALLED:
-    from pydantic import WithJsonSchema
+    from pydantic import (
+        GetCoreSchemaHandler,
+        TypeAdapter,
+        WithJsonSchema,
+    )
+    from pydantic_core import core_schema
 
     from flagsmith_schemas.pydantic_types import (
         ValidateDecimalAsFloat,
@@ -13,6 +28,7 @@
         ValidateStrAsISODateTime,
         ValidateStrAsUUID,
     )
+    from flagsmith_schemas.utils import json_gzip
 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.
@@ -26,6 +42,39 @@ def WithJsonSchema(_: object) -> object:
     ValidateStrAsISODateTime = ...
     ValidateStrAsUUID = ...
 
+T = TypeVar("T")
+
+
+class DynamoBinary(SupportsBytes):
+    """boto3's wrapper type for bytes stored in DynamoDB."""
+
+    value: bytes | bytearray
+
+
+class JsonGzipped(DynamoBinary, Generic[T]):
+    """A gzipped JSON blob representing a value of type `T`."""
+
+    if PYDANTIC_INSTALLED:
+
+        @classmethod
+        def __get_pydantic_core_schema__(
+            cls,
+            source_type: "type[JsonGzipped[T]]",
+            handler: GetCoreSchemaHandler,
+        ) -> core_schema.CoreSchema:
+            _adapter: TypeAdapter[T] = TypeAdapter(get_args(source_type)[0])
+
+            def _validate_json_gzipped(data: Any) -> bytes:
+                return json_gzip(_adapter.validate_python(data))
+
+            # We're returning bytes here for two reasons:
+            # 1. boto3.dynamodb seems to expect bytes as input for Binary columns.
+            # 2. We want to avoid having boto3 as a dependency.
+            return core_schema.no_info_before_validator_function(
+                _validate_json_gzipped,
+                core_schema.bytes_schema(strict=False),
+            )
+
 
 DynamoInt: TypeAlias = Annotated[Decimal, ValidateDecimalAsInt]
 """An integer value stored in DynamoDB.

src/flagsmith_schemas/utils.py

@@ -0,0 +1,15 @@
+import gzip
+import typing
+
+import simplejson as json
+
+
+def json_gzip(value: typing.Any) -> bytes:
+    return gzip.compress(
+        json.dumps(
+            value,
+            separators=(",", ":"),
+            sort_keys=True,
+        ).encode("utf-8"),
+        mtime=0,
+    )

tests/integration/flagsmith_schemas/data/flagsmith_environments.json

@@ -257,28 +257,28 @@
                     "multivariate_feature_option": {
                         "value": "baz"
                     },
-                    "percentage_allocation": 30
+                    "percentage_allocation": 30.0
                 },
                 {
                     "id": 3402,
                     "multivariate_feature_option": {
                         "value": "bar"
                     },
-                    "percentage_allocation": 30
+                    "percentage_allocation": 30.0
                 },
                 {
                     "id": 3405,
                     "multivariate_feature_option": {
                         "value": 1
                     },
-                    "percentage_allocation": 0
+                    "percentage_allocation": 0.0
                 },
                 {
                     "id": 3406,
                     "multivariate_feature_option": {
                         "value": true
                     },
-                    "percentage_allocation": 0
+                    "percentage_allocation": 0.0
                 }
             ],
             "django_id": 78986,