allow use of $or queries in modify's

foursquare · Feb 1, 2012 · 3316b30 · 3316b30
1 parent 6b50300
commit 3316b30
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/src/main/scala/com/foursquare/rogue/Rogue.scala b/src/main/scala/com/foursquare/rogue/Rogue.scala
@@ -71,8 +71,8 @@ trait Rogue {
   /* A couple of implicit conversions that take a query builder, and convert it to a modify. This allows
    * users to write "RecordType.where(...).modify(...)".
    */
-  implicit def queryBuilderToModifyQuery[M <: MongoRecord[M]]
-      (query: AbstractQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause])
+  implicit def queryBuilderToModifyQuery[M <: MongoRecord[M], Or <: MaybeHasOrClause]
+      (query: AbstractQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, Or])
       : AbstractModifyQuery[M] = {
     query match {
       case q: BaseEmptyQuery[_, _, _, _, _, _, _] => new EmptyModifyQuery[M]
@@ -83,8 +83,8 @@ trait Rogue {
     }
   }
 
-  implicit def queryBuilderToFindAndModifyQuery[M <: MongoRecord[M], R, Ord <: MaybeOrdered, Sel <: MaybeSelected]
-      (query: AbstractQuery[M, R, Ord, Sel, Unlimited, Unskipped, HasNoOrClause])
+  implicit def queryBuilderToFindAndModifyQuery[M <: MongoRecord[M], R, Ord <: MaybeOrdered, Sel <: MaybeSelected, Or <: MaybeHasOrClause]
+      (query: AbstractQuery[M, R, Ord, Sel, Unlimited, Unskipped, Or])
       : AbstractFindAndModifyQuery[M, R] = {
     query match {
       case q: BaseEmptyQuery[_, _, _, _, _, _, _] => new EmptyFindAndModifyQuery[M, R]

diff --git a/src/test/scala/com/foursquare/rogue/QueryTest.scala b/src/test/scala/com/foursquare/rogue/QueryTest.scala
@@ -417,6 +417,12 @@ class QueryTest extends SpecsMatchers {
         _.where(_.mayor eqs 2))
       .toString() must_== """db.venues.find({ "$or" : [ { "legid" : 1 , "$or" : [ { "closed" : true} , { "closed" : { "$exists" : false}}]} , { "mayor" : 2}]})"""
 
+    // $or with modify
+    Venue.or(
+        _.where(_.legacyid eqs 1),
+        _.where(_.mayor eqs 2))
+      .modify(_.userid setTo 1).toString() must_== """db.venues.update({ "$or" : [ { "legid" : 1} , { "mayor" : 2}]}, { "$set" : { "userid" : 1}}, false, false)"""
+
     // OrQuery syntax
     val q1 = Venue.where(_.legacyid eqs 1)
     val q2 = Venue.where(_.legacyid eqs 2)