added sync segment local class

Author: chillaq

splitio/sync/segment.py

@@ -8,6 +8,7 @@
 from splitio.tasks.util import workerpool
 from splitio.models import segments
 from splitio.util.backoff import Backoff
+from splitio.optional.loaders import asyncio, aiofiles
 from splitio.sync import util
 
 _LOGGER = logging.getLogger(__name__)
@@ -195,27 +196,57 @@ def segment_exist_in_storage(self, segment_name):
         """
         return self._segment_storage.get(segment_name) != None
 
-class LocalSegmentSynchronizer(object):
-    """Localhost mode segment synchronizer."""
+class LocalSegmentSynchronizerBase(object):
+    """Localhost mode segment base synchronizer."""
 
     _DEFAULT_SEGMENT_TILL = -1
 
-    def __init__(self, segment_folder, split_storage, segment_storage):
+    def _sanitize_segment(self, parsed):
+        """
+        Sanitize json elements.
+
+        :param parsed: segment dict
+        :type parsed: Dict
+
+        :return: sanitized segment structure dict
+        :rtype: Dict
+        """
+        if 'name' not in parsed or parsed['name'] is None:
+            _LOGGER.warning("Segment does not have [name] element, skipping")
+            raise Exception("Segment does not have [name] element")
+        if parsed['name'].strip() == '':
+            _LOGGER.warning("Segment [name] element is blank, skipping")
+            raise Exception("Segment [name] element is blank")
+
+        for element in [('till', -1, -1, None, None, [0]),
+                        ('added', [], None, None, None, None),
+                        ('removed', [], None, None, None, None)
+                        ]:
+            parsed = util._sanitize_object_element(parsed, 'segment', element[0], element[1], lower_value=element[2], upper_value=element[3], in_list=None, not_in_list=element[5])
+        parsed = util._sanitize_object_element(parsed, 'segment', 'since', parsed['till'], -1, parsed['till'], None, [0])
+
+        return parsed
+
+
+class LocalSegmentSynchronizer(LocalSegmentSynchronizerBase):
+    """Localhost mode segment synchronizer."""
+
+    def __init__(self, segment_folder, feature_flag_storage, segment_storage):
         """
         Class constructor.
 
         :param segment_folder: patch to the segment folder
         :type segment_folder: str
 
-        :param split_storage: Feature flag Storage.
-        :type split_storage: splitio.storage.InMemorySplitStorage
+        :param feature_flag_storage: Feature flag Storage.
+        :type feature_flag_storage: splitio.storage.InMemorySplitStorage
 
         :param segment_storage: Segment storage reference.
         :type segment_storage: splitio.storage.SegmentStorage
 
         """
         self._segment_folder = segment_folder
-        self._split_storage = split_storage
+        self._feature_flag_storage = feature_flag_storage
         self._segment_storage = segment_storage
         self._segment_sha = {}
 
@@ -231,7 +262,7 @@ def synchronize_segments(self, segment_names = None):
         """
         _LOGGER.info('Synchronizing segments now.')
         if segment_names is None:
-            segment_names = self._split_storage.get_segment_names()
+            segment_names = self._feature_flag_storage.get_segment_names()
 
         return_flag = True
         for segment_name in segment_names:
@@ -295,33 +326,118 @@ def _read_segment_from_json_file(self, filename):
         except Exception as exc:
             raise ValueError("Error parsing file %s. Make sure it's readable." % filename) from exc
 
-    def _sanitize_segment(self, parsed):
+    def segment_exist_in_storage(self, segment_name):
         """
-        Sanitize json elements.
+        Check if a segment exists in the storage
 
-        :param parsed: segment dict
-        :type parsed: Dict
+        :param segment_name: Name of the segment
+        :type segment_name: str
 
-        :return: sanitized segment structure dict
-        :rtype: Dict
+        :return: True if segment exist. False otherwise.
+        :rtype: bool
         """
-        if 'name' not in parsed or parsed['name'] is None:
-            _LOGGER.warning("Segment does not have [name] element, skipping")
-            raise Exception("Segment does not have [name] element")
-        if parsed['name'].strip() == '':
-            _LOGGER.warning("Segment [name] element is blank, skipping")
-            raise Exception("Segment [name] element is blank")
+        return self._segment_storage.get(segment_name) != None
 
-        for element in [('till', -1, -1, None, None, [0]),
-                        ('added', [], None, None, None, None),
-                        ('removed', [], None, None, None, None)
-                        ]:
-            parsed = util._sanitize_object_element(parsed, 'segment', element[0], element[1], lower_value=element[2], upper_value=element[3], in_list=None, not_in_list=element[5])
-        parsed = util._sanitize_object_element(parsed, 'segment', 'since', parsed['till'], -1, parsed['till'], None, [0])
 
-        return parsed
+class LocalSegmentSynchronizerAsync(LocalSegmentSynchronizerBase):
+    """Localhost mode segment async synchronizer."""
+
+    def __init__(self, segment_folder, feature_flag_storage, segment_storage):
+        """
+        Class constructor.
+
+        :param segment_folder: patch to the segment folder
+        :type segment_folder: str
+
+        :param feature_flag_storage: Feature flag Storage.
+        :type feature_flag_storage: splitio.storage.InMemorySplitStorage
+
+        :param segment_storage: Segment storage reference.
+        :type segment_storage: splitio.storage.SegmentStorage
 
-    def segment_exist_in_storage(self, segment_name):
+        """
+        self._segment_folder = segment_folder
+        self._feature_flag_storage = feature_flag_storage
+        self._segment_storage = segment_storage
+        self._segment_sha = {}
+
+    async def synchronize_segments(self, segment_names = None):
+        """
+        Loop through given segment names and synchronize each one.
+
+        :param segment_names: Optional, array of segment names to update.
+        :type segment_name: {str}
+
+        :return: True if no error occurs. False otherwise.
+        :rtype: bool
+        """
+        _LOGGER.info('Synchronizing segments now.')
+        if segment_names is None:
+            segment_names = await self._feature_flag_storage.get_segment_names()
+
+        return_flag = True
+        for segment_name in segment_names:
+            if not await self.synchronize_segment(segment_name):
+                return_flag = False
+
+        return return_flag
+
+    async def synchronize_segment(self, segment_name, till=None):
+        """
+        Update a segment from queue
+
+        :param segment_name: Name of the segment to update.
+        :type segment_name: str
+
+        :param till: ChangeNumber received.
+        :type till: int
+
+        :return: True if no error occurs. False otherwise.
+        :rtype: bool
+        """
+        try:
+            fetched = await self._read_segment_from_json_file(segment_name)
+            fetched_sha = util._get_sha(json.dumps(fetched))
+            if not await self.segment_exist_in_storage(segment_name):
+                    self._segment_sha[segment_name] = fetched_sha
+                    await self._segment_storage.put(segments.from_raw(fetched))
+                    _LOGGER.debug("segment %s is added to storage", segment_name)
+                    return True
+
+            if fetched_sha == self._segment_sha[segment_name]:
+                return True
+
+            self._segment_sha[segment_name] = fetched_sha
+            if await self._segment_storage.get_change_number(segment_name) > fetched['till'] and fetched['till'] != self._DEFAULT_SEGMENT_TILL:
+                return True
+
+            await self._segment_storage.update(segment_name, fetched['added'], fetched['removed'], fetched['till'])
+            _LOGGER.debug("segment %s is updated", segment_name)
+        except Exception as e:
+            _LOGGER.error("Could not fetch segment: %s \n" + str(e), segment_name)
+            return False
+
+        return True
+
+    async def _read_segment_from_json_file(self, filename):
+        """
+        Parse a segment and store in segment storage.
+
+        :param filename: Path of the file containing Feature flag
+        :type filename: str.
+
+        :return: Sanitized segment structure
+        :rtype: Dict
+        """
+        try:
+            async with aiofiles.open(os.path.join(self._segment_folder, "%s.json" % filename), 'r') as flo:
+                parsed = json.loads(await flo.read())
+            santitized_segment = self._sanitize_segment(parsed)
+            return santitized_segment
+        except Exception as exc:
+            raise ValueError("Error parsing file %s. Make sure it's readable." % filename) from exc
+
+    async def segment_exist_in_storage(self, segment_name):
         """
         Check if a segment exists in the storage
 
@@ -331,4 +447,4 @@ def segment_exist_in_storage(self, segment_name):
         :return: True if segment exist. False otherwise.
         :rtype: bool
         """
-        return self._segment_storage.get(segment_name) != None
+        return await self._segment_storage.get(segment_name) != None

tests/sync/test_segments_synchronizer.py

@@ -6,9 +6,10 @@
 from splitio.api import APIException
 from splitio.api.commons import FetchOptions
 from splitio.storage import SplitStorage, SegmentStorage
-from splitio.storage.inmemmory import InMemorySegmentStorage, InMemorySplitStorage
-from splitio.sync.segment import SegmentSynchronizer, LocalSegmentSynchronizer
+from splitio.storage.inmemmory import InMemorySegmentStorage, InMemorySegmentStorageAsync, InMemorySplitStorage, InMemorySplitStorageAsync
+from splitio.sync.segment import SegmentSynchronizer, LocalSegmentSynchronizer, LocalSegmentSynchronizerAsync
 from splitio.models.segments import Segment
+from splitio.optional.loaders import aiofiles
 
 import pytest
 
@@ -356,3 +357,124 @@ def test_json_elements_sanitization(self, mocker):
         segment3["till"] = 12
         segment2 = {"name": 'seg', "added": [], "removed": [], "since": 20, "till": 12}
         assert(segment_synchronizer._sanitize_segment(segment2) == segment3)
+
+
+class LocalSegmentsSynchronizerTests(object):
+    """Segments synchronizer test cases."""
+
+    @pytest.mark.asyncio
+    async def test_synchronize_segments_error(self, mocker):
+        """On error."""
+        split_storage = mocker.Mock(spec=SplitStorage)
+        async def get_segment_names():
+            return ['segmentA', 'segmentB', 'segmentC']
+        split_storage.get_segment_names = get_segment_names
+
+        storage = mocker.Mock(spec=SegmentStorage)
+        async def get_change_number():
+            return -1
+        storage.get_change_number = get_change_number
+
+        segments_synchronizer = LocalSegmentSynchronizerAsync('/,/,/invalid folder name/,/,/', split_storage, storage)
+        assert not await segments_synchronizer.synchronize_segments()
+
+    @pytest.mark.asyncio
+    async def test_synchronize_segments(self, mocker):
+        """Test the normal operation flow."""
+        split_storage = mocker.Mock(spec=InMemorySplitStorage)
+        async def get_segment_names():
+            return ['segmentA', 'segmentB', 'segmentC']
+        split_storage.get_segment_names = get_segment_names
+
+        storage = InMemorySegmentStorageAsync()
+
+        segment_a = {'name': 'segmentA', 'added': ['key1', 'key2', 'key3'], 'removed': [],
+                        'since': -1, 'till': 123}
+        segment_b = {'name': 'segmentB', 'added': ['key4', 'key5', 'key6'], 'removed': [],
+                        'since': -1, 'till': 123}
+        segment_c = {'name': 'segmentC', 'added': ['key7', 'key8', 'key9'], 'removed': [],
+                        'since': -1, 'till': 123}
+        blank = {'added': [], 'removed': [], 'since': 123, 'till': 123}
+
+        async def read_segment_from_json_file(*args, **kwargs):
+            if args[0] == 'segmentA':
+                return segment_a
+            if args[0] == 'segmentB':
+                return segment_b
+            if args[0] == 'segmentC':
+                return segment_c
+            return blank
+
+        segments_synchronizer = LocalSegmentSynchronizerAsync('segment_path', split_storage, storage)
+        segments_synchronizer._read_segment_from_json_file = read_segment_from_json_file
+        assert await segments_synchronizer.synchronize_segments()
+
+        segment = await storage.get('segmentA')
+        assert segment.name == 'segmentA'
+        assert segment.contains('key1')
+        assert segment.contains('key2')
+        assert segment.contains('key3')
+
+        segment = await storage.get('segmentB')
+        assert segment.name == 'segmentB'
+        assert segment.contains('key4')
+        assert segment.contains('key5')
+        assert segment.contains('key6')
+
+        segment = await storage.get('segmentC')
+        assert segment.name == 'segmentC'
+        assert segment.contains('key7')
+        assert segment.contains('key8')
+        assert segment.contains('key9')
+
+        # Should sync when changenumber is not changed
+        segment_a['added'] = ['key111']
+        await segments_synchronizer.synchronize_segments(['segmentA'])
+        segment = await storage.get('segmentA')
+        assert segment.contains('key111')
+
+        # Should not sync when changenumber below till
+        segment_a['till'] = 122
+        segment_a['added'] = ['key222']
+        await segments_synchronizer.synchronize_segments(['segmentA'])
+        segment = await storage.get('segmentA')
+        assert not segment.contains('key222')
+
+        # Should sync when changenumber above till
+        segment_a['till'] = 124
+        await segments_synchronizer.synchronize_segments(['segmentA'])
+        segment = await storage.get('segmentA')
+        assert segment.contains('key222')
+
+        # Should sync when till is default (-1)
+        segment_a['till'] = -1
+        segment_a['added'] = ['key33']
+        await segments_synchronizer.synchronize_segments(['segmentA'])
+        segment = await storage.get('segmentA')
+        assert segment.contains('key33')
+
+        # verify remove keys
+        segment_a['added'] = []
+        segment_a['removed'] = ['key111']
+        segment_a['till'] = 125
+        await segments_synchronizer.synchronize_segments(['segmentA'])
+        segment = await storage.get('segmentA')
+        assert not segment.contains('key111')
+
+    @pytest.mark.asyncio
+    async def test_reading_json(self, mocker):
+        """Test reading json file."""
+        async with aiofiles.open("./segmentA.json", "w") as f:
+            await f.write('{"name": "segmentA", "added": ["key1", "key2", "key3"], "removed": [],"since": -1, "till": 123}')
+        split_storage = mocker.Mock(spec=InMemorySplitStorageAsync)
+        storage = InMemorySegmentStorageAsync()
+        segments_synchronizer = LocalSegmentSynchronizerAsync('.', split_storage, storage)
+        assert await segments_synchronizer.synchronize_segments(['segmentA'])
+
+        segment = await storage.get('segmentA')
+        assert segment.name == 'segmentA'
+        assert segment.contains('key1')
+        assert segment.contains('key2')
+        assert segment.contains('key3')
+
+        os.remove("./segmentA.json")