Allow sharding on keys other than symbol, default BSONStore to shard on _id (pandas-dev#374)

* Allow sharding on keys other than symbol, default BSONStore to shard on _id

Author: achamayou

arctic/_util.py

@@ -23,7 +23,24 @@ def are_equals(o1, o2, **kwargs):
         return False
 
 
-def enable_sharding(arctic, library_name, hashed=True):
+def enable_sharding(arctic, library_name, hashed=True, key='symbol'):
+    """
+    Enable sharding on a library
+
+    Parameters:
+    -----------
+    arctic: `arctic.Arctic` Arctic class
+
+    library_name: `basestring` library name
+
+    hashed: `bool` if True, use hashed sharding, if False, use range sharding
+            See https://docs.mongodb.com/manual/core/hashed-sharding/,
+            as well as https://docs.mongodb.com/manual/core/ranged-sharding/ for details.
+
+    key: `basestring` key to be used for sharding. Defaults to 'symbol', applicable to
+         all of Arctic's built-in stores except for BSONStore, which typically uses '_id'.
+         See https://docs.mongodb.com/manual/core/sharding-shard-key/ for details.
+    """
     c = arctic._conn
     lib = arctic[library_name]._arctic_lib
     dbname = lib._db.name
@@ -34,8 +51,8 @@ def enable_sharding(arctic, library_name, hashed=True):
         if not 'already enabled' in str(e):
             raise
     if not hashed:
-        logger.info("Range sharding 'symbol' on: " + dbname + '.' + library_name)
-        c.admin.command('shardCollection', dbname + '.' + library_name, key={'symbol': 1})
+        logger.info("Range sharding '" + key + "' on: " + dbname + '.' + library_name)
+        c.admin.command('shardCollection', dbname + '.' + library_name, key={key: 1})
     else:
-        logger.info("Hash sharding 'symbol' on: " + dbname + '.' + library_name)
-        c.admin.command('shardCollection', dbname + '.' + library_name, key={'symbol': 'hashed'})
+        logger.info("Hash sharding '" + key + "' on: " + dbname + '.' + library_name)
+        c.admin.command('shardCollection', dbname + '.' + library_name, key={key: 'hashed'})

arctic/store/bson_store.py

@@ -30,7 +30,10 @@ def __init__(self, arctic_lib):
     def initialize_library(cls, arctic_lib, hashed=True, **kwargs):
         logger.info("Trying to enable sharding...")
         try:
-            enable_sharding(arctic_lib.arctic, arctic_lib.get_name(), hashed=hashed)
+            if not hashed:
+                logger.warning("Ignored hashed=False when enabling sharding, only hashed=True "
+                               " makes sense when they key is an ObjectId")
+            enable_sharding(arctic_lib.arctic, arctic_lib.get_name(), hashed=True, key='_id')
         except OperationFailure as exception:
             logger.warning(("Library created, but couldn't enable sharding: "
                             "%s. This is OK if you're not 'admin'"), exception)

tests/unit/store/test_bson_store.py

@@ -11,8 +11,11 @@ def test_initialize_library():
     with patch('arctic.store.bson_store.enable_sharding', autospec=True) as enable_sharding:
         arctic_lib.get_top_level_collection.return_value.database.create_collection.__name__ = 'some_name'
         arctic_lib.get_top_level_collection.return_value.database.collection_names.__name__ = 'some_name'
-        BSONStore.initialize_library(arctic_lib, hashed=sentinel.hashed)
-        assert enable_sharding.call_args_list == [call(arctic_lib.arctic, arctic_lib.get_name(), hashed=sentinel.hashed)]
+        BSONStore.initialize_library(arctic_lib, hashed=True)
+        BSONStore.initialize_library(arctic_lib, hashed=False)
+        # Check we always set the sharding to be hashed, regarless of user input
+        assert enable_sharding.call_args_list == [call(arctic_lib.arctic, arctic_lib.get_name(), hashed=True, key='_id'),
+                                                  call(arctic_lib.arctic, arctic_lib.get_name(), hashed=True, key='_id')]
 
 
 def test_find():