Expire escaped cursors and retry schema scans after resource failures

datasette/connection_pool.py

@@ -75,16 +75,17 @@ class LeasedConnection:
     C functions that require a real connection object (for example the
     *target* argument of ``Connection.backup()``) reject the proxy.
 
-    Not enforceable: a cursor created inside the callback and used after it
-    returns, and the raw connection reached through ``cursor.connection``.
+    Cursors and saved bound methods share the lease. All cursors are closed
+    on the owning thread before the connection is returned to the pool.
     """
 
-    __slots__ = ("_conn", "_db_name", "_kind")
+    __slots__ = ("_conn", "_db_name", "_kind", "_cursors")
 
     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", [])
 
     def _live(self):
         conn = self._conn
@@ -99,15 +100,32 @@ class LeasedConnection:
         return conn
 
     def _expire(self):
+        for cursor in self._cursors:
+            _close_cursor(cursor)
+        self._cursors.clear()
         object.__setattr__(self, "_conn", None)
 
+    def _wrap(self, result):
+        if isinstance(result, sqlite3.Cursor):
+            self._cursors.append(result)
+            return LeasedCursor(result, self)
+        return result
+
     # isinstance(conn, sqlite3.Connection) checks keep working
     @property
     def __class__(self):
         return sqlite3.Connection
 
     def __getattr__(self, name):
-        return getattr(self._live(), name)
+        value = getattr(self._live(), name)
+        if getattr(value, "__self__", None) is self._conn:
+            # Resolve the method again when called: caching a bound method
+            # must not bypass expiry or keep the raw connection alive.
+            def call(*args, **kwargs):
+                return self._wrap(getattr(self._live(), name)(*args, **kwargs))
+
+            return call
+        return value
 
     def __setattr__(self, name, value):
         setattr(self._live(), name, value)
@@ -129,10 +147,10 @@ class LeasedConnection:
 
     # Hot paths, to skip __getattr__
     def execute(self, *args, **kwargs):
-        return self._live().execute(*args, **kwargs)
+        return self._wrap(self._live().execute(*args, **kwargs))
 
     def cursor(self, *args, **kwargs):
-        return self._live().cursor(*args, **kwargs)
+        return self._wrap(self._live().cursor(*args, **kwargs))
 
     def close(self):
         conn = self._live()
@@ -154,6 +172,74 @@ class LeasedConnection:
         return f"<LeasedConnection {self._kind} database={self._db_name!r} {state}>"
 
 
+class LeasedCursor:
+    """A cursor cannot outlive the callback that owns its connection.
+
+    Metadata remains readable after closure for execute_write_many() and
+    callbacks returning insertion cursors. cursor.connection never exposes
+    the raw SQLite connection.
+    """
+
+    __slots__ = ("_cursor", "_lease")
+
+    def __init__(self, cursor, lease):
+        object.__setattr__(self, "_cursor", cursor)
+        object.__setattr__(self, "_lease", lease)
+
+    @property
+    def __class__(self):
+        return sqlite3.Cursor
+
+    @property
+    def connection(self):
+        return self._lease
+
+    def _live(self):
+        self._lease._live()
+        return self._cursor
+
+    def __getattr__(self, name):
+        if name in ("rowcount", "lastrowid", "description"):
+            return getattr(self._cursor, name)
+        value = getattr(self._live(), name)
+        if getattr(value, "__self__", None) is self._cursor:
+
+            def call(*args, **kwargs):
+                result = getattr(self._live(), name)(*args, **kwargs)
+                return self if result is self._cursor else result
+
+            return call
+        return value
+
+    def __setattr__(self, name, value):
+        setattr(self._live(), name, value)
+
+    def __iter__(self):
+        self._live()
+        return self
+
+    def __next__(self):
+        return next(self._live())
+
+    def close(self):
+        self._live().close()
+
+
+def _contains_cursor(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):
+        return True
+    if isinstance(value, dict):
+        return any(_contains_cursor(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 False
+
+
 def close_cursors(value):
     """Close the cursors in a callback's result (a cursor, or a tuple or
     list holding one), in a generator's locals, or - for an exception - in
@@ -274,7 +360,7 @@ class ReadConnectionPool:
         rejected = False
         try:
             result = fn(lease)
-            if isinstance(result, sqlite3.Cursor) or _generator_uses_connection(result):
+            if _contains_cursor(result) or _generator_uses_connection(result):
                 rejected = True
                 close_cursors(result)
         except BaseException as e:
@@ -285,7 +371,7 @@ class ReadConnectionPool:
             self.release(entry)
         if rejected:
             raise ConnectionLeaseError(
-                f"An execute_fn() callback for database {db.name!r} returned a {type(result).__name__}, which "
+                f"An execute_fn() callback for database {db.name!r} returned a {result.__class__.__name__} containing database resources, which "
                 "would keep using the pooled connection after the callback "
                 "returned. Fetch the rows inside the callback instead"
             )

datasette/schema_watcher.py

@@ -74,6 +74,7 @@ open.
 """
 
 import asyncio
+import errno
 import json
 import logging
 import os
@@ -250,9 +251,22 @@ def header_schema_version(path):
 
 
 def _is_transient(error):
+    if isinstance(error, OSError):
+        return error.errno in (errno.EMFILE, errno.ENFILE, errno.ENOMEM, errno.EAGAIN)
     if isinstance(error, sqlite3.OperationalError):
-        message = str(error)
-        return "locked" in message or "busy" in message
+        code = getattr(error, "sqlite_errorcode", 0) & 0xFF
+        if code in (sqlite3.SQLITE_CANTOPEN, sqlite3.SQLITE_NOMEM):
+            return True
+        message = str(error).lower()
+        return any(
+            text in message
+            for text in (
+                "locked",
+                "busy",
+                "unable to open database file",
+                "too many open files",
+            )
+        )
     return False
 
 
@@ -1041,11 +1055,12 @@ class SchemaWatcher:
                     handled.add(state)
                     continue
                 if _is_transient(r["error"]):
+                    state.error = None
                     # "database is locked" / "database table is locked" (a
                     # shared-cache memory database being filled from another
                     # connection): not a property of the database, retry
                     logger.info(
-                        "Schema of database %r is locked, will retry: %s",
+                        "Transient error reading schema of database %r, will retry: %s",
                         state.name,
                         r["error"],
                     )

tests/test_connection_recovery.py

@@ -0,0 +1,112 @@
+import errno
+import sqlite3
+
+import pytest
+
+from datasette.app import Datasette
+from datasette.connection_pool import ConnectionLeaseError
+
+
+@pytest.fixture
+def database_file(tmp_path):
+    path = tmp_path / "data.db"
+    conn = sqlite3.connect(path)
+    conn.executescript("create table t(id); insert into t values (1), (2)")
+    conn.close()
+    return str(path)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("kind", ["read", "write", "isolated"])
+async def test_saved_methods_and_cursors_cannot_escape_callback(database_file, kind):
+    ds = Datasette([database_file], settings={"schema_watch_interval_ms": 0})
+    db = ds.get_database("data")
+    saved = {}
+
+    def callback(conn):
+        saved["commit"] = conn.commit
+        saved["cursor"] = conn.execute("select * from t")
+        saved["fetch"] = saved["cursor"].fetchone
+        saved["connection"] = saved["cursor"].connection
+        assert saved["fetch"]()[0] == 1
+
+    method = {
+        "read": db.execute_fn,
+        "write": db.execute_write_fn,
+        "isolated": db.execute_isolated_fn,
+    }[kind]
+    try:
+        await method(callback)
+        for call in (saved["commit"], lambda: saved["connection"].execute("select 1")):
+            with pytest.raises(ConnectionLeaseError):
+                call()
+        for call in (saved["fetch"], lambda: next(saved["cursor"])):
+            with pytest.raises((ConnectionLeaseError, sqlite3.ProgrammingError)):
+                call()
+        assert (await db.execute("select count(*) from t")).single_value() == 2
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+async def test_read_result_containing_cursor_is_rejected(database_file):
+    ds = Datasette([database_file])
+    db = ds.get_database("data")
+    try:
+        with pytest.raises(ConnectionLeaseError):
+            await db.execute_fn(lambda c: {"nested": [c.execute("select * from t")]})
+        result = await db.execute_write_fn(
+            lambda c: {"nested": [c.execute("insert into t values (3)")]}
+        )
+        assert result["nested"][0].rowcount == 1
+        with pytest.raises((ConnectionLeaseError, sqlite3.ProgrammingError)):
+            result["nested"][0].fetchone()
+        assert (await db.execute("select count(*) from t")).single_value() == 3
+    finally:
+        ds.close()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "failure",
+    [
+        sqlite3.OperationalError("unable to open database file"),
+        OSError(errno.EMFILE, "Too many open files"),
+    ],
+)
+async def test_schema_recovers_after_resource_error_without_file_change(
+    database_file, monkeypatch, failure
+):
+    ds = Datasette([database_file], settings={"schema_watch_interval_ms": 0})
+    try:
+        await ds.invoke_startup()
+        db = ds.get_database("data")
+        conn = sqlite3.connect(database_file)
+        conn.execute("create table added_externally(id)")
+        conn.close()
+        watcher = ds._schema_watcher
+        original = watcher._connect
+
+        def fail(database, prepare):
+            if database is db and prepare:
+                raise failure
+            return original(database, prepare)
+
+        with monkeypatch.context() as patch:
+            patch.setattr(watcher, "_connect", fail)
+            await ds.refresh_schemas()
+        state = db._watch_state
+        assert state.needs_scan
+        # Move past the retry backoff and fingerprint race window without
+        # touching the file or relying on timing-sensitive sleeps.
+        state.retry_at = 0
+        state.fp_racy = False
+        await watcher.sweep(background=True)
+        tables = await ds.get_internal_database().execute(
+            "select table_name from catalog_tables where database_name=?", ["data"]
+        )
+        assert {row[0] for row in tables.rows} == {"t", "added_externally"}
+        assert state.error is None
+        assert not state.needs_scan
+    finally:
+        ds.close()

tests/test_read_pool.py

@@ -337,7 +337,7 @@ async def test_leased_connection_discarded_on_release_after_invalidation(tmp_pat
     assert await db.execute_fn(callback) == 1
     assert db._read_pool_state.open == 0
     assert pool.stats["discarded_stale"] == 1
-    with pytest.raises(sqlite3.ProgrammingError):
+    with pytest.raises(ConnectionLeaseError):
         seen["raw"].execute("select 1")
     assert (await db.execute("select count(*) from t")).single_value() == 1
     ds.close()