Expire SQLite blob handles and dump iterators with their leases

tests/test_connection_resource_leases.py

@@ -0,0 +1,326 @@
+"""SQLite resources must not retain locks or run after their callback ends."""
+
+import operator
+import sqlite3
+
+import pytest
+
+from datasette.app import Datasette
+from datasette.connection_pool import ConnectionLeaseError
+
+
+@pytest.fixture
+def resource_db(tmp_path):
+    path = tmp_path / "resources.db"
+    conn = sqlite3.connect(path)
+    conn.executescript(
+        "create table t(id integer primary key, data blob);"
+        "insert into t(data) values (x'616263646566'), (zeroblob(6)), (zeroblob(6));"
+        "create table another(id);"
+    )
+    conn.close()
+    return str(path)
+
+
+def callback_method(db, kind):
+    return {
+        "read": db.execute_fn,
+        "write": db.execute_write_fn,
+        "isolated": db.execute_isolated_fn,
+    }[kind]
+
+
+def check_no_read_lock(path):
+    # A retained cursor/blob would block this commit in rollback journal mode.
+    conn = sqlite3.connect(path, timeout=0)
+    try:
+        conn.execute("insert into another values (1)")
+        conn.commit()
+    finally:
+        conn.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("num_sql_threads", [0, 2])
+@pytest.mark.parametrize("kind", ["read", "write", "isolated", "immutable_isolated"])
+@pytest.mark.parametrize("started", [False, True])
+async def test_stashed_dump_iterator_expires_and_releases_lock(
+    resource_db, num_sql_threads, kind, started
+):
+    ds = Datasette(
+        **{("immutables" if kind == "immutable_isolated" else "files"): [resource_db]},
+        settings={"num_sql_threads": num_sql_threads},
+    )
+    db = ds.get_database("resources")
+    saved = {}
+
+    def callback(conn):
+        saved["dump"] = dump = conn.iterdump()
+        saved["next"] = dump.__next__
+        saved["send"] = dump.send
+        saved["close"] = dump.close
+        assert iter(dump) is dump
+        if started:
+            for statement in dump:
+                if statement.startswith('INSERT INTO "t"'):
+                    break
+            else:
+                pytest.fail("dump did not reach a data row")
+        return "done"
+
+    try:
+        assert (
+            await callback_method(db, kind.replace("immutable_", ""))(callback)
+            == "done"
+        )
+        for call in (
+            saved["next"],
+            lambda: iter(saved["dump"]),
+            lambda: saved["send"](None),
+            lambda: saved["dump"].throw(ValueError),
+            saved["close"],
+        ):
+            with pytest.raises(ConnectionLeaseError):
+                call()
+        if kind != "immutable_isolated":
+            check_no_read_lock(resource_db)
+        assert (await db.execute("select count(*) from t")).single_value() == 3
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("num_sql_threads", [0, 2])
+@pytest.mark.parametrize("kind", ["read", "write", "isolated"])
+async def test_dump_can_be_consumed_inside_callback(resource_db, num_sql_threads, kind):
+    ds = Datasette([resource_db], settings={"num_sql_threads": num_sql_threads})
+    try:
+        statements = await callback_method(ds.get_database("resources"), kind)(
+            lambda conn: list(conn.iterdump())
+        )
+        restored = sqlite3.connect(":memory:")
+        try:
+            restored.executescript("\n".join(statements))
+            assert (
+                restored.execute("select data from t where id=1").fetchone()[0]
+                == b"abcdef"
+            )
+        finally:
+            restored.close()
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("nested", [False, True])
+async def test_returning_dump_from_read_callback_is_rejected(resource_db, nested):
+    ds = Datasette([resource_db])
+
+    def callback(conn):
+        dump = conn.iterdump()
+        return {"dump": [dump]} if nested else dump
+
+    try:
+        with pytest.raises(ConnectionLeaseError, match="database resources"):
+            await ds.get_database("resources").execute_fn(callback)
+        check_no_read_lock(resource_db)
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+async def test_generator_over_dump_from_read_callback_is_rejected(resource_db):
+    ds = Datasette([resource_db])
+    try:
+        with pytest.raises(ConnectionLeaseError, match="database resources"):
+            await ds.get_database("resources").execute_fn(
+                lambda conn: (line for line in conn.iterdump())
+            )
+        check_no_read_lock(resource_db)
+    finally:
+        ds.close()
+
+
+requires_blob = pytest.mark.skipif(
+    not hasattr(sqlite3.Connection, "blobopen"),
+    reason="Incremental BLOB access requires Python 3.11",
+)
+
+
+@requires_blob
+@pytest.mark.asyncio
+@pytest.mark.parametrize("num_sql_threads", [0, 2])
+@pytest.mark.parametrize("kind", ["read", "write", "isolated", "immutable_isolated"])
+async def test_stashed_blob_expires_and_releases_lock(
+    resource_db, num_sql_threads, kind
+):
+    ds = Datasette(
+        **{("immutables" if kind == "immutable_isolated" else "files"): [resource_db]},
+        settings={"num_sql_threads": num_sql_threads},
+    )
+    saved = {}
+
+    def callback(conn):
+        saved["blob"] = blob = conn.blobopen("t", "data", 1, readonly=True)
+        saved["read"] = blob.read
+        saved["close"] = blob.close
+        assert isinstance(blob, sqlite3.Blob)
+        assert len(blob) == 6
+        assert blob[0] == ord("a")
+        assert blob[1:3] == b"bc"
+        assert blob.read(2) == b"ab"
+        assert blob.tell() == 2
+        blob.seek(-1, 2)
+        assert saved["read"]() == b"f"
+
+    try:
+        await callback_method(
+            ds.get_database("resources"), kind.replace("immutable_", "")
+        )(callback)
+        blob = saved["blob"]
+        for call in (
+            saved["read"],
+            saved["close"],
+            lambda: blob.write(b"x"),
+            lambda: blob.tell(),
+            lambda: blob.seek(0),
+            lambda: len(blob),
+            lambda: blob[0],
+            lambda: blob[:],
+            lambda: operator.setitem(blob, 0, 120),
+            lambda: operator.setitem(blob, slice(0, 1), b"x"),
+            blob.__enter__,
+            lambda: blob.__exit__(None, None, None),
+        ):
+            with pytest.raises(ConnectionLeaseError):
+                call()
+        if kind != "immutable_isolated":
+            check_no_read_lock(resource_db)
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("kind", ["write", "isolated"])
+@pytest.mark.parametrize(
+    "resource", ["dump", pytest.param("blob", marks=requires_blob)]
+)
+async def test_returned_write_resource_is_expired(resource_db, kind, resource):
+    ds = Datasette([resource_db])
+
+    def callback(conn):
+        if resource == "dump":
+            return conn.iterdump()
+        return conn.blobopen("t", "data", 1, readonly=True)
+
+    try:
+        result = await callback_method(ds.get_database("resources"), kind)(callback)
+        with pytest.raises(ConnectionLeaseError):
+            if resource == "dump":
+                next(result)
+            else:
+                result.read()
+        check_no_read_lock(resource_db)
+    finally:
+        ds.close()
+
+
+@requires_blob
+@pytest.mark.asyncio
+@pytest.mark.parametrize("num_sql_threads", [0, 2])
+async def test_writable_blob_closes_before_transaction_commits(
+    resource_db, num_sql_threads
+):
+    ds = Datasette([resource_db], settings={"num_sql_threads": num_sql_threads})
+    db = ds.get_database("resources")
+    saved = {}
+
+    def write(conn):
+        saved["blob"] = blob = conn.blobopen("t", "data", 1)
+        blob.write(b"xy")
+        blob[2] = ord("z")
+        blob[3:] = b"123"
+        # Leaving this BLOB open on SQLite makes the enclosing COMMIT fail
+        # with "SQL statements in progress"; expiry must close it first.
+        return "written"
+
+    try:
+        assert await db.execute_write_fn(write) == "written"
+        assert (
+            await db.execute("select data from t where id=1")
+        ).single_value() == b"xyz123"
+        with pytest.raises(ConnectionLeaseError):
+            saved["blob"].read()
+    finally:
+        ds.close()
+
+
+@requires_blob
+@pytest.mark.asyncio
+async def test_blob_context_manager_returns_lease(resource_db):
+    ds = Datasette([resource_db])
+    saved = {}
+
+    def write(conn):
+        with conn.blobopen("t", "data", 1) as blob:
+            saved["blob"] = blob
+            blob.write(b"xyz")
+        with pytest.raises(sqlite3.ProgrammingError):
+            blob.read()
+
+    try:
+        await ds.get_database("resources").execute_write_fn(write)
+        with pytest.raises(ConnectionLeaseError):
+            saved["blob"].read()
+    finally:
+        ds.close()
+
+
+@requires_blob
+@pytest.mark.asyncio
+@pytest.mark.parametrize("nested", [False, True])
+async def test_returning_blob_from_read_callback_is_rejected(resource_db, nested):
+    ds = Datasette([resource_db])
+
+    def callback(conn):
+        blob = conn.blobopen("t", "data", 1, readonly=True)
+        return {"blob": [blob]} if nested else blob
+
+    try:
+        with pytest.raises(ConnectionLeaseError, match="database resources"):
+            await ds.get_database("resources").execute_fn(callback)
+        check_no_read_lock(resource_db)
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("kind", ["read", "write", "isolated"])
+@pytest.mark.parametrize(
+    "resource", ["dump", pytest.param("blob", marks=requires_blob)]
+)
+async def test_resources_close_when_callback_raises(resource_db, kind, resource):
+    ds = Datasette([resource_db])
+    saved = {}
+
+    def fail(conn):
+        if resource == "dump":
+            saved["resource"] = value = conn.iterdump()
+            for statement in value:
+                if statement.startswith('INSERT INTO "t"'):
+                    break
+        else:
+            saved["resource"] = conn.blobopen("t", "data", 1, readonly=True)
+        raise ValueError("callback failed")
+
+    try:
+        with pytest.raises(ValueError, match="callback failed"):
+            await callback_method(ds.get_database("resources"), kind)(fail)
+        with pytest.raises(ConnectionLeaseError):
+            if resource == "dump":
+                next(saved["resource"])
+            else:
+                saved["resource"].read()
+        check_no_read_lock(resource_db)
+    finally:
+        ds.close()