208 lines
		
	
	
		
			6.8 KiB
		
	
	
	
		
			Nim
		
	
	
	
	
	
			
		
		
	
	
			208 lines
		
	
	
		
			6.8 KiB
		
	
	
	
		
			Nim
		
	
	
	
	
	
## DB Migrate
 | 
						|
## ==========
 | 
						|
##
 | 
						|
## Simple tool to manage database migrations.
 | 
						|
 | 
						|
import algorithm, json, times, os, strutils, docopt, db_postgres, sets,
 | 
						|
  sequtils
 | 
						|
 | 
						|
type 
 | 
						|
  LogLevel = enum quiet, normal, verbose, very_verbose
 | 
						|
  DbMigrateConfig* = tuple[ driver, sqlDir, connectionString: string, logLevel: LogLevel ]
 | 
						|
 | 
						|
proc createMigrationsTable(conn: DbConn): void =
 | 
						|
  conn.exec(sql("""
 | 
						|
CREATE TABLE IF NOT EXISTS migrations (
 | 
						|
  id SERIAL PRIMARY KEY,
 | 
						|
  name VARCHAR NOT NULL,
 | 
						|
  run_at TIMESTAMP NOT NULL DEFAULT NOW())"""))
 | 
						|
 | 
						|
proc loadConfig*(filename: string): DbMigrateConfig =
 | 
						|
  ## Load DbMigrateConfig from a file. 
 | 
						|
  let cfg = json.parseFile(filename)
 | 
						|
 | 
						|
  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)
 | 
						|
 | 
						|
proc createMigration*(config: DbMigrateConfig, migrationName: string): seq[string] =
 | 
						|
  ## Create a new set of database migration files.
 | 
						|
  let timestamp = getTime().getLocalTime().format("yyyyMMddHHmmss")
 | 
						|
  let filenamePrefix =  timestamp & "-" & migrationName
 | 
						|
 | 
						|
  let upFilename = joinPath(config.sqlDir, filenamePrefix & "-up.sql")
 | 
						|
  let downFilename = joinPath(config.sqlDir, filenamePrefix & "-down.sql")
 | 
						|
 | 
						|
  let scriptDesc = migrationName & " (" & timestamp & ")"
 | 
						|
 | 
						|
  let upFile = open(upFilename, fmWrite)
 | 
						|
  let downFile = open(downFilename, fmWrite)
 | 
						|
 | 
						|
  upFile.writeLine "-- UP script for " & scriptDesc
 | 
						|
  downFile.writeLine "-- DOWN script for " & scriptDesc
 | 
						|
 | 
						|
  upFile.close()
 | 
						|
  downFile.close()
 | 
						|
 | 
						|
  return @[upFilename, downFilename]
 | 
						|
 | 
						|
proc diffMigrations*(pgConn: DbConn, config: DbMigrateConfig):
 | 
						|
  tuple[ run, notRun, missing: seq[string] ] =
 | 
						|
 | 
						|
  # Query the database to find out what migrations have been run.
 | 
						|
  var migrationsRun = initSet[string]()
 | 
						|
 | 
						|
  for row in pgConn.fastRows(sql"SELECT * FROM migrations ORDER BY name", @[]):
 | 
						|
    migrationsRun.incl(row[1])
 | 
						|
 
 | 
						|
  # Inspect the filesystem to see what migrations are available.
 | 
						|
  var migrationsAvailable = initSet[string]()
 | 
						|
  for filePath in walkFiles(joinPath(config.sqlDir, "*.sql")):
 | 
						|
    var migrationName = filePath.extractFilename
 | 
						|
    migrationName.removeSuffix("-up.sql")
 | 
						|
    migrationName.removeSuffix("-down.sql")
 | 
						|
    migrationsAvailable.incl(migrationName)
 | 
						|
 | 
						|
  # Diff with the list of migrations that we have in our migrations
 | 
						|
  # directory.
 | 
						|
  let migrationsInOrder =
 | 
						|
    toSeq(migrationsAvailable.items).sorted(system.cmp)
 | 
						|
 | 
						|
  var migrationsNotRun: seq[string] = @[]
 | 
						|
  var missingMigrations: seq[string] = @[]
 | 
						|
 | 
						|
  for migration in migrationsInOrder:
 | 
						|
    if not migrationsRun.contains(migration):
 | 
						|
      migrationsNotRun.add(migration)
 | 
						|
 | 
						|
    # if we've already seen some migrations that have not been run, but this
 | 
						|
    # one has been, that means we have a gap and are missing migrations
 | 
						|
    elif migrationsNotRun.len > 0:
 | 
						|
      missingMigrations.add(migrationsNotRun)
 | 
						|
      migrationsNotRun = @[]
 | 
						|
 
 | 
						|
  return (run: toSeq(migrationsRun.items), notRun: migrationsNotRun, missing: missingMigrations)
 | 
						|
 | 
						|
proc up*(config: DbMigrateConfig, count: int): seq[string] =
 | 
						|
  let pgConn = open("", "", "", config.connectionString)
 | 
						|
  pgConn.createMigrationsTable
 | 
						|
 | 
						|
  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:
 | 
						|
    # 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:
 | 
						|
      let filename = joinPath(config.sqlDir, migration & "-up.sql")
 | 
						|
 | 
						|
      if not filename.fileExists:
 | 
						|
        rollbackWithErr "Can not find up file for " & migration &
 | 
						|
          ". Expected '" & filename & "'."
 | 
						|
 | 
						|
      let migrationSql = filename.readFile
 | 
						|
      try:
 | 
						|
        pgConn.exec(sql(migrationSql))
 | 
						|
        migrationsRun.add(migration)
 | 
						|
      except DbError:
 | 
						|
        rollbackWithErr "Migration '" & migration & "'failed:\n\t" &
 | 
						|
          getCurrentExceptionMsg()
 | 
						|
 | 
						|
 | 
						|
  return migrationsRun
 | 
						|
    
 | 
						|
when isMainModule:
 | 
						|
  let doc = """
 | 
						|
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)
 | 
						|
 | 
						|
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).
 | 
						|
 | 
						|
  -V --version                Print the tools version information.
 | 
						|
"""
 | 
						|
 | 
						|
  # Parse arguments
 | 
						|
  let args = docopt(doc, version = "db-migrate 0.2.0")
 | 
						|
 | 
						|
  let exitErr = proc(msg: string): void =
 | 
						|
    stderr.writeLine("db_migrate: " & msg)
 | 
						|
    quit(QuitFailure)
 | 
						|
 | 
						|
  # Load configuration file
 | 
						|
  let configFilename =
 | 
						|
    if args["--config"]: $args["<config-file>"]
 | 
						|
    else: "database.json"
 | 
						|
 | 
						|
  var config: DbMigrateConfig
 | 
						|
  try:
 | 
						|
    config = loadConfig(configFilename)
 | 
						|
  except IOError:
 | 
						|
    exitErr "Cannot open config file: " & configFilename
 | 
						|
  except:
 | 
						|
    exitErr "Error parsing config file: " & configFilename & "\L\t" & getCurrentExceptionMsg()
 | 
						|
 | 
						|
  # Check for migrations directory
 | 
						|
  if not existsDir config.sqlDir:
 | 
						|
    try:
 | 
						|
      echo "SQL directory '" & config.sqlDir &
 | 
						|
        "' does not exist and will be created."
 | 
						|
      createDir config.sqlDir
 | 
						|
    except IOError:
 | 
						|
      exitErr "Unable to create directory: " & config.sqlDir & ":\L\T" & getCurrentExceptionMsg()
 | 
						|
 | 
						|
  # Execute commands
 | 
						|
  if args["create"]:
 | 
						|
    try:
 | 
						|
      let filesCreated = createMigration(config, $args["<migration-name>"])
 | 
						|
      echo "Created new migration files:"
 | 
						|
      for filename in filesCreated: echo "\t" & filename
 | 
						|
    except IOError:
 | 
						|
      exitErr "Unable to create migration scripts: " & getCurrentExceptionMsg()
 | 
						|
 | 
						|
  elif 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."
 | 
						|
    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["init"]: discard
 | 
						|
 |