add setSlaveOk to Query

Allow callers to specify a value for slaveOk if they wish to override
the behavior they've configured their driver with for a particular
query.

src/main/scala/com/foursquare/rogue/MongoHelpers.scala

@@ -2,7 +2,7 @@
 
 package com.foursquare.rogue
 
-import com.mongodb.{BasicDBObjectBuilder, DBObject, DBCursor, WriteConcern}
+import com.mongodb.{BasicDBObjectBuilder, Bytes, DBObject, DBCursor, WriteConcern}
 import net.liftweb.mongodb._
 import net.liftweb.mongodb.record._
 import scala.collection.immutable.ListMap
@@ -272,9 +272,19 @@ object MongoHelpers {
         MongoDB.useCollection(queryClause.meta.mongoIdentifier, queryClause.meta.collectionName) {
           coll =>
           try {
-            val cursor = coll.find(cnd, sel).limit(queryClause.lim getOrElse 0)
-                .skip(queryClause.sk getOrElse 0)
-            ord.foreach(cursor sort _)
+            val cursor = coll.find(cnd, sel)
+            queryClause.lim.foreach(cursor.limit _)
+            queryClause.sk.foreach(cursor.skip _)
+            ord.foreach(cursor.sort _)
+            queryClause.slaveOk.foreach(so => {
+              if (so) {
+                // Use bitwise-or to add in slave-ok
+                cursor.setOptions(cursor.getOptions | Bytes.QUERYOPTION_SLAVEOK)
+              } else {
+                // Remove slave-ok from options
+                cursor.setOptions(cursor.getOptions & ~Bytes.QUERYOPTION_SLAVEOK)
+              }
+            })
             queryClause.maxScan.foreach(cursor addSpecial("$maxScan", _))
             queryClause.comment.foreach(cursor addSpecial("$comment", _))
             hnt.foreach(cursor hint _)

src/main/scala/com/foursquare/rogue/Query.scala

@@ -65,27 +65,27 @@ class IndexEnforcerBuilder[M <: MongoRecord[M]](meta: M with MongoMetaRecord[M]
   type MetaM = M with MongoMetaRecord[M] with IndexedRecord[M]
 
   def useIndex[F1 <: Field[_, M]](i: MongoIndex1[M, F1, _]): IndexEnforcer1[M, NoIndexInfo, F1, HasntUsedIndex] = {
-    new IndexEnforcer1[M, NoIndexInfo, F1, HasntUsedIndex](meta, new BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](meta, None, None, None, None, None, AndCondition(Nil, None), None, None))
+    new IndexEnforcer1[M, NoIndexInfo, F1, HasntUsedIndex](meta, new BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](meta, None, None, None, None, None, AndCondition(Nil, None), None, None, None))
   }
 
   def useIndex[F1 <: Field[_, M], F2 <: Field[_, M]](i: MongoIndex2[M, F1, _, F2, _]): IndexEnforcer2[M, NoIndexInfo, F1, F2, HasntUsedIndex] = {
-    new IndexEnforcer2[M, NoIndexInfo, F1, F2, HasntUsedIndex](meta, new BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](meta, None, None, None, None, None, AndCondition(Nil, None), None, None))
+    new IndexEnforcer2[M, NoIndexInfo, F1, F2, HasntUsedIndex](meta, new BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](meta, None, None, None, None, None, AndCondition(Nil, None), None, None, None))
   }
 
   def useIndex[F1 <: Field[_, M], F2 <: Field[_, M], F3 <: Field[_, M]](i: MongoIndex3[M, F1, _, F2, _, F3, _]): IndexEnforcer3[M, NoIndexInfo, F1, F2, F3, HasntUsedIndex] = {
-    new IndexEnforcer3[M, NoIndexInfo, F1, F2, F3, HasntUsedIndex](meta, new BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](meta, None, None, None, None, None, AndCondition(Nil, None), None, None))
+    new IndexEnforcer3[M, NoIndexInfo, F1, F2, F3, HasntUsedIndex](meta, new BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](meta, None, None, None, None, None, AndCondition(Nil, None), None, None, None))
   }
 
   def useIndex[F1 <: Field[_, M], F2 <: Field[_, M], F3 <: Field[_, M], F4 <: Field[_, M]](i: MongoIndex4[M, F1, _, F2, _, F3, _, F4, _]): IndexEnforcer4[M, NoIndexInfo, F1, F2, F3, F4, HasntUsedIndex] = {
-    new IndexEnforcer4[M, NoIndexInfo, F1, F2, F3, F4, HasntUsedIndex](meta, new BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](meta, None, None, None, None, None, AndCondition(Nil, None), None, None))
+    new IndexEnforcer4[M, NoIndexInfo, F1, F2, F3, F4, HasntUsedIndex](meta, new BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](meta, None, None, None, None, None, AndCondition(Nil, None), None, None, None))
   }
 
   def useIndex[F1 <: Field[_, M], F2 <: Field[_, M], F3 <: Field[_, M], F4 <: Field[_, M], F5 <: Field[_, M]](i: MongoIndex5[M, F1, _, F2, _, F3, _, F4, _, F5, _]): IndexEnforcer5[M, NoIndexInfo, F1, F2, F3, F4, F5, HasntUsedIndex] = {
-    new IndexEnforcer5[M, NoIndexInfo, F1, F2, F3, F4, F5, HasntUsedIndex](meta, new BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](meta, None, None, None, None, None, AndCondition(Nil, None), None, None))
+    new IndexEnforcer5[M, NoIndexInfo, F1, F2, F3, F4, F5, HasntUsedIndex](meta, new BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](meta, None, None, None, None, None, AndCondition(Nil, None), None, None, None))
   }
 
   def useIndex[F1 <: Field[_, M], F2 <: Field[_, M], F3 <: Field[_, M], F4 <: Field[_, M], F5 <: Field[_, M], F6 <: Field[_, M]](i: MongoIndex6[M, F1, _, F2, _, F3, _, F4, _, F5, _, F6, _]): IndexEnforcer6[M, NoIndexInfo, F1, F2, F3, F4, F5, F6, HasntUsedIndex] = {
-    new IndexEnforcer6[M, NoIndexInfo, F1, F2, F3, F4, F5, F6, HasntUsedIndex](meta, new BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](meta, None, None, None, None, None, AndCondition(Nil, None), None, None))
+    new IndexEnforcer6[M, NoIndexInfo, F1, F2, F3, F4, F5, F6, HasntUsedIndex](meta, new BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](meta, None, None, None, None, None, AndCondition(Nil, None), None, None, None))
   }
 }
 
@@ -499,6 +499,19 @@ trait AbstractQuery[M <: MongoRecord[M], R,
 
   def comment(c: String): AbstractQuery[M, R, Ord, Sel, Lim, Sk, Or]
 
+  /**
+   * Set a flag to indicate whether this query may hit secondaries. This only
+   * really makes sense if you're using replica sets. If this field is
+   * unspecified, rogue will leave the option untouched, so you'll use
+   * secondaries or not depending on how you configure the mongo java driver.
+   * Also, this only works if you're doing a query -- findAndModify, updates,
+   * and deletes always go to the primaries.
+   *
+   * For more info, see
+   * http://www.mongodb.org/display/DOCS/Querying#Querying-slaveOk%28QueryingSecondaries%29.
+   */
+  def setSlaveOk(b: Boolean): AbstractQuery[M, R, Ord, Sel, Lim, Sk, Or]
+
   /**
    * Adds a select clause to the query. The use of this method constrains the type
    * signature of the query to force the "Sel" field to be type "Selected".
@@ -644,7 +657,8 @@ case class BaseQuery[M <: MongoRecord[M], R,
     hint: Option[ListMap[String, Any]],
     condition: AndCondition,
     order: Option[MongoOrder],
-    select: Option[MongoSelect[R, M]]) extends AbstractQuery[M, R, Ord, Sel, Lim, Sk, Or] {
+    select: Option[MongoSelect[R, M]],
+    slaveOk: Option[Boolean]) extends AbstractQuery[M, R, Ord, Sel, Lim, Sk, Or] {
 
   // The meta field on the MongoMetaRecord (as an instance of MongoRecord)
   // points to the master MongoMetaRecord. This is here in case you have a
@@ -842,6 +856,8 @@ case class BaseQuery[M <: MongoRecord[M], R,
 
   override def comment(c: String): AbstractQuery[M, R, Ord, Sel, Lim, Sk, Or] = this.copy(comment = Some(c))
 
+  override def setSlaveOk(b: Boolean): AbstractQuery[M, R, Ord, Sel, Lim, Sk, Or] = this.copy(slaveOk = Some(b))
+
   override def hint(index: MongoIndex[M]) = this.copy(hint = Some(index.asListMap))
 
   override def select[F1](f: M => SelectField[F1, M])
@@ -1149,6 +1165,8 @@ class BaseEmptyQuery[M <: MongoRecord[M], R,
 
   override def comment(c: String) = this
 
+  override def setSlaveOk(b: Boolean) = this
+
   override def hint(index: MongoIndex[M]) = this
 
   override def select[F1](f: M => SelectField[F1, M])(implicit ev: Sel =:= Unselected) =

src/main/scala/com/foursquare/rogue/Rogue.scala

@@ -39,7 +39,7 @@ trait Rogue {
         val orCondition = QueryHelpers.orConditionFromQueries(q :: qs)
         BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasOrClause](
           q.meta, None, None, None, None, None,
-          AndCondition(Nil, Some(orCondition)), None, None)
+          AndCondition(Nil, Some(orCondition)), None, None, None)
       }
     }
   }
@@ -54,7 +54,7 @@ trait Rogue {
   implicit def metaRecordToQueryBuilder[M <: MongoRecord[M]]
       (rec: M with MongoMetaRecord[M]): BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause] =
     BaseQuery[M, M, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](
-      rec, None, None, None, None, None, AndCondition(Nil, None), None, None)
+      rec, None, None, None, None, None, AndCondition(Nil, None), None, None, None)
 
   implicit def metaRecordToModifyQuery[M <: MongoRecord[M]](rec: M with MongoMetaRecord[M]): AbstractModifyQuery[M] =
       BaseModifyQuery(metaRecordToQueryBuilder(rec), MongoModify(Nil))
@@ -75,7 +75,7 @@ trait Rogue {
       case q: BaseEmptyQuery[_, _, _, _, _, _, _] => new EmptyModifyQuery[M]
       case q: BaseQuery[_, _, _, _, _, _, _] =>
           BaseModifyQuery[M](q.asInstanceOf[BaseQuery[M, M, Unordered, Unselected, Unlimited,
-                                                      Unskipped, HasNoOrClause]], 
+                                                      Unskipped, HasNoOrClause]],
                              MongoModify(Nil))
     }
   }

src/test/scala/com/foursquare/rogue/EndToEndTest.scala

@@ -172,4 +172,17 @@ class EndToEndTest extends SpecsMatchers {
     // This assertion is what we want, and it fails.
     // subuseridsAndStatuses must_== List((Full(List(1234, 5678)), Full(List(ClaimStatus.pending, ClaimStatus.approved))))
   }
+
+  @Test
+  def testSlaveOk: Unit = {
+    // Note: this isn't a real test of slaveok because the test mongo setup
+    // doesn't have replicas. This basically just makes sure that slaveok
+    // doesn't break everything.
+    val v = baseTestVenue().save
+
+    // eqs
+    Venue.where(_._id eqs v.id).fetch().map(_.id)                         must_== List(v.id)
+    Venue.where(_._id eqs v.id).setSlaveOk(true).fetch().map(_.id)        must_== List(v.id)
+    Venue.where(_._id eqs v.id).setSlaveOk(false).fetch().map(_.id)       must_== List(v.id)
+  }
 }

src/test/scala/com/foursquare/rogue/QueryExecutorTest.scala

@@ -18,7 +18,7 @@ class QueryExecutorTest extends SpecsMatchers {
   @Test
   def testExeptionInRunCommandIsDecorated {
     val query = BaseQuery[Dummy, Dummy, Unordered, Unselected, Unlimited, Unskipped, HasNoOrClause](
-      Dummy, None, None, None, None, None, AndCondition(Nil, None), None, None)
+      Dummy, None, None, None, None, None, AndCondition(Nil, None), None, None, None)
     (QueryExecutor.runCommand("hello", query){
       throw new RuntimeException("bang")
       "hi"

src/test/scala/com/foursquare/rogue/QueryTest.scala

@@ -538,6 +538,16 @@ class QueryTest extends SpecsMatchers {
     maybeLimit(1, Some(5)) toString() must_== """db.venues.find({ "legid" : 1}).limit(5)"""
   }
 
+  @Test
+  def testSetSlaveOk: Unit = {
+    type Q = BaseQuery[Venue, Venue, _, _, _, _, _]
+
+    Venue.where(_.mayor eqs 2).asInstanceOf[Q].slaveOk must_== None
+    Venue.where(_.mayor eqs 2).setSlaveOk(true).asInstanceOf[Q].slaveOk must_== Some(true)
+    Venue.where(_.mayor eqs 2).setSlaveOk(false).asInstanceOf[Q].slaveOk must_== Some(false)
+    Venue.where(_.mayor eqs 2).setSlaveOk(true).setSlaveOk(false).asInstanceOf[Q].slaveOk must_== Some(false)
+  }
+
   @Test
   def thingsThatShouldntCompile {
     val compiler = new Compiler