Merge pull request pandas-dev#597 from dimosped/fix-pymongo-warning-collection-names

fix failing integration tests on travis due to list_collection_names deprecation

Author: dimosped

README.md

@@ -111,7 +111,7 @@ Arctic storage implementations are **pluggable**.  VersionStore is the default.
 Arctic currently works with:
 
  * Python 2.7, 3.4, 3.5, 3.6
- * pymongo >= 3.0
+ * pymongo >= 3.6
  * Pandas
  * MongoDB >= 2.4.x
 
@@ -137,6 +137,7 @@ It wouldn't be possible without the work of the AHL Data Engineering Team includ
  * [Wilfred Hughes](https://github.com/wilfred)
  * [Edward Easton](https://github.com/eeaston)
  * [Bryant Moscon](https://github.com/bmoscon)
+ * [Dimosthenis Pediaditakis](https://github.com/dimosped)
  * ... and many others ...
 
 Contributions welcome!

arctic/_util.py

@@ -1,10 +1,22 @@
 from pandas import DataFrame
 from pandas.util.testing import assert_frame_equal
-from pymongo.errors import OperationFailure
 import logging
+import pymongo
+
 
 logger = logging.getLogger(__name__)
 
+# Avoid import-time extra logic
+_use_new_count_api = None
+
+
+def _detect_new_count_api():
+    try:
+        mongo_v = [int(v) for v in pymongo.version.split('.')]
+        return mongo_v[0] >= 3 and mongo_v[1] >= 7
+    except:
+        return False
+
 
 def indent(s, num_spaces):
     s = s.split('\n')
@@ -47,7 +59,7 @@ def enable_sharding(arctic, library_name, hashed=True, key='symbol'):
     library_name = lib.get_top_level_collection().name
     try:
         c.admin.command('enablesharding', dbname)
-    except OperationFailure as e:
+    except pymongo.errors.OperationFailure as e:
         if 'already enabled' not in str(e):
             raise
     if not hashed:
@@ -56,3 +68,15 @@ def enable_sharding(arctic, library_name, hashed=True, key='symbol'):
     else:
         logger.info("Hash sharding '" + key + "' on: " + dbname + '.' + library_name)
         c.admin.command('shardCollection', dbname + '.' + library_name, key={key: 'hashed'})
+
+
+def mongo_count(collection, filter=None, **kwargs):
+    filter = {} if filter is None else filter
+    global _use_new_count_api
+    _use_new_count_api = _detect_new_count_api() if _use_new_count_api is None else _use_new_count_api
+    # This is a temporary compatibility fix for compatibility with pymongo>=3.7, and also avoid deprecation warnings
+    if _use_new_count_api:
+        # Projection is ignored for count_documents
+        return collection.count_documents(filter=filter, **kwargs)
+    else:
+        return collection.count(filter=filter, **kwargs)

arctic/arctic.py

@@ -181,11 +181,11 @@ def list_libraries(self):
         libs = []
         for db in self._conn.list_database_names():
             if db.startswith(self.DB_PREFIX + '_'):
-                for coll in self._conn[db].collection_names():
+                for coll in self._conn[db].list_collection_names():
                     if coll.endswith(self.METADATA_COLL):
                         libs.append(db[len(self.DB_PREFIX) + 1:] + "." + coll[:-1 * len(self.METADATA_COLL) - 1])
             elif db == self.DB_PREFIX:
-                for coll in self._conn[db].collection_names():
+                for coll in self._conn[db].list_collection_names():
                     if coll.endswith(self.METADATA_COLL):
                         libs.append(coll[:-1 * len(self.METADATA_COLL) - 1])
         return libs
@@ -212,9 +212,9 @@ def initialize_library(self, library, lib_type=VERSION_STORE, **kwargs):
         # check that we don't create too many namespaces
         # can be disabled check_library_count=False
         check_library_count = kwargs.pop('check_library_count', True)
-        if len(self._conn[l.database_name].collection_names()) > 5000 and check_library_count:
+        if len(self._conn[l.database_name].list_collection_names()) > 5000 and check_library_count:
             raise ArcticException("Too many namespaces %s, not creating: %s" %
-                                  (len(self._conn[l.database_name].collection_names()), library))
+                                  (len(self._conn[l.database_name].list_collection_names()), library))
         l.set_library_type(lib_type)
         LIBRARY_TYPES[lib_type].initialize_library(l, **kwargs)
         # Add a 10G quota just in case the user is calling this with API.
@@ -233,11 +233,11 @@ def delete_library(self, library):
         """
         l = ArcticLibraryBinding(self, library)
         colname = l.get_top_level_collection().name
-        if not [c for c in l._db.collection_names(False) if re.match(r"^{}([\.].*)?$".format(colname), c)]:
+        if not [c for c in l._db.list_collection_names(False) if re.match(r"^{}([\.].*)?$".format(colname), c)]:
             logger.info('Nothing to delete. Arctic library %s does not exist.' % colname)
         logger.info('Dropping collection: %s' % colname)
         l._db.drop_collection(colname)
-        for coll in l._db.collection_names():
+        for coll in l._db.list_collection_names():
             if coll.startswith(colname + '.'):
                 logger.info('Dropping collection: %s' % coll)
                 l._db.drop_collection(coll)
@@ -352,7 +352,7 @@ def rename_library(self, from_lib, to_lib):
 
         logger.info('Renaming collection: %s' % colname)
         l._db[colname].rename(to_colname)
-        for coll in l._db.collection_names():
+        for coll in l._db.list_collection_names():
             if coll.startswith(colname + '.'):
                 l._db[coll].rename(coll.replace(colname, to_colname))
 

arctic/chunkstore/chunkstore.py

@@ -11,7 +11,7 @@
 from pymongo.errors import OperationFailure
 
 from ..decorators import mongo_retry
-from .._util import indent
+from .._util import indent, mongo_count
 from ..serialization.numpy_arrays import FrametoArraySerializer, DATA, METADATA, COLUMNS
 from .date_chunker import DateChunker, START, END
 from .passthrough_chunker import PassthroughChunker
@@ -87,7 +87,7 @@ def _check_invalid_segment(self):
         # Issue 442
         # for legacy data that was incorectly marked with segment start of -1
         for symbol in self.list_symbols():
-            if self._collection.find({SYMBOL: symbol, SEGMENT: -1}).count() > 1:
+            if mongo_count(self._collection, filter={SYMBOL: symbol, SEGMENT: -1}) > 1:
                 logger.warning("Symbol %s has malformed segments. Data must be rewritten or fixed with chunkstore segment_id_repair tool" % symbol)
 
     @mongo_retry
@@ -153,7 +153,7 @@ def delete(self, symbol, chunk_range=None, audit=None):
                 # update symbol metadata (rows and chunk count)
                 sym = self._get_symbol_info(symbol)
                 sym[LEN] -= row_adjust
-                sym[CHUNK_COUNT] = self._collection.count({SYMBOL: symbol})
+                sym[CHUNK_COUNT] = mongo_count(self._collection, filter={SYMBOL: symbol})
                 self._symbols.replace_one({SYMBOL: symbol}, sym)
 
         else:
@@ -458,7 +458,7 @@ def __update(self, sym, item, metadata=None, combine_method=None, chunk_range=No
             meta = data[METADATA]
 
             chunk_count = int(len(data[DATA]) / MAX_CHUNK_SIZE + 1)
-            seg_count = self._collection.count({SYMBOL: symbol, START: start, END: end})
+            seg_count = mongo_count(self._collection, filter={SYMBOL: symbol, START: start, END: end})
             # remove old segments for this chunk in case we now have less
             # segments than we did before
             if seg_count > chunk_count:

arctic/fixtures/arctic.py

@@ -324,6 +324,12 @@ def user_library(arctic, user_library_name):
 def overlay_library(arctic, overlay_library_name):
     """ Overlay library fixture, returns a pair of libs, read-write: ${name} and read-only: ${name}_RAW
     """
+    # Call _create_overlay_library to avoid:
+    #  RemovedInPytest4Warning: Fixture overlay_library called directly. Fixtures are not meant to be called directly
+    return _overlay_library(arctic, overlay_library)
+
+
+def _overlay_library(arctic, overlay_library_name):
     rw_name = overlay_library_name
     ro_name = '{}_RAW'.format(overlay_library_name)
     arctic.initialize_library(rw_name, m.VERSION_STORE, segment='year')
@@ -333,6 +339,12 @@ def overlay_library(arctic, overlay_library_name):
 
 @pytest.fixture(scope="function")
 def tickstore_lib(arctic, library_name):
+    # Call _create_overlay_library to avoid:
+    #  RemovedInPytest4Warning: Fixture overlay_library called directly. Fixtures are not meant to be called directly
+    return _tickstore_lib(arctic, library_name)
+
+
+def _tickstore_lib(arctic, library_name):
     arctic.initialize_library(library_name, TICK_STORE_TYPE)
     return arctic.get_library(library_name)
 

arctic/store/_ndarray_store.py

@@ -8,6 +8,7 @@
 import pymongo
 from pymongo.errors import OperationFailure, DuplicateKeyError
 
+from arctic._util import mongo_count
 from ..decorators import mongo_retry
 from ..exceptions import UnhandledDtypeException, DataIntegrityException
 from ._version_store_utils import checksum, version_base_or_id, _fast_check_corruption
@@ -455,8 +456,7 @@ def check_written(self, collection, symbol, version):
         parent_id = version_base_or_id(version)
 
         # Check all the chunks are in place
-        seen_chunks = collection.find({'symbol': symbol, 'parent': parent_id},
-                                      ).count()
+        seen_chunks = mongo_count(collection, filter={'symbol': symbol, 'parent': parent_id})
 
         if seen_chunks != version['segment_count']:
             segments = [x['segment'] for x in collection.find({'symbol': symbol, 'parent': parent_id},

arctic/store/_version_store_utils.py

@@ -11,6 +11,8 @@
 from pandas.compat import pickle_compat
 from pymongo.errors import OperationFailure
 
+from arctic._util import mongo_count
+
 
 def _split_arrs(array_2d, slices):
     """
@@ -113,7 +115,7 @@ def analyze_symbol(l, sym, from_ver, to_ver, do_reads=False):
             matching = 0
         else:
             spec = {'symbol': sym, 'parent': v.get('base_version_id', v['_id']), 'segment': {'$lt': v.get('up_to', 0)}}
-            matching = l._collection.find(spec).count() if not is_deleted else 0
+            matching = mongo_count(l._collection, filter=spec) if not is_deleted else 0
 
         base_id = v.get('base_version_id')
         snaps = ['/'.join((str(x), str(x.generation_time))) for x in v.get('parent')] if v.get('parent') else None
@@ -129,7 +131,21 @@ def analyze_symbol(l, sym, from_ver, to_ver, do_reads=False):
         corrupted = not is_deleted and (is_corrupted(l, sym, v) if do_reads else fast_is_corrupted(l, sym, v))
 
         logging.info(
-            "v{: <6} {: <6} {: <5} ({: <20}):   expected={: <6} found={: <6} last_row={: <10} new_rows={: <10} append count={: <10} append_size={: <10} type={: <14} {: <14} base={: <24}/{: <28} snap={: <30}[{:.1f} mins delayed] {: <20} {: <20}".format(
+            "v{: <6} "
+            "{: <6} "
+            "{: <5} "
+            "({: <20}):   "
+            "expected={: <6} "
+            "found={: <6} "
+            "last_row={: <10} "
+            "new_rows={: <10} "
+            "append count={: <10} "
+            "append_size={: <10} "
+            "type={: <14} {: <14} "
+            "base={: <24}/{: <28} "
+            "snap={: <30}[{:.1f} mins delayed] "
+            "{: <20} "
+            "{: <20}".format(
                 n,
                 prev_v_diff,
                 'DEL' if is_deleted else 'ALIVE',
@@ -144,7 +160,7 @@ def analyze_symbol(l, sym, from_ver, to_ver, do_reads=False):
                 'meta-same' if meta_match_with_prev else 'meta-changed',
                 str(base_id),
                 str(base_id.generation_time) if base_id else '',
-                snaps,
+                str(snaps),
                 delta_snap_creation,
                 'PREV_MISSING' if prev_n < n - 1 else '',
                 'CORRUPTED VERSION' if corrupted else '')
@@ -159,7 +175,7 @@ def analyze_symbol(l, sym, from_ver, to_ver, do_reads=False):
             hashlib.sha1(seg['sha']).hexdigest(),
             seg.get('segment'),
             'compressed' if seg.get('compressed', False) else 'raw',
-            [str(p) for p in seg.get('parent', [])]
+            str([str(p) for p in seg.get('parent', [])])
         ))
 
 
@@ -193,7 +209,7 @@ def _fast_check_corruption(collection, sym, v, check_count, check_last_segment,
         #     collection.find_one(spec, {'segment': 1}, sort=[('segment', pymongo.DESCENDING)])
 
         if check_count:
-            total_segments = collection.find(spec, {'segment': 1}).count()
+            total_segments = mongo_count(collection, filter=spec)
             # Quick check: compare segment count
             if total_segments != v.get('segment_count', 0):
                 return True  # corrupted, don't proceed with fetching from mongo the first hit

arctic/store/bson_store.py

@@ -1,7 +1,7 @@
 import logging
 from pymongo.errors import OperationFailure
 from ..decorators import mongo_retry
-from .._util import enable_sharding
+from .._util import enable_sharding, mongo_count
 
 logger = logging.getLogger(__name__)
 
@@ -161,7 +161,7 @@ def count(self, filter, **kwargs):
         """
         See http://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.count
         """
-        return self._collection.count(filter, **kwargs)
+        return mongo_count(self._collection, filter=filter, **kwargs)
 
     @mongo_retry
     def aggregate(self, pipeline, **kwargs):

arctic/store/version_store.py

@@ -6,7 +6,7 @@
 import pymongo
 from pymongo.errors import OperationFailure, AutoReconnect, DuplicateKeyError
 
-from .._util import indent, enable_sharding
+from .._util import indent, enable_sharding, mongo_count
 from ..date import mktz, datetime_to_ms, ms_to_datetime
 from ..decorators import mongo_retry
 from ..exceptions import NoDataFoundException, DuplicateSnapshotException, \
@@ -41,7 +41,7 @@ class VersionStore(object):
     def initialize_library(cls, arctic_lib, hashed=True, **kwargs):
         c = arctic_lib.get_top_level_collection()
 
-        if '%s.changes' % c.name not in mongo_retry(c.database.collection_names)():
+        if '%s.changes' % c.name not in mongo_retry(c.database.list_collection_names)():
             # 32MB buffer for change notifications
             mongo_retry(c.database.create_collection)('%s.changes' % c.name, capped=True, size=32 * 1024 * 1024)
 
@@ -88,7 +88,7 @@ def _reset(self):
         self._snapshots = self._collection.snapshots
         self._versions = self._collection.versions
         self._version_nums = self._collection.version_nums
-        self._publish_changes = '%s.changes' % self._collection.name in self._collection.database.collection_names()
+        self._publish_changes = '%s.changes' % self._collection.name in self._collection.database.list_collection_names()
         if self._publish_changes:
             self._changes = self._collection.changes
 
@@ -1034,7 +1034,7 @@ def _cleanup_orphaned_chunks(self, dry_run):
             if len(leaked_versions):
                 logger.info("%s leaked %d versions" % (symbol, len(leaked_versions)))
             for x in leaked_versions:
-                chunk_count = chunks_coll.find({'symbol': symbol, 'parent': x}).count()
+                chunk_count = mongo_count(chunks_coll, filter={'symbol': symbol, 'parent': x})
                 logger.info("%s: Missing Version %s (%s) ; %s chunks ref'd" % (symbol,
                                                                                x.generation_time,
                                                                                x,
@@ -1078,7 +1078,7 @@ def _cleanup_orphaned_versions(self, dry_run):
         if len(leaked_snaps):
             logger.info("leaked %d snapshots" % (len(leaked_snaps)))
         for x in leaked_snaps:
-            ver_count = versions_coll.find({'parent': x}).count()
+            ver_count = mongo_count(versions_coll, filter={'parent': x})
             logger.info("Missing Snapshot %s (%s) ; %s versions ref'd" % (x.generation_time,
                                                                           x,
                                                                           ver_count

tests/integration/chunkstore/test_chunkstore.py

@@ -2,6 +2,8 @@
 from datetime import datetime as dt
 from datetime import timedelta
 from pandas.util.testing import assert_frame_equal, assert_series_equal
+
+from arctic._util import mongo_count
 from arctic.date import DateRange
 from arctic.exceptions import NoDataFoundException
 import pandas as pd
@@ -943,7 +945,7 @@ def test_delete_range_segment(chunkstore_lib):
     chunkstore_lib.delete('test_df', chunk_range=pd.date_range(dt(2016, 1, 1), dt(2016, 1, 1)))
     read_df = chunkstore_lib.read('test_df')
     assert(read_df.equals(dg))
-    assert(chunkstore_lib._collection.count({'sy': 'test_df'}) == 1)
+    assert(mongo_count(chunkstore_lib._collection, {'sy': 'test_df'}) == 1)
 
 
 def test_size_chunk_update(chunkstore_lib):
@@ -960,7 +962,7 @@ def test_size_chunk_update(chunkstore_lib):
     read_df = chunkstore_lib.read('test_df')
 
     assert_frame_equal(dh, read_df)
-    assert(chunkstore_lib._collection.count({'sy': 'test_df'}) == 1)
+    assert mongo_count(chunkstore_lib._collection, filter={'sy': 'test_df'}) == 1
 
 
 def test_size_chunk_multiple_update(chunkstore_lib):
@@ -976,7 +978,7 @@ def test_size_chunk_multiple_update(chunkstore_lib):
     expected = pd.concat([df_large, df_small]).reset_index(drop=True)
 
     assert_frame_equal(expected, read_df)
-    assert(chunkstore_lib._collection.count({'sy': 'test_df'}) == 3)
+    assert mongo_count(chunkstore_lib._collection, filter={'sy': 'test_df'}) == 3
 
 
 def test_get_chunk_range(chunkstore_lib):