Index Support

.gitignore

@@ -6,3 +6,7 @@ Mongothon.egg-info
 build
 dist
 env
+#*#
+*#
+.#*
+*~

mongothon/__init__.py

@@ -2,7 +2,7 @@
 from inflection import camelize
 from document import Document
 from model import Model, NotFoundException
-from schema import Schema
+from schema import Schema, IndexSpec
 from schemer import Mixed, ValidationException, Array
 
 

mongothon/model.py

@@ -7,7 +7,7 @@
 from .exceptions import NotFoundException
 from .events import EventHandlerRegistrar
 from .scopes import STANDARD_SCOPES
-
+from .schema import IndexSpec
 
 OBJECTIDEXPR = re.compile(r"^[a-fA-F0-9]{24}$")
 
@@ -109,6 +109,46 @@ def apply_defaults(self):
         self.schema.apply_defaults(self)
         self.emit('did_apply_defaults')
 
+    @classmethod
+    def apply_index(cls, index):
+        index.apply_to(cls.get_collection())
+
+    @classmethod
+    def apply_indexes(cls):
+        for index in cls.schema.indexes:
+            index.apply_to(cls.get_collection())
+
+    @classmethod
+    def _existing_indexes(cls):
+        """
+          >>> db.<col>.index_information()
+          {u'_id_': {u'key': [(u'_id', 1)]},
+           u'x_1': {u'unique': True, u'key': [(u'x', 1)]}}
+        """
+        info = cls.get_collection().index_information()
+        indexes = []
+        for k, v in info.iteritems():
+            if k == '_id_': # this is the primary key index, not interesting
+                continue
+            index = IndexSpec(k, v['key'])
+            for arg, val in v.iteritems():
+                if arg == 'key':
+                    continue
+                index.kwargs[arg] == val
+            index.validate()
+            indexes.append(index)
+        return indexes
+
+    @classmethod
+    def applied_indexes(cls):
+        return [i.name for i in cls._existing_indexes()]
+
+    @classmethod
+    def unapplied_indexes(cls):
+        existing_names = set([i.name for i in cls._existing_indexes()])
+        expected_names = [i.name for i in cls.schema.indexes]
+        return [name for name in expected_names if name not in existing_names]
+
     @classmethod
     def get_collection(cls):
         if not hasattr(cls, '_collection'):

mongothon/schema.py

@@ -1,13 +1,39 @@
 from bson.objectid import ObjectId
+import pymongo
 import schemer
 
 
+class IndexSpec(object):
+    def __init__(self, name, key_spec, **kwargs):
+        self.name = name
+        self.key_spec = key_spec
+        self.kwargs = kwargs
+
+    def apply_to(self, collection):
+        collection.create_index(self.key_spec, name=self.name, **self.kwargs)
+
+    def validate(self):
+        if not self.name:
+            raise ValueError("Must specify a non-nil name for every index")
+        if not self.key_spec:
+            raise ValueError("Must specify the actual index for {}".format(self.name))
+        for name, index_type in self.key_spec:
+            if index_type not in {pymongo.ASCENDING, pymongo.DESCENDING, pymongo.HASHED}:
+                raise ValueError('Unsupported Index Type {} for {}'.format(index_type, self.name))
+        return self
+
+
+
 class Schema(schemer.Schema):
     """A Schema encapsulates the structure and constraints of a Mongo document."""
 
-    def __init__(self, doc_spec, **kwargs):
+    indexes = []
+
+    def __init__(self, doc_spec, indexes=[], **kwargs):
         super(Schema, self).__init__(doc_spec, **kwargs)
 
         # Every mongothon schema should expect an ID field.
         if '_id' not in self._doc_spec:
             self._doc_spec['_id'] = {"type": ObjectId}
+
+        self.indexes = [i.validate() for i in indexes]

tests/mongothon/model_test.py

@@ -2,7 +2,7 @@
 from pickle import dumps, loads
 from unittest import TestCase
 from mock import Mock, ANY, call, NonCallableMock
-from mongothon import Document, Schema, NotFoundException, Array
+from mongothon import Document, Schema, NotFoundException, Array, IndexSpec
 from mongothon.validators import one_of
 from mongothon.scopes import STANDARD_SCOPES
 from bson import ObjectId
@@ -22,7 +22,8 @@
         "diameter":         {"type": int}
     }))},
     "options":              {"type": Array(basestring)}
-})
+},
+                    indexes=[IndexSpec('make_1', [('make', 1)])])
 
 
 doc = {
@@ -68,6 +69,7 @@ class TestModel(TestCase):
     def setUp(self):
         self.mock_collection = Mock()
         self.mock_collection.name = "car"
+        self.mock_collection.index_information = Mock(return_value={})
         self.Car = create_model(car_schema, self.mock_collection)
         self.CarOffline = create_model_offline(car_schema, lambda: self.mock_collection, 'Car')
         self.car = self.Car(doc)
@@ -126,6 +128,14 @@ def test_constructor_with_kwargs_and_initial_state(self):
     def test_instantiate(self):
         self.assert_predicates(self.car, is_new=True)
 
+    def test_indexes(self):
+        self.assertEqual(['make_1'], self.Car.unapplied_indexes())
+        self.Car.apply_indexes()
+        self.mock_collection.create_index.assert_called_once_with([('make', 1)], name='make_1')
+        self.mock_collection.index_information.return_value = {'make_1': {'key': [('make', 1)]}}
+        self.assertEqual([], self.Car.unapplied_indexes())
+        self.assertEqual(['make_1'], self.Car.applied_indexes())
+
     def test_validation_of_valid_doc(self):
         self.car.validate()
 

tests/mongothon/schema_test.py

@@ -0,0 +1,7 @@
+from mongothon.schema import Schema, IndexSpec
+from mock import Mock
+import unittest
+
+class TestSchema(unittest.TestCase):
+    def test_indexes(self):
+        Schema({}, indexes=[IndexSpec('myindex', [('key', 1)])])