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

Lease proxies for write and isolated callbacks; close() never frees a busy connection

Browse source 
Callbacks passed to execute_write_fn() - and the internal callbacks of
execute_write(), execute_write_many() and execute_write_script(), plus
write_wrapper hooks - and execute_isolated_fn() now receive a
LeasedConnection that raises ConnectionLeaseError once the callback has
returned, in threaded and non-threaded mode. "with conn as c" yields the
proxy, not the raw connection. Write callbacks may still return a cursor
(lambda conn: conn.execute(...) is common, and the write_many/script
methods return theirs).

The threaded execute_isolated_fn() path for immutable databases now runs
as a tracked executor future, so Database.close() waits for it instead
of closing its connection mid-query (the segfault in the review, P03).
close() also leaves alone the connections of a write thread that is
still running after the 10s join; that thread closes them itself.

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:51:51 +00:00
commit 5c30d14d3b
No known key found for this signature in database
3 changed files with 257 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

81
datasette/database.py
View file

@ -16,7 +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 .connection_pool import ConnectionLeaseError, LeasedConnection # noqa: F401
from .inspect import inspect_hash
from .telemetry import (
callback_name,
@ -127,6 +127,8 @@ class Database:
# Write threads that have deregistered (idle exit) but may still be
# closing their connection. close() joins them.
self._retiring_write_threads = set()
# Connections currently held by write threads (see _WriteConnection)
self._write_thread_connections = set()
# Number of write threads started over the lifetime of this Database,
# and how often a thread's idle timeout fired but it found a newly
# queued task once it held the lock (so it kept running)
@ -341,13 +343,18 @@ class Database:
retiring_threads = tuple(self._retiring_write_threads)
if write_thread is not None:
self._write_queue.put(_SHUTDOWN)
still_running = False
for thread in ((write_thread,) if write_thread else ()) + retiring_threads:
thread.join(timeout=10)
if thread.is_alive():
still_running = True
sys.stderr.write(
f"Datasette: write thread for {self.name!r} did not exit within 10s\n"
)
sys.stderr.flush()
# A write thread that is still running owns its connection and closes
# it when it finishes; closing it here could free it mid-query
in_use = set(self._write_thread_connections) if still_running else set()
for future in pending_execute_futures:
try:
future.result()
@ -361,6 +368,8 @@ class Database:
read_pool.close_database(self)
# Close anything still tracked in _all_connections
for connection in list(self._all_connections):
if connection in in_use:
continue
try:
connection.close()
except Exception: # noqa: BLE001, S110
@ -495,18 +504,16 @@ class Database:
write = self.is_mutable
def _run():
# Always closed here, by the thread using it. close() waits for
# this call before closing tracked connections (the threaded
# immutable case runs as a tracked executor future)
isolated_connection = self.connect(write=write)
try:
return fn(isolated_connection)
return _call_with_lease(fn, isolated_connection, self.name, "isolated")
finally:
if write:
self._schema_check(isolated_connection)
isolated_connection.close()
try:
self._all_connections.remove(isolated_connection)
except ValueError:
# May already have been cleared by close().
pass
self._forget_connection(isolated_connection)
with tracer.start_as_current_span(DB_QUERY, kind=DB_QUERY.kind) as span:
span.set_attribute(DB_SYSTEM, "sqlite")
@ -522,11 +529,9 @@ class Database:
if not write:
# Immutable database - no writes can ever occur, so there
# is no write queue to block; run against a fresh
# read-only connection
ctx = contextvars.copy_context()
return await asyncio.get_running_loop().run_in_executor(
self.ds.executor, ctx.run, _run
)
# read-only connection on the SQL thread pool, tracked so
# close() waits for it
return await self._run_in_executor(_run)
# Threaded mode - send to write thread
result = await self._send_to_write_thread(fn, isolated_connection=True)
await self._after_write()
@ -573,9 +578,9 @@ class Database:
if transaction:
with conn:
conn.execute("BEGIN IMMEDIATE")
result = fn(conn)
result = _call_with_lease(fn, conn, self.name, "write")
else:
result = fn(conn)
result = _call_with_lease(fn, conn, self.name, "write")
finally:
self._schema_check(conn)
if not block:
@ -810,11 +815,19 @@ class Database:
)
span.set_attribute(TRANSACTION, task.transaction)
isolated_connection = self.connect(write=True)
# Owned by this thread: close() must not close it
# if this thread outlives close()'s join timeout
self._write_thread_connections.add(isolated_connection)
try:
result = task.fn(isolated_connection)
result = _call_with_lease(
task.fn, isolated_connection, self.name, "isolated"
)
finally:
self._schema_check(isolated_connection, task.loop)
self._forget_connection(isolated_connection)
self._write_thread_connections.discard(
isolated_connection
)
except Exception as e: # noqa: BLE001
# Write thread must survive any task failure or the database wedges
sys.stderr.write(f"{e}\n")
@ -831,13 +844,20 @@ class Database:
task.isolated_connection,
)
span.set_attribute(TRANSACTION, task.transaction)
# The callback gets a lease proxy that stops
# working when it returns: the connection belongs
# to this thread and is closed when it idles out
try:
if task.transaction:
with conn:
conn.execute("BEGIN IMMEDIATE")
result = task.fn(conn)
result = _call_with_lease(
task.fn, conn, self.name, "write"
)
else:
result = task.fn(conn)
result = _call_with_lease(
task.fn, conn, self.name, "write"
)
finally:
self._schema_check(conn, task.loop)
except Exception as e: # noqa: BLE001
@ -880,13 +900,19 @@ class Database:
def in_thread():
return read_pool.run(self, fn)
return await self._run_in_executor(in_thread)
async def _run_in_executor(self, fn):
"""Run fn() on the shared SQL thread pool. The future is tracked in
_pending_execute_futures so close() waits for it before closing any
connection."""
with self._pending_execute_futures_lock:
self._check_not_closed()
# Run in a copy of the caller's context so spans created in the
# thread have the correct parent. This needs a fresh copy for
# each submit, since a Context cannot be entered concurrently.
ctx = contextvars.copy_context()
future = self.ds.executor.submit(ctx.run, in_thread)
future = self.ds.executor.submit(ctx.run, fn)
self._pending_execute_futures.add(future)
future.add_done_callback(self._remove_pending_execute_future)
return await asyncio.wrap_future(future)
@ -1333,6 +1359,21 @@ def _apply_write_wrapper(fn, wrapper_factory, track_event):
return wrapped
def _call_with_lease(fn, conn, db_name, kind):
"""Call fn with a LeasedConnection for conn that expires when it returns.
Unlike read callbacks, write callbacks may return a cursor (for example
``lambda conn: conn.execute(...)``, and execute_write_many() and
execute_write_script() return theirs); reading attributes such as
``lastrowid`` from it afterwards is fine, fetching rows is not.
"""
lease = LeasedConnection(conn, db_name, kind)
try:
return fn(lease)
finally:
lease._expire()
class _WriteConnection:
"""The write thread's connection to its database.
@ -1359,6 +1400,7 @@ class _WriteConnection:
self.exception = None
try:
self.conn = db.connect(write=True)
db._write_thread_connections.add(self.conn)
# Threads do not inherit the caller's context, so any spans
# created by prepare_connection hooks here are root spans
db.ds._prepare_connection(self.conn, db.name)
@ -1376,6 +1418,7 @@ class _WriteConnection:
# Also drops it from _all_connections, so that list does not
# grow by one closed connection per write thread restart
self.db._forget_connection(conn)
self.db._write_thread_connections.discard(conn)
class WriteTask:

189
tests/test_connection_leases.py Normal file
View file

@ -0,0 +1,189 @@
"""
Connections handed to execute_write_fn(), the SQL write methods and
execute_isolated_fn() callbacks are lease proxies, like read connections:
they stop working once the callback returns.
"""
import asyncio
import sqlite3
import threading
import time
import pytest
import sqlite_utils
from datasette.app import Datasette
from datasette.database import ConnectionLeaseError
def _make_db(tmp_path, name="leases"):
path = str(tmp_path / f"{name}.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()
return path
@pytest.mark.asyncio
@pytest.mark.parametrize("num_sql_threads", [0, 2])
@pytest.mark.parametrize("transaction", [True, False])
async def test_write_callback_connection_expires(
tmp_path, num_sql_threads, transaction
):
ds = Datasette([_make_db(tmp_path)], settings={"num_sql_threads": num_sql_threads})
db = ds.get_database("leases")
stash = {}
def write(conn):
stash["conn"] = conn
assert isinstance(conn, sqlite3.Connection)
conn.execute("insert into t (v) values ('y')")
if not transaction:
# transaction=False DML is not committed for you (unchanged)
conn.commit()
return conn.execute("select count(*) from t").fetchone()[0]
assert await db.execute_write_fn(write, transaction=transaction) == 2
with pytest.raises(ConnectionLeaseError, match="execute_write_fn"):
stash["conn"].execute("select 1")
with pytest.raises(ConnectionLeaseError):
stash["conn"].row_factory = None
# The write connection itself is fine for the next write
await db.execute_write("insert into t (v) values ('z')")
assert (await db.execute("select count(*) from t")).single_value() == 3
ds.close()
@pytest.mark.asyncio
@pytest.mark.parametrize("num_sql_threads", [0, 2])
async def test_write_callback_may_return_cursor(tmp_path, num_sql_threads):
# A common plugin pattern: lambda conn: conn.execute(...). The returned
# cursor's attributes stay readable; only the connection is leased.
ds = Datasette([_make_db(tmp_path)], settings={"num_sql_threads": num_sql_threads})
db = ds.get_database("leases")
cursor = await db.execute_write_fn(
lambda conn: conn.execute("insert into t (v) values ('y')")
)
assert cursor.lastrowid == 2
assert cursor.rowcount == 1
many = await db.execute_write_many("insert into t (v) values (?)", [["a"], ["b"]])
assert many.rowcount == 2
await db.execute_write_script("insert into t (v) values ('c');")
assert (await db.execute("select count(*) from t")).single_value() == 5
ds.close()
@pytest.mark.asyncio
@pytest.mark.parametrize("num_sql_threads", [0, 2])
async def test_with_conn_inside_write_callback_yields_proxy(tmp_path, num_sql_threads):
ds = Datasette([_make_db(tmp_path)], settings={"num_sql_threads": num_sql_threads})
db = ds.get_database("leases")
stash = {}
def write(conn):
with conn as inner:
stash["inner"] = inner
inner.execute("insert into t (v) values ('y')")
await db.execute_write_fn(write, transaction=False)
with pytest.raises(ConnectionLeaseError):
stash["inner"].execute("select 1")
assert (await db.execute("select count(*) from t")).single_value() == 2
ds.close()
@pytest.mark.asyncio
@pytest.mark.parametrize("num_sql_threads", [0, 2])
async def test_sqlite_utils_on_write_lease(tmp_path, num_sql_threads):
ds = Datasette([_make_db(tmp_path)], settings={"num_sql_threads": num_sql_threads})
db = ds.get_database("leases")
def write(conn):
sdb = sqlite_utils.Database(conn)
sdb["dogs"].insert_all([{"id": 1, "name": "Cleo"}], pk="id")
sdb["dogs"].transform(rename={"name": "dog_name"})
return sdb["dogs"].columns_dict
assert await db.execute_write_fn(write) == {"id": int, "dog_name": str}
ds.close()
@pytest.mark.asyncio
@pytest.mark.parametrize("num_sql_threads", [0, 2])
@pytest.mark.parametrize("immutable", [False, True])
async def test_isolated_callback_connection_expires(
tmp_path, num_sql_threads, immutable
):
path = _make_db(tmp_path)
if immutable:
ds = Datasette(immutables=[path], settings={"num_sql_threads": num_sql_threads})
else:
ds = Datasette([path], settings={"num_sql_threads": num_sql_threads})
db = ds.get_database("leases")
stash = {}
def isolated(conn):
stash["conn"] = conn
return conn.execute("select count(*) from t").fetchone()[0]
assert await db.execute_isolated_fn(isolated) == 1
with pytest.raises(ConnectionLeaseError, match="execute_isolated_fn"):
stash["conn"].execute("select 1")
ds.close()
@pytest.mark.asyncio
async def test_write_wrapper_receives_lease(tmp_path):
from datasette import hookimpl
from datasette.plugins import pm
stash = {}
class WrapperPlugin:
__name__ = "WrapperPlugin"
@hookimpl
def write_wrapper(self, datasette, database, request, transaction):
if database != "leases":
return None
def wrapper(conn):
stash["conn"] = conn
yield
return wrapper
pm.register(WrapperPlugin(), name="test_write_wrapper_receives_lease")
try:
ds = Datasette([_make_db(tmp_path)])
db = ds.get_database("leases")
await db.execute_write("insert into t (v) values ('y')")
with pytest.raises(ConnectionLeaseError):
stash["conn"].execute("select 1")
ds.close()
finally:
pm.unregister(name="test_write_wrapper_receives_lease")
@pytest.mark.asyncio
async def test_close_waits_for_immutable_isolated_fn(tmp_path):
# On main this ran on the executor untracked, so close() could close
# the connection while the callback was still stepping it (a segfault
# with a progress handler installed)
ds = Datasette(immutables=[_make_db(tmp_path)], settings={"num_sql_threads": 2})
db = ds.get_database("leases")
started = threading.Event()
def slow(conn):
started.set()
time.sleep(0.3)
return conn.execute("select count(*) from t").fetchone()[0]
task = asyncio.ensure_future(db.execute_isolated_fn(slow))
assert await asyncio.to_thread(started.wait, 5)
assert len(db._pending_execute_futures) == 1
db.close()
assert await task == 1
ds.close()

7
tests/test_internals_database.py
View file

@ -1319,9 +1319,14 @@ async def test_close_releases_memory_connections(num_sql_threads, named):
read_connection = await db.execute_fn(
lambda conn: conn.execute("select 1").connection
)
write_connection = await db.execute_write_fn(lambda conn: conn)
leased_write = await db.execute_write_fn(lambda conn: conn)
write_connection = await db.execute_write_fn(
lambda conn: conn.execute("select 1").connection
)
with pytest.raises(ConnectionLeaseError):
leased.execute("select 1")
with pytest.raises(ConnectionLeaseError):
leased_write.execute("select 1")
ds.close()
for conn in (read_connection, write_connection):
with pytest.raises(sqlite3.ProgrammingError, match="closed"):