DbMigrator.scala

package hoffrocket


import _root_.net.liftweb.mapper._
import _root_.java.sql.{PreparedStatement,Types, Connection,DriverManager}
import scala.collection.mutable.HashMap
import _root_.net.liftweb.util._
import Helpers._

/**
Usage Example:

val fromDbMan = new DbMigrator.SimpleConnManager("com.mysql.jdbc.Driver","jdbc:mysql://localhost/mydb",Full("user"),Full("pass"))
val toDbMan = new DbMigrator.SimpleConnManager("org.h2.Driver","jdbc:h2:~/test;DB_CLOSE_DELAY=-1",None,None)
val batchSize = 1000
DbMigrator.migrate(fromDbMan, toDbMan, batchSize, User, Dog, Employee, Etc)

println("The counts should be equal: " + User.countDb(DbMigrator.FromConnectionIdentifier) == User.countDb(DbMigrator.ToConnectionIdentifier))
*/
object DbMigrator {
    
    def metaCast[T <: MetaMapper[T]](s:Any):T = s.asInstanceOf[T]
    def findSlice[T <: Mapper[T]](mapper: KeyedMetaMapper[_,T], start:Long, maxRows:Long):List[Mapper[T]] = mapper.findAllDb(FromConnectionIdentifier, StartAt(start), MaxRows(maxRows), OrderBy(mapper.primaryKeyField, Ascending))
    
    def migrate(fromManager:ConnectionManager, toManager:ConnectionManager, batchSize:Long, tables:KeyedMetaMapper[_,_]*) {
        
        DB.defineConnectionManager(FromConnectionIdentifier,  fromManager)
        DB.defineConnectionManager(ToConnectionIdentifier, toManager)
        
        Schemifier.schemify(true, Log.infoF _, ToConnectionIdentifier, tables:_*)
        for (table <- tables){
            val mappedColumnsMethod = table.getClass.getMethod("net$liftweb$mapper$MetaMapper$$mappedColumns")
            val mappedColumns = mappedColumnsMethod.invoke(table).asInstanceOf[HashMap[String,java.lang.reflect.Method]]
            val query = "INSERT INTO "+table.dbTableName+" (" + mappedColumns.map(_._1).mkString(",") + ") VALUES ("+ mappedColumns.map(p=>"?").mkString(",") + ")"
            Log.info("INSERTING with " + query)

            def batchCopy(start:Long) {
                findSlice(metaCast(table), start, batchSize) match {
                    case Nil => 
                    case records => {
                        DB.use(ToConnectionIdentifier){
                            conn=> {
                                    def prepStat(record: Any,st: PreparedStatement) {
                                          var colNum = 1

                                          for (col <- mappedColumns) {
                                              val colVal = col._2.invoke(record).asInstanceOf[MappedField[AnyRef, _]]
                                              colVal.targetSQLType(col._1) match {
                                                case Types.VARCHAR =>
                                                  st.setString(colNum, colVal.jdbcFriendly(col._1).asInstanceOf[String])

                                                case _ => st.setObject(colNum, colVal.jdbcFriendly(col._1), conn.driverType.columnTypeMap(colVal.targetSQLType(col._1)))                    
                                              }
                                              colNum = colNum + 1
                                          }
                                    }
                                    DB.prepareStatement(query, conn){ stmt =>
                                       for(record <- records){
                                          prepStat(record, stmt)
                                          stmt.addBatch
                                       }
                                       Log.info("\tExecuting a batch")
                                       stmt.executeBatch
                                   }
                            }
                        }
                        
                        if (records.length == batchSize) {
                            batchCopy(start + batchSize)
                        }
                    }
                }
            }
            
            batchCopy(0L)
        }
        
    }
    
    class SimpleConnManager(driver:String, url:String, username:Box[String],password:Box[String]) extends ConnectionManager {
      def newConnection(name: ConnectionIdentifier): Box[Connection] = {
        try {
          Class.forName(driver)
          val dm = (username, password) match {
            case (Full(user), Full(pwd)) => DriverManager.getConnection(url, user, pwd)
            case _ => DriverManager.getConnection(url)
          }
          Full(dm)
        } catch {
          case e : Exception => e.printStackTrace; Empty
        }
      }
      def releaseConnection(conn: Connection) {conn.close}
    }
    
    case class NamedConnectionIdentifier(jndiName:String) extends ConnectionIdentifier
    case object ToConnectionIdentifier extends NamedConnectionIdentifier("migrateTo")
    case object FromConnectionIdentifier extends NamedConnectionIdentifier("migrateFrom")
}