Fix initial segment fetch

Author: chillaq

splitio/client/factory.py

@@ -564,7 +564,7 @@ def _build_in_memory_factory(api_key, cfg, sdk_url=None, events_url=None,  # pyl
 
     synchronizers = SplitSynchronizers(
         SplitSynchronizer(apis['splits'], storages['splits'], storages['rule_based_segments']),
-        SegmentSynchronizer(apis['segments'], storages['splits'], storages['segments']),
+        SegmentSynchronizer(apis['segments'], storages['splits'], storages['segments'], storages['rule_based_segments']),
         ImpressionSynchronizer(apis['impressions'], storages['impressions'],
                                cfg['impressionsBulkSize']),
         EventSynchronizer(apis['events'], storages['events'], cfg['eventsBulkSize']),

splitio/storage/inmemmory.py

@@ -233,7 +233,7 @@ def contains(self, segment_names):
 
     def fetch_many(self, segment_names):
         return {rb_segment_name: self.get(rb_segment_name) for rb_segment_name in segment_names}
-
+        
 class InMemoryRuleBasedSegmentStorageAsync(RuleBasedSegmentsStorage):
     """InMemory implementation of a feature flag storage base."""    
     def __init__(self):

splitio/sync/segment.py

@@ -10,6 +10,7 @@
 from splitio.util.backoff import Backoff
 from splitio.optional.loaders import asyncio, aiofiles
 from splitio.sync import util
+from splitio.util.storage_helper import get_standard_segment_names_in_rbs_storage
 from splitio.optional.loaders import asyncio
 
 _LOGGER = logging.getLogger(__name__)
@@ -22,7 +23,7 @@
 
 
 class SegmentSynchronizer(object):
-    def __init__(self, segment_api, feature_flag_storage, segment_storage):
+    def __init__(self, segment_api, feature_flag_storage, segment_storage, rule_based_segment_storage):
         """
         Class constructor.
 
@@ -39,6 +40,7 @@ def __init__(self, segment_api, feature_flag_storage, segment_storage):
         self._api = segment_api
         self._feature_flag_storage = feature_flag_storage
         self._segment_storage = segment_storage
+        self._rule_based_segment_storage = rule_based_segment_storage
         self._worker_pool = workerpool.WorkerPool(_MAX_WORKERS, self.synchronize_segment)
         self._worker_pool.start()
         self._backoff = Backoff(
@@ -182,8 +184,11 @@ def synchronize_segments(self, segment_names = None, dont_wait = False):
         """
         if segment_names is None:
             segment_names = self._feature_flag_storage.get_segment_names()
+            segment_names.update(get_standard_segment_names_in_rbs_storage(self._rule_based_segment_storage))
 
         for segment_name in segment_names:
+            _LOGGER.debug("Adding segment name to sync worker")
+            _LOGGER.debug(segment_name)
             self._worker_pool.submit_work(segment_name)
         if (dont_wait):
             return True

splitio/sync/split.py

@@ -139,7 +139,7 @@ def _fetch_until(self, fetch_options, till=None, rbs_till=None):
             rbs_segment_list = update_rule_based_segment_storage(self._rule_based_segment_storage, fetched_rule_based_segments, feature_flag_changes.get('rbs')['t'], self._api.clear_storage)
 
             fetched_feature_flags = [(splits.from_raw(feature_flag)) for feature_flag in feature_flag_changes.get('ff').get('d', [])]
-            segment_list = update_feature_flag_storage(self._feature_flag_storage, fetched_feature_flags, feature_flag_changes.get('ff')['t'], self._api.clear_storage)
+            segment_list.update(update_feature_flag_storage(self._feature_flag_storage, fetched_feature_flags, feature_flag_changes.get('ff')['t'], self._api.clear_storage))
             segment_list.update(rbs_segment_list)
 
             if feature_flag_changes.get('ff')['t'] == feature_flag_changes.get('ff')['s'] and feature_flag_changes.get('rbs')['t'] == feature_flag_changes.get('rbs')['s']:

splitio/util/storage_helper.py

@@ -70,6 +70,21 @@ def update_rule_based_segment_storage(rule_based_segment_storage, rule_based_seg
 def _get_segment_names(excluded_segments):
     return [excluded_segment.name for excluded_segment in excluded_segments]
 
+def get_standard_segment_names_in_rbs_storage(rule_based_segment_storage):
+    """
+    Retrieve a list of all standard segments names.
+
+    :return: Set of segment names.
+    :rtype: Set(str)
+    """
+    segment_list = set()
+    for rb_segment in rule_based_segment_storage.get_segment_names():
+        rb_segment_obj = rule_based_segment_storage.get(rb_segment)
+        segment_list.update(set(_get_segment_names(rb_segment_obj.excluded.get_excluded_segments())))
+        segment_list.update(rb_segment_obj.get_condition_segment_names())
+        
+    return segment_list
+    
 async def update_feature_flag_storage_async(feature_flag_storage, feature_flags, change_number, clear_storage=False):
     """
     Update feature flag storage from given list of feature flags while checking the flag set logic