initial commit

tjulien · Jan 15, 2011 · c56a4b2 · c56a4b2
commit c56a4b2
Show file tree

Hide file tree

Showing 10 changed files with 994 additions and 0 deletions.
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,9 @@
+.DS_Store
+.idea
+Main/
+lib_managed
+project/boot/*
+project/build/target/*
+sbtlib
+target
+*~
diff --git a/project/build.properties b/project/build.properties
@@ -0,0 +1,8 @@
+#Project properties
+#Sat Jan 15 02:15:19 UTC 2011
+project.organization=com.foursquare
+project.name=rogue
+sbt.version=0.7.4
+project.version=1.0
+build.scala.versions=2.8.0
+project.initialize=false
diff --git a/project/build/RogueProject.scala b/project/build/RogueProject.scala
@@ -0,0 +1,32 @@
+import sbt._
+import sbt.Process._
+import java.lang.System
+import java.util.concurrent.{Callable, Executors}
+
+class RogueProject(info: ProjectInfo) extends DefaultProject(info) {
+  val liftVer = "2.2"
+  val scalaVer = buildScalaVersion
+
+  // Lift Libraries
+  val liftCommon      = "net.liftweb" %% "lift-common"         % liftVer % "compile" withSources()
+  val liftJson        = "net.liftweb" %% "lift-json"           % liftVer % "compile" withSources()
+  val liftJsonExt     = "net.liftweb" %% "lift-json-ext"       % liftVer % "compile" withSources()
+  val liftRecord      = "net.liftweb" %% "lift-record"         % liftVer % "compile" withSources()
+  val liftUtil        = "net.liftweb" %% "lift-util"           % liftVer % "compile" withSources()
+  val liftMongo       = "net.liftweb" %% "lift-mongodb"        % liftVer % "compile" withSources()
+  val liftMongoRecord = "net.liftweb" %% "lift-mongodb-record" % liftVer % "compile" withSources()
+
+  // Scala Libraries
+  val scalaCompiler     = "org.scala-lang"           % "scala-compiler"     % scalaVer % "test"
+  val specs             = "org.scala-tools.testing" %% "specs"              % "1.6.5"  % "test" withSources()
+
+  // Java Libraries
+  val mongo        = "org.mongodb"         % "mongo-java-driver" % "2.4"
+  val jodaTime     = "joda-time"           % "joda-time"         % "1.6" withSources()
+  val commons      = "org.apache.commons"  % "commons-math"      % "2.1"
+  val junit        = "junit"               % "junit"             % "4.8.2" % "test" withSources()
+
+  val bryanjswift = "Bryan J Swift Repository" at "http://repos.bryanjswift.com/maven2/" 
+  val junitInterface = "com.novocode" % "junit-interface" % "0.5" % "test->default"
+  override def testFrameworks = super.testFrameworks ++ List(new TestFramework("com.novocode.junit.JUnitFrameworkNoMarker")) 
+}
diff --git a/src/main/scala/com/foursquare/rogue/MongoHelpers.scala b/src/main/scala/com/foursquare/rogue/MongoHelpers.scala
@@ -0,0 +1,142 @@
+// Copyright 2011 Foursquare Labs Inc. All Rights Reserved.
+
+package com.foursquare.rogue
+
+import com.mongodb.{BasicDBObjectBuilder, DBObject}
+import net.liftweb.mongodb._
+import net.liftweb.mongodb.record._
+
+object MongoHelpers {
+  sealed abstract class MongoCondition
+  case class AndCondition(clauses: List[QueryClause[_]]) extends MongoCondition
+  case class OrCondition(conditions: List[MongoCondition]) extends MongoCondition
+
+  sealed case class MongoOrder(terms: List[(String, Boolean)])
+  sealed case class MongoModify(clauses: List[ModifyClause[_]])
+  sealed case class MongoSelect[R](fields: List[SelectField[_, _]], transformer: List[_] => R)
+
+  object MongoBuilder {
+    def buildCondition(q: MongoCondition): DBObject = q match {
+      case AndCondition(clauses) =>
+        val builder = BasicDBObjectBuilder.start
+        (clauses.groupBy(_.fieldName)
+                .toList
+                .sortBy { case (fieldName, _) => -clauses.findIndexOf(_.fieldName == fieldName) }
+                .foreach { case (name, cs) =>
+          // Equality clauses look like { a : 3 }
+          // but all other clauses look like { a : { $op : 3 }}
+          // and can be chained like { a : { $gt : 2, $lt: 6 }}.
+          // So if there is any equality clause, apply it (only) to the builder;
+          // otherwise, chain the clauses.
+          cs.filter(_.isInstanceOf[EqClause[_]]).headOption.map(_.extend(builder)).getOrElse {
+            builder.push(name)
+            cs.foreach(_.extend(builder))
+            builder.pop
+          }
+        })
+        builder.get
+
+      case OrCondition(conditions) => 
+        // Room for optimization here by manipulating the AST, e.g.,
+        // { $or : [ { a : 1 }, { a : 2 } ] }  ==>  { a : { $in : [ 1, 2 ] }}
+        BasicDBObjectBuilder.start("$or", QueryHelpers.list(conditions.map(buildCondition))).get
+    }
+
+    def buildOrder(o: MongoOrder): DBObject = {
+      val builder = BasicDBObjectBuilder.start
+      o.terms.reverse.foreach { case (field, ascending) => builder.add(field, if (ascending) 1 else -1) }
+      builder.get
+    }
+
+    def buildModify(m: MongoModify): DBObject = {
+      val builder = BasicDBObjectBuilder.start
+      m.clauses.groupBy(_.operator).foreach{ case (op, cs) => {
+        builder.push(op.toString)
+        cs.foreach(_.extend(builder))
+        builder.pop
+      }}
+      builder.get
+    }
+
+    def buildSelect[R](s: MongoSelect[R]): DBObject = {
+      val builder = BasicDBObjectBuilder.start
+      s.fields.foreach(f => builder.add(f.field.name, 1))
+      builder.get
+    }
+
+    def buildString[R](query: BaseQuery[_, R, _, _, _, _],
+                       modify: Option[MongoModify]): String = {
+      val sb = new StringBuilder
+      sb.append(buildCondition(query.condition).toString)
+      query.order.foreach(o => sb.append(" order by " + buildOrder(o).toString))
+      query.select.foreach(s => sb.append(" select " + buildSelect(s).toString))
+      query.lim.foreach(l => sb.append(" limit " + l))
+      query.sk.foreach(s => sb.append(" skip " + s))
+      modify.foreach(m => sb.append(" modify with " + buildModify(m)))
+      sb.toString
+    }
+  }
+
+  object QueryExecutor {
+
+    import QueryHelpers._
+    import MongoHelpers.MongoBuilder._
+
+    def condition[M <: MongoRecord[M], T](operation: String,
+                                            query: BaseQuery[M, _, _, _, _, _])
+                                           (f: DBObject => T): T = {
+      val start = System.currentTimeMillis
+      val collection = query.meta.collectionName
+      val cnd = buildCondition(query.condition)
+      try {
+        f(cnd)
+      } finally {
+        logger.log("Mongo %s.%s (%s)" format (collection, operation, cnd), System.currentTimeMillis - start)
+      }
+    }
+
+    def modify[M <: MongoRecord[M], T](operation: String,
+                                         mod: ModifyQuery[M])
+                                        (f: (DBObject, DBObject) => T): T = {
+      val start = System.currentTimeMillis
+      val collection = mod.query.meta.collectionName
+      val q = buildCondition(mod.query.condition)
+      val m = buildModify(mod.modify)
+      try {
+        f(q, m)
+      } finally {
+        logger.log("Mongo %s.%s (%s, %s)" format (collection, operation, q, m), System.currentTimeMillis - start)
+      }
+    }
+
+    def query[M <: MongoRecord[M]](operation: String,
+                                     query: BaseQuery[M, _, _, _, _, _])
+                                    (f: DBObject => Unit): Unit = {
+      val start = System.currentTimeMillis
+      MongoDB.useCollection(query.meta.mongoIdentifier, query.meta.collectionName) { coll =>
+        val collection = coll.getName
+        val cnd = buildCondition(query.condition)
+        val ord = query.order.map(buildOrder)
+        val sel = query.select.map(buildSelect)
+        lazy val empty = BasicDBObjectBuilder.start.get
+        try {
+          val cursor = coll.find(cnd, sel getOrElse empty).limit(query.lim getOrElse 0).skip(query.sk getOrElse 0)
+          ord.foreach(cursor sort _)
+          while (cursor.hasNext)
+            f(cursor.next)
+        } finally {
+          logger.log( {
+            val str = new StringBuilder("Mongo " + collection +"." + operation)
+            str.append("("+cnd)
+            sel.foreach(s => str.append(", "+s))
+            str.append(")")
+            ord.foreach(o => str.append(".sort("+o+")"))
+            query.sk.foreach(sk => str.append(".skip("+sk+")"))
+            query.lim.foreach(l => str.append(".limit("+l+")"))
+            str.toString
+          }, System.currentTimeMillis - start)
+        }
+      }
+    }
+  }
+}
diff --git a/src/main/scala/com/foursquare/rogue/Query.scala b/src/main/scala/com/foursquare/rogue/Query.scala
@@ -0,0 +1,193 @@
+// Copyright 2011 Foursquare Labs Inc. All Rights Reserved.
+
+package com.foursquare.rogue
+
+import com.foursquare.rogue.MongoHelpers._
+import com.mongodb.DBObject
+import net.liftweb.mongodb.record._
+
+/////////////////////////////////////////////////////////////////////////////
+// Phantom types
+/////////////////////////////////////////////////////////////////////////////
+
+abstract sealed class Ordered
+abstract sealed class Unordered
+abstract sealed class Selected
+abstract sealed class Unselected
+abstract sealed class Limited
+abstract sealed class Unlimited
+abstract sealed class Skipped
+abstract sealed class Unskipped
+
+/////////////////////////////////////////////////////////////////////////////
+// Builders
+/////////////////////////////////////////////////////////////////////////////
+
+class BaseQuery[M <: MongoRecord[M], R, Ord, Sel, Lim, Sk](
+    val meta: M with MongoMetaRecord[M],
+    val lim: Option[Int],
+    val sk: Option[Int],
+    val condition: AndCondition,
+    val order: Option[MongoOrder],
+    val select: Option[MongoSelect[R]]) {
+
+  // The meta field on the MongoMetaRecord (as an instance of MongoRecord)
+  // points to the master MongoMetaRecord
+  lazy val master = meta.meta
+
+  def where[F](clause: M => QueryClause[F]) = {
+    clause(meta) match {
+      case cl: EmptyQueryClause[_] => new BaseEmptyQuery[M, R, Ord, Sel, Lim, Sk]
+      case cl => new BaseQuery[M, R, Ord, Sel, Lim, Sk](meta, lim, sk, AndCondition(cl :: condition.clauses), order, select)
+    }
+  }
+  def and[F](clause: M => QueryClause[F]) = where(clause)
+  def orderAsc[V](field: M => QueryField[V, M])(implicit ev: Ord =:= Unordered) =
+    new BaseQuery[M, R, Ordered, Sel, Lim, Sk](meta, lim, sk, condition, Some(MongoOrder(List((field(meta).field.name, true)))), select)
+  def orderDesc[V](field: M => QueryField[V, M])(implicit ev: Ord =:= Unordered) =
+    new BaseQuery[M, R, Ordered, Sel, Lim, Sk](meta, lim, sk, condition, Some(MongoOrder(List((field(meta).field.name, false)))), select)
+  def andAsc[V](field: M => QueryField[V, M])(implicit ev: Ord =:= Ordered) =
+    new BaseQuery[M, R, Ordered, Sel, Lim, Sk](meta, lim, sk, condition, Some(MongoOrder((field(meta).field.name, true) :: order.get.terms)), select)
+  def andDesc[V](field: M => QueryField[V, M])(implicit ev: Ord =:= Ordered) =
+    new BaseQuery[M, R, Ordered, Sel, Lim, Sk](meta, lim, sk, condition, Some(MongoOrder((field(meta).field.name, false) :: order.get.terms)), select)
+
+  def limit(n: Int)(implicit ev: Lim =:= Unlimited) =
+    new BaseQuery[M, R, Ord, Sel, Limited, Sk](meta, Some(n), sk, condition, order, select)
+  def limitOpt(n: Option[Int])(implicit ev: Lim =:= Unlimited) =
+    new BaseQuery[M, R, Ord, Sel, Limited, Sk](meta, n, sk, condition, order, select)
+  def skip(n: Int)(implicit ev: Sk =:= Unskipped) =
+    new BaseQuery[M, R, Ord, Sel, Lim, Skipped](meta, lim, Some(n), condition, order, select)
+
+  private def extract(f: R => Unit): DBObject => Unit = select match {
+    case Some(MongoSelect(fields, transformer)) => (dbo) => {
+      val inst = meta.createRecord
+      f(transformer(fields.map(f => f(inst.fieldByName(f.field.name).open_!.setFromAny(dbo.get(f.field.name))))))
+    }
+    case None => 
+      dbo => f(meta.fromDBObject(dbo).asInstanceOf[R])
+  }
+
+  private def collect[T](f: (T => Unit) => Unit): List[T] = {
+    val buf = new scala.collection.mutable.ListBuffer[T]
+    f{ t: T => buf += t }
+    buf.toList
+  }
+
+  def count()(implicit ev1: Lim =:= Unlimited, ev2: Sk =:= Unskipped): Long =
+    QueryExecutor.condition("count", this)(meta.count(_))
+  def countDistinct[V](field: M => QueryField[V, M])(implicit ev1: Lim =:= Unlimited, ev2: Sk =:= Unskipped): Long =
+    QueryExecutor.condition("countDistinct", this)(meta.countDistinct(field(meta).field.name, _))
+  def foreach(f: R => Unit): Unit =
+    QueryExecutor.query("find", this)(extract(f))
+  def fetch(): List[R] =
+    collect[R](f => QueryExecutor.query("find", this)(extract(f)))
+  def fetch(limit: Int)(implicit ev: Lim =:= Unlimited): List[R] =
+    this.limit(limit).fetch()
+  def get()(implicit ev: Lim =:= Unlimited): Option[R] =
+    fetch(1).headOption
+  def paginate(countPerPage: Int)(implicit ev1: Lim =:= Unlimited, ev2: Sk =:= Unskipped) = {
+    val q = new BaseQuery[M, R, Ord, Sel, Unlimited, Unskipped](meta, lim, sk, condition, order, select)
+    new BasePaginatedQuery(q, countPerPage)
+  }
+  def modify[F](clause: M => ModifyClause[F])(implicit ev1: Sel =:= Unselected, ev2: Lim =:= Unlimited, ev3: Sk =:= Unskipped) =
+    new ModifyQuery(this, MongoModify(List(clause(meta))))
+
+  // Always do modifications against master (not meta, which could point to slave)
+  def bulkDelete_!!()(implicit ev1: Sel =:= Unselected, ev2: Lim =:= Unlimited, ev3: Sk =:= Unskipped): Unit =
+    QueryExecutor.condition("bulkDelete", this)(master.bulkDelete_!!(_))
+  override def toString: String =
+    MongoBuilder.buildString(this, None)
+
+  def select[F1](f: M => SelectField[F1, M])(implicit ev: Sel =:= Unselected): BaseQuery[M, F1, Ord, Selected, Lim, Sk] = {
+    val fields = List(f(meta))
+    val transformer = (xs: List[_]) => xs.head.asInstanceOf[F1]
+    new BaseQuery(meta, lim, sk, condition, order, Some(MongoSelect(fields, transformer)))
+  }
+
+  def select[F1, F2](f1: M => SelectField[F1, M], f2: M => SelectField[F2, M])(implicit ev: Sel =:= Unselected): BaseQuery[M, (F1, F2), Ord, Selected, Lim, Sk] = {
+    val fields = List(f1(meta), f2(meta))
+    val transformer = (xs: List[_]) => (xs(0).asInstanceOf[F1], xs(1).asInstanceOf[F2])
+    new BaseQuery(meta, lim, sk, condition, order, Some(MongoSelect(fields, transformer)))
+  }
+
+  def select[F1, F2, F3](f1: M => SelectField[F1, M], f2: M => SelectField[F2, M], f3: M => SelectField[F3, M])(implicit ev: Sel =:= Unselected): BaseQuery[M, (F1, F2, F3), Ord, Selected, Lim, Sk] = {
+    val fields = List(f1(meta), f2(meta), f3(meta))
+    val transformer = (xs: List[_]) => (xs(0).asInstanceOf[F1], xs(1).asInstanceOf[F2], xs(2).asInstanceOf[F3])
+    new BaseQuery(meta, lim, sk, condition, order, Some(MongoSelect(fields, transformer)))
+  }
+
+  def select[F1, F2, F3, F4](f1: M => SelectField[F1, M], f2: M => SelectField[F2, M], f3: M => SelectField[F3, M], f4: M => SelectField[F4, M])(implicit ev: Sel =:= Unselected): BaseQuery[M, (F1, F2, F3, F4), Ord, Selected, Lim, Sk] = {
+    val fields = List(f1(meta), f2(meta), f3(meta), f4(meta))
+    val transformer = (xs: List[_]) => (xs(0).asInstanceOf[F1], xs(1).asInstanceOf[F2], xs(2).asInstanceOf[F3], xs(3).asInstanceOf[F4])
+    new BaseQuery(meta, lim, sk, condition, order, Some(MongoSelect(fields, transformer)))
+  }
+
+  def fetchBatch[T](batchSize: Int)(f: List[R] => List[T])(implicit ev1: Lim =:= Unlimited, ev2: Sk =:= Unskipped): List[T] = {
+    val paginatedQuery = paginate(batchSize)
+    (1 to paginatedQuery.numPages).toList.flatMap(page => f(paginatedQuery.setPage(page).fetch))
+  }
+}
+
+class BaseEmptyQuery[M <: MongoRecord[M], R, Ord, Sel, Lim, Sk] extends BaseQuery[M, R, Ord, Sel, Lim, Sk](null.asInstanceOf[M with MongoMetaRecord[M]], None, None, null.asInstanceOf[AndCondition], None, None) {
+  override def where[F](clause: M => QueryClause[F]) = this
+  override def and[F](clause: M => QueryClause[F]) = this
+  override def orderAsc[V](field: M => QueryField[V, M])(implicit ev: Ord =:= Unordered) = new BaseEmptyQuery[M, R, Ordered, Sel, Lim, Sk]
+  override def orderDesc[V](field: M => QueryField[V, M])(implicit ev: Ord =:= Unordered) = new BaseEmptyQuery[M, R, Ordered, Sel, Lim, Sk]
+  override def andAsc[V](field: M => QueryField[V, M])(implicit ev: Ord =:= Ordered) = new BaseEmptyQuery[M, R, Ordered, Sel, Lim, Sk]
+  override def andDesc[V](field: M => QueryField[V, M])(implicit ev: Ord =:= Ordered) = new BaseEmptyQuery[M, R, Ordered, Sel, Lim, Sk]
+  override def limit(n: Int)(implicit ev: Lim =:= Unlimited) = new BaseEmptyQuery[M, R, Ord, Sel, Limited, Sk]
+  override def limitOpt(n: Option[Int])(implicit ev: Lim =:= Unlimited) = new BaseEmptyQuery[M, R, Ord, Sel, Limited, Sk]
+  override def skip(n: Int)(implicit ev: Sk =:= Unskipped) = new BaseEmptyQuery[M, R, Ord, Sel, Lim, Skipped]
+
+  override def count()(implicit ev1: Lim =:= Unlimited, ev2: Sk =:= Unskipped): Long = 0
+  override def countDistinct[V](field: M => QueryField[V, M])(implicit ev1: Lim =:= Unlimited, ev2: Sk =:= Unskipped): Long = 0
+  override def foreach(f: R => Unit): Unit = ()
+  override def fetch(): List[R] = Nil
+  override def fetch(limit: Int)(implicit ev: Lim =:= Unlimited): List[R] = Nil
+  override def get()(implicit ev: Lim =:= Unlimited): Option[R] = None
+  override def paginate(countPerPage: Int)(implicit ev1: Lim =:= Unlimited, ev2: Sk =:= Unskipped) = {
+    val emptyQuery = new BaseEmptyQuery[M, R, Ord, Sel, Unlimited, Unskipped]
+    new BasePaginatedQuery(emptyQuery, countPerPage)
+  }
+
+  override def modify[F](clause: M => ModifyClause[F])(implicit ev1: Sel =:= Unselected, ev2: Lim =:= Unlimited, ev3: Sk =:= Unskipped) = new NoopModifyQuery
+  override def bulkDelete_!!()(implicit ev1: Sel =:= Unselected, ev2: Lim =:= Unlimited, ev3: Sk =:= Unskipped): Unit = ()
+  override def toString = "empty query"
+
+  override def select[F1](f: M => SelectField[F1, M])(implicit ev: Sel =:= Unselected) = new BaseEmptyQuery[M, F1, Ord, Selected, Lim, Sk]
+  override def select[F1, F2](f1: M => SelectField[F1, M], f2: M => SelectField[F2, M])(implicit ev: Sel =:= Unselected) = new BaseEmptyQuery[M, (F1, F2), Ord, Selected, Lim, Sk]
+  override def select[F1, F2, F3](f1: M => SelectField[F1, M], f2: M => SelectField[F2, M], f3: M => SelectField[F3, M])(implicit ev: Sel =:= Unselected) = new BaseEmptyQuery[M, (F1, F2, F3), Ord, Selected, Lim, Sk]
+  override def select[F1, F2, F3, F4](f1: M => SelectField[F1, M], f2: M => SelectField[F2, M], f3: M => SelectField[F3, M], f4: M => SelectField[F4, M])(implicit ev: Sel =:= Unselected) = new BaseEmptyQuery[M, (F1, F2, F3, F4), Ord, Selected, Lim, Sk]
+
+  override def fetchBatch[T](batchSize: Int)(f: List[R] => List[T])(implicit ev1: Lim =:= Unlimited, ev2: Sk =:= Unskipped): List[T] = Nil
+}
+
+class ModifyQuery[M <: MongoRecord[M]](val query: BaseQuery[M, _, _, _, _, _],
+                                       val modify: MongoModify) {
+  def and[F](clause: M => ModifyClause[F]) = new ModifyQuery(query, MongoModify(clause(query.meta) :: modify.clauses))
+
+  // Always do modifications against master (not query.meta, which could point to slave)
+  def updateMulti(): Unit = QueryExecutor.modify("updateMulti", this)(query.master.updateMulti(_, _))
+  def updateOne(): Unit = QueryExecutor.modify("updateOne", this)(query.master.update(_, _))
+  def upsertOne(): Unit = QueryExecutor.modify("upsertOne", this)(query.master.upsert(_, _))
+
+  override def toString =
+    MongoBuilder.buildString(query, Some(modify))
+}
+
+class NoopModifyQuery[M <: MongoRecord[M]] extends ModifyQuery[M](
+  null.asInstanceOf[BaseQuery[M, _, _, _, _, _]],
+  null.asInstanceOf[MongoModify]) {
+  override def and[F](clause: M => ModifyClause[F]) = this
+  override def updateMulti(): Unit = ()
+  override def updateOne(): Unit = ()
+  override def upsertOne(): Unit = ()
+}
+
+class BasePaginatedQuery[M <: MongoRecord[M], R](q: BaseQuery[M, R, _, _, Unlimited, Unskipped], val countPerPage: Int, val pageNum: Int = 1) {
+  def copy() = new BasePaginatedQuery(q, countPerPage, pageNum)
+  def setPage(p: Int) = if (p == pageNum) this else new BasePaginatedQuery(q, countPerPage, p)
+  lazy val countAll: Long = q.count
+  def fetch(): List[R] = q.skip(countPerPage * (pageNum - 1)).limit(countPerPage).fetch()
+  def numPages = math.ceil(countAll.toDouble / countPerPage.toDouble).toInt max 1
+}