diff --git a/service_capacity_modeling/models/org/netflix/entity.py b/service_capacity_modeling/models/org/netflix/entity.py index b52720f..4bc311a 100644 --- a/service_capacity_modeling/models/org/netflix/entity.py +++ b/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), )