import time,psycopg2

NUNIQUE = 10000
NDUPES  = 1

def bench( f, *args, **kwargs ):
    t = time.time()
    try:
        return f( *args, **kwargs )
    finally:
        dt = time.time()-t
        c = db.cursor()
        c.execute( "SELECT pg_relation_size( 'testins' )")
        size, = c.fetchone()
        c.execute( "SELECT count(*) FROM testins")
        rows, = c.fetchone()
        print("%s: %.03fs %d rows %10.03f kB" % (f.__name__, dt, rows, size/1024))

def insert_only( ):
    c = db.cursor()
    for id in range( NUNIQUE ):
        for n in range( NDUPES ):
            try:
                c.execute("INSERT INTO testins (id,t) VALUES (%s,'hello, world')", (id,) )
                c.execute("COMMIT")
            except Exception as e:
                # print(e)
                c.execute("ROLLBACK")

def select_then_insert( ):
    c = db.cursor()
    for id in range( NUNIQUE ):
        for n in range( NDUPES ):
            try:
                c.execute("SELECT * FROM testins WHERE id=%s", (id,) )
                if not c.fetchall():
                    c.execute("INSERT INTO testins (id,t) VALUES (%s,'hello, world')", (id,) )
                c.execute("COMMIT")
            except Exception as e:
                # print(e)
                c.execute("ROLLBACK")

def insert_select( ):
    c = db.cursor()
    for id in range( NUNIQUE ):
        for n in range( NDUPES ):
            try:
                c.execute("""INSERT INTO testins (id,t) SELECT %s,'hello, world'
                    WHERE NOT EXISTS( SELECT FROM testins WHERE id=%s ) RETURNING id
                    """, (id,id) )
            except Exception as e:
                # print(e)
                c.execute("ROLLBACK")

def on_conflict( ):
    c = db.cursor()
    for id in range( NUNIQUE ):
        for n in range( NDUPES ):
            try:
                c.execute("INSERT INTO testins (id,t) VALUES (%s,'hello, world') ON CONFLICT (id) DO NOTHING RETURNING id", (id,) )
                c.execute("COMMIT")
            except Exception as e:
                # print(e)
                c.execute("ROLLBACK")



if __name__ == "__main__":
    db = psycopg2.connect("dbname=test")
    c = db.cursor()
    try:
        c.execute( "CREATE TABLE testins( id INTEGER PRIMARY KEY, t TEXT )");
        c.execute("COMMIT")
    except Exception as e:
        print(e)
        c.execute("ROLLBACK")

    c.execute( "TRUNCATE testins;" )
    r = bench( on_conflict )

    c.execute( "TRUNCATE testins;" )
    r = bench( insert_only )

    c.execute( "TRUNCATE testins;" )
    r = bench( select_then_insert )

    c.execute( "TRUNCATE testins;" )
    r = bench( insert_select )