#!/usr/bin/env python

import os, sys, gettext, getopt, re, time, errno, stat

_ = gettext.gettext

LOCALEDIR	= "/usr/share/locale"

PG_VERSION	= "7.4.2"
BLCKSZ		= 8192
progname	= None

archiveDestDir	= None
dataDir		= None

XLogDir		= ""
XLogArchiveDir	= ""

nextrlogId	= -1L
nextrlogSeg	= 0L


XLogSegSize	= 16 * 1024 * 1024
XLogSegsPerFile	= (16 ** 8 - 1) / XLogSegSize

try:
    O_BINARY = os.O_BINARY
except:
    O_BINARY = 0

def XLogArchiveXLogs(rlogdir):
    """ XLogArchiveXLogs

	Return name of the oldest xlog file that has not yet been archived,
	setting notification that file archiving is now in progress. 
	It is important that we return the oldest, so that we archive xlogs
	in order that they were written, for two reasons: 
	1) to maintain the sequential chain of xlogs required for recovery 
	2) because the oldest ones will sooner become candidates for 
	recycling by checkpoint backend.
    """
    """ implementation:
	if first call, open XLogArchive directory and read through list of 
	rlogs that have the .full suffix, looking for earliest file. 
	Decode xlog part of rlog filename back to 
	log/seg values, then increment, so we can predict next rlog.
	If not first call, use remembered next rlog value in call to stat,
	to see if that file is available yet. If not, return empty handed.
	If so, set rlog file to .busy, increment rlog value again and then 
	return name of available file to allow copy to archive to begin.
    """
    global nextrlogId, nextrlogSeg
    
    if nextrlogId == -1:
	if not os.path.isdir(rlogdir):
	    print >> sys.stderr, _("%s could not open rlog directory" % progname)
	    sys.exit(1)

	print _("\n%s firstcall: scanning rlogdir..." % progname)
	is_file_found = False
	newxlog = ""
	for filename in os.listdir(rlogdir):
	    _re = re.compile(r'^[0-9A-F]{16}\.full$')
	    if _re.match(filename):
		if not is_file_found:
		    is_file_found = True
		    newxlog = filename
		elif filename < newxlog:
		    newxlog = filename

	# strip off the suffix to get xlog name
	if len(newxlog) == 21:
	    newxlog = newxlog[:16]

	print _("%s closing rlogdir..." % progname)

	if not is_file_found:
	    print _("%s no .full rlogs found..." % progname)
	    return (False, None)
	
	print _("%s found...%s" % (progname,  newxlog))

	# decode xlog back to LogId and SegId, so we can increment
	nextlogstr = newxlog[:8]
	try:
	    nextrlogId = long(nextlogstr, 16)
	except:
	    print >> sys.stderr, _("%s decode xlog logID error" % progname)
	    sys.exit(1)

	nextlogstr = newxlog[8:]
	try:
	    nextrlogSeg = long(nextlogstr, 16)
	except:
	    print >> sys.stderr, _("%s decode xlog logSeg error" % progname)
	    sys.exit(1)
	xlog = newxlog

	# set the rlog to .busy until XLogArchiveComplete is called
	rlogfull = "%s/%s.full" % (rlogdir, xlog)
	rlogbusy = "%s/%s.busy" % (rlogdir, xlog)
	try:
	    os.rename (rlogfull, rlogbusy)
	except OSError:
	    print >> sys.stderr, _("%s XLogArchiveXLogs could not rename %s to %s" %
			   (progname, rlogfull, rlogbusy))
	    return (False, None)
    else:
	nextlogstr = "%08X%08X" % (nextrlogId, nextrlogSeg)
	rlogfull = "%s/%s.full" % (rlogdir, nextlogstr)

	# if .full file is not there...that's OK...we wait until it is
	if not os.path.isfile(rlogfull):
	    # Good error checking required here, otherwise we might loop
	    # forever, slowly!
	    print _("%s %s not found yet..." % (progname, nextlogstr))
	    return (False, None)

	# set the xlog that will be archived next
	xlog = "%08X%08X" % (nextrlogId, nextrlogSeg)

	# set the rlog to .busy until XLogArchiveComplete is called
	rlogbusy = "%s/%s.busy" % (rlogdir, xlog)
	try:
	    os.rename (rlogfull, rlogbusy)
	except OSError:
	    print >> sys.stderr, _("%s XLogArchiveComplete could not rename %s to %s" % 
			   (progname, rlogfull, rlogbusy))
	    return (False, None)
    
    # increment onto the next rlog
#    NextLogSeg(nextrlogId, nextrlogSeg)

    if (nextrlogSeg >= XLogSegsPerFile - 1):
	nextrlogId += 1
	nextrlogSeg = 0
    else:
	nextrlogSeg += 1

    # we have an xlog to archive...
    return (True, xlog)


def XLogArchiveComplete(xlog, rlogdir):
    """ XLogArchiveComplete
	Write notification that an xlog has now been successfully archived
    """
    """ implementation:
	stat the notification file as xlog filename with .busy suffix
	Rename the notification file to a suffix of .done
    """
    rlogbusy = "%s/%s.busy" % (rlogdir, xlog)
    if not os.path.isfile (rlogbusy):
	print >> sys.stderr, _("%s XLogArchiveComplete could not locate %s" % progname, rlogbusy)
	return False

    rlogdone = "%s/%s.done" % (rlogdir, xlog)
    try:
        os.rename (rlogbusy, rlogdone)
    except OSError:
	print >> sys.stderr, _("%s XLogArchiveComplete could not rename %s to %s" %
			(progname, rlogbusy, rlogdone))
	return False

    return True


def CopyXLogtoArchive(xlog):
    """ CopyXLogtoArchive

	Copy transaction log from the pg_xlog directory of a PostgreSQL instance identified
	by the DATADIR parameter through to an archive destination, ARCHIVEDESTDIR

	Should ignore signals during this section, to allow archive to complete
    """
    """ Implementation:
	We open the archive file using O_SYNC to make sure no mistakes
	writing data in buffers equal to the blocksize, so we will
	always have at least a partially consistent set of data to recover from
	...then we check filesize of written file to ensure we did it right
    """
    xlogpath = "%s/%s" % (XLogDir, xlog)
    print _("%s xlogpath= %s" % (progname, xlogpath))
    if not os.path.isfile(xlogpath):
	print >> sys.stderr, _("%s xlog does not exist" % progname)
	return False

#    xlogpath = "%s/%s" % (XLogDir, xlog)
    xlogfd = -1
    try:
	xlogfd = os.open(xlogpath, os.O_RDONLY)
    except OSError, e:
	if (e.errno != errno.ENOENT):
	    if e.errno == errno.EACCES:
		print >> sys.stderr, _("%s EACCES" % progname)
    if xlogfd < 0:
        return False
    print >> sys.stderr, _("%s xlog file opened" % progname)

    archpath = "%s/%s" % (archiveDestDir, xlog)
    try:
	flags = os.O_RDWR | os.O_CREAT | os.O_EXCL | os.O_SYNC | O_BINARY
	mode = stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IWGRP
	archfd = os.open(archpath, flags, mode)
    except OSError, e:
	if e.errno == os.EEXIST:
	    print >> sys.stderr, _("%s archive file %s already exists in %s" % (progname, xlog, archiveDestDir))
	    sys.exit(1)
	return False

    print >> sys.stderr, _("%s archive file opened" % progname)

    n = os.read(xlogfd, BLCKSZ)
    while n:
	if (os.write(archfd, n) != len(n)):
	    return False
	n = os.read(xlogfd, BLCKSZ)
    if n != "":
        return False

    print >> sys.stderr, _("%s archive written..." % progname)
    try:
	os.close(archfd)
    except OSError:
	return False

    # Should stat the archpath, to check filesize == XLogFileSize

    # Reset the file date/time on the xlog, to maintain the original
    # timing of the xlog final write by PostgreSQL

    try:
	os.close(xlogfd)
    except OSError:
	return False

    return True


def main(_argc, _argv):
    """
    """
    gettext.bindtextdomain("pg_arch", LOCALEDIR)
    gettext.textdomain("pg_arch")

    global progname
    progname = os.path.basename(_argv[0])

    try:
	opts, args = getopt.getopt(_argv[1:], 
		"snt:?V", ("single", "noarchive", "timeout=", "help", "version"))
    except getopt.GetoptError:
	usage()
	sys.exit(2)

    loop = True
    archive = True
    loopTimeout = 30

    for o, a in opts:
	if o in ("--help", "-?"):
	    usage()
	    sys.exit(0)
	elif o in ("--version", "-V"):
	    print "pg_arch.py (PostgreSQL)", PG_VERSION
	    sys.exit(0)
	elif o in ("--single", "-s"):
	    loop = False
	elif o in ("--noarchive", "-n"):
	    archive = False
	elif o in ("--timeout", "-t"):
	    try:
		loopTimeout = long(a)
	    except:
		print >> sys.stderr, _("%s invalid argument for option -t" % progname)
		print >> sys.stderr, _('Try "%s --help" for more information.' % progname)
		sys.exit(1)
	    if loopTimeout < 1:
		print >> sys.stderr, _("%s wait time (-t) must be > 0" % progname)
		sys.exit(1)
	    if loopTimeout > 999:
		print >> sys.stderr, _("%s wait time (-t) must be < 1000" % progname)
		sys.exit(1)
	else:
	    print >> sys.stderr, _('Try "%s --help" for more information.' % progname)
	    sys.exit(1)

    _argv = _argv[1 + len(opts):]

    if len(_argv) == 0:
	print >> sys.stderr, _("%s no archive directory specified" % progname)
	print >> sys.stderr, _('Try "%s --help" for more information.' % progname)
	sys.exit(1)

    global archiveDestDir
    archiveDestDir = _argv[0]
    global dataDir
    
    if len(_argv) == 1:
	dataDir = os.getenv("PGDATA")
    else:
	dataDir = _argv[1]

    if not dataDir:
	print >> sys.stderr, _("%s no data directory specified" % progname)
	print >> sys.stderr, _('Try "%s --help" for more information.' % progname)
	sys.exit(1)

    global XLogDir
    XLogDir = "%s/pg_xlog" % dataDir

    print >> sys.stderr, _("%s Archiving transaction logs from %s to %s" % 
	(progname, XLogDir, archiveDestDir))

    global XLogArchiveDir
    XLogArchiveDir = "%s/pg_rlog" % dataDir

# File/Directory Permissions required:
# pg_arch should run as either:
#  i) database-owning userid i.e. postgres
# ii) another user in the same group as database-owning userid
#
# Permissions required are:
#  XLogDir		r
#  XLogArchiveDir	rw
#  ArchiveDestDir	w
# 
# Security options are:
#  i) add XLogArchiveDir under DataDir
#  	allow access to ArchiveDestDir
# ii) 	chmod 760 DataDir
#      	chmod 760 XLogArchiveDir
# 	chmod 740 XLogDir
# 
# Let's test our access rights to these directories now. At this stage
# all of these directories may be empty, or not, without error.

# check directory XLogDir

# check directory ArchiveDestDir 
# Directory must NOT have World read rights - security hole

# XLogArchive environment creation & connection to PostgreSQL
#
# Currently, there isn't any. If there was, it would go here

    # Main Loop

    is_firstrun = True
    while (loop or is_firstrun):
	is_firstrun = False
	rc, xlog = XLogArchiveXLogs(XLogArchiveDir)
	if rc:
	    print _("%s archive starting for transaction log %s" % (progname, xlog))
	    if (not archive or (archive and CopyXLogtoArchive(xlog))):
		if XLogArchiveComplete(xlog, XLogArchiveDir):
		    print >> sys.stderr, _("%s archive complete for transaction log %s \n" %
				  (progname, xlog))
		else:
		    print >> sys.stderr, _("%s XLogArchiveComplete error" % progname)
		    sys.exit(1)
	    else:
		print >> sys.stderr, _("%s archive copy error" % progname)
		sys.exit(1)

	# if we have copied one file, we do not wait:
	# immediately loop back round and check to see if another is there.
	# If we're too quick....then we wait
	else:
	    print _("%s sleeping..." % progname)
	    time.sleep(loopTimeout)
	    print _("%s .....awake" % progname)

    print _("%s ending" % progname)
    return 0    


def usage():
    print _("%s copies PostgreSQL transaction log files to an archive directory.\n" % progname)
    print _("Usage:\n  %s [OPTIONS]... ARCHIVEDESTDIR [DATADIR]\n" % progname)
    print _("Options:")
    print _("  -t              wait time (secs) between checks for xlogs to archive")
    print _("  -n              no archival, just show xlog file names (for testing)")
    print _("  -s              single execution - archive all full xlogs then stop")
    print _("  --help          show this help, then exit")
    print _("  --version       output version information, then exit")
    print
    print _("If no data directory is specified, the environment variable PGDATA")
    print _("is used. An archive destination must be specified. Default wait=30 secs.")
    print
    print _("Report bugs to <pgsql-bugs@postgresql.org>.")


if __name__ == '__main__':
    main(len(sys.argv), sys.argv)
