Keep named memory databases alive during isolated callbacks

datasette/database.py

@@ -620,6 +620,10 @@ class Database:
                                     None,
                                 )
                                 self._forget_connection(old)
+                            if self.memory_name:
+                                # The isolated connection must not be the last
+                                # connection keeping a named memory DB alive.
+                                self._non_threaded_write_connection()
                             result = _run()
                     else:
                         result = self.ds._read_pool.run(self, fn, isolated=True)
@@ -919,7 +923,13 @@ class Database:
                 and not self.is_memory
             ):
                 writer.close()
-            if not task.isolated_connection and (
+            # File-backed isolated tasks only need their temporary connection.
+            # Named memory databases also need a retained keeper, including
+            # when their very first operation is isolated.
+            needs_retained_connection = not task.isolated_connection or bool(
+                self.memory_name
+            )
+            if needs_retained_connection and (
                 writer.generation != self._conn_generation or not writer.usable()
             ):
                 # The file was replaced or deleted since this connection was
@@ -952,7 +962,7 @@ class Database:
                     **write_span_kwargs,
                 ).end(end_time=dequeued_at_ns)
                 record_write_queue_wait(self.name, dequeued_at_ns - task.enqueued_at_ns)
-                if writer.exception is not None and not task.isolated_connection:
+                if writer.exception is not None and needs_retained_connection:
                     exception = writer.exception
                 elif task.isolated_connection:
                     try:

tests/test_isolated_memory_keeper.py

@@ -0,0 +1,116 @@
+import sqlite3
+import uuid
+
+import pytest
+
+from datasette.app import Datasette
+from datasette.write_budget import DatabaseResourceError
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("threads", [0, 2])
+@pytest.mark.parametrize("raise_after_commit", [False, True])
+async def test_first_isolated_memory_write_keeps_committed_data(
+    threads, raise_after_commit
+):
+    ds = Datasette(settings={"num_sql_threads": threads, "schema_watch_interval_ms": 0})
+    db = ds.add_memory_database(uuid.uuid4().hex)
+    try:
+
+        def create(conn):
+            conn.execute("create table t(id)")
+            conn.execute("insert into t values (123)")
+            conn.commit()
+            if raise_after_commit:
+                raise ValueError("after commit")
+            return conn.execute("select id from t").fetchone()[0]
+
+        if raise_after_commit:
+            with pytest.raises(ValueError, match="after commit"):
+                await db.execute_isolated_fn(create)
+        else:
+            assert await db.execute_isolated_fn(create) == 123
+        # Only the keeper remains after the isolated callback, on both
+        # success and failure. Do not open a reader before checking this.
+        assert len(db._all_connections) == 1
+        assert (await db.execute("select id from t")).single_value() == 123
+        assert (
+            await db.execute_isolated_fn(
+                lambda conn: conn.execute("select id from t").fetchone()[0]
+            )
+            == 123
+        )
+        await db.execute_write("insert into t values (456)")
+        assert (await db.execute("select count(*) from t")).single_value() == 2
+        connections = list(db._all_connections)
+    finally:
+        ds.close()
+    for conn in connections:
+        with pytest.raises(sqlite3.ProgrammingError):
+            conn.execute("select 1")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("threads", [0, 2])
+async def test_isolated_memory_keeper_occupies_writer_slot(threads, tmp_path):
+    path = tmp_path / "file.db"
+    conn = sqlite3.connect(path)
+    conn.execute("create table t(id)")
+    conn.close()
+    ds = Datasette(
+        [str(path)],
+        settings={
+            "num_sql_threads": threads,
+            "max_write_connections": 1,
+            "write_queue_timeout_ms": 50,
+            "schema_watch_interval_ms": 0,
+        },
+    )
+    memory = ds.add_memory_database(uuid.uuid4().hex)
+    file_db = ds.get_database("file")
+    try:
+        await memory.execute_isolated_fn(
+            lambda conn: conn.execute("create table t(id)")
+        )
+        with pytest.raises(DatabaseResourceError):
+            await file_db.execute_write("insert into t values (1)")
+        assert "t" in await memory.table_names()
+        memory.close()
+        await file_db.execute_write("insert into t values (1)")
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("threads", [0, 2])
+async def test_keeper_preparation_failure_does_not_run_isolated_callback(
+    threads, monkeypatch
+):
+    ds = Datasette(
+        settings={
+            "num_sql_threads": threads,
+            "max_write_connections": 1,
+            "schema_watch_interval_ms": 0,
+        }
+    )
+    db = ds.add_memory_database(uuid.uuid4().hex)
+    original = ds._prepare_connection
+    ran = []
+
+    def prepare(conn, name):
+        if name == db.name:
+            raise ValueError("keeper failed")
+        return original(conn, name)
+
+    try:
+        with monkeypatch.context() as patch:
+            patch.setattr(ds, "_prepare_connection", prepare)
+            with pytest.raises(ValueError, match="keeper failed"):
+                await db.execute_isolated_fn(lambda conn: ran.append(True))
+        assert not ran
+        assert not db._all_connections
+        db.close()
+        another = ds.add_memory_database(uuid.uuid4().hex)
+        await another.execute_write("create table t(id)")
+    finally:
+        ds.close()