Added logging, implemented diff, up, and down.

2016-02-07 07:37:06 -06:00 · 2016-02-07 07:37:06 -06:00 · 02e326e661
commit 02e326e661
parent efc5dd2ed9
1 changed files with 127 additions and 66 deletions
--- a/db_migrate.nim
+++ b/db_migrate.nim
@ -4,11 +4,10 @@
 ## Simple tool to manage database migrations.

 import algorithm, json, times, os, strutils, docopt, db_postgres, sets,
-  sequtils
+  sequtils, logging

 type
-  LogLevel = enum quiet, normal, verbose, very_verbose
-  DbMigrateConfig* = tuple[ driver, sqlDir, connectionString: string, logLevel: LogLevel ]
+  DbMigrateConfig* = tuple[ driver, sqlDir, connectionString: string, logLevel: Level ]

 proc createMigrationsTable(conn: DbConn): void =
  conn.exec(sql("""
@ -17,15 +16,24 @@ CREATE TABLE IF NOT EXISTS migrations (
  name VARCHAR NOT NULL,
  run_at TIMESTAMP NOT NULL DEFAULT NOW())"""))

+proc rollbackWithErr (pgConn: DbConn, errMsg: string): void =
+  pgConn.exec(sql"ROLLBACK")
+  dbError(errMsg)
+
 proc loadConfig*(filename: string): DbMigrateConfig =
  ## Load DbMigrateConfig from a file.
  let cfg = json.parseFile(filename)

+  var logLevel: Level
+  if cfg.hasKey("logLevel"):
+    let idx = find(LevelNames, cfg["logLevel"].getStr.toUpper)
+    logLevel = if idx == -1: lvlInfo else: (Level)(idx)
+
  return (
    driver: if cfg.hasKey("driver"): cfg["driver"].getStr else: "postres",
    sqlDir: if cfg.hasKey("sqlDir"): cfg["sqlDir"].getStr else: "migrations",
    connectionString: cfg["connectionString"].getStr,
-    logLevel: if cfg.hasKey("logLevel"): parseEnum[LogLevel](cfg["logLevel"].getStr) else: normal)
+    logLevel: logLevel)

 proc createMigration*(config: DbMigrateConfig, migrationName: string): seq[string] =
  ## Create a new set of database migration files.
@ -70,8 +78,8 @@ proc diffMigrations*(pgConn: DbConn, config: DbMigrateConfig):
  let migrationsInOrder =
    toSeq(migrationsAvailable.items).sorted(system.cmp)

-  var migrationsNotRun: seq[string] = @[]
-  var missingMigrations: seq[string] = @[]
+  var migrationsNotRun = newSeq[string]()
+  var missingMigrations = newSeq[string]()

  for migration in migrationsInOrder:
    if not migrationsRun.contains(migration):
@ -81,49 +89,75 @@ proc diffMigrations*(pgConn: DbConn, config: DbMigrateConfig):
    # one has been, that means we have a gap and are missing migrations
    elif migrationsNotRun.len > 0:
      missingMigrations.add(migrationsNotRun)
-      migrationsNotRun = @[]
+      migrationsNotRun = newSeq[string]()

-  return (run: toSeq(migrationsRun.items), notRun: migrationsNotRun, missing: missingMigrations)
+  return (run: toSeq(migrationsRun.items).sorted(system.cmp),
+          notRun: migrationsNotRun,
+          missing: missingMigrations)

-proc up*(config: DbMigrateConfig, count: int): seq[string] =
-  let pgConn = open("", "", "", config.connectionString)
-  pgConn.createMigrationsTable
+proc readStatements*(filename: string): seq[SqlQuery] =
+  let migrationSql = filename.readFile
+  result = migrationSql.split(';').
+    filter(proc(st: string): bool = st.strip.len > 0 and not st.startsWith("--")).
+    map(proc(st: string): SqlQuery = sql(st & ";"))

-  let (run, toRun, missing) = diffMigrations(pgConn, config)
- 
-  # Make sure we have no gaps (database is in an unknown state)
-  if missing.len > 0:
-    dbError("Database is in an inconsistent state. Migrations have been " &
-      "run that are not sequential.")
-
-  var migrationsRun: seq[string] = @[]
-  if toRun.len > 0:
+proc up*(pgConn: DbConn, config: DbMigrateConfig, toRun: seq[string]): seq[string] =
+  var migrationsRun = newSeq[string]()
  # Begin a transaction.
  pgConn.exec(sql"BEGIN")

-    let rollbackWithErr = proc (errMsg: string): void =
-      pgConn.exec(sql"ROLLBACK")
-      dbError(errMsg)
-
  # Apply each of the migrations.
  for migration in toRun:
+    info migration
    let filename = joinPath(config.sqlDir, migration & "-up.sql")

    if not filename.fileExists:
-        rollbackWithErr "Can not find up file for " & migration &
+      pgConn.rollbackWithErr "Can not find UP file for " & migration &
        ". Expected '" & filename & "'."

-      let migrationSql = filename.readFile
+    let statements = filename.readStatements
+
    try:
-        pgConn.exec(sql(migrationSql))
-        migrationsRun.add(migration)
+      for statement in statements: pgConn.exec(statement)
+      pgConn.exec(sql"INSERT INTO migrations (name) VALUES (?);", migration)
    except DbError:
-        rollbackWithErr "Migration '" & migration & "'failed:\n\t" &
+      pgConn.rollbackWithErr "Migration '" & migration & "' failed:\n\t" &
        getCurrentExceptionMsg()

+    migrationsRun.add(migration)
+
+  pgConn.exec(sql"COMMIT")

  return migrationsRun

+proc down*(pgConn: DbConn, config: DbMigrateConfig, migrationsToDown: seq[string]): seq[string] =
+  var migrationsDowned = newSeq[string]()
+
+  pgConn.exec(sql"BEGIN")
+
+  for migration in migrationsToDown:
+    info migration
+
+    let filename = joinPath(config.sqlDir, migration & "-down.sql")
+
+    if not filename.fileExists:
+      pgConn.rollbackWithErr "Can not find DOWN file for " & migration &
+        ". Expected '" & filename & "'."
+
+    let statements = filename.readStatements
+
+    try:
+      for statement in statements: pgConn.exec(statement)
+      pgConn.exec(sql"DELETE FROM migrations WHERE name = ?;", migration)
+    except DbError:
+      pgConn.rollbackWithErr "Migration '" & migration & "' failed:\n\t" &
+        getCurrentExceptionMsg()
+
+    migrationsDowned.add(migration)
+
+  pgConn.exec(sql"COMMIT")
+  return migrationsDowned
+
 when isMainModule:
  let doc = """
 Usage:
@ -138,8 +172,11 @@ Options:
  -c --config <config-file>   Use the given configuration file (defaults to
                              "database.json").

-  -v --verbose                Print detailed log information (use -vv to
-                              print even more details).
+  -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.
 """
@ -148,7 +185,7 @@ Options:
  let args = docopt(doc, version = "db-migrate 0.2.0")

  let exitErr = proc(msg: string): void =
-    stderr.writeLine("db_migrate: " & msg)
+    fatal("db_migrate: " & msg)
    quit(QuitFailure)

  # Load configuration file
@ -164,10 +201,17 @@ Options:
  except:
    exitErr "Error parsing config file: " & configFilename & "\L\t" & getCurrentExceptionMsg()

+
+  logging.addHandler(newConsoleLogger())
+  if args["--quiet"]: logging.setLogFilter(lvlError)
+  elif args["--very-verbose"]: logging.setLogFilter(lvlAll)
+  elif args["--verbose"]: logging.setlogFilter(lvlDebug)
+  else: logging.setLogFilter(config.logLevel)
+
  # Check for migrations directory
  if not existsDir config.sqlDir:
    try:
-      echo "SQL directory '" & config.sqlDir &
+      warn "SQL directory '" & config.sqlDir &
        "' does not exist and will be created."
      createDir config.sqlDir
    except IOError:
@ -177,31 +221,48 @@ Options:
  if args["create"]:
    try:
      let filesCreated = createMigration(config, $args["<migration-name>"])
-      echo "Created new migration files:"
-      for filename in filesCreated: echo "\t" & filename
+      info "Created new migration files:"
+      for filename in filesCreated: info "\t" & filename
    except IOError:
      exitErr "Unable to create migration scripts: " & getCurrentExceptionMsg()

-  elif args["up"]:
+  else:
+    let pgConn: DbConn =
+      try: open("", "", "", config.connectionString)
+      except DbError:
+        exitErr "Unable to connect to the database."
+        (DbConn)(nil)
+
+    pgConn.createMigrationsTable
+
+    let (run, notRun, missing) = diffMigrations(pgConn, config)
+
+    # Make sure we have no gaps (database is in an unknown state)
+    if missing.len > 0:
+      exitErr "Database is in an inconsistent state. Migrations have been " &
+        "run that are not sequential."
+
+    if args["up"]:
      try:
-      let migrationsRun = up(config, if args["<count>"]: parseInt($args["<count>"]) else: 1)
-      for migration in migrationsRun:
-        echo migration
-      echo "Up to date."
+        let count = if args["<count>"]: parseInt($args["<count>"]) else: high(int)
+        let toRun = if count < notRun.len: notRun[0..<count] else: notRun
+        let migrationsRun = pgConn.up(config, toRun)
+        if count < high(int): info "Went up " & $(migrationsRun.len) & "."
      except DbError:
        exitErr "Unable to migrate database: " & getCurrentExceptionMsg()

-  elif args["down"]: discard
-    # Query the database to find out what migrations have been run.
-
-    # Find how many we need to go down (default to 1 if the count parameter was
-    # not passed.
-
-    # Begin transaction
-
-      # Apply each down script
-      # If any fail, roll back the transaction
-      # Otherwise report success
+    elif args["down"]:
+      try:
+        let count = if args["<count>"]: parseInt($args["<count>"]) else: 1
+        let toRun = if count < run.len: run.reversed[0..<count] else: run.reversed
+        let migrationsRun = pgConn.down(config, toRun)
+        info "Went down " & $(migrationsRun.len) & "."
+      except DbError:
+        exitErr "Unable to migrate database: " & getCurrentExceptionMsg()

    elif args["init"]: discard

+    let newResults = diffMigrations(pgConn, config)
+    if newResults.notRun.len > 0:
+      info "Database is behind by " & $(newResults.notRun.len) & " migrations."
+    else: info "Database is up to date."