psycopg2/sandbox/leak.test.py
2008-05-27 17:40:19 +02:00

84 lines
2.1 KiB
Python

"""
script: test_leak.py
This script attempts to repeatedly insert the same list of rows into
the database table, causing a duplicate key error to occur. It will
then roll back the transaction and try again.
Database table schema:
-- CREATE TABLE t (foo TEXT PRIMARY KEY);
There are two ways to run the script, which will launch one of the
two functions:
# leak() will cause increasingly more RAM to be used by the script.
$ python <script_nam> leak
# noleak() does not have the RAM usage problem. The only difference
# between it and leak() is that 'rows' is created once, before the loop.
$ python <script_name> noleak
Use Control-C to quit the script.
"""
import sys
import psycopg2
DB_NAME = 'test'
connection = psycopg2.connect(database=DB_NAME)
cursor = connection.cursor()
# Uncomment the following if table 't' does not exist
create_table = """CREATE TABLE t (foo TEXT PRIMARY KEY)"""
cursor.execute(create_table)
insert = """INSERT INTO t VALUES (%(foo)s)"""
def leak():
"""rows created in each loop run"""
count = 0
while 1:
try:
rows = []
for i in range(1, 100):
row = {'foo': i}
rows.append(row)
count += 1
print "loop count:", count
cursor.executemany(insert, rows)
connection.commit()
except psycopg2.IntegrityError:
connection.rollback()
def noleak():
"""rows created once, before the loop"""
rows = []
for i in range(1, 100):
row = {'foo': i}
rows.append(row)
count = 0
while 1:
try:
count += 1
print "loop count:", count
cursor.executemany(insert, rows)
connection.commit()
except psycopg2.IntegrityError:
connection.rollback()
usage = "%s requires one argument: 'leak' or 'noleak'" % sys.argv[0]
try:
if 'leak' == sys.argv[1]:
run_function = leak
elif 'noleak' == sys.argv[1]:
run_function = noleak
else:
print usage
sys.exit()
except IndexError:
print usage
sys.exit()
# Run leak() or noleak(), whichever was indicated on the command line
run_function()