PageRenderTime 43ms CodeModel.GetById 16ms RepoModel.GetById 0ms app.codeStats 0ms

/commons/src/main/scala/io/prediction/commons/appdata/mongodb/MongoItems.scala

https://github.com/eddieliu/PredictionIO
Scala | 140 lines | 110 code | 21 blank | 9 comment | 2 complexity | 33229cfee825720460ce75c5b380c0c8 MD5 | raw file
  1. package io.prediction.commons.appdata.mongodb
  2. import io.prediction.commons.MongoUtils.{ emptyObj, mongoDbListToListOfString, idWithAppid }
  3. import io.prediction.commons.MongoUtils.{ attributesToMongoDBObject, getAttributesFromDBObject }
  4. import io.prediction.commons.appdata.{ Item, Items }
  5. import com.mongodb.casbah.Imports._
  6. import com.mongodb.casbah.commons.conversions.scala._
  7. import com.github.nscala_time.time.Imports._
  8. /** MongoDB implementation of Items. */
  9. class MongoItems(db: MongoDB) extends Items {
  10. private val itemColl = db("items")
  11. /** Indices and hints. */
  12. val starttimeIndex = MongoDBObject("starttime" -> -1)
  13. val lnglatIndex = MongoDBObject("lnglat" -> "2d")
  14. itemColl.ensureIndex(starttimeIndex)
  15. itemColl.ensureIndex(lnglatIndex)
  16. RegisterJodaTimeConversionHelpers()
  17. def insert(item: Item) = {
  18. val id = MongoDBObject("_id" -> idWithAppid(item.appid, item.id))
  19. val appid = MongoDBObject("appid" -> item.appid)
  20. val ct = MongoDBObject("ct" -> item.ct)
  21. val itypes = MongoDBObject("itypes" -> item.itypes)
  22. val starttime = item.starttime map { s => MongoDBObject("starttime" -> s) } getOrElse emptyObj
  23. val endtime = item.endtime map { e => MongoDBObject("endtime" -> e) } getOrElse emptyObj
  24. val price = item.price map { p => MongoDBObject("price" -> p) } getOrElse emptyObj
  25. val profit = item.profit map { p => MongoDBObject("profit" -> p) } getOrElse emptyObj
  26. val lnglat = item.latlng map { l => MongoDBObject("lnglat" -> MongoDBList(l._2, l._1)) } getOrElse emptyObj
  27. val inactive = item.inactive map { i => MongoDBObject("inactive" -> i) } getOrElse emptyObj
  28. //val attributes = item.attributes map { a => MongoDBObject("attributes" -> a) } getOrElse emptyObj
  29. // add "ca_" prefix for custom attributes
  30. val attributes = item.attributes map { a => attributesToMongoDBObject(a) } getOrElse emptyObj
  31. itemColl.save(id ++ appid ++ ct ++ itypes ++ starttime ++ endtime ++ price ++ profit ++ lnglat ++ inactive ++ attributes)
  32. }
  33. def get(appid: Int, id: String) = {
  34. itemColl.findOne(MongoDBObject("_id" -> idWithAppid(appid, id))) map { dbObjToItem(_) }
  35. }
  36. def getByAppid(appid: Int) = new MongoItemsIterator(itemColl.find(MongoDBObject("appid" -> appid)))
  37. def getByAppidAndLatlng(appid: Int, latlng: Tuple2[Double, Double], within: Option[Double], unit: Option[String]) = {
  38. val earthRadiusInKm = 6371
  39. val earthRadiusInMiles = 3959
  40. val nearSphereObj = MongoDBObject("$nearSphere" -> MongoDBList(latlng._2, latlng._1))
  41. val maxDistObj = within map { maxDist =>
  42. unit match {
  43. case Some("km") => MongoDBObject("$maxDistance" -> maxDist / earthRadiusInKm)
  44. case Some("mi") => MongoDBObject("$maxDistance" -> maxDist / earthRadiusInMiles)
  45. case _ => MongoDBObject("$maxDistance" -> maxDist / earthRadiusInKm)
  46. }
  47. } getOrElse emptyObj
  48. new MongoItemsIterator(itemColl.find(MongoDBObject("appid" -> appid, "lnglat" -> (nearSphereObj ++ maxDistObj))))
  49. }
  50. def getByAppidAndItypes(appid: Int, itypes: Seq[String]): Iterator[Item] = {
  51. getByAppidAndItypesAndTime(appid, optItypes = Some(itypes), optTime = None)
  52. }
  53. def getByAppidAndItypesAndTime(appid: Int,
  54. optItypes: Option[Seq[String]] = None,
  55. optTime: Option[DateTime] = None): Iterator[Item] = {
  56. val queryObj = ("appid" $eq appid)
  57. if (!optItypes.isEmpty)
  58. queryObj ++= ("itypes" $in optItypes.get)
  59. if (!optTime.isEmpty) {
  60. val time = optTime.get
  61. queryObj ++= ("starttime" $lt time)
  62. // Need to create an ugly DBObject as casbah don't have scala-ble support
  63. // for $or
  64. queryObj ++= DBObject("$or" -> MongoDBList(
  65. ("endtime" $exists false),
  66. ("endtime" $gt time)))
  67. }
  68. new MongoItemsIterator(itemColl.find(queryObj))
  69. }
  70. def getByIds(appid: Int, ids: Seq[String]) = {
  71. itemColl.find(MongoDBObject("_id" -> MongoDBObject("$in" -> ids.map(idWithAppid(appid, _))))).toList map { dbObjToItem(_) }
  72. }
  73. def getRecentByIds(appid: Int, ids: Seq[String]) = {
  74. itemColl.find(MongoDBObject("_id" -> MongoDBObject("$in" -> ids.map(idWithAppid(appid, _))))).sort(starttimeIndex).toList map { dbObjToItem(_) }
  75. }
  76. def update(item: Item) = {
  77. val id = MongoDBObject("_id" -> idWithAppid(item.appid, item.id))
  78. val appid = MongoDBObject("appid" -> item.appid)
  79. val ct = MongoDBObject("ct" -> item.ct)
  80. val itypes = MongoDBObject("itypes" -> item.itypes)
  81. val starttime = item.starttime map { s => MongoDBObject("starttime" -> s) } getOrElse emptyObj
  82. val endtime = item.endtime map { e => MongoDBObject("endtime" -> e) } getOrElse emptyObj
  83. val price = item.price map { p => MongoDBObject("price" -> p) } getOrElse emptyObj
  84. val profit = item.profit map { p => MongoDBObject("profit" -> p) } getOrElse emptyObj
  85. val lnglat = item.latlng map { l => MongoDBObject("lnglat" -> MongoDBList(l._2, l._1)) } getOrElse emptyObj
  86. val inactive = item.inactive map { i => MongoDBObject("inactive" -> i) } getOrElse emptyObj
  87. //val attributes = item.attributes map { a => MongoDBObject("attributes" -> a) } getOrElse emptyObj
  88. // add "ca_" prefix for custom attributes
  89. val attributes = item.attributes map { a => attributesToMongoDBObject(a) } getOrElse emptyObj
  90. itemColl.update(id, id ++ appid ++ ct ++ itypes ++ starttime ++ endtime ++ price ++ profit ++ lnglat ++ inactive ++ attributes)
  91. }
  92. def delete(appid: Int, id: String) = itemColl.remove(MongoDBObject("_id" -> idWithAppid(appid, id)))
  93. def delete(item: Item) = delete(item.appid, item.id)
  94. def deleteByAppid(appid: Int): Unit = {
  95. itemColl.remove(MongoDBObject("appid" -> appid))
  96. }
  97. def countByAppid(appid: Int): Long = itemColl.count(MongoDBObject("appid" -> appid))
  98. private def dbObjToItem(dbObj: DBObject) = {
  99. val appid = dbObj.as[Int]("appid")
  100. Item(
  101. id = dbObj.as[String]("_id").drop(appid.toString.length + 1),
  102. appid = appid,
  103. ct = dbObj.as[DateTime]("ct"),
  104. itypes = mongoDbListToListOfString(dbObj.as[MongoDBList]("itypes")),
  105. starttime = dbObj.getAs[DateTime]("starttime"),
  106. endtime = dbObj.getAs[DateTime]("endtime"),
  107. price = dbObj.getAs[Double]("price"),
  108. profit = dbObj.getAs[Double]("profit"),
  109. latlng = dbObj.getAs[MongoDBList]("lnglat") map { lnglat => (lnglat(1).asInstanceOf[Double], lnglat(0).asInstanceOf[Double]) },
  110. inactive = dbObj.getAs[Boolean]("inactive"),
  111. //attributes = dbObj.getAs[DBObject]("attributes") map { dbObjToMap(_) }
  112. attributes = Option(getAttributesFromDBObject(dbObj)).filter(!_.isEmpty)
  113. )
  114. }
  115. class MongoItemsIterator(it: MongoCursor) extends Iterator[Item] {
  116. def next = dbObjToItem(it.next)
  117. def hasNext = it.hasNext
  118. }
  119. }