scripts: add an abstraction for updating the configs. Will be used to update configs dynamically.

Author: guyf-starkware

scripts/prod/restart_all_nodes_together.py

@@ -7,6 +7,7 @@
 import urllib.request
 from update_config_and_restart_nodes_lib import (
     ApolloArgsParserBuilder,
+    ConstConfigValuesUpdater,
     NamespaceAndInstructionArgs,
     RestartStrategy,
     Service,
@@ -137,7 +138,7 @@ def main():
     )
 
     update_config_and_restart_nodes(
-        config_overrides,
+        ConstConfigValuesUpdater(config_overrides),
         namespace_and_instruction_args,
         Service.Core,
         restarter,

scripts/prod/set_node_revert_mode.py

@@ -6,6 +6,7 @@
 import urllib.parse
 from update_config_and_restart_nodes_lib import (
     ApolloArgsParserBuilder,
+    ConstConfigValuesUpdater,
     NamespaceAndInstructionArgs,
     RestartStrategy,
     Service,
@@ -65,7 +66,7 @@ def set_revert_mode(
     )
 
     update_config_and_restart_nodes(
-        config_overrides,
+        ConstConfigValuesUpdater(config_overrides),
         namespace_and_instruction_args,
         Service.Core,
         restarter,

scripts/prod/update_config_and_restart_nodes.py

@@ -8,6 +8,7 @@
 from update_config_and_restart_nodes_lib import (
     ApolloArgsParserBuilder,
     Colors,
+    ConstConfigValuesUpdater,
     RestartStrategy,
     NamespaceAndInstructionArgs,
     Service,
@@ -165,7 +166,10 @@ def main():
     )
 
     update_config_and_restart_nodes(
-        config_overrides, namespace_and_instruction_args, args.service, restarter
+        ConstConfigValuesUpdater(config_overrides),
+        namespace_and_instruction_args,
+        args.service,
+        restarter,
     )
 
 

scripts/prod/update_config_and_restart_nodes_lib.py

@@ -274,6 +274,63 @@ def get_context_list_from_args(
         ]
 
 
+class ConfigValuesUpdater(ABC):
+    """Abstract class for updating configuration values for different service instances."""
+
+    def get_updated_config(self, orig_config_yaml: str, instance_index: int) -> str:
+        """Get updated configuration YAML for a specific instance.
+
+        Args:
+            orig_config_yaml: Original configuration as YAML string
+            instance_index: Index of the instance to update configuration for
+
+        Returns:
+            Updated configuration as YAML string
+        """
+        config, config_data = parse_config_from_yaml(orig_config_yaml)
+        updated_config_data = self.get_updated_config_for_instance(config_data, instance_index)
+        return serialize_config_to_yaml(config, updated_config_data)
+
+    @abstractmethod
+    def get_updated_config_for_instance(
+        self, config_data: dict[str, any], instance_index: int
+    ) -> dict[str, any]:
+        """Get updated configuration data for a specific instance.
+
+        Args:
+            config_data: Current configuration data dictionary
+            instance_index: Index of the instance to update configuration for
+
+        Returns:
+            Updated configuration data dictionary
+        """
+        pass
+
+
+class ConstConfigValuesUpdater(ConfigValuesUpdater):
+    """Concrete implementation that applies constant configuration overrides."""
+
+    def __init__(self, config_overrides: dict[str, any]):
+        """Initialize with configuration overrides.
+
+        Args:
+            config_overrides: Dictionary of configuration keys and values to override
+        """
+        self.config_overrides = config_overrides
+
+    def get_updated_config_for_instance(
+        self, config_data: dict[str, any], instance_index: int
+    ) -> dict[str, any]:
+        """Apply the same configuration overrides to the config data for each instance."""
+        updated_config = config_data.copy()
+
+        for key, value in self.config_overrides.items():
+            print_colored(f"  Overriding config: {key} = {value}")
+            updated_config[key] = value
+
+        return updated_config
+
+
 class ServiceRestarter(ABC):
     """Abstract class for restarting service instances."""
 
@@ -541,22 +598,6 @@ def represent_literal_str(dumper, data):
     return result
 
 
-def update_config_values(
-    config_content: str,
-    config_overrides: dict[str, Any] = None,
-) -> str:
-    """Update configuration values in the YAML content and return the updated YAML"""
-    # Parse the configuration
-    config, config_data = parse_config_from_yaml(config_content)
-
-    for key, value in config_overrides.items():
-        print_colored(f"  Overriding config: {key} = {value}")
-        config_data[key] = value
-
-    # Serialize back to YAML
-    return serialize_config_to_yaml(config, config_data)
-
-
 def normalize_config(config_content: str) -> str:
     """Normalize configuration by parsing and re-serializing without changes.
 
@@ -633,12 +674,12 @@ def apply_configmap(
 
 
 def update_config_and_restart_nodes(
-    config_overrides: dict[str, Any],
+    config_values_updater: ConfigValuesUpdater,
     namespace_and_instruction_args: NamespaceAndInstructionArgs,
     service: Service,
     restarter: ServiceRestarter,
 ) -> None:
-    assert config_overrides is not None, "config_overrides must be provided"
+    assert config_values_updater is not None, "config_values_updater must be provided"
     assert namespace_and_instruction_args.namespace_list is not None, "namespaces must be provided"
 
     if not namespace_and_instruction_args.cluster_list:
@@ -669,7 +710,7 @@ def update_config_and_restart_nodes(
         )
 
         # Update config
-        updated_config = update_config_values(original_config, config_overrides)
+        updated_config = config_values_updater.get_updated_config(original_config, index)
 
         # Store configs
         configs.append({"original": original_config, "updated": updated_config})