Test multi-statement TX behaviour / latency of psycopg driver

psycopg_latency_test.py

import psycopg
import random
import time

CONNSTR = "postgresql://james:[email protected]:5432/postgres?sslmode=require"
SCALE = 10
BASE_SQL = """
BEGIN;
SELECT abalance FROM pgbench_accounts WHERE aid = {aid};
SELECT abalance FROM pgbench_accounts WHERE aid = {aid2};
SELECT abalance FROM pgbench_accounts WHERE aid = {aid3};
COMMIT;
"""
LOOPS = 10
print(f"Running SQL {LOOPS} times: \n{BASE_SQL}")

def get_sql():
    return BASE_SQL.format(aid=random.randint(1, int(SCALE*1e5)), aid2=random.randint(1, int(SCALE*1e5)), aid3=random.randint(1, int(SCALE*1e5)))

response_times = []

with psycopg.connect(CONNSTR, autocommit=True) as conn:
# with psycopg.connect(CONNSTR) as conn:  # Adds 1 round-trip

    with conn.cursor() as cur:

        for loop in range(LOOPS):
            sql = get_sql()
            t1 = time.time()
            cur.execute(sql)
            t2 = time.time()
            
            response_times.append(t2 - t1)

print("Avg. response time:", sum(response_times) / len(response_times))

select-only-multi.pgbench

-- select-only: <builtin: select only>
\set aid random(1, 100000 * :scale)
\set aid2 random(1, 100000 * :scale)
\set aid3 random(1, 100000 * :scale)
BEGIN;
SELECT abalance FROM pgbench_accounts WHERE aid = :aid;
SELECT abalance FROM pgbench_accounts WHERE aid = :aid2;
SELECT abalance FROM pgbench_accounts WHERE aid = :aid3;
COMMIT;