From 94d1e566c018bc26875b1cbc585ba2215d88f80d Mon Sep 17 00:00:00 2001 From: Ross Lawley Date: Tue, 23 Apr 2013 14:44:17 +0000 Subject: [PATCH] Added SequenceField.set_next_value(value) helper (#159) --- docs/changelog.rst | 1 + mongoengine/fields.py | 19 +++++++++++++++---- tests/fields/fields.py | 40 +++++++++++++++++++++++++--------------- 3 files changed, 41 insertions(+), 19 deletions(-) diff --git a/docs/changelog.rst b/docs/changelog.rst index fb9c35da..d22fc600 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -4,6 +4,7 @@ Changelog Changes in 0.8.X ================ +- Added SequenceField.set_next_value(value) helper (#159) - Updated .only() behaviour - now like exclude it is chainable (#202) - Added with_limit_and_skip support to count() (#235) - Removed __len__ from queryset (#247) diff --git a/mongoengine/fields.py b/mongoengine/fields.py index 690e7ace..4fc65c7e 100644 --- a/mongoengine/fields.py +++ b/mongoengine/fields.py @@ -1408,13 +1408,13 @@ class SequenceField(BaseField): COLLECTION_NAME = 'mongoengine.counters' VALUE_DECORATOR = int - def __init__(self, collection_name=None, db_alias=None, - sequence_name=None, value_decorator=None, *args, **kwargs): + def __init__(self, collection_name=None, db_alias=None, sequence_name=None, + value_decorator=None, *args, **kwargs): self.collection_name = collection_name or self.COLLECTION_NAME self.db_alias = db_alias or DEFAULT_CONNECTION_NAME self.sequence_name = sequence_name self.value_decorator = (callable(value_decorator) and - value_decorator or self.VALUE_DECORATOR) + value_decorator or self.VALUE_DECORATOR) return super(SequenceField, self).__init__(*args, **kwargs) def generate(self): @@ -1430,6 +1430,17 @@ class SequenceField(BaseField): upsert=True) return self.value_decorator(counter['next']) + def set_next_value(self, value): + """Helper method to set the next sequence value""" + sequence_name = self.get_sequence_name() + sequence_id = "%s.%s" % (sequence_name, self.name) + collection = get_db(alias=self.db_alias)[self.collection_name] + counter = collection.find_and_modify(query={"_id": sequence_id}, + update={"$set": {"next": value}}, + new=True, + upsert=True) + return self.value_decorator(counter['next']) + def get_sequence_name(self): if self.sequence_name: return self.sequence_name @@ -1438,7 +1449,7 @@ class SequenceField(BaseField): return owner._get_collection_name() else: return ''.join('_%s' % c if c.isupper() else c - for c in owner._class_name).strip('_').lower() + for c in owner._class_name).strip('_').lower() def __get__(self, instance, owner): value = super(SequenceField, self).__get__(instance, owner) diff --git a/tests/fields/fields.py b/tests/fields/fields.py index 9a7b82f7..ade44b8d 100644 --- a/tests/fields/fields.py +++ b/tests/fields/fields.py @@ -2164,18 +2164,21 @@ class FieldTest(unittest.TestCase): Person.drop_collection() for x in xrange(10): - p = Person(name="Person %s" % x) - p.save() + Person(name="Person %s" % x).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)) + self.assertEqual(ids, xrange(1, 11)) c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'}) self.assertEqual(c['next'], 10) + Person.id.set_next_value(1000) + c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'}) + self.assertEqual(c['next'], 1000) + def test_sequence_field_sequence_name(self): class Person(Document): id = SequenceField(primary_key=True, sequence_name='jelly') @@ -2185,8 +2188,7 @@ class FieldTest(unittest.TestCase): Person.drop_collection() for x in xrange(10): - p = Person(name="Person %s" % x) - p.save() + Person(name="Person %s" % x).save() c = self.db['mongoengine.counters'].find_one({'_id': 'jelly.id'}) self.assertEqual(c['next'], 10) @@ -2197,6 +2199,10 @@ class FieldTest(unittest.TestCase): c = self.db['mongoengine.counters'].find_one({'_id': 'jelly.id'}) self.assertEqual(c['next'], 10) + Person.id.set_next_value(1000) + c = self.db['mongoengine.counters'].find_one({'_id': 'jelly.id'}) + self.assertEqual(c['next'], 1000) + def test_multiple_sequence_fields(self): class Person(Document): id = SequenceField(primary_key=True) @@ -2207,21 +2213,28 @@ class FieldTest(unittest.TestCase): Person.drop_collection() for x in xrange(10): - p = Person(name="Person %s" % x) - p.save() + Person(name="Person %s" % x).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)) + self.assertEqual(ids, xrange(1, 11)) counters = [i.counter for i in Person.objects] - self.assertEqual(counters, range(1, 11)) + self.assertEqual(counters, xrange(1, 11)) c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'}) self.assertEqual(c['next'], 10) + Person.id.set_next_value(1000) + c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'}) + self.assertEqual(c['next'], 1000) + + Person.counter.set_next_value(999) + c = self.db['mongoengine.counters'].find_one({'_id': 'person.counter'}) + self.assertEqual(c['next'], 999) + def test_sequence_fields_reload(self): class Animal(Document): counter = SequenceField() @@ -2230,8 +2243,7 @@ class FieldTest(unittest.TestCase): self.db['mongoengine.counters'].drop() Animal.drop_collection() - a = Animal(name="Boi") - a.save() + a = Animal(name="Boi").save() self.assertEqual(a.counter, 1) a.reload() @@ -2261,10 +2273,8 @@ class FieldTest(unittest.TestCase): Person.drop_collection() for x in xrange(10): - a = Animal(name="Animal %s" % x) - a.save() - p = Person(name="Person %s" % x) - p.save() + Animal(name="Animal %s" % x).save() + Person(name="Person %s" % x).save() c = self.db['mongoengine.counters'].find_one({'_id': 'person.id'}) self.assertEqual(c['next'], 10)