Tweaked SequenceField so that it doesn't increment on creation.

[refs #238]

AUTHORS

@@ -4,3 +4,4 @@ Deepak Thukral <iapain@yahoo.com>
 Florian Schlachter <flori@n-schlachter.de>
 Steve Challis <steve@stevechallis.com>
 Ross Lawley <ross.lawley@gmail.com>
+Wilson Júnior <wilsonpjunior@gmail.com>

docs/changelog.rst

@@ -5,6 +5,7 @@ Changelog
 Changes in dev
 ==============
 
+- Added SequenceField - for creating sequential counters
 - Added update() convenience method to a document
 - Added cascading saves - so changes to Referenced documents are saved on .save()
 - Added select_related() support

mongoengine/base.py

@@ -587,7 +587,8 @@ class BaseDocument(object):
 
         # Set any get_fieldname_display methods
         self.__set_field_display()
-
+        # Flag initialised
+        self._initialised = True
         signals.post_init.send(self.__class__, document=self)
 
     def validate(self):

mongoengine/fields.py

@@ -880,13 +880,27 @@ class GeoPointField(BaseField):
 
 
 class SequenceField(IntField):
+    """Provides a sequental counter.
+
+    ..note:: Although traditional databases often use increasing sequence
+             numbers for primary keys. In MongoDB, the preferred approach is to
+             use Object IDs instead.  The concept is that in a very large
+             cluster of machines, it is easier to create an object ID than have
+             global, uniformly increasing sequence numbers.
+
+    .. versionadded:: 0.5
+    """
+    def __init__(self, collection_name=None, *args, **kwargs):
+        self.collection_name = collection_name or 'mongoengine.counters'
+        return super(SequenceField, self).__init__(*args, **kwargs)
+
     def generate_new_value(self):
         """
-        Generate and Increment counter
+        Generate and Increment the counter
         """
         sequence_id = "{0}.{1}".format(self.owner_document._get_collection_name(),
                                        self.name)
-        collection = _get_db()['mongoengine.counters']
+        collection = _get_db()[self.collection_name]
         counter = collection.find_and_modify(query={"_id": sequence_id},
                                              update={"$inc": {"next": 1}},
                                              new=True,
@@ -894,16 +908,18 @@ class SequenceField(IntField):
         return counter['next']
 
     def __get__(self, instance, owner):
-        if not instance._data:
-            return
 
         if instance is None:
             return self
-
+        if not instance._data:
+            return
         value = instance._data.get(self.name)
-        
-        if not value:
+        if not value and instance._initialised:
             value = self.generate_new_value()
             instance._data[self.name] = value
-
+        return value
+
+    def to_python(self, value):
+        if value is None:
+            value = self.generate_new_value()
         return value

tests/fields.py

@@ -1379,20 +1379,7 @@ class FieldTest(unittest.TestCase):
         self.assertEqual(d2.data, {})
         self.assertEqual(d2.data2, {})
 
-
     def test_sequence_field(self):
-        class Person(Document):
-            id = SequenceField(primary_key=True)
-
-        self.db['mongoengine.counters'].drop()
-        Person.drop_collection()
-        p = Person()
-        p.save()
-
-        p = Person.objects.first()
-        self.assertEqual(p.id, 1)
-
-    def test_multiple_sequence_field(self):
         class Person(Document):
             id = SequenceField(primary_key=True)
             name = StringField()
@@ -1404,18 +1391,76 @@ class FieldTest(unittest.TestCase):
             p = Person(name="Person %s" % x)
             p.save()
 
+        c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
+        self.assertEqual(c['next'], 10)
+
         ids = [i.id for i in Person.objects]
         self.assertEqual(ids, range(1, 11))
 
+        c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
+        self.assertEqual(c['next'], 10)
+
+    def test_multiple_sequence_fields(self):
+        class Person(Document):
+            id = SequenceField(primary_key=True)
+            counter = SequenceField()
+            name = StringField()
+
+        self.db['mongoengine.counters'].drop()
+        Person.drop_collection()
+
         for x in xrange(10):
             p = Person(name="Person %s" % x)
             p.save()
 
-        ids = [i.id for i in Person.objects]
-        self.assertEqual(ids, range(1, 21))
+        c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
+        self.assertEqual(c['next'], 10)
+
+        ids = [i.id for i in Person.objects]
+        self.assertEqual(ids, range(1, 11))
+
+        counters = [i.counter for i in Person.objects]
+        self.assertEqual(counters, range(1, 11))
+
+        c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
+        self.assertEqual(c['next'], 10)
+
+    def test_multiple_sequence_fields_on_docs(self):
+
+        class Animal(Document):
+            id = SequenceField(primary_key=True)
+
+        class Person(Document):
+            id = SequenceField(primary_key=True)
+
+        self.db['mongoengine.counters'].drop()
+        Animal.drop_collection()
+        Person.drop_collection()
+
+        for x in xrange(10):
+            a = Animal(name="Animal %s" % x)
+            a.save()
+            p = Person(name="Person %s" % x)
+            p.save()
+
+        c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
+        self.assertEqual(c['next'], 10)
+
+        c = self.db['mongoengine.counters'].find_one({'_id': 'animal.id'})
+        self.assertEqual(c['next'], 10)
+
+        ids = [i.id for i in Person.objects]
+        self.assertEqual(ids, range(1, 11))
+
+        id = [i.id for i in Animal.objects]
+        self.assertEqual(id, range(1, 11))
+
+        c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
+        self.assertEqual(c['next'], 10)
+
+        c = self.db['mongoengine.counters'].find_one({'_id': 'animal.id'})
+        self.assertEqual(c['next'], 10)
 
-        counter = self.db['mongoengine.counters'].find_one({'_id': 'person.id'})
-        self.assertEqual(counter['next'], 20)
 
 if __name__ == '__main__':
     unittest.main()