db-migrate/src/main/groovy/com/jdblabs/dbmigrate/DbMigrate.groovy

package com.jdblabs.dbmigrate

import groovy.sql.Sql

import com.jdbernard.util.AnsiEscapeCodeSequence as ANSI
import com.zaxxer.hikari.HikariConfig
import com.zaxxer.hikari.HikariDataSource
import java.text.SimpleDateFormat
import org.docopt.Docopt
import org.slf4j.Logger
import org.slf4j.LoggerFactory

public class DbMigrate {

  public static final VERSION = "0.2.1"

  public static final def DOC = """\
db-migrate.groovy v${VERSION}

Usage:

  db_migrate [options] create <migration-name>
  db_migrate [options] up [<count>]
  db_migrate [options] down [<count>]
  db_migrate [options] init <schema-name>
  db_migrate (-V | --version)
  db_migrate (-h | --help)

Options:

  -c --config <config-file>   Use the given configuration file (defaults to
                              "database.properties").

  -q --quiet                  Suppress log information.

  -v --verbose                Print detailed log information.

     --very-verbose           Print very detailed log information.

  -V --version                Print the tools version information.

  -h --help                   Print this usage information.
"""

  private static sdf = new SimpleDateFormat('yyyyMMddHHmmss')
  private static Logger LOGGER = LoggerFactory.getLogger(DbMigrate)

  Sql sql
  File migrationsDir

  public static void main(String[] args) {

    // Parse arguments
    def opts = new Docopt(DOC).withVersion("wdiwtlt v$VERSION").parse(args)

    if (opts['--version']) {
      println "db-migrate.groovy v$VERSION"
      System.exit(0) }

    if (opts['--help']) { println DOC; System.exit(0) }

    // TODO: Setup logging & output levels
    Logger clilog = LoggerFactory.getLogger("db-migrate.cli")

    // Load the configuration file
    def givenCfg = new Properties()
    File cfgFile
    if (opts["--config"]) cfgFile = new File(opts["--config"])

    if (!cfgFile || !cfgFile.exists() || !cfgFile.isFile())
      cfgFile = new File("database.properties")

    if (!cfgFile.exists() || !cfgFile.isFile()) {
      clilog.warn("Config file '{}' does not exist or is not a regular file.",
        cfgFile.canonicalPath) }

    if (cfgFile.exists() && cfgFile.isFile()) {
      try { cfgFile.withInputStream { givenCfg.load(it) } }
      catch (Exception e) {
        clilog.error("Could not read configuration file.", e)
        givenCfg.clear() } }

    // Check for migrations directory
    File migrationsDir = new File(givenCfg["migrations.dir"])
    if (!migrationsDir.exists() || !migrationsDir.isDirectory()) {
      clilog.error("'{}' does not exist or is not a directory.",
        migrationsDir.canonicalPath)
      System.exit(1) }

    // Create the datasource
    HikariConfig hcfg = new HikariConfig(givenCfg)
    HikariDataSource hds = new HikariDataSource(hcfg)

    // Instantiate the DbMigrate instance
    DbMigrate dbmigrate = new DbMigrate(new Sql(hds))

    // Execute the appropriate command.
    if (opts['create']) {
      try {
        List<File> files = dbmigrate.createMigration(opts['<migration-name>'])
        clilog.info("Creted new migration files:\n\t${files.name.join('\n\t')}") }
      catch (Exception e) {
        clilog.error('Unable to create migration scripts.', e) } }

    else if (opts['up']) dbmigrate.up(opts['<count>'])
    else if (opts['down']) dbmigrate.down(opts['<count>'] ?: 1) }

  public File createMigration(String migrationName) {
    String timestamp = sdf.format(new Date())

    File upFile = new File(migrationsDir, "$timestamp-$migrationName-up.sql")
    File downFiles = new File(migrationsDir, "$timestamp-$migrationName-down.sql")

    upFile.text = "-- UP script for $migrationName ($timestamp)"
    downFile.text = "-- DOWN script for $migrationName ($timestamp)"

    return [upFile, downFile] }

  public def createMigrationsTable() {
    sql.execute('''
CREATE TABLE IF NOT EXISTS migrations (
  id SERIAL PRIMARY KEY,
  name VARCHAR NOT NULL,
  run_at TIMESTAMP NOT NULL DEFAULT NOW())''') }

  public def diffMigrations() {
    def results = [notRun: [], missing: []]

    results.run = sql.rows('SELECT name FROM migrations ORDER BY name')
      .collect { it.name }.sort()

    SortedSet<String> available = new TreeSet<>()
    migrationsDir.eachFile { f ->
      available << f.name.replaceAll(/-(up|down).sql$/, '') }

    available.each { migrationName ->
      if (!results.run.contains(migrationName))
        results.notRun << migrationName

      // If we've already seen some migrations that have not been run but this
      // one has been run, that means  we have a gap and are missing migrations.
      else if (results.notRun.size() > 0) {
        results.missing += reults.notRun
        results.notRun = [] } }

    return results }

  public List<String> up(Integer count = null) {
    createMigrationsTable()
    def diff = diffMigrations()

    List<String> toRun = count < diff.notRun.size() ?
      notRun[0..<count] : notRun

    return runMigrations(toRun, true) }

  public List<File> down(Integer count = 1) {
    createMigrationsTable()
    def diff = diffMigrations()

    List<String> toRun = count < diff.notRun.size() ?
      notRun.reverse()[0..<count] : notRun.reverse()

    return runMigrations(toRun, false) }

  private List<String> runMigrations(List<String> toRun, boolean up = true) {
    List<String> migrationsRun = []

    try {
      sql.execute('BEGIN')

      toRun.each { migrationName ->
        File migrationFile = new File(migrationsDir,
          "$migrationName-${up ? 'up' : 'down'}.sql")

        if (!migrationFile.exists() || !migrationFile.isFile())
          throw new FileNotFoundException(migrationFile.canonicalPath +
            "does not exist or is not a regular file.")

        List<String> statements = migrationFile.text.split(/;/)
          .findAll { it.trim().length() > 0 && !it.startsWith('--') }
          .collect { "$it;" }

        statements.each { sql.execute(it) }
        if (up) sql.executeInsert(
          'INSERT INTO migrations (name) VALUES (?);', migrationName)
        else sql.execute(
          'DELETE FROM migrations WHERE name = ?;', migrationName)

        migrationsRun << migrationName }

      sql.execute('COMMIT') }
    catch (Exception e) { sql.execute('ROLLBACK'); }

    return migrationsRun }
}