mirror of
https://github.com/simonw/datasette.git
synced 2026-10-18 15:26:57 +02:00
Expire SQLite blob handles and dump iterators with their leases
This commit is contained in:
parent
162029b98f
commit
2e4bc5bed0
3 changed files with 442 additions and 15 deletions
|
|
@ -75,17 +75,20 @@ class LeasedConnection:
|
|||
C functions that require a real connection object (for example the
|
||||
*target* argument of ``Connection.backup()``) reject the proxy.
|
||||
|
||||
Cursors and saved bound methods share the lease. All cursors are closed
|
||||
on the owning thread before the connection is returned to the pool.
|
||||
Cursors, blobs, dump iterators and saved bound methods share the lease.
|
||||
Resources are closed on the owning thread before the connection is
|
||||
returned to the pool or the write transaction commits.
|
||||
"""
|
||||
|
||||
__slots__ = ("_conn", "_cursors", "_db_name", "_kind")
|
||||
__slots__ = ("_blobs", "_conn", "_cursors", "_db_name", "_iterators", "_kind")
|
||||
|
||||
def __init__(self, conn, db_name, kind="read"):
|
||||
object.__setattr__(self, "_conn", conn)
|
||||
object.__setattr__(self, "_db_name", db_name)
|
||||
object.__setattr__(self, "_kind", kind)
|
||||
object.__setattr__(self, "_cursors", [])
|
||||
object.__setattr__(self, "_blobs", [])
|
||||
object.__setattr__(self, "_iterators", [])
|
||||
|
||||
def _live(self):
|
||||
conn = self._conn
|
||||
|
|
@ -94,12 +97,18 @@ class LeasedConnection:
|
|||
f"A {self._kind} connection to database {self._db_name!r} was used "
|
||||
f"after the {_CALLBACK_NAMES.get(self._kind, 'callback')} callback "
|
||||
"it was passed to had returned. Connections are only valid inside "
|
||||
"the callback: do not store the connection (or a cursor) for "
|
||||
"the callback: do not store the connection or its resources for "
|
||||
"later use"
|
||||
)
|
||||
return conn
|
||||
|
||||
def _expire(self):
|
||||
for iterator in self._iterators:
|
||||
close_cursors(iterator)
|
||||
self._iterators.clear()
|
||||
for blob in self._blobs:
|
||||
_close_cursor(blob)
|
||||
self._blobs.clear()
|
||||
for cursor in self._cursors:
|
||||
_close_cursor(cursor)
|
||||
self._cursors.clear()
|
||||
|
|
@ -109,6 +118,17 @@ class LeasedConnection:
|
|||
if isinstance(result, sqlite3.Cursor):
|
||||
self._cursors.append(result)
|
||||
return LeasedCursor(result, self)
|
||||
# Incremental BLOB access was added in Python 3.11.
|
||||
blob_type = getattr(sqlite3, "Blob", None)
|
||||
if blob_type is not None and isinstance(result, blob_type):
|
||||
self._blobs.append(result)
|
||||
return LeasedBlob(result, self)
|
||||
if inspect.isgenerator(result):
|
||||
# iterdump() holds a raw connection, then a raw cursor once it
|
||||
# starts. Track even unstarted or stashed iterators so neither
|
||||
# the iterator nor its read lock can outlive this callback.
|
||||
self._iterators.append(result)
|
||||
return LeasedIterator(result, self)
|
||||
return result
|
||||
|
||||
# isinstance(conn, sqlite3.Connection) checks keep working
|
||||
|
|
@ -225,18 +245,94 @@ class LeasedCursor:
|
|||
self._live().close()
|
||||
|
||||
|
||||
def _contains_cursor(value, seen=None):
|
||||
class LeasedBlob:
|
||||
"""Incremental BLOB access lasts only as long as its connection lease."""
|
||||
|
||||
__slots__ = ("_blob", "_lease")
|
||||
|
||||
def __init__(self, blob, lease):
|
||||
self._blob = blob
|
||||
self._lease = lease
|
||||
|
||||
@property
|
||||
def __class__(self):
|
||||
return sqlite3.Blob
|
||||
|
||||
def _live(self):
|
||||
self._lease._live()
|
||||
return self._blob
|
||||
|
||||
def __getattr__(self, name):
|
||||
value = getattr(self._live(), name)
|
||||
if getattr(value, "__self__", None) is self._blob:
|
||||
|
||||
def call(*args, **kwargs):
|
||||
return getattr(self._live(), name)(*args, **kwargs)
|
||||
|
||||
return call
|
||||
return value
|
||||
|
||||
def __len__(self):
|
||||
return len(self._live())
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self._live()[key]
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
self._live()[key] = value
|
||||
|
||||
def __enter__(self):
|
||||
self._live().__enter__()
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
return self._live().__exit__(*args)
|
||||
|
||||
|
||||
class LeasedIterator:
|
||||
"""An iterator returned by SQLite cannot resume after its lease ends."""
|
||||
|
||||
__slots__ = ("_iterator", "_lease")
|
||||
|
||||
def __init__(self, iterator, lease):
|
||||
self._iterator = iterator
|
||||
self._lease = lease
|
||||
|
||||
def _live(self):
|
||||
self._lease._live()
|
||||
return self._iterator
|
||||
|
||||
def __iter__(self):
|
||||
self._live()
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
return next(self._live())
|
||||
|
||||
def send(self, value):
|
||||
return self._live().send(value)
|
||||
|
||||
def throw(self, *args):
|
||||
return self._live().throw(*args)
|
||||
|
||||
def close(self):
|
||||
return self._live().close()
|
||||
|
||||
|
||||
def _contains_database_resource(value, seen=None):
|
||||
"""Inspect ordinary containers without invoking arbitrary user iterators."""
|
||||
seen = set() if seen is None else seen
|
||||
if id(value) in seen:
|
||||
return False
|
||||
seen.add(id(value))
|
||||
if isinstance(value, sqlite3.Cursor):
|
||||
if isinstance(value, (sqlite3.Cursor, LeasedBlob, LeasedIterator)):
|
||||
return True
|
||||
if isinstance(value, dict):
|
||||
return any(_contains_cursor(v, seen) for pair in value.items() for v in pair)
|
||||
return any(
|
||||
_contains_database_resource(v, seen) for pair in value.items() for v in pair
|
||||
)
|
||||
if isinstance(value, (tuple, list, set, frozenset)):
|
||||
return any(_contains_cursor(v, seen) for v in value)
|
||||
return any(_contains_database_resource(v, seen) for v in value)
|
||||
return False
|
||||
|
||||
|
||||
|
|
@ -303,7 +399,10 @@ def _generator_uses_connection(result):
|
|||
values = inspect.getgeneratorlocals(result).values()
|
||||
except Exception: # noqa: BLE001
|
||||
return False
|
||||
return any(isinstance(v, (sqlite3.Connection, sqlite3.Cursor)) for v in values)
|
||||
return any(
|
||||
isinstance(v, (sqlite3.Connection, sqlite3.Cursor, LeasedBlob, LeasedIterator))
|
||||
for v in values
|
||||
)
|
||||
|
||||
|
||||
class _Entry:
|
||||
|
|
@ -364,7 +463,9 @@ class ReadConnectionPool:
|
|||
try:
|
||||
with callback_scope():
|
||||
result = fn(lease)
|
||||
if _contains_cursor(result) or _generator_uses_connection(result):
|
||||
if _contains_database_resource(result) or _generator_uses_connection(
|
||||
result
|
||||
):
|
||||
rejected = True
|
||||
close_cursors(result)
|
||||
except BaseException as e:
|
||||
|
|
|
|||
|
|
@ -2082,15 +2082,15 @@ Connections are only valid inside your function
|
|||
|
||||
The ``conn`` passed to an ``execute_fn()``, ``execute_write_fn()`` or ``execute_isolated_fn()`` function - and to :ref:`plugin_hook_write_wrapper` hooks - is lent to that function for the duration of the call. It behaves like a ``sqlite3.Connection`` (``isinstance(conn, sqlite3.Connection)`` is ``True`` and it works with ``sqlite_utils.Database(conn)``), but once the function returns it stops working: any further use raises ``datasette.database.ConnectionLeaseError``. Do not store the connection on an object, in a global or in a closure for later use. Do everything you need inside the function and return plain Python values.
|
||||
|
||||
``execute_fn()`` also raises ``ConnectionLeaseError`` if your function returns a cursor, or a generator that holds the connection or a cursor (for example ``(row for row in conn.execute(...))``), since those would keep reading from the connection after it had been returned to the pool. A generator over rows already fetched inside the function is fine.
|
||||
``execute_fn()`` also raises ``ConnectionLeaseError`` if your function returns a cursor, a blob handle, a dump iterator, or a generator that holds the connection or one of those resources (for example ``(row for row in conn.execute(...))``), since those would keep reading from the connection after it had been returned to the pool. This also checks resources in ordinary containers such as lists and dictionaries. A generator over rows already fetched inside the function is fine.
|
||||
|
||||
The write connection belongs to Datasette: calling ``conn.close()`` inside an ``execute_write_fn()`` function raises ``ConnectionLeaseError``. Closing the connection passed to an ``execute_fn()`` or ``execute_isolated_fn()`` function is allowed.
|
||||
|
||||
Datasette cannot detect every misuse. These are unsupported and may fail unpredictably, or read from a connection that is now being used by someone else:
|
||||
Cursors, blob handles from ``conn.blobopen()`` (Python 3.11 and later), and iterators from ``conn.iterdump()`` share the connection's lifetime. Using any of them after the function returns raises ``ConnectionLeaseError``, including through a saved method such as ``read = blob.read``. ``cursor.connection`` returns the same leased connection. Datasette closes these resources on the owning thread before returning a read connection to the pool or committing the write transaction, even if the function stores them elsewhere or raises an exception. This releases their locks and prevents an unfinished dump iterator or open blob from blocking later writes.
|
||||
|
||||
- Iterating over or fetching from a **cursor** created inside the function after the function has returned. Write functions may return a cursor (``lambda conn: conn.execute(...)`` is a common pattern) and reading attributes such as ``.lastrowid`` and ``.rowcount`` from it afterwards is fine; fetching rows from it is not. A cursor returned by a write or isolated function - including the ones returned by ``execute_write_many()`` and ``execute_write_script()`` - is closed before it is handed back, on the thread that owns the connection, and so are cursors left in the local variables of a function that raised an exception. (On Python versions before 3.12, a cursor that was freed on another thread could reset the same statement while the connection was running it again, making that query fail with "bad parameter or other API misuse".)
|
||||
- Reaching the underlying connection through ``cursor.connection``.
|
||||
- Relying on state set up on the connection by an earlier call (see above).
|
||||
Write functions may return a cursor (``lambda conn: conn.execute(...)`` is a common pattern), and reading attributes such as ``.lastrowid``, ``.rowcount`` and ``.description`` afterwards is fine; fetching rows from it is not. This includes cursors returned by ``execute_write_many()`` and ``execute_write_script()``. Consume dump iterators and read or modify blobs inside the function, then return ordinary Python values.
|
||||
|
||||
Do not rely on state set up on the connection by an earlier call (see above). Lease proxies enforce normal API use; they are not a sandbox for untrusted Python code.
|
||||
|
||||
The lent connection cannot be used as the *target* of another connection's ``backup()`` method, because that C function needs the real connection object. Using it as the source, ``conn.backup(other)``, works.
|
||||
|
||||
|
|
|
|||
326
tests/test_connection_resource_leases.py
Normal file
326
tests/test_connection_resource_leases.py
Normal file
|
|
@ -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()
|
||||
Loading…
Add table
Add a link
Reference in a new issue