drop-test.py

text/x-python

Download

#!/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 'connection string'"
		sys.exit(1)
	
	ntables = int(sys.argv[1])
	nlimit  = int(sys.argv[2])
	connstr = str(sys.argv[3])
	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,))
		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 '  all tables dropped 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,))
		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 '  all tables dropped in %s seconds' % ((end-start).total_seconds(),)
	
	conn.close()