update the entity model to provision crdb+kv

service_capacity_modeling/models/org/netflix/entity.py

@@ -2,7 +2,6 @@
 from typing import Callable
 from typing import Dict
 from typing import Optional
-from typing import Sequence
 from typing import Tuple
 
 from .stateless_java import nflx_java_app_capacity_model
@@ -60,6 +59,16 @@ def extra_model_arguments_schema() -> Dict[str, Any]:
     def compose_with(
         user_desires: CapacityDesires, extra_model_arguments: Dict[str, Any]
     ) -> Tuple[Tuple[str, Callable[[CapacityDesires], CapacityDesires]], ...]:
+        def _modify_crdb_desires(
+            user_desires: CapacityDesires,
+        ) -> CapacityDesires:
+            relaxed = user_desires.copy(deep=True)
+            item_count = relaxed.data_shape.estimated_state_item_count
+            if item_count is None:
+                item_count = relaxed.data_shape.estimated_state_size_gib.scale(1000)
+            relaxed.data_shape.estimated_state_size_gib = item_count.scale(256)
+            return relaxed
+
         def _modify_elasticsearch_desires(
             user_desires: CapacityDesires,
         ) -> CapacityDesires:
@@ -70,7 +79,8 @@ def _modify_elasticsearch_desires(
             return relaxed
 
         return (
-            ("org.netflix.cassandra", lambda x: x),
+            ("org.netflix.cockroachdb", _modify_crdb_desires),
+            ("org.netflix.key-value", lambda x: x),
             ("org.netflix.elasticsearch", _modify_elasticsearch_desires),
         )