Optional name argument for create_index and ensure_index

Author: carletes

pymongo/collection.py

@@ -490,7 +490,7 @@ def _gen_index_name(self, keys):
         """
         return u"_".join([u"%s_%s" % item for item in keys])
 
-    def create_index(self, key_or_list, unique=False, ttl=300):
+    def create_index(self, key_or_list, unique=False, ttl=300, name=None):
         """Creates an index on this collection.
 
         Takes either a single key or a list of (key, direction) pairs.
@@ -519,6 +519,8 @@ def create_index(self, key_or_list, unique=False, ttl=300):
             this index will be recognized by subsequent calls to
             :meth:`ensure_index` - see documentation for
             :meth:`ensure_index` for details
+          - `name` (optional): name for the index. If none given, a name
+            will be generated.
 
         .. seealso:: :meth:`ensure_index`
 
@@ -528,7 +530,7 @@ def create_index(self, key_or_list, unique=False, ttl=300):
         """
         keys = helpers._index_list(key_or_list)
         index_doc = helpers._index_document(keys)
-        name = self._gen_index_name(keys)
+        name = name is not None and name or self._gen_index_name(keys)
         to_save = SON()
         to_save["name"] = name
         to_save["ns"] = self.__full_name
@@ -542,7 +544,7 @@ def create_index(self, key_or_list, unique=False, ttl=300):
                                               check_keys=False)
         return to_save["name"]
 
-    def ensure_index(self, key_or_list, unique=False, ttl=300):
+    def ensure_index(self, key_or_list, unique=False, ttl=300, name=None):
         """Ensures that an index exists on this collection.
 
         Takes either a single key or a list of (key, direction) pairs.
@@ -577,17 +579,20 @@ def ensure_index(self, key_or_list, unique=False, ttl=300):
           - `ttl` (optional): time window (in seconds) during which
             this index will be recognized by subsequent calls to
             :meth:`ensure_index`
+          - `name` (optional): name for the index. If none given, a name
+            will be generated.
 
         .. seealso:: :meth:`create_index`
         """
         if not isinstance(key_or_list, (str, unicode, list)):
             raise TypeError("key_or_list must either be a single key or a list of (key, direction) pairs")
 
         keys = helpers._index_list(key_or_list)
-        name = self._gen_index_name(keys)
+        name = name is not None and name or self._gen_index_name(keys)
         if self.__database.connection._cache_index(self.__database.name,
                                                    self.__name, name, ttl):
-            return self.create_index(key_or_list, unique=unique, ttl=ttl)
+            return self.create_index(key_or_list, unique=unique, ttl=ttl,
+                                     name=name)
         return None
 
     def drop_indexes(self):

test/test_collection.py

@@ -91,6 +91,11 @@ def test_create_index(self):
             count += 1
         self.assertEqual(count, 3)
 
+        db.test.drop_indexes()
+        ix = db.test.create_index([("hello", DESCENDING),
+                                   ("world", ASCENDING)], name="hello_world")
+        self.assertEquals(ix, "hello_world")
+
         db.test.drop_indexes()
         self.assertEqual(db.system.indexes.find({"ns": u"pymongo_test.test"})
                          .count(), 1)
@@ -127,6 +132,11 @@ def test_ensure_index(self):
                          db.test.ensure_index("goodbye"))
         self.assertEqual(None, db.test.ensure_index("goodbye"))
 
+        db.test.drop_indexes()
+        self.assertEqual("foo",
+                         db.test.ensure_index("goodbye", name="foo"))
+        self.assertEqual(None, db.test.ensure_index("goodbye", name="foo"))
+
         db.test.drop_indexes()
         self.assertEqual("goodbye_1",
                          db.test.ensure_index("goodbye"))