refactor so defaults are done in UserConfig

manage_arkime/commands/cluster_create.py

@@ -22,8 +22,7 @@
 from core.price_report import PriceReport
 import core.compatibility as compat
 from core.capacity_planning import (get_capture_node_capacity_plan, get_viewer_node_capacity_plan, get_ecs_sys_resource_plan, get_os_domain_plan,
-                                    ClusterPlan, VpcPlan, MINIMUM_TRAFFIC, DEFAULT_SPI_DAYS, DEFAULT_REPLICAS, get_capture_vpc_plan,
-                                    S3Plan, DEFAULT_S3_STORAGE_CLASS, DEFAULT_S3_STORAGE_DAYS, DEFAULT_HISTORY_DAYS,
+                                    ClusterPlan, VpcPlan, get_capture_vpc_plan, S3Plan, DEFAULT_S3_STORAGE_CLASS,
                                     CaptureNodesPlan, ViewerNodesPlan, DataNodesPlan, EcsSysResourcePlan, MasterNodesPlan, OSDomainPlan,
                                     get_viewer_vpc_plan)
 import core.versioning as ver
@@ -156,12 +155,6 @@ def _get_previous_user_config(cluster_name: str, aws_provider: AwsClientProvider
     except ssm_ops.ParamDoesNotExist:
         return UserConfig(None, None, None, None, None)
 
-def _not_none(s, d):
-    if s is None:
-        return d
-    else:
-        return s
-
 def _get_next_user_config(cluster_name: str, expected_traffic: float, spi_days: int, history_days: int, replicas: int,
                           pcap_days: int, viewer_prefix_list: str, aws_provider: AwsClientProvider) -> UserConfig:
     # At least one parameter isn't defined
@@ -173,8 +166,11 @@ def _get_next_user_config(cluster_name: str, expected_traffic: float, spi_days:
                 "userConfig",
                 aws_provider
             )
+
+            # Load UserConfig from what was in SSM
             user_config = UserConfig(**stored_config_json)
 
+            # Now replace what was in SSM with any provided new values
             if expected_traffic is not None:
                 user_config.expectedTraffic = expected_traffic
             if spi_days is not None:
@@ -187,18 +183,11 @@ def _get_next_user_config(cluster_name: str, expected_traffic: float, spi_days:
                 user_config.pcapDays = pcap_days
             if viewer_prefix_list is not None:
                 user_config.viewerPrefixList = viewer_prefix_list
-
             return user_config
 
         # Existing configuration doesn't exist, use defaults
         except ssm_ops.ParamDoesNotExist:
-            return UserConfig(_not_none(expected_traffic, MINIMUM_TRAFFIC),
-                              _not_none(spi_days, DEFAULT_SPI_DAYS),
-                              _not_none(history_days, DEFAULT_HISTORY_DAYS),
-                              _not_none(replicas, DEFAULT_REPLICAS),
-                              _not_none(pcap_days,DEFAULT_S3_STORAGE_DAYS),
-                              viewer_prefix_list
-                             )
+            return UserConfig(expected_traffic, spi_days, history_days, replicas, pcap_days, viewer_prefix_list)
     # All of the parameters defined
     else:
         return UserConfig(expected_traffic, spi_days, history_days, replicas, pcap_days, viewer_prefix_list)

manage_arkime/core/user_config.py

@@ -2,6 +2,8 @@
 import logging
 from typing import Dict
 
+from core.capacity_planning import (MINIMUM_TRAFFIC, DEFAULT_SPI_DAYS, DEFAULT_REPLICAS, DEFAULT_S3_STORAGE_DAYS, DEFAULT_HISTORY_DAYS)
+
 logger = logging.getLogger(__name__)
 
 @dataclass
@@ -13,6 +15,30 @@ class UserConfig:
     pcapDays: int
     viewerPrefixList: str = None
 
+    def __init__(self, expectedTraffic: float, spiDays: int, historyDays: int, replicas: int, pcapDays: int, viewerPrefixList: str = None):
+        if (expectedTraffic is None):
+            expectedTraffic = MINIMUM_TRAFFIC
+
+        if (spiDays is None):
+            spiDays = DEFAULT_SPI_DAYS
+
+        if (historyDays is None):
+            historyDays = DEFAULT_HISTORY_DAYS
+
+        if (replicas is None):
+            replicas = DEFAULT_REPLICAS
+
+        if (pcapDays is None):
+            pcapDays = DEFAULT_S3_STORAGE_DAYS
+
+
+        self.expectedTraffic = expectedTraffic
+        self.spiDays = spiDays
+        self.historyDays = historyDays
+        self.replicas = replicas
+        self.pcapDays = pcapDays
+        self.viewerPrefixList = viewerPrefixList
+
     """ Only process fields we still need, this allows us to ignore config no longer used """
     @classmethod
     def from_dict(cls, d):