ensure_index

Author: Mike Dirolf

pymongo/collection.py

@@ -321,7 +321,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, direction=None, unique=False):
+    def create_index(self, key_or_list, direction=None, unique=False, ttl=300):
         """Creates an index on this collection.
 
         Takes either a single key and a direction, or a list of (key, direction)
@@ -331,27 +331,76 @@ def create_index(self, key_or_list, direction=None, unique=False):
 
         :Parameters:
           - `key_or_list`: a single key or a list of (key, direction) pairs
-            specifying the index to ensure
+            specifying the index to create
           - `direction` (optional): must be included if key_or_list is a single
             key, otherwise must be None
           - `unique` (optional): should this index guarantee uniqueness?
+          - `ttl` (optional): time window (in seconds) during which this index
+            will be recognized by subsequent calls to `ensure_index` - see
+            documentation for `ensure_index` for details
         """
         to_save = SON()
         keys = pymongo._index_list(key_or_list, direction)
-        to_save["name"] = self._gen_index_name(keys)
+        name = self._gen_index_name(keys)
+        to_save["name"] = name
         to_save["ns"] = self.full_name()
         to_save["key"] = pymongo._index_document(keys)
         to_save["unique"] = unique
 
-        self.__database.system.indexes.save(to_save, False)
+        self.database().connection()._cache_index(self.__database.name(),
+                                                  self.name(),
+                                                  name, ttl)
+
+        self.database().system.indexes.save(to_save, False)
         return to_save["name"]
 
+    def ensure_index(self, key_or_list, direction=None, unique=False, ttl=300):
+        """Ensures that an index exists on this collection.
+
+        Takes either a single key and a direction, or a list of (key, direction)
+        pairs. The key(s) must be an instance of (str, unicode), and the
+        direction(s) must be one of (`pymongo.ASCENDING`, `pymongo.DESCENDING`).
+
+        Unlike `create_index`, which attempts to create an index
+        unconditionally, `ensure_index` takes advantage of some caching within
+        the driver such that it only attempts to create indexes that might
+        not already exist. When an index is created (or ensured) by PyMongo
+        it is "remembered" for `ttl` seconds. Repeated calls to `ensure_index`
+        within that time limit will be lightweight - they will not attempt to
+        actually create the index.
+
+        Care must be taken when the database is being accessed through multiple
+        connections at once. If an index is created using PyMongo and then
+        deleted using another connection any call to `ensure_index` within the
+        cache window will fail to re-create the missing index.
+
+        Returns the name of the created index if an index is actually created.
+        Returns None if the index already exists.
+
+        :Parameters:
+          - `key_or_list`: a single key or a list of (key, direction) pairs
+            specifying the index to ensure
+          - `direction` (optional): must be included if key_or_list is a single
+            key, otherwise must be None
+          - `unique` (optional): should this index guarantee uniqueness?
+          - `ttl` (optional): time window (in seconds) during which this index
+            will be recognized by subsequent calls to `ensure_index`
+        """
+        keys = pymongo._index_list(key_or_list, direction)
+        name = 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, direction, unique, ttl)
+        return None
+
     def drop_indexes(self):
         """Drops all indexes on this collection.
 
         Can be used on non-existant collections or collections with no indexes.
         Raises OperationFailure on an error.
         """
+        self.database().connection()._purge_index(self.database().name(), self.name())
         self.drop_index(u"*")
 
     def drop_index(self, index_or_name):
@@ -373,6 +422,7 @@ def drop_index(self, index_or_name):
         if not isinstance(name, types.StringTypes):
             raise TypeError("index_or_name must be an index name or list")
 
+        self.database().connection()._purge_index(self.database().name(), self.name(), name)
         self.__database._command(SON([("deleteIndexes",
                                        self.__collection_name),
                                       ("index", name)]),

pymongo/connection.py

@@ -22,6 +22,7 @@
 import threading
 import random
 import errno
+import datetime
 
 from errors import ConnectionFailure, InvalidName, OperationFailure, ConfigurationError
 from database import Database
@@ -108,6 +109,9 @@ def __init__(self, host=None, port=None, pool_size=None,
         self.__sockets = [None for _ in range(self.__pool_size)]
         self.__currently_resetting = False
 
+        # cache of existing indexes used by ensure_index ops
+        self.__index_cache = {}
+
         if _connect:
             self.__find_master()
 
@@ -181,6 +185,59 @@ def __master(self, sock):
                 port = int(strings[1])
             return (strings[0], port)
 
+    def _cache_index(self, database_name, collection_name, index_name, ttl):
+        """Add an index to the index cache for ensure_index operations.
+
+        Return True if the index has been newly cached or if the index had
+        expired and is being re-cached.
+
+        Return False if the index exists and is valid.
+        """
+        now = datetime.datetime.utcnow()
+        expire = datetime.timedelta(seconds=ttl) + now
+
+        if database_name not in self.__index_cache:
+            self.__index_cache[database_name] = {}
+            self.__index_cache[database_name][collection_name] = {}
+            self.__index_cache[database_name][collection_name][index_name] = expire
+            return True
+
+        if collection_name not in self.__index_cache[database_name]:
+            self.__index_cache[database_name][collection_name] = {}
+            self.__index_cache[database_name][collection_name][index_name] = expire
+            return True
+
+        if index_name in self.__index_cache[database_name][collection_name]:
+            if now < self.__index_cache[database_name][collection_name][index_name]:
+                return False
+
+        self.__index_cache[database_name][collection_name][index_name] = expire
+        return True
+
+    def _purge_index(self, database_name, collection_name=None, index_name=None):
+        """Purge an index from the index cache.
+
+        If `index_name` is None purge an entire collection.
+
+        If `collection_name` is None purge an entire database.
+        """
+        if not database_name in self.__index_cache:
+            return
+
+        if collection_name is None:
+            del self.__index_cache[database_name]
+            return
+
+        if not collection_name in self.__index_cache[database_name]:
+            return
+
+        if index_name is None:
+            del self.__index_cache[database_name][collection_name]
+            return
+
+        if index_name in self.__index_cache[database_name][collection_name]:
+            del self.__index_cache[database_name][collection_name][index_name]
+
     def host(self):
         """Get the connection's current host.
         """
@@ -574,6 +631,7 @@ def drop_database(self, name_or_database):
         if not isinstance(name, types.StringTypes):
             raise TypeError("name_or_database must be an instance of (Database, str, unicode)")
 
+        self._purge_index(name)
         self[name]._command({"dropDatabase": 1})
 
     def __iter__(self):

pymongo/database.py

@@ -196,6 +196,8 @@ def drop_collection(self, name_or_collection):
         if not isinstance(name, types.StringTypes):
             raise TypeError("name_or_collection must be an instance of (Collection, str, unicode)")
 
+        self.connection()._purge_index(self.name(), name)
+
         if name not in self.collection_names():
             return
 

pymongo/master_slave_connection.py

@@ -207,3 +207,9 @@ def __iter__(self):
 
     def next(self):
         raise TypeError("'MasterSlaveConnection' object is not iterable")
+
+    def _cache_index(self, database_name, collection_name, index_name, ttl):
+        return self.__master._cache_index(database_name, collection_name, index_name, ttl)
+
+    def _purge_index(self, database_name, collection_name=None, index_name=None):
+        return self.__master._purge_index(database_name, collection_name, index_name)

test/test_collection.py

@@ -15,6 +15,7 @@
 """Test the collection module."""
 import unittest
 import re
+import time
 import sys
 sys.path[0:0] = [""]
 
@@ -29,7 +30,8 @@
 
 class TestCollection(unittest.TestCase):
     def setUp(self):
-        self.db = get_connection().pymongo_test
+        self.connection = get_connection()
+        self.db = self.connection.pymongo_test
 
     def test_collection(self):
         self.assertRaises(TypeError, Collection, self.db, 5)
@@ -92,6 +94,49 @@ def test_create_index(self):
                           (u"key", SON([(u"hello", -1),
                                         (u"world", 1)]))]) in list(db.system.indexes.find({"ns": u"pymongo_test.test"})))
 
+    def test_ensure_index(self):
+        db = self.db
+
+        db.test.drop_indexes()
+        self.assertEqual("hello_1", db.test.create_index("hello", ASCENDING))
+        self.assertEqual("hello_1", db.test.create_index("hello", ASCENDING))
+
+        self.assertEqual("goodbye_1", db.test.ensure_index("goodbye", ASCENDING))
+        self.assertEqual(None, db.test.ensure_index("goodbye", ASCENDING))
+
+        db.test.drop_indexes()
+        self.assertEqual("goodbye_1", db.test.ensure_index("goodbye", ASCENDING))
+        self.assertEqual(None, db.test.ensure_index("goodbye", ASCENDING))
+
+        db.test.drop_index("goodbye_1")
+        self.assertEqual("goodbye_1", db.test.ensure_index("goodbye", ASCENDING))
+        self.assertEqual(None, db.test.ensure_index("goodbye", ASCENDING))
+
+        db.drop_collection("test")
+        self.assertEqual("goodbye_1", db.test.ensure_index("goodbye", ASCENDING))
+        self.assertEqual(None, db.test.ensure_index("goodbye", ASCENDING))
+
+        db_name = self.db.name()
+        self.connection.drop_database(self.db.name())
+        self.assertEqual("goodbye_1", db.test.ensure_index("goodbye", ASCENDING))
+        self.assertEqual(None, db.test.ensure_index("goodbye", ASCENDING))
+
+        db.test.drop_index("goodbye_1")
+        self.assertEqual("goodbye_1", db.test.create_index("goodbye", ASCENDING))
+        self.assertEqual(None, db.test.ensure_index("goodbye", ASCENDING))
+
+        db.test.drop_index("goodbye_1")
+        self.assertEqual("goodbye_1", db.test.ensure_index("goodbye", ASCENDING,
+                                                           ttl=1))
+        time.sleep(1.1)
+        self.assertEqual("goodbye_1", db.test.ensure_index("goodbye", ASCENDING))
+
+        db.test.drop_index("goodbye_1")
+        self.assertEqual("goodbye_1", db.test.create_index("goodbye", ASCENDING,
+                                                           ttl=1))
+        time.sleep(1.1)
+        self.assertEqual("goodbye_1", db.test.ensure_index("goodbye", ASCENDING))
+
     def test_index_on_binary(self):
         db = self.db
         db.drop_collection("test")