User defined metadata for tickstore

arctic/tickstore/tickstore.py

@@ -74,6 +74,8 @@
 COUNT = 'c'
 VERSION = 'v'
 
+META = 'md'
+
 CHUNK_VERSION_NUMBER = 3
 
 
@@ -90,6 +92,8 @@ def _ensure_index(self):
                                  (START, pymongo.ASCENDING)], background=True)
         collection.create_index([(START, pymongo.ASCENDING)], background=True)
 
+        self._metadata.create_index([(SYMBOL, pymongo.ASCENDING)], background=True, unique=True)
+
     def __init__(self, arctic_lib, chunk_size=100000):
         """
         Parameters
@@ -110,6 +114,7 @@ def __init__(self, arctic_lib, chunk_size=100000):
     def _reset(self):
         # The default collections
         self._collection = self._arctic_lib.get_top_level_collection()
+        self._metadata = self._collection.metadata
 
     def __getstate__(self):
         return {'arctic_lib': self._arctic_lib}
@@ -144,6 +149,9 @@ def delete(self, symbol, date_range=None):
             assert date_range.start and date_range.end
             query[START] = {'$gte': date_range.start}
             query[END] = {'$lte': date_range.end}
+        else:
+            # delete metadata on complete deletion
+            self._metadata.delete_one({SYMBOL: symbol})
         return self._collection.delete_many(query)
 
     def list_symbols(self, date_range=None):
@@ -238,7 +246,7 @@ def _read_preference(self, 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):
+             _target_tick_count=0, metadata=False):
         """
         Read data for the named symbol.  Returns a VersionedItem object with
         a data and metdata element (as passed into write).
@@ -258,10 +266,13 @@ def read(self, symbol, date_range=None, columns=None, include_images=False, allo
             `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
+        metadata: `bool`
+            if `True` read and return user defined metadata for the symbol
+            
 
         Returns
         -------
-        pandas.DataFrame of data
+        pandas.DataFrame of data, or tuple of (data, metadata) if metadata is `True`
         """
         perf_start = dt.now()
         rtn = {}
@@ -343,6 +354,8 @@ def read(self, symbol, date_range=None, columns=None, include_images=False, allo
         if date_range:
             # FIXME: support DateRange.interval...
             rtn = rtn.loc[date_range.start:date_range.end]
+        if metadata:
+            return rtn, self._metadata.find_one({SYMBOL: symbol})[META]
         return rtn
 
     def _pad_and_fix_dtypes(self, cols, column_dtypes):
@@ -514,7 +527,7 @@ def _assert_nonoverlapping_data(self, symbol, start, end):
                 raise OverlappingDataException("Document already exists with start:{} end:{} in the range of our start:{} end:{}".format(
                                                             doc[START], doc[END], start, end))
 
-    def write(self, symbol, data, initial_image=None):
+    def write(self, symbol, data, initial_image=None, metadata=None):
         """
         Writes a list of market data events.
 
@@ -530,6 +543,8 @@ def write(self, symbol, data, initial_image=None):
         initial_image : dict
             Dict of the initial image at the start of the document. If this contains a 'index' entry it is
             assumed to be the time of the timestamp of the index
+        metadata: dict
+            optional user defined metadata - one per symbol
         """
         pandas = False
         # Check for overlapping data
@@ -550,6 +565,11 @@ def write(self, symbol, data, initial_image=None):
             buckets = self._to_buckets(data, symbol, initial_image)
         self._write(buckets)
 
+        if metadata:
+            ret = self._metadata.replace_one({SYMBOL: symbol},
+                                             {SYMBOL: symbol, META: metadata},
+                                             upsert=True)
+
     def _write(self, buckets):
         start = dt.now()
         mongo_retry(self._collection.insert_many)(buckets)

tests/integration/tickstore/test_ts_read.py

@@ -605,3 +605,10 @@ def test_read_with_image(tickstore_lib):
     assert set(df.columns) == set(['c'])
     assert_array_equal(df['c'].values, np.array([2]))
     assert df.index[0] == dt(2013, 1, 1, 10, tzinfo=mktz('Europe/London'))
+
+
+def test_read_with_metadata(tickstore_lib):
+    metadata = {'metadata': 'important data'}
+    tickstore_lib.write('test', [{'index': dt(2013, 6, 1, 13, 00, tzinfo=mktz('Europe/London')), 'price': 100.50, 'ticker': 'QQQ'}], metadata=metadata)
+    data, m = tickstore_lib.read('test', metadata=True)
+    assert(metadata == m)