MetadataStore (pandas-dev#392)

* Initial commit on MetadataStore and preliminary unit tests

* Fix bugs and add integration tests

* Remove possible duplicate in write_metadata_history and update strings and names

* Move _collection to top_level_collection.metadata

* Fix unit tests

* Push query methods to BSONStore

* Change names of class methods

* Remove unused _insert()

* Fix magical return in read() for readability

* Add prepend() and read_history() and change input for write_history() to agree with the output of read_history()

* Fix typo and add integrity check for end_time in tests

* Add mongo_retry at various places

Author: kraphtuos

arctic/arctic.py

@@ -9,26 +9,28 @@
 from .decorators import mongo_retry
 from .exceptions import LibraryNotFoundException, ArcticException, QuotaExceededException
 from .hooks import get_mongodb_uri
-from .store import version_store, bson_store
+from .store import version_store, bson_store, metadata_store
 from .tickstore import tickstore, toplevel
 from .chunkstore import chunkstore
 from six import string_types
 
 
-__all__ = ['Arctic', 'VERSION_STORE', 'TICK_STORE', 'CHUNK_STORE', 'register_library_type']
+__all__ = ['Arctic', 'VERSION_STORE', 'METADATA_STORE', 'TICK_STORE', 'CHUNK_STORE', 'register_library_type']
 
 logger = logging.getLogger(__name__)
 
 # Default Arctic application name: 'arctic'
 APPLICATION_NAME = 'arctic'
 VERSION_STORE = version_store.VERSION_STORE_TYPE
+METADATA_STORE = metadata_store.METADATA_STORE_TYPE
 TICK_STORE = tickstore.TICK_STORE_TYPE
 CHUNK_STORE = chunkstore.CHUNK_STORE_TYPE
 LIBRARY_TYPES = {version_store.VERSION_STORE_TYPE: version_store.VersionStore,
                  tickstore.TICK_STORE_TYPE: tickstore.TickStore,
                  toplevel.TICK_STORE_TYPE: toplevel.TopLevelTickStore,
                  chunkstore.CHUNK_STORE_TYPE: chunkstore.ChunkStore,
-                 bson_store.BSON_STORE_TYPE: bson_store.BSONStore
+                 bson_store.BSON_STORE_TYPE: bson_store.BSONStore,
+                 metadata_store.METADATA_STORE_TYPE: metadata_store.MetadataStore
                  }
 
 
@@ -53,6 +55,7 @@ class Arctic(object):
                                 (other Python types are pickled)
        - arctic.TICK_STORE - Tick specific library. Supports 'snapshots', efficiently
                              stores updates, not versioned.
+       - arctic.METADATA_STORE - Stores metadata with timestamps
 
     Arctic and ArcticLibrary are responsible for Connection setup, authentication,
     dispatch to the appropriate library implementation, and quotas.

arctic/fixtures/arctic.py

@@ -96,3 +96,9 @@ def tickstore_lib(arctic, library_name):
 def chunkstore_lib(arctic, library_name):
     arctic.initialize_library(library_name, CHUNK_STORE_TYPE)
     return arctic.get_library(library_name)
+
+
+@pytest.fixture(scope="function")
+def ms_lib(arctic, library_name):
+    arctic.initialize_library(library_name, m.METADATA_STORE)
+    return arctic.get_library(library_name)

arctic/store/bson_store.py

@@ -58,6 +58,13 @@ def find(self, *args, **kwargs):
         """
         return self._collection.find(*args, **kwargs)
 
+    @mongo_retry
+    def find_one(self, *args, **kwargs):
+        """
+        See http://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.find_one
+        """
+        return self._collection.find_one(*args, **kwargs)
+
     @mongo_retry
     def insert_one(self, document, **kwargs):
         """
@@ -119,6 +126,20 @@ def find_one_and_replace(self, filter, replacement, **kwargs):
         self._arctic_lib.check_quota()
         return self._collection.find_one_and_replace(filter, replacement, **kwargs)
 
+    @mongo_retry
+    def find_one_and_update(self, filter, update, **kwargs):
+        """
+        See http://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.find_one_and_update
+        """
+        self._arctic_lib.check_quota()
+        return self._collection.find_one_and_update(filter, update, **kwargs)
+
+    def find_one_and_delete(self, filter, **kwargs):
+        """
+        See http://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.find_one_and_delete
+        """
+        return self._collection.find_one_and_delete(filter, **kwargs)
+
     @mongo_retry
     def bulk_write(self, requests, **kwargs):
         """

arctic/store/metadata_store.py

@@ -0,0 +1,235 @@
+from datetime import datetime as dt
+import logging
+
+import pandas as pd
+
+import bson
+import pymongo
+from ..decorators import mongo_retry
+from ..exceptions import NoDataFoundException
+
+from .bson_store import BSONStore
+
+logger = logging.getLogger(__name__)
+
+METADATA_STORE_TYPE = 'MetadataStore'
+
+class MetadataStore(BSONStore):
+    """
+    Metadata Store. This stores metadata with timestamps to allow temporal queries.
+
+    Entries are stored in the following format:
+        'symbol': symbol name
+        'metadata': metadata to be persisted
+        'start_time': when entry becomes effective
+        'end_time': (Optional) when entry expires. If not set, it is still in effect
+
+    For each symbol end_time of a entry should match start_time of the next one except for the current entry.
+    """
+
+    @classmethod
+    def initialize_library(cls, arctic_lib, hashed=True, **kwargs):
+        MetadataStore(arctic_lib)._ensure_index()
+        BSONStore.initialize_library(arctic_lib, hashed, **kwargs)
+
+    @mongo_retry
+    def _ensure_index(self):
+        self.create_index([('symbol', pymongo.ASCENDING), ('start_time', pymongo.DESCENDING)],
+                          unique=True, background=True)
+
+    def __init__(self, arctic_lib):
+        self._arctic_lib = arctic_lib
+        self._reset()
+
+    def _reset(self):
+        self._collection = self._arctic_lib.get_top_level_collection().metadata
+
+    @mongo_retry
+    def list_symbols(self):
+        return self.distinct('symbol')
+
+    @mongo_retry
+    def has_symbol(self, symbol):
+        return self.find_one({'symbol': symbol}) is not None
+
+    @mongo_retry
+    def read_history(self, symbol):
+        """
+        Return all metadata saved for `symbol`
+
+        Parameters
+        ----------
+        symbol : `str`
+            symbol name for the item
+
+        Returns
+        -------
+        pandas.DateFrame containing timestamps and metadata entries
+        """
+        find = self.find({'symbol': symbol}, sort=[('start_time', pymongo.ASCENDING)])
+        times = []
+        entries = []
+        for item in find:
+            times.append(item['start_time'])
+            entries.append(item['metadata'])
+        return pd.DataFrame({symbol: entries}, times)
+
+    @mongo_retry
+    def read(self, symbol):
+        """
+        Return current metadata saved for `symbol`
+
+        Parameters
+        ----------
+        symbol : `str`
+            symbol name for the item
+
+        Returns
+        -------
+        metadata
+        """
+        res = self.find_one({'symbol': symbol}, sort=[('start_time', pymongo.DESCENDING)])
+        return res['metadata'] if res is not None else None
+
+
+    def write_history(self, collection):
+        """
+        Manually overwrite entire metadata history for symbols in `collection`
+
+        Parameters
+        ----------
+        collection : `list of pandas.DataFrame`
+            with symbol names as headers and timestamps as indices
+            (the same format as output of read_history)
+            Example:
+                [pandas.DataFrame({'symbol': [{}]}, [datetime.datetime.utcnow()])]
+        """
+        documents = []
+        for dataframe in collection:
+            if len(dataframe.columns) != 1:
+                raise ValueError('More than one symbol found in a DataFrame')
+            symbol = dataframe.columns[0]
+            times = dataframe.index
+            entries = dataframe[symbol].values
+            if self.has_symbol(symbol):
+                self.purge(symbol)
+            doc = {'symbol': symbol, 'metadata': entries[0], 'start_time': times[0]}
+            for metadata, start_time in zip(entries[1:], times[1:]):
+                if metadata == doc['metadata']:
+                    continue
+                doc['end_time'] = start_time
+                documents.append(doc)
+                doc = {'symbol': symbol, 'metadata': metadata, 'start_time': start_time}
+            documents.append(doc)
+
+        self.insert_many(documents)
+
+    def append(self, symbol, metadata, start_time=None):
+        """
+        Update metadata entry for `symbol`
+
+        Parameters
+        ----------
+        symbol : `str`
+            symbol name for the item
+        metadata : `dict`
+            to be persisted
+        start_time : `datetime.datetime`
+            when metadata becomes effective
+            Default: datetime.datetime.utcnow()
+        """
+        if start_time is None:
+            start_time = dt.utcnow()
+        old_metadata = self.find_one({'symbol': symbol}, sort=[('start_time', pymongo.DESCENDING)])
+        if old_metadata is not None:
+            if old_metadata['start_time'] >= start_time:
+                raise ValueError('start_time={} is earlier than the last metadata @{}'.format(start_time,
+                                                                                              old_metadata['start_time']))
+            if old_metadata['metadata'] == metadata:
+                logger.warning('No change to metadata')
+                return metadata
+        elif metadata is None:
+            return
+
+        self.find_one_and_update({'symbol': symbol}, {'$set': {'end_time': start_time}},
+                                  sort=[('start_time', pymongo.DESCENDING)])
+        document = {'_id': bson.ObjectId(), 'symbol': symbol, 'metadata': metadata, 'start_time': start_time}
+        mongo_retry(self.insert_one)(document)
+
+        logger.debug('Finished writing metadata for %s', symbol)
+        return document
+
+    def prepend(self, symbol, metadata, start_time=None):
+        """
+        Prepend a metadata entry for `symbol`
+
+        Parameters
+        ----------
+        symbol : `str`
+            symbol name for the item
+        metadata : `dict`
+            to be persisted
+        start_time : `datetime.datetime`
+            when metadata becomes effective
+            Default: datetime.datetime.min
+        """
+        if metadata is None:
+            return
+        if start_time is None:
+            start_time = dt.min
+        old_metadata = self.find_one({'symbol': symbol}, sort=[('start_time', pymongo.ASCENDING)])
+        if old_metadata is not None:
+            if old_metadata['start_time'] <= start_time:
+                raise ValueError('start_time={} is later than the first metadata @{}'.format(start_time,
+                                                                                             old_metadata['start_time']))
+            if old_metadata['metadata'] == metadata:
+                self.find_one_and_update({'symbol': symbol}, {'$set': {'start_time': start_time}},
+                                         sort=[('start_time', pymongo.ASCENDING)])
+                return metadata
+            end_time = old_metadata.get('start_time')
+        else:
+            end_time = None
+
+        document = {'_id': bson.ObjectId(), 'symbol': symbol, 'metadata': metadata, 'start_time': start_time}
+        if end_time is not None:
+            document['end_time'] = end_time
+        mongo_retry(self.insert_one)(document)
+
+        logger.debug('Finished writing metadata for %s', symbol)
+        return document
+
+    def pop(self, symbol):
+        """
+        Delete current metadata of `symbol`
+
+        Parameters
+        ----------
+        symbol : `str`
+            symbol name to delete
+
+        Returns
+        -------
+        Deleted metadata
+        """
+        last_metadata = self.find_one({'symbol': symbol}, sort=[('start_time', pymongo.DESCENDING)])
+        if last_metadata is None:
+            raise NoDataFoundException('No metadata found for symbol {}'.format(symbol))
+
+        self.find_one_and_delete({'symbol': symbol}, sort=[('start_time', pymongo.DESCENDING)])
+        mongo_retry(self.find_one_and_update)({'symbol': symbol}, {'$unset': {'end_time': ''}},
+                                              sort=[('start_time', pymongo.DESCENDING)])
+
+        return last_metadata
+
+    @mongo_retry
+    def purge(self, symbol):
+        """
+        Delete all metadata of `symbol`
+
+        Parameters
+        ----------
+        symbol : `str`
+            symbol name to delete
+        """
+        logger.warning("Deleting entire metadata history for %r from %r" % (symbol, self._arctic_lib.get_name()))
+        self.delete_many({'symbol': symbol})