Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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")
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment