Liquibase integration for Play 2.0

SchemaGen.java

package com.benmccann.example.schema;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.PrintStream;
import java.io.PrintWriter;
import java.io.Writer;
import java.net.URL;
import java.util.Arrays;
import java.util.List;
import java.util.Set;

import javax.xml.parsers.ParserConfigurationException;

import liquibase.Liquibase;
import liquibase.database.Database;
import liquibase.diff.Diff;
import liquibase.diff.DiffResult;
import liquibase.diff.DiffStatusListener;
import liquibase.exception.DatabaseException;
import liquibase.exception.LiquibaseException;
import liquibase.integration.commandline.CommandLineResourceAccessor;
import liquibase.integration.commandline.CommandLineUtils;
import liquibase.resource.CompositeResourceAccessor;
import liquibase.resource.FileSystemResourceAccessor;
import liquibase.snapshot.DatabaseSnapshot;
import liquibase.snapshot.DatabaseSnapshotGeneratorFactory;

import org.apache.commons.io.FileUtils;

import play.Application;
import play.Configuration;
import play.Logger;

import com.avaje.ebean.Ebean;
import com.avaje.ebean.EbeanServer;
import com.avaje.ebean.SqlRow;
import com.avaje.ebean.config.ServerConfig;
import com.avaje.ebean.config.dbplatform.MySqlPlatform;
import com.avaje.ebeaninternal.api.SpiEbeanServer;
import com.avaje.ebeaninternal.server.ddl.DdlGenerator;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Lists;
import com.typesafe.config.ConfigFactory;

public class SchemaGen {

  private static final String NEW_LINE = System.getProperty("line.separator");
  private static final String EMPTY_DIFF = "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>" + NEW_LINE
      + "<databaseChangeLog xmlns=\"http://www.liquibase.org/xml/ns/dbchangelog\" xmlns:xsi=\"http://www.w3.org/2001/XMLSchema-instance\" xsi:schemaLocation=\"http://www.liquibase.org/xml/ns/dbchangelog http://www.liquibase.org/xml/ns/dbchangelog/dbchangelog-2.0.xsd\"/>" + NEW_LINE;
  private static final DiffStatusListener STATUS_LISTENER = new OutDiffStatusListener();
  private static final Set<DiffStatusListener> STATUS_LISTENERS = ImmutableSet.of(STATUS_LISTENER);
  
  public static void dropCreateDB(Application application)
      throws IOException, DatabaseException {
    String serverName = "default";
    EbeanServer server = Ebean.getServer(serverName);
    ServerConfig config = new ServerConfig();
    config.isDdlGenerate();
    config.setDdlRun(true);

    DdlGenerator ddl = new DdlGenerator((SpiEbeanServer) server, new MySqlPlatform(), config);
    String createSql = ddl.generateCreateDdl();

    // Write the SQL for dropping and creating the DB to disk.
    // If it's different than last time then execute the SQL.
    File dropCreateFile = application.getFile(
        "conf/evolutions/" + serverName + "/create.sql");
    if (!dropCreateFile.exists()
        || !createSql.equals(FileUtils.readFileToString(dropCreateFile))
        || getNumMySqlTables(server) == 0) {
      FileUtils.writeStringToFile(dropCreateFile, createSql);
      dropLocalDb();
      ddl.runScript(false, createSql);
    }
  }

  private static void dropLocalDb() throws DatabaseException {
    Database referenceDatabase = getDevDatabase();
    String schema = referenceDatabase.getConnection().getCatalog();

    Logger.warn("Dropping database " + schema);
    referenceDatabase.dropDatabaseObjects(schema);
  }
  
  private static Integer getNumMySqlTables(EbeanServer server) {
    SqlRow row = server.createSqlQuery(
        "select count(*) from information_schema.TABLES where TABLE_SCHEMA=DATABASE()")
        .findList().get(0);
    return row.getInteger(row.keys().next());
  }

  private static String generateLiquibaseXmlDiff()
      throws DatabaseException, ParserConfigurationException, IOException {
    Database referenceDatabase = getDevDatabase();
    DatabaseSnapshot referenceSnapshot = generateDatabaseSnapshot(referenceDatabase);

    Database targetDatabase = getProdDatabase();
    DatabaseSnapshot targetSnapshot = generateDatabaseSnapshot(targetDatabase);

    Diff diff = new Diff(referenceSnapshot, targetSnapshot);
    diff.addStatusListener(STATUS_LISTENER);
    DiffResult diffResult = diff.compare();

    ByteArrayOutputStream changeLogStream = new ByteArrayOutputStream();
    diffResult.printChangeLog(new PrintStream(changeLogStream), targetDatabase);
    return changeLogStream.toString("UTF-8");
  }

  private static void generateSqlFromLiquibaseXml(int evolutionNum)
      throws ParserConfigurationException, IOException, LiquibaseException {
    String sqlFile = getSqlFileName(evolutionNum);
    PrintWriter printWriter = new PrintWriter(sqlFile);
    generateLiquibaseSql(
        getLiquibaseChangeLogFileName(evolutionNum),
        getProdDatabase(),
        printWriter);

    // Read the file and delete out the liquibase stuff
    File file = new File(sqlFile);
    String fileContents = FileUtils.readFileToString(file);
    List<String> lines = Lists.newArrayList(fileContents.split("\n"));
    int lineNum = lines.size() - 1;
    while (lineNum > 0) {
      String line = lines.get(lineNum);
      if (line.contains("DATABASECHANGELOG")) {
        lines.remove(lineNum);        
      }
      lineNum--;
    }
    FileUtils.writeLines(file, lines);
  }
  
  private static void generateLiquibaseSql(
      String changeLogFile, Database targetDatabase, Writer writer)
      throws ParserConfigurationException, IOException, LiquibaseException {
    FileSystemResourceAccessor fsOpener = new FileSystemResourceAccessor();
    CommandLineResourceAccessor clOpener = new CommandLineResourceAccessor(SchemaGen.class.getClassLoader());
    CompositeResourceAccessor fileOpener = new CompositeResourceAccessor(fsOpener, clOpener);

    Liquibase liquibase = new Liquibase(changeLogFile, fileOpener, targetDatabase);
    liquibase.update(null, writer);
  }

  private static Database getDevDatabase() throws DatabaseException {
    return getDatabase("application.conf", null);
  }
  
  private static Database getProdDatabase() throws DatabaseException {
    return getDatabase("prod.conf", "application.conf");
  }

  private static Database getDatabase(String confName, String fallbackConfName)
      throws DatabaseException {
    URL confUrl = SchemaGen.class.getClassLoader().getResource(confName);
    Configuration config = new Configuration(
        new play.api.Configuration(ConfigFactory.parseURL(confUrl)));
    if (fallbackConfName == null) {
      return getDatabase(config, null);      
    }

    URL fallbackConfUrl = SchemaGen.class.getClassLoader().getResource(fallbackConfName);
    Configuration fallbackConfig = new Configuration(
        new play.api.Configuration(ConfigFactory.parseURL(fallbackConfUrl)));
    return getDatabase(config, fallbackConfig);
  }

  private static String readConfigProperty(String configProperty,
      Configuration config, Configuration fallbackConfig) {
    String value = config.getString(configProperty);
    if (Strings.isNullOrEmpty(value) && fallbackConfig != null) {
      value = fallbackConfig.getString(configProperty);
    }
    if (Strings.isNullOrEmpty(value)) {
      throw new IllegalArgumentException("Could not find property " + configProperty);
    }
    return value;
  }

  private static Database getDatabase(Configuration config, Configuration fallbackConfig) throws DatabaseException {
    return CommandLineUtils.createDatabaseObject(
        SchemaGen.class.getClassLoader(),
        readConfigProperty("db.default.url", config, fallbackConfig),
        readConfigProperty("db.default.user", config, fallbackConfig),
        readConfigProperty("db.default.password", config, fallbackConfig),
        readConfigProperty("db.default.driver", config, fallbackConfig),
        null, null, null);
  }

  private static DatabaseSnapshot generateDatabaseSnapshot(Database database)
      throws DatabaseException {
    return DatabaseSnapshotGeneratorFactory.getInstance()
        .createSnapshot(database, null, STATUS_LISTENERS);
  }

  private static String getLiquibaseChangeLogFileName(int evolutionNum) {
    return "conf/evolutions/default/" + evolutionNum + ".xml";
  }

  private static String getSqlFileName(int evolutionNum) {
    return "conf/evolutions/default/" + evolutionNum + ".sql";
  }

  private static void writeStringToFile(int evolutionNum, String fileContents) throws IOException {
    FileUtils.writeStringToFile(new File(getLiquibaseChangeLogFileName(evolutionNum)), fileContents);
  }
  
  private static class OutDiffStatusListener implements DiffStatusListener {
    public void statusUpdate(String message) {
      Logger.info(message);
    }
  }

  public static void printHelp() {
    StringBuilder sb = new StringBuilder()
        .append("Usage:").append(NEW_LINE)
        .append("SchemaGen <evolution_number>").append(NEW_LINE)
        .append("    Pass the evolution number you want to auto-generate.").append(NEW_LINE)
        .append("    If there is an XML file already present for that evolution number then only the SQL will be regenerated.").append(NEW_LINE)
        .append("SchemaGen --check_equality").append(NEW_LINE)
        .append("    Prints whether the local and remote schemas are equal.").append(NEW_LINE);
    System.out.println(sb.toString());
  }

  private static void handleCmdLine(List<String> args) throws Exception {
    if (args.size() != 1) {
      throw new IllegalArgumentException("There must be exactly one argument.");
    }

    if (args.get(0).equals("--check_equality")) {
      String content = SchemaGen.generateLiquibaseXmlDiff();
      if (!content.equals(EMPTY_DIFF)) {
        System.err.println("Schemas are not equal.");
        System.exit(1);
      }
      System.err.println("Schemas are equal.");
      System.exit(0);
    }

    int evolutionNum = Integer.parseInt(args.get(0));
    if (!new File(getLiquibaseChangeLogFileName(evolutionNum)).exists()) {
      String content = SchemaGen.generateLiquibaseXmlDiff();
      writeStringToFile(evolutionNum, content);
    }
    SchemaGen.generateSqlFromLiquibaseXml(evolutionNum);
  }

  public static void main(String[] args) throws Exception {
    try {
      handleCmdLine(Arrays.asList(args));
    } catch (IllegalArgumentException e) {
      printHelp();
      throw e;
    }    
  }

}