/tools/migration/0.5/appdata/src/main/scala/io/prediction/tools/migration/Appdata.scala
Scala | 135 lines | 97 code | 32 blank | 6 comment | 5 complexity | 7ec309769fa1a9d1185f1714f68551b9 MD5 | raw file
- package io.prediction.tools.migration
- //import com.typesafe.config._
- import io.prediction.commons.{Config}
- import com.mongodb.casbah.query.Imports._
- import com.mongodb.casbah.Imports.{MongoConnection}
- import com.mongodb.WriteResult
- object Appdata {
- //val config = ConfigFactory.load()
- val config = new Config()
- val appdataDbName = config.appdataDbName
- val appdataTrainingDbName = config.appdataTrainingDbName
- val appdataValidationDbName = config.appdataValidationDbName
- val appdataTestDbName = config.appdataTestDbName
- val appdataDb = config.appdataMongoDb.get
- val trainingDb = config.appdataTrainingMongoDb.get
- val validationDb = config.appdataValidationMongoDb.get
- val testDb = config.appdataTestMongoDb.get
- val dbSeq = Seq(
- appdataDbName -> appdataDb,
- appdataTrainingDbName -> trainingDb,
- appdataValidationDbName -> validationDb,
- appdataTestDbName -> testDb )
- def main(args: Array[String]) {
- println("PredictionIO 0.4 to 0.5 appdata Migration")
- println()
- println()
- val emptyObj = MongoDBObject()
- for ( (name, db) <- dbSeq ) {
-
- println(s"\nChecking ${name} DB...")
- val usersColl = db("users")
- val itemsColl = db("items")
- val u2iActionsColl = db("u2iActions")
- // users
- print("Updating user records with custom attributes...")
- val usersWithAttributes = usersColl.find(MongoDBObject("attributes" -> MongoDBObject("$exists" -> true)))
- var userCount = 0
- while (usersWithAttributes.hasNext) {
- val dbObj = usersWithAttributes.next
- userCount += 1
- val attributes = dbObj.getAs[DBObject]("attributes") map { dbObjToMap(_) }
- val flattenedObj = attributes map { a => attributesToMongoDBObject(a) } getOrElse emptyObj
-
- usersColl.update(MongoDBObject("_id" -> dbObj.as[String]("_id")), MongoDBObject("$set" -> flattenedObj, "$unset" -> MongoDBObject("attributes" -> 1)))
- }
- println(s"${userCount} users were updated.")
- // items
- print("Updating item records with custom attributes...")
- val itemsWithAttributes = itemsColl.find(MongoDBObject("attributes" -> MongoDBObject("$exists" -> true)))
- var itemCount = 0
- while (itemsWithAttributes.hasNext) {
- val dbObj = itemsWithAttributes.next
- itemCount += 1
- val attributes = dbObj.getAs[DBObject]("attributes") map { dbObjToMap(_) }
- val flattenedObj = attributes map { a => attributesToMongoDBObject(a) } getOrElse emptyObj
-
- itemsColl.update(MongoDBObject("_id" -> dbObj.as[String]("_id")), MongoDBObject("$set" -> flattenedObj, "$unset" -> MongoDBObject("attributes" -> 1)))
- }
- println(s"${itemCount} items were updated.")
-
- // u2iActions
- val rate: Int = 0
- val likeDislike: Int = 1 //(v==1 => like, v==0 => dislike)
- val view: Int = 2
- val viewDetails: Int = 3
- val conversion: Int = 4
- def checkWriteResult(wr: WriteResult) = {
- if (wr.getLastError().ok()) {
- println(s"${wr.getN} records were updated.")
- } else {
- wr.getLastError().throwOnError()
- }
- }
- print("Updating rate actions to new action name (may take a while)...")
- val rateResult = u2iActionsColl.update(MongoDBObject("action" -> rate), MongoDBObject("$set" -> MongoDBObject("action" -> "rate")), false, true)
- checkWriteResult(rateResult)
- print("Updating like actions to new action name (may take a while)...")
- val likeResult = u2iActionsColl.update(MongoDBObject("action" -> likeDislike, "v" -> 1), MongoDBObject("$set" -> MongoDBObject("action" -> "like"), "$unset" -> MongoDBObject("v" -> 1)), false, true)
- checkWriteResult(likeResult)
- print("Updating dislike actions to new action name (may take a while)...")
- val dislikeResult = u2iActionsColl.update(MongoDBObject("action" -> likeDislike, "v" -> 0), MongoDBObject("$set" -> MongoDBObject("action" -> "dislike"), "$unset" -> MongoDBObject("v" -> 1)), false, true)
- checkWriteResult(dislikeResult)
- print("Updating view actions to new action name (may take a while)...")
- val viewResult = u2iActionsColl.update(MongoDBObject("action" -> view), MongoDBObject("$set" -> MongoDBObject("action" -> "view")), false, true)
- checkWriteResult(viewResult)
- print("Updating viewDetails actions to new action name (may take a while)...")
- val viewDetailsResult = u2iActionsColl.update(MongoDBObject("action" -> viewDetails), MongoDBObject("$set" -> MongoDBObject("action" -> "viewDetails")), false, true)
- checkWriteResult(viewDetailsResult)
- print("Updating conversion actions to new action name (may take a while)...")
- val conversionResult = u2iActionsColl.update(MongoDBObject("action" -> conversion), MongoDBObject("$set" -> MongoDBObject("action" -> "conversion")), false, true)
- checkWriteResult(conversionResult)
-
- }
- println("Done.")
- }
- /** Modified from Salat. */
- def dbObjToMap(dbObj: DBObject): Map[String, Any] = {
- val builder = Map.newBuilder[String, Any]
- dbObj foreach {
- case (field, value) =>
- builder += field -> value
- }
- builder.result()
- }
- def attributesToMongoDBObject(attributes: Map[String, Any]) = {
- MongoDBObject( (attributes map { case (k, v) => ("ca_" + k, v) }).toList )
- }
- }