Apply incremental updates to a PostgreSQL database.

update-schema.rb

#!/usr/bin/ruby

require 'optparse'

options = {}
OptionParser.new do |opts|
  opts.banner = "Usage: update-schema [options]"

  opts.on("-h", "--host HOSTNAME", "Database host to connect to") do |value|
    options[:host] = value
  end
  opts.on("-d", "--database DBNAME", "Name of database to connect to") do |value|
    options[:database] = value
  end
  opts.on("-u", "--user USERNAME", "Username to connect to database as") do |value|
    options[:user] = value
  end
  opts.on("-p", "--print", "Print current schema version") do |q|
    options[:print] = q
  end
  opts.on("-n", "--new NAME", "Create a new version") do |name|
    options[:new] = name
  end
  opts.on("-a", "--apply VERSION", "Apply specified version") do |v|
    options[:apply] = v
  end
  opts.on("-c", "--cat", "Concatenates all updates to standard output") do |c|
    options[:cat] = c
  end
  opts.on("-r", "--reset VERSION", "Reset current schema version to specified version") do |r|
    options[:reset] = r
  end
  opts.on("-v", "--[no-]verbose", "Run verbosely") do |v|
    options[:verbose] = v
  end
end.parse!

class UnknownSchemaError < Exception
end

class SchemaUpgradeError < Exception
end

class Database

  def host=(h)
    @host = h
  end

  def user=(u)
    @user = u
  end

  def database=(d)
    @database = d
  end

  def psql_cmd
    u = (@user.nil? ? "" : "-U #{@user}")
    h = (@host.nil? ? "" : "-h #{@host}")
    d = (@database.nil? ? "" : "#{@database}")
    "psql #{u} #{h} #{d}"
  end

  def can_connect?
    begin
      system("#{psql_cmd} -c 'select 1'")
    rescue Exception => e
      puts e
      false
    end
  end

  def read_version
    begin
      v = `#{psql_cmd} -t -c 'select version from schema_version'`.chomp
      raise UnknownSchemaError unless v != ''
      v.to_i
    rescue Exception
      raise UnknownSchemaError
    end
  end

  def available_versions
    files = Dir.glob("*.sql")
    files.collect {|file| file.split('_')[0].to_i }.sort.uniq
  end

  def needed_versions
    v = 0
    begin
      v = read_version
    rescue UnknownSchemaError
      v = 0
    end

    puts "Schema version is: #{v}".blue
    available = available_versions
    needed = available.select {|avail| avail > v }
    puts "Available updates: #{needed.join(',')}".blue unless needed.empty?
    needed
  end

  def upgrade_to(version)
    apply_version(version)
    set_current_version(version)
  end

  def apply_version(version)
    puts "Applying version: #{version}".blue
    filenames = version_filenames(version)
    filenames.each do |f|
      puts "Running: #{f}".blue
      raise SchemaUpgradeError unless system("#{psql_cmd} -1 --set ON_ERROR_STOP=1 -f '#{f}'")
    end
  end

  def set_current_version(version)
    begin
      v = read_version
    rescue UnknownSchemaError
      puts "Cannot read current DB version; creating schema_version table.".blue
      `#{psql_cmd} -c 'create table schema_version (version text)'`
      `#{psql_cmd} -c 'insert into schema_version values (0)'`
    end
    `#{psql_cmd} -c 'update schema_version set version=#{version}'`
  end

  def upgrade_schema
    needed = needed_versions
    v = 0
    begin
      needed.each {|v| upgrade_to(v) }
      puts
      puts "Your schema is up-to-date.".green
    rescue SchemaUpgradeError
      puts "Error applying version #{v}. Aborting.".red
      raise
    end
  end

  def version_filenames(version)
    names = Dir.glob("*.sql").select {|file| file.split('_')[0].to_i == version}.sort
    names = ["#{version}.sql"] if names.empty?
    names
  end

  def cat_version(version)
    version_filenames(version).each do |f|
      puts "-- File begins: #{f}"
      puts
      system("cat '#{f}'")
    end
  end

  def create_version(name)
    last = available_versions[-1]
    new = last + 1
    filename = "#{new}_#{name}.sql"
    system("touch #{filename}")
    puts "Your new version is: #{filename}".green
  end

  def reset(version = 0)
    set_current_version(version)
  end
end

class String
  def red; colorize(self, "\e[1m\e[31m"); end
  def green; colorize(self, "\e[1m\e[32m"); end
  def dark_green; colorize(self, "\e[32m"); end
  def yellow; colorize(self, "\e[1m\e[33m"); end
  def blue; colorize(self, "\e[1m\e[34m"); end
  def dark_blue; colorize(self, "\e[34m"); end
  def pur; colorize(self, "\e[1m\e[35m"); end
  def colorize(text, color_code)  "#{color_code}#{text}\e[0m" end
end

Dir.chdir("schema")
db = Database.new
db.database = options[:database]
db.host = options[:host]
db.user = options[:user]
if options[:cat]
  db.available_versions.each do |v|
    puts "-- Version #{v}"
    puts
    db.cat_version(v)
    puts
  end
elsif options[:print]
  needed = db.needed_versions
  if needed.empty?
    puts "Your schema is up-to-date.".green
  else
    needed.each do |v|
      puts " + #{db.version_filenames(v).join(' ')}"
    end
  end
elsif options[:new]
  db.create_version(options[:new])
elsif options[:apply]
  db.apply_version(options[:apply].to_i)
elsif options[:reset]
  v = options[:reset].to_i
  db.reset(v)
  puts "Current schema version set to #{v}".green
else
  unless db.can_connect?
    puts "Can't connect!".red
    exit 1 
  end
  db.upgrade_schema
end