#!/usr/bin/env python

import argparse
import datetime
import psycopg2
import psycopg2.extras

from random import shuffle
from time import sleep

def get_attributes(conn, relname):
	'fetch attributes for a relation (excluding system ones)'

	cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
	cur.execute('SELECT attname FROM pg_attribute WHERE attrelid = %(relname)s::regclass AND attnum > 0', {'relname' : relname})

	attrs = cur.fetchall()
	cur.close()

	return [a['attname'] for a in attrs]


def shuffle_attrs(conn, relname, attrs):
	'shuffle the lognums (randomly)'

	shuffle(attrs)

	cur = conn.cursor()

	# conn.begin()

	for lognum in range(len(attrs)):
		cur.execute('UPDATE pg_attribute SET attlognum = %(lognum)s WHERE attrelid = %(relname)s::regclass AND attname = %(name)s', {'lognum' : (lognum+1), 'relname' : relname, 'name' : attrs[lognum]})

	conn.commit()

	cur.close()

	return attrs


def fetch_attrs(conn, tables):

	rels = {}
	for t in tables:
		rels.update({t : get_attributes(conn, t)})

	return rels


def parse_args():
	''
	parser = argparse.ArgumentParser(description='attlognum randomizer')
	parser.add_argument('-t', '--tables', metavar='T', nargs='+', required=True, help='name of table to randomize')
	parser.add_argument('-q', '--query',  metavar='QUERY', required=False, nargs='*', type=str, help='query to execute')

	parser.add_argument('--host', dest='host', default='localhost', help='hostname of the database')
	parser.add_argument('--db',   dest='dbname', default='test', help='name of the database')

	parser.add_argument('-i', '--iterations',  metavar='COUNT', type=int, default=100, help='number of iterations')
	parser.add_argument('-T', '--time',  metavar='SECONDS', type=int, default=0, help='number of seconds')

	parser.add_argument('--init-script', type=str, help='initialization script (create tables etc.)')
	parser.add_argument('--test-script', type=str, help='test script (may be multiple queries)')

	return parser.parse_args()

if __name__ == '__main__':
	
	args = parse_args()

	iter_count = 1
	start_time = datetime.datetime.now()

	while True:

		print "========== iteration",iter_count,"=========="

		# recreate the database (drop if exists)
		conn = psycopg2.connect('host=%s dbname=postgres' % (args.host,))
		conn.autocommit = True
		cur = conn.cursor()
		cur.execute('DROP DATABASE IF EXISTS %s' % (args.dbname,))
		cur.execute('CREATE DATABASE %s' % (args.dbname,))
		cur.close()
		conn.close()

		# now, run the queries (if supplied), or the test script
		try:

			# connect to the proper database
			conn = psycopg2.connect('host=%s dbname=%s' % (args.host, args.dbname))
			conn.autocommit = True

			# run init script (if supplied)
			if args.init_script:
				with open(args.init_script, 'r') as init_script:
					data = init_script.read()
					cur = conn.cursor()
					cur.execute(data)
					cur.close()

			# fetch attributes for the tables
			rels = fetch_attrs(conn, args.tables)

			# randomize the attlognums
			for t in rels:
				attrs = shuffle_attrs(conn, t, rels[t])
				print t,':',attrs

			cur = conn.cursor()

			if args.query:
				for q in args.query:
					print "running query:",q
					cur.execute(q)

			if args.test_script:
				with open(args.test_script, 'r') as test_script:
					data = test_script.read()
					cur = conn.cursor()
					cur.execute(data)
					cur.close()

			cur.close()

		except Exception as ex:
			print "FAILED: ", str(ex)
			# sleep for a few seconds, in case we need a recovery
			sleep(10)

		finally:
			conn.close()

		iter_count += 1
		duration = (datetime.datetime.now() - start_time).total_seconds()

		if (args.iterations > 0) and (iter_count > args.iterations):
			break

		if (args.time > 0) and (duration > args.time):
			break