My Swiss-Army Script, Redux

prt.pl

#!/usr/bin/env perl

use warnings;
use strict;

my ($home, $psql, $pg_ctl, $pg_dump, $code_root, $mig_dir, $replication,
    $localdb, $rvm, $rubyver, $gemset_name, $gemset, $ruby_bin, $padrino,
    $rep_port, $local_port);
my ($SUPER, $OWNER, $PRODDB, $DEVDB, $TESTDB, $TESTTMPLDB);

$home        = $ENV{HOME}         || '~';
$psql        = $ENV{PSQL};
$pg_ctl      = $ENV{PG_CTL};
$pg_dump     = $ENV{PG_DUMP};
$code_root   = $ENV{CODEROOT}     || "$home/Sources/project";
$mig_dir     = $ENV{MIG_DIR}      || "$code_root/MIGRATIONS";
$replication = $ENV{REPLICATION}  || "$home/.replication";
$localdb     = $ENV{LOCALDB}      || "$home/.live";
$rvm         = $ENV{RVM_ROOT}     || "$home/.rvm";
$rubyver     = $ENV{RUBY_VERSION} || "ruby-1.9.3-p385";
$gemset_name = $ENV{GEMSET_NAME}  || "project";
$gemset      = $ENV{GEMSET}       || "$rubyver\@$gemset_name";
$ruby_bin    = $ENV{RUBY_BIN}     || "$rvm/wrappers/$gemset/ruby";
$padrino     = $ENV{PADDRINO_BIN} || "$rvm/gems/$gemset/bin/padrino";
$rep_port    = $ENV{REP_PORT}     || 5433;
$local_port  = $ENV{LOCAL_PORT}   || 5432;
$SUPER       = $ENV{SUPER}        || "pgsql";
$OWNER       = $ENV{OWNER}        || "owner";
$PRODDB      = $ENV{PRODDB}       || "prod";
$DEVDB       = $ENV{DEVDB}        || "dev";
$TESTDB      = $ENV{TESTDB}       || "test";
$TESTTMPLDB  = $ENV{TESTTMPLDB}   || "test-template";

unless ($psql) {
    $psql = `which psql`;
    chomp $psql;
}

unless ($pg_ctl) {
    $pg_ctl = `which pg_ctl`;
    chomp $pg_ctl;
}

unless ($pg_dump) {
    $pg_dump = `which pg_dump`;
    chomp $pg_dump;
}

sub cmd($;$\$);
sub strip(\$);
sub quote_path($);
sub replication_is_running();
sub localdb_is_running();
sub shutdown_replication();
sub shutdown_localdb();
sub wipe_localdb();
sub copy_replication();
sub sub_configs();
sub startup_replication();
sub startup_localdb();
sub rebuild_dev();
sub rebuild_test();
sub migrate($;$);
sub current_version($;$);
sub register_command($$&;$);
sub commands();
sub commands_with_args();
sub help_for($);
sub arg_help_for($);
sub run_command(@);
sub usage();

# here's where the logic lives
register_commands: {
    register_command usage => "Display this help", sub {
        print "Usage: $0 [command] [args]\n\n";

        print "Where command is one of:\n";
        foreach my $cmd (commands) {
            printf("  %10s   -- %s\n", $cmd, help_for($cmd));
        }

        foreach my $cmd (commands_with_args) {
            printf("\n  %10s accepts arguments:\n    %s\n", $cmd,
                   arg_help_for($cmd));
        }
    };

    register_command startall => "Start all DB servers", sub {
        startup_replication;
        startup_localdb;
    };

    register_command start => "Start specific server", sub {
        my ($serv) = @_;
        die unless $serv;

        if (($serv eq 'replication') || ($serv =~ 'r')) {
            if (replication_is_running) {
                print "Replication is already running.\n";
            }
            else {
                startup_replication;
            }
        }
        elsif (($serv eq 'live') || ($serv =~ 'l')) {
            if (localdb_is_running) {
                print "Local live DB is already running.\n";
            }
            else {
                startup_localdb;
            }
        }
        elsif (($serv eq 'web') || ($serv =~ 'w')) {
            my ($path) = quote_path $code_root;
            my ($cmd)  = "(cd $path && padrino start)";

            system($cmd);
            exit(0);
        }
        else {
            print "Unknown server.\n";
            die;
        }
    }, "Argument can be either 'replication'|'r' or 'live'|'l'";

    register_command term => "Connect to a server via cmdline", sub {
        local ($_);

        my ($serv, $db) = @_;
        my ($path) = quote_path($psql);
        my ($port, $name, $cmd);
        die unless ($db and $serv);

        $port = do {
            if    ($serv =~ /^r(ep)?$/ ) { $rep_port   }
            elsif ($serv =~ /^l(ive)?$/) { $local_port }
            else { die "Unknown server $serv" }
        };

        $name = do {
            if    ($db eq 'prod') { $PRODDB }
            elsif ($db eq 'dev' ) { $DEVDB  }
            else { die "Unknown database $db" }
        };

        $cmd = "$path -p $port -d $name -U $SUPER";
        system($cmd);
        exit(0);
    }, "1st arg is 'rep'|'r' or 'live'|'l', second arg is 'prod' or 'dev'";

    register_command wd => "Startup the web development server", sub {
        my ($dir)  = quote_path($code_root);
        my ($rbin) = quote_path($ruby_bin);
        my ($pbin) = quote_path($padrino);
        my ($cmd) = "(cd $dir && $rbin $pbin start)";

        print "Starting web server, switching off control.\n";
        system($cmd);
        exit(0);
    };

    register_command taillog => "Tail the logs of a particular server", sub {
        my ($serv) = @_;
        my ($path, $cmd);
        die unless $serv;

        if (($serv eq 'replication') || ($serv =~ 'r')) {
            $path = quote_path "$replication/pg_log/postgresql.log";
        }
        elsif (($serv eq 'live') || ($serv =~ 'l')) {
            $path = quote_path "$localdb/pg_log/postgresql.log";
        }
        else {
            print "Unknown server.\n";
            die;
        }

        $cmd  = "tail -f $path";
        system($cmd);
        exit(0);
    }, "Argument can be either 'replication'|'r' or 'live'|'l'";

    register_command backup => "Backup a running server", sub {
        my ($serv, $db) = @_;
        my ($port, $date, $backpath, $path, $cmd, @date);

        if ($serv =~ /^r(ep)?$/) {
            unless (replication_is_running) {
                print "Replication server is not running.\n";
                print "  Hint: Start it with: prt start rep\n";
                exit(1);
            }

            $port = $rep_port;
        }
        elsif ($serv =~ /^l(ive)?$/) {
            unless (localdb_is_running) {
                print "Live DB copy is not running.\n";
                print "  Hint: Start it with: prt start live\n";
                exit(1);
            }

            $path = $local_port;
        }
        else {
            print "Unknown server.\n";
            die;
        }

        if ($db eq 'prod') {
            $db = $PRODDB;
        }
        elsif ($db eq 'dev') {
            $db = $DEVDB;
        }
        else {
            print "Unknown database.\n";
            die;
        }

        @date = localtime;
        $date = sprintf('%04d-%02d-%02d', $date[5] + 1900, $date[4],
                        $date[3]);

        $backpath = quote_path "./db_backup_$date.sql";
        $path     = quote_path $pg_dump;
        $cmd      = "$path --file=$backpath --format=p --superuser=$SUPER " .
                    "--quote-all-identifiers --username=$SUPER --port=$port ".
                    $db;

        cmd $cmd, "Backing up database";
    }, "1st arg is 'rep'|'r' or 'live'|'l', second arg is 'prod' or 'dev'";

    register_command swap => "Swap running server", sub {
        if (replication_is_running) {
            shutdown_replication;
            startup_localdb;
        }
        elsif (localdb_is_running) {
            shutdown_localdb;
            startup_replication;
        }
        else {
            print "No servers are currently running to swap.\n";
        }
    };

    register_command stopall => "Stop all DB servers", sub {
        my ($stopping) = 0;

        if (replication_is_running) {
            shutdown_replication;
            $stopping = 1;
        }
        if (localdb_is_running) {
            shutdown_localdb;
            $stopping = 1;
        }

        unless ($stopping) {
            print "All servers are already stopped.\n";
        }
    };

    register_command status => "Check the status of the servers", sub {
        my ($rep) = replication_is_running ? 'running' : 'not running';
        my ($loc) = localdb_is_running     ? 'running' : 'not running';

        print "Replication is $rep\n";
        print "Local live DB copy is $loc\n";
    };

    register_command reclone => "Clone local live DB from replication", sub {
        my ($loc) = quote_path $localdb;
        my ($cmd) = "sleep 3; ls $loc/recovery.done";

        run_command 'stopall';
        wipe_localdb;
        copy_replication;
        sub_configs;
        run_command 'start', 'l';

        cmd $cmd, "Verifying local DB copy is fully live";
    };

    register_command builddev => "Rebuild dev DB from prod", sub {
        rebuild_dev
    };

    register_command buildtest => "Rebuild test DB from test template", sub {
        rebuild_test
    };

    register_command version => "Display current database versions", sub {
        my ($running) = 0;

        if (replication_is_running) {
            print "Replicated development schema is at " .
                  current_version('dev', 'replication') . "\n";
            print "Replicated production schema is at " .
                  current_version('prod', 'replication') . "\n";

            $running = 1;
        }

        if (localdb_is_running) {
            print "Local DB development schema is at " .
                  current_version('dev', 'live') . "\n";
            print "Local DB production schema is at " .
                  current_version('prod', 'live') . "\n";

            $running = 1;
        }
    };

    register_command migrate => "Upgrade DB using SQL migrations", sub {
        my ($db, $v) = @_;
        die unless ($db);

        unless (localdb_is_running) {
            print "Local live DB is not running, cannot migrate.\n";
            print "  HINT, try this first: $0 start live\n";
            exit(0);
        }

        unless ($db =~ /^(p(rod)?|d(ev)?)$/) {
            print "Unknown database.\n";
        }

        migrate $db, $v;
    }, "1st arg is DB ('dev'|'d' or 'prod'|'p'), 2nd arg is optional " .
       "migration #";
}

run_script: {
    unless (run_command @ARGV) {
        run_command 'usage';
        exit(1);
    }

    exit(0);
}

##############################################################################
# everything below this point is just method calls to do the actual work.

sub cmd($;$\$) {
    my ($cmd, $message, $output) = @_;
    print "$message..." if ($message);

    eval {
        # print "Running '$cmd'";
        if ($output) {
            $$output = `($cmd) 2>/dev/null`;
        }
        else {
            `($cmd) 2>&1 1>/dev/null`;
        }
    };
    if ($?) {
        print "failure!\n" if $message;
        0
    }
    else {
        print "success!\n" if $message;
        1
    }
}

sub strip(\$) {
    my ($strip) = @_;
    $$strip =~ s/^\s*(.*?)\s*$/$1/;
    $$strip;
}

sub quote_path($) {
    my ($path) = shift;
    $path =~ s/([^a-z0-9_\-\.\\\/])/\\$1/ig;
    $path;
}

sub replication_is_running() {
    my ($rep) = quote_path $replication;
    my ($cmd);

    $cmd = "ls $rep/postmaster.pid";
    return 0 unless (cmd $cmd);

    $cmd = "kill -0 \$(head -n1 $rep/postmaster.pid)";
    return cmd($cmd);
}

sub localdb_is_running() {
    my ($loc) = quote_path $localdb;
    my ($cmd);

    $cmd = "ls $loc/postmaster.pid";
    return 0 unless (cmd $cmd);

    $cmd = "kill -0 \$(head -n1 $loc/postmaster.pid)";
    return cmd($cmd);
}

sub shutdown_replication() {
    my ($ctl) = quote_path $pg_ctl;
    my ($rep) = quote_path $replication;
    my ($cmd) = "$ctl -D $rep stop";

    cmd $cmd, "Shutting down replication server";
}

sub shutdown_localdb() {
    my ($ctl) = quote_path $pg_ctl;
    my ($loc) = quote_path $localdb;
    my ($cmd) = "$ctl -D $loc stop";

    cmd $cmd, "Shutting down local database copy";
}

sub wipe_localdb() {
    my ($loc) = quote_path $localdb;
    my ($run) = localdb_is_running;
    my ($cmd) = "rm -r $loc";

    if (localdb_is_running) {
        print "Cannot execute: local DB copy is still running\n";
        exit(1);
    }
    else {
        cmd $cmd, "Wiping existing local live DB copy";
    }
}

sub copy_replication() {
    my ($loc) = quote_path $localdb;
    my ($rep) = quote_path $replication;
    my ($cmd) = "cp -rp $rep $loc";

    cmd $cmd, "Copying current replication over to new live DB copy";
}

sub sub_configs() {
    my ($loc) = quote_path $localdb;
    my ($qrep, $qloc);
    my ($cmd);

    $cmd = "ln=\$(grep -n \"port = $rep_port\" $loc/postgresql.conf | cut " .
           "-d\\: -f1); perl -pe \"s/.*/port = $local_port/ if \\\$. == " .
           "\$ln\" < $loc/postgresql.conf > $loc/postgresql.new.conf && mv " .
           "$loc/postgresql.new.conf $loc/postgresql.conf";
    cmd $cmd, "Replacing port number in configuration";

    $qrep = $replication;
    $qloc = $localdb;

    $qrep =~ s/([^a-z0-9_])/\\$1/ig;
    $qloc =~ s/([^a-z0-9_])/\\$1/ig;
    $cmd = "ln=\$(grep -n \"trigger_file\" $loc/recovery.conf | cut -d\\: " .
           "-f1); perl -pe \"s/$qrep/$qloc/ if \\\$. == \$ln\" < " .
           "$loc/recovery.conf > $loc/recovery.new.conf && mv " .
           "$loc/recovery.new.conf $loc/recovery.conf";
    cmd $cmd, "Replacing trigger file in recovery configuration";

    $cmd = "touch $loc/promote_to_master";
    cmd $cmd, "Creating trigger file";
}

sub startup_replication() {
    my ($ctl) = quote_path $pg_ctl;
    my ($rep) = quote_path $replication;
    my ($cmd) = "$ctl -D $rep start; sleep 2; pid=\$(head -n1 " .
                "$rep/postmaster.pid); kill -0 \$pid";

    cmd $cmd, "Starting up replication server";
}

sub startup_localdb() {
    my ($ctl) = quote_path $pg_ctl;
    my ($loc) = quote_path $localdb;
    my ($cmd) = "$ctl -D $loc start; sleep 2; pid=\$(head -n1 " .
                "$loc/postmaster.pid); kill -0 \$pid";

    cmd $cmd, "Starting up local database copy";
}

sub rebuild_dev() {
    my ($prompt) = quote_path $psql;
    my ($sql, $cmd);

    $sql = "DROP DATABASE \\\"$DEVDB\\\"";
    $cmd = "$prompt -d $TESTDB -U $SUPER -c \"$sql\"";
    cmd $cmd, "Deleting old dev database";

    $sql = "CREATE DATABASE \\\"$DEVDB\\\" WITH OWNER \\\"$OWNER\\\" " .
           "TEMPLATE \\\"$PRODDB\\\"";
    $cmd = "$prompt -d $TESTDB -U $SUPER -c \"$sql\"";
    cmd $cmd, "Creating new dev database from prod";
}

sub rebuild_test() {
    my ($prompt) = quote_path $psql;
    my ($sql, $cmd);

    $sql = "DROP DATABASE \\\"$TESTDB\\\"";
    $cmd = "$prompt -d $DEVDB -U $SUPER -c \"$sql\"";
    cmd $cmd, "Deleting old test database";

    $sql = "CREATE DATABASE \\\"$TESTDB\\\" WITH OWNER \\\"$OWNER\\\"" .
           "TEMPLATE \\\"$TESTTMPLDB\\\"";
    $cmd = "$prompt -d $DEVDB -U $SUPER -c \"$sql\"";
    cmd $cmd, "Creating new test database from test template";
}

sub migrate($;$) {
    my ($db, $v) = @_;
    my ($psql)   = quote_path $psql;
    my ($cur)    = current_version $db;
    my ($mig_glob, $latest, $sql, $cmd, @migrations);

    # resolve database name
    if ($db =~ /^p(rod)?$/) {
        $db = $PRODDB;
    }
    elsif ($db =~ /^d(ev)?$/) {
        $db = $DEVDB;
    }
    else {
        die "Unknown database.\n";
    }

    # ensure $v is defined, if only to false
    $v ||= '';

    # if we have a version and it's less than or at our current, bail out.
    if ($v && ($v < $cur)) {
        return print "Current version is already past that migration.\n";
    }
    elsif ($v && ($v == $cur)) {
        return print "No migrations necessary, already at $v\n";
    }

    # grab all potentially valid migrations
    $mig_glob = "$mig_dir/???_*.sql";
    @migrations = sort grep {
        (/\/(\d{3})_[^\/]+\.sql$/) && ($1 > $cur) && ((!$v) || ($1 <= $v))
    } glob($mig_glob);

    # no migrations at all were returned
    unless (@migrations) {
        return print "No migrations to run.\n";
    }

    # check that our latest migration found is at least what we requested,
    # if we requested anything
    ($latest) = $migrations[-1] =~ /\/(\d{3})_[^\/]+\.sql$/;
    if ($v && ($latest < $v)) {
        return print "Migrations are not written up to that version.\n";
    }

    # apply each migration we found in turn, as safely as possible
    foreach my $migration (@migrations) {
        ($latest) = $migration =~ /\/(\d{3})_[^\/]+\.sql$/;
        $latest += 0; # cast to int

        # skip .psqlrc, do in one transaction, stop on first failure, specify
        # object ownership, use correct DB/user and give the migration file.
        # that's a lot of flags.
        $cmd = "$psql -X -1 -v ON_ERROR_STOP=1 -v objects_owner=\"$OWNER\" " .
               "-d $db -U $SUPER -f $migration";
               print "command: $cmd\n";
        unless (cmd $cmd, "Migrating to version $latest") {
            die "Could not migrate past version " . ($latest - 1) . "\n";
        }
        

        # move the DB version up
        $sql = "ALTER DATABASE \\\"$db\\\" SET version.schema = $latest";
        $cmd = "$psql -d $db -U $SUPER -c \"$sql\"";
        unless (cmd $cmd, "Marking database as version $latest") {
            warn "Could not mark database version!\n";
        }
    }
}

sub current_version($;$) {
    my ($db, $serv) = @_;
    my ($prompt)    = quote_path $psql;
    my ($port)      = '';
    my ($sql, $cmd, $v);

    if ($db =~ /^p(rod)?$/) {
        $db = $PRODDB;
    }
    elsif ($db =~ /^d(ev)?$/) {
        $db = $DEVDB;
    }
    else {
        die "Unknown database.\n";
    }

    if ($serv && ($serv =~ /^r(eplication)?$/)) {
        $port = "-p $rep_port";
    }
    elsif ($serv && ($serv =~ /^l(ive)?$/)) {
        $port = "-p $local_port";
    }
    elsif ($serv) {
        die "Unknown server $serv.\n";
    }

    $sql = "SELECT public.schema_version()";
    $cmd = "$prompt -t $port -d $db -U $SUPER -c \"$sql\"";
    cmd $cmd, "Retrieving $db schema version", $v;
    strip $v;
}

command_block: {
    my (%commands);
    my (%help);
    my (%arghelp);

    sub register_command($$&;$) {
        my ($command, $help, $block, $arghelp) = @_;
        $commands{$command} = $block;
        $help{$command}     = $help;

        if ($arghelp) { $arghelp{$command}  = $arghelp }
    }

    sub commands() {
        sort keys %commands;
    }

    sub commands_with_args() {
        sort keys %arghelp;
    }

    sub run_command(@) {
        my ($command, @args) = @_;
        my ($block);
    
        return unless $command;

        $block = $commands{$command};
        return unless $block;

        eval { $block->(@args) };
        # if ($@) {
        #     print ".....Oh noooooooo: $@\n";
        # }
        return $@ ? 0 : 1;
    }

    sub help_for($) { $help{shift()} }
    sub arg_help_for($) { $arghelp{shift()} }
}