async def test_min_size_max_size(dsn):
async with pool.AsyncConnectionPool(dsn, min_size=2) as p:
assert p.min_size == p.max_size == 2
async with pool.AsyncConnectionPool(dsn, min_size=2, max_size=4) as p:
assert p.min_size == 2
assert p.max_size == 4
with pytest.raises(ValueError):
pool.AsyncConnectionPool(dsn, min_size=4, max_size=2)
async def test_wait_ready(dsn, monkeypatch):
delay_connection(monkeypatch, 0.1)
with pytest.raises(pool.PoolTimeout):
async with pool.AsyncConnectionPool(dsn, min_size=4,
num_workers=1) as p:
await p.wait(0.3)
async with pool.AsyncConnectionPool(dsn, min_size=4, num_workers=1) as p:
await p.wait(0.5)
async with pool.AsyncConnectionPool(dsn, min_size=4, num_workers=2) as p:
await p.wait(0.3)
await p.wait(0.0001) # idempotent
async def test_stats_usage(dsn, retries):
async def worker(n):
try:
async with p.connection(timeout=0.3) as conn:
await conn.execute("select pg_sleep(0.2)")
except pool.PoolTimeout:
pass
async for retry in retries:
with retry:
async with pool.AsyncConnectionPool(dsn, min_size=3) as p:
await p.wait(2.0)
ts = [create_task(worker(i)) for i in range(7)]
await asyncio.gather(*ts)
stats = p.get_stats()
assert stats["requests_num"] == 7
assert stats["requests_queued"] == 4
assert 850 <= stats["requests_wait_ms"] <= 950
assert stats["requests_errors"] == 1
assert 1150 <= stats["usage_ms"] <= 1250
assert stats.get("returns_bad", 0) == 0
async with p.connection() as conn:
await conn.close()
await p.wait()
stats = p.pop_stats()
assert stats["requests_num"] == 8
assert stats["returns_bad"] == 1
async with p.connection():
pass
assert p.get_stats()["requests_num"] == 1
async def test_putconn_no_pool(dsn):
async with pool.AsyncConnectionPool(dsn, min_size=1) as p:
conn = await psycopg.AsyncConnection.connect(dsn)
with pytest.raises(ValueError):
await p.putconn(conn)
await conn.close()
async def test_queue_size(dsn):
async def worker(t, ev=None):
try:
async with p.connection():
if ev:
ev.set()
await asyncio.sleep(t)
except pool.TooManyRequests as e:
errors.append(e)
else:
success.append(True)
errors = []
success = []
async with pool.AsyncConnectionPool(dsn, min_size=1, max_waiting=3) as p:
await p.wait()
ev = asyncio.Event()
create_task(worker(0.3, ev))
await ev.wait()
ts = [create_task(worker(0.1)) for i in range(4)]
await asyncio.gather(*ts)
assert len(success) == 4
assert len(errors) == 1
assert isinstance(errors[0], pool.TooManyRequests)
assert p.name in str(errors[0])
assert str(p.max_waiting) in str(errors[0])
assert p.get_stats()["requests_errors"] == 1
async def test_queue_timeout(dsn, retries):
async def worker(n):
t0 = time()
try:
async with p.connection() as conn:
cur = await conn.execute(
"select pg_backend_pid() from pg_sleep(0.2)")
(pid, ) = await cur.fetchone() # type: ignore[misc]
except pool.PoolTimeout as e:
t1 = time()
errors.append((n, t1 - t0, e))
else:
t1 = time()
results.append((n, t1 - t0, pid))
async for retry in retries:
with retry:
results: List[Tuple[int, float, int]] = []
errors: List[Tuple[int, float, Exception]] = []
async with pool.AsyncConnectionPool(dsn, min_size=2,
timeout=0.1) as p:
ts = [create_task(worker(i)) for i in range(4)]
await asyncio.gather(*ts)
assert len(results) == 2
assert len(errors) == 2
for e in errors:
assert 0.1 < e[1] < 0.15
async def test_setup_no_timeout(dsn, proxy):
with pytest.raises(pool.PoolTimeout):
async with pool.AsyncConnectionPool(proxy.client_dsn,
min_size=1,
num_workers=1) as p:
await p.wait(0.2)
async with pool.AsyncConnectionPool(proxy.client_dsn,
min_size=1,
num_workers=1) as p:
await asyncio.sleep(0.5)
assert not p._pool
proxy.start()
async with p.connection() as conn:
await conn.execute("select 1")
async def test_grow(dsn, monkeypatch, retries):
delay_connection(monkeypatch, 0.1)
async def worker(n):
t0 = time()
async with p.connection() as conn:
await conn.execute("select 1 from pg_sleep(0.2)")
t1 = time()
results.append((n, t1 - t0))
async for retry in retries:
with retry:
async with pool.AsyncConnectionPool(dsn,
min_size=2,
max_size=4,
num_workers=3) as p:
await p.wait(1.0)
ts = []
results = []
ts = [create_task(worker(i)) for i in range(6)]
await asyncio.gather(*ts)
want_times = [0.2, 0.2, 0.3, 0.4, 0.4, 0.4]
times = [item[1] for item in results]
for got, want in zip(times, want_times):
assert got == pytest.approx(want, 0.1), times
async def test_queue_timeout_override(dsn, retries):
async def worker(n):
t0 = time()
timeout = 0.25 if n == 3 else None
try:
async with p.connection(timeout=timeout) as conn:
cur = await conn.execute(
"select pg_backend_pid() from pg_sleep(0.2)")
(pid, ) = await cur.fetchone()
except pool.PoolTimeout as e:
t1 = time()
errors.append((n, t1 - t0, e))
else:
t1 = time()
results.append((n, t1 - t0, pid))
async for retry in retries:
with retry:
results = []
errors = []
async with pool.AsyncConnectionPool(dsn, min_size=2,
timeout=0.1) as p:
ts = [create_task(worker(i)) for i in range(4)]
await asyncio.gather(*ts)
assert len(results) == 3
assert len(errors) == 1
for e in errors:
assert 0.1 < e[1] < 0.15
async def test_defaults(dsn):
async with pool.AsyncConnectionPool(dsn) as p:
assert p.min_size == p.max_size == 4
assert p.timeout == 30
assert p.max_idle == 10 * 60
assert p.max_lifetime == 60 * 60
assert p.num_workers == 3
async def test_shrink(dsn, monkeypatch):
from psycopg_pool.pool_async import ShrinkPool
results = []
async def run_hacked(self, pool):
n0 = pool._nconns
await orig_run(self, pool)
n1 = pool._nconns
results.append((n0, n1))
orig_run = ShrinkPool._run
monkeypatch.setattr(ShrinkPool, "_run", run_hacked)
async def worker(n):
async with p.connection() as conn:
await conn.execute("select pg_sleep(0.1)")
async with pool.AsyncConnectionPool(dsn,
min_size=2,
max_size=4,
max_idle=0.2) as p:
await p.wait(5.0)
assert p.max_idle == 0.2
ts = [create_task(worker(i)) for i in range(4)]
await asyncio.gather(*ts)
await asyncio.sleep(1)
assert results == [(4, 4), (4, 3), (3, 2), (2, 2), (2, 2)]
async def test_reset(dsn):
resets = 0
async def setup(conn):
async with conn.transaction():
await conn.execute("set timezone to '+1:00'")
async def reset(conn):
nonlocal resets
resets += 1
async with conn.transaction():
await conn.execute("set timezone to utc")
async with pool.AsyncConnectionPool(dsn, min_size=1, reset=reset) as p:
async with p.connection() as conn:
assert resets == 0
await conn.execute("set timezone to '+2:00'")
await p.wait()
assert resets == 1
async with p.connection() as conn:
cur = await conn.execute("show timezone")
assert (await cur.fetchone()) == ("UTC", )
await p.wait()
assert resets == 2
async def test_stats_measures(dsn):
async def worker(n):
async with p.connection() as conn:
await conn.execute("select pg_sleep(0.2)")
async with pool.AsyncConnectionPool(dsn, min_size=2, max_size=4) as p:
await p.wait(2.0)
stats = p.get_stats()
assert stats["pool_min"] == 2
assert stats["pool_max"] == 4
assert stats["pool_size"] == 2
assert stats["pool_available"] == 2
assert stats["requests_waiting"] == 0
ts = [create_task(worker(i)) for i in range(3)]
await asyncio.sleep(0.1)
stats = p.get_stats()
await asyncio.gather(*ts)
assert stats["pool_min"] == 2
assert stats["pool_max"] == 4
assert stats["pool_size"] == 3
assert stats["pool_available"] == 0
assert stats["requests_waiting"] == 0
await p.wait(2.0)
ts = [create_task(worker(i)) for i in range(7)]
await asyncio.sleep(0.1)
stats = p.get_stats()
await asyncio.gather(*ts)
assert stats["pool_min"] == 2
assert stats["pool_max"] == 4
assert stats["pool_size"] == 4
assert stats["pool_available"] == 0
assert stats["requests_waiting"] == 3
async def test_configure(dsn):
inits = 0
async def configure(conn):
nonlocal inits
inits += 1
async with conn.transaction():
await conn.execute("set default_transaction_read_only to on")
async with pool.AsyncConnectionPool(dsn, min_size=1,
configure=configure) as p:
await p.wait(timeout=1.0)
async with p.connection() as conn:
assert inits == 1
res = await conn.execute("show default_transaction_read_only")
assert (await res.fetchone())[0] == "on"
async with p.connection() as conn:
assert inits == 1
res = await conn.execute("show default_transaction_read_only")
assert (await res.fetchone())[0] == "on"
await conn.close()
async with p.connection() as conn:
assert inits == 2
res = await conn.execute("show default_transaction_read_only")
assert (await res.fetchone())[0] == "on"
async def test_fail_rollback_close(dsn, caplog, monkeypatch):
async with pool.AsyncConnectionPool(dsn, min_size=1) as p:
conn = await p.getconn()
async def bad_rollback():
conn.pgconn.finish()
await orig_rollback()
# Make the rollback fail
orig_rollback = conn.rollback
monkeypatch.setattr(conn, "rollback", bad_rollback)
pid = conn.pgconn.backend_pid
with pytest.raises(psycopg.ProgrammingError):
await conn.execute("wat")
assert conn.pgconn.transaction_status == TransactionStatus.INERROR
await p.putconn(conn)
async with p.connection() as conn2:
assert conn2.pgconn.backend_pid != pid
assert conn2.pgconn.transaction_status == TransactionStatus.IDLE
recs = [
r for r in caplog.records
if r.name.startswith("psycopg") and r.levelno >= logging.WARNING
]
assert len(recs) == 3
assert "INERROR" in recs[0].message
assert "OperationalError" in recs[1].message
assert "BAD" in recs[2].message
async def test_connection_class(dsn):
class MyConn(psycopg.AsyncConnection):
pass
async with pool.AsyncConnectionPool(dsn,
connection_class=MyConn,
min_size=1) as p:
async with p.connection() as conn:
assert isinstance(conn, MyConn)
async def test_connection_not_lost(dsn):
async with pool.AsyncConnectionPool(dsn, min_size=1) as p:
with pytest.raises(ZeroDivisionError):
async with p.connection() as conn:
pid = conn.pgconn.backend_pid
1 / 0
async with p.connection() as conn2:
assert conn2.pgconn.backend_pid == pid
async def test_reopen(dsn):
p = pool.AsyncConnectionPool(dsn)
async with p.connection() as conn:
await conn.execute("select 1")
await p.close()
assert p._sched_runner is None
with pytest.raises(psycopg.OperationalError, match="cannot be reused"):
await p.open()
async def test_closed_putconn(dsn):
p = pool.AsyncConnectionPool(dsn, min_size=1)
async with p.connection() as conn:
pass
assert not conn.closed
async with p.connection() as conn:
await p.close()
assert conn.closed
async def test_close_no_tasks(dsn):
p = pool.AsyncConnectionPool(dsn)
assert not p._sched_runner.done()
for t in p._workers:
assert not t.done()
await p.close()
assert p._sched_runner.done()
for t in p._workers:
assert t.done()
async def test_max_lifetime(dsn):
async with pool.AsyncConnectionPool(dsn, min_size=1,
max_lifetime=0.2) as p:
await asyncio.sleep(0.1)
pids = []
for i in range(5):
async with p.connection() as conn:
pids.append(conn.pgconn.backend_pid)
await asyncio.sleep(0.2)
assert pids[0] == pids[1] != pids[4], pids
async def test_uniform_use(dsn, retries):
async for retry in retries:
with retry:
async with pool.AsyncConnectionPool(dsn, min_size=4) as p:
counts = Counter()
for i in range(8):
async with p.connection() as conn:
await asyncio.sleep(0.1)
counts[id(conn)] += 1
assert len(counts) == 4
assert set(counts.values()) == set([2])
async def test_its_really_a_pool(dsn):
async with pool.AsyncConnectionPool(dsn, min_size=2) as p:
async with p.connection() as conn:
cur = await conn.execute("select pg_backend_pid()")
(pid1, ) = await cur.fetchone()
async with p.connection() as conn2:
cur = await conn2.execute("select pg_backend_pid()")
(pid2, ) = await cur.fetchone()
async with p.connection() as conn:
assert conn.pgconn.backend_pid in (pid1, pid2)
async def test_open_context(dsn):
p = pool.AsyncConnectionPool(dsn, open=False)
assert p.closed
async with p:
assert not p.closed
async with p.connection() as conn:
cur = await conn.execute("select 1")
assert await cur.fetchone() == (1, )
assert p.closed
async def test_closed_getconn(dsn):
p = pool.AsyncConnectionPool(dsn, min_size=1)
assert not p.closed
async with p.connection():
pass
await p.close()
assert p.closed
with pytest.raises(pool.PoolClosed):
async with p.connection():
pass
async def test_broken_reconnect(dsn):
async with pool.AsyncConnectionPool(dsn, min_size=1) as p:
async with p.connection() as conn:
cur = await conn.execute("select pg_backend_pid()")
(pid1, ) = await cur.fetchone()
await conn.close()
async with p.connection() as conn2:
cur = await conn2.execute("select pg_backend_pid()")
(pid2, ) = await cur.fetchone()
assert pid1 != pid2
async def test_configure_badstate(dsn, caplog):
caplog.set_level(logging.WARNING, logger="psycopg.pool")
async def configure(conn):
await conn.execute("select 1")
async with pool.AsyncConnectionPool(dsn, min_size=1,
configure=configure) as p:
with pytest.raises(pool.PoolTimeout):
await p.wait(timeout=0.5)
assert caplog.records
assert "INTRANS" in caplog.records[0].message
async def test_open_no_op(dsn):
p = pool.AsyncConnectionPool(dsn)
try:
assert not p.closed
await p.open()
assert not p.closed
async with p.connection() as conn:
cur = await conn.execute("select 1")
assert await cur.fetchone() == (1, )
finally:
await p.close()
async def test_close_no_tasks(dsn):
p = pool.AsyncConnectionPool(dsn)
assert p._sched_runner and not p._sched_runner.done()
assert p._workers
workers = p._workers[:]
for t in workers:
assert not t.done()
await p.close()
assert p._sched_runner is None
assert not p._workers
for t in workers:
assert t.done()
async def test_spike(dsn, monkeypatch):
# Inspired to https://github.com/brettwooldridge/HikariCP/blob/dev/
# documents/Welcome-To-The-Jungle.md
delay_connection(monkeypatch, 0.15)
async def worker():
async with p.connection():
await asyncio.sleep(0.002)
async with pool.AsyncConnectionPool(dsn, min_size=5, max_size=10) as p:
await p.wait()
ts = [create_task(worker()) for i in range(50)]
await asyncio.gather(*ts)
await p.wait()
assert len(p._pool) < 7
