hal/src/zutil/db/DBUpgradeHandler.java

package zutil.db;

import zutil.db.handler.ListSQLResult;
import zutil.db.handler.SimpleSQLResult;
import zutil.log.LogUtil;

import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * This class will take a reference DB and alter
 * the source DB to have the same structure.
 * NOTE: only works with sqllite
 *
 * Created by Ziver
 */
public class DBUpgradeHandler {
    private static final Logger logger = LogUtil.getLogger();

    private DBConnection reference;
    private DBConnection target;


    public DBUpgradeHandler(DBConnection reference){
        this.reference = reference;
    }

    public void setTargetDB(DBConnection db){
        this.target = db;
    }

    public void upgrade() throws SQLException {
        /*
        - beginTransaction
        - run a table creation with if not exists (we are doing an upgrade, so the table might not exists yet, it will fail alter and drop)
        - put in a list the existing columns List<String> columns = DBUtils.GetColumns(db, TableName);
        - backup table (ALTER table " + TableName + " RENAME TO 'temp_"                    + TableName)
        - create new table (the newest table creation schema)
        - get the intersection with the new columns, this time columns taken from the upgraded table (columns.retainAll(DBUtils.GetColumns(db, TableName));)
        - restore data (String cols = StringUtils.join(columns, ",");
                    db.execSQL(String.format(
                            "INSERT INTO %s (%s) SELECT %s from temp_%s",
                            TableName, cols, cols, TableName));
        )
        - remove backup table (DROP table 'temp_" + TableName)
        - setTransactionSuccessful
         */
        try {
            logger.fine("Starting upgrade transaction...");
            target.exec("BEGIN IMMEDIATE TRANSACTION");

            upgradeCreateTabels();
            upgradeDropTables();
            upgradeAlterTables();

            logger.fine("Committing upgrade transaction...");
            target.exec("COMMIT TRANSACTION");
        } catch(SQLException e){
            try {
                target.exec("ROLLBACK TRANSACTION");
            } catch (SQLException secondary_e) {
                logger.log(Level.SEVERE, null, secondary_e);
            }
            throw e;
        }
    }


    private void upgradeCreateTabels() throws SQLException {
        List<String> refTables = reference.exec("SELECT name FROM sqlite_master WHERE type='table';", new ListSQLResult<String>());
        List<String> targetTables = reference.exec("SELECT name FROM sqlite_master WHERE type='table';", new ListSQLResult<String>());

        for(String table : refTables){
            if(!targetTables.contains(table)){
                logger.fine("Creating new table: "+ table);
                // Get reference create sql
                PreparedStatement stmt = reference.getPreparedStatement("SELECT sql FROM sqlite_master WHERE name == ?");
                stmt.setString(1, table);
                String sql = DBConnection.exec(stmt, new SimpleSQLResult<String>());
                // Execute sql on target
                target.exec(sql);
            }
        }
    }

    private void upgradeDropTables() throws SQLException {
        List<String> refTables = reference.exec("SELECT name FROM sqlite_master WHERE type='table';", new ListSQLResult<String>());
        List<String> targetTables = reference.exec("SELECT name FROM sqlite_master WHERE type='table';", new ListSQLResult<String>());

        for(String table : targetTables){
            if(!refTables.contains(table)){
                logger.fine("Dropping table: ");
                PreparedStatement stmt = target.getPreparedStatement("DROP TABLE ?");
                stmt.setString(1, table);
                DBConnection.exec(stmt);
            }
        }
    }

    private void upgradeAlterTables() throws SQLException {
        List<String> refTables = reference.exec("SELECT name FROM sqlite_master WHERE type='table';", new ListSQLResult<String>());
        List<String> targetTables = reference.exec("SELECT name FROM sqlite_master WHERE type='table';", new ListSQLResult<String>());

        PreparedStatement stmt;
        for(String table : targetTables){
            if(refTables.contains(table)){
                // Get reference structure
                List<DBColumn> refStruct = reference.exec("PRAGMA table_info("+table+")",
                        new TableStructureResultHandler());
                // Get target structure
                List<DBColumn> targetStruct = target.exec("PRAGMA table_info("+table+")",
                        new TableStructureResultHandler());

                // Check existing columns
                for(DBColumn column : refStruct) {
                    if(!targetStruct.contains(column)) {
                        logger.fine("Adding column '" + column.name + "' to table: " + table);
                        stmt = target.getPreparedStatement("ALTER TABLE ? ADD COLUMN ? ? ?");
                        stmt.setString(1, table); // Table name
                        stmt.setString(2, column.name); // Column name
                        stmt.setString(2, column.type); // Column type
                        stmt.setString(2, ""+// Column constraints
                                (column.defaultValue != null ? " DEFAULT '"+column.defaultValue+"'" : "")+
                                (column.notNull ? " NOT NULL" : "")+
                                (column.publicKey ? " PRIMARY KEY" : "")
                        );
                        DBConnection.exec(stmt);
                    }
                }
                // Check unnecessary columns
                for(DBColumn column : targetStruct) {
                    if(!refStruct.contains(column)) {
                        logger.warning("Unable to drop column: '" + column.name + "' from table: "+ table +" (SQLite does not support dropping columns)");
                    }
                }
            }
        }
    }

    private static class DBColumn{
        String name;
        String type;
        boolean notNull;
        String defaultValue;
        boolean publicKey;

        public boolean equals(Object obj){
            return obj instanceof DBColumn &&
                    name.equals(((DBColumn)obj).name);
        }
    }

    private static class TableStructureResultHandler implements SQLResultHandler<List<DBColumn>>{

        @Override
        public List<DBColumn> handleQueryResult(Statement stmt, ResultSet result) throws SQLException {
            ArrayList<DBColumn> list = new ArrayList<>();
            while (result.next()){
                DBColumn column = new DBColumn();
                column.name = result.getString("name");
                column.type = result.getString("type");
                column.notNull = result.getBoolean("notnull");
                column.defaultValue = result.getString("dflt_value");
                column.publicKey = result.getBoolean("pk");
                list.add(column);
            }
            return list;
        }
    }
}