#!/bin/env python

import datetime
import psycopg2
import sys

if __name__ == '__main__':

	if len(sys.argv) < 4:
		print "ERORR: not enough parameters"
		print "HINT:  drop-test.py num-of-tables drop-num num-of-indexes 'connection string'"
		sys.exit(1)

	ntables = int(sys.argv[1])
	nlimit  = int(sys.argv[2])
	nindex  = int(sys.argv[3])
	connstr = str(sys.argv[4])
	debug   = False

	conn = psycopg2.connect(connstr)
	cur  = conn.cursor()

	# print 'creating %s tables' % (ntables,)
	start = datetime.datetime.now()
	for i in range(ntables):
		cur.execute('CREATE TABLE tab_%s (id INT)' % (i,))
		for j in range(nindex):
			cur.execute('CREATE INDEX idx_%s_%s ON tab_%s(id)' % (i,j,i))
		conn.commit();
		if (i % 1000 == 0) and debug:
			print '  tables created: %s' % (i,)
	conn.commit()
	end = datetime.datetime.now()
	# print '  all tables created in %s seconds' % ((end-start).total_seconds(),)

	# set to autocommit mode
	conn.autocommit = True
	start = datetime.datetime.now()
	# print 'dropping %s tables one by one ...' % (ntables,)
	for i in range(ntables):
		cur.execute('DROP TABLE tab_%s' % (i,))
		if (i % 1000 == 0) and debug:
			print '  tables dropped: %s' % (i,)

	end = datetime.datetime.now()
	print 'dropped one-by-one in %s seconds' % ((end-start).total_seconds(),)

	# cancel the autocommit mode
	conn.autocommit = False

	# recreate tables
	# print 'creating %s tables' % (ntables,)
	start = datetime.datetime.now()
	for i in range(ntables):
		cur.execute('CREATE TABLE tab_%s (id INT)' % (i,))
		for j in range(nindex):
			cur.execute('CREATE INDEX idx_%s_%s ON tab_%s(id)' % (i,j,i))
		conn.commit();
		if (i % 1000 == 0) and debug:
			print '  tables created: %s' % (i,)

	conn.commit()
	end = datetime.datetime.now()
	# print '  all tables created in %s seconds' % ((end-start).total_seconds(),)

	# drop the tables in batches
	# print 'dropping %s tables in batches by %s...' % (ntables,nlimit)
	start = datetime.datetime.now()
	for i in range(ntables):
		cur.execute('DROP TABLE tab_%s' % (i,))
		if i % nlimit == 0:
			conn.commit()
			if debug:
				print '  tables dropped: %s' % (i,)
	conn.commit()
	end = datetime.datetime.now()
	print 'dropped in a TX in %s seconds' % ((end-start).total_seconds(),)

	conn.close()