updated push splitworker

Author: chillaq

splitio/push/workers.py

@@ -12,7 +12,7 @@
 from splitio.models.telemetry import UpdateFromSSE
 from splitio.push.parser import UpdateType
 from splitio.optional.loaders import asyncio
-
+from splitio.util.storage_helper import update_feature_flag_storage, update_feature_flag_storage_async
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -218,17 +218,13 @@ def _run(self):
             try:
                 if self._check_instant_ff_update(event):
                     try:
-                        new_split = from_raw(json.loads(self._get_feature_flag_definition(event)))
-                        if new_split.status == Status.ACTIVE:
-                            self._feature_flag_storage.put(new_split)
-                            _LOGGER.debug('Feature flag %s is updated', new_split.name)
-                            for segment_name in new_split.get_segment_names():
-                                if self._segment_storage.get(segment_name) is None:
-                                    _LOGGER.debug('Fetching new segment %s', segment_name)
-                                    self._segment_handler(segment_name, event.change_number)
-                        else:
-                            self._feature_flag_storage.remove(new_split.name)
-                        self._feature_flag_storage.set_change_number(event.change_number)
+                        new_feature_flag = from_raw(json.loads(self._get_feature_flag_definition(event)))
+                        segment_list = update_feature_flag_storage(self._feature_flag_storage, [new_feature_flag], event.change_number)
+                        for segment_name in segment_list:
+                            if self._segment_storage.get(segment_name) is None:
+                                _LOGGER.debug('Fetching new segment %s', segment_name)
+                                self._segment_handler(segment_name, event.change_number)
+
                         self._telemetry_runtime_producer.record_update_from_sse(UpdateFromSSE.SPLIT_UPDATE)
                         continue
                     except Exception as e:
@@ -318,17 +314,13 @@ async def _run(self):
             try:
                 if await self._check_instant_ff_update(event):
                     try:
-                        new_split = from_raw(json.loads(self._get_feature_flag_definition(event)))
-                        if new_split.status == Status.ACTIVE:
-                            await self._feature_flag_storage.put(new_split)
-                            _LOGGER.debug('Feature flag %s is updated', new_split.name)
-                            for segment_name in new_split.get_segment_names():
-                                if await self._segment_storage.get(segment_name) is None:
-                                    _LOGGER.debug('Fetching new segment %s', segment_name)
-                                    await self._segment_handler(segment_name, event.change_number)
-                        else:
-                            await self._feature_flag_storage.remove(new_split.name)
-                        await self._feature_flag_storage.set_change_number(event.change_number)
+                        new_feature_flag = from_raw(json.loads(self._get_feature_flag_definition(event)))
+                        segment_list = await update_feature_flag_storage_async(self._feature_flag_storage, [new_feature_flag], event.change_number)
+                        for segment_name in segment_list:
+                            if await self._segment_storage.get(segment_name) is None:
+                                _LOGGER.debug('Fetching new segment %s', segment_name)
+                                await self._segment_handler(segment_name, event.change_number)
+
                         await self._telemetry_runtime_producer.record_update_from_sse(UpdateFromSSE.SPLIT_UPDATE)
                         continue
                     except Exception as e: