#!/bin/bash

# Copyright (c) 2007-2016 credativ GmbH
#
# AUTHORS:
#    Peter Eisentraut <peter.eisentraut@credativ.de>
#    Bernd Helmle <bernd.helmle@credativ.de>
#    Christoph Berg <christoph.berg@credativ.de>
#    Arnd Hannemann <arnd.hannemann@credativ.de>
#    Adrian Vondendriesch <adrian.vondendriesch@credativ.de>
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is furnished to do
# so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE

##
## pg_backup_ctl - PostgreSQL transaction log archival backup control program
##
## Version: 0.8

set -e
umask 077
# This is required so that sorted directory listings are consistent.
export LC_COLLATE=C

me=$(basename $0)
hdrline=$(printf "%080d"|tr "0" "-")

PG_BACKUP_LOCK_TIMEOUT=60
LOCK=""

## If operating on non-Linux, make sure the GNU tools are used
## On Solaris, they're usually available via the 'g' prefix...

if echo "$OSTYPE" | grep darwin > /dev/null || echo "$OSTYPE" | grep solaris > /dev/null;
then
    SED="gsed"
    GREP="ggrep"
    FIND="gfind"
    TAR="gtar"
    TAIL="gtail"
else
    SED="sed"
    GREP="grep"
    FIND="find"
    TAR="tar"
    TAIL="tail"
fi

print_help() {
	cat <<EOF
PostgreSQL transaction log archival backup control program
Supports PostgreSQL 8.3 and above

Usage: $me -A ARCHIVEDIR [OPTION...] MODE

Modes:
  basebackup     perform a base backup
  cleanup [ FILENAME | XLOG | +[1-9]* ]
                 remove old WAL files after new base backup (run from cron job)

                 It is possible to specify the base backup filename FILENAME,
                 which WAL files should be kept at least or the WAL file XLOG.
                 If a positive number greater than zero is specified, the
                 cleanup command will treat it as its retention policy and keep
                 at least this number of base backup files. Please note that the
                 latter form of the cleanup command will delete all outdated
                 base backups as well whereas the two forms of cleanup invoked
                 with filenames will delete the WAL files only.

                 If no argument is specified, cleanup will remove all WAL
                 files except those which are required by the latest base
                 backup.

  create-lvmsnapshot
                 create an LVM snapshot for an external backup command
                 (requires -L -M -n -N)

  currentbackup  backup the current WAL file (run from cron job)

  ls[+]          Lists available base backups and their size in the current archive.
                 When issued with +, the ls command will examine the WAL archive
                 and the minimum WAL segment file, required to use the backup
                 to perform a full recovery.

  lvmbasebackup  perform a base backup using LVM snapshot
                 (requires -L -M -n -N)
  rsyncbackup
                 Perform a basebackup with rsync. This also saves backup space
                 in case multiple rsync basebackups are used by hardlinking
                 unchanged files. If a retention policy > 1 is used, then
                 unchanged files are just allocated once in the base backup
                 repository. See the --link-dest parameter in the rsync
                 documentation for details.

  setup          prepare server for transaction log archival

  streambackup   perform a streaming basebackup.

                 This command requires PostgreSQL 9.1 and above and pg_basebackup
                 accessible via PATH. The server should be configured to allow
                 streaming replication connections.

  remove-lvmsnapshot
                 remove an LVM snapshot created with create-lvmsnapshot

  restore BASEBACKUP
                 Restores the specified basebackup into the specified
                 directory by the -D parameter. 
                 The directory must already exist and be empty.
                 The destination directory will also contain a generated
                 recovery.conf, suitable to start a PostgreSQL instance for
                 recovery immediately.

Options:
  -A ARCHIVEDIR archival target directory (required)
  -D DATADIR    database system directory (required for specific commands)
  -T TABLESPACE Target directory for tablespace location during restore
                (replaces the original symlinks in the base backup, but places
                all tablespaces into one directory)
  -m            archive old log files before deleting them
  -z            use gzip to compress archived WAL segments
  -l            Place pg_backup_ctl lock file in the specified directory
                (default is ARCHIVEDIR)
LVM snapshot control:
  -L LVM SIZE   determines the buffer size for an LVM snapshot
  -M VOLUME     LVM volume identifier to create the snapshot on
  -n SNAPSHOT   LVM snapshot volume name
  -N LVMDATADIR PostgreSQL DATADIR relative to partition (i.e. the path
                to DATADIR inside the LVM snapshot)
  -o MOUNTOPTS  additional options passed to LVM snapshot mount
  -t FS TYPE    filesystem type to mount LVM snapshot
Server connection control:
  -h HOSTNAME   server host name
  -p PORT       server port
  -U USERNAME   server user name
EOF
	exit 0
}

if [ "${1:---help}" = "--help" ]; then
    print_help
fi

set -- $(getopt A:D:h:l:L:M:mn:N:o:p:U:t:T:z "$@")

while :; do
    case $1 in
        -A) archivedir=$(readlink -f $2); shift;;
        -D) datadir=$(readlink -f $2); shift;;
        -h) export PGHOST=$2; shift;;
        -l) LOCK=$2; shift;;
	-L) lvm_size=$2; shift;;
	-n) lvm_snap_name=$2; shift;;
	-M) lvm_vol=$2; shift;;
        -m) cleanup_move=yes;;
	-n) lvm_snap_name=$2; shift;;
	-N) lvmdatadir=$2; shift;;
        -o) mountopts=$2; shift;;
        -p) export PGPORT=$2; shift;;
        -t) lvm_fstype=$2; shift;;
        -T) pgtblspc_replace_dir=$(readlink -f $2); shift;;
        -U) export PGUSER=$2; shift;;
        -z) gzip=yes;;
        --) shift; break;;
    esac
    shift
done

export PGDATABASE=postgres

mode=$1
shift

# functions

warn() {
    echo "$me WARNING:" "$@" 1>&2
}

error() {
    echo "$me ERROR:" "$@" 1>&2
    exit 1
}

notice() {
    echo "$me NOTICE: " "$@" 1>&2
}

current_setting() {
    psql -tA -c "SELECT current_setting('$1');"
}

timestamp() {
    # A detailed timestamp is good to uniquely identify a backup.
    date +'%Y-%m-%dT%H%M'
}

check_psql_dep() {
    # check if psql binary is available
    if ! command -v psql >/dev/null; then
        error "cannot find psql executable"
    fi
}

check_pg_basebackup_dep() {
    # check if pg_basebackup can be found somewhere
    if ! command -v pg_basebackup >/dev/null; then
        error "cannot find pg_basebackup executable (requires PostgreSQL >= 9.1)"
    fi
}

check_rsync_dep() {
    # check if needed programs are there
    if ! command -v rsync >/dev/null; then
	error "cannot find rsync executable"
    fi;
}

##
## Runs a command within a exclusive locking. This functions is required to run
## a command only once at a time.
##
run_with_lock()
{
    local rc command="$*"

    ## Special handling for non-Linux systems.
    ## We use mkdir instead if the flock() API for
    ## locking, since flock is not avaiable e.g on
    ## OSX and Solaris.

    if echo "$OSTYPE" | grep darwin > /dev/null || echo "$OSTYPE" | grep solaris > /dev/null;
    then
        local PROG=$( basename $0 )

        if mkdir "$LOCK" >/dev/null 2>&1
        then
            trap "rm -rf '$LOCK'" INT TRAP QUIT ABRT TERM EXIT
            chmod 0 "$LOCK" # discourage anyone from messing with it else the rmdir might fail
        else
            echo >&2 "Lock ($LOCK) exists. exiting"
            exit 1
        fi

        ## run the command
        $command
	status=$?

        ## ..and we're done
        ## rmdir "$LOCK"
    else
        set +e  # disable immediate exit on errors for flock()

        (
            flock -x -w $PG_BACKUP_LOCK_TIMEOUT 284
            rc=$?

            # check if the lock was aquired
            if [ $rc -eq 0 ]; then
                $command
            else
                error "failed to run \"$command\" could not aquire exclusive lock
            $LOCK"
            fi

        ) 284>$LOCK

	status=$?
        set -e
    fi
    return $status
}

##
## Returns the PostgreSQL major version of the specified PGDATA directory
## as an integer. That is, the first two digits are concatenated to a unified
## number so it can be compared arithmetically (e.g. 9.0.4 will be returned
## as 90).
##
## Caller should make sure PGDATA really exists
##
get_pg_major_version() {

    local PGMAJOR=$(cat "$datadir"/PG_VERSION | awk -F'.' '{print $1$2;}')
    echo $PGMAJOR

}

## Checks the specified archive directory (-A)
## Args: $1: extra directory to create (current, lvm_snapshot)
check_archivedir() {

    if [ -z "$archivedir" ]; then
        error "no archive directory specified"
    fi

    if [ ! -d "$archivedir" ]; then
        error "archive directory \"$archivedir\" does not exist"
    fi

    # Archive directory writeable?
    if [ ! -w "$archivedir" ]; then
        error "no write access to archive directory \"$archivedir\""
    fi

    ## Force LOCK file to be in $archivedir. This fixes an issue on
    ## some Linux systems, where /var/lock/ isn't writeable by daemons.
    ## $archivedir must be writable by the postgres user anyway, so we
    ## assume we can use it for the lock file, too.
    LOCK=${LOCK:-"$archivedir/pg_backup_ctl.base.lock"}

}

## Checks the specified data directory (-D, if present) or
## gets the setting from a possible running PostgreSQL instance.
check_datadir() {

    # Especially when calling currentbackup, specifying a data directory
    # directly is advisable to avoid frequent database connections.
    if [ -z "$datadir" ]; then
        datadir="$(current_setting data_directory)"
    fi

    if [ ! -r "$datadir"/PG_VERSION ]; then
        error "cannot read data directory (permissions?)"
    fi

}

check_lvm_params() {

	## check LVM parameters

	if [ -z "$lvm_size" ]; then
		error "no LVM snapshot size (-L) specified"
	fi

	## specified volume exists?
	if [ -z "$lvm_vol" ]; then
		error "no LVM volume (-M) specified"
        else
                if ! sudo /sbin/lvdisplay $lvm_vol > /dev/null 2>&1; then
                    error "logical volume \"$lvm_vol\" not found"
                else
		    local RC=$(sudo /sbin/lvdisplay -c $lvm_vol | awk -F ':' '{print $4;}')
		    if [ "$RC" -ne 1 ]; then
                        error "\"$lvm_vol\" is not a valid LVM volume"
                    fi
                fi
	fi

	if [ -z "$lvm_snap_name" ]; then
		error "no LVM backup name (-n) specified"
	fi

	if [ -z "$lvmdatadir" ]; then
		error "LVM snapshots requires the datadir relative to $archivedir/lvm_snapshot (-N)"
	fi

        if [ ! -z "lvm_fstype" ]; then
            lvm_fstype="-t $lvm_fstype"
        fi

        if [ ! -z "$mountopts" ]; then
            mountopts="-o $mountopts"
        fi
}

do_setup() {


    ## XXX: We don't care if the archive directory already exists, since
    ##      it might just be possible that the user wants to reconfigure
    ##      PostgreSQL only.
    local x

    for x in "$archivedir"/{current,base,log,lvm_snapshot} ${1:+"$archivedir/$1"}; do
        if [ ! -d "$x" ]; then
            mkdir -p "$x"
        fi
    done

    if [ $(ls "$datadir"/pg_tblspc | wc -l) -gt 0 ]; then
        warn "clusters uses tablespaces, this is not supported by basebackup and lvm snapshot commands"
    fi

    local cmd
    if [ "$gzip" = yes ]; then
        cmd="test ! -f \\\\'$archivedir/log/%f.gz\\\\' && gzip -c \\\\'%p\\\\' > \\\\'$archivedir/log/%f.gz\\\\'"
    else
        cmd="test ! -f \\\\'$archivedir/log/%f\\\\' && cp \\\\'%p\\\\' \\\\'$archivedir/log/%f\\\\'"
    fi
    local cf=$(current_setting config_file)

    if $GREP -q archive_command "$cf"; then
        $SED -i -e "/archive_command/ c archive_command = '$cmd'" "$cf"
    else
        (echo; echo "# automatically added by $me"; echo "archive_command = '$cmd'") >>$cf
        echo "Added archive_command to postgresql.conf"
    fi

    ## Activate archiving. This requires a server restart if not yet set. Hint the
    ## user, if we changed anything here.
    local archivemode=$(current_setting archive_mode)

    if $GREP -q archive_mode "$cf"; then
        if [ ! -z "$archivemode" ] && [ "$archivemode" = "off" ]; then
            $SED -i -e '/archive_mode/ c archive_mode = on' "$cf"
            echo "HINT: activated archive_mode, you need to restart the server to get archiving activated"
        else
            echo "HINT: archive_mode already activated, adjusted archive_command only"
        fi
    else
        (echo; echo "# automatically added by $me"; echo "archive_mode = on") >>$cf
        echo "Added archive_mode to postgresql.conf"
    fi

    local pg_major_version=$(get_pg_major_version)

    ## Starting with PostgreSQL 9.0, we define different wal levels to adjust the traffic into
    ## the transaction log. For archiving, we need at least wal_level = archive to be set, otherwise
    ## PostgreSQL will refuse to start.

    if [ $pg_major_version -ge 90 ]; then
        local wallevel=$(current_setting wal_level)

        if $GREP -q wal_level "$cf"; then
            if [ ! -z "$wallevel" ] && [ "$wallevel" != "archive" ] \
                && [ "$wallevel" != "hot_standby" ] && [ "$wallevel" != "logical" ]; then
                $SED -i -e '/wal_level/ c wal_level = archive' "$cf"
                echo "HINT: set wal_level to 'archive', you need to restart the server to get this setting into effect"
            else
                echo "HINT: wal_level already set to \"$wallevel\", no action required"
            fi
        else
            (echo; echo "# automatically added by $me"; echo "wal_level = archive") >>$cf
            echo "Added wal_level to postgresql.conf"
        fi
    fi

    kill -HUP $(head -1 "$datadir/postmaster.pid")

    # warn if rsync is missing
    if [ $pg_major_version -lt 90 ]; then
        if ! command -v rsync >/dev/null; then
            echo "HINT: rsync is not installed. For PostgreSQL version $pg_major_version, you should use the "currentbackup" functionality of pg_backup_ctl which needs rsync"
        fi
    fi
}


do_currentbackup() {

    local files x y

    # identify unarchived files
    for x in $(ls -r "$datadir"/pg_xlog/ | $GREP -E '^[0-9A-F]{24}$'); do
        if [ ! -f "$datadir"/pg_xlog/archive_status/$x.done ]; then
            files="$files $x"
        fi
    done

    # copy unarchived files
    (cd "$datadir"/pg_xlog/ && rsync $files "$archivedir"/current/)

    # remove previously copied files that are now archived
    for x in $(cd "$archivedir"/current && ls); do
        for y in $files; do
            [ $x = $y ] && break 2
        done
        rm "$archivedir"/current/$x
    done
}

do_streambackup() {

    local status=0
    local ts="$(timestamp)"

    local tmpfn="$archivedir/base/.streaming_backup_$ts.in-progress"
    local fn="$archivedir/base/streaming_backup_$ts"

    ##
    ## XXX: Take care for trap here, since run_with_lock() might already have
    ##      installed it's cleanup code to remove the lock file for non-Linux
    ##      platforms. We must not overwrite its cleanup action, so make
    ##      sure to keep it.
    ##      
    trap "rm -rf $tmpfn && rmdir '$LOCK' > /dev/null 2>&1" INT TRAP QUIT ABRT TERM EXIT # clean up on error

    [ -d "$fn" ] && error "Streaming basebackup $fn already exists"
    # NOTE: We don't want to see any NOTICE messages here...
    PGOPTIONS='--client-min-messages=WARNING' pg_basebackup -z -l "$fn" -Ft -D "$tmpfn" || status=$?

    if [ $status -eq 0 ] ; then
	mv "$tmpfn" "$fn"
	return 0
    else
	rm -rf "$tmpfn"
	return $status
    fi
}

_get_last_rsync_backup() {

    local last_one="$(ls -td1 $archivedir/base/rsync_backup* 2> /dev/null | head -n1)"
    echo "$last_one"

}

do_rsync_basebackup() {

    local status=0
    local ts="$(timestamp)"

    local tmpfn="$archivedir/base/.rsync_backup_$ts.in-progress"
    local fn="$archivedir/base/rsync_backup_$ts"
    local last_one=$(_get_last_rsync_backup)
    local rsync_cmd="rsync -q -a --delete  --exclude postmaster.pid --exclude pg_xlog"
    local rsync_cmd_tblspc="$rsync_cmd"
    local tblspc_oids=""

    ##
    ## XXX: Take care for trap here, since run_with_lock() might already have
    ##      installed it's cleanup code to remove the lock file for non-Linux
    ##      platforms. We must not overwrite its cleanup action, so make
    ##      sure to keep it.
    ##
    trap "rm -rf $tmpfn && rmdir '$LOCK' > /dev/null 2>&1" INT TRAP QUIT ABRT TERM EXIT # clean up on error

    [ -d "$fn" ] && error "rsync basebackup $fn already exists"

    ## Check if we have an old backup, which can be used to
    ## use --link-dest of rsync. That will save at least transmits
    ## for relfilenodes which haven't changed since.
    ##
    ## rsync will hardlink all unchanged relfilenodes.
    if [ ! -d "$last_one" ] || [ -z "$last_one" ]; then
        echo "no previous rsync backup found, need to rsync all files"
    else
       echo "using backup \"$last_one\" to hardlink unchanged files"
       rsync_cmd="$rsync_cmd --link-dest=$last_one"
       rsync_cmd_tblspc="$rsync_cmd_tblspc --link-dest=$last_one/.rsync_tblspc_backup"
    fi

    psql -c "SELECT pg_start_backup('$fn');" >/dev/null

    ## Catch non-zero exit code explicitely and don't rely on set -e. We need
    ## to issue a pg_stop_backup(), otherwise this would be hanging around
    ## forever.
    $rsync_cmd "$datadir/" "$tmpfn/" || status=$?

    ## non-zero exit status means failed rsync
    if [ "$status" -ne 0 ]; then
        # NOTE: We don't want to see any NOTICE messages here, but error messages...
        PGOPTIONS='--client-min-messages=WARNING' psql -c "SELECT pg_stop_backup();" 1>/dev/null
        error "error creating base backup, you might need to do a manual cleanup"
    fi

    ## The basebackup for PGDATA should be completed now, we need to
    ## check if there are any remaining tablespaces to be synced. We do
    ## this by just examining the contents of the current basebackup/pg_tblspc
    ## and check for any OIDs present. So the overall sync method for
    ## any tablespaces found is:
    ##
    ## - Create a special directory in $tmpfn, .rsync_tblspc_backup,
    ##   we assume this directory to be always present.
    ## - Check OID symlinks in $tmpfn/pg_tblspc
    ##   We use the already synced base backup to check existing
    ##   tablespace locations.
    ## - For any symlink, loop and sync the tablespace location into
    ##   $tmpfn/.rsync_tblspc_backup/<OID>
    ## - lather, rinse, repeat
    tblspc_oids=$(ls "$tmpfn"/pg_tblspc)

    ## Create the private directory to hold tablespace relfilenodes
    mkdir "$tmpfn"/.rsync_tblspc_backup 2>/dev/null || error "could not create directory for tablespace backups"

    for i in $tblspc_oids; do
        echo "syncing tablespace OID \"$i\""

        ## Create the new target directory

        mkdir "$tmpfn"/.rsync_tblspc_backup/"$i"

        ## Take care if this tablespace OID is new and not already present
        ## in the last backup, if any

        if [ -d "$last_one/.rsync_tblspc_backup/$i" ]; then
            $rsync_cmd_tblspc/$i "$(readlink $datadir/pg_tblspc/$i)/" "$tmpfn"/.rsync_tblspc_backup/"$i"/
        else
            echo "tablespace OID \"$i\" is new, full sync required"
            $rsync_cmd_tblspc "$(readlink $datadir/pg_tblspc/$i)/" "$tmpfn"/.rsync_tblspc_backup/"$i"/
        fi

    done

    # NOTE: We don't want to see any NOTICE messages here, but error messages...
    PGOPTIONS='--client-min-messages=WARNING' psql -c "SELECT pg_stop_backup();" 1>/dev/null

    if [ $status -eq 0 ] ; then
	mv "$tmpfn" "$fn"
	return 0
    else
	rm -rf "$tmpfn"
	return $status
    fi

}

do_basebackup() {

    local status=0
    local ts="$(timestamp)"
    local tmpfn="$archivedir/base/.basebackup_$ts.tar.gz.in-progress"
    local fn="$archivedir/base/basebackup_$ts.tar.gz"

    ## Since do_basebackup() currently doesn't support tablespaces,
    ## error out in case the datadir uses them.
    if [ $(ls "$datadir"/pg_tblspc | wc -l) -gt 0 ]; then
        error "tablespaces not supported by basebackup command."
    fi

    ##
    ## XXX: Take care for trap here, since run_with_lock() might already have
    ##      installed it's cleanup code to remove the lock file for non-Linux
    ##      platforms. We must not overwrite its cleanup action, so make
    ##      sure to keep it.
    ##
    trap "rm -f $tmpfn && rmdir '$LOCK' > /dev/null 2>&1" INT TRAP QUIT ABRT TERM EXIT # clean up on error

    [ -f "$fn" ] && error "Basebackup $fn already exists"
    psql -c "SELECT pg_start_backup('$fn');" >/dev/null
    $TAR --force-local -C "$datadir" -c -z -f "$tmpfn" --exclude=postmaster.pid --exclude='pg_xlog/*' ./backup_label . || status=$?

    # NOTE: We don't want to see any NOTICE messages here, but error messages...
    PGOPTIONS='--client-min-messages=WARNING' psql -c "SELECT pg_stop_backup();" 1>/dev/null

    if [ $status -eq 0 ] || [ $status -eq 1 ]; then # exit 1 is "some files changed"
	mv "$tmpfn" "$fn"
	return 0
    else
	rm -f "$tmpfn"
	return $status
    fi
}

##
## Recover the specified basebackup into the specified
## data directory.
##
do_restore_backup() {

    local status=0
    local fn="$archivedir/base/$1"
    local dn="$datadir"

    # sanity checks first
    if [ -z "$dn" ]; then
	error "The restore command requires the -D command line argument"
    fi

    # check if requested basebackup exists
    if [ ! -f "$fn" ] && [ ! -d "$fn" ]; then
	error "requested basebackup \"$fn\" does not exist"
    fi

    # File found, proceed. Check if the specified
    # recovery target directory exists.
    if [ ! -d "$dn" ]; then
	error "target directory does not exist"
    fi

    # So far, so good, check if an alternate
    # tablespace restore directory was specified and
    # if it exists
    if [ ! -z "$pgtblspc_replace_dir" ] && [ ! -d "$pgtblspc_replace_dir" ]; then
        error "tablespace directory \"$pgtblspc_replace_dir\" does not exist"
    fi

    ## Relocated tablespace directory writable by us?
    if [ ! -z "$pgtblspc_replace_dir" ] && [ ! -w "$pgtblspc_replace_dir" ]; then
        error "no write access to relocated tablespace directory \"$pgtblspc_replace_dir\""
    fi

    # Target directory empty?
    if [ "$($FIND "$dn" -type f | wc -l)" -gt 2 ]; then
	## does it contain a running PostgreSQL instance?
	if [ -e "$dn"/PG_VERSION ]; then
	    local ver=$(get_pg_major_version)
	    error "target directory contains a PostgreSQL database cluster \"$ver\""
	else
	    error "target directory not empty"
	fi
    fi

    # Make sure, target directory has 0700 permissions
    chmod 0700 $dn

    echo "restoring base archive ${fn}"

    ## Unpack the base backup into the destination
    ## PGDATA directory
    case "$(basename ${fn})" in
        rsync_backup*)

            rsync -a --exclude '.rsync_tblspc_backup' "${fn}/" "${dn}/"

            ## Get all OIDs for tablespaces to be restored, call _restore_tablespace
            ## for each OID found...
            ## The easiest way to do this is to get the OID symlinks from the already
            ## extracted base.tar file, pass it down to _restore_tablespace which
            ## does all the remaining leg work.

            local tblspc_oid=$(ls "$dn/pg_tblspc")
            for i in $tblspc_oid; do
                _restore_tablespace "$i" "$dn" "$fn"
            done

            ;;
        streaming_backup*)

            $TAR --force-local -C "$dn" -x -z -f "$fn"/base.tar.gz && echo "successfully restored base.tar.gz"
            echo "checking for tablespaces"

            ## Get all OIDs for tablespaces to be restored, call _restore_tablespace
            ## for each OID found...
            ## The easiest way to do this is to get the OID symlinks from the already
            ## extracted base.tar file, pass it down to _restore_tablespace which
            ## does all the remaining leg work.

            local tblspc_oid=$(ls "$dn/pg_tblspc")
            for i in $tblspc_oid; do
                _restore_tablespace "$i" "$dn" "$fn"
            done

            ;;
        basebackup_*.tar*)
            # Unpack TAR archive into destination directory
            $TAR --force-local -C "$dn" -x -z -f "$fn" && echo "successfully restored \"$fn\""
            ;;
        *)
            error "Unknown basebackup: \"$fn\""
            ;;
    esac

    # create pg_xlog/archive_status directory (newer PostgreSQL
    # releases will do this automatically, but it doesn't no harm
    # if done here...)
    mkdir -p "$dn"/pg_xlog/archive_status

    ## place the recovery.conf
    _gen_recovery_conf "$dn"

    [ $status -eq 0 ] || [ $status -eq 1 ]

}

##
## Restores a tablespace belonging to a specific
## base backup. The base backup files are required to be
## restored previously, so this function is intended
## to be called by do_restore_backup() only.
##
## Arguments required are:
## $1: Tablespace OID
## $2: PGDATA directory
## $3: Archive directory
##
_restore_tablespace() {

    local tblspc_oid="$1"
    local dn="$2"
    local fn="$3"
    local lnk=""
    local adj_symlink

    case "$tblspc_oid" in
        [0-9]*)

            ## Get the link directory, but honor an eventually specified
            ## -T option; In this case we replace all tablespace target
            ## directories with the directory provided by -T...
            if [ -z "$pgtblspc_replace_dir" ]; then
                lnk=$(readlink "$dn"/pg_tblspc/"$tblspc_oid")
            else
                lnk="$pgtblspc_replace_dir/$tblspc_oid"
                mkdir -p "$lnk"
                adj_symlink=1
            fi

            ## link target directory exists?
            if [ -d "$lnk" ]; then

                ## target directory empty?
                if [ $(ls "$lnk" | wc -l) -gt 0 ]; then
                    error "directory \"$lnk\" for tablespace OID \"$tblspc_oid\" not empty"
                fi

                case "$fn" in
                    *rsync_backup_*)
                        rsync -a "$fn"/.rsync_tblspc_backup/"$tblspc_oid"/ "$lnk"
                        ;;
                    *streaming_backup_*)
                        $TAR --force-local -C "$lnk" -x -z -f "$fn"/"$tblspc_oid".tar.gz
                        ;;
                esac

                echo "successfully restored tablespace ${tblspc_oid} to \"$lnk\""

                ## In case we have a -T option to replace tablespace directories,
                ## we need to adjust the symlinks located in PGDATA/pg_tblspc
                ## as well...
                if [ ! -z "$adj_symlink" ]; then
                    local pg_major_ver
                    echo Adjusting symlink "$dn"/pg_tblspc/"$tblspc_oid" to "$lnk"
                    ln -sf -t "$dn"/pg_tblspc/ "$lnk"

                    pg_major_ver=$(get_pg_major_version "$dn")
                    if [ $pg_major_ver -lt 92 ]; then
                        echo "Adjusted tablespace; be sure to adjust pg_tablespace"
                        echo "You must update the location of tablespace OID \"tblspc_oid\":"
                        echo ""
                        echo "UPDATE pg_tablespace SET spclocation='$lnk' WHERE oid = $tblspc_oid"
                        echo ""
                    fi
                fi

            else
                error "directory \"$lnk\" for tablespace OID \"$tblspc_oid\" does not exist"
            fi
            ;;
        *)
            error "restore tablespace \"$tblspc_oid\": not a valid OID"
            ;;
    esac
}

##
## Generates a recovery.conf file and places
## it withtin the specified target directory
##
## Caller is responsible to place a valid
## target directory within $1...
##
_gen_recovery_conf() {

    local dn="$1"
    local cmd=""

    if [ "$gzip" = yes ]; then
        cmd="gzip -d -c \"$archivedir/log/%f.gz\" > \"%p\""
    else
        cmd="cp \"$archivedir/log/%f\"  \"%p\""
    fi

    echo "restore_command='$cmd'" > "$dn"/recovery.conf

}

##
## Create an LVM snapshot, pass $2 as a backup label
## to pg_start_backup() if required, overriding the default
## label "LVM SNAPSHOT YYYY-MM-DDTHHMM"
##
_create_snapshot_LVM() {

    local status=0
    local fn="$1"
    local label

    if [ ! -z "$2" ]; then
        label="$fn"
    else
        ##
        ## XXX: We obtain a separate timestamp for labeling the LVM snapshot.
        ##      Do this only in case no backup_label was specified (which,
        ##      effectively would have generated the label and basebackup
        ##      filename already).
        ##
        label="LVM SNAPSHOT $(timestamp)"
    fi

    psql -c "SELECT pg_start_backup('$label');" >/dev/null
    sudo lvcreate -s -n "$lvm_snap_name" -L "$lvm_size" "$lvm_vol" || status=$?

    psql -c "SELECT pg_stop_backup();" >/dev/null
    [ $status -ne 0 ] && echo "LVM snapshot failed with status $status" && return $status

    ## mount the snapshot in the archive directory LVM snapshot mountpoint
    local LVMSNAPSHOTMNT="$(dirname $lvm_vol)/$lvm_snap_name"
    sudo /bin/mount $lvm_fstype $mountopts $LVMSNAPSHOTMNT "$archivedir"/lvm_snapshot || error "could not mount snapshot in $archivedir/lvm_snapshot"

}

##
## Remove an LVM snapshot previously created with _create_snapshot_LVM
##
_remove_snapshot_LVM() {

    local status=0
    local LVMSNAPSHOTMNT="$(dirname $lvm_vol)/$lvm_snap_name"

    ## unmount the snapshot
    sudo /bin/umount "$archivedir"/lvm_snapshot
    echo "removing LVM snapshot $LVMSNAPSHOTMNT"
    sudo /sbin/lvremove -f -t $LVMSNAPSHOTMNT > /dev/null 2>&1 && sudo /sbin/lvremove -f $LVMSNAPSHOTMNT 1> /dev/null || status=$?

    return $status
}

##
## Create an LVM snapshot and perform a basebackup via tar.
##
## This calls _create_snapshot_LVM() internally, which does
## the legwork of creating the necessary LVM snapshot and
## prepares the mount.
##
## Please note that do_create-lvmsnapshot never performs a real basebackup,
## instead it creates an archive file with just the backup_label included. It
## is up to an external backup software to do the actual base backup.
##
do_create-lvmsnapshot() {
    check_lvm_params

    local status=0
    local fn="$archivedir"/base/basebackup_$(timestamp).tar.gz

    ## do the backup, fall through all errors!
    if _create_snapshot_LVM "$fn"; then
        $TAR --force-local -C "$archivedir/lvm_snapshot/$lvmdatadir/" -c -z -f "$fn" ./backup_label || status=$?
    fi

    [ $status -eq 0 ] || [ $status -eq 1 ]
}

##
## Remove an LVM snapshot previously created via _create_snapshot_LVM().
##
do_remove-lvmsnapshot() {
    local status=0

    _remove_snapshot_LVM || status=$?

    [ $status -eq 0 ] || [ $status -eq 1 ]
}

##
## Creates a full basebackup from an LVM snapshot.
##
do_lvmbasebackup() {
    check_lvm_params

    set -e
    local status=0

    local fn="$archivedir"/base/basebackup_$(timestamp).tar.gz
    _create_snapshot_LVM "$fn" "$fn"

    ## do the backup, fall through all errors!
    $TAR --force-local -C "$archivedir/lvm_snapshot/$lvmdatadir" -c -z -f "$fn" --exclude=postmaster.pid --exclude=pg_xlog . || status=$?

    ## unmount the snapshot
    _remove_snapshot_LVM || status=$?

    [ $status -eq 0 ] || [ $status -eq 1 ]
}

start_wal_location() {
    $SED -n -r '/^START WAL LOCATION:/s/^.*file ([0-9A-F]{24}).*$/\1/p'
}


stop_wal_location() {
    $SED -n -r '/^STOP WAL LOCATION:/s/^.*file ([0-9A-F]{24}).*$/\1/p'
}


do_cleanup() {
    local first_wal last_wal x

    local fn=$(basename "$1")
    case $fn in
        rsync_backup_*)
            ## Fall through, same handler for streaming base backups...
            ;&
        streaming_backup_*)

            ##
            ## If the basebackup is not there anymore, print a notice and
            ## go further. We delete the archive logs anyways, since they don't
            ## have any connection to an existing basebackup.
            ##
            ## Make sure we check for the correct backup_label file. tar basebackups
            ## include a separate backup_label file to make sure it occurs in the archive
            ## first.
            ##
            if [ -e "${archivedir}/base/${fn}/base.tar.gz" ]; then
                first_wal=$(gunzip -c -f "${archivedir}/base/${fn}/base.tar.gz" | $TAR -f - -x --occurrence=1 -O backup_label | start_wal_location)
            elif [ -e "${archivedir}/base/${fn}/PG_VERSION" ]; then
                first_wal=$(cat "${archivedir}/base/${fn}/backup_label" | start_wal_location)
            else
                echo "Basebackup ${fn} does not exist, use cleanup <HISTORYFILE> to tidy up"
            fi
            ;;
        basebackup_*.tar*)

            ##
            ## If the basebackup is not there anymore, print a notice and
            ## go further. We delete the archive logs anyways, since they don't
            ## have any connection to an existing basebackup.
            ##
            if [ -e "$archivedir/base/${fn}" ]; then
                first_wal=$(gunzip -c -f "$archivedir/base/${fn}" | $TAR -f - -x --occurrence=1 -O ./backup_label | start_wal_location)
            else
                echo "Basebackup $fn does not exist, use cleanup <HISTORYFILE> to tidy up"
            fi
            ;;
        [0-9A-F]*[0-9A-F].[0-9A-F]*[0-9A-F].backup)
            first_wal=$(cat "$archivedir/log/$1" 2>/dev/null | start_wal_location)
            last_wal=$(cat "$archivedir/log/$1" 2>/dev/null | stop_wal_location)
            if [ ! -f "$archivedir"/log/$last_wal ]; then
                error "base backup \"$1\" is not complete; log segment \"$last_wal\" is not archived yet"
            fi
            ;;
        [0-9A-F]*[0-9A-F].[0-9A-F]*[0-9A-F].backup.gz)
            first_wal=$(gunzip -c -f "$archivedir/log/$1" | start_wal_location)
            last_wal=$(gunzip -c -f "$archivedir/log/$1" | stop_wal_location)
            if [ ! -f "$archivedir"/log/"$last_wal.gz" ]; then
                error "base backup \"$1\" is not complete; log segment \"$last_wal\" is not archived yet"
            fi
            ;;
        '')
            local last_backup
            for x in $(ls -r "$archivedir"/log/*.backup* | head -1); do
                last_wal=$(gunzip -c -f "$x" | stop_wal_location)
                if [ -f "$archivedir"/log/$last_wal ] || [ -f "$archivedir"/log/$last_wal.gz ] ; then
                    last_backup=$x
                    break
                fi
            done
            if [ -z "$last_backup" ]; then
                error "no complete backup found"
            fi
            first_wal=$(gunzip -c -f "$last_backup" | start_wal_location)
            ;;
	+[1-9]*)
	    local retention_policy=$1

            ## NOTE: We need to consider both, traditional tar files and
            ## streamed backups here.
	    local base_backups=$(ls -tdr "$archivedir"/base/*backup_* 2>/dev/null)
	    local num_backups=$(echo "$base_backups" | wc -l)

	    if [ $num_backups -le $retention_policy ]; then
		echo "no base backups to delete found"
		return 1
	    else
		local candidates
		local newest_candidate

		candidates=$(echo "$base_backups" | head -n $(($num_backups - $retention_policy)))
		oldest_keep=$(echo "$base_backups" | head -n $(($num_backups - $retention_policy + 1)) | $TAIL -n1)
		newest_candidate=$(echo "$candidates" | $TAIL -n1)

		## This recurses into cleanup and removes all WAL files
		## belonging to the list of candidates. Those aren't required
		## anymore, since we are going to delete base backup candidates
		## afterwards.
		## NOTE:
		##
		## We need to specify the oldest base backup to *keep*, since
		## do_cleanup() will keep all WAL files belonging to the specified
		## base backup.
                ##
                ## XXX: Currently we are performing the cleanup based on selecting
                ##      the base backup files and directories from $archivedir/base.
                ##      In the future, we might change this code to select the base
                ##      backups according to the backup history files found there.
		do_cleanup $oldest_keep

	    fi

	    ## exit immediately after cleaning base backups. We *must* return
	    ## to the caller, otherwise do_cleanup() will proceed and delete all WAL
	    ## files older than the current base backup.
	    return 0
	    ;;
        *)
            error "invalid cleanup target specification"
            ;;
    esac

    if [ -f "$archivedir"/log/$first_wal ] || [ -f "$archivedir"/log/$first_wal.gz ]; then
        local old_files=$(ls "$archivedir"/log | $SED -n "/$first_wal/q;p")
       
        (
        cd "$archivedir"/log

        local old_base_backups=$(echo $old_files | $SED -e 's/ /\n/g' | $GREP -P '^[0-9A-Z]{24}.*.backup*' |
	    if [ -f "$archivedir"/log/$first_wal ]; then
		xargs cat
	    else
		xargs gunzip -d -c
	    fi | backup_label
	)

        if [ "$cleanup_move" = yes ]; then
            $TAR --force-local -c -z -f "$archivedir"/base/oldwals-$(timestamp).tar.gz $old_files
        fi

        ## Delete old base backups first. We loop through the list,
        ## since we assume not that many base backup files in the archive that
        ## would eat much execution time.
        ##
        ## NOTE: Since the base backup is fully qualified
        ##       in the backup label, we don't need to explicitely add
        ##       the absolute archive path.
        ##
        ##       LVM snapshot needs special handling here, since
        ##       they aren't stored inline the backup catalog.
        for i in $old_base_backups; do
            local rc
            basebackup_exists $i
            rc=$?

            if [ $rc -eq 0 ]; then
                notice "deleting base backup \"$i\""
                rm -r "$i"
            elif [ $rc -lt 1 ]; then
                warn "LVM snapshot \"$i\" skipped"
            else
                warn "base backup \"$i\" does not exist, ignored"
            fi

        done

        ## now cleanup transaction log files
        [ ! -z "$old_files" ] && echo $old_files | xargs rm
        )
    fi
}

##
## Checks wether the specified base backup exists in the catalog.
## Returns 0 in case of true, 1 in case of not existant.
## A return code 2 indicates a LVM snapshot not stored inline
## in the backup catalog.
basebackup_exists() {

    local fn="$(basename $1)"
    local rc=1

    case ${fn} in
        LVM*SNAPSHOT*)
            ## LVM SNAPSHOT is not an internally stored base backup,
            rc=2
            ;;
        *)
            test -e "$archivedir/base/${fn}"
            rc=$?
            ;;
    esac

    return $rc

}

backup_label() {
    $SED -n -e '/^LABEL:/s/^LABEL://;s/^\s//p'
}

## Returns the size of the specified file or directory
##
## In case the specified file does not exist, 0 is returned
_get_size() {

    local fn="$1"
    local basesize=0

    if [ -e "$fn" ]; then
        basesize=$(du -sh "$fn" | awk '{print $1;}')
    else
        basesize="N/A"
    fi

    echo -n $basesize
}

##
## Returns a list of tablespaces OIDs and paths included in the specified
## base backup. The returned pathname assigned to the tablespace OID
## is directly retrieved from the specified basebackup and reflects
## the real physical tablespace path from the origin pgsql instance.
##
_get_backup_tblspc_oids() {

    local fn="$1"
    local _archive="$archivedir/base/$fn"

    case "$fn" in

        rsync_backup_*)
            for i in $(ls "$_archive"/.rsync_tblspc_backup 2>/dev/null); do
                echo "$i->$(readlink $_archive/pg_tblspc/$i)"
            done
            ;;
        streaming_backup_*)
            for i in $(ls "$_archive/"/[0-9]*.tar.gz 2>/dev/null); do
                echo $(basename $i .tar.gz)
            done
            ;;
    esac

    echo ""

}

##
## Checks if the specified base backups contains tablespaces.
## 0 if true, otherwise 1 is returned.
##
_backup_has_tablespaces() {

    local fn="$1"

    case "$fn" in

        rsync_backup_*)

            ## RSYNC base backups store tablespaces within its
            ## private directory, just check if there are any subdirs
            ## and we know that there are tablespaces present

            [ $(ls -1 "$archivedir/base/${fn}"/.rsync_tblspc_backup/ 2>/dev/null | wc -l) -gt 0 ] && return 0

            ;;

        streaming_backup_*)

            ## A streaming basebackup has tablespaces within separate tar.gz archives for each
            ## tablespace.

            [ $(ls -1 "$archivedir/base/${fn}"/[0-9]*.tar.gz 2>/dev/null | wc -l) -gt 0 ] && return 0
            ;;

        *)
            ## Unsupported backup type or invalid name, but treat it
            ## as a simple "no, this is a backup type without tablespaces",
            ## the error should have been catched before.
            return 1
            ;;
    esac

    ## No tablespaces present
    return 1
}

##
## do_ls()
##
## Accepts an argument literal '+' to make the output more verbose.
## This verbose output also verifies wether the required WAL files
## for a specific base backup are still present or not. Since the
## verbose output calculates these information according to the
## existing backup history files, we should protect any calls
## to do_ls '+' with an filesystem lock...see the function
## run_with_lock() for details.
do_ls() {

    local backups basesize total_basesize wallogsize walcount

    backups=$(ls -t $archivedir/log/[0-9A-F]*[0-9A-F].[0-9A-F]*[0-9A-F].backup* 2> /dev/null || return 0)

    if [ -z "$1" ]; then
        printf "%-51s\t%-9s\n" "Basebackup Filename" "Size"
    else
        printf "%-33s %-9s %-24s %-9s\n" "Basebackup Filename" "Size" "Required WAL" "Available"
    fi

    echo $hdrline

    for x in $backups; do

        local fn=$(basename "$(cat "$x" | gunzip -f -c | backup_label)")

        case $fn in

            basebackup_*.tar.gz)
                fn="$(basename $fn)"
                basesize=$(_get_size "$archivedir/base/$fn")
                ;;
            *SNAPSHOT*)
                basesize=0
                ;;
            rsync_backup_*)
                ;&
            streaming_backup_*)
                # This is an archive directory containing a streaming base backup
                fn="$(basename $fn)"
                basesize=$(_get_size "$archivedir/base/$fn")
                ;;
            *)
                continue
                ;;

        esac

        if [ -z "$1" ]; then
            printf "%-51s\t%-9s\n" "$fn" $basesize
        elif [ "$1" = "+" ]; then

            local wal_avail='YES'
            local has_tablespaces="NO"

            ## XXX: a backup history file might be compressed, too
            first_wal=$(cat "$x" | gunzip -c -f | start_wal_location)

            (test -f "$archivedir/log/$first_wal" || test -f "$archivedir/log/$first_wal.gz") || export wal_avail='NO'
            _backup_has_tablespaces "$fn" && export has_tablespaces="YES"
            printf "%-33s %-9s %-24s %-9s\n" "$fn" $basesize $first_wal $wal_avail
            printf "\`- %-24s\n" $(basename $x)
            if [ x"$has_tablespaces" = x"YES" ]; then
                printf "\`- Tablespaces\n"
                printf "   \`- %-9s\n" $(_get_backup_tblspc_oids "$fn")
            fi
        fi
    done

    ## sum up archive log size
    wallogsize=$(du -sh "$archivedir/log/" | awk '{print $1;}')
    walcount=$(ls -1 "$archivedir/log/" | $GREP -E '^[0-9A-Z]{24}\.?[gz]*$' | wc -l)
    total_basesize="$(_get_size $archivedir/base)"
    echo $hdrline
    printf "Total size of base backups: %-9s\n" $total_basesize
    printf "Total size occupied by %d WAL segments in archive: %-9s\n" $walcount $wallogsize

    return 0
}

# main

case $mode in
    help)
	print_help;;
    setup)
        check_psql_dep
	check_archivedir
	check_datadir
        do_setup;;
    currentbackup)
        check_psql_dep
        check_rsync_dep
	check_archivedir current
	check_datadir
        run_with_lock do_currentbackup
        ;;
    streambackup)
        check_psql_dep
        check_pg_basebackup_dep
        check_archivedir
        run_with_lock do_streambackup
        ;;
    rsyncbackup)
        check_psql_dep
	check_archivedir
	check_datadir
        run_with_lock do_rsync_basebackup
        ;;
    basebackup)
        check_psql_dep
	check_archivedir
	check_datadir
        run_with_lock do_basebackup
        ;;
    lvmbasebackup)
        check_psql_dep
	check_archivedir lvm_snapshot
	check_datadir
        run_with_lock do_lvmbasebackup
        ;;
    create-lvmsnapshot)
        check_psql_dep
	check_archivedir lvm_snapshot
        run_with_lock do_create-lvmsnapshot
        ;;
    remove-lvmsnapshot)
        # NOTE: doesn't require psql, only remove the LVM snapshot via lvremove
	check_archivedir lvm_snapshot
        run_with_lock do_remove-lvmsnapshot
        ;;
    ls)
        ## ls command without any arguments accesses the
        ## filesystem only, thus no lock required
        do_ls
        ;;
    ls+)
        check_archivedir
        run_with_lock do_ls '+'
        ;;
    cleanup)
	check_archivedir
        run_with_lock do_cleanup "$@"
        ;;
    restore)
	check_archivedir
        ## NOTE: do_restore_basebackup() does its
        ##       own sanity checks on $datadir
        run_with_lock do_restore_backup "$1" "$datadir"
	    ;;
    *)
        error "invalid mode: \"$mode\"";;
esac
