Add ObjectIdListField to Slashem

foursquare · May 21, 2012 · d107aee · d107aee
1 parent 242ad44
commit d107aee
Show file tree

Hide file tree

Showing 3 changed files with 79 additions and 1 deletion.
diff --git a/src/main/scala/com/foursquare/slashem/Schema.scala b/src/main/scala/com/foursquare/slashem/Schema.scala
@@ -984,6 +984,22 @@ class SlashemLongListField[T <: Record[T]](owner: T) extends LongListField[T](ow
   def nin(lst: List[Long]) = Clause[Long](queryName, groupWithOr(lst.map({l: Long => Phrase(l)})),false)
 }
 
+class SlashemObjectIdListField[T <: Record[T]](owner: T) extends ObjectIdListField[T](owner) with SlashemField[List[ObjectId], T] {
+  import Helpers._
+  override def valueBoxFromAny(a: Any) = objectIdBoxFromAny(a)
+  def contains(item: ObjectId) = {
+    Clause[ObjectId](queryName, Phrase(item))
+  }
+  /**
+   * See if this list has any elements in that list.
+   * @param List[ObjectId] the list to check for any intersections.
+   */
+  def in(lst: List[ObjectId]) = Clause[ObjectId](queryName, groupWithOr(lst.map({oid: ObjectId => Phrase(oid)})))
+  def nin(lst: List[ObjectId]) = Clause[ObjectId](queryName, groupWithOr(lst.map({oid: ObjectId => Phrase(oid)})),false)
+}
+
+
+
 class SlashemPointField[T <: Record[T]](owner: T) extends PointField[T](owner) with SlashemField[Pair[Double,Double], T] {
   def geoDistance(geolat: Double, geolng: Double) = {
     GeoDist(this.name,geolat,geolng)
@@ -1075,7 +1091,7 @@ class ObjectIdField[T <: Record[T]](override val owner: T) extends Field[ObjectI
 
   def valueBoxFromAny(a: Any): Box[ObjectId] = objectIdBoxFromAny(a)
 
-  override def setFromAny(a: Any) ={
+  override def setFromAny(a: Any) = {
     val vb = valueBoxFromAny(a)
     vb.map(set(_))
   }
@@ -1191,6 +1207,49 @@ class LongListField[T <: Record[T]](override val owner: T) extends Field[List[Lo
   override def valueBox() = e
 }
 
+class ObjectIdListField[T <: Record[T]](override val owner: T) extends Field[List[ObjectId], T] {
+  type ValueType = List[ObjectId]
+  var e: Box[ValueType] = Empty
+
+  def setFromString(s: String) = {
+    Full(set(s.split(" ").map(x => new ObjectId(x)).toList))
+  }
+  // NOTE(simon): See notes for ObjectIdField#objectIdBoxFromAny
+  def objectIdBoxFromAny(a: Any): Box[List[ObjectId]] = {
+    try {
+      a match {
+        case "" => Empty
+        case ar: Array[ObjectId] => Full(ar.toList)
+        case ar: Array[String] => Full(ar.toList.map(x => new ObjectId(x)))
+        case s: String => Full(s.split(" ").map(x => new ObjectId(x)).toList)
+        case _ => Empty
+      }
+    } catch {
+      case _ => Empty
+    }
+  }
+  def valueBoxFromAny(a: Any): Box[List[ObjectId]] = objectIdBoxFromAny(a)
+  override def setFromAny(a: Any) = {
+    val vb = valueBoxFromAny(a)
+    vb.map(set(_))
+  }
+  override def setFromJValue(jv: net.liftweb.json.JsonAST.JValue) = Empty
+  override def liftSetFilterToBox(a: Box[ValueType]) = Empty
+  override def toBoxMyType(a: ValueType) = Empty
+  override def defaultValueBox = Empty
+  override def toValueType(a: Box[MyType]) = null.asInstanceOf[ValueType]
+  override def asJValue() = net.liftweb.json.JsonAST.JNothing
+  override def asJs() = net.liftweb.http.js.JE.JsNull
+  override def toForm = Empty
+  override def set(a: ValueType) = {e = Full(a)
+                                    a.asInstanceOf[ValueType]}
+  override def get() = e.get
+  override def is() = e.get
+  def value() = e getOrElse Nil
+  override def valueBox() = e
+}
+
+
 class StringListField[T <: Record[T]](override val owner: T) extends Field[List[String], T] {
   type ValueType = List[String]
   var e: Box[ValueType] = Empty

diff --git a/src/test/scala/com/foursquare/slashem/ElasticQueryTest.scala b/src/test/scala/com/foursquare/slashem/ElasticQueryTest.scala
@@ -314,6 +314,7 @@ class ElasticQueryTest extends SpecsMatchers with ScalaCheckMatchers {
     val response7 = ESimplePanda where (_.hugenums contains 1L) fetch()
     val response8 = ESimplePanda where (_.hugenums contains 9L) fetch()
     val response9 = ESimplePanda where (_.hugenums contains 9001L) fetch()
+    val response10 = ESimplePanda where (_.favvenueids contains new ObjectId("4daf213893a0096fbaaef003")) fetch()
     Assert.assertEquals(response1.response.results.length, 2)
     Assert.assertEquals(response2.response.results.length, 1)
     Assert.assertEquals(response3.response.results.length, 2)
@@ -323,16 +324,19 @@ class ElasticQueryTest extends SpecsMatchers with ScalaCheckMatchers {
     Assert.assertEquals(response7.response.results.length, 2)
     Assert.assertEquals(response8.response.results.length, 1)
     Assert.assertEquals(response9.response.results.length, 0)
+    Assert.assertEquals(response10.response.results.length, 1)
   }
 
   @Test
   def testListFieldIn {
     val response1 = ESimplePanda where (_.favnums in List(2, 3, 4, 5)) fetch()
     val response2 = ESimplePanda where (_.favnums in List(99)) fetch()
     val response3 = ESimplePanda where (_.termsfield in List("termhit", "lol")) fetch()
+    val response4 = ESimplePanda where (_.favvenueids in List(new ObjectId("4daf213893a0096fbaaef003"))) fetch()
     Assert.assertEquals(response1.response.results.length, 2)
     Assert.assertEquals(response2.response.results.length, 0)
     Assert.assertEquals(response3.response.results.length, 1)
+    Assert.assertEquals(response4.response.results.length, 1)
   }
 
   @Test
@@ -351,6 +355,18 @@ class ElasticQueryTest extends SpecsMatchers with ScalaCheckMatchers {
     Assert.assertEquals(response2.response.results.length, 8)
   }
 
+  @Test
+  def testObjectIdListFieldEmptyIn {
+    val response1 = ESimplePanda where (_.favvenueids in List()) fetch()
+    Assert.assertEquals(response1.response.results.length, 0)
+  }
+
+ @Test
+  def testObjectIdListFieldEmptyNin {
+    val response1 = ESimplePanda where (_.favvenueids nin List()) fetch()
+    Assert.assertEquals(response1.response.results.length, 8)
+  }  
+
   @Test
   def testListFieldNin {
     val idsWithFavNums = List(new ObjectId("4c809f4251ada1cdc3790b10"),
@@ -443,6 +459,7 @@ class ElasticQueryTest extends SpecsMatchers with ScalaCheckMatchers {
     val hugenums1 = List(1L, 2L, 3L).asJava
     val hugenums2 = List(1L, 9L, 8L).asJava
     val hugenums3 = List(7L, 10L, 13L).asJava
+    val venueids1 = List("4daf213893a0096fbaaef003", "49ee02e9f964a52010681fe3", "42b21280f964a5206d251fe3").asJava
     val r = client.prepareIndex(ESimplePanda.meta.indexName,ESimplePanda.meta.docType,"4c809f4251ada1cdc3790b10").setSource(jsonBuilder()
                                                                           .startObject()
                                                                           .field("name","lolerskates")
@@ -451,6 +468,7 @@ class ElasticQueryTest extends SpecsMatchers with ScalaCheckMatchers {
                                                                           .field("nicknames", nicknames1)
                                                                           .field("hugenums", hugenums1)
                                                                           .field("termsfield", terms1)
+                                                                          .field("favvenueids", venueids1)
                                                                           .endObject()
       ).execute()
     .actionGet();

diff --git a/src/test/scala/com/foursquare/slashem/ElasticTest.scala b/src/test/scala/com/foursquare/slashem/ElasticTest.scala
@@ -19,6 +19,7 @@ class ESimplePanda extends ElasticSchema[ESimplePanda] {
   object nicknames extends SlashemStringListField(this)
   object hugenums extends SlashemLongListField(this)
   object termsfield extends SlashemUnanalyzedStringField(this)
+  object favvenueids extends SlashemObjectIdListField(this)
 }
 
 object ESimpleGeoPanda extends ESimpleGeoPanda with ElasticMeta[ESimpleGeoPanda] {