#!/usr/bin/perl -wT

# pg_restorecluster: restore from a pg_backupcluster backup
#
# Copyright (C) 2021 Christoph Berg <myon@debian.org>
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.

use strict;
use warnings;
use Cwd 'abs_path';
use Getopt::Long;
use PgCommon;

my ($version, $cluster);

# untaint environment
$ENV{'PATH'} = '/sbin:/bin:/usr/sbin:/usr/bin';
delete @ENV{'IFS', 'CDPATH', 'ENV', 'BASH_ENV'};
umask 022;

sub help () {
    print "Syntax: $0 [options] <version> <cluster> <backup>
Options:
  -d --datadir DIR  Data directory for restored cluster (default per createcluster.conf)
  -p --port PORT    Use port PORT for restored cluster (default is next free port)
  -s --start        Start cluster after restoring (default for restore from dump)
     --archive      Configure recovery from WAL archive
     --pitr TIMEST  Configure point-in-time recovery to TIMESTAMP from WAL archive
     --wal-archive DIR  Read WAL from archive DIR (default <backup>/../wal)
";
}

my $createclusterconf = "$PgCommon::common_confdir/createcluster.conf";
my ($datadir, $port, $start, $recovery_target_time, $archive_recovery, $wal_archive);

exit 1 unless GetOptions (
    'd|datadir=s' => \$datadir,
    'p|port=s' => \$port,
    's|start' => \$start,
    'archive' => \$archive_recovery,
    'pitr=s' => \$recovery_target_time,
    'recovery-target-time=s' => \$recovery_target_time,
    'wal-archive=s' => \$wal_archive,
    'createclusterconf=s' => \$createclusterconf,
);
if ($recovery_target_time) {
    ($recovery_target_time) = $recovery_target_time =~ /(.*)/; # untaint
}

# accept both "version cluster" and "version[-/]cluster"
if (@ARGV >= 2 and $ARGV[0] =~ m!^(\d+\.?\d)[-/]([^/]+)$!) {
    ($version, $cluster) = ($1, $2);
    shift @ARGV;
} elsif (@ARGV >= 3 and $ARGV[0] =~ /^(\d+\.?\d)$/) {
    $version = $1;
    ($cluster) = ($ARGV[1]) =~ m!^([^/]+)$!;
    shift @ARGV;
    shift @ARGV;
} else {
    help();
    exit 1;
}

error "cluster $version $cluster already exists" if cluster_exists $version, $cluster;

my %defaultconf = PgCommon::read_conf_file ($createclusterconf);

# functions to be run as root

sub create_cluster_directories($$$$) {
    my ($owneruid, $ownergid, $configdir, $datadir) = @_;
    my @postgres_user = getpwnam 'postgres';
    my ($pg_uid, $pg_gid) = (getpwnam 'postgres')[2, 3];

    for my $pgdir ("/etc/postgresql", "/etc/postgresql/$version", "/var/lib/postgresql", "/var/lib/postgresql/$version") {
        if (! -e $pgdir) {
            mkdir $pgdir or error "mkdir $pgdir: $!";
            chown $pg_uid, $pg_gid, $pgdir or error "chown $pgdir: $!";
        }
    }
    mkdir $configdir or error "mkdir $configdir: $!";
    chown $owneruid, $ownergid, $configdir or error "chown $configdir: $!";
    mkdir $datadir, 0700 or error "mkdir $datadir: $!";
    chown $owneruid, $ownergid, $datadir or error "chown $datadir: $!";
}

sub create_cluster($$$$$) {
    my ($backup, $owneruid, $ownergid, $configdir, $datadir) = @_;

    my @createclusteropts = ();
    if (-f "$backup/createcluster.opts") {
        open my $fh, "$backup/createcluster.opts" or error "$backup/createcluster.opts: $!";
        local $/; # slurp mode
        my ($opts) = <$fh> =~ /(.*)/; # untaint
        @createclusteropts = split /\s+/, $opts;
        close $fh;
    }

    system_or_error "pg_createcluster",
        "--datadir", $datadir,
        "--user", $owneruid, "--group", $ownergid,
        $version, $cluster,
        @createclusteropts;
}

sub start_cluster() {
    print "Starting cluster $version $cluster ...\n";
    system_or_error "pg_ctlcluster", $version, $cluster, "start";
}

sub switch_to_cluster_owner($$) {
    my ($owneruid, $ownergid) = @_;
    change_ugid $owneruid, $ownergid;
}

# restore functions

sub unpack_tar($$$) {
    my ($backup, $tar, $dir) = @_;

    if (-f "$backup/$tar.gz") {
        $tar = "$tar.gz";
    } elsif (-f "$backup/$tar.bz2") {
        $tar = "$tar.bz2";
    } elsif (-f "$backup/$tar.xz") {
        $tar = "$tar.xz";
    } elsif (-f "$backup/$tar") {
        # do nothing
    } else {
        error "$backup/config.tar* is missing";
    }

    print "Restoring $backup/$tar to $dir ...\n";
    system_or_error "tar", "-C", $dir, "-xf", "$backup/$tar";
}

sub restore_config($$) {
    my ($backup, $configdir) = @_;
    unpack_tar($backup, "config.tar", $configdir);
}

sub update_config($$$) {
    my ($configdir, $datadir, $port) = @_;
    my %settings = (
        data_directory => $datadir,
        hba_file => "$configdir/pg_hba.conf",
        ident_file => "$configdir/pg_ident.conf",
        external_pid_file => "/var/run/postgresql/$version-$cluster.pid",
        port => $port,
    );
    $settings{cluster_name} = "$version/$cluster" if ($version >= 9.5);
    $settings{stats_temp_directory} = "/var/run/postgresql/$version-$cluster.pg_stat_tmp" if ($version < 15);
    my %config = PgCommon::read_conf_file("$configdir/postgresql.conf");
    for my $guc (sort keys %settings) {
        if (not exists $config{$guc} or $config{$guc} ne $settings{$guc}) {
            print "Setting $guc = $settings{$guc}\n";
            PgCommon::set_conffile_value("$configdir/postgresql.conf", $guc, $settings{$guc});
        }
    }
}

sub restore_basebackup($$) {
    my ($backup, $datadir) = @_;
    unpack_tar($backup, "base.tar", $datadir);
}

sub restore_wal($$) {
    my ($backup, $datadir) = @_;
    return if ($version < 10); # WAL contained in base.tar.gz in PG 9.x
    unpack_tar($backup, "pg_wal.tar", "$datadir/pg_wal");
}

sub archive_recovery_options($$$) {
    my ($backup, $datadir, $wal_archive) = @_;

    $wal_archive = abs_path($wal_archive) or error "$wal_archive: $!";
    -d $wal_archive or error "$wal_archive is not a directory";

    print "Setting archive recovery options";
    my $recovery_options = "restore_command = '/usr/share/postgresql-common/pg_getwal $wal_archive/%f %p'\n";
    if ($recovery_target_time) {
        $recovery_options .= "recovery_target_time = '$recovery_target_time'\n";
        $recovery_options .= "recovery_target_action = 'promote'\n";
    }

    if ($version >= 12) {
        my $autoconf = "$datadir/postgresql.auto.conf";
        open my $fh, ">>", $autoconf or error "$autoconf: $!";
        print $fh $recovery_options or error "$autoconf: $!";
        close $fh or error "$autoconf: $!";

        my $recoverysignal = "$datadir/recovery.signal";
        open my $fh2, ">", $recoverysignal or error "$recoverysignal: $!";
        close $fh2 or error "$recoverysignal: $!";

    } else {
        my $recoveryconf = "$datadir/recovery.conf";
        open my $fh, ">>", $recoveryconf or error "$recoveryconf: $!";
        print $fh $recovery_options or error "$recoveryconf: $!";
        close $fh or error "$recoveryconf: $!";
    }
}

sub reset_archive_recovery_options() {
    if ($version >= 12) {
        system_or_error "psql",
            "--cluster", "$version/$cluster",
            "-XAtqc", "ALTER SYSTEM RESET restore_command";
        system_or_error "psql",
            "--cluster", "$version/$cluster",
            "-XAtqc", "ALTER SYSTEM RESET recovery_target_time"
                if ($recovery_target_time);
        system_or_error "psql",
            "--cluster", "$version/$cluster",
            "-XAtqc", "ALTER SYSTEM RESET recovery_target_action"
                if ($recovery_target_time);
    }
}

sub restore_globals($$) {
    my ($backup, $owneruid) = @_;
    my $owner = (getpwuid $owneruid)[0] or error "UID $owneruid has no name";

    print "Restoring $backup/globals.sql ...\n";
    open my $globals, "$backup/globals.sql" or error "$backup/globals.sql: $!";
    open my $psql, "|-", "psql", "--cluster", "$version/$cluster", "-vON_ERROR_STOP=1", "-Xq" or error "psql: $!";
    while (my $line = <$globals>) {
        next if ($line eq "CREATE ROLE $owner;\n");
        print $psql $line or error "psql: $!";
    }
    close $globals;
    close $psql;
    error "psql failed" if ($?);
}

sub create_databases($) {
    my ($backup) = @_;

    print "Creating databases from $backup/databases.sql ...\n";
    system_or_error "psql", "--cluster", "$version/$cluster", "-vON_ERROR_STOP=1", "-Xqf", "$backup/databases.sql";
}

sub restore_dumps($) {
    my ($backup) = @_;

    for my $dump (sort glob "$backup/*.dump") {
        $dump =~ m!(.*/([^/]*).dump)$!;
        $dump = $1; # untaint
        my $db = $2;
        print "Restoring $dump to database $db ...\n";
        system_or_error "pg_restore", "--cluster", "$version/$cluster", "-d", $db, $dump;
    }
}

sub wait_for_recovery() {
    my $sleep = 1;

    print "Waiting for end of recovery ...\n";
    while (1) {
        open my $psql, "-|", "psql", "--cluster", "$version/$cluster", "-XAtc", "SELECT pg_is_in_recovery()" or error "psql: $!";
        my $status = <$psql>;
        error "psql: $!" unless (defined $status);
        close $psql or error "psql: $!";
        last if ($status eq "f\n");
        sleep($sleep++);
    }
}

sub analyze() {
    system_or_error "vacuumdb",
        "--cluster", "$version/$cluster",
        "--analyze-only",
        ($version >= 9.4 ? "--analyze-in-stages" : ()),
        "--all";
}

sub lscluster() {
    system_or_error "pg_lsclusters", $version, $cluster;
}

# main

my ($backup) = $ARGV[0] =~ /(.*)/; # untaint
error "$backup is not a directory" unless (-d $backup);
$backup =~ s/\/$//; # strip trailing slash
my ($owneruid, $ownergid) = (stat $backup)[4, 5];
my $configdir = "/etc/postgresql/$version/$cluster";
$datadir //= replace_v_c($defaultconf{data_directory} // "/var/lib/postgresql/%v/%c", $version, $cluster);
($datadir) = $datadir =~ /(.*)/; # untaint
$wal_archive //= "$backup/../wal";
$port //= next_free_port();

if ($backup =~ /\.backup$/) {
    create_cluster_directories($owneruid, $ownergid, $configdir, $datadir);
    if (fork == 0) {
        switch_to_cluster_owner($owneruid, $ownergid);
        restore_config($backup, $configdir);
        update_config($configdir, $datadir, $port);
        restore_basebackup($backup, $datadir);
        restore_wal($backup, $datadir);
        archive_recovery_options($backup, $datadir, $wal_archive)
            if ($archive_recovery or $recovery_target_time);
        exit(0);
    }
    wait;
    exit(1) if ($?);
    if ($start) {
        print "\n";
        start_cluster();
        switch_to_cluster_owner($owneruid, $ownergid);
        wait_for_recovery();
        reset_archive_recovery_options() if ($archive_recovery or $recovery_target_time);
        analyze();
    }
    print "\n";
    lscluster();

} elsif ($backup =~ /\.dump$/) {
    create_cluster($backup, $owneruid, $ownergid, $configdir, $datadir);
    print "\n";
    if (fork == 0) {
        switch_to_cluster_owner($owneruid, $ownergid);
        restore_config($backup, $configdir);
        update_config($configdir, $datadir, $port);
        exit(0);
    }
    wait;
    exit(1) if ($?);
    start_cluster();
    switch_to_cluster_owner($owneruid, $ownergid);
    restore_globals($backup, $owneruid);
    create_databases($backup);
    restore_dumps($backup);
    analyze();
    print "\n";
    lscluster();

} else {
    error "$backup must end in either .backup or .dump";
}

__END__

=head1 NAME

pg_restorecluster - Restore from a pg_backupcluster backup

=head1 SYNOPSIS

B<pg_restorecluster> [I<options>] I<version> I<cluster> I<backup>

=head1 DESCRIPTION

B<pg_restorecluster> restores a PostgreSQL cluster from a backup created by
B<pg_backupcluster>. The cluster will be newly created in the system using the
name provided on the command line; this allows renaming a cluster on restore.
The restored cluster configuration will be updated to reflect the new name and
location.

The I<backup> name passed must end in either B<.basebackup> or B<.dump>;
usually this will be the full path to a backup directory in
C</var/backups/postgresql/version-cluster/> as reported by
B<pg_backupcluster ... list>.

Basebackups are restored as-is. For dumps, B<pg_createcluster> is used to
create a new cluster, and schema and data are restored via B<pg_restore>.

=head1 OPTIONS

=over 4

=item B<-d --datadir> I<DIR>

Use I<DIR> as data directory for the restored cluster (default per
createcluster.conf, by default /var/lib/postgresql/I<version>/I<cluster>).

=item B<-p --port> I<N>

Use port I<N> for restored cluster (default is next free port).

=item B<-s --start>

Start cluster after restoring (default for restore from dump; off for
basebackup restores).

After the cluster has been started, B<ANALYZE> is run on all databases.

=item B<--archive>

Configure cluster for recovery from WAL archive. This sets B<restore_command>
to retrieve WAL files from I<backup>B</../wal>.

=item B<--pitr> I<TIMESTAMP>

=item B<--recovery-target-time> I<TIMESTAMP>

Additionally to setting B<restore_command>, set B<recovery_target_time> to
I<TIMESTAMP> for point-in-time recovery. Also sets
B<recovery_target_action='promote'>.

=item B<--wal-archive> I<DIR>

For archive recovery, read WAL from archive I<DIR> (default is
I<backup>B</../wal>).

=back

=head1 FILES

=over 4

=item /var/backups

Default root directory for cluster backup directories.

=back

See L<pg_backupcluster(1)> for a description of files.

=head1 SEE ALSO

L<pg_backupcluster(1)>, L<pg_restore(1)>, L<vacuumdb(1)>.

=head1 AUTHOR

Christoph Berg L<E<lt>myon@debian.orgE<gt>>