From 80c85733bf5dfd82576b8b1bab3152f3014efe0b Mon Sep 17 00:00:00 2001 From: Neil Sanchala Date: Tue, 10 Jan 2012 17:38:16 -0500 Subject: [PATCH] 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. --- .../com/foursquare/rogue/MongoHelpers.scala | 18 ++++++++--- .../scala/com/foursquare/rogue/Query.scala | 32 +++++++++++++++---- .../scala/com/foursquare/rogue/Rogue.scala | 6 ++-- .../com/foursquare/rogue/EndToEndTest.scala | 13 ++++++++ .../foursquare/rogue/QueryExecutorTest.scala | 2 +- .../com/foursquare/rogue/QueryTest.scala | 10 ++++++ 6 files changed, 66 insertions(+), 15 deletions(-) diff --git a/src/main/scala/com/foursquare/rogue/MongoHelpers.scala b/src/main/scala/com/foursquare/rogue/MongoHelpers.scala index 65ad65c..6388d88 100644 --- a/src/main/scala/com/foursquare/rogue/MongoHelpers.scala +++ b/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 _) diff --git a/src/main/scala/com/foursquare/rogue/Query.scala b/src/main/scala/com/foursquare/rogue/Query.scala index 9225f78..b68665c 100644 --- a/src/main/scala/com/foursquare/rogue/Query.scala +++ b/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) = diff --git a/src/main/scala/com/foursquare/rogue/Rogue.scala b/src/main/scala/com/foursquare/rogue/Rogue.scala index a98016c..d10da37 100644 --- a/src/main/scala/com/foursquare/rogue/Rogue.scala +++ b/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)) } } diff --git a/src/test/scala/com/foursquare/rogue/EndToEndTest.scala b/src/test/scala/com/foursquare/rogue/EndToEndTest.scala index 0e4a93a..a22524e 100644 --- a/src/test/scala/com/foursquare/rogue/EndToEndTest.scala +++ b/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) + } } diff --git a/src/test/scala/com/foursquare/rogue/QueryExecutorTest.scala b/src/test/scala/com/foursquare/rogue/QueryExecutorTest.scala index bc55ef5..abc9282 100644 --- a/src/test/scala/com/foursquare/rogue/QueryExecutorTest.scala +++ b/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" diff --git a/src/test/scala/com/foursquare/rogue/QueryTest.scala b/src/test/scala/com/foursquare/rogue/QueryTest.scala index 46f8172..1e7d66a 100644 --- a/src/test/scala/com/foursquare/rogue/QueryTest.scala +++ b/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