From 791ce7a3d9fcf136ee3c3b0580d97d48692603a1 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 4 Oct 2026 01:26:01 +0000 Subject: [PATCH] Spike: leased read-connection pool with global cap and idle reaper Replace the per-thread cached read connections with a pool. Each execute_fn()/execute() callback checks a connection out for its duration and gets a LeasedConnection proxy that raises ConnectionLeaseError if used after the callback returns. New settings: max_open_connections (global LRU cap, default 128), connection_idle_timeout (seconds, default 60, enforced by one process-wide reaper thread plus opportunistic reaping on checkout), connection_pool_wait_ms (soft cap: wait then exceed) and pool_read_connections (off = fresh connection per query, a control). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01EwkLo2rJWidpkDVT3HkvNv --- datasette/app.py | 43 +++ datasette/connection_pool.py | 464 +++++++++++++++++++++++++++++++ datasette/database.py | 29 +- docs/cli-reference.rst | 13 + docs/internals.rst | 2 + docs/settings.rst | 66 +++++ tests/test_api.py | 4 + tests/test_internals_database.py | 10 +- tests/test_read_pool.py | 284 +++++++++++++++++++ 9 files changed, 900 insertions(+), 15 deletions(-) create mode 100644 datasette/connection_pool.py create mode 100644 tests/test_read_pool.py diff --git a/datasette/app.py b/datasette/app.py index 3e4c5acf..66d5b0ba 100644 --- a/datasette/app.py +++ b/datasette/app.py @@ -230,6 +230,26 @@ SETTINGS = ( 3, "Number of threads in the thread pool for executing SQLite queries", ), + Setting( + "max_open_connections", + 128, + "Maximum number of pooled read connections open across all databases - least recently used idle connections are closed first, 0 for no limit", + ), + Setting( + "connection_idle_timeout", + 60, + "Close pooled read connections that have been idle for this many seconds - 0 to keep them open", + ), + Setting( + "connection_pool_wait_ms", + 0, + "When max_open_connections is reached and every connection is in use, wait this long for one to be returned before opening one more", + ), + Setting( + "pool_read_connections", + True, + "Reuse read connections between queries - turn off to open a new connection for every query", + ), Setting("sql_time_limit_ms", 1000, "Time limit for a SQL query in milliseconds"), Setting( "default_facet_size", 30, "Number of values to return for requested facets" @@ -991,6 +1011,9 @@ class Datasette: # Collect the first failure and re-raise after every close() has run if first_exception is None: first_exception = e + read_pool = self._read_pool_or_none + if read_pool is not None: + read_pool.close() if self.executor is not None: try: self.executor.shutdown(wait=True, cancel_futures=True) @@ -1003,6 +1026,26 @@ class Datasette: def setting(self, key): return self._settings.get(key, None) + @property + def _read_pool(self): + pool = self.__dict__.get("_read_pool_instance") + if pool is None: + from .connection_pool import ReadConnectionPool + + pool = ReadConnectionPool( + self._prepare_connection, + max_open=self.setting("max_open_connections") or 0, + idle_timeout=float(self.setting("connection_idle_timeout") or 0), + wait_ms=self.setting("connection_pool_wait_ms") or 0, + enabled=bool(self.setting("pool_read_connections")), + ) + pool = self.__dict__.setdefault("_read_pool_instance", pool) + return pool + + @property + def _read_pool_or_none(self): + return self.__dict__.get("_read_pool_instance") + def settings_dict(self): # Returns a fully resolved settings dictionary, useful for templates return {option.name: self.setting(option.name) for option in SETTINGS} diff --git a/datasette/connection_pool.py b/datasette/connection_pool.py new file mode 100644 index 00000000..aa6fb9da --- /dev/null +++ b/datasette/connection_pool.py @@ -0,0 +1,464 @@ +"""Leased read-connection pool. + +Read connections are no longer cached per worker thread. Instead every +``execute_fn()`` / ``execute()`` call checks a connection out of a pool for +the duration of one callback and returns it afterwards. + +* One ``ReadConnectionPool`` per ``Datasette`` instance. It enforces a global + cap (``max_open_connections``) on open pooled read connections across all + file-backed databases. When a new connection is needed and the cap has been + reached, the least recently used *idle* connection (from any database) is + closed first. +* Each database keeps a LIFO stack of its idle connections, so the most + recently used (warmest) connection is reused and surplus ones age out. +* Idle connections are closed after ``connection_idle_timeout`` seconds. This + is done opportunistically on every checkout and by one process-wide daemon + reaper thread, so it works with or without a running event loop + (non-threaded mode, sync callers, closed loops in tests). +* When the cap has been reached and every connection is checked out the pool + waits up to ``connection_pool_wait_ms`` for one to come back and then opens + a connection anyway (the cap is soft). Each checked-out connection occupies + one executor thread, so the overflow can never exceed ``num_sql_threads``. +* Callbacks receive a ``LeasedConnection`` proxy. Once the callback returns + the proxy is invalidated and any further use raises + ``ConnectionLeaseError``. +* In-memory databases are leased the same way but do not count against the + cap and are never reaped: they use no file descriptors, and closing every + connection to a named in-memory database would destroy its contents. +""" + +import collections +import inspect +import threading +import time +import weakref + +from .utils import sqlite3 + + +class ConnectionLeaseError(RuntimeError): + """A pooled read connection was used outside the callback it was lent to.""" + + +class LeasedConnection: + """Forwarding proxy handed to read callbacks. + + Every attribute access, method call and attribute assignment is forwarded + to the underlying ``sqlite3.Connection`` while the lease is active. After + the callback returns the proxy is expired and any use raises + ``ConnectionLeaseError``. Each lease gets a fresh proxy, so a stale + reference fails even if the same underlying connection has since been + lent to another callback. + + ``isinstance(proxy, sqlite3.Connection)`` is True (via ``__class__``), but + C functions that require a real connection object (for example the + *target* argument of ``Connection.backup()``) reject the proxy. + """ + + __slots__ = ("_conn", "_db_name") + + def __init__(self, conn, db_name): + object.__setattr__(self, "_conn", conn) + object.__setattr__(self, "_db_name", db_name) + + def _live(self): + conn = self._conn + if conn is None: + raise ConnectionLeaseError( + f"A read connection to database {self._db_name!r} was used after the " + "execute_fn() callback it was passed to had returned. Read " + "connections are pooled and only valid inside the callback: do " + "not store the connection (or a cursor) for later use" + ) + return conn + + def _expire(self): + object.__setattr__(self, "_conn", None) + + # isinstance(conn, sqlite3.Connection) checks keep working + @property + def __class__(self): + return sqlite3.Connection + + def __getattr__(self, name): + return getattr(self._live(), name) + + def __setattr__(self, name, value): + setattr(self._live(), name, value) + + def __delattr__(self, name): + delattr(self._live(), name) + + # Dunder methods are looked up on the type, so forward them explicitly + def __enter__(self): + return self._live().__enter__() + + def __exit__(self, *args): + return self._live().__exit__(*args) + + def __call__(self, *args, **kwargs): + return self._live()(*args, **kwargs) + + # Hot paths, to skip __getattr__ + def execute(self, *args, **kwargs): + return self._live().execute(*args, **kwargs) + + def cursor(self, *args, **kwargs): + return self._live().cursor(*args, **kwargs) + + def close(self): + # Closing a pooled connection is allowed; the pool notices on release + # and discards it instead of returning it to the idle stack. + self._live().close() + + def __repr__(self): + state = "expired" if self._conn is None else "active" + return f"" + + +class _Entry: + __slots__ = ("conn", "last_used", "state") + + def __init__(self, conn, state): + self.conn = conn + self.state = state + self.last_used = 0.0 + + +class _DatabaseState: + """Per-database bookkeeping, stored on the Database as _read_pool_state.""" + + __slots__ = ("closed", "counted", "db", "idle", "leased", "open") + + def __init__(self, db): + self.db = db + self.idle = [] # LIFO stack of _Entry + self.open = 0 # idle + leased + self.leased = 0 + self.closed = False + # In-memory databases use no file descriptors and are never evicted + self.counted = not db.is_memory + + +class ReadConnectionPool: + def __init__( + self, + prepare_connection, + max_open=128, + idle_timeout=60.0, + wait_ms=0, + enabled=True, + ): + self._prepare_connection = prepare_connection + self.max_open = max_open # 0 = no limit + self.idle_timeout = idle_timeout # seconds, 0 = never + self.wait_s = wait_ms / 1000.0 + self.enabled = enabled + self._lock = threading.Lock() + self._cond = threading.Condition(self._lock) + # Idle, counted entries, least recently used first + self._lru = collections.OrderedDict() + self._open = 0 # counted open connections, idle + leased + self._waiters = 0 + self._closed = False + self.stats = collections.Counter() + + # -- public API --------------------------------------------------------- + + def run(self, db, fn): + """Call fn(conn) with a pooled connection leased for the duration.""" + entry = self.acquire(db) + lease = LeasedConnection(entry.conn, db.name) + try: + result = fn(lease) + finally: + lease._expire() + self.release(entry) + if isinstance(result, sqlite3.Cursor) or inspect.isgenerator(result): + raise ConnectionLeaseError( + f"An execute_fn() callback for database {db.name!r} returned a {type(result).__name__}, which " + "would keep using the pooled connection after the callback " + "returned. Fetch the rows inside the callback instead" + ) + return result + + def acquire(self, db): + state = self._state(db) + to_close = [] + entry = None + with self._lock: + if self._closed or state.closed: + from .database import DatasetteClosedError + + raise DatasetteClosedError(f"Database {db.name!r} has been closed") + now = time.monotonic() + self._collect_expired_locked(now, to_close) + deadline = None + while True: + if state.idle: + entry = state.idle.pop() + if state.counted: + del self._lru[entry] + state.leased += 1 + self.stats["reused"] += 1 + break + if not state.counted or not self.max_open or self._open < self.max_open: + break + if self._lru: + # Make room by closing the least recently used idle + # connection, which may belong to any database + victim, _ = self._lru.popitem(last=False) + victim.state.idle.remove(victim) + victim.state.open -= 1 + self._open -= 1 + to_close.append(victim) + self.stats["evicted_lru"] += 1 + continue + # Cap reached and every connection is checked out + if self.wait_s > 0: + if deadline is None: + deadline = now + self.wait_s + self.stats["waits"] += 1 + remaining = deadline - time.monotonic() + if remaining > 0: + self._waiters += 1 + t0 = time.monotonic() + try: + self._cond.wait(remaining) + finally: + self._waiters -= 1 + self.stats["wait_ms_total"] += int( + (time.monotonic() - t0) * 1000 + ) + continue + self.stats["wait_timeouts"] += 1 + # Soft cap: open one more. Bounded by the number of threads + # that can hold a lease at once (num_sql_threads). + self.stats["exceeded_cap"] += 1 + break + if entry is None: + state.open += 1 + state.leased += 1 + self.stats["opened"] += 1 + if state.counted: + self._open += 1 + self.stats["peak_open"] = max(self.stats["peak_open"], self._open) + self._close_entries(to_close) + if entry is not None: + return entry + # Open and prepare outside the lock + conn = None + try: + conn = db.connect() + self._prepare_connection(conn, db.name) + except BaseException: + with self._lock: + state.open -= 1 + state.leased -= 1 + self.stats["opened"] -= 1 + self.stats["open_failed"] += 1 + if state.counted: + self._open -= 1 + if self._waiters: + self._cond.notify() + if conn is not None: + self._close_conn(state, conn) + raise + return _Entry(conn, state) + + def release(self, entry): + state = entry.state + conn = entry.conn + discard = False + try: + # Hygiene: a callback that began a transaction and did not end it + # must not hand that transaction to the next borrower + rolled_back = False + if conn.in_transaction: + conn.rollback() + rolled_back = True + except Exception: # noqa: BLE001 + # e.g. ProgrammingError: the callback closed the connection + discard = True + notify_reaper = None + with self._lock: + if discard: + self.stats["discarded_broken"] += 1 + elif rolled_back: + self.stats["rolled_back"] += 1 + state.leased -= 1 + over_cap = state.counted and self.max_open and self._open > self.max_open + if discard or over_cap or self._closed or state.closed or not self.enabled: + state.open -= 1 + if state.counted: + self._open -= 1 + discard = True + if over_cap: + self.stats["closed_over_cap"] += 1 + else: + entry.last_used = time.monotonic() + state.idle.append(entry) + if state.counted: + self._lru[entry] = None + if self.idle_timeout > 0: + notify_reaper = entry.last_used + self.idle_timeout + if self._waiters: + self._cond.notify() + if discard: + self._close_conn(state, conn) + if notify_reaper is not None: + _reaper.wake(self, notify_reaper) + + def reap_expired(self): + """Close expired idle connections. Returns the next expiry time + (time.monotonic() based) or None if nothing is idle.""" + to_close = [] + with self._lock: + self._collect_expired_locked(time.monotonic(), to_close) + next_expiry = None + if self._lru and self.idle_timeout > 0: + head = next(iter(self._lru)) + next_expiry = head.last_used + self.idle_timeout + self._close_entries(to_close) + return next_expiry + + def close_database(self, db): + """Close idle connections for db; leased ones are closed on release.""" + state = getattr(db, "_read_pool_state", None) + if state is None: + return + with self._lock: + state.closed = True + to_close = state.idle + state.idle = [] + for entry in to_close: + if state.counted: + self._lru.pop(entry, None) + self._open -= 1 + state.open -= 1 + self._close_entries(to_close) + + def close(self): + with self._lock: + self._closed = True + to_close = list(self._lru) + self._lru.clear() + for entry in to_close: + entry.state.idle.remove(entry) + entry.state.open -= 1 + self._open -= 1 + self._cond.notify_all() + self._close_entries(to_close) + + def snapshot(self): + with self._lock: + return { + "open": self._open, + "idle": len(self._lru), + "max_open": self.max_open, + "idle_timeout": self.idle_timeout, + "stats": dict(self.stats), + } + + # -- internals ---------------------------------------------------------- + + def _state(self, db): + state = db._read_pool_state + if state is None: + with self._lock: + state = db._read_pool_state + if state is None: + state = db._read_pool_state = _DatabaseState(db) + return state + + def _collect_expired_locked(self, now, to_close): + if self.idle_timeout <= 0 or not self._lru: + return + cutoff = now - self.idle_timeout + lru = self._lru + while lru: + head = next(iter(lru)) + if head.last_used > cutoff: + break + lru.popitem(last=False) + head.state.idle.remove(head) + head.state.open -= 1 + self._open -= 1 + to_close.append(head) + self.stats["expired"] += 1 + + def _close_entries(self, entries): + for entry in entries: + self._close_conn(entry.state, entry.conn) + + def _close_conn(self, state, conn): + try: + conn.close() + except Exception: # noqa: BLE001, S110 + pass + with self._lock: + self.stats["closed"] += 1 + try: + state.db._all_connections.remove(conn) + except ValueError: + # May already have been cleared by Database.close() + pass + + +class _Reaper: + """One daemon thread per process that closes expired idle connections + for every live pool. Pools are held by weak reference.""" + + def __init__(self): + self._cond = threading.Condition() + self._pools = weakref.WeakSet() + self._thread = None + self._dirty = False + # monotonic time of the next planned scan, inf while sleeping + # indefinitely or while scanning + self.wake_at = float("inf") + + def wake(self, pool, expiry): + # Unlocked pre-check: most releases do not need to wake the reaper + if expiry >= self.wake_at and pool in self._pools: + return + with self._cond: + self._pools.add(pool) + if self._thread is None or not self._thread.is_alive(): + self._thread = threading.Thread( + target=self._run, name="datasette-read-pool-reaper", daemon=True + ) + self._thread.start() + if expiry < self.wake_at: + self._dirty = True + self._cond.notify() + + def _run(self): + while True: + with self._cond: + self.wake_at = float("inf") + self._dirty = False + pools = list(self._pools) + next_wake = None + for pool in pools: + try: + expiry = pool.reap_expired() + except Exception: # noqa: BLE001 + expiry = None + if expiry is not None and (next_wake is None or expiry < next_wake): + next_wake = expiry + del pools + with self._cond: + if self._dirty: + continue + if next_wake is None: + self.wake_at = float("inf") + self._cond.wait() + else: + self.wake_at = next_wake + # Small slack so a batch of connections that expire close + # together is handled in one pass + self._cond.wait(max(0.0, next_wake - time.monotonic()) + 0.05) + + +_reaper = _Reaper() diff --git a/datasette/database.py b/datasette/database.py index 542b3012..68ab37e0 100644 --- a/datasette/database.py +++ b/datasette/database.py @@ -16,6 +16,7 @@ import sqlite_utils from opentelemetry import context as otel_context_api from opentelemetry.trace import Status, StatusCode +from .connection_pool import ConnectionLeaseError # noqa: F401 from .inspect import inspect_hash from .telemetry import ( callback_name, @@ -127,9 +128,11 @@ class Database: self._closed = False self._pending_execute_futures = set() self._pending_execute_futures_lock = threading.Lock() - # These are used when in non-threaded mode: + # Used when in non-threaded mode (reads go through ds._read_pool): self._read_connection = None self._write_connection = None + # Bookkeeping for ds._read_pool, created on first read + self._read_pool_state = None # Track file and memory connections, including reads on worker threads, # so close() can release all of them from the calling thread. self._all_connections = [] @@ -242,8 +245,13 @@ class Database: except Exception: # noqa: BLE001, S110 # Shutdown teardown - a failed pending write must not block close() pass + # Close idle pooled read connections; any still leased are closed + # by the pool when they are released + read_pool = getattr(self.ds, "_read_pool_or_none", None) + if read_pool is not None: + read_pool.close_database(self) # Close anything still tracked in _all_connections - for connection in self._all_connections: + for connection in list(self._all_connections): try: connection.close() except Exception: # noqa: BLE001, S110 @@ -686,21 +694,14 @@ class Database: async def _execute_fn(self, fn): self._check_not_closed() + read_pool = self.ds._read_pool if self.ds.executor is None: - # non-threaded mode - if self._read_connection is None: - self._read_connection = self.connect() - self.ds._prepare_connection(self._read_connection, self.name) - return fn(self._read_connection) + # non-threaded mode: lease on the event loop thread + return read_pool.run(self, fn) - # threaded mode + # threaded mode: lease on the executor thread for one callback def in_thread(): - conn = getattr(connections, self._thread_local_id, None) - if not conn: - conn = self.connect() - self.ds._prepare_connection(conn, self.name) - setattr(connections, self._thread_local_id, conn) - return fn(conn) + return read_pool.run(self, fn) with self._pending_execute_futures_lock: self._check_not_closed() diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst index db903e8e..00e52867 100644 --- a/docs/cli-reference.rst +++ b/docs/cli-reference.rst @@ -251,6 +251,19 @@ These can be passed to ``datasette serve`` using ``datasette serve --setting nam this limit (default=2097152) num_sql_threads Number of threads in the thread pool for executing SQLite queries (default=3) + max_open_connections Maximum number of pooled read connections open + across all databases - least recently used idle + connections are closed first, 0 for no limit + (default=128) + connection_idle_timeout Close pooled read connections that have been idle + for this many seconds - 0 to keep them open + (default=60) + connection_pool_wait_ms When max_open_connections is reached and every + connection is in use, wait this long for one to + be returned before opening one more (default=0) + pool_read_connections Reuse read connections between queries - turn off + to open a new connection for every query + (default=True) sql_time_limit_ms Time limit for a SQL query in milliseconds (default=1000) default_facet_size Number of values to return for requested facets diff --git a/docs/internals.rst b/docs/internals.rst index 6ecf5b73..6aa3f56c 100644 --- a/docs/internals.rst +++ b/docs/internals.rst @@ -2174,6 +2174,8 @@ Example usage: version = await db.execute_fn(get_version) +The connection is borrowed from a pool for the duration of the call (see :ref:`setting_max_open_connections`) and is only valid inside the function. Using it after the function has returned - for example by storing it on an object - raises a ``datasette.database.ConnectionLeaseError``. Returning a cursor or a generator from the function raises the same error: fetch the rows you need before returning. + The call is traced as a ``db.query`` OpenTelemetry span carrying ``datasette.callback`` (the function's qualified name) rather than ``db.query.text``, since the SQL is whatever the function chooses to run - see :ref:`internals_telemetry`. Passing a named function gives the span a readable identity; a lambda reports ````. .. _database_execute_write: diff --git a/docs/settings.rst b/docs/settings.rst index f3e6636d..facc0bf8 100644 --- a/docs/settings.rst +++ b/docs/settings.rst @@ -214,6 +214,72 @@ Maximum number of threads in the thread pool Datasette uses to execute SQLite qu Setting this to 0 turns off threaded SQL queries entirely - useful for environments that do not support threading such as `Pyodide `__. +.. _setting_max_open_connections: + +max_open_connections +~~~~~~~~~~~~~~~~~~~~ + +.. [[[cog + setting_default(cog, "max_open_connections") +.. ]]] + +Default: ``128`` + +.. [[[end]]] + +Read queries borrow a connection from a pool for the duration of one query and return it afterwards. This setting caps how many pooled read connections can be open at once across all attached file databases. When a new connection is needed and the cap has been reached, the least recently used idle connection is closed first. + +If every connection is in use the pool waits up to :ref:`setting_connection_pool_wait_ms` and then opens one more anyway, so the cap can be exceeded by at most :ref:`setting_num_sql_threads` connections. Set this to 0 for no limit. + +:: + + datasette *.db --setting max_open_connections 64 + +.. _setting_connection_idle_timeout: + +connection_idle_timeout +~~~~~~~~~~~~~~~~~~~~~~~ + +.. [[[cog + setting_default(cog, "connection_idle_timeout") +.. ]]] + +Default: ``60`` + +.. [[[end]]] + +Pooled read connections that have not been used for this many seconds are closed. Set to 0 to keep idle connections open until they are evicted by :ref:`setting_max_open_connections`. + +.. _setting_connection_pool_wait_ms: + +connection_pool_wait_ms +~~~~~~~~~~~~~~~~~~~~~~~ + +.. [[[cog + setting_default(cog, "connection_pool_wait_ms") +.. ]]] + +Default: ``0`` + +.. [[[end]]] + +How long a query waits for a connection to be returned when :ref:`setting_max_open_connections` has been reached and every pooled connection is in use, before opening one more connection anyway. + +.. _setting_pool_read_connections: + +pool_read_connections +~~~~~~~~~~~~~~~~~~~~~ + +.. [[[cog + setting_default(cog, "pool_read_connections") +.. ]]] + +Default: ``on`` + +.. [[[end]]] + +Reuse read connections between queries. Turn this off to open, prepare and close a new connection for every read query. + .. _setting_allow_facet: allow_facet diff --git a/tests/test_api.py b/tests/test_api.py index 690cb080..808b8248 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -759,6 +759,10 @@ async def test_settings_json(ds_client): "suggest_facets": True, "default_cache_ttl": 5, "num_sql_threads": 1, + "max_open_connections": 128, + "connection_idle_timeout": 60, + "connection_pool_wait_ms": 0, + "pool_read_connections": True, "cache_size_kb": 0, "allow_csv_stream": True, "max_csv_mb": 100, diff --git a/tests/test_internals_database.py b/tests/test_internals_database.py index 9398afb9..619894cd 100644 --- a/tests/test_internals_database.py +++ b/tests/test_internals_database.py @@ -13,6 +13,7 @@ from opentelemetry import context as otel_context_api from datasette.app import Datasette from datasette.database import ( + ConnectionLeaseError, Database, DatasetteClosedError, ExecuteWriteResult, @@ -1312,8 +1313,15 @@ async def test_database_close_is_idempotent(tmpdir): async def test_close_releases_memory_connections(num_sql_threads, named): ds = Datasette(memory=True, settings={"num_sql_threads": num_sql_threads}) db = ds.add_memory_database(uuid.uuid4().hex) if named else ds.get_database() - read_connection = await db.execute_fn(lambda conn: conn) + leased = await db.execute_fn(lambda conn: conn) + # Read connections are leased: the proxy is unusable once the callback + # returns, so reach the underlying connection through a cursor + read_connection = await db.execute_fn( + lambda conn: conn.execute("select 1").connection + ) write_connection = await db.execute_write_fn(lambda conn: conn) + with pytest.raises(ConnectionLeaseError): + leased.execute("select 1") ds.close() for conn in (read_connection, write_connection): with pytest.raises(sqlite3.ProgrammingError, match="closed"): diff --git a/tests/test_read_pool.py b/tests/test_read_pool.py new file mode 100644 index 00000000..d1dcfd0e --- /dev/null +++ b/tests/test_read_pool.py @@ -0,0 +1,284 @@ +import asyncio +import sqlite3 +import threading +import time + +import pytest +import sqlite_utils + +from datasette.app import Datasette +from datasette.connection_pool import ConnectionLeaseError + + +def _make_dbs(tmp_path, n): + paths = [] + for i in range(n): + path = tmp_path / f"db{i}.db" + conn = sqlite3.connect(path) + conn.execute("create table t (id integer primary key, v text)") + conn.execute("insert into t (v) values ('x')") + conn.commit() + conn.close() + paths.append(str(path)) + return paths + + +def _file_conns(db): + return [c for c in db._all_connections] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", [0, 2]) +async def test_lease_expires_after_callback(tmp_path, num_sql_threads): + (path,) = _make_dbs(tmp_path, 1) + ds = Datasette([path], settings={"num_sql_threads": num_sql_threads}) + db = ds.get_database("db0") + stash = {} + + def callback(conn): + stash["conn"] = conn + assert isinstance(conn, sqlite3.Connection) + return conn.execute("select count(*) from t").fetchone()[0] + + assert await db.execute_fn(callback) == 1 + with pytest.raises(ConnectionLeaseError, match="db0"): + stash["conn"].execute("select 1") + with pytest.raises(ConnectionLeaseError): + stash["conn"].row_factory = None + # A second lease of the same underlying connection gets a new proxy, so + # the stale one still fails while another callback is using it + seen = {} + + def callback2(conn): + with pytest.raises(ConnectionLeaseError): + stash["conn"].execute("select 1") + seen["ok"] = True + + await db.execute_fn(callback2) + assert seen["ok"] + ds.close() + + +@pytest.mark.asyncio +async def test_returning_cursor_raises(tmp_path): + (path,) = _make_dbs(tmp_path, 1) + ds = Datasette([path]) + db = ds.get_database("db0") + with pytest.raises(ConnectionLeaseError, match="returned a Cursor"): + await db.execute_fn(lambda conn: conn.execute("select * from t")) + # Pool is still usable afterwards + assert (await db.execute("select count(*) from t")).single_value() == 1 + ds.close() + + +@pytest.mark.asyncio +async def test_leased_connection_compatibility(tmp_path): + (path,) = _make_dbs(tmp_path, 1) + ds = Datasette([path]) + db = ds.get_database("db0") + + def callback(conn): + out = {} + out["tables"] = sqlite_utils.Database(conn).table_names() + out["columns"] = list(sqlite_utils.Database(conn)["t"].columns_dict) + target = sqlite3.connect(":memory:") + conn.backup(target) + out["backup"] = target.execute("select count(*) from t").fetchone()[0] + conn.create_function("double", 1, lambda x: x * 2) + out["fn"] = conn.execute("select double(21)").fetchone()[0] + cursor = conn.cursor() + cursor.execute("select v from t") + out["cursor"] = cursor.fetchall()[0][0] + return out + + assert await db.execute_fn(callback) == { + "tables": ["t"], + "columns": ["id", "v"], + "backup": 1, + "fn": 42, + "cursor": "x", + } + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", [0, 3]) +async def test_global_cap_evicts_least_recently_used(tmp_path, num_sql_threads): + paths = _make_dbs(tmp_path, 6) + ds = Datasette( + paths, + settings={"num_sql_threads": num_sql_threads, "max_open_connections": 2}, + ) + pool = ds._read_pool + for i in range(6): + await ds.get_database(f"db{i}").execute("select 1") + assert pool.snapshot()["open"] <= 2 + # The two most recently used databases are the ones still open + open_dbs = { + name + for name in (f"db{i}" for i in range(6)) + if ds.get_database(name)._read_pool_state.open + } + assert open_dbs == {"db4", "db5"} + assert pool.stats["evicted_lru"] >= 4 + for i in range(6): + db = ds.get_database(f"db{i}") + assert len(_file_conns(db)) == db._read_pool_state.open + ds.close() + + +@pytest.mark.asyncio +async def test_prepare_connection_runs_for_every_pooled_connection(tmp_path): + paths = _make_dbs(tmp_path, 3) + ds = Datasette(paths, settings={"max_open_connections": 1}) + calls = [] + original = ds._prepare_connection + + def counting(conn, database): + calls.append(database) + return original(conn, database) + + ds._prepare_connection = counting + # Pool may have been created already with the original bound method + ds._read_pool._prepare_connection = counting + for _ in range(2): + for i in range(3): + assert ( + await ds.get_database(f"db{i}").execute("select v from t") + ).single_value() == "x" + # Cap of 1 means every switch of database reopens and re-prepares + assert calls == ["db0", "db1", "db2"] * 2 + ds.close() + + +def test_idle_reaper_closes_connections_without_event_loop(tmp_path): + paths = _make_dbs(tmp_path, 3) + ds = Datasette(paths, settings={"num_sql_threads": 0}) + pool = ds._read_pool + pool.idle_timeout = 0.2 + + async def query_all(): + for i in range(3): + await ds.get_database(f"db{i}").execute("select 1") + + asyncio.run(query_all()) + assert pool.snapshot()["open"] == 3 + # No event loop is running now - the reaper thread does the work + deadline = time.monotonic() + 5 + while pool.snapshot()["open"] and time.monotonic() < deadline: + time.sleep(0.05) + assert pool.snapshot()["open"] == 0 + assert pool.stats["expired"] == 3 + for i in range(3): + assert _file_conns(ds.get_database(f"db{i}")) == [] + ds.close() + + +@pytest.mark.asyncio +async def test_pooling_disabled_opens_connection_per_query(tmp_path): + (path,) = _make_dbs(tmp_path, 1) + ds = Datasette([path], settings={"pool_read_connections": False}) + db = ds.get_database("db0") + for _ in range(5): + await db.execute("select 1") + pool = ds._read_pool + assert pool.stats["opened"] == 5 + assert pool.snapshot()["open"] == 0 + assert _file_conns(db) == [] + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("wait_ms,expect_exceeded", [(0, True), (10000, False)]) +async def test_cap_exhausted_policy(tmp_path, wait_ms, expect_exceeded): + paths = _make_dbs(tmp_path, 3) + ds = Datasette( + paths, + settings={ + "num_sql_threads": 3, + "max_open_connections": 1, + "connection_pool_wait_ms": wait_ms, + }, + ) + barrier = threading.Barrier(3, timeout=0.5) + + def slow(conn): + try: + # With a hard cap only one callback can run at a time + barrier.wait() + except threading.BrokenBarrierError: + pass + time.sleep(0.05) + return conn.execute("select 1").fetchone()[0] + + results = await asyncio.gather( + *[ds.get_database(f"db{i}").execute_fn(slow) for i in range(3)] + ) + assert results == [1, 1, 1] + pool = ds._read_pool + if expect_exceeded: + assert pool.stats["exceeded_cap"] >= 1 + assert pool.stats["peak_open"] > 1 + else: + assert pool.stats["exceeded_cap"] == 0 + assert pool.stats["peak_open"] == 1 + # Back under the cap once everything has been returned + assert pool.snapshot()["open"] <= 1 + ds.close() + + +@pytest.mark.asyncio +async def test_memory_databases_not_counted(tmp_path): + ds = Datasette(memory=True, settings={"max_open_connections": 1}) + db = ds.add_memory_database("pool_mem_test") + await db.execute_write("create table t (id integer)") + await db.execute_write("insert into t values (1)") + await ds.get_database("_memory").execute("select 1") + assert (await db.execute("select count(*) from t")).single_value() == 1 + assert ds._read_pool.snapshot()["open"] == 0 + ds.close() + + +@pytest.mark.asyncio +async def test_database_close_with_leased_connection(tmp_path): + (path,) = _make_dbs(tmp_path, 1) + ds = Datasette([path], settings={"num_sql_threads": 0}) + db = ds.get_database("db0") + await db.execute("select 1") + + def closes_database(conn): + conn.execute("select 1") + db.close() + return "done" + + assert await db.execute_fn(closes_database) == "done" + assert _file_conns(db) == [] + assert ds._read_pool.snapshot()["open"] == 0 + ds.close() + + +@pytest.mark.asyncio +async def test_callback_that_closes_connection_is_discarded(tmp_path): + (path,) = _make_dbs(tmp_path, 1) + ds = Datasette([path]) + db = ds.get_database("db0") + await db.execute_fn(lambda conn: conn.close()) + assert ds._read_pool.stats["discarded_broken"] == 1 + assert (await db.execute("select count(*) from t")).single_value() == 1 + ds.close() + + +@pytest.mark.asyncio +async def test_open_transaction_rolled_back_on_release(tmp_path): + (path,) = _make_dbs(tmp_path, 1) + ds = Datasette([path], settings={"num_sql_threads": 1}) + db = ds.get_database("db0") + + def begin(conn): + conn.execute("begin") + conn.execute("select * from t").fetchall() + + await db.execute_fn(begin) + assert ds._read_pool.stats["rolled_back"] == 1 + assert not await db.execute_fn(lambda conn: conn.in_transaction) + ds.close()