mirror of
https://github.com/simonw/datasette.git
synced 2026-10-19 15:56:58 +02:00
791ce7a3d9Replace the per-thread cached read connections with a pool. Each execute_fn()/execute() callback checks a connection out for its duration and gets a LeasedConnection proxy that raises ConnectionLeaseError if used after the callback returns. New settings: max_open_connections (global LRU cap, default 128), connection_idle_timeout (seconds, default 60, enforced by one process-wide reaper thread plus opportunistic reaping on checkout), connection_pool_wait_ms (soft cap: wait then exceed) and pool_read_connections (off = fresh connection per query, a control). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EwkLo2rJWidpkDVT3HkvNv
284 lines
9 KiB
Python
284 lines
9 KiB
Python
import asyncio
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
|
|
import pytest
|
|
import sqlite_utils
|
|
|
|
from datasette.app import Datasette
|
|
from datasette.connection_pool import ConnectionLeaseError
|
|
|
|
|
|
def _make_dbs(tmp_path, n):
|
|
paths = []
|
|
for i in range(n):
|
|
path = tmp_path / f"db{i}.db"
|
|
conn = sqlite3.connect(path)
|
|
conn.execute("create table t (id integer primary key, v text)")
|
|
conn.execute("insert into t (v) values ('x')")
|
|
conn.commit()
|
|
conn.close()
|
|
paths.append(str(path))
|
|
return paths
|
|
|
|
|
|
def _file_conns(db):
|
|
return [c for c in db._all_connections]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("num_sql_threads", [0, 2])
|
|
async def test_lease_expires_after_callback(tmp_path, num_sql_threads):
|
|
(path,) = _make_dbs(tmp_path, 1)
|
|
ds = Datasette([path], settings={"num_sql_threads": num_sql_threads})
|
|
db = ds.get_database("db0")
|
|
stash = {}
|
|
|
|
def callback(conn):
|
|
stash["conn"] = conn
|
|
assert isinstance(conn, sqlite3.Connection)
|
|
return conn.execute("select count(*) from t").fetchone()[0]
|
|
|
|
assert await db.execute_fn(callback) == 1
|
|
with pytest.raises(ConnectionLeaseError, match="db0"):
|
|
stash["conn"].execute("select 1")
|
|
with pytest.raises(ConnectionLeaseError):
|
|
stash["conn"].row_factory = None
|
|
# A second lease of the same underlying connection gets a new proxy, so
|
|
# the stale one still fails while another callback is using it
|
|
seen = {}
|
|
|
|
def callback2(conn):
|
|
with pytest.raises(ConnectionLeaseError):
|
|
stash["conn"].execute("select 1")
|
|
seen["ok"] = True
|
|
|
|
await db.execute_fn(callback2)
|
|
assert seen["ok"]
|
|
ds.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_returning_cursor_raises(tmp_path):
|
|
(path,) = _make_dbs(tmp_path, 1)
|
|
ds = Datasette([path])
|
|
db = ds.get_database("db0")
|
|
with pytest.raises(ConnectionLeaseError, match="returned a Cursor"):
|
|
await db.execute_fn(lambda conn: conn.execute("select * from t"))
|
|
# Pool is still usable afterwards
|
|
assert (await db.execute("select count(*) from t")).single_value() == 1
|
|
ds.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_leased_connection_compatibility(tmp_path):
|
|
(path,) = _make_dbs(tmp_path, 1)
|
|
ds = Datasette([path])
|
|
db = ds.get_database("db0")
|
|
|
|
def callback(conn):
|
|
out = {}
|
|
out["tables"] = sqlite_utils.Database(conn).table_names()
|
|
out["columns"] = list(sqlite_utils.Database(conn)["t"].columns_dict)
|
|
target = sqlite3.connect(":memory:")
|
|
conn.backup(target)
|
|
out["backup"] = target.execute("select count(*) from t").fetchone()[0]
|
|
conn.create_function("double", 1, lambda x: x * 2)
|
|
out["fn"] = conn.execute("select double(21)").fetchone()[0]
|
|
cursor = conn.cursor()
|
|
cursor.execute("select v from t")
|
|
out["cursor"] = cursor.fetchall()[0][0]
|
|
return out
|
|
|
|
assert await db.execute_fn(callback) == {
|
|
"tables": ["t"],
|
|
"columns": ["id", "v"],
|
|
"backup": 1,
|
|
"fn": 42,
|
|
"cursor": "x",
|
|
}
|
|
ds.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("num_sql_threads", [0, 3])
|
|
async def test_global_cap_evicts_least_recently_used(tmp_path, num_sql_threads):
|
|
paths = _make_dbs(tmp_path, 6)
|
|
ds = Datasette(
|
|
paths,
|
|
settings={"num_sql_threads": num_sql_threads, "max_open_connections": 2},
|
|
)
|
|
pool = ds._read_pool
|
|
for i in range(6):
|
|
await ds.get_database(f"db{i}").execute("select 1")
|
|
assert pool.snapshot()["open"] <= 2
|
|
# The two most recently used databases are the ones still open
|
|
open_dbs = {
|
|
name
|
|
for name in (f"db{i}" for i in range(6))
|
|
if ds.get_database(name)._read_pool_state.open
|
|
}
|
|
assert open_dbs == {"db4", "db5"}
|
|
assert pool.stats["evicted_lru"] >= 4
|
|
for i in range(6):
|
|
db = ds.get_database(f"db{i}")
|
|
assert len(_file_conns(db)) == db._read_pool_state.open
|
|
ds.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_prepare_connection_runs_for_every_pooled_connection(tmp_path):
|
|
paths = _make_dbs(tmp_path, 3)
|
|
ds = Datasette(paths, settings={"max_open_connections": 1})
|
|
calls = []
|
|
original = ds._prepare_connection
|
|
|
|
def counting(conn, database):
|
|
calls.append(database)
|
|
return original(conn, database)
|
|
|
|
ds._prepare_connection = counting
|
|
# Pool may have been created already with the original bound method
|
|
ds._read_pool._prepare_connection = counting
|
|
for _ in range(2):
|
|
for i in range(3):
|
|
assert (
|
|
await ds.get_database(f"db{i}").execute("select v from t")
|
|
).single_value() == "x"
|
|
# Cap of 1 means every switch of database reopens and re-prepares
|
|
assert calls == ["db0", "db1", "db2"] * 2
|
|
ds.close()
|
|
|
|
|
|
def test_idle_reaper_closes_connections_without_event_loop(tmp_path):
|
|
paths = _make_dbs(tmp_path, 3)
|
|
ds = Datasette(paths, settings={"num_sql_threads": 0})
|
|
pool = ds._read_pool
|
|
pool.idle_timeout = 0.2
|
|
|
|
async def query_all():
|
|
for i in range(3):
|
|
await ds.get_database(f"db{i}").execute("select 1")
|
|
|
|
asyncio.run(query_all())
|
|
assert pool.snapshot()["open"] == 3
|
|
# No event loop is running now - the reaper thread does the work
|
|
deadline = time.monotonic() + 5
|
|
while pool.snapshot()["open"] and time.monotonic() < deadline:
|
|
time.sleep(0.05)
|
|
assert pool.snapshot()["open"] == 0
|
|
assert pool.stats["expired"] == 3
|
|
for i in range(3):
|
|
assert _file_conns(ds.get_database(f"db{i}")) == []
|
|
ds.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pooling_disabled_opens_connection_per_query(tmp_path):
|
|
(path,) = _make_dbs(tmp_path, 1)
|
|
ds = Datasette([path], settings={"pool_read_connections": False})
|
|
db = ds.get_database("db0")
|
|
for _ in range(5):
|
|
await db.execute("select 1")
|
|
pool = ds._read_pool
|
|
assert pool.stats["opened"] == 5
|
|
assert pool.snapshot()["open"] == 0
|
|
assert _file_conns(db) == []
|
|
ds.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("wait_ms,expect_exceeded", [(0, True), (10000, False)])
|
|
async def test_cap_exhausted_policy(tmp_path, wait_ms, expect_exceeded):
|
|
paths = _make_dbs(tmp_path, 3)
|
|
ds = Datasette(
|
|
paths,
|
|
settings={
|
|
"num_sql_threads": 3,
|
|
"max_open_connections": 1,
|
|
"connection_pool_wait_ms": wait_ms,
|
|
},
|
|
)
|
|
barrier = threading.Barrier(3, timeout=0.5)
|
|
|
|
def slow(conn):
|
|
try:
|
|
# With a hard cap only one callback can run at a time
|
|
barrier.wait()
|
|
except threading.BrokenBarrierError:
|
|
pass
|
|
time.sleep(0.05)
|
|
return conn.execute("select 1").fetchone()[0]
|
|
|
|
results = await asyncio.gather(
|
|
*[ds.get_database(f"db{i}").execute_fn(slow) for i in range(3)]
|
|
)
|
|
assert results == [1, 1, 1]
|
|
pool = ds._read_pool
|
|
if expect_exceeded:
|
|
assert pool.stats["exceeded_cap"] >= 1
|
|
assert pool.stats["peak_open"] > 1
|
|
else:
|
|
assert pool.stats["exceeded_cap"] == 0
|
|
assert pool.stats["peak_open"] == 1
|
|
# Back under the cap once everything has been returned
|
|
assert pool.snapshot()["open"] <= 1
|
|
ds.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_memory_databases_not_counted(tmp_path):
|
|
ds = Datasette(memory=True, settings={"max_open_connections": 1})
|
|
db = ds.add_memory_database("pool_mem_test")
|
|
await db.execute_write("create table t (id integer)")
|
|
await db.execute_write("insert into t values (1)")
|
|
await ds.get_database("_memory").execute("select 1")
|
|
assert (await db.execute("select count(*) from t")).single_value() == 1
|
|
assert ds._read_pool.snapshot()["open"] == 0
|
|
ds.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_database_close_with_leased_connection(tmp_path):
|
|
(path,) = _make_dbs(tmp_path, 1)
|
|
ds = Datasette([path], settings={"num_sql_threads": 0})
|
|
db = ds.get_database("db0")
|
|
await db.execute("select 1")
|
|
|
|
def closes_database(conn):
|
|
conn.execute("select 1")
|
|
db.close()
|
|
return "done"
|
|
|
|
assert await db.execute_fn(closes_database) == "done"
|
|
assert _file_conns(db) == []
|
|
assert ds._read_pool.snapshot()["open"] == 0
|
|
ds.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_callback_that_closes_connection_is_discarded(tmp_path):
|
|
(path,) = _make_dbs(tmp_path, 1)
|
|
ds = Datasette([path])
|
|
db = ds.get_database("db0")
|
|
await db.execute_fn(lambda conn: conn.close())
|
|
assert ds._read_pool.stats["discarded_broken"] == 1
|
|
assert (await db.execute("select count(*) from t")).single_value() == 1
|
|
ds.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_open_transaction_rolled_back_on_release(tmp_path):
|
|
(path,) = _make_dbs(tmp_path, 1)
|
|
ds = Datasette([path], settings={"num_sql_threads": 1})
|
|
db = ds.get_database("db0")
|
|
|
|
def begin(conn):
|
|
conn.execute("begin")
|
|
conn.execute("select * from t").fetchall()
|
|
|
|
await db.execute_fn(begin)
|
|
assert ds._read_pool.stats["rolled_back"] == 1
|
|
assert not await db.execute_fn(lambda conn: conn.in_transaction)
|
|
ds.close()
|