Merge pull request pandas-dev#568 from dimosped/fix-metadata-piclestore

fix the pickle store read corruption after calling write_metadata()

Author: dimosped

CHANGES.md

@@ -1,6 +1,9 @@
 ## Changelog
 
-### 1.66
+### 1.67
+  * Bugfix: #561 Fix PickleStore read corruption after write_metadata
+
+### 1.66 (2018-05-21)
   * Bugfix: #168 Do not allow empty string as a column name
   * Bugfix: #483 Remove potential floating point error from datetime_to_ms
   * Bugfix: #271 Log when library doesnt exist on delete

arctic/store/_ndarray_store.py

@@ -8,7 +8,7 @@
 
 from ..decorators import mongo_retry
 from ..exceptions import UnhandledDtypeException, DataIntegrityException
-from ._version_store_utils import checksum
+from ._version_store_utils import checksum, version_base_or_id
 
 from .._compression import compress_array, decompress
 from six.moves import xrange
@@ -39,20 +39,25 @@ def _promote(type1, type2):
 def _attempt_update_unchanged(symbol, unchanged_segment_ids, collection, version, previous_version):
     if not unchanged_segment_ids or not collection or not version:
         return
+
+    # Currenlty it is called only from _concat_and_rewrite, with "base_version_id" always empty
+    # Use version_base_or_id() instead, to make the method safe going forward, called form anywhere
+    parent_id = version_base_or_id(version)
+
     # Update the parent set of the unchanged/compressed segments
     result = collection.update_many({
                                         'symbol': symbol,  # hit only the right shard
                                                            # update_many is a broadcast query otherwise
                                         '_id': {'$in': [x['_id'] for x in unchanged_segment_ids]}
                                     },
-                                    {'$addToSet': {'parent': version['_id']}})
+                                    {'$addToSet': {'parent': parent_id}})
     # Fast check for success without extra query
     if result.matched_count == len(unchanged_segment_ids):
         return
     # update_many is tricky sometimes wrt matched_count across replicas when balancer runs. Check based on _id.
     unchanged_ids = set([x['_id'] for x in unchanged_segment_ids])
     spec = {'symbol': symbol,
-            'parent': version['_id'],
+            'parent': parent_id,
             'segment': {'$lte': unchanged_segment_ids[-1]['segment']}}
     matched_segments_ids = set([x['_id'] for x in collection.find(spec)])
     if unchanged_ids != matched_segments_ids:
@@ -240,7 +245,7 @@ def _do_read(self, collection, version, symbol, index_range=None):
         segment_count = None
 
         spec = {'symbol': symbol,
-                'parent': version.get('base_version_id', version['_id']),
+                'parent': version_base_or_id(version),
                 'segment': {'$lt': to_index}
                 }
         if from_index is not None:
@@ -333,7 +338,7 @@ def _do_append(self, collection, version, symbol, item, previous_version, dirty_
         #_CHUNK_SIZE is probably too big if we're only appending single rows of data - perhaps something smaller,
         #or also look at number of appended segments?
         if not dirty_append and version['append_count'] < _APPEND_COUNT and version['append_size'] < _APPEND_SIZE:
-            version['base_version_id'] = previous_version.get('base_version_id', previous_version['_id'])
+            version['base_version_id'] = version_base_or_id(previous_version)
 
             if len(item) > 0:
 
@@ -375,7 +380,7 @@ def _concat_and_rewrite(self, collection, version, symbol, item, previous_versio
 
         # Figure out which is the last 'full' chunk
         spec = {'symbol': symbol,
-                'parent': previous_version.get('base_version_id', previous_version['_id']),
+                'parent': version_base_or_id(previous_version),
                 'segment': {'$lt': previous_version['up_to']}}
 
         read_index_range = [0, None]
@@ -423,17 +428,21 @@ def _concat_and_rewrite(self, collection, version, symbol, item, previous_versio
             self.check_written(collection, symbol, version)
 
     def check_written(self, collection, symbol, version):
+        # Currently only called from methods which guarantee 'base_version_id' is not populated.
+        # Make it nonetheless safe for the general case.
+        parent_id = version_base_or_id(version)
+
         # Check all the chunks are in place
-        seen_chunks = collection.find({'symbol': symbol, 'parent': version['_id']},
+        seen_chunks = collection.find({'symbol': symbol, 'parent': parent_id},
                                       ).count()
 
         if seen_chunks != version['segment_count']:
-            segments = [x['segment'] for x in collection.find({'symbol': symbol, 'parent': version['_id']},
+            segments = [x['segment'] for x in collection.find({'symbol': symbol, 'parent': parent_id},
                                                               projection={'segment': 1},
                                                               )]
             raise pymongo.errors.OperationFailure("Failed to write all the Chunks. Saw %s expecting %s"
                                                   "Parent: %s \n segments: %s" %
-                                                  (seen_chunks, version['segment_count'], version['_id'], segments))
+                                                  (seen_chunks, version['segment_count'], parent_id, segments))
 
     def checksum(self, item):
         sha = hashlib.sha1()

arctic/store/_pickle_store.py

@@ -7,7 +7,7 @@
 from .._compression import decompress, compress_array
 import pymongo
 
-from ._version_store_utils import checksum, pickle_compat_load
+from ._version_store_utils import checksum, pickle_compat_load, version_base_or_id
 from ..exceptions import UnsupportedPickleStoreVersion
 
 
@@ -38,12 +38,12 @@ def read(self, mongoose_lib, version, symbol, **kwargs):
             if blob == _MAGIC_CHUNKEDV2:
                 collection = mongoose_lib.get_top_level_collection()
                 data = b''.join(decompress(x['data']) for x in collection.find({'symbol': symbol,
-                                                                                'parent': version['_id']},
+                                                                                'parent': version_base_or_id(version)},
                                                                                sort=[('segment', pymongo.ASCENDING)]))
             elif blob == _MAGIC_CHUNKED:
                 collection = mongoose_lib.get_top_level_collection()
                 data = b''.join(x['data'] for x in collection.find({'symbol': symbol,
-                                                                    'parent': version['_id']},
+                                                                    'parent': version_base_or_id(version)},
                                                                    sort=[('segment', pymongo.ASCENDING)]))
                 data = decompress(data)
             else:

arctic/store/_version_store_utils.py

@@ -63,6 +63,10 @@ def cleanup(arctic_lib, symbol, version_ids):
     collection.delete_one({'symbol':  symbol, 'parent': []})
 
 
+def version_base_or_id(version):
+    return version.get('base_version_id', version['_id'])
+
+
 def _define_compat_pickle_load():
     """Factory function to initialise the correct Pickle load function based on
     the Pandas version.
@@ -71,6 +75,7 @@ def _define_compat_pickle_load():
         return pickle.load
     return functools.partial(pickle_compat.load, compat=True)
 
+
 # Initialise the pickle load function and delete the factory function.
 pickle_compat_load = _define_compat_pickle_load()
 del _define_compat_pickle_load

tests/integration/store/test_pickle_store.py

@@ -110,3 +110,12 @@ def test_prune_previous_doesnt_kill_other_objects(library):
     library._delete_version('BLOB', 2)
     assert library._collection.count() == 0
     assert library._collection.versions.count() == 1
+
+
+def test_write_metadata(library):
+    blob = {'foo': dt(2015, 1, 1), 'object': Arctic}
+    library.write(symbol='symX', data=blob, metadata={'key1': 'value1'})
+    library.write_metadata(symbol='symX', metadata={'key2': 'value2'})
+    v = library.read('symX')
+    assert v.data == blob
+    assert v.metadata == {'key2': 'value2'}

tests/unit/store/test_pickle_store.py

@@ -65,6 +65,23 @@ def test_read_object_2():
     assert PickleStore.read(self, arctic_lib, version, sentinel.symbol) == object
     assert coll.find.call_args_list == [call({'symbol': sentinel.symbol, 'parent': sentinel._id}, sort=[('segment', 1)])]
 
+
+def test_read_with_base_version_id():
+    self = create_autospec(PickleStore)
+    version = {'_id': sentinel._id,
+               'base_version_id': sentinel.base_version_id,
+               'blob': '__chunked__'}
+    coll = Mock()
+    arctic_lib = Mock()
+    coll.find.return_value = [{'data': Binary(compressHC(cPickle.dumps(object))),
+                               'symbol': 'sentinel.symbol'}
+                              ]
+    arctic_lib.get_top_level_collection.return_value = coll
+
+    assert PickleStore.read(self, arctic_lib, version, sentinel.symbol) == object
+    assert coll.find.call_args_list == [call({'symbol': sentinel.symbol, 'parent': sentinel.base_version_id}, sort=[('segment', 1)])]
+
+
 @pytest.mark.xfail(sys.version_info >= (3,),
                    reason="lz4 data written with python2 not compatible with python3")
 def test_read_backward_compatibility():
@@ -153,5 +170,3 @@ def test_pickle_store_future_version():
     with pytest.raises(UnsupportedPickleStoreVersion) as e:
         ps.read(arctic_lib, version, sentinel.symbol)
     assert('unsupported version of pickle store' in str(e))
-
-

tests/unit/store/test_version_store_utils.py

@@ -1,8 +1,10 @@
-import pytest
-import numpy as np
 import binascii
 
-from arctic.store._version_store_utils import _split_arrs, checksum
+import numpy as np
+import pytest
+from mock import sentinel
+
+from arctic.store._version_store_utils import _split_arrs, checksum, version_base_or_id
 
 
 def test_split_arrs_empty():
@@ -30,3 +32,12 @@ def test_checksum_handles_p3strs_and_binary():
     expected = b'4O11 ;<[email protected](JRB1.?D[ZEN!8'
     assert binascii.b2a_uu(digest).strip() == expected
 
+
+def test_version_base_or_id():
+    with pytest.raises(KeyError):
+        version_base_or_id({})
+    assert version_base_or_id({'_id': sentinel._id}) == sentinel._id
+    assert version_base_or_id({
+        '_id': sentinel._id,
+        'base_version_id': sentinel.base_version_id
+    }) == sentinel.base_version_id