import pg
import sys
import time
import traceback
import popen2
import popen2
import os
import exceptions

doc = """DumpAll <agl@imperialviolet.org>
-------

This script attempts to contact a postgres DB and dump all the tables within it
using pg_dump

It reads the following varables. They can either be set via a command line
option, or by setting the given env variable. The command line option takes
priority. Given an option called xyz the option to set the value would be
'--xyz=value'

Name		| Env Var		| Default	
--------------------------------------------------------
dbname		| PGDATABASE		
passwd		| PGPASSWORD
host		| PGHOST
port		| PGPORT
user		| PGUSER
path		|			| "." 
pgdumpopt	| 			| ""
logfile		| 			| "dumpall_output"
excfile		| 			| "dumpall_exc"
permissive	|			| false
silent		|			| false
skip		|			| []

@path is the directory where the dumps will be placed. @pgdumpopt is added to
the pg_dump command line. 

Note that this script *only sets* the output options for pg_dump. It doesn't
pass on it's idea of the database host to pg_dump. Also, if pg_dump asks for a
password the whole thing will hang on that, so you may wish to set PGPASSWORD
to something so that doesn't happen.

Note dbname exists because the script needs to connect to a database inorder to
get a list of existing databases.

The script will dump the databases are tar files in the @path directory and
then call bzip2 -9 to compress them.

Output is logged to stdout and to a file called @logfile. The file @excfile is
unlinked at the start of the script and created if an error occurs. Thus if
@excfile exists after the script completes then something went wrong.

If @permissive is set (via '--permissive') then unknown command line options
don't cause an error and the @excfile is never created (it is still deleted
however).

If @silent is set the program produces on output on stdout during normal
operations

Any number of '--skip=value' options can be passed to skip the dumping of the
named databases.

Note (from the pg_dump manpage):
Members of tar archives are limited to a size less than 8 GB.  (This is an
inherent limitation of the tar file format.) Therefore  this  format cannot  be
used  if the textual representation of a table exceeds that size.  """

permissive = 0
silent = 0

def env (env, default):
	try:
		r = os.environ[env]
	except:
		r = default
	return r
		

dbname = env ("PGDATABASE", None)
passwd = env ("PGPASSWORD", None) 
host = env ("PGHOST", None)
port = int (env ("PGPORT", "-1"))
user = env ("PGUSER", None)
path = "."
pgdumpopt = ""
logname = "dumpall_output"
excfile = "dumpall_exc"
skip = []

opts = ["port", "path", "dbname", "user", "host", "passwd", "pgdumpopt", "logname", "excfile"]

try:
	os.unlink ("dumpall_exc")
except:
	pass

db = 0

def flag_exc ():
	global excfile
	if (permissive == 0):
		open (excfile, "w+").close ()

def log (level, string):
	string = time.strftime ("%d/%m/%y %H:%M:%S", time.localtime()) + " " + string 
	if (level == 'EXC'):
		string = "*** " + string
	if (silent == 0):
		print string
	try:
		logfile.write (string + "\n")
	except:
		pass
	if (level == 'WARNING' and permissive == 0):
		try:
			logfile.write (string + "\n")
		except:
			pass
		print "Warning in non-permissive mode. Fatal."
		flag_exc ()
		sys.exit (1)
	logfile.flush ()

def do_arguments():
	global opts,skip,permissive,silent
	for x in sys.argv[1:]:
		if (x == '--permissive'):
			permissive = 1
			continue
		if (x == '--silent'):
			silent = 1
			continue
		if (x == '--help'):
			print doc
			sys.exit (1)
		if (x == '-h'):
			print doc
			sys.exit (1)
	
		i = x.find ('=')
		if (i == -1):
			log ('WARNING', "Option %s not processed" % x)
			continue
		option = x[2:i]
		value = x[i + 1:]
		if (option == 'skip'):
			skip.append (value)
			continue
		if (not option in opts):
			log ('WARNING', "Option %s not recognised" % option)
		if (option == "port"):
			globals()[option] = int (value)
		else:
			globals()[option] = value

def do_connect ():
	global passwd, db
	db = pg.DB(dbname = dbname, passwd = passwd, host = host, port = port, user = user)

	passwd = ""

def main():
	global logfile, db, path, pgdumpopt, opts, skip
	do_arguments ()
	logfile = open ("dumpall_output", "w+")
	logfile.write ("dumpall starting at " + time.strftime ("%d/%m/%y %H:%M", time.localtime()) + "\n\n")
	do_connect ()

	for x in opts:
		if (x != "passwd"):
			log ('INFO', "Option %s=%s" % (x, str(globals()[x])))
	
	dbs = db.get_databases ()
	for x in skip:
		try:
			dbs.remove (x)
		except:
			pass
	log ('INFO', "Dumping the following databases")
	log ('INFO', str(dbs))
	log ('INFO', "")

	for d in dbs:
		pear_shaped = 0
		log ('INFO', "Now dumping %s" % d)
		dump = popen2.Popen3 ("pg_dump -b -Ft -f %s %s %s" % (pgdumpopt, path + "/" + d + ".tar", d), 1)
		a = dump.wait ()
		if (a != 0):
			log ('EXC', "pg_dump returned %d" % a)
			pear_shaped = 1
			flag_exc ()
		err = dump.childerr.read()
		if (len (err) != 0):
			pear_shaped = 1
			log ('EXC', "pg_dump sent the following to stderr:")
			log ('EXC', err)
			flag_exc ()
		del err
		del dump
		if (pear_shaped == 1):
			log ('INFO', "Unlinking the tar file")
			os.unlink (path + "/" + d + ".tar")
			continue
		log ('INFO', "Compressing %s" % d)
		try:
			os.unlink (path + "/" + d + ".tar.bz2")
		except:
			pass
		bz = popen2.Popen3 ("bzip2 -9 %s" % (path + "/" + d + ".tar"), 1)
		a = bz.wait ()
		if (a != 0):
			log ('EXC', "bzip2 returned %d" % a)
			flag_exc ()
		err = bz.childerr.read ()
		if (len (err) != 0):
			log ('EXC', "bzip2 sent the following to stderr:")
			log ('EXC', err)
			flag_exc ()
		del bz
		del err
		

if (__name__ == "__main__"):
	try:
		main ()
	except:
		if (sys.exc_info()[0] == exceptions.SystemExit):
			raise
		flag_exc ()
		for x in sys.exc_info ():
			print x
		print traceback.print_tb(sys.exc_info()[2])
		for x in sys.exc_info ():
			logfile.write (str(x) + "\n")
		traceback.print_tb(sys.exc_info()[2], file=logfile)