PageRenderTime 45ms CodeModel.GetById 19ms RepoModel.GetById 1ms app.codeStats 0ms

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

https://github.com/eddieliu/PredictionIO
Scala | 97 lines | 74 code | 19 blank | 4 comment | 0 complexity | c9103ab8da207a0f82f02befe18388cd MD5 | raw file
  1. package io.prediction.commons.settings.mongodb
  2. import io.prediction.commons.MongoUtils
  3. import io.prediction.commons.settings.{ Engine, Engines }
  4. import com.mongodb.casbah.Imports._
  5. /** MongoDB implementation of Engines. */
  6. class MongoEngines(db: MongoDB) extends Engines {
  7. private val engineColl = db("engines")
  8. private val seq = new MongoSequences(db)
  9. private def dbObjToEngine(dbObj: DBObject) = {
  10. /** Transparent upgrade. Remove in next minor version. */
  11. dbObj.getAs[DBObject]("settings") map { settings =>
  12. val e = Engine(
  13. id = dbObj.as[Int]("_id"),
  14. appid = dbObj.as[Int]("appid"),
  15. name = dbObj.as[String]("name"),
  16. infoid = dbObj.as[String]("infoid"),
  17. itypes = dbObj.getAs[MongoDBList]("itypes") map { MongoUtils.mongoDbListToListOfString(_) },
  18. params = MongoUtils.dbObjToMap(settings),
  19. trainingdisabled = dbObj.getAs[Boolean]("trainingdisabled"),
  20. trainingschedule = dbObj.getAs[String]("trainingschedule"))
  21. update(e)
  22. e
  23. } getOrElse {
  24. Engine(
  25. id = dbObj.as[Int]("_id"),
  26. appid = dbObj.as[Int]("appid"),
  27. name = dbObj.as[String]("name"),
  28. infoid = dbObj.as[String]("infoid"),
  29. itypes = dbObj.getAs[MongoDBList]("itypes") map { MongoUtils.mongoDbListToListOfString(_) },
  30. params = MongoUtils.dbObjToMap(dbObj.as[DBObject]("params")),
  31. trainingdisabled = dbObj.getAs[Boolean]("trainingdisabled"),
  32. trainingschedule = dbObj.getAs[String]("trainingschedule"))
  33. }
  34. }
  35. class MongoEngineIterator(it: MongoCursor) extends Iterator[Engine] {
  36. def next = dbObjToEngine(it.next)
  37. def hasNext = it.hasNext
  38. }
  39. def insert(engine: Engine) = {
  40. val id = seq.genNext("engineid")
  41. // required fields
  42. val obj = MongoDBObject(
  43. "_id" -> id,
  44. "appid" -> engine.appid,
  45. "name" -> engine.name,
  46. "infoid" -> engine.infoid,
  47. "params" -> engine.params
  48. )
  49. // optional fields
  50. val optObj = engine.itypes.map(x => MongoDBObject("itypes" -> x)).getOrElse(MongoUtils.emptyObj) ++
  51. engine.trainingdisabled.map(x => MongoDBObject("trainingdisabled" -> x)).getOrElse(MongoUtils.emptyObj) ++
  52. engine.trainingschedule.map(x => MongoDBObject("trainingschedule" -> x)).getOrElse(MongoUtils.emptyObj)
  53. engineColl.insert(obj ++ optObj)
  54. id
  55. }
  56. def get(id: Int) = engineColl.findOne(MongoDBObject("_id" -> id)) map { dbObjToEngine(_) }
  57. def getAll() = new MongoEngineIterator(engineColl.find())
  58. def getByAppid(appid: Int) = new MongoEngineIterator(engineColl.find(MongoDBObject("appid" -> appid)).sort(MongoDBObject("name" -> 1)))
  59. def getByAppidAndName(appid: Int, name: String) = engineColl.findOne(MongoDBObject("appid" -> appid, "name" -> name)) map { dbObjToEngine(_) }
  60. def getByIdAndAppid(id: Int, appid: Int): Option[Engine] = engineColl.findOne(MongoDBObject("_id" -> id, "appid" -> appid)) map { dbObjToEngine(_) }
  61. def update(engine: Engine, upsert: Boolean = false) = {
  62. val idObj = MongoDBObject("_id" -> engine.id)
  63. val nameObj = MongoDBObject("name" -> engine.name)
  64. val appidObj = MongoDBObject("appid" -> engine.appid)
  65. val infoidObj = MongoDBObject("infoid" -> engine.infoid)
  66. val itypesObj = engine.itypes.map(x => MongoDBObject("itypes" -> x)).getOrElse(MongoUtils.emptyObj)
  67. val paramsObj = MongoDBObject("params" -> engine.params)
  68. val trainingdisabledObj = engine.trainingdisabled.map(x => MongoDBObject("trainingdisabled" -> x)).getOrElse(MongoUtils.emptyObj)
  69. val trainingscheduleObj = engine.trainingschedule.map(x => MongoDBObject("trainingschedule" -> x)).getOrElse(MongoUtils.emptyObj)
  70. engineColl.update(
  71. idObj,
  72. idObj ++ appidObj ++ nameObj ++ infoidObj ++ itypesObj ++ paramsObj ++ trainingdisabledObj ++ trainingscheduleObj,
  73. upsert
  74. )
  75. }
  76. def deleteByIdAndAppid(id: Int, appid: Int) = engineColl.remove(MongoDBObject("_id" -> id, "appid" -> appid))
  77. def existsByAppidAndName(appid: Int, name: String) = engineColl.findOne(MongoDBObject("name" -> name, "appid" -> appid)) map { _ => true } getOrElse false
  78. }