Skip to content

Instantly share code, notes, and snippets.

@lucky-verma

lucky-verma/benchmark_pg.py

Forked from pwtail/benchmark_pg.py
Created May 16, 2024 23:30
Show Gist options
  • Save lucky-verma/7f9f0cd4890615fdd56d5c814f5cf787 to your computer and use it in GitHub Desktop.
Save lucky-verma/7f9f0cd4890615fdd56d5c814f5cf787 to your computer and use it in GitHub Desktop.
Download ZIP
Benchmark pg
Raw
benchmark_pg.py
import asyncio
import concurrent
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
import sys
import random
import time
import uvloop
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
table = """
CREATE TABLE customer (
id SERIAL NOT NULL,
name VARCHAR(255),
description VARCHAR(255),
q INTEGER,
p INTEGER,
x INTEGER,
y INTEGER,
z INTEGER,
PRIMARY KEY (id)
)
"""
drop = "DROP TABLE IF EXISTS customer"
ids = list(range(10_000))
tests = 10000
data = [
dict(
id=i,
name="c%d" % i,
description="c%d" % i,
q=i * 10,
p=i * 20,
x=i * 30,
y=i * 40,
)
for i in ids
]
insert = """
INSERT INTO customer (id, name, description, q, p, x, y) VALUES
(%(id)s, %(name)s, %(description)s, %(q)s, %(p)s, %(x)s, %(y)s)
"""
select = """
SELECT customer.id, customer.name, customer.description, customer.q,
customer.p, customer.x, customer.y, customer.z
FROM customer
WHERE customer.id = %(id)s
"""
url = "postgresql://postgres:postgres@localhost/test"
@contextmanager
def time_log(message):
start = time.monotonic()
result = {'time': None}
yield result
end = time.monotonic()
result['time'] = end-start
print(f"Run {message} in {end-start} s")
def run_psycopg2():
print("Running psycopg2")
import psycopg2
with psycopg2.connect(url) as conn:
with conn.cursor() as cursor:
cursor.execute(drop)
cursor.execute(table)
cursor.executemany(insert, data)
conn.commit()
def run(_i):
print(f"Insert done. Running {tests} queries")
to_query = random.choices(ids, k=tests)
with psycopg2.connect(url) as conn:
with time_log("psycopg2") as res:
for id_ in to_query:
with conn.cursor() as cursor:
cursor.execute(select, {"id": id_})
cursor.fetchall()
# conn.rollback()
return res['time']
with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
results = executor.map(run, range(5))
print(f'Results: {list(results)}')
with psycopg2.connect(url) as conn:
with conn.cursor() as cursor:
cursor.execute(drop)
conn.commit()
def run_psycopg():
print("Running psycopg sync")
import psycopg
with psycopg.connect(url) as conn:
with conn.cursor() as cursor:
cursor.execute(drop)
cursor.execute(table)
cursor.executemany(insert, data)
conn.commit()
def run(_i):
print(f"Insert done. Running {tests} queries")
to_query = random.choices(ids, k=tests)
with psycopg.connect(url) as conn:
with time_log("psycopg") as res:
for id_ in to_query:
cursor = conn.execute(select, {"id": id_})
cursor.fetchall()
cursor.close()
# conn.rollback()
return res['time']
with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
results = executor.map(run, range(5))
print(f'Results: {list(results)}')
with psycopg.connect(url) as conn:
with conn.cursor() as cursor:
cursor.execute(drop)
conn.commit()
async def run_psycopg_async():
print("Running psycopg async")
import psycopg
conn: psycopg.AsyncConnection
async with await psycopg.AsyncConnection.connect(url) as conn:
async with conn.cursor() as cursor:
await cursor.execute(drop)
await cursor.execute(table)
await cursor.executemany(insert, data)
await conn.commit()
async def run():
print(f"Insert done. Running {tests} queries")
to_query = random.choices(ids, k=tests)
async with await psycopg.AsyncConnection.connect(url) as conn:
with time_log("psycopg_async") as res:
for id_ in to_query:
cursor = await conn.execute(select, {"id": id_})
await cursor.fetchall()
await cursor.close()
# await conn.rollback()
return res['time']
tasks = [run() for i in range(5)]
results = await asyncio.gather(*tasks)
print(results)
async with await psycopg.AsyncConnection.connect(url) as conn:
async with conn.cursor() as cursor:
await cursor.execute(drop)
await conn.commit()
async def run_asyncpg(tests=tests):
print("Running asyncpg")
import asyncpg
places = dict(
id="$1", name="$2", description="$3", q="$4", p="$5", x="$6", y="$7"
)
a_insert = insert % places
a_select = select % {"id": "$1"}
conn: asyncpg.Connection = await asyncpg.connect(url)
async with conn.transaction():
await conn.execute(drop)
await conn.execute(table)
await conn.executemany(a_insert, [tuple(d.values()) for d in data])
await conn.close()
async def run():
print(f"Insert done. Running {tests} queries")
to_query = random.choices(ids, k=tests)
conn: asyncpg.Connection = await asyncpg.connect(url)
with time_log("asyncpg") as res:
for id_ in to_query:
# tr = conn.transaction()
# await tr.start()
await conn.fetch(a_select, id_)
# await tr.rollback()
await conn.close()
return res['time']
tasks = [run() for i in range(5)]
results = await asyncio.gather(*tasks)
print(results)
conn: asyncpg.Connection = await asyncpg.connect(url)
async with conn.transaction():
await conn.execute(drop)
await conn.close()
for name in sys.argv[1:]:
if name == "psycopg2":
run_psycopg2()
elif name == "psycopg":
run_psycopg()
elif name == "psycopg_async":
if hasattr(asyncio, "WindowsSelectorEventLoopPolicy"):
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
asyncio.run(run_psycopg_async())
elif name == "asyncpg":
asyncio.run(run_asyncpg())
else:
print("unknown driver", sys.argv[1])
@lucky-verma
Copy link
Author 

$ python benchmark_pg.py psycopg2 psycopg psycopg_async asyncpg
Running psycopg2
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Run psycopg2 in 3.357030812000062 s
Run psycopg2 in 3.3637375739999698 s
Run psycopg2 in 3.3649347389999775 s
Run psycopg2 in 3.3681958050000276 s
Run psycopg2 in 3.3674273870000206 s
Results: [3.3637375739999698, 3.3681958050000276, 3.3649347389999775, 3.3674273870000206, 3.357030812000062]
Running psycopg sync
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Run psycopg in 16.48907318099998 s
Run psycopg in 16.499815033000004 s
Run psycopg in 16.52971814 s
Run psycopg in 16.533924019000096 s
Run psycopg in 16.57415224099998 s
Results: [16.499815033000004, 16.48907318099998, 16.57415224099998, 16.52971814, 16.533924019000096]
Running psycopg async
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Run psycopg_async in 4.8395525170000155 s
Run psycopg_async in 4.847907569999961 s
Run psycopg_async in 4.847376985999972 s
Run psycopg_async in 4.849904216000027 s
Run psycopg_async in 4.857496874000049 s
[4.847907569999961, 4.849904216000027, 4.8395525170000155, 4.857496874000049, 4.847376985999972]
Running asyncpg
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Insert done. Running 10000 queries
Run asyncpg in 2.7885446159999674 s
Run asyncpg in 2.8417371619999585 s
Run asyncpg in 2.9000188850000086 s
Run asyncpg in 2.9369143659999963 s
Run asyncpg in 2.9621391139999105 s
[2.9621391139999105, 2.8417371619999585, 2.9369143659999963, 2.7885446159999674, 2.9000188850000086]

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment