#!/usr/bin/env python
#
# pg_restorebackup.py
#
# PoC of tool to build a ful backup starting from an incremental
# backup and one or more parent bakcups.
#
# The -T option is mandatory if there was any tablespace defined in
# the PostgreSQL instance when the incremental_backup was taken

from __future__ import print_function

import os
import shutil
import sys
from optparse import OptionParser

usage = """usage: %prog [options] dest_dir backup_1 backup_2 [backup_3 ...] 

The -T option is mandatory if there was any tablespace defined in the
PostgreSQL instance when the incremental_backup was taken
"""
parser = OptionParser(usage=usage)
parser.add_option("-T", "--tablespace-mapping",
                  dest="tablespace_mapping",
                  default=[],
                  action="append",
                  help="Relocate the tablespace present in incremental_backup from directory olddir to newdir",
                  metavar="olddir=newdir")
(options, args) = parser.parse_args()

# build a tablespace_mapping dictionary from options
tablespace_mapping = {}
for mapping in options.tablespace_mapping:
    try:
        olddir, newdir = mapping.split('=')
    except:
        print("error: invalid tablespace mapping (%s)" % mapping, file=sys.stderr)
        sys.exit(1)
    tablespace_mapping[olddir]=newdir

# need at least 3 arguments
if len(args) < 3:
    parser.print_help()
    sys.exit(1)

# the first argument is the destination directory
dest=args[0]
# last argument is the backup we are restoring
incr=args[-1]
# all other arguments are parent backups in the correct order
base=args[1:-1]

# verify that all the backup are in the right order
last_lsn = None
last_backup = None
for backup in base + [incr]:
    try:
        profile=open(os.path.join(backup, 'backup_profile'), 'r')
    except IOError as e:
        print("error: invalid backup directory '%s': " % (backup, e), file=sys.stderr)
        sys.exit(1)
    start_lsn = None
    incremental_lsn = None
    for line in profile:
        if line.startswith('START WAL LOCATION:'):
            start_lsn = line.split()[3]
        if line.startswith('INCREMENTAL FROM LOCATION:'):
            incremental_lsn = line.split()[3]
        if line.strip() == 'FILE LIST':
            break
    if last_lsn is not None:
        if incremental_lsn is None:
            print("error: invalid backup sequence: '%s' does not contains an incremental backup"
                  % backup, file=sys.stderr)
            sys.exit(1)
        if last_lsn != incremental_lsn:
            print("error: invalid backup sequence: '%s' is not incremental from '%s' location %s"
                  % (backup, last_backup, last_lsn), file=sys.stderr)
            sys.exit(1)
    # save parent for the next cycle
    last_lsn = start_lsn
    last_backup = backup

# 1: start copying the incremental backup to destination directory
if os.path.exists(dest):
    print("error: destination '%s' must not exist" % dest, file=sys.stderr)
    sys.exit(1)
shutil.copytree(incr, dest, symlinks=True)

# 2: apply the tablespaces relocation as required by -T options
incr_tblspc = os.path.join(incr, 'pg_tblspc')
dest_tblspc = os.path.join(dest, 'pg_tblspc')
for tblspc in os.listdir(incr_tblspc):
    incr_file = os.path.join(incr_tblspc, tblspc)
    dest_file = os.path.join(dest_tblspc, tblspc)
    if not os.path.islink(incr_file):
        print("error: illegal file in source pg_tblspc directory (%s)"
              % incr_file, file=sys.stderr)
        sys.exit(1)
    old_target = os.readlink(incr_file)
    if old_target not in tablespace_mapping:
        print("error: missing tablespace mapping for '%s'" % old_target, file=sys.stderr)
        sys.exit(1)
    new_target = tablespace_mapping[old_target]
    os.unlink(dest_file)
    os.symlink(new_target, dest_file)
    # tablespace destinations must not exist
    if os.path.exists(new_target):
        print("error: tablespace destination directory '%s' must not exist"
              % new_target, file=sys.stderr)
        sys.exit(1)
    shutil.copytree(old_target, new_target)

# 3: parse the backup profile, skipping the backup label
profile=open(os.path.join(incr, 'backup_profile'), 'r')
for line in profile:
    if line.strip() == 'FILE LIST':
        break

# 4: for every file copy it from ancestors if needed
for line in profile:
    tblspc, lsn, sent, date, size, path = line.strip().split('\t')
    # If lsn is '\N' the file is included in the backup
    if lsn == '\\N':
        continue
    # If the file is in a tablespace prepend the tablespace path
    if tblspc != '\\N':
        path = os.path.join('pg_tblspc', tblspc, path)
    # Search the file in all the ancestors
    found = False
    for backup in reversed(base):
        base_file = os.path.join(backup, path)
        if os.path.exists(base_file):
            found = True
            break
    if not found:
        print("error: the file '%s' is not present in any ancestor backup"
              % path, file=sys.stderr)
        sys.exit(1)
    dest_file = os.path.join(dest, path)
    shutil.copy2(base_file, dest_file)