remote/datasette
1
0
Fork 0
mirror of https://github.com/simonw/datasette.git synced 2026-10-19 07:47:00 +02:00
Code Issues Projects Releases Packages Wiki Activity

Merge spike/read-pool into spike/integrated

Browse source 
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EwkLo2rJWidpkDVT3HkvNv
This commit is contained in:
Claude 2026-10-04 04:40:21 +00:00
commit 25803c23e8
9 changed files with 920 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

43
datasette/app.py
View file

@ -230,6 +230,26 @@ SETTINGS = (
3,
"Number of threads in the thread pool for executing SQLite queries",
),
Setting(
"max_open_connections",
128,
"Maximum number of pooled read connections open across all databases - least recently used idle connections are closed first, 0 for no limit",
),
Setting(
"connection_idle_timeout",
60,
"Close pooled read connections that have been idle for this many seconds - 0 to keep them open",
),
Setting(
"connection_pool_wait_ms",
0,
"When max_open_connections is reached and every connection is in use, wait this long for one to be returned before opening one more",
),
Setting(
"pool_read_connections",
True,
"Reuse read connections between queries - turn off to open a new connection for every query",
),
Setting("sql_time_limit_ms", 1000, "Time limit for a SQL query in milliseconds"),
Setting(
"default_facet_size", 30, "Number of values to return for requested facets"
@ -991,6 +1011,9 @@ class Datasette:
# Collect the first failure and re-raise after every close() has run
if first_exception is None:
first_exception = e
read_pool = self._read_pool_or_none
if read_pool is not None:
read_pool.close()
if self.executor is not None:
try:
self.executor.shutdown(wait=True, cancel_futures=True)
@ -1003,6 +1026,26 @@ class Datasette:
def setting(self, key):
return self._settings.get(key, None)
@property
def _read_pool(self):
pool = self.__dict__.get("_read_pool_instance")
if pool is None:
from .connection_pool import ReadConnectionPool
pool = ReadConnectionPool(
self._prepare_connection,
max_open=self.setting("max_open_connections") or 0,
idle_timeout=float(self.setting("connection_idle_timeout") or 0),
wait_ms=self.setting("connection_pool_wait_ms") or 0,
enabled=bool(self.setting("pool_read_connections")),
)
pool = self.__dict__.setdefault("_read_pool_instance", pool)
return pool
@property
def _read_pool_or_none(self):
return self.__dict__.get("_read_pool_instance")
def settings_dict(self):
# Returns a fully resolved settings dictionary, useful for templates
return {option.name: self.setting(option.name) for option in SETTINGS}

473
datasette/connection_pool.py Normal file
View file

@ -0,0 +1,473 @@
"""Leased read-connection pool.
Read connections are no longer cached per worker thread. Instead every
``execute_fn()`` / ``execute()`` call checks a connection out of a pool for
the duration of one callback and returns it afterwards.
* One ``ReadConnectionPool`` per ``Datasette`` instance. It enforces a global
cap (``max_open_connections``) on open pooled read connections across all
file-backed databases. When a new connection is needed and the cap has been
reached, the least recently used *idle* connection (from any database) is
closed first.
* Each database keeps a LIFO stack of its idle connections, so the most
recently used (warmest) connection is reused and surplus ones age out.
* Idle connections are closed after ``connection_idle_timeout`` seconds. This
is done opportunistically on every checkout and by one process-wide daemon
reaper thread, so it works with or without a running event loop
(non-threaded mode, sync callers, closed loops in tests).
* When the cap has been reached and every connection is checked out the pool
waits up to ``connection_pool_wait_ms`` for one to come back and then opens
a connection anyway (the cap is soft). Each checked-out connection occupies
one executor thread, so the overflow can never exceed ``num_sql_threads``.
* Callbacks receive a ``LeasedConnection`` proxy. Once the callback returns
the proxy is invalidated and any further use raises
``ConnectionLeaseError``.
* In-memory databases are leased the same way but do not count against the
cap and are never reaped: they use no file descriptors, and closing every
connection to a named in-memory database would destroy its contents.
"""
import collections
import inspect
import threading
import time
import weakref
from .utils import sqlite3
class ConnectionLeaseError(RuntimeError):
"""A pooled read connection was used outside the callback it was lent to."""
class LeasedConnection:
"""Forwarding proxy handed to read callbacks.
Every attribute access, method call and attribute assignment is forwarded
to the underlying ``sqlite3.Connection`` while the lease is active. After
the callback returns the proxy is expired and any use raises
``ConnectionLeaseError``. Each lease gets a fresh proxy, so a stale
reference fails even if the same underlying connection has since been
lent to another callback.
``isinstance(proxy, sqlite3.Connection)`` is True (via ``__class__``), but
C functions that require a real connection object (for example the
*target* argument of ``Connection.backup()``) reject the proxy.
"""
__slots__ = ("_conn", "_db_name")
def __init__(self, conn, db_name):
object.__setattr__(self, "_conn", conn)
object.__setattr__(self, "_db_name", db_name)
def _live(self):
conn = self._conn
if conn is None:
raise ConnectionLeaseError(
f"A read connection to database {self._db_name!r} was used after the "
"execute_fn() callback it was passed to had returned. Read "
"connections are pooled and only valid inside the callback: do "
"not store the connection (or a cursor) for later use"
)
return conn
def _expire(self):
object.__setattr__(self, "_conn", None)
# isinstance(conn, sqlite3.Connection) checks keep working
@property
def __class__(self):
return sqlite3.Connection
def __getattr__(self, name):
return getattr(self._live(), name)
def __setattr__(self, name, value):
setattr(self._live(), name, value)
def __delattr__(self, name):
delattr(self._live(), name)
# Dunder methods are looked up on the type, so forward them explicitly
def __enter__(self):
return self._live().__enter__()
def __exit__(self, *args):
return self._live().__exit__(*args)
def __call__(self, *args, **kwargs):
return self._live()(*args, **kwargs)
# Hot paths, to skip __getattr__
def execute(self, *args, **kwargs):
return self._live().execute(*args, **kwargs)
def cursor(self, *args, **kwargs):
return self._live().cursor(*args, **kwargs)
def close(self):
# Closing a pooled connection is allowed; the pool notices on release
# and discards it instead of returning it to the idle stack.
self._live().close()
def __repr__(self):
state = "expired" if self._conn is None else "active"
return f"<LeasedConnection database={self._db_name!r} {state}>"
class _Entry:
__slots__ = ("conn", "last_used", "state")
def __init__(self, conn, state):
self.conn = conn
self.state = state
self.last_used = 0.0
class _DatabaseState:
"""Per-database bookkeeping, stored on the Database as _read_pool_state."""
__slots__ = ("closed", "counted", "db", "idle", "leased", "open")
def __init__(self, db):
self.db = db
self.idle = [] # LIFO stack of _Entry
self.open = 0 # idle + leased
self.leased = 0
self.closed = False
# In-memory databases use no file descriptors and are never evicted
self.counted = not db.is_memory
class ReadConnectionPool:
def __init__(
self,
prepare_connection,
max_open=128,
idle_timeout=60.0,
wait_ms=0,
enabled=True,
):
self._prepare_connection = prepare_connection
self.max_open = max_open # 0 = no limit
self.idle_timeout = idle_timeout # seconds, 0 = never
self.wait_s = wait_ms / 1000.0
self.enabled = enabled
self._lock = threading.Lock()
self._cond = threading.Condition(self._lock)
# Idle, counted entries, least recently used first
self._lru = collections.OrderedDict()
self._open = 0 # counted open connections, idle + leased
self._waiters = 0
self._closed = False
self.stats = collections.Counter()
# -- public API ---------------------------------------------------------
def run(self, db, fn):
"""Call fn(conn) with a pooled connection leased for the duration."""
entry = self.acquire(db)
lease = LeasedConnection(entry.conn, db.name)
try:
result = fn(lease)
finally:
lease._expire()
self.release(entry)
if isinstance(result, sqlite3.Cursor) or inspect.isgenerator(result):
raise ConnectionLeaseError(
f"An execute_fn() callback for database {db.name!r} returned a {type(result).__name__}, which "
"would keep using the pooled connection after the callback "
"returned. Fetch the rows inside the callback instead"
)
return result
def acquire(self, db):
state = self._state(db)
to_close = []
entry = None
with self._lock:
if self._closed or state.closed:
from .database import DatasetteClosedError
raise DatasetteClosedError(f"Database {db.name!r} has been closed")
now = time.monotonic()
self._collect_expired_locked(now, to_close)
deadline = None
while True:
if deadline is not None:
self.stats["wait_ms_max"] = max(
self.stats["wait_ms_max"],
int((time.monotonic() - now) * 1000),
)
if state.idle:
entry = state.idle.pop()
if state.counted:
del self._lru[entry]
state.leased += 1
self.stats["reused"] += 1
break
if not state.counted or not self.max_open or self._open < self.max_open:
break
if self._lru:
# Make room by closing the least recently used idle
# connection, which may belong to any database
victim, _ = self._lru.popitem(last=False)
victim.state.idle.remove(victim)
victim.state.open -= 1
self._open -= 1
to_close.append(victim)
self.stats["evicted_lru"] += 1
continue
# Cap reached and every connection is checked out
if self.wait_s > 0:
if deadline is None:
deadline = now + self.wait_s
self.stats["waits"] += 1
remaining = deadline - time.monotonic()
if remaining > 0:
self._waiters += 1
t0 = time.monotonic()
try:
self._cond.wait(remaining)
finally:
self._waiters -= 1
waited_ms = int((time.monotonic() - t0) * 1000)
self.stats["wait_ms_total"] += waited_ms
continue
self.stats["wait_timeouts"] += 1
if deadline is not None:
self.stats["wait_ms_max"] = max(
self.stats["wait_ms_max"],
int((time.monotonic() - now) * 1000),
)
# Soft cap: open one more. Bounded by the number of threads
# that can hold a lease at once (num_sql_threads).
self.stats["exceeded_cap"] += 1
break
if entry is None:
state.open += 1
state.leased += 1
self.stats["opened"] += 1
if state.counted:
self._open += 1
self.stats["peak_open"] = max(self.stats["peak_open"], self._open)
self._close_entries(to_close)
if entry is not None:
return entry
# Open and prepare outside the lock
conn = None
try:
conn = db.connect()
self._prepare_connection(conn, db.name)
except BaseException:
with self._lock:
state.open -= 1
state.leased -= 1
self.stats["opened"] -= 1
self.stats["open_failed"] += 1
if state.counted:
self._open -= 1
if self._waiters:
self._cond.notify()
if conn is not None:
self._close_conn(state, conn)
raise
return _Entry(conn, state)
def release(self, entry):
state = entry.state
conn = entry.conn
discard = False
try:
# Hygiene: a callback that began a transaction and did not end it
# must not hand that transaction to the next borrower
rolled_back = False
if conn.in_transaction:
conn.rollback()
rolled_back = True
except Exception: # noqa: BLE001
# e.g. ProgrammingError: the callback closed the connection
discard = True
notify_reaper = None
with self._lock:
if discard:
self.stats["discarded_broken"] += 1
elif rolled_back:
self.stats["rolled_back"] += 1
state.leased -= 1
over_cap = state.counted and self.max_open and self._open > self.max_open
if discard or over_cap or self._closed or state.closed or not self.enabled:
state.open -= 1
if state.counted:
self._open -= 1
discard = True
if over_cap:
self.stats["closed_over_cap"] += 1
else:
entry.last_used = time.monotonic()
state.idle.append(entry)
if state.counted:
self._lru[entry] = None
if self.idle_timeout > 0:
notify_reaper = entry.last_used + self.idle_timeout
if self._waiters:
self._cond.notify()
if discard:
self._close_conn(state, conn)
if notify_reaper is not None:
_reaper.wake(self, notify_reaper)
def reap_expired(self):
"""Close expired idle connections. Returns the next expiry time
(time.monotonic() based) or None if nothing is idle."""
to_close = []
with self._lock:
self._collect_expired_locked(time.monotonic(), to_close)
next_expiry = None
if self._lru and self.idle_timeout > 0:
head = next(iter(self._lru))
next_expiry = head.last_used + self.idle_timeout
self._close_entries(to_close)
return next_expiry
def close_database(self, db):
"""Close idle connections for db; leased ones are closed on release."""
state = getattr(db, "_read_pool_state", None)
if state is None:
return
with self._lock:
state.closed = True
to_close = state.idle
state.idle = []
for entry in to_close:
if state.counted:
self._lru.pop(entry, None)
self._open -= 1
state.open -= 1
self._close_entries(to_close)
def close(self):
with self._lock:
self._closed = True
to_close = list(self._lru)
self._lru.clear()
for entry in to_close:
entry.state.idle.remove(entry)
entry.state.open -= 1
self._open -= 1
self._cond.notify_all()
self._close_entries(to_close)
def snapshot(self):
with self._lock:
return {
"open": self._open,
"idle": len(self._lru),
"max_open": self.max_open,
"idle_timeout": self.idle_timeout,
"stats": dict(self.stats),
}
# -- internals ----------------------------------------------------------
def _state(self, db):
state = db._read_pool_state
if state is None:
with self._lock:
state = db._read_pool_state
if state is None:
state = db._read_pool_state = _DatabaseState(db)
return state
def _collect_expired_locked(self, now, to_close):
if self.idle_timeout <= 0 or not self._lru:
return
cutoff = now - self.idle_timeout
lru = self._lru
while lru:
head = next(iter(lru))
if head.last_used > cutoff:
break
lru.popitem(last=False)
head.state.idle.remove(head)
head.state.open -= 1
self._open -= 1
to_close.append(head)
self.stats["expired"] += 1
def _close_entries(self, entries):
for entry in entries:
self._close_conn(entry.state, entry.conn)
def _close_conn(self, state, conn):
try:
conn.close()
except Exception: # noqa: BLE001, S110
pass
with self._lock:
self.stats["closed"] += 1
try:
state.db._all_connections.remove(conn)
except ValueError:
# May already have been cleared by Database.close()
pass
class _Reaper:
"""One daemon thread per process that closes expired idle connections
for every live pool. Pools are held by weak reference."""
def __init__(self):
self._cond = threading.Condition()
self._pools = weakref.WeakSet()
self._thread = None
self._dirty = False
# monotonic time of the next planned scan, inf while sleeping
# indefinitely or while scanning
self.wake_at = float("inf")
def wake(self, pool, expiry):
# Unlocked pre-check: most releases do not need to wake the reaper
if expiry >= self.wake_at and pool in self._pools:
return
with self._cond:
self._pools.add(pool)
if self._thread is None or not self._thread.is_alive():
self._thread = threading.Thread(
target=self._run, name="datasette-read-pool-reaper", daemon=True
)
self._thread.start()
if expiry < self.wake_at:
self._dirty = True
self._cond.notify()
def _run(self):
while True:
with self._cond:
self.wake_at = float("inf")
self._dirty = False
pools = list(self._pools)
next_wake = None
for pool in pools:
try:
expiry = pool.reap_expired()
except Exception: # noqa: BLE001
expiry = None
if expiry is not None and (next_wake is None or expiry < next_wake):
next_wake = expiry
del pools
with self._cond:
if self._dirty:
continue
if next_wake is None:
self.wake_at = float("inf")
self._cond.wait()
else:
self.wake_at = next_wake
# Small slack so a batch of connections that expire close
# together is handled in one pass
self._cond.wait(max(0.0, next_wake - time.monotonic()) + 0.05)
_reaper = _Reaper()

42
datasette/database.py
View file

@ -16,6 +16,7 @@ import sqlite_utils
from opentelemetry import context as otel_context_api
from opentelemetry.trace import Status, StatusCode
from .connection_pool import ConnectionLeaseError # noqa: F401
from .inspect import inspect_hash
from .telemetry import (
callback_name,
@ -127,9 +128,11 @@ class Database:
self._closed = False
self._pending_execute_futures = set()
self._pending_execute_futures_lock = threading.Lock()
# These are used when in non-threaded mode:
# Used when in non-threaded mode (reads go through ds._read_pool):
self._read_connection = None
self._write_connection = None
# Bookkeeping for ds._read_pool, created on first read
self._read_pool_state = None
# Track file and memory connections, including reads on worker threads,
# so close() can release all of them from the calling thread.
self._all_connections = []
@ -242,8 +245,13 @@ class Database:
except Exception: # noqa: BLE001, S110
# Shutdown teardown - a failed pending write must not block close()
pass
# Close idle pooled read connections; any still leased are closed
# by the pool when they are released
read_pool = getattr(self.ds, "_read_pool_or_none", None)
if read_pool is not None:
read_pool.close_database(self)
# Close anything still tracked in _all_connections
for connection in self._all_connections:
for connection in list(self._all_connections):
try:
connection.close()
except Exception: # noqa: BLE001, S110
@ -686,21 +694,14 @@ class Database:
async def _execute_fn(self, fn):
self._check_not_closed()
read_pool = self.ds._read_pool
if self.ds.executor is None:
# non-threaded mode
if self._read_connection is None:
self._read_connection = self.connect()
self.ds._prepare_connection(self._read_connection, self.name)
return fn(self._read_connection)
# non-threaded mode: lease on the event loop thread
return read_pool.run(self, fn)
# threaded mode
# threaded mode: lease on the executor thread for one callback
def in_thread():
conn = getattr(connections, self._thread_local_id, None)
if not conn:
conn = self.connect()
self.ds._prepare_connection(conn, self.name)
setattr(connections, self._thread_local_id, conn)
return fn(conn)
return read_pool.run(self, fn)
with self._pending_execute_futures_lock:
self._check_not_closed()
@ -742,6 +743,7 @@ class Database:
) as execute_span:
try:
with sqlite_timelimit(conn, time_limit_ms):
cursor = None
try:
cursor = conn.cursor()
cursor.execute(sql, params if params is not None else {})
@ -755,6 +757,7 @@ class Database:
else:
rows = cursor.fetchall()
truncated = False
description = cursor.description
except (sqlite3.OperationalError, sqlite3.DatabaseError) as e:
if e.args == ("interrupted",):
raise QueryInterrupted(e, sql, params)
@ -764,6 +767,13 @@ class Database:
)
sys.stderr.flush()
raise
finally:
# Release the statement now: a traceback that
# keeps this frame alive would otherwise keep it
# unfinalized, and closing the pooled connection
# later would leave a zombie holding its fds
if cursor is not None:
cursor.close()
except QueryInterrupted as e:
if not timeout_expected:
execute_span.record_exception(e)
@ -776,10 +786,10 @@ class Database:
raise
if truncate:
return Results(rows, truncated, cursor.description)
return Results(rows, truncated, description)
else:
return Results(rows, False, cursor.description)
return Results(rows, False, description)
with trace( # noqa: SIM117
"sql", database=self.name, sql=sql.strip(), params=params

13
docs/cli-reference.rst
View file

@ -251,6 +251,19 @@ These can be passed to ``datasette serve`` using ``datasette serve --setting nam
this limit (default=2097152)
num_sql_threads Number of threads in the thread pool for
executing SQLite queries (default=3)
max_open_connections Maximum number of pooled read connections open
across all databases - least recently used idle
connections are closed first, 0 for no limit
(default=128)
connection_idle_timeout Close pooled read connections that have been idle
for this many seconds - 0 to keep them open
(default=60)
connection_pool_wait_ms When max_open_connections is reached and every
connection is in use, wait this long for one to
be returned before opening one more (default=0)
pool_read_connections Reuse read connections between queries - turn off
to open a new connection for every query
(default=True)
sql_time_limit_ms Time limit for a SQL query in milliseconds
(default=1000)
default_facet_size Number of values to return for requested facets

2
docs/internals.rst
View file

@ -2174,6 +2174,8 @@ Example usage:
version = await db.execute_fn(get_version)
The connection is borrowed from a pool for the duration of the call (see :ref:`setting_max_open_connections`) and is only valid inside the function. Using it after the function has returned - for example by storing it on an object - raises a ``datasette.database.ConnectionLeaseError``. Returning a cursor or a generator from the function raises the same error: fetch the rows you need before returning.
The call is traced as a ``db.query`` OpenTelemetry span carrying ``datasette.callback`` (the function's qualified name) rather than ``db.query.text``, since the SQL is whatever the function chooses to run - see :ref:`internals_telemetry`. Passing a named function gives the span a readable identity; a lambda reports ``<lambda>``.
.. _database_execute_write:

66
docs/settings.rst
View file

@ -214,6 +214,72 @@ Maximum number of threads in the thread pool Datasette uses to execute SQLite qu
Setting this to 0 turns off threaded SQL queries entirely - useful for environments that do not support threading such as `Pyodide <https://pyodide.org/>`__.
.. _setting_max_open_connections:
max_open_connections
~~~~~~~~~~~~~~~~~~~~
.. [[[cog
setting_default(cog, "max_open_connections")
.. ]]]
Default: ``128``
.. [[[end]]]
Read queries borrow a connection from a pool for the duration of one query and return it afterwards. This setting caps how many pooled read connections can be open at once across all attached file databases. When a new connection is needed and the cap has been reached, the least recently used idle connection is closed first.
If every connection is in use the pool waits up to :ref:`setting_connection_pool_wait_ms` and then opens one more anyway, so the cap can be exceeded by at most :ref:`setting_num_sql_threads` connections. Set this to 0 for no limit.
::
datasette *.db --setting max_open_connections 64
.. _setting_connection_idle_timeout:
connection_idle_timeout
~~~~~~~~~~~~~~~~~~~~~~~
.. [[[cog
setting_default(cog, "connection_idle_timeout")
.. ]]]
Default: ``60``
.. [[[end]]]
Pooled read connections that have not been used for this many seconds are closed. Set to 0 to keep idle connections open until they are evicted by :ref:`setting_max_open_connections`.
.. _setting_connection_pool_wait_ms:
connection_pool_wait_ms
~~~~~~~~~~~~~~~~~~~~~~~
.. [[[cog
setting_default(cog, "connection_pool_wait_ms")
.. ]]]
Default: ``0``
.. [[[end]]]
How long a query waits for a connection to be returned when :ref:`setting_max_open_connections` has been reached and every pooled connection is in use, before opening one more connection anyway.
.. _setting_pool_read_connections:
pool_read_connections
~~~~~~~~~~~~~~~~~~~~~
.. [[[cog
setting_default(cog, "pool_read_connections")
.. ]]]
Default: ``on``
.. [[[end]]]
Reuse read connections between queries. Turn this off to open, prepare and close a new connection for every read query.
.. _setting_allow_facet:
allow_facet

4
tests/test_api.py
View file

@ -759,6 +759,10 @@ async def test_settings_json(ds_client):
"suggest_facets": True,
"default_cache_ttl": 5,
"num_sql_threads": 1,
"max_open_connections": 128,
"connection_idle_timeout": 60,
"connection_pool_wait_ms": 0,
"pool_read_connections": True,
"cache_size_kb": 0,
"allow_csv_stream": True,
"max_csv_mb": 100,

10
tests/test_internals_database.py
View file

@ -13,6 +13,7 @@ from opentelemetry import context as otel_context_api
from datasette.app import Datasette
from datasette.database import (
ConnectionLeaseError,
Database,
DatasetteClosedError,
ExecuteWriteResult,
@ -1312,8 +1313,15 @@ async def test_database_close_is_idempotent(tmpdir):
async def test_close_releases_memory_connections(num_sql_threads, named):
ds = Datasette(memory=True, settings={"num_sql_threads": num_sql_threads})
db = ds.add_memory_database(uuid.uuid4().hex) if named else ds.get_database()
read_connection = await db.execute_fn(lambda conn: conn)
leased = await db.execute_fn(lambda conn: conn)
# Read connections are leased: the proxy is unusable once the callback
# returns, so reach the underlying connection through a cursor
read_connection = await db.execute_fn(
lambda conn: conn.execute("select 1").connection
)
write_connection = await db.execute_write_fn(lambda conn: conn)
with pytest.raises(ConnectionLeaseError):
leased.execute("select 1")
ds.close()
for conn in (read_connection, write_connection):
with pytest.raises(sqlite3.ProgrammingError, match="closed"):

284
tests/test_read_pool.py Normal file
View file

@ -0,0 +1,284 @@
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()