feat: Support gzip compression for DynamoDB environment documents (#6816)

Author: khvn26

api/app/settings/common.py

@@ -986,21 +986,21 @@
 DEFAULT_ORG_STORE_TRAITS_VALUE = env.bool("DEFAULT_ORG_STORE_TRAITS_VALUE", True)
 
 # DynamoDB table name for storing environment
-ENVIRONMENTS_TABLE_NAME_DYNAMO = env.str("ENVIRONMENTS_TABLE_NAME_DYNAMO", None)
+ENVIRONMENTS_TABLE_NAME_DYNAMO = env.str("ENVIRONMENTS_TABLE_NAME_DYNAMO", "")
 
 # V2 was created to improve storage over overrides data.
-ENVIRONMENTS_V2_TABLE_NAME_DYNAMO = env.str("ENVIRONMENTS_V2_TABLE_NAME_DYNAMO", None)
+ENVIRONMENTS_V2_TABLE_NAME_DYNAMO = env.str("ENVIRONMENTS_V2_TABLE_NAME_DYNAMO", "")
 
 # DynamoDB table name for storing identities
-IDENTITIES_TABLE_NAME_DYNAMO = env.str("IDENTITIES_TABLE_NAME_DYNAMO", None)
+IDENTITIES_TABLE_NAME_DYNAMO = env.str("IDENTITIES_TABLE_NAME_DYNAMO", "")
 
 # DynamoDB table name for storing environment api keys
 ENVIRONMENTS_API_KEY_TABLE_NAME_DYNAMO = env.str(
-    "ENVIRONMENTS_API_KEY_TABLE_NAME_DYNAMO", None
+    "ENVIRONMENTS_API_KEY_TABLE_NAME_DYNAMO", ""
 )
 
 # DynamoDB table name for storing project metadata(currently only used for identity migration)
-PROJECT_METADATA_TABLE_NAME_DYNAMO = env.str("PROJECT_METADATA_TABLE_NAME_DYNAMO", None)
+PROJECT_METADATA_TABLE_NAME_DYNAMO = env.str("PROJECT_METADATA_TABLE_NAME_DYNAMO", "")
 
 # Front end environment variables
 API_URL = env("API_URL", default="/api/v1/")

api/environments/dynamodb/constants.py

@@ -8,3 +8,30 @@
 
 DYNAMODB_MAX_BATCH_WRITE_ITEM_COUNT = 25
 IDENTITIES_PAGINATION_LIMIT = 1000
+
+# DynamoDB max item size is 400 KB (409,600 bytes).
+DOCUMENT_SIZE_HISTOGRAM_BUCKETS = (
+    1_000,
+    5_000,
+    10_000,
+    25_000,
+    50_000,
+    100_000,
+    200_000,
+    300_000,
+    409_600,
+)
+
+COMPRESSION_RATIO_HISTOGRAM_BUCKETS = (
+    0.05,
+    0.1,
+    0.15,
+    0.2,
+    0.25,
+    0.3,
+    0.4,
+    0.5,
+    0.6,
+    0.8,
+    1.0,
+)

api/environments/dynamodb/utils.py

@@ -1,3 +1,33 @@
+import json
+from collections.abc import Mapping
+from decimal import Decimal
+from typing import Any
+
+
+def estimate_document_size(document: Mapping[str, Any]) -> int:
+    """Estimate the size of a DynamoDB item in bytes.
+
+    Uses compact JSON serialization with type-aware encoding
+    for Decimal (numeric) and bytes (binary) values.
+    """
+    return len(
+        json.dumps(document, default=_json_default, separators=(",", ":")).encode()
+    )
+
+
+def _json_default(obj: object) -> Any:
+    if isinstance(obj, bytes):
+        # Binary values: use raw byte length as a placeholder string
+        # to approximate DynamoDB Binary (B) type storage.
+        return "0" * len(obj)
+    if isinstance(obj, Decimal):
+        # Convert to native numeric types for compact JSON serialization.
+        if obj == obj.to_integral_value():
+            return int(obj)
+        return float(obj)
+    return str(obj)
+
+
 def get_environments_v2_identity_override_document_key(
     feature_id: int | None = None,
     identity_uuid: str | None = None,

api/environments/dynamodb/wrappers/base.py

@@ -17,15 +17,13 @@
 
     DynamoDBOutput = QueryOutputTableTypeDef | ScanOutputTableTypeDef
 
-    P = typing.ParamSpec("P")
-
 # Avoid `decimal.Rounded` when reading large numbers
 # See https://github.com/boto/boto3/issues/2500
 boto3.dynamodb.types.DYNAMODB_CONTEXT = Context(prec=100)
 
 
 class BaseDynamoWrapper:
-    table_name: str = None  # type: ignore[assignment]
+    table_name: str = ""
 
     def __init__(self) -> None:
         self._table: typing.Optional["Table"] = None
@@ -39,29 +37,30 @@ def table(self) -> typing.Optional["Table"]:
     def get_table_name(self) -> str:
         return self.table_name
 
-    def get_table(self) -> typing.Optional["Table"]:  # type: ignore[return]
+    def get_table(self) -> "Table | None":
         if table_name := self.get_table_name():
             return boto3.resource("dynamodb", config=Config(tcp_keepalive=True)).Table(
                 table_name
             )
+        return None
 
     @property
     def is_enabled(self) -> bool:
         return self.table is not None
 
-    def _iter_all_items(  # type: ignore[valid-type]
+    def _iter_all_items(
         self,
-        response_getter_method: "typing.Callable[[P], DynamoDBOutput]",  # type: ignore[valid-type]
-        **kwargs: "P.kwargs",
+        response_getter_method: "typing.Callable[..., DynamoDBOutput]",
+        **kwargs: typing.Any,
     ) -> typing.Generator[dict[str, "TableAttributeValueTypeDef"], None, None]:
         response_getter = partial(response_getter_method, **kwargs)
         set_context(
             "dynamodb",
-            {"table_name": self.table_name, **kwargs},
+            {"table_name": self.get_table_name(), **kwargs},
         )
 
         while True:
-            query_response = response_getter()  # type: ignore[call-arg]
+            query_response = response_getter()
 
             for item in query_response["Items"]:
                 yield item
@@ -73,17 +72,19 @@ def _iter_all_items(  # type: ignore[valid-type]
             response_getter.keywords["ExclusiveStartKey"] = last_evaluated_key
             set_context(
                 "dynamodb",
-                {"table_name": self.table_name, **response_getter.keywords},
+                {"table_name": self.get_table_name(), **response_getter.keywords},
             )
 
     def scan_iter_all_items(
         self,
         **kwargs: typing.Any,
     ) -> typing.Generator[dict[str, "TableAttributeValueTypeDef"], None, None]:
-        return self._iter_all_items(self.table.scan, **kwargs)  # type: ignore[arg-type,union-attr]
+        assert self.table is not None
+        return self._iter_all_items(self.table.scan, **kwargs)
 
     def query_iter_all_items(
         self,
         **kwargs: typing.Any,
     ) -> typing.Generator[dict[str, "TableAttributeValueTypeDef"], None, None]:
-        return self._iter_all_items(self.table.query, **kwargs)  # type: ignore[arg-type,union-attr]
+        assert self.table is not None
+        return self._iter_all_items(self.table.query, **kwargs)

api/environments/dynamodb/wrappers/environment_api_key_wrapper.py

@@ -13,7 +13,7 @@
 
 
 class DynamoEnvironmentAPIKeyWrapper(BaseDynamoWrapper):
-    table_name = settings.ENVIRONMENTS_API_KEY_TABLE_NAME_DYNAMO  # type: ignore[assignment]
+    table_name = settings.ENVIRONMENTS_API_KEY_TABLE_NAME_DYNAMO
 
     def write_api_key(self, api_key: "EnvironmentAPIKey"):  # type: ignore[no-untyped-def]
         self.write_api_keys([api_key])

api/environments/dynamodb/wrappers/environment_wrapper.py

@@ -1,9 +1,12 @@
+import abc
 import typing
 from typing import Any, Iterable
 
+import structlog
 from boto3.dynamodb.conditions import Key
 from django.conf import settings
 from django.core.exceptions import ObjectDoesNotExist
+from django.db.models import prefetch_related_objects
 
 from environments.dynamodb.constants import (
     DYNAMODB_MAX_BATCH_WRITE_ITEM_COUNT,
@@ -12,9 +15,17 @@
 )
 from environments.dynamodb.types import IdentityOverridesV2Changeset
 from environments.dynamodb.utils import (
+    estimate_document_size,
     get_environments_v2_identity_override_document_key,
 )
+from environments.metrics import (
+    flagsmith_dynamo_environment_document_compression_ratio,
+    flagsmith_dynamo_environment_document_size_bytes,
+)
+from integrations.flagsmith.client import get_client
 from util.mappers import (
+    map_environment_to_compressed_environment_document,
+    map_environment_to_compressed_environment_v2_document,
     map_environment_to_environment_document,
     map_environment_to_environment_v2_document,
     map_identity_override_to_identity_override_document,
@@ -27,26 +38,82 @@
     from mypy_boto3_dynamodb.type_defs import QueryInputRequestTypeDef
 
     from environments.models import Environment
+    from util.dataclasses import CompressedEnvironmentDocument
+
+logger = structlog.get_logger("dynamodb")
 
 
-class BaseDynamoEnvironmentWrapper(BaseDynamoWrapper):
+class BaseDynamoEnvironmentWrapper(BaseDynamoWrapper, abc.ABC):
     def write_environment(self, environment: "Environment") -> None:
         self.write_environments([environment])
 
     def write_environments(self, environments: Iterable["Environment"]) -> None:
-        raise NotImplementedError()
+        self._write_environments(environments)
+
+    @abc.abstractmethod
+    def _map_environment_document(
+        self,
+        environment: "Environment",
+    ) -> dict[str, Any]: ...
+
+    @abc.abstractmethod
+    def _map_compressed_environment_document(
+        self,
+        environment: "Environment",
+    ) -> "CompressedEnvironmentDocument": ...
+
+    def _write_environments(self, environments: Iterable["Environment"]) -> None:
+        flagsmith_client = get_client("local", local_eval=True)
+        prefetch_related_objects(
+            environments,
+            "project__organisation",
+            "project__organisation__subscription",
+        )
+
+        assert self.table
+        with self.table.batch_writer() as writer:
+            for environment in environments:
+                organisation = environment.project.organisation
+                if flagsmith_client.get_identity_flags(
+                    organisation.flagsmith_identifier,
+                    traits=organisation.flagsmith_on_flagsmith_api_traits,
+                ).is_feature_enabled("compress_dynamo_documents"):
+                    result = self._map_compressed_environment_document(environment)
+                    writer.put_item(Item=result.document)
+
+                    flagsmith_dynamo_environment_document_size_bytes.labels(
+                        table=self.get_table_name(),
+                        compressed="true",
+                    ).observe(result.compressed_size_bytes)
+                    flagsmith_dynamo_environment_document_compression_ratio.labels(
+                        table=self.get_table_name(),
+                    ).observe(result.compression_ratio)
+                    logger.info(
+                        "environment-document-compressed",
+                        environment_id=environment.id,
+                        environment_api_key=environment.api_key,
+                    )
+                else:
+                    item = self._map_environment_document(environment)
+                    writer.put_item(Item=item)
+
+                    flagsmith_dynamo_environment_document_size_bytes.labels(
+                        table=self.get_table_name(),
+                        compressed="false",
+                    ).observe(estimate_document_size(item))
 
 
 class DynamoEnvironmentWrapper(BaseDynamoEnvironmentWrapper):
     def get_table_name(self) -> str | None:  # type: ignore[override]
         return settings.ENVIRONMENTS_TABLE_NAME_DYNAMO
 
-    def write_environments(self, environments: Iterable["Environment"]):  # type: ignore[no-untyped-def]
-        with self.table.batch_writer() as writer:  # type: ignore[union-attr]
-            for environment in environments:
-                writer.put_item(
-                    Item=map_environment_to_environment_document(environment),
-                )
+    def _map_environment_document(self, environment: "Environment") -> dict[str, Any]:
+        return map_environment_to_environment_document(environment)
+
+    def _map_compressed_environment_document(
+        self, environment: "Environment"
+    ) -> "CompressedEnvironmentDocument":
+        return map_environment_to_compressed_environment_document(environment)
 
     def get_item(self, api_key: str) -> dict:  # type: ignore[type-arg]
         try:
@@ -59,7 +126,7 @@ def delete_environment(self, api_key: str) -> None:
 
 
 class DynamoEnvironmentV2Wrapper(BaseDynamoEnvironmentWrapper):
-    def get_table_name(self) -> str | None:  # type: ignore[override]
+    def get_table_name(self) -> str:
         return settings.ENVIRONMENTS_V2_TABLE_NAME_DYNAMO
 
     def get_identity_overrides_by_environment_id(
@@ -122,12 +189,14 @@ def update_identity_overrides(
                         ),
                     )
 
-    def write_environments(self, environments: Iterable["Environment"]) -> None:
-        with self.table.batch_writer() as writer:  # type: ignore[union-attr]
-            for environment in environments:
-                writer.put_item(
-                    Item=map_environment_to_environment_v2_document(environment),
-                )
+    def _map_environment_document(self, environment: "Environment") -> dict[str, Any]:
+        return map_environment_to_environment_v2_document(environment)
+
+    def _map_compressed_environment_document(
+        self,
+        environment: "Environment",
+    ) -> "CompressedEnvironmentDocument":
+        return map_environment_to_compressed_environment_v2_document(environment)
 
     def delete_environment(self, environment_id: int):  # type: ignore[no-untyped-def]
         environment_id = str(environment_id)  # type: ignore[assignment]

api/environments/metrics.py

@@ -1,5 +1,10 @@
 import prometheus_client
 
+from environments.dynamodb.constants import (
+    COMPRESSION_RATIO_HISTOGRAM_BUCKETS,
+    DOCUMENT_SIZE_HISTOGRAM_BUCKETS,
+)
+
 CACHE_HIT = "CACHE_HIT"
 CACHE_MISS = "CACHE_MISS"
 
@@ -8,3 +13,17 @@
     "Results of cache retrieval for environment document. `result` label is either `hit` or `miss`.",
     ["result"],
 )
+
+flagsmith_dynamo_environment_document_size_bytes = prometheus_client.Histogram(
+    "flagsmith_dynamo_environment_document_size_bytes",
+    "Size of environment documents written to DynamoDB.",
+    ["table", "compressed"],
+    buckets=DOCUMENT_SIZE_HISTOGRAM_BUCKETS,
+)
+
+flagsmith_dynamo_environment_document_compression_ratio = prometheus_client.Histogram(
+    "flagsmith_dynamo_environment_document_compression_ratio",
+    "Compression ratio (compressed_size / uncompressed_size) of environment documents.",
+    ["table"],
+    buckets=COMPRESSION_RATIO_HISTOGRAM_BUCKETS,
+)