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

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

tests/test_connection_leases.py

@@ -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()

tests/test_internals_database.py

@@ -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"):