Merge pull request pandas-dev#155 from manahl/tickstore_allow_secondary

TickStore should respect the allow_secondary flag on Arctic

Author: jamesblackburn

.travis.yml

@@ -2,17 +2,17 @@ sudo: false
 language: python
 python:
   - "2.7"
-  - "3.4"
-  - "3.5"
+  - "3.4.4"
+  - "3.5.1"
 install:
   - pip install --upgrade pip
-  - pip install python-dateutil --upgrade 
-  - pip install pytz --upgrade 
-  - pip install tzlocal --upgrade 
+  - pip install python-dateutil --upgrade
+  - pip install pytz --upgrade
+  - pip install tzlocal --upgrade
   - pip install cython --upgrade
   - pip install pymongo --upgrade
   - pip install numpy --upgrade
-  - pip install pandas --upgrade 
+  - pip install pandas --upgrade
   - pip install decorator --upgrade
   - pip install enum34 --upgrade
   - pip install lz4 --upgrade
@@ -24,4 +24,6 @@ install:
   - pip install pytest-timeout --upgrade
   - pip install pytest-xdist --upgrade
   - pip install setuptools-git --upgrade
-script: python setup.py test
+script:
+  - pip freeze
+  - python setup.py test

CHANGES.md

@@ -3,6 +3,7 @@
 ### 1.26
 
   * Bugfix: Faster TickStore querying for multiple symbols simultaneously
+  * Bugfix: TickStore.read now respects `allow_secondary=True`
   * Bugfix: #147 Add get_info method to ChunkStore
 
 ### 1.25 (2016-05-23)

arctic/tickstore/tickstore.py

@@ -9,6 +9,7 @@
 import pandas as pd
 from pandas.core.frame import _arrays_to_mgr
 import pymongo
+from pymongo import ReadPreference
 from pymongo.errors import OperationFailure
 from six import iteritems, string_types
 
@@ -230,7 +231,14 @@ def _symbol_query(self, symbol):
             query = {}
         return query
 
-    def read(self, symbol, date_range=None, columns=None, include_images=False, _target_tick_count=0):
+    def _read_preference(self, allow_secondary):
+        """ Return the mongo read preference given an 'allow_secondary' argument
+        """
+        allow_secondary = self._allow_secondary if allow_secondary is None else allow_secondary
+        return ReadPreference.NEAREST if allow_secondary else ReadPreference.PRIMARY
+
+    def read(self, symbol, date_range=None, columns=None, include_images=False, allow_secondary=None,
+             _target_tick_count=0):
         """
         Read data for the named symbol.  Returns a VersionedItem object with
         a data and metdata element (as passed into write).
@@ -245,6 +253,12 @@ def read(self, symbol, date_range=None, columns=None, include_images=False, _tar
             Columns (fields) to return from the tickstore
         include_images : `bool`
             Should images (/snapshots) be included in the read
+        allow_secondary : `bool` or `None`
+            Override the default behavior for allowing reads from secondary members of a cluster:
+            `None` : use the settings from the top-level `Arctic` object used to query this version store.
+            `True` : allow reads from secondary members
+            `False` : only allow reads from primary members
+
         Returns
         -------
         pandas.DataFrame of data
@@ -277,7 +291,8 @@ def read(self, symbol, date_range=None, columns=None, include_images=False, _tar
 
         column_dtypes = {}
         ticks_read = 0
-        for b in self._collection.find(query, projection=projection).sort([(START, pymongo.ASCENDING)],):
+        data_coll = self._collection.with_options(read_preference=self._read_preference(allow_secondary))
+        for b in data_coll.find(query, projection=projection).sort([(START, pymongo.ASCENDING)],):
             data = self._read_bucket(b, column_set, column_dtypes,
                                      multiple_symbols or (columns is not None and 'SYMBOL' in columns),
                                      include_images, columns)

tests/integration/tickstore/test_ts_read.py

@@ -1,10 +1,11 @@
 from datetime import datetime as dt
-from mock import patch
+from mock import patch, call
 import numpy as np
 from numpy.testing.utils import assert_array_equal
 from pandas.util.testing import assert_frame_equal
 import pandas as pd
 from pandas.tseries.index import DatetimeIndex
+from pymongo import ReadPreference
 import pytest
 import pytz
 
@@ -44,6 +45,44 @@ def test_read(tickstore_lib):
     assert tickstore_lib._collection.find_one()['c'] == 2
 
 
+def test_read_allow_secondary(tickstore_lib):
+    data = [{'ASK': 1545.25,
+                  'ASKSIZE': 1002.0,
+                  'BID': 1545.0,
+                  'BIDSIZE': 55.0,
+                  'CUMVOL': 2187387.0,
+                  'DELETED_TIME': 0,
+                  'INSTRTYPE': 'FUT',
+                  'PRICE': 1545.0,
+                  'SIZE': 1.0,
+                  'TICK_STATUS': 0,
+                  'TRADEHIGH': 1561.75,
+                  'TRADELOW': 1537.25,
+                  'index': 1185076787070},
+                 {'CUMVOL': 354.0,
+                  'DELETED_TIME': 0,
+                  'PRICE': 1543.75,
+                  'SIZE': 354.0,
+                  'TRADEHIGH': 1543.75,
+                  'TRADELOW': 1543.75,
+                  'index': 1185141600600}]
+    tickstore_lib.write('FEED::SYMBOL', data)
+
+
+    with patch('pymongo.collection.Collection.find', side_effect=tickstore_lib._collection.find) as find:
+        with patch('pymongo.collection.Collection.with_options', side_effect=tickstore_lib._collection.with_options) as with_options:
+            with patch.object(tickstore_lib, '_read_preference', side_effect=tickstore_lib._read_preference) as read_pref:
+                df = tickstore_lib.read('FEED::SYMBOL', columns=['BID', 'ASK', 'PRICE'], allow_secondary=True)
+    assert read_pref.call_args_list == [call(True)]
+    assert with_options.call_args_list == [call(read_preference=ReadPreference.NEAREST)]
+    assert find.call_args_list == [call({'sy': 'FEED::SYMBOL'}, sort=[('s', 1)], projection={'s': 1, '_id': 0}),
+                                   call({'sy': 'FEED::SYMBOL', 's': {'$lte': dt(2007, 8, 21, 3, 59, 47, 70000)}}, 
+                                        projection={'sy': 1, 'cs.PRICE': 1, 'i': 1, 'cs.BID': 1, 's': 1, 'im': 1, 'v': 1, 'cs.ASK': 1})]
+
+    assert_array_equal(df['ASK'].values, np.array([1545.25, np.nan]))
+    assert tickstore_lib._collection.find_one()['c'] == 2
+
+
 def test_read_symbol_as_column(tickstore_lib):
     data = [{'ASK': 1545.25,
                   'index': 1185076787070},
@@ -168,7 +207,7 @@ def test_date_range(tickstore_lib):
     tickstore_lib._chunk_size = 3
     tickstore_lib.write('SYM', DUMMY_DATA)
 
-    with patch.object(tickstore_lib._collection, 'find', side_effect=tickstore_lib._collection.find) as f:
+    with patch('pymongo.collection.Collection.find', side_effect=tickstore_lib._collection.find) as f:
         df = tickstore_lib.read('SYM', date_range=DateRange(20130101, 20130103), columns=None)
         assert_array_equal(df['b'].values, np.array([2., 3., 5.]))
         assert tickstore_lib._collection.find(f.call_args_list[-1][0][0]).count() == 1

tests/unit/tickstore/test_tickstore.py

@@ -2,6 +2,7 @@
 from functools import partial
 from mock import create_autospec, sentinel, call
 import pytest
+from pymongo import ReadPreference
 from pymongo.collection import Collection
 import numpy as np
 import pandas as pd
@@ -164,3 +165,23 @@ def test_tickstore_pandas_to_bucket_image():
     assert bucket[SYMBOL] == symbol
     assert bucket[IMAGE_DOC] == {IMAGE: initial_image,
                                  IMAGE_TIME: initial_image['index']}
+
+
+def test__read_preference__allow_secondary_true():
+    self = create_autospec(TickStore)
+    assert TickStore._read_preference(self, True) == ReadPreference.NEAREST
+
+
+def test__read_preference__allow_secondary_false():
+    self = create_autospec(TickStore)
+    assert TickStore._read_preference(self, False) == ReadPreference.PRIMARY
+
+
+def test__read_preference__default_true():
+    self = create_autospec(TickStore, _allow_secondary=True)
+    assert TickStore._read_preference(self, None) == ReadPreference.NEAREST
+
+
+def test__read_preference__default_false():
+    self = create_autospec(TickStore, _allow_secondary=False)
+    assert TickStore._read_preference(self, None) == ReadPreference.PRIMARY