diff --git a/datasette/database.py b/datasette/database.py index 204d53ec..42f1f528 100644 --- a/datasette/database.py +++ b/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: diff --git a/tests/test_isolated_memory_keeper.py b/tests/test_isolated_memory_keeper.py new file mode 100644 index 00000000..5195225c --- /dev/null +++ b/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()