Merge pull request pandas-dev#416 from manahl/bugfix/metadata_store_append_inconsistent_return

Bugfix/metadata store append inconsistent return

Author: kraphtuos

CHANGES.md

@@ -4,6 +4,7 @@
   * Perf:    #408 Improve memory performance of version store's serializer
   * Bugfix   #394 Multi symbol read in chunkstore
   * Bugfix:  #407 Fix segment issue on appends in chunkstore
+  * Bugfix:  Inconsistent returns on MetadataStore.append
 
 ### 1.51 (2017-08-21)
   * Bugfix:  #397 Remove calls to deprecated methods in pymongo

arctic/store/metadata_store.py

@@ -2,12 +2,12 @@
 import logging
 
 import pandas as pd
-
 import bson
 import pymongo
+
+from .._util import indent
 from ..decorators import mongo_retry
 from ..exceptions import NoDataFoundException
-
 from .bson_store import BSONStore
 
 logger = logging.getLogger(__name__)
@@ -44,6 +44,18 @@ def __init__(self, arctic_lib):
     def _reset(self):
         self._collection = self._arctic_lib.get_top_level_collection().metadata
 
+    def __getstate__(self):
+        return {'arctic_lib': self._arctic_lib}
+
+    def __setstate__(self, state):
+        return MetadataStore.__init__(self, state['arctic_lib'])
+
+    def __str__(self):
+        return """<%s at %s>\n%s""" % (self.__class__.__name__, hex(id(self)), indent(str(self._arctic_lib), 4))
+
+    def __repr__(self):
+        return str(self)
+
     @mongo_retry
     def list_symbols(self):
         return self.distinct('symbol')
@@ -151,7 +163,7 @@ def append(self, symbol, metadata, start_time=None):
                 raise ValueError('start_time={} is earlier than the last metadata @{}'.format(start_time,
                                                                                               old_metadata['start_time']))
             if old_metadata['metadata'] == metadata:
-                return metadata
+                return old_metadata
         elif metadata is None:
             return
 
@@ -189,7 +201,8 @@ def prepend(self, symbol, metadata, start_time=None):
             if old_metadata['metadata'] == metadata:
                 self.find_one_and_update({'symbol': symbol}, {'$set': {'start_time': start_time}},
                                          sort=[('start_time', pymongo.ASCENDING)])
-                return metadata
+                old_metadata['start_time'] = start_time
+                return old_metadata
             end_time = old_metadata.get('start_time')
         else:
             end_time = None

tests/integration/store/test_metadata_store.py

@@ -1,3 +1,7 @@
+try:
+    import cPickle as pickle
+except ImportError:
+    import pickle
 from datetime import datetime as dt
 import pandas as pd
 from pandas.util.testing import assert_frame_equal
@@ -32,6 +36,16 @@ def integrity_check(ms_lib, symbol):
         metadata = item['metadata']
     assert start_time == None, 'end_time of the last entry should be unset'
 
+
+def test_pickle(ms_lib):
+    buff = pickle.dumps(ms_lib)
+    mnew = pickle.loads(buff)
+    assert ms_lib._arctic_lib.get_name() == mnew._arctic_lib.get_name()
+
+    assert "arctic_test.TEST" in str(ms_lib)
+    assert str(ms_lib) == repr(ms_lib)
+
+
 def test_has_symbol(ms_lib):
     assert not ms_lib.has_symbol(symbol1)
     ms_lib.append(symbol1, metadata1)
@@ -72,19 +86,25 @@ def test_write_history(ms_lib):
 
 
 def test_append(ms_lib):
-    ms_lib.append(symbol1, None)
+    ret1 = ms_lib.append(symbol1, None)
     assert not ms_lib.has_symbol(symbol1)
+    assert ret1 is None
 
-    ms_lib.append(symbol1, metadata1, start_time1)
+    ret2 = ms_lib.append(symbol1, metadata1, start_time1)
     assert ms_lib.read(symbol1) == metadata1
+    assert ret2['symbol'] == symbol1
+    assert ret2['start_time'] == start_time1
+    assert ret2['metadata'] == metadata1
 
     # ensure writing same metadata does not create new entry
-    ms_lib.append(symbol1, metadata1, start_time2)
+    ret3 = ms_lib.append(symbol1, metadata1, start_time2)
     assert ms_lib.read(symbol1) == metadata1
     assert_frame_equal(ms_lib.read_history(symbol1), dataframe1)
+    assert ret3 == ret2
 
-    ms_lib.append(symbol1, metadata2, start_time2)
+    ret4 = ms_lib.append(symbol1, metadata2, start_time2)
     assert_frame_equal(ms_lib.read_history(symbol1), dataframe3)
+    assert ret4['metadata'] == metadata2
 
     with pytest.raises(ValueError):
         ms_lib.append(symbol1, metadata1, start_time1)
@@ -93,19 +113,26 @@ def test_append(ms_lib):
 
 
 def test_prepend(ms_lib):
-    ms_lib.prepend(symbol1, None)
+    ret1 = ms_lib.prepend(symbol1, None)
     assert not ms_lib.has_symbol(symbol1)
+    assert ret1 is None
 
-    ms_lib.prepend(symbol1, metadata2, start_time2)
+    ret2 = ms_lib.prepend(symbol1, metadata2, start_time2)
     assert ms_lib.read(symbol1) == metadata2
     assert_frame_equal(ms_lib.read_history(symbol1), dataframe4)
+    assert ret2['symbol'] == symbol1
+    assert ret2['start_time'] == start_time2
+    assert ret2['metadata'] == metadata2
 
-    ms_lib.prepend(symbol1, metadata1, start_time1)
+    ret3 = ms_lib.prepend(symbol1, metadata1, start_time1)
     assert_frame_equal(ms_lib.read_history(symbol1), dataframe3)
+    assert ret3['metadata'] == metadata1
 
     # ensure writing same metadata does not create new entry
-    ms_lib.prepend(symbol1, metadata1, start_time0)
+    ret4 = ms_lib.prepend(symbol1, metadata1, start_time0)
     assert_frame_equal(ms_lib.read_history(symbol1), dataframe5)
+    ret3['start_time'] = start_time0
+    assert ret4 == ret3
 
     with pytest.raises(ValueError):
         ms_lib.append(symbol1, metadata2, start_time2)