Adding .upsert and .multi to update expressions, better error message…

…s for .one and tests for update expressions which actually write to the database

mongoalchemy/query.py

@@ -112,11 +112,13 @@ def one(self):
         '''Execute the query and return one result.  If more than one result 
             is returned, raises a ``BadResultException``
         '''
-        try:
-            [the_one] = self
-        except ValueError:
-            raise BadResultException('Too many results for .one()')
-        return the_one
+        count = -1
+        for count, result in enumerate(self):
+            if count > 0:
+                raise BadResultException('Too many results for .one()')
+        if count == -1:
+            raise BadResultException('Too few results for .one()')
+        return result
 
     def first(self):
         '''Execute the query and return the first result.  Unlike ``one``, if

mongoalchemy/update_expression.py

@@ -31,6 +31,19 @@ class UpdateExpression(object):
     def __init__(self, query):
         self.query = query
         self.update_data = {}
+        self.__upsert = False
+        self.__multi = False
+
+    def upsert(self):
+        ''' If a document matching the query doesn't exist, create one '''
+        self.__upsert = True
+        return self
+
+    def multi(self):
+        ''' Update multiple documents.  The Mongo default is to only update 
+            the first matching document '''
+        self.__multi = True
+        return self
 
     def set(self, qfield, value):
         ''' Atomically set ``qfield`` to ``value``'''
@@ -127,7 +140,7 @@ def execute(self):
         collection = self.query.db[self.query.type.get_collection_name()]
         for index in self.query.type.get_indexes():
             index.ensure(collection)
-        collection.update(self.query.query, self.update_data)
+        collection.update(self.query.query, self.update_data, upsert=self.__upsert, multi=self.__multi)
 
 class UpdateException(Exception):
     ''' Base class for exceptions related to updates '''

test/test_update_expressions.py

@@ -33,10 +33,22 @@ def update_test_setup():
 
 def set_test():
     q = update_test_setup()
-    assert q.set(T.f.i, 5).update_data == {
-        '$set' : { 'i' : 5 }
+    assert q.set(T.f.i, 5).set(T.f.j, 7).update_data == {
+        '$set' : { 'i' : 5, 'j' : 7 }
     }
 
+def set_db_test():
+    q = update_test_setup()
+
+    q.set(T.f.i, 5).set(T.f.j, 6).upsert().execute()
+    t = q.one()
+    assert t.i == 5 and t.j == 6
+
+    q.filter(T.f.i == 5).set(T.f.j, 7).execute()
+    t = q.one()
+    assert t.i == 5 and t.j == 7
+
+
 # UNSET
 
 def unset_test():
@@ -45,22 +57,44 @@ def unset_test():
         '$unset' : { 'i' : True }
     }
 
+def unset_db_test():
+    q = update_test_setup()
+
+    q.set(T.f.i, 5).set(T.f.j, 6).upsert().execute()
+    t = q.one()
+    assert t.i == 5 and t.j == 6
+
+    q.filter(T.f.i == 5).unset(T.f.j).execute()
+    t = q.one()
+    assert t.i == 5 and not hasattr(t, 'j')
+
+
 # INC
 
 def inc_test():
     q = update_test_setup()
     assert q.inc(T.f.i, 4).update_data == {
         '$inc' : { 'i' : 4 }
     }
-    assert q.inc(T.f.i).update_data == {
-        '$inc' : { 'i' : 1 }
+    assert q.inc(T.f.i).inc(T.f.j).update_data == {
+        '$inc' : { 'i' : 1, 'j' : 1 }
     }
 
 @raises(InvalidModifierException)
 def inc_invalid_test():
     q = update_test_setup()
     q.inc(T.f.s, 1)
 
+def inc_db_test():
+    q = update_test_setup()
+
+    q.inc(T.f.i, 5).inc(T.f.j, 6).upsert().execute()
+    t = q.one()
+    assert t.i == 5 and t.j == 6
+    q.inc(T.f.j, 6).inc(T.f.i, 5).execute()
+    t = q.one()
+    assert t.i == 10 and t.j == 12
+
 # APPEND
 
 def append_test():
@@ -74,6 +108,18 @@ def append_invalid_test():
     q = update_test_setup()
     q.append(T.f.s, 1)
 
+def append_db_test():
+    q = update_test_setup()
+
+    q.set(T.f.i, 5).append(T.f.l, 6).upsert().execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [6]
+
+    q.append(T.f.l, 5).execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [6, 5]
+
+
 # EXTEND
 
 def extend_test():
@@ -82,6 +128,16 @@ def extend_test():
         '$pushAll' : { 'l' : (1, 2, 3) }
     }
 
+def extend_db_test():
+    q = update_test_setup()
+
+    q.set(T.f.i, 5).extend(T.f.l, 6, 5).upsert().execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [6, 5]
+
+    q.extend(T.f.l, 4, 3).execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [6, 5, 4, 3]
 
 @raises(InvalidModifierException)
 def extend_invalid_test():
@@ -101,6 +157,18 @@ def remove_invalid_test():
     q = update_test_setup()
     q.remove(T.f.s, 1)
 
+def remove_db_test():
+    q = update_test_setup()
+
+    q.set(T.f.i, 5).extend(T.f.l, 6, 5, 4, 3).upsert().execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [6, 5, 4, 3]
+
+    q.remove(T.f.l, 4).execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [6, 5, 3]
+
+
 # REMOVE ALL
 
 def remove_all_test():
@@ -114,6 +182,18 @@ def remove_all_invalid_test():
     q = update_test_setup()
     q.remove_all(T.f.s, 1, 2, 3)
 
+def remove_all_db_test():
+    q = update_test_setup()
+
+    q.set(T.f.i, 5).extend(T.f.l, 6, 5, 4, 3).upsert().execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [6, 5, 4, 3]
+
+    q.remove_all(T.f.l, 6, 5, 4).execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [3]
+
+
 # ADD TO SET
 
 def add_to_set_test():
@@ -122,6 +202,17 @@ def add_to_set_test():
         '$addToSet' : { 'l' : 1 }
     }
 
+def add_to_set_db_test():
+    q = update_test_setup()
+
+    q.set(T.f.i, 5).extend(T.f.l, 6, 5, 4, 3).upsert().execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [6, 5, 4, 3]
+
+    q.add_to_set(T.f.l, 6).add_to_set(T.f.l, 2).execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [6, 5, 4, 3, 2]
+
 @raises(InvalidModifierException)
 def add_to_set_invalid_test():
     q = update_test_setup()
@@ -140,6 +231,18 @@ def pop_first_invalid_test():
     q = update_test_setup()
     q.pop_last(T.f.s)
 
+def pop_first_db_test():
+    q = update_test_setup()
+
+    q.set(T.f.i, 5).extend(T.f.l, 6, 5, 4, 3).upsert().execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [6, 5, 4, 3]
+
+    q.pop_first(T.f.l).execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [5, 4, 3]
+
+
 # POP LAST
 
 def pop_last_test():
@@ -154,3 +257,13 @@ def pop_last_invalid_test():
     q = update_test_setup()
     q.pop_first(T.f.s)
 
+def pop_last_db_test():
+    q = update_test_setup()
+
+    q.set(T.f.i, 5).extend(T.f.l, 6, 5, 4, 3).upsert().execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [6, 5, 4, 3]
+
+    q.pop_last(T.f.l).execute()
+    t = q.one()
+    assert t.i == 5 and t.l == [6, 5, 4]