Ticket #9964: bench.py

from psycopg2 import connect as _pg
from sqlite3 import connect as _lite
from MySQLdb import connect as _mysql

def mysql():
    con = _mysql()
    con.select_db('shai')
    return con

def sqlite():
    return _lite('sqlite.db')

def pg():
    return _pg("")

def prepare(connection):
    cursor = connection.cursor()
    cursor.execute("""create table test(
                            testid int primary key, 
                            test2 char(10) not null, 
                            test3 varchar(50) not null
                            )""")
def fill_up(connection):
    cursor = connection.cursor()
    cursor.execute("delete from test where 1=1")
    s = '0123456789'
    t = 'abc'
    for i in range(100000):
        stmt = "insert into test values(%d, '%s', '%s')" % \
                (i, s, t)
        cursor.execute(stmt)
        s = s[-1:]+s[:-1] # rotate
        t = (t+ 'x'+ str(i))[-49:]
    connection.commit()

def warm_up(connection):
    cursor = connection.cursor()
    for i in range(1, 100000, 1017):
        cursor.execute("select * from test where testid=%d" % (i,))
        cursor.fetchall()
    connection.rollback()
    for i in range(1, 100000, 1017):
        cursor.execute("select * from test where testid=%d" % (i,))
        cursor.fetchall()
    connection.commit()
    
def bench(connection, end_tran):
    cursor = connection.cursor()
    for i in range(7, 99990, 69):
        cursor.execute("select * from test where testid=%d" % (i,))
        cursor.fetchall()
        cursor.execute("select * from test where testid=%d" % (i+1,))
        cursor.fetchall()
        cursor.execute("select * from test where testid=%d" % (i+6,))
        cursor.fetchall()
        end_tran()

def leave_open(connection):
    def noop(): pass
    bench(connection,  noop)

def commit(connection):
    bench(connection,  connection.commit)

def rollback(connection):
    bench(connection,  connection.rollback)
    
if __name__=='__main__':
    fill_up(mysql())
    fill_up(sqlite())
    fill_up(pg())
    print "styles:"
    print "pg", _pg.__module__.paramstyle
    print "mysql", _mysql.__module__.paramstyle
    print "sqlite", _lite.__module__.paramstyle