pub/mongoengine
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Added SequenceField.set_next_value(value) helper (#159)

Browse Source
This commit is contained in:
Ross Lawley 2013-04-23 14:44:17 +00:00
parent a692316293
commit 94d1e566c0
3 changed files with 41 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
docs/changelog.rst
View File

@ -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)

15
mongoengine/fields.py
View File

@ -1408,8 +1408,8 @@ 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
@ -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

40
tests/fields/fields.py
View File

@ -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)