From 0797c341ff8f9bb7ccd33ddb6b5dd60e23f3ccca Mon Sep 17 00:00:00 2001 From: Mike Dirolf Date: Thu, 14 May 2009 15:50:31 -0400 Subject: [PATCH] ensure_index --- pymongo/collection.py | 58 +++++++++++++++++++++++++++--- pymongo/connection.py | 58 ++++++++++++++++++++++++++++++ pymongo/database.py | 2 ++ pymongo/master_slave_connection.py | 6 ++++ test/test_collection.py | 47 +++++++++++++++++++++++- 5 files changed, 166 insertions(+), 5 deletions(-) diff --git a/pymongo/collection.py b/pymongo/collection.py index 3bc988ebf..6ea1d3cd9 100644 --- a/pymongo/collection.py +++ b/pymongo/collection.py @@ -321,7 +321,7 @@ class Collection(object): """ 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 @@ class Collection(object): :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 @@ class Collection(object): 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)]), diff --git a/pymongo/connection.py b/pymongo/connection.py index 51c5573bb..22a94c192 100644 --- a/pymongo/connection.py +++ b/pymongo/connection.py @@ -22,6 +22,7 @@ import logging import threading import random import errno +import datetime from errors import ConnectionFailure, InvalidName, OperationFailure, ConfigurationError from database import Database @@ -108,6 +109,9 @@ class Connection(object): 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 @@ class Connection(object): 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 @@ class Connection(object): 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): diff --git a/pymongo/database.py b/pymongo/database.py index 032ff09f0..6ce516140 100644 --- a/pymongo/database.py +++ b/pymongo/database.py @@ -196,6 +196,8 @@ class Database(object): 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 diff --git a/pymongo/master_slave_connection.py b/pymongo/master_slave_connection.py index 9e337e7d2..4f8dd2b33 100644 --- a/pymongo/master_slave_connection.py +++ b/pymongo/master_slave_connection.py @@ -207,3 +207,9 @@ class MasterSlaveConnection(object): 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) diff --git a/test/test_collection.py b/test/test_collection.py index 0dc5afc0d..95217b207 100644 --- a/test/test_collection.py +++ b/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 @@ from pymongo.son import SON 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 @@ class TestCollection(unittest.TestCase): (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")