mirror of
https://github.com/simonw/datasette.git
synced 2026-10-19 07:47:00 +02:00
378 lines
12 KiB
Python
378 lines
12 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})
|
|
pool = ds._read_pool
|
|
# Below the 4 x num_sql_threads clamp, so set the private attribute
|
|
pool.max_open = 2
|
|
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)
|
|
ds._read_pool.max_open = 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.parametrize(
|
|
"num_sql_threads,configured,expected",
|
|
[
|
|
(3, 128, 128),
|
|
(3, 5, 5),
|
|
(8, 20, 20),
|
|
(0, 1, 1),
|
|
(3, 0, 0),
|
|
],
|
|
)
|
|
def test_max_open_connections_respects_configured_limit(
|
|
num_sql_threads, configured, expected
|
|
):
|
|
ds = Datasette(
|
|
settings={
|
|
"num_sql_threads": num_sql_threads,
|
|
"max_open_connections": configured,
|
|
}
|
|
)
|
|
assert ds._read_pool.max_open == expected
|
|
# /-/settings.json reports the configured value
|
|
assert ds.setting("max_open_connections") == configured
|
|
ds.close()
|
|
|
|
|
|
def test_connection_idle_timeout_ms_setting():
|
|
ds = Datasette(settings={"connection_idle_timeout_ms": 1500})
|
|
assert ds._read_pool.idle_timeout == 1.5
|
|
ds.close()
|
|
ds = Datasette(settings={"connection_idle_timeout_ms": 0})
|
|
assert ds._read_pool.idle_timeout == 0
|
|
ds.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_hard_cap_when_every_connection_is_leased(tmp_path):
|
|
paths = _make_dbs(tmp_path, 3)
|
|
ds = Datasette(paths, settings={"num_sql_threads": 3, "max_open_connections": 1})
|
|
entered = threading.Event()
|
|
release = threading.Event()
|
|
|
|
def slow(conn):
|
|
entered.set()
|
|
assert release.wait(5)
|
|
return conn.execute("select 1").fetchone()[0]
|
|
|
|
tasks = [
|
|
asyncio.create_task(ds.get_database(f"db{i}").execute_fn(slow))
|
|
for i in range(3)
|
|
]
|
|
try:
|
|
assert await asyncio.to_thread(entered.wait, 5)
|
|
await asyncio.sleep(0.05)
|
|
assert ds._read_pool.snapshot()["open"] == 1
|
|
assert not any(task.done() for task in tasks)
|
|
release.set()
|
|
assert await asyncio.gather(*tasks) == [1, 1, 1]
|
|
assert ds._read_pool.stats["peak_open"] == 1
|
|
finally:
|
|
release.set()
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
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()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("num_sql_threads", [0, 2])
|
|
async def test_invalidate_database_discards_stale_connections(
|
|
tmp_path, num_sql_threads
|
|
):
|
|
paths = _make_dbs(tmp_path, 2)
|
|
ds = Datasette(paths, settings={"num_sql_threads": num_sql_threads})
|
|
db0, db1 = ds.get_database("db0"), ds.get_database("db1")
|
|
pool = ds._read_pool
|
|
await db0.execute("select 1")
|
|
await db1.execute("select 1")
|
|
old = db0._read_pool_state.idle[-1].conn
|
|
assert db0._read_pool_state.open == 1
|
|
# What the SchemaWatcher does when it sees db0's file replaced
|
|
db0._invalidate_connections()
|
|
# Idle connections for db0 are closed at once; db1 is untouched
|
|
assert db0._read_pool_state.open == 0
|
|
assert db1._read_pool_state.open == 1
|
|
assert old not in db0._all_connections
|
|
with pytest.raises(sqlite3.ProgrammingError):
|
|
old.execute("select 1")
|
|
# The next read opens a new connection
|
|
await db0.execute("select 1")
|
|
assert db0._read_pool_state.idle[-1].conn is not old
|
|
assert pool.stats["discarded_stale"] == 1
|
|
ds.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_leased_connection_discarded_on_release_after_invalidation(tmp_path):
|
|
(path,) = _make_dbs(tmp_path, 1)
|
|
ds = Datasette([path], settings={"num_sql_threads": 2})
|
|
db = ds.get_database("db0")
|
|
pool = ds._read_pool
|
|
seen = {}
|
|
|
|
def callback(conn):
|
|
seen["raw"] = conn.execute("select 1").connection
|
|
# The file is replaced while this callback holds the connection:
|
|
# it must not be closed under us, only discarded on release
|
|
db._invalidate_connections()
|
|
return conn.execute("select count(*) from t").fetchone()[0]
|
|
|
|
assert await db.execute_fn(callback) == 1
|
|
assert db._read_pool_state.open == 0
|
|
assert pool.stats["discarded_stale"] == 1
|
|
with pytest.raises(ConnectionLeaseError):
|
|
seen["raw"].execute("select 1")
|
|
assert (await db.execute("select count(*) from t")).single_value() == 1
|
|
ds.close()
|
|
|
|
|
|
def test_reaper_unavailable_falls_back_to_checkout_reaping(tmp_path, monkeypatch):
|
|
# Pyodide cannot start threads: idle connections must then be reaped by
|
|
# the next checkout instead, and reads must keep working
|
|
from datasette import connection_pool
|
|
|
|
reaper = connection_pool._Reaper()
|
|
monkeypatch.setattr(connection_pool, "_reaper", reaper)
|
|
|
|
def no_threads(self):
|
|
raise RuntimeError("can't start new thread")
|
|
|
|
monkeypatch.setattr(threading.Thread, "start", no_threads)
|
|
paths = _make_dbs(tmp_path, 2)
|
|
ds = Datasette(paths, settings={"num_sql_threads": 0})
|
|
pool = ds._read_pool
|
|
pool.idle_timeout = 0.05
|
|
|
|
async def read(name):
|
|
return (await ds.get_database(name).execute("select 1")).single_value()
|
|
|
|
assert asyncio.run(read("db0")) == 1
|
|
assert reaper.unavailable
|
|
assert pool.snapshot()["open"] == 1
|
|
time.sleep(0.1)
|
|
assert asyncio.run(read("db1")) == 1
|
|
# db0's expired connection was closed during db1's checkout
|
|
assert pool.stats["expired"] == 1
|
|
assert ds.get_database("db0")._read_pool_state.open == 0
|
|
ds.close()
|