PageRenderTime 50ms CodeModel.GetById 26ms RepoModel.GetById 1ms app.codeStats 0ms

/commons/src/main/scala/io/prediction/commons/settings/mongodb/MongoOfflineTunes.scala

https://github.com/eddieliu/PredictionIO
Scala | 91 lines | 67 code | 22 blank | 2 comment | 0 complexity | bd09c14a1c8270ef5e681c88b63abdef MD5 | raw file
  1. package io.prediction.commons.settings.mongodb
  2. import io.prediction.commons.settings.{ OfflineTune, OfflineTunes }
  3. import com.mongodb.casbah.Imports._
  4. import com.mongodb.casbah.commons.conversions.scala.RegisterJodaTimeConversionHelpers
  5. import com.github.nscala_time.time.Imports._
  6. class MongoOfflineTunes(db: MongoDB) extends OfflineTunes {
  7. private val emptyObj = MongoDBObject()
  8. private val offlineTuneColl = db("offlineTunes")
  9. private val seq = new MongoSequences(db)
  10. private def genNextId = seq.genNext("offlineTuneid")
  11. private val getFields = MongoDBObject( // fields to be read
  12. "engineid" -> 1,
  13. "loops" -> 1,
  14. "createtime" -> 1,
  15. "starttime" -> 1,
  16. "endtime" -> 1
  17. )
  18. RegisterJodaTimeConversionHelpers()
  19. private def dbObjToOfflineTune(dbObj: DBObject) = {
  20. OfflineTune(
  21. id = dbObj.as[Int]("_id"),
  22. engineid = dbObj.as[Int]("engineid"),
  23. loops = dbObj.as[Int]("loops"),
  24. createtime = dbObj.getAs[DateTime]("createtime"),
  25. starttime = dbObj.getAs[DateTime]("starttime"),
  26. endtime = dbObj.getAs[DateTime]("endtime")
  27. )
  28. }
  29. class MongoOfflineTuneIterator(it: MongoCursor) extends Iterator[OfflineTune] {
  30. def next = dbObjToOfflineTune(it.next)
  31. def hasNext = it.hasNext
  32. }
  33. def insert(offlineTune: OfflineTune): Int = {
  34. val id = genNextId
  35. val obj = MongoDBObject(
  36. "_id" -> id,
  37. "engineid" -> offlineTune.engineid,
  38. "loops" -> offlineTune.loops
  39. )
  40. // optional fields
  41. val createtimeObj = offlineTune.createtime.map(x => MongoDBObject("createtime" -> x)).getOrElse(emptyObj)
  42. val starttimeObj = offlineTune.starttime.map(x => MongoDBObject("starttime" -> x)).getOrElse(emptyObj)
  43. val endtimeObj = offlineTune.endtime.map(x => MongoDBObject("endtime" -> x)).getOrElse(emptyObj)
  44. val optObj = createtimeObj ++ starttimeObj ++ endtimeObj
  45. offlineTuneColl.insert(obj ++ optObj)
  46. id
  47. }
  48. def get(id: Int): Option[OfflineTune] = {
  49. offlineTuneColl.findOne(MongoDBObject("_id" -> id), getFields) map { dbObjToOfflineTune(_) }
  50. }
  51. def getAll() = new MongoOfflineTuneIterator(offlineTuneColl.find())
  52. def getByEngineid(engineid: Int): Iterator[OfflineTune] = new MongoOfflineTuneIterator(offlineTuneColl.find(MongoDBObject("engineid" -> engineid), getFields))
  53. def update(offlineTune: OfflineTune, upsert: Boolean = false) = {
  54. val obj = MongoDBObject(
  55. "_id" -> offlineTune.id,
  56. "engineid" -> offlineTune.engineid,
  57. "loops" -> offlineTune.loops
  58. )
  59. // optional fields
  60. val createtimeObj = offlineTune.createtime.map(x => MongoDBObject("createtime" -> x)).getOrElse(emptyObj)
  61. val starttimeObj = offlineTune.starttime.map(x => MongoDBObject("starttime" -> x)).getOrElse(emptyObj)
  62. val endtimeObj = offlineTune.endtime.map(x => MongoDBObject("endtime" -> x)).getOrElse(emptyObj)
  63. val optObj = createtimeObj ++ starttimeObj ++ endtimeObj
  64. offlineTuneColl.update(MongoDBObject("_id" -> offlineTune.id), obj ++ optObj, upsert)
  65. }
  66. def delete(id: Int) = {
  67. offlineTuneColl.remove(MongoDBObject("_id" -> id))
  68. }
  69. }