diff --git a/datasette/app.py b/datasette/app.py index b2cb24b7..a332846b 100644 --- a/datasette/app.py +++ b/datasette/app.py @@ -8,6 +8,7 @@ if TYPE_CHECKING: from datasette.permissions import Resource from datasette.tokens import TokenRestrictions import collections +import copy import dataclasses import datetime import functools @@ -49,6 +50,11 @@ from .events import Event from .plugins import DEFAULT_PLUGINS, get_plugins, pm from .renderer import json_renderer from .resources import DatabaseResource, TableResource +from .schema_watcher import ( + FILES_DEFAULT, + SchemaWatcher, + store_closing_fingerprints, +) from .telemetry import ( TelemetryMiddleware, _in_datasette_client, @@ -108,7 +114,9 @@ from .utils.asgi import ( asgi_send_redirect, asgi_static, ) -from .utils.internal_db import init_internal_db, populate_schema_tables +from .utils.catalog import all_derived_table_dependencies +from .utils.inflight import InFlight, wait_for_concurrent +from .utils.internal_db import init_internal_db from .utils.sqlite import ( sqlite3, using_pysqlite3, @@ -230,6 +238,33 @@ SETTINGS = ( 3, "Number of threads in the thread pool for executing SQLite queries", ), + Setting( + "max_open_connections", + 32, + "Maximum number of pooled read connections open across all databases - 0 for no limit", + ), + Setting("max_write_connections", 8, "Maximum concurrent user database writers"), + Setting( + "max_pending_writes", 256, "Maximum queued user writes across all databases" + ), + Setting( + "write_queue_timeout_ms", 5000, "Maximum wait before a queued write starts" + ), + Setting( + "connection_idle_timeout_ms", + 30000, + "Close read connections, and stop write threads, that have been idle for this many milliseconds - 0 to keep them open", + ), + Setting( + "schema_watch_interval_ms", + 1000, + "How often to check external database files for schema changes - 0 to disable polling", + ), + Setting( + "default_schema_watch", + "external", + "Schema watch mode for database files and named in-memory databases without a mode of their own: external (poll for changes made outside Datasette) or owned (only Datasette changes them)", + ), 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" @@ -466,32 +501,38 @@ class Datasette: self._column_types = {} # .invoke_startup() will populate this self._setup_db_done = False self._suppress_background_tasks = False - try: - self._refresh_schemas_lock = asyncio.Lock() - self._startup_lock = asyncio.Lock() - except RuntimeError as rex: - # Workaround for intermittent test failure, see: - # https://github.com/simonw/datasette/issues/1802 - if "There is no current event loop in thread" in str(rex): - loop = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - self._refresh_schemas_lock = asyncio.Lock() - self._startup_lock = asyncio.Lock() - else: - raise + # A threading.Lock only ever acquired without blocking: one + # Datasette can be driven by several event loops, and an + # asyncio.Lock binds to whichever loop first waits on it + self._refresh_schemas_lock = threading.Lock() + # invoke_startup() in progress (datasette.utils.inflight.InFlight), + # claimed under this lock: concurrent callers on any event loop wait + # for it. Not an asyncio.Lock, which binds to one event loop + self._startup_running = None + self._startup_state_lock = threading.Lock() self._background_tasks = BackgroundTaskSupervisor(self) self.crossdb = crossdb self.nolock = nolock + self.internal_db_created = False + self._schema_watcher = SchemaWatcher(self) + # (catalog generation, derived-table dependency map) - see + # datasette.utils.catalog.all_derived_table_dependencies() + self._catalog_derived_cache = None + # SQLite version/extension details for /-/versions, computed once + self._sqlite_versions_info = None if memory or crossdb or not self.files: self.add_database( Database(self, is_mutable=False, is_memory=True), name="_memory" ) for file in self.files: + is_mutable = file not in self.immutables self.add_database( - Database(self, file, is_mutable=file not in self.immutables) + Database(self, file, is_mutable=is_mutable), + # Files named on the command line or in files= use the + # default_schema_watch setting (external unless configured) + schema_watch=FILES_DEFAULT if is_mutable else "immutable", ) - self.internal_db_created = False if internal is None: self._internal_database = Database(self, is_temp_disk=True) else: @@ -595,6 +636,13 @@ class Datasette: self.executor = futures.ThreadPoolExecutor( max_workers=self.setting("num_sql_threads") ) + for name in ( + "max_write_connections", + "max_pending_writes", + "write_queue_timeout_ms", + ): + if self.setting(name) <= 0: + raise StartupError(f"{name} must be greater than zero") self.max_returned_rows = self.setting("max_returned_rows") self.sql_time_limit_ms = self.setting("sql_time_limit_ms") self.page_size = self.setting("default_page_size") @@ -651,6 +699,8 @@ class Datasette: self.root_enabled = False self.default_deny = default_deny self.client = DatasetteClient(self) + # ds.config is available now: resolve per-database schema_watch modes + self._schema_watcher.configure() # Last, so metric callbacks never see a partially initialized instance register_datasette(self) @@ -713,66 +763,48 @@ class Datasette: return None async def refresh_schemas(self, *, force=False): - # Throttle schema refreshes to at most once per second + """Bring the _internal catalog up to date with every attached database. + + Uses the SchemaWatcher stat() prefilter, so it only opens connections + to databases whose files changed. Not called on the request path any + more - the watcher polls external databases in the background and + the write path catches schema changes made through Datasette. + """ + # Throttle non-forced refreshes to at most once per second if ( not force and time.monotonic() - getattr(self, "_last_schema_refresh", 0) < 1.0 ): return self._last_schema_refresh = time.monotonic() - if self._refresh_schemas_lock.locked(): - return - async with self._refresh_schemas_lock: + if not self._refresh_schemas_lock.acquire(blocking=False): + if not force: + return + # A refresh is already running, maybe on another event loop. + # A forced refresh must see changes made before it was called, + # so run another sweep now rather than wait for that one (sweeps + # are safe to overlap: catalog writes go through the internal + # database's write queue) await self._refresh_schemas() + return + try: + await self._refresh_schemas() + finally: + self._refresh_schemas_lock.release() - async def _refresh_schemas(self): + async def _refresh_schemas(self, *, background=False): internal_db = self.get_internal_database() if not self.internal_db_created: + if background: + return await init_internal_db(internal_db) await self.apply_metadata_json() self.internal_db_created = True - current_schema_versions = { - row["database_name"]: row["schema_version"] - for row in await internal_db.execute( - "select database_name, schema_version from catalog_databases" - ) - } - catalog_table_names = ( - "catalog_columns", - "catalog_foreign_keys", - "catalog_indexes", - "catalog_views", - "catalog_tables", - "catalog_databases", - ) - # Delete stale entries for databases that are no longer attached - catalog_database_names = set(current_schema_versions.keys()) - for table in catalog_table_names[:-1]: - catalog_database_names.update( - row["database_name"] - for row in await internal_db.execute( - f"select distinct database_name from {table}" - ) - if row["database_name"] is not None - ) - stale_databases = catalog_database_names - set(self.databases.keys()) - if stale_databases: - - def delete_stale_database_catalog(conn): - for stale_db_name in stale_databases: - for table in catalog_table_names: - conn.execute( - f"DELETE FROM {table} WHERE database_name = ?", - [stale_db_name], - ) - - await internal_db.execute_write_fn(delete_stale_database_catalog) - for database_name, db in self.databases.items(): - schema_version = (await db.execute("PRAGMA schema_version")).first()[0] - # Compare schema versions to see if we should skip it - if schema_version == current_schema_versions.get(database_name): - continue - await populate_schema_tables(internal_db, db, schema_version) + # Full catalog build (skips files whose persisted fingerprint + # shows they are unchanged when --internal is a real file) + await self._schema_watcher.initial_scan() + return + await self._schema_watcher.sweep(background=background) @property def urls(self): @@ -790,9 +822,41 @@ class Datasette: return pm async def invoke_startup(self): - # This must be called for Datasette to be in a usable state - if self._startup_invoked: + # This must be called for Datasette to be in a usable state. + # Several event loops (threads) may call it at once: one runs + # startup, the others wait for it to finish - and try again + # themselves if it failed or its event loop went away. + while not self._startup_invoked: + loop = asyncio.get_running_loop() + task = asyncio.current_task() + with self._startup_state_lock: + running = self._startup_running + if running is not None and (running.done() or running.abandoned()): + running.finish() + running = self._startup_running = None + if running is None: + running = self._startup_running = InFlight(loop, task) + mine = True + elif running.task is task: + # Called from inside startup itself (a startup hook + # making a request through datasette.client) + return + else: + mine = False + if not mine: + await wait_for_concurrent(running.future) + continue + try: + await self._invoke_startup() + finally: + with self._startup_state_lock: + if self._startup_running is running: + self._startup_running = None + running.finish() return + + async def _invoke_startup(self): + self._warn_connection_headroom() # Group spans created during startup under a single parent span with tracer.start_as_current_span(STARTUP): # Register event classes @@ -940,7 +1004,16 @@ class Datasette: name = next(iter(self.databases.keys())) return self.databases[name] - def add_database(self, db, name=None, route=None): + def add_database(self, db, name=None, route=None, schema_watch=None): + """Attach a database. + + ``schema_watch`` is how the _internal catalog is kept current for it: + ``"owned"`` (default - schema changes come through Datasette's write + methods), ``"external"`` (other processes may change the file, poll + it) or ``"immutable"`` (scan once). ``databases..schema_watch`` + in datasette.yaml overrides it. Databases that are not mutable are + always ``"immutable"``. + """ new_databases = self.databases.copy() if name is None: # Pick a unique name for this database @@ -957,18 +1030,42 @@ class Datasette: new_databases[name] = db # don't mutate! that causes race conditions with live import self.databases = new_databases + self._schema_watcher.register( + db, schema_watch or getattr(db, "schema_watch", None) + ) + self._crossdb_attachments_changed(db) return db + def _crossdb_attachments_changed(self, db): + # --crossdb: each pooled _memory connection ATTACHed the databases + # that existed when it was opened. Discard them so the next query + # sees databases added since, and none removed since + if not self.crossdb or db.is_memory: + return + memory = self.databases.get("_memory") + if memory is not None and memory is not db: + memory._invalidate_connections() + def add_memory_database(self, memory_name, name=None, route=None): return self.add_database( Database(self, memory_name=memory_name), name=name, route=route ) def remove_database(self, name): - self.get_database(name).close() + db = self._detach_database(name) + db.close() + + def _detach_database(self, name): + # remove_database() without the close() + db = self.get_database(name) new_databases = self.databases.copy() new_databases.pop(name) self.databases = new_databases + # Deletes this database's catalog rows (explicitly - there is no + # periodic stale-catalog scan any more) + self._schema_watcher.unregister(name) + self._crossdb_attachments_changed(db) + return db def close(self): """Release all resources held by this Datasette instance. @@ -980,17 +1077,45 @@ class Datasette: if self._closed: return self._closed = True + self._schema_watcher.stop() # Stop reporting metrics before closing databases unregister_datasette(self) first_exception = None - dbs = list(self.databases.values()) + [self._internal_database] - for db in dbs: + databases = list(self.databases.values()) + # Release every user database's idle or pinned writer before + # waiting for accepted work to obtain the global writer budget. + # The internal catalog remains available until user drains finish. + for db in databases: + db._begin_close() + for db in databases: try: db.close() except Exception as e: # noqa: BLE001 # Collect the first failure and re-raise after every close() has run if first_exception is None: first_exception = e + internal = self._internal_database + closing_fingerprints = [] + if self.internal_db_created and not internal.is_temp_disk: + # Owned databases are closed now: record their final + # fingerprints so a restart does not rescan them + try: + closing_fingerprints = self._schema_watcher.closing_fingerprints() + except Exception: + logger.exception("Could not fingerprint databases at close") + try: + internal.close() + except Exception as e: # noqa: BLE001 + if first_exception is None: + first_exception = e + if closing_fingerprints: + try: + store_closing_fingerprints(internal.path, closing_fingerprints) + except Exception: + logger.exception("Could not store fingerprints at close") + 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 +1128,70 @@ class Datasette: def setting(self, key): return self._settings.get(key, None) + @property + def _write_budget(self): + from .write_budget import WriteBudget + + budget = self.__dict__.get("_write_budget_instance") + if budget is None: + budget = WriteBudget( + self.setting("max_write_connections"), + self.setting("max_pending_writes"), + self.setting("write_queue_timeout_ms"), + ) + budget = self.__dict__.setdefault("_write_budget_instance", budget) + return budget + + @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._effective_max_open_connections(), + idle_timeout=self._connection_idle_timeout_s(), + ) + pool = self.__dict__.setdefault("_read_pool_instance", pool) + return pool + + def _warn_connection_headroom(self): + try: + import resource + + limit, _ = resource.getrlimit(resource.RLIMIT_NOFILE) + except (ImportError, OSError, ValueError): + return + reads = self._effective_max_open_connections() + writes = self.setting("max_write_connections") + attachments = min(10, len(self.databases)) if self.crossdb else 0 + estimate = 64 + 3 * (reads * (1 + attachments) + writes + 2) + if limit != resource.RLIM_INFINITY and (not reads or estimate > limit): + logging.getLogger(__name__).warning( + "Connection settings may exceed the file descriptor limit (%s): " + "read limit %s, writer limit %s, estimated descriptors including " + "WAL, attachments and headroom %s. Reduce connection limits or " + "increase the OS limit. This estimate excludes plugin allocations.", + limit, + reads, + writes, + estimate, + ) + + def _effective_max_open_connections(self): + """Configured read cap; do not silently increase it with thread count.""" + return max(self.setting("max_open_connections") or 0, 0) + + def _connection_idle_timeout_s(self): + "connection_idle_timeout_ms in seconds; 0 means never close idle ones." + timeout_ms = self.setting("connection_idle_timeout_ms") or 0 + return max(timeout_ms, 0) / 1000 + + @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} @@ -1549,7 +1738,7 @@ class Datasette: separator = "&" if "?" in url else "?" return url + separator + urllib.parse.urlencode({"_hash": hash_value}) - def _prepare_connection(self, conn, database): + def _prepare_connection(self, conn, database, *, crossdb=True): conn.row_factory = sqlite3.Row conn.text_factory = lambda x: str(x, "utf-8", "replace") if self.sqlite_extensions and database != INTERNAL_DB_NAME: @@ -1581,7 +1770,7 @@ class Datasette: if database != INTERNAL_DB_NAME: pm.hook.prepare_connection(conn=conn, database=database, datasette=self) # If self.crossdb and this is _memory, connect the first SQLITE_LIMIT_ATTACHED databases - if self.crossdb and database == "_memory": + if crossdb and self.crossdb and database == "_memory": count = 0 for db_name, db in self.databases.items(): if count >= SQLITE_LIMIT_ATTACHED or db.is_memory: @@ -1806,19 +1995,34 @@ class Datasette: parent, include_is_private, ): - databases = ( - [(parent, self.databases[parent])] - if parent in self.databases - else ([] if parent is not None else list(self.databases.items())) - ) - dependency_maps = dict( - zip( - (name for name, _ in databases), - await asyncio.gather( - *(db.derived_table_dependencies() for _, db in databases) - ), + if parent is not None: + dependency_maps = ( + {parent: await self.databases[parent].derived_table_dependencies()} + if parent in self.databases + else {} ) - ) + else: + # Every database: read the dependencies from the catalog (cached + # until it changes) rather than asking each database in turn - + # that was a query against every attached database for every + # page of allowed_resources(). The listing being filtered comes + # from the same catalog rows. Only databases the SchemaWatcher + # does not track fall back to live introspection. + dependency_maps = dict(await all_derived_table_dependencies(self)) + unwatched = [ + (name, db) + for name, db in self.databases.items() + if db._watch_state is None + ] + if unwatched: + dependency_maps.update( + zip( + (name for name, _ in unwatched), + await asyncio.gather( + *(db.derived_table_dependencies() for _, db in unwatched) + ), + ) + ) dependencies = [ (database_name, child, source) for database_name, dependency_map in dependency_maps.items() @@ -1968,7 +2172,29 @@ ORDER BY allowed.parent, allowed.child # Validate and cap limit limit = min(max(1, limit), 1000) + return await self._allowed_resources_page( + action, + actor, + parent=parent, + include_is_private=include_is_private, + include_reasons=include_reasons, + limit=limit, + next=next, + ) + async def _allowed_resources_page( + self, + action, + actor, + *, + parent, + include_is_private, + include_reasons, + limit, + next, + ): + """allowed_resources() without the limit cap: ``limit=None`` returns + every remaining resource in one query (PaginatedResources.all()).""" # Get base SQL query query, params = await self.allowed_resources_sql( action=action, @@ -1998,15 +2224,16 @@ ORDER BY allowed.parent, allowed.child # Add LIMIT (fetch limit+1 to detect if there are more results) # Note: query from allowed_resources_sql() already includes ORDER BY parent, child - query = f"{query} LIMIT :limit" - params["limit"] = limit + 1 + if limit is not None: + query = f"{query} LIMIT :limit" + params["limit"] = limit + 1 # Execute query result = await self.get_internal_database().execute(query, params) rows = list(result.rows) # Check if truncated (got more than limit rows) - truncated = len(rows) > limit + truncated = limit is not None and len(rows) > limit if truncated: rows = rows[:limit] # Remove the extra row @@ -2353,7 +2580,7 @@ ORDER BY allowed.parent, allowed.child url = "https://" + url[len("http://") :] return url - def _connected_databases(self): + def _connected_databases(self, names=None): return [ { "name": d.name, @@ -2365,23 +2592,58 @@ ORDER BY allowed.parent, allowed.child "hash": d.hash, } for name, d in self.databases.items() + if names is None or name in names ] async def _connected_databases_for_actor(self, actor): page = await self.allowed_resources("view-database", actor) allowed_names = {resource.parent async for resource in page.all()} - return [ - database - for database in self._connected_databases() - if database["name"] in allowed_names - ] + # Only stat (size) and hash (a full read of an immutable file, the + # first time) the databases this actor may see + return self._connected_databases(allowed_names) async def _databases_data(self, request): return {"databases": await self._connected_databases_for_actor(request.actor)} def _versions(self): + if self._sqlite_versions_info is None: + self._sqlite_versions_info = self._sqlite_versions() + sqlite_info, pysqlite3_version = self._sqlite_versions_info + datasette_version = {"version": __version__} + if self.version_note: + datasette_version["note"] = self.version_note + + try: + # Optional import to avoid breaking Pyodide + # https://github.com/simonw/datasette/issues/1733#issuecomment-1115268245 + import uvicorn + + uvicorn_version = uvicorn.__version__ + except ImportError: + uvicorn_version = None + info = { + "python": { + "version": ".".join(map(str, sys.version_info[:3])), + "full": sys.version, + }, + "datasette": datasette_version, + "asgi": "3.0", + "uvicorn": uvicorn_version, + "sqlite": copy.deepcopy(sqlite_info), + } + if pysqlite3_version is not None: + info["pysqlite3"] = pysqlite3_version + return info + + def _sqlite_versions(self): + """SQLite version, extensions, FTS versions and compile options, as + seen by a fully prepared connection. They cannot change while the + process runs, so /-/versions computes this once per instance rather + than preparing (loading extensions into) a connection per request. + crossdb=False: ATTACHing every --crossdb database is not needed to + report versions.""" conn = sqlite3.connect(":memory:") - self._prepare_connection(conn, "_memory") + self._prepare_connection(conn, "_memory", crossdb=False) sqlite_version = conn.execute("select sqlite_version()").fetchone()[0] sqlite_extensions = {"json1": detect_json1(conn)} for extension, testsql, hasversion in ( @@ -2415,44 +2677,24 @@ ORDER BY allowed.parent, allowed.child fts_versions.append(fts) except sqlite3.OperationalError: continue - datasette_version = {"version": __version__} - if self.version_note: - datasette_version["note"] = self.version_note - - try: - # Optional import to avoid breaking Pyodide - # https://github.com/simonw/datasette/issues/1733#issuecomment-1115268245 - import uvicorn - - uvicorn_version = uvicorn.__version__ - except ImportError: - uvicorn_version = None - info = { - "python": { - "version": ".".join(map(str, sys.version_info[:3])), - "full": sys.version, - }, - "datasette": datasette_version, - "asgi": "3.0", - "uvicorn": uvicorn_version, - "sqlite": { - "version": sqlite_version, - "fts_versions": fts_versions, - "extensions": sqlite_extensions, - "compile_options": [ - r[0] for r in conn.execute("pragma compile_options;").fetchall() - ], - }, + sqlite_info = { + "version": sqlite_version, + "fts_versions": fts_versions, + "extensions": sqlite_extensions, + "compile_options": [ + r[0] for r in conn.execute("pragma compile_options;").fetchall() + ], } + conn.close() + pysqlite3_version = None if using_pysqlite3: for package in ("pysqlite3", "pysqlite3-binary"): try: - info["pysqlite3"] = importlib.metadata.version(package) + pysqlite3_version = importlib.metadata.version(package) break except importlib.metadata.PackageNotFoundError: pass - conn.close() - return info + return sqlite_info, pysqlite3_version def _plugins(self, request=None, all=False): ps = list(get_plugins()) @@ -3068,17 +3310,13 @@ ORDER BY allowed.parent, allowed.child """ if self._startup_invoked and self._setup_db_done: return - async with self._startup_lock: - if self._startup_invoked and self._setup_db_done: - return - if not self._setup_db_done: - # First time server starts up, calculate table counts for - # immutable databases - for database in self.databases.values(): - if not database.is_mutable: - await database.table_counts(limit=60 * 60 * 1000) - self._setup_db_done = True - await self.invoke_startup() + # Immutable databases used to have their table counts computed + # here, opening every one of them at startup. They are now computed + # the first time a page needs them (see Database.table_counts()) and + # cached from then on. invoke_startup() makes concurrent callers - + # on any event loop - wait for the one running it. + self._setup_db_done = True + await self.invoke_startup() def add_background_task(self, func, name=None) -> BackgroundTask: """Register a piece of supervised background work, typically from @@ -3114,6 +3352,7 @@ ORDER BY allowed.parent, allowed.child """ await self.invoke_startup() await self._background_tasks.launch_all() + await self._schema_watcher.start() async def _launch_background_tasks(self): """Idempotently launch every registered background task. Private: @@ -3136,6 +3375,7 @@ ORDER BY allowed.parent, allowed.child if self._suppress_background_tasks: return await self._background_tasks.launch_all() + await self._schema_watcher.start() async def invoke_shutdown(self): """Run the graceful teardown sequence: plugin ``shutdown`` hooks, @@ -3151,6 +3391,7 @@ ORDER BY allowed.parent, allowed.child except Exception: logging.getLogger("datasette").exception("shutdown hook failed") await self._background_tasks.cancel_all(grace=5.0) + await self._schema_watcher.astop() self.close() def app(self): diff --git a/datasette/cli.py b/datasette/cli.py index e83de93a..bd3f14a4 100644 --- a/datasette/cli.py +++ b/datasette/cli.py @@ -896,9 +896,23 @@ pm.hook.register_commands(cli=cli) async def check_databases(ds): - # Run check_connection against every connected database - # to confirm they are all usable - for database in list(ds.databases.values()): + """Confirm every connected database is usable before serving. + + Building the _internal catalog already opens every database with a + prepared connection and reads every table's columns, which is what + check_connection() checks. So build it first, then re-check only: + + * databases whose catalog scan failed, to turn the failure into a + helpful error (SpatiaLite not loaded, not a database...) + * databases restored from a persisted catalog without being opened that + contain virtual tables, whose modules may no longer be loaded + + instead of opening every database a second time. + + """ + await ds._refresh_schemas() + to_check = await _databases_needing_check(ds) + for database in to_check: try: await database.execute_fn(check_connection) except SpatialiteConnectionProblem: @@ -926,3 +940,37 @@ async def check_databases(ds): ): msg = f"Warning: --crossdb only works with the first {SQLITE_LIMIT_ATTACHED} attached databases" click.echo(click.style(msg, bold=True, fg="yellow"), err=True) + + +async def _databases_needing_check(ds): + restored_with_virtual_tables = set() + databases = [(name, db) for name, db in list(ds.databases.items())] + restored = [ + name + for name, db in databases + if db._watch_state is not None + and db._watch_state.error is None + and db._watch_state.stats["scans"] == 0 + and db._watch_state.catalog_version is not None + ] + if restored: + rows = await ds.get_internal_database().execute( + """ + select distinct database_name from catalog_tables + where coalesce(rootpage, 0) = 0 + and database_name in (select value from json_each(:names)) + """, + {"names": json.dumps(restored)}, + ) + restored_with_virtual_tables = {row[0] for row in rows.rows} + to_check = [] + for name, db in databases: + state = db._watch_state + if ( + state is None + or state.error is not None + or name in restored_with_virtual_tables + or (state.catalog_version is None and not state.missing) + ): + to_check.append(db) + return to_check diff --git a/datasette/connection_pool.py b/datasette/connection_pool.py new file mode 100644 index 00000000..14235757 --- /dev/null +++ b/datasette/connection_pool.py @@ -0,0 +1,796 @@ +"""Leased read-connection pool, and the lease proxy every callback receives. + +Read connections are not cached per worker thread. 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 (the ``max_open_connections`` setting) 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_ms``. 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). Where threads + cannot be started at all (Pyodide) only the opportunistic reaping runs. +* Reusable connections are opened with ``Database.connect()`` and prepared + with ``Datasette._prepare_connection()`` exactly once, outside the pool + lock. Immutable isolated callbacks use fresh, unprepared connections + counted against the same cap and closed when the callback finishes. +* Each pooled connection remembers the database's ``_conn_generation`` when + it was opened. The SchemaWatcher bumps that generation when a file is + replaced or deleted and calls ``invalidate_database()``: idle connections + are closed at once, and leased ones are discarded by their own thread when + they are returned. A connection is never closed while another thread may + be stepping it (that segfaults). +* The cap is hard. When every connection is leased, admission waits for + release instead of exceeding it. Retiring connections remain counted + until they have actually closed. +* Callbacks receive a ``LeasedConnection`` proxy. Once the callback returns + the proxy is expired and any further use raises ``ConnectionLeaseError``. + Write callbacks get the same proxy (see ``Database._execute_writes``). +* 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. The + exception is ``_memory`` under ``--crossdb``: each of its connections + ATTACHes up to ten database files, so those are counted and reaped like + file connections. +""" + +import collections +import inspect +import threading +import time +import weakref + +from .utils import sqlite3 +from .write_budget import DatabaseResourceError, callback_scope + + +class ConnectionLeaseError(RuntimeError): + """A connection was used outside the callback it was lent to.""" + + +_CALLBACK_NAMES = { + "read": "execute_fn()", + "write": "execute_write_fn()", + "isolated": "execute_isolated_fn()", +} + + +class LeasedConnection: + """Forwarding proxy handed to read, write and isolated 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. + + 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__ = ("_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 + if conn is None: + raise ConnectionLeaseError( + 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 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() + object.__setattr__(self, "_conn", None) + + def _wrap(self, result): + 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 + @property + def __class__(self): + return sqlite3.Connection + + def __getattr__(self, 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) + + def __delattr__(self, name): + delattr(self._live(), name) + + # Dunder methods are looked up on the type, so forward them explicitly + def __enter__(self): + self._live().__enter__() + # "with conn as c" must not hand out the raw connection + return self + + 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._wrap(self._live().execute(*args, **kwargs)) + + def cursor(self, *args, **kwargs): + return self._wrap(self._live().cursor(*args, **kwargs)) + + def close(self): + conn = self._live() + if self._kind == "write": + # The write connection belongs to the database's write thread, + # which keeps using it for every later write + raise ConnectionLeaseError( + f"Cannot close the write connection to database " + f"{self._db_name!r} from inside an execute_write_fn() " + "callback: Datasette owns it and closes it when it is idle" + ) + # Closing a pooled read connection is allowed; the pool notices on + # release and discards it instead of returning it to the idle stack. + # An isolated connection is closed afterwards anyway. + conn.close() + + def __repr__(self): + state = "expired" if self._conn is None else "active" + return f"" + + +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() + + +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, LeasedBlob, LeasedIterator)): + return True + if isinstance(value, dict): + 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_database_resource(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 + the locals of the frames it was raised through. + + Call it on the thread that owns the connection, before the connection + can be used by anyone else. A cursor keeps a reference to the + connection's cached statement for its SQL; on Python before 3.12, + deallocating the cursor resets that statement. If the cursor dies on + another thread (the caller's event loop) while the connection is running + the same SQL again - the write thread's next execute_write_many(), or + another worker that has leased the connection - the statement is reset + mid-step and that query fails with "bad parameter or other API misuse". + A closed cursor keeps its rowcount, lastrowid and description. + """ + if isinstance(value, BaseException): + tb = value.__traceback__ + while tb is not None: + try: + values = list(tb.tb_frame.f_locals.values()) + except Exception: # noqa: BLE001 + values = [] + for item in values: + if isinstance(item, sqlite3.Cursor): + _close_cursor(item) + tb = tb.tb_next + return + if isinstance(value, sqlite3.Cursor): + _close_cursor(value) + elif isinstance(value, (tuple, list)): + for item in value: + if isinstance(item, sqlite3.Cursor): + _close_cursor(item) + elif inspect.isgenerator(value): + try: + values = list(inspect.getgeneratorlocals(value).values()) + except Exception: # noqa: BLE001 + values = [] + for item in values: + if isinstance(item, sqlite3.Cursor): + _close_cursor(item) + value.close() + + +def _close_cursor(cursor): + try: + cursor.close() + except Exception: # noqa: BLE001, S110 + # Its connection is closed already + pass + + +def _generator_uses_connection(result): + """True for a generator that would step the leased connection after + the callback returned: one holding the connection (a closure over + ``conn``) or a live cursor (``(r for r in conn.execute(...))``). A + generator over rows already fetched inside the callback is fine.""" + if not inspect.isgenerator(result): + return False + try: + values = inspect.getgeneratorlocals(result).values() + except Exception: # noqa: BLE001 + return False + return any( + isinstance(v, (sqlite3.Connection, sqlite3.Cursor, LeasedBlob, LeasedIterator)) + for v in values + ) + + +class _Entry: + __slots__ = ("conn", "generation", "last_used", "state") + + def __init__(self, conn, state, generation): + self.conn = conn + self.state = state + self.generation = generation + 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 + # (closing every connection to a named one would destroy it) - + # except _memory under --crossdb, whose connections each ATTACH up + # to ten database files and are private, so closing one loses + # nothing + self.counted = not db.is_memory or ( + db.name == "_memory" + and not db.memory_name + and bool(getattr(db.ds, "crossdb", False)) + ) + + +class ReadConnectionPool: + def __init__(self, prepare_connection, max_open=32, idle_timeout=30.0): + self._prepare_connection = prepare_connection + self.max_open = max_open # 0 = no limit + self.idle_timeout = idle_timeout # seconds, 0 = never + self._lock = threading.Condition() + # Idle, counted entries, least recently used first + self._lru = collections.OrderedDict() + self._open = 0 # counted open connections, idle + leased + self._closed = False + self.stats = collections.Counter() + + # -- public API --------------------------------------------------------- + + def run(self, db, fn, isolated=False): + """Call fn(conn) with a pooled connection leased for the duration.""" + entry = self.acquire(db, fresh=isolated) + lease = LeasedConnection( + entry.conn, db.name, "isolated" if isolated else "read" + ) + rejected = False + discard = isolated + try: + with callback_scope(): + result = fn(lease) + if _contains_database_resource(result) or _generator_uses_connection( + result + ): + rejected = True + close_cursors(result) + except BaseException as e: + close_cursors(e) + discard = True + raise + finally: + lease._expire() + self.release(entry, discard=discard) + if rejected: + raise ConnectionLeaseError( + 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" + ) + return result + + def acquire(self, db, fresh=False): + state = self._state(db) + deadline = time.monotonic() + 5 + while True: + to_close = [] + entry = None + reserved = False + with self._lock: + if self._closed or state.closed or db._closed: + from .database import DatasetteClosedError + + raise DatasetteClosedError(f"Database {db.name!r} has been closed") + self._collect_expired_locked(time.monotonic(), to_close) + if to_close: + pass # Close before spending the capacity it releases. + elif state.idle and not fresh: + entry = state.idle.pop() + if state.counted: + del self._lru[entry] + if entry.generation != db._conn_generation: + to_close.append(entry) + self.stats["discarded_stale"] += 1 + entry = None + else: + state.leased += 1 + self.stats["reused"] += 1 + elif state.counted and self.max_open and self._open >= self.max_open: + if self._lru: + victim, _ = self._lru.popitem(last=False) + victim.state.idle.remove(victim) + to_close.append(victim) + self.stats["evicted_lru"] += 1 + else: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise DatabaseResourceError( + "Timed out waiting for a read connection" + ) + self._lock.wait(min(remaining, 0.1)) + else: + generation = db._conn_generation + 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 + ) + reserved = True + if to_close: + self._close_entries(to_close) + if entry is not None: + return entry + if reserved: + break + conn = None + try: + conn = db.connect() + if not fresh: + with callback_scope(): + self._prepare_connection(conn, db.name) + except BaseException: + if conn is not None: + self._close_conn(state, conn) + with self._lock: + if conn is None: + state.open -= 1 + if state.counted: + self._open -= 1 + state.leased -= 1 + self.stats["opened"] -= 1 + self.stats["open_failed"] += 1 + self._lock.notify_all() + raise + return _Entry(conn, state, generation) + + def evict_idle(self): + """Reclaim idle file connections after external descriptor pressure.""" + with self._lock: + entries = list(self._lru) + self._lru.clear() + for entry in entries: + entry.state.idle.remove(entry) + self._close_entries(entries) + return len(entries) + + def release(self, entry, discard=False): + state = entry.state + conn = entry.conn + rolled_back = False + try: + # Hygiene: a callback that began a transaction and did not end it + # must not hand that transaction to the next borrower + 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 + stale = entry.generation != state.db._conn_generation + over_cap = state.counted and self.max_open and self._open > self.max_open + if discard or stale or over_cap or self._closed or state.closed: + discard = True + if over_cap: + self.stats["closed_over_cap"] += 1 + if stale: + self.stats["discarded_stale"] += 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 + self._lock.notify_all() + 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 invalidate_database(self, db): + """db's file was replaced or deleted and its ``_conn_generation`` has + been bumped: close its idle connections now. Leased connections are + closed by the thread using them, when they are released.""" + state = getattr(db, "_read_pool_state", None) + if state is None: + return + with self._lock: + stale = [e for e in state.idle if e.generation != db._conn_generation] + for entry in stale: + state.idle.remove(entry) + if state.counted: + self._lru.pop(entry, None) + self.stats["discarded_stale"] += 1 + self._close_entries(stale) + + 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._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) + 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) + 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: + state.open -= 1 + if state.counted: + self._open -= 1 + self.stats["closed"] += 1 + self._lock.notify_all() + 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 + # Set when this platform cannot start threads (Pyodide): pools then + # rely on reaping during checkout only + self.unavailable = 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 self.unavailable or (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(): + thread = threading.Thread( + target=self._run, name="datasette-read-pool-reaper", daemon=True + ) + try: + thread.start() + except RuntimeError: + # "can't start new thread": no threads on this platform + self.unavailable = True + return + self._thread = thread + 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) + # In a helper so that no reference to a pool (and through its + # _prepare_connection, the whole Datasette) survives in this + # long-lived frame while it waits + next_wake = self._reap(pools) + 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) + + @staticmethod + def _reap(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 + return next_wake + + +_reaper = _Reaper() diff --git a/datasette/database.py b/datasette/database.py index 542b3012..3096ac87 100644 --- a/datasette/database.py +++ b/datasette/database.py @@ -1,7 +1,8 @@ import asyncio -import atexit +import contextlib import contextvars import inspect +import logging import os import queue import sys @@ -9,6 +10,7 @@ import tempfile import threading import time import uuid +import weakref from collections import namedtuple from pathlib import Path @@ -16,6 +18,11 @@ import sqlite_utils from opentelemetry import context as otel_context_api from opentelemetry.trace import Status, StatusCode +from .connection_pool import ( + ConnectionLeaseError, + LeasedConnection, + close_cursors, +) from .inspect import inspect_hash from .telemetry import ( callback_name, @@ -51,6 +58,7 @@ from .telemetry_registry import ( ) from .tracer import trace from .utils import ( + await_me_maybe, call_with_supported_arguments, detect_fts, detect_primary_keys, @@ -64,12 +72,21 @@ from .utils import ( table_column_details, table_columns, ) +from .utils.catalog import SPATIALITE_HIDDEN_TABLES from .utils.sql_analysis import SQLAnalysis, analyze_sql_tables from .utils.sqlite import sqlite_derived_table_dependencies, sqlite_hidden_table_names +from .write_budget import ( + RETIRE_WRITER, + DatabaseResourceError, + callback_scope, + check_reentrancy, + connect_with_retry, +) -connections = threading.local() +logger = logging.getLogger(__name__) EXECUTE_WRITE_RETURNING_LIMIT = 10 +WRITE_SHUTDOWN_TIMEOUT = 10 AttachedDatabase = namedtuple("AttachedDatabase", ("seq", "name", "file")) @@ -79,12 +96,17 @@ class DatasetteClosedError(RuntimeError): _SHUTDOWN = object() +_write_connection_owner = contextvars.ContextVar("write_connection_owner", default=None) class Database: # For table counts stop at this many rows: count_limit = 10000 - _thread_local_id_counter = 1 + # Counts for immutable databases are computed once, the first time a + # page needs them, with this per-table time limit, then cached (they + # cannot change). They used to be computed for every immutable database + # at startup. + immutable_count_time_limit_ms = 60 * 60 * 1000 def __init__( self, @@ -97,8 +119,6 @@ class Database: is_temp_disk=False, ): self.name = None - self._thread_local_id = f"x{self._thread_local_id_counter}" - Database._thread_local_id_counter += 1 self.route = None self.ds = ds self.path = path @@ -115,7 +135,13 @@ class Database: self.is_mutable = True self.mode = "rwc" self._wal_enabled = False - atexit.register(self._cleanup_temp_file) + # Unlinked by close(), when this Database is garbage collected, + # or at exit. The finalizer holds only the path: an + # atexit.register() of a bound method would keep this Database - + # and through it the whole Datasette - alive until exit + self._temp_file_finalizer = weakref.finalize( + self, _unlink_temp_files, temp_path + ) else: self._wal_enabled = False self.cached_hash = None @@ -124,18 +150,59 @@ class Database: self._cached_derived_table_dependencies = None self._write_thread = None self._write_queue = None + # Guards the write thread lifecycle. A task is only ever enqueued, + # and a write thread only ever decides to exit, while holding this + # lock - see _send_to_write_thread() and _execute_writes(). + self._write_thread_lock = threading.Lock() + # Write threads that have deregistered (idle exit) but may still be + # closing their connection. close() joins them. + self._retiring_write_threads = set() + # Connections currently held by write threads (see _WriteConnection) + self._write_thread_connections = set() + # Number of write threads started over the lifetime of this Database, + # and how often a thread's idle timeout fired but it found a newly + # queued task once it held the lock (so it kept running) + self._write_threads_started = 0 + self._write_thread_exits_averted = 0 self._closed = False + self._close_lock = threading.RLock() + self._close_complete = False + self._closing_execute_futures = () + self._cleanup_temp_file_on_writer_exit = False self._pending_execute_futures = set() self._pending_execute_futures_lock = threading.Lock() - # These are used when in non-threaded mode: - self._read_connection = None + # The write connection in non-threaded mode (reads go through + # ds._read_pool in both modes) self._write_connection = None + self._write_connection_generation = 0 + # 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 = [] + # Set by SchemaWatcher.register(); None for the internal database + self._watch_state = None + # Bumped when the file behind this database is replaced or deleted, + # so pooled read connections and the write connection still pointing + # at the old inode are discarded and reopened + self._conn_generation = 0 + # True once Datasette has seen this database's file exist (a + # connection opened, or the SchemaWatcher stat()ed it). From then on + # write connections never create the file - see connect() + self._file_seen = False + # Short-lived connections that are not in _all_connections (the + # SchemaWatcher's scans and pragma checks): counted, so that deleting + # the file can wait until none is open - see + # _untracked_connection_opened() + self._untracked_open = 0 + self._untracked_cond = threading.Condition() if not is_temp_disk: self.mode = mode + @property + def _budgeted_writer(self): + return self is not getattr(self.ds, "_internal_database", None) + def _check_not_closed(self): if self._closed: raise DatasetteClosedError(f"Database {self.name!r} has been closed") @@ -144,6 +211,133 @@ class Database: with self._pending_execute_futures_lock: self._pending_execute_futures.discard(future) + def _invalidate_connections(self): + """The file behind this database was replaced or deleted. + + Bumps the connection generation: pooled read connections are closed + (idle ones now, leased ones when their callback returns) and the + write connection is reopened before the next write task. Nothing is + closed from this thread while another thread may be using it. + """ + with self._write_thread_lock: + self._conn_generation += 1 + read_pool = getattr(self.ds, "_read_pool_or_none", None) + if read_pool is not None: + read_pool.invalidate_database(self) + + def _forget_connection(self, conn): + try: + conn.close() + except Exception: # noqa: BLE001, S110 + pass + try: + self._all_connections.remove(conn) + except ValueError: + pass + + def _connect_write_thread(self): + # Publish the connection only after marking its owning writer. A + # timed-out close() must never see a half-registered write connection + # and close it before the callback can use it. + token = _write_connection_owner.set(self) + try: + return self.connect(write=True) + finally: + _write_connection_owner.reset(token) + + def _track_connection(self, conn): + if _write_connection_owner.get() is self: + self._write_thread_connections.add(conn) + self._all_connections.append(conn) + + def _untracked_connection_opened(self): + """Call before opening a connection with ``connect(track=False)``. + + Raises DatasetteClosedError once close() has started, so after + close() no new untracked connection can be opened and + _wait_for_untracked_connections() only has to wait for the ones + already counted.""" + with self._untracked_cond: + if self._closed: + raise DatasetteClosedError(f"Database {self.name!r} has been closed") + self._untracked_open += 1 + + def _untracked_connection_closed(self): + with self._untracked_cond: + self._untracked_open -= 1 + if self._untracked_open <= 0: + self._untracked_cond.notify_all() + + def _wait_for_untracked_connections(self, timeout=None): + """Wait until every untracked connection has been closed. Only + meaningful after close(). Returns False on timeout.""" + with self._untracked_cond: + return self._untracked_cond.wait_for( + lambda: self._untracked_open <= 0, timeout=timeout + ) + + def _schema_check(self, conn, loop=None): + # Runs on the thread that owns conn, right after a write task + state = self._watch_state + if state is not None: + self.ds._schema_watcher.check_after_write(state, conn, loop) + + async def _execute_on_write_connection(self, fn): + """Run fn(conn) on the write connection, serialized with writes, + without write_wrapper hooks, events or a transaction. Used by the + SchemaWatcher for shared-cache memory databases, where a second + connection reading sqlite_master would make concurrent DDL fail + with "database table is locked" (shared-cache locks do not wait).""" + self._check_not_closed() + if self.ds.executor is None: + with ( + self.ds._write_budget.inline(self) + if self._budgeted_writer + else contextlib.nullcontext() + ): + return fn(self._non_threaded_write_connection()) + return await self._send_to_write_thread(fn, transaction=False) + + def _non_threaded_write_connection(self): + """num_sql_threads=0: the single write connection, used inline on + the event loop. Reopened if the file was replaced or deleted. Unlike + the threaded write connection it is not closed when idle - there is + no thread to time it out (Pyodide has none at all).""" + conn = self._write_connection + if ( + conn is not None + and self._write_connection_generation != self._conn_generation + ): + self._write_connection = None + self._forget_connection(conn) + conn = None + if conn is None: + generation = self._conn_generation + conn = self.connect(write=True) + try: + _pin_write_connection_pragmas(conn) + with callback_scope(): + self.ds._prepare_connection(conn, self.name) + except BaseException: + # Never cache a half-prepared connection + self._forget_connection(conn) + raise + self._write_connection = conn + self._write_connection_generation = generation + return conn + + async def _after_write(self): + # Read-your-writes for the catalog: if that write changed the schema, + # wait for the catalog rows for this database to be rebuilt + state = self._watch_state + if state is not None and (state.needs_scan or state.scan_future is not None): + try: + await self.ds._schema_watcher.after_write(state) + except DatabaseResourceError as error: + raise RuntimeError( + "Write completed, but catalog refresh failed; do not retry the write automatically" + ) from error + @property def cached_table_counts(self): if self._cached_table_counts is not None: @@ -172,7 +366,11 @@ class Database: else: return "db" - def connect(self, write=False): + def connect(self, write=False, track=True): + """Open a new connection. ``track=False`` leaves it out of + ``_all_connections``: the caller closes it itself and close() must + never close it from another thread while it is in use (the + SchemaWatcher's short-lived scan connections).""" extra_kwargs = {} if write: extra_kwargs["isolation_level"] = "IMMEDIATE" @@ -183,11 +381,13 @@ class Database: ) if not write: conn.execute("PRAGMA query_only=1") - self._all_connections.append(conn) + if track: + self._track_connection(conn) return conn if self.is_memory: conn = sqlite3.connect(":memory:", uri=True, check_same_thread=False) - self._all_connections.append(conn) + if track: + self._track_connection(conn) return conn # mode=ro or immutable=1? @@ -199,13 +399,26 @@ class Database: qs = "?immutable=1" assert not (write and not self.is_mutable) if write: - qs = "" + # The first write connection may create the file (a database + # added with add_database() before its file exists). Once + # Datasette has seen the file exist, writes use mode=rw: if it + # is deleted from under us later, writes fail instead of + # silently creating a new empty database in its place. + qs = "?mode=rw" if self._file_seen else "" if self.mode is not None: qs = f"?mode={self.mode}" - conn = sqlite3.connect( - f"file:{self.path}{qs}", uri=True, check_same_thread=False, **extra_kwargs + conn = connect_with_retry( + self, + lambda: sqlite3.connect( + f"file:{self.path}{qs}", + uri=True, + check_same_thread=False, + **extra_kwargs, + ), ) - self._all_connections.append(conn) + self._file_seen = True + if track: + self._track_connection(conn) if self.is_temp_disk and not self._wal_enabled: conn.execute("PRAGMA journal_mode=WAL") self._wal_enabled = True @@ -215,52 +428,78 @@ class Database: """Release all resources held by this database. Idempotent. After close() further calls to execute()/execute_fn()/ - execute_write()/execute_write_fn() raise DatasetteClosedError. + execute_write()/execute_write_fn() raise DatasetteClosedError. Accepted + writes drain through normal admission, with their existing deadlines. + Budgeted writes still awaiting admission after ten seconds are + cancelled; running writers retain their connections until finished. """ - if self._closed: - return - with self._pending_execute_futures_lock: + with self._close_lock: + if self._close_complete: + return + self._begin_close() + self._finish_close() + self._close_complete = True + + def _begin_close(self): + # Datasette begins closing every user database before waiting for + # any one of them, releasing pinned writers needed by queued work. + with self._close_lock: if self._closed: return - self._closed = True - pending_execute_futures = tuple(self._pending_execute_futures) - # Shut down the write thread, if any, via a sentinel. The thread - # drains any writes already queued before the sentinel and then - # closes its own write connection and returns. - write_thread = self._write_thread - if write_thread is not None and self._write_queue is not None: - self._write_queue.put(_SHUTDOWN) - write_thread.join(timeout=10) - if write_thread.is_alive(): - sys.stderr.write( - f"Datasette: write thread for {self.name!r} did not exit within 10s\n" - ) - sys.stderr.flush() - for future in pending_execute_futures: + with self._pending_execute_futures_lock, self._untracked_cond: + self._closed = True + self._closing_execute_futures = tuple(self._pending_execute_futures) + with self._write_thread_lock: + # Include writers awaiting admission: the scheduler must + # drain accepted tasks before consuming this sentinel. + if self._write_queue is not None: + self._write_queue.put(_SHUTDOWN) + + def _finish_close(self): + budget = self.ds.__dict__.get("_write_budget_instance") + still_running = False + if self._budgeted_writer and budget is not None: + still_running = not budget.drain_database(self, WRITE_SHUTDOWN_TIMEOUT) + with self._write_thread_lock: + write_thread = self._write_thread + retiring_threads = tuple(self._retiring_write_threads) + if not self._budgeted_writer: + deadline = time.monotonic() + WRITE_SHUTDOWN_TIMEOUT + for thread in ((write_thread,) if write_thread else ()) + retiring_threads: + thread.join(timeout=max(0, deadline - time.monotonic())) + still_running = still_running or thread.is_alive() + if still_running: + sys.stderr.write( + f"Datasette: write thread for {self.name!r} did not exit within " + f"{WRITE_SHUTDOWN_TIMEOUT:g}s\n" + ) + sys.stderr.flush() + # A write thread that is still running owns its connection and closes + # it when it finishes; closing it here could free it mid-query + # (membership tests only: copying the set could race the thread) + in_use = self._write_thread_connections if still_running else () + for future in self._closing_execute_futures: try: future.result() except Exception: # noqa: BLE001, S110 # Shutdown teardown - a failed pending write must not block close() pass + self._closing_execute_futures = () + # 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): + if connection in in_use: + continue try: connection.close() except Exception: # noqa: BLE001, S110 pass self._all_connections = [] - # Drop per-thread cached read connections we can reach - try: - delattr(connections, self._thread_local_id) - except AttributeError: - pass - # Close non-threaded-mode cached connections if still open - if self._read_connection is not None: - try: - self._read_connection.close() - except Exception: # noqa: BLE001, S110 - pass - self._read_connection = None + # Close the non-threaded-mode write connection if still open if self._write_connection is not None: try: self._write_connection.close() @@ -268,15 +507,17 @@ class Database: pass self._write_connection = None if self.is_temp_disk: - self._cleanup_temp_file() + with self._write_thread_lock: + if ( + self._write_thread is not None and self._write_thread.is_alive() + ) or self._retiring_write_threads: + self._cleanup_temp_file_on_writer_exit = True + else: + self._cleanup_temp_file() def _cleanup_temp_file(self): if self.is_temp_disk and self.path: - for suffix in ("", "-wal", "-shm"): - try: - os.unlink(self.path + suffix) - except OSError: - pass + self._temp_file_finalizer() async def execute_write( self, @@ -383,22 +624,23 @@ class Database: return results async def execute_isolated_fn(self, fn): + check_reentrancy() self._check_not_closed() # Open a new connection just for the duration of this function, # blocking the write queue to avoid any writes occurring during it write = self.is_mutable def _run(): + # Always closed here, by the thread using it. close() waits for + # this call before closing tracked connections (the threaded + # immutable case runs as a tracked executor future) isolated_connection = self.connect(write=write) try: - return fn(isolated_connection) + return _call_with_lease(fn, isolated_connection, self.name, "isolated") finally: - isolated_connection.close() - try: - self._all_connections.remove(isolated_connection) - except ValueError: - # May already have been cleared by close(). - pass + if write: + self._schema_check(isolated_connection) + self._forget_connection(isolated_connection) with tracer.start_as_current_span(DB_QUERY, kind=DB_QUERY.kind) as span: span.set_attribute(DB_SYSTEM, "sqlite") @@ -408,17 +650,42 @@ class Database: with record_operation_duration(self.name, "write" if write else "read"): if self.ds.executor is None: # non-threaded mode - return _run() + if write: + with ( + self.ds._write_budget.inline(self) + if self._budgeted_writer + else contextlib.nullcontext() + ): + if ( + self._write_connection is not None + and not self.is_memory + ): + old, self._write_connection = ( + self._write_connection, + 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) + await self._after_write() + return result if not write: # Immutable database - no writes can ever occur, so there # is no write queue to block; run against a fresh - # read-only connection - ctx = contextvars.copy_context() - return await asyncio.get_running_loop().run_in_executor( - self.ds.executor, ctx.run, _run + # read-only connection on the SQL thread pool, tracked so + # close() waits for it + return await self._run_in_executor( + lambda: self.ds._read_pool.run(self, fn, isolated=True) ) # Threaded mode - send to write thread - return await self._send_to_write_thread(fn, isolated_connection=True) + result = await self._send_to_write_thread(fn, isolated_connection=True) + await self._after_write() + return result async def analyze_sql(self, sql, params=None) -> SQLAnalysis: self._check_not_closed() @@ -447,6 +714,7 @@ class Database: ) async def _execute_write_fn(self, fn, block=True, transaction=True, request=None): + check_reentrancy() self._check_not_closed() pending_events = [] @@ -456,15 +724,21 @@ class Database: fn = self._wrap_fn_with_hooks(fn, request, transaction, track_event) if self.ds.executor is None: # non-threaded mode - if self._write_connection is None: - self._write_connection = self.connect(write=True) - self.ds._prepare_connection(self._write_connection, self.name) - if transaction: - with self._write_connection: - self._write_connection.execute("BEGIN IMMEDIATE") - result = fn(self._write_connection) - else: - result = fn(self._write_connection) + with ( + self.ds._write_budget.inline(self) + if self._budgeted_writer + else contextlib.nullcontext() + ): + conn = self._non_threaded_write_connection() + try: + if transaction: + with conn: + conn.execute("BEGIN IMMEDIATE") + result = _call_with_lease(fn, conn, self.name, "write") + else: + result = _call_without_transaction(fn, conn, self.name) + finally: + self._schema_check(conn) if not block: # There is no write thread here, so the write has already # finished. Hand back the same (task_id, reply_future) shape @@ -479,8 +753,14 @@ class Database: fn, block=block, transaction=transaction ) if block: - for event in pending_events: - await self.ds.track_event(event) + await self._after_write() + try: + for event in pending_events: + await self.ds.track_event(event) + except DatabaseResourceError as error: + raise RuntimeError( + "Write completed, but event dispatch failed; do not retry the write automatically" + ) from error else: # For non-blocking writes, spawn a background task to # dispatch events after the write thread completes @@ -489,10 +769,19 @@ class Database: async def _dispatch_events_after_write(): try: await reply_future - except Exception: # noqa: BLE001 - # The write failed; skip success events regardless of why - # if the write failed, don't emit success events + except Exception as error: + logger.exception( + "Nonblocking write %s failed", + task_id, + extra={ + "database": self.name, + "task_id": str(task_id), + "error_code": getattr(error, "code", None), + }, + ) + await self._notify_write_completed(task_id, error) return + await self._notify_write_completed(task_id, None) for event in pending_events: await self.ds.track_event(event) @@ -500,6 +789,20 @@ class Database: result = task_id return result + async def _notify_write_completed(self, task_id, exception): + from .plugins import pm + + try: + for result in pm.hook.write_task_completed( + datasette=self.ds, + database=self.name, + task_id=task_id, + exception=exception, + ): + await await_me_maybe(result) + except Exception: + logger.exception("write_task_completed hook failed for %s", task_id) + def _wrap_fn_with_hooks(self, fn, request, transaction, track_event): from .plugins import pm @@ -537,20 +840,22 @@ class Database: async def _send_to_write_thread( self, fn, block=True, isolated_connection=False, transaction=True ): - if self._write_queue is None: - self._write_queue = queue.Queue() - if self._write_thread is None: - self._write_thread = threading.Thread( - target=self._execute_writes, daemon=True - ) - self._write_thread.name = f"_execute_writes for database {self.name}" - self._write_thread.start() + check_reentrancy() task_id = uuid.uuid4() loop = asyncio.get_running_loop() reply_future = loop.create_future() - # Capture the OpenTelemetry context and enqueue time for the write thread - self._write_queue.put( - WriteTask( + with self._write_thread_lock: + # Enqueueing and checking for a live write thread happen under + # one lock. A write thread only decides to exit while holding + # this lock and after seeing an empty queue, so a task can never + # be left in the queue with no thread to run it: either the + # thread sees this task and keeps going, or it has already + # deregistered and we start a new one. + if self._closed: + raise DatasetteClosedError(f"Database {self.name!r} has been closed") + if self._write_queue is None: + self._write_queue = queue.Queue() + task = WriteTask( fn, task_id, loop, @@ -561,35 +866,137 @@ class Database: time.time_ns(), block, ) - ) + if self._budgeted_writer: + self.ds._write_budget.submit(self, task) + else: + if self._write_thread is None: + self._start_write_thread() + self._write_queue.put(task) if block: return await reply_future else: return task_id, reply_future + def _start_write_thread(self): + # Caller must hold self._write_thread_lock + thread = threading.Thread( + target=self._execute_writes, + daemon=True, + name=f"_execute_writes for database {self.name}", + ) + thread.start() + self._write_thread = thread + self._write_threads_started += 1 + + def _write_thread_idle_timeout(self): + """Seconds the write thread may wait for a task before exiting. + + None means wait forever (the thread lives until close()). + """ + if self.is_memory: + # Closing the last connection to a memory database discards + # its contents, so its write connection must stay open + return None + get_timeout = getattr(self.ds, "_connection_idle_timeout_s", None) + timeout_s = get_timeout() if get_timeout is not None else 0 + return timeout_s or None + def _execute_writes(self): - # Infinite looping thread that protects the single write connection - # to this database - conn_exception = None - conn = None + # Thread that protects the single write connection to this database. + # It runs until close(), or until it has waited idle_timeout seconds + # with no task, in which case it closes its connection and exits. + # The next write starts a new thread. + writer = _WriteConnection(self) + deregistered = False try: - conn = self.connect(write=True) - # Threads do not inherit the caller's context, so any spans - # created by prepare_connection hooks here are root spans - self.ds._prepare_connection(conn, self.name) - except Exception as e: # noqa: BLE001 - # Stored and re-raised to whoever queues the next write - conn_exception = e + self._process_write_queue(writer) + deregistered = True + finally: + # Close outside _write_thread_lock: closing the last connection + # to a WAL database runs a checkpoint, which can be slow + writer.close() + current = threading.current_thread() + with self._write_thread_lock: + self._retiring_write_threads.discard(current) + if self._budgeted_writer: + if self._write_thread is current: + self._write_thread = None + self.ds._write_budget.wake() + elif not deregistered and self._write_thread is current: + # Died from an unexpected BaseException. Deregister so + # the next write starts a fresh thread, and start one + # now if tasks are already waiting. + self._write_thread = None + if not self._closed and not self._write_queue.empty(): + self._start_write_thread() + if self._closed and self._write_thread is current: + self._write_thread = None + if ( + self._cleanup_temp_file_on_writer_exit + and self._write_thread is None + and not self._retiring_write_threads + ): + self._cleanup_temp_file() + + def _process_write_queue(self, writer): + configured_idle_timeout = self._write_thread_idle_timeout() + + def current_idle_timeout(): + if writer.exception is not None: + # The connect error has been reported to every write already + # queued: exit as soon as the queue is empty, so the next + # write starts a new thread that tries to connect again + return 0 + return configured_idle_timeout + + idle_timeout = current_idle_timeout() while True: - task = self._write_queue.get() - if task is _SHUTDOWN: - if conn is not None: - try: - conn.close() - except Exception: # noqa: BLE001, S110 - # Best-effort close as the write thread exits - pass + try: + task = self._write_queue.get(timeout=idle_timeout) + except queue.Empty: + with self._write_thread_lock: + if not self._write_queue.empty(): + # A task was enqueued after the timeout fired + self._write_thread_exits_averted += 1 + continue + # Deregister while holding the lock: any write enqueued + # from now on will start a new thread + current = threading.current_thread() + if self._write_thread is current: + self._write_thread = None + self._retiring_write_threads.add(current) return + if task is _SHUTDOWN: + return + if task is RETIRE_WRITER: + if self._closed: + continue + return + if self._budgeted_writer and not self.ds._write_budget.take(self, task): + continue + if ( + task.isolated_connection + and writer.conn is not None + and not self.is_memory + ): + writer.close() + # 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 + # opened (the SchemaWatcher bumps the generation): reopen by + # path. connect() uses mode=rw for a file Datasette has seen, + # so a deleted file is reported, not silently recreated. + # Also reopened if something closed it (a callback reaching + # the raw connection through cursor.connection). + writer.reopen() + idle_timeout = current_idle_timeout() # block=True: the caller awaits the result, so the write spans # are children of the caller's span. The token must be detached # in the finally block or the context leaks into later writes. @@ -612,8 +1019,8 @@ 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 conn_exception is not None: - exception = conn_exception + if writer.exception is not None and needs_retained_connection: + exception = writer.exception elif task.isolated_connection: try: with tracer.start_as_current_span( @@ -624,22 +1031,24 @@ class Database: task.isolated_connection, ) span.set_attribute(TRANSACTION, task.transaction) - isolated_connection = self.connect(write=True) + isolated_connection = self._connect_write_thread() + # Owned by this thread: close() must not close it + # if this thread outlives close()'s join timeout try: - result = task.fn(isolated_connection) + result = _call_with_lease( + task.fn, isolated_connection, self.name, "isolated" + ) finally: - isolated_connection.close() - try: - self._all_connections.remove(isolated_connection) - except ValueError: - # May already have been cleared by close(). - pass - except Exception as e: # noqa: BLE001 + self._schema_check(isolated_connection, task.loop) + self._forget_connection(isolated_connection) + self._write_thread_connections.discard( + isolated_connection + ) + except BaseException as e: # noqa: BLE001 # Write thread must survive any task failure or the database wedges - sys.stderr.write(f"{e}\n") - sys.stderr.flush() - exception = e + exception = self._write_task_failed(e) else: + conn = writer.conn try: with tracer.start_as_current_span( DB_WRITE_EXECUTE, **write_span_kwargs @@ -649,20 +1058,52 @@ class Database: task.isolated_connection, ) span.set_attribute(TRANSACTION, task.transaction) - if task.transaction: - with conn: - conn.execute("BEGIN IMMEDIATE") - result = task.fn(conn) - else: - result = task.fn(conn) - except Exception as e: # noqa: BLE001 - sys.stderr.write(f"{e}\n") - sys.stderr.flush() - exception = e + # The callback gets a lease proxy that stops + # working when it returns: the connection belongs + # to this thread and is closed when it idles out + try: + if task.transaction: + with conn: + conn.execute("BEGIN IMMEDIATE") + result = _call_with_lease( + task.fn, conn, self.name, "write" + ) + else: + result = _call_without_transaction( + task.fn, conn, self.name + ) + finally: + self._schema_check(conn, task.loop) + except BaseException as e: # noqa: BLE001 + exception = self._write_task_failed(e) _deliver_write_result(task, result, exception) finally: if token is not None: otel_context_api.detach(token) + if self._budgeted_writer: + self.ds._write_budget.completed(self) + if self._budgeted_writer and self.ds._write_budget.should_yield(self): + return + + def _write_task_failed(self, e): + """The exception to deliver to the caller of a write task that + raised ``e``. An Exception is delivered as it is. Any other + BaseException (SystemExit, or a plugin's own BaseException subclass) + is delivered wrapped in a RuntimeError: re-raised by ``await`` on the + caller's event loop, SystemExit or KeyboardInterrupt would stop that + loop. The write thread keeps running (``with conn`` has rolled the + transaction back). Before this the thread died and the caller + waited forever.""" + if isinstance(e, Exception): + sys.stderr.write(f"{e}\n") + sys.stderr.flush() + return e + wrapped = RuntimeError( + f"A write callback for database {self.name!r} raised " + f"{type(e).__name__}: {e}" + ) + wrapped.__cause__ = e + return wrapped async def execute_fn(self, fn): """Run `fn(conn)` on a read connection, traced as a `db.query` span. @@ -685,30 +1126,30 @@ class Database: return await self._execute_fn(fn_in_execute_span) async def _execute_fn(self, fn): + check_reentrancy() 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) + return await self._run_in_executor(in_thread) + + async def _run_in_executor(self, fn): + """Run fn() on the shared SQL thread pool. The future is tracked in + _pending_execute_futures so close() waits for it before closing any + connection.""" with self._pending_execute_futures_lock: self._check_not_closed() # Run in a copy of the caller's context so spans created in the # thread have the correct parent. This needs a fresh copy for # each submit, since a Context cannot be entered concurrently. ctx = contextvars.copy_context() - future = self.ds.executor.submit(ctx.run, in_thread) + future = self.ds.executor.submit(ctx.run, fn) self._pending_execute_futures.add(future) future.add_done_callback(self._remove_pending_execute_future) return await asyncio.wrap_future(future) @@ -742,6 +1183,7 @@ class Database: ) as execute_span: try: with sqlite_timelimit(conn, time_limit_ms): + cursor = None try: cursor = conn.cursor() cursor.execute(sql, params if params is not None else {}) @@ -755,6 +1197,7 @@ class Database: else: rows = cursor.fetchall() truncated = False + description = cursor.description except (sqlite3.OperationalError, sqlite3.DatabaseError) as e: if e.args == ("interrupted",): raise QueryInterrupted(e, sql, params) @@ -764,6 +1207,13 @@ class Database: ) sys.stderr.flush() raise + finally: + # Release the statement now: a traceback that + # keeps this frame alive would otherwise keep it + # unfinalized, and closing the pooled connection + # later would leave a zombie holding its fds + if cursor is not None: + cursor.close() except QueryInterrupted as e: if not timeout_expected: execute_span.record_exception(e) @@ -776,10 +1226,10 @@ class Database: raise if truncate: - return Results(rows, truncated, cursor.description) + return Results(rows, truncated, description) else: - return Results(rows, False, cursor.description) + return Results(rows, False, description) with trace( # noqa: SIM117 "sql", database=self.name, sql=sql.strip(), params=params @@ -843,7 +1293,12 @@ class Database: elif self.is_memory: return 0 elif self.is_mutable: - return Path(self.path).stat().st_size + try: + return Path(self.path).stat().st_size + except FileNotFoundError: + # Deleted from under Datasette: report it as empty rather + # than failing every page that lists databases + return 0 elif self.ds.inspect_data and self.ds.inspect_data.get(self.name): self.cached_size = self.ds.inspect_data[self.name]["size"] return self.cached_size @@ -852,8 +1307,12 @@ class Database: return self.cached_size async def table_counts(self, limit=10): - if not self.is_mutable and self.cached_table_counts is not None: - return self.cached_table_counts + if not self.is_mutable: + if self.cached_table_counts is not None: + return self.cached_table_counts + # The result is cached for the life of the process, so do not + # let a caller's short limit cache timeouts (None) forever + limit = max(limit, self.immutable_count_time_limit_ms) # Try to get counts for each table, $limit timeout for each count counts = {} for table in await self.table_names(): @@ -1004,20 +1463,7 @@ class Database: has_spatialite = await self.execute_fn(detect_spatialite) if has_spatialite: # Also hide Spatialite internal tables - hidden_tables += [ - "ElementaryGeometries", - "SpatialIndex", - "geometry_columns", - "spatial_ref_sys", - "spatialite_history", - "sql_statements_log", - "sqlite_sequence", - "views_geometry_columns", - "virts_geometry_columns", - "data_licenses", - "KNN", - "KNN2", - ] + [ + hidden_tables += list(SPATIALITE_HIDDEN_TABLES) + [ r[0] for r in (await self.execute(""" select name from sqlite_master where name like "idx_%" @@ -1029,7 +1475,27 @@ class Database: async def derived_table_dependencies(self): """Return implementation tables and the tables they derive from.""" + state = self._watch_state + if state is not None and self.ds._schema_watcher.catalog_is_current(state): + # From the catalog rows (computed by the SchemaWatcher when it + # scanned this database, or read back from _internal), so a + # permission check on one of this database's tables does not + # open it and agrees with what allowed_resources() lists. Only + # when the catalog provably matches the file (a stat and a + # header read): this decides whether a table that derives from + # another - an external-content FTS table - may be shown, and a + # catalog that lags an out-of-band change would allow a table + # it has not seen yet. + from .utils.catalog import catalog_derived_table_dependencies + + return (await catalog_derived_table_dependencies(self.ds, [self]))[ + self.name + ] schema_version = (await self.execute("PRAGMA schema_version")).first()[0] + if state is not None: + # Found the catalog behind: rebuild it now rather than at the + # next write (owned) or poll (external) + self.ds._schema_watcher.note_live_version(state, schema_version) if ( self._cached_derived_table_dependencies is None or self._cached_derived_table_dependencies[0] != schema_version @@ -1131,6 +1597,127 @@ def _apply_write_wrapper(fn, wrapper_factory, track_event): return wrapped +def _pin_write_connection_pragmas(conn): + """Connection state Datasette's writes rely on, set explicitly on every + new write connection (before prepare_connection hooks, which may change + it). Write connections are reopened after an idle period, so anything + that used to "drift" into place on a long-lived connection must be set + here instead. + + recursive_triggers: sqlite-utils turns it on whenever it wraps a + connection (core's JSON write API does that), and the triggers that keep + sqlite-utils FTS indexes in sync need it for INSERT OR REPLACE to remove + the replaced row's index entry. + """ + conn.execute("PRAGMA recursive_triggers = on") + + +def _call_without_transaction(fn, conn, db_name): + try: + result = _call_with_lease(fn, conn, db_name, "write") + except BaseException: + if conn.in_transaction: + conn.rollback() + raise + if conn.in_transaction: + conn.rollback() + raise ConnectionLeaseError( + "transaction=False callbacks must commit or roll back before returning; " + "the unfinished transaction was rolled back" + ) + return result + + +def _call_with_lease(fn, conn, db_name, kind): + """Call fn with a LeasedConnection for conn that expires when it returns. + + Unlike read callbacks, write callbacks may return a cursor (for example + ``lambda conn: conn.execute(...)``, and execute_write_many() and + execute_write_script() return theirs); it is closed before it leaves + this thread, so reading attributes such as ``lastrowid`` from it + afterwards is fine, fetching rows is not. + """ + lease = LeasedConnection(conn, db_name, kind) + try: + with callback_scope(): + result = fn(lease) + except BaseException as e: + close_cursors(e) + if isinstance(e, DatabaseResourceError): + raise RuntimeError( # noqa: TRY004 - execution failure, not invalid input type + "Database callback started before a resource error; do not retry the write automatically" + ) from e + raise + else: + # A returned cursor is closed here, on the thread that owns the + # connection (see close_cursors()); its rowcount and lastrowid can + # still be read + close_cursors(result) + return result + finally: + lease._expire() + + +class _WriteConnection: + """The write thread's connection to its database. + + Reopened by the write thread when the file is replaced or deleted + (``Database._conn_generation`` moves). A failure to open or prepare the + connection is kept in ``exception`` and handed to every write that runs + while it is set. + """ + + __slots__ = ("conn", "db", "exception", "generation") + + def __init__(self, db): + self.db = db + self.conn = None + self.exception = None + self.generation = None + + def open(self): + db = self.db + # Read the generation before connecting: if it moves while we + # connect, the next task reopens - one connection too many, never + # a stale one + self.generation = db._conn_generation + self.exception = None + try: + self.conn = db._connect_write_thread() + _pin_write_connection_pragmas(self.conn) + # Threads do not inherit the caller's context, so any spans + # created by prepare_connection hooks here are root spans + with callback_scope(): + db.ds._prepare_connection(self.conn, db.name) + except BaseException as e: # noqa: BLE001 + self.exception = db._write_task_failed(e) + self.close() + + def reopen(self): + self.close() + self.open() + + def usable(self): + """False if the connection has been closed behind the write thread's + back. A failed open (``exception`` set) counts as usable: that error + is reported to each write until the thread exits and retries.""" + conn = self.conn + if conn is None: + return self.exception is not None + try: + return conn.total_changes >= 0 + except sqlite3.ProgrammingError: + return False + + def close(self): + conn, self.conn = self.conn, None + if conn is not None: + # Also drops it from _all_connections, so that list does not + # grow by one closed connection per write thread restart + self.db._forget_connection(conn) + self.db._write_thread_connections.discard(conn) + + class WriteTask: __slots__ = ( "block", @@ -1186,6 +1773,14 @@ def _deliver_write_result(task, result, exception): pass +def _unlink_temp_files(path): + for suffix in ("", "-wal", "-shm"): + try: + os.unlink(path + suffix) + except OSError: + pass + + class QueryInterrupted(Exception): def __init__(self, e, sql, params): self.e = e diff --git a/datasette/handle_exception.py b/datasette/handle_exception.py index ef6c7b7e..231961b4 100644 --- a/datasette/handle_exception.py +++ b/datasette/handle_exception.py @@ -9,6 +9,7 @@ from .utils.asgi import ( Base400, ) from .views.base import DatasetteError +from .write_budget import DatabaseResourceError # Debugger imports are deliberate - they back the "pdb" setting, which drops # into a debugger on unhandled exceptions @@ -34,7 +35,11 @@ def handle_exception(datasette, request, exception): title = None plain_message = None - if isinstance(exception, Base400): + if isinstance(exception, DatabaseResourceError): + status = 503 + info = {"code": exception.code, "execution_started": False} + message = str(exception) + elif isinstance(exception, Base400): status = exception.status info = {} message = exception.args[0] @@ -52,7 +57,9 @@ def handle_exception(datasette, request, exception): message = str(exception) traceback.print_exc() templates = [f"{status}.html", "error.html"] - headers = {} + headers = ( + {"Retry-After": "1"} if isinstance(exception, DatabaseResourceError) else {} + ) if datasette.cors: add_cors_headers(headers) if request.path.split("?")[0].endswith(".json"): diff --git a/datasette/hookspecs.py b/datasette/hookspecs.py index 49d8e8ea..78a9756b 100644 --- a/datasette/hookspecs.py +++ b/datasette/hookspecs.py @@ -267,3 +267,8 @@ def write_wrapper(datasette, database, request, transaction): Return ``None`` to skip wrapping. """ + + +@hookspec +def write_task_completed(datasette, database, task_id, exception): + """Completion of an accepted block=False write, including admission expiry.""" diff --git a/datasette/schema_watcher.py b/datasette/schema_watcher.py new file mode 100644 index 00000000..c7dfd8f5 --- /dev/null +++ b/datasette/schema_watcher.py @@ -0,0 +1,1391 @@ +""" +Keeps the ``_internal`` catalog (``catalog_databases``, ``catalog_tables``, +``catalog_columns`` ...) current without checking every attached database on +the request path. + +Every attached database has a *schema watch mode*: + +``owned`` + Every schema change goes through Datasette. After each write task the + write thread reads ``PRAGMA schema_version`` on the connection it just + used (microseconds, same thread). If the value moved, the catalog for + that one database is rebuilt before the write call returns, so callers + get read-your-writes on the catalog. Owned databases are never polled. + +``external`` + The file can also be changed by other processes (sqlite-utils, cron + jobs...). On top of the write-path check, a background task stats the + database file, its ``-wal`` and its ``-journal`` every + ``schema_watch_interval_ms``. Only a database whose fingerprint + ``(st_dev, st_ino, st_size, st_mtime_ns, st_ctime_ns)`` changed - or + whose previous fingerprint was *racily clean* - gets a + ``PRAGMA schema_version`` check on a short-lived connection, and only a + changed ``schema_version`` (or a replaced or deleted file) triggers a + catalog rebuild for that database. + +``immutable`` + Scanned once at startup and never again. + +Defaults: database files and named in-memory databases use the +``default_schema_watch`` setting (``external`` unless set to ``owned``), +whether they were passed to ``Datasette(files=...)`` / the CLI or added +later with ``datasette.add_database()``. Immutable files are ``immutable`` +and private ``:memory:`` databases ``owned``. Override per database with +``databases: {name: {schema_watch: owned|external|immutable}}`` in +``datasette.yaml`` or ``add_database(..., schema_watch=...)``. A named +in-memory database in ``external`` mode is polled with ``PRAGMA +schema_version`` on its write connection (shared-cache table locks do not +wait, so reading it from another connection could fail Datasette's writes). + +Several event loops +------------------- +One Datasette can be driven by several event loops at once, each in its own +thread. Scans in flight are tracked with ``concurrent.futures`` futures that +any loop can wait for; a scan is only taken over when the loop running it +has closed or stopped. There is at most one live polling task. State that a +write thread marks (``needs_scan``, ``_pending``) is set under a lock, before +the write's result is delivered, so nothing depends on a notification +reaching a loop that may have gone away. + +Replaced and deleted files +-------------------------- +A sweep that sees a new inode or a missing file calls +``Database._invalidate_connections()``: the database's connection generation +is bumped, idle pooled read connections are closed, leased ones are +discarded when their callback returns, and the write thread reopens its +connection before the next write. Once Datasette has seen a file exist, +write connections open it with ``mode=rw``, so a deleted file is never +silently recreated. + +Racily clean fingerprints +------------------------- +Like git's index: a fingerprint whose newest timestamp is within +``RACY_WINDOW_NS`` of the moment it was taken could be followed by another +write inside the same timestamp tick that leaves size, inode and mtime +identical. Such a fingerprint is flagged ``racy`` and the next sweep runs +the ``PRAGMA schema_version`` check even though the fingerprint looks +unchanged. The stat is always taken *before* the pragma, so a write landing +between the two either shows up in the pragma or changes the next stat. + +All catalog reads use one short-lived connection per database inside a +single read transaction, so ``schema_version`` and the rows written to the +catalog describe the same snapshot. The watcher never keeps a connection +open. +""" + +import asyncio +import errno +import json +import logging +import os +import threading +import time + +from .utils import sqlite3 +from .utils.catalog import remember_derived_table_dependencies +from .utils.inflight import InFlight, wait_for_concurrent +from .utils.internal_db import CATALOG_TABLES, collect_schema, write_catalog_entries + +logger = logging.getLogger("datasette.schema_watcher") + +MODES = ("owned", "external", "immutable") +# Valid values for the default_schema_watch setting +DEFAULT_MODES = ("external", "owned") +# Passed by Datasette.__init__ for files=/CLI databases: resolved to the +# default_schema_watch setting once settings are available +FILES_DEFAULT = "_files_default" +# Timestamps closer than this to the time of the stat are treated as racy. +# 2s covers coarse filesystems (FAT: 2s, HFS+: 1s) and clock skew on +# network filesystems; ext4/xfs/tmpfs on Linux are far finer than this. +RACY_WINDOW_NS = 2_000_000_000 +# Databases per thread hop / per internal-DB transaction for bulk scans +SCAN_CHUNK = 64 +SCAN_CONCURRENCY = 1 +# Directories with at least this many watched files are listed once per +# sweep with os.scandir() instead of probing -wal/-journal with stat() +SCANDIR_THRESHOLD = 8 +PRAGMA_BUSY_TIMEOUT_MS = 200 +# Read the schema cookie from the file header instead of opening a +# connection when no -wal/-journal has content +USE_HEADER_CHECK = True +# Seconds before a database whose catalog scan or catalog write failed for a +# transient reason (internal database locked, cancelled) is scanned again by +# a request or sweep; doubled after each consecutive failure up to the max +RETRY_DELAY_S = 1.0 +RETRY_DELAY_MAX_S = 30.0 + + +class WatchState: + __slots__ = ( + "catalog_version", + "db", + "error", + "fp", + "fp_racy", + "missing", + "mode", + "name", + "needs_scan", + "notified_version", + "removed", + "requested_mode", + "retry_at", + "retry_delay", + "scan_future", + "stats", + ) + + def __init__(self, db, name, requested_mode): + self.db = db + self.name = name + self.requested_mode = requested_mode + self.mode = None + self.fp = None + self.fp_racy = True + self.catalog_version = None + self.notified_version = None + self.needs_scan = True + # An InFlight (datasette.utils.inflight) while a catalog scan of + # this database runs + self.scan_future = None + # time.monotonic() before which a failed scan is not retried, and + # the delay used for the next failure (doubles, see RETRY_DELAY_S) + self.retry_at = 0.0 + self.retry_delay = RETRY_DELAY_S + self.missing = False + self.removed = False + self.error = None + self.stats = {"checks": 0, "scans": 0, "write_detections": 0} + + @property + def is_file(self): + db = self.db + return bool(db.path) and not db.is_memory + + def __repr__(self): + return f"" + + +def _fp_main(st): + return (st.st_dev, st.st_ino, st.st_size, st.st_mtime_ns, st.st_ctime_ns) + + +def _fp_side(st): + # No ctime for -wal/-journal: when running as root SQLite fchown()s + # these files to match the database every time a connection opens + # them, so their ctime moves on every open - including our own checks + return (st.st_ino, st.st_size, st.st_mtime_ns) + + +def fingerprint(path, present=None): + """(main, wal, journal) stat tuples; None for a file that does not exist. + + ``present`` is an optional set of file names in the directory (from + scandir) used to skip stat() calls for -wal/-journal files that are not + there. + """ + path = os.fspath(path) + try: + main = _fp_main(os.stat(path)) + except FileNotFoundError: + main = None + parts = [main] + base = os.path.basename(path) if present is not None else None + for suffix in ("-wal", "-journal"): + if present is not None and (base + suffix) not in present: + parts.append(None) + continue + try: + st = os.stat(path + suffix) + except FileNotFoundError: + parts.append(None) + continue + # An empty -wal holds no frames (and an empty -journal nothing to + # roll back): the database is exactly the main file, same as when + # the file is absent. Read-only opens of a WAL database create an + # empty -wal, which must not look like a change. + parts.append(_fp_side(st) if st.st_size else None) + return tuple(parts) + + +def newest_timestamp(fp): + main = fp[0] + newest = max(main[3], main[4]) if main is not None else 0 + for part in fp[1:]: + if part is not None: + newest = max(newest, part[2]) + return newest + + +def is_racy(fp, t_ns): + return newest_timestamp(fp) >= t_ns - RACY_WINDOW_NS + + +SQLITE_HEADER = b"SQLite format 3\x00" + + +def header_schema_version(path): + """The schema cookie (offset 40 of the database header) read straight + from the file - what PRAGMA schema_version returns - without opening a + SQLite connection or taking any lock. + + Only meaningful when there is no non-empty -wal or -journal: then the + main file is the whole database. The caller must have taken the + fingerprint *before* calling this; a commit in flight at that moment has + already bumped the mtime it saw, so the fingerprint is racy and the next + sweep looks again. + """ + try: + fd = os.open(path, os.O_RDONLY) + except OSError: + return None + try: + header = os.pread(fd, 100, 0) + except OSError: + return None + finally: + os.close(fd) + if len(header) < 100 or not header.startswith(SQLITE_HEADER): + return None + return int.from_bytes(header[40:44], "big") + + +def _is_transient(error): + from .write_budget import DatabaseResourceError + + if isinstance(error, DatabaseResourceError): + return True + if isinstance(error, OSError): + return error.errno in (errno.EMFILE, errno.ENFILE, errno.ENOMEM, errno.EAGAIN) + if isinstance(error, sqlite3.OperationalError): + code = getattr(error, "sqlite_errorcode", None) + if code is not None and code & 0xFF in ( + getattr(sqlite3, "SQLITE_CANTOPEN", None), + getattr(sqlite3, "SQLITE_NOMEM", None), + ): + 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", + "out of memory", + ) + ) + return False + + +def _fp_to_json(fp, t_ns, closed=False): + data = {"fp": fp, "t": t_ns} + if closed: + # Taken by Datasette.close() after every connection to an owned + # database had been closed: nothing in this process can change the + # file after the stat, so the racily-clean rule does not apply + data["closed"] = True + return json.dumps(data) + + +def _fp_from_json(value): + """(fp, t_ns, closed) - (None, None, False) if unreadable.""" + try: + data = json.loads(value) + fp = tuple(tuple(p) if p is not None else None for p in data["fp"]) + return fp, int(data["t"]), bool(data.get("closed")) + except Exception: # noqa: BLE001 + return None, None, False + + +def store_closing_fingerprints(internal_path, rows): + """Write SchemaWatcher.closing_fingerprints() to a persistent internal + database that has already been closed (Datasette.close() is synchronous + and the internal database's write thread is gone by then). Only rows + whose stored schema_version still matches are updated.""" + conn = sqlite3.connect(internal_path, isolation_level=None) + try: + conn.execute("BEGIN IMMEDIATE") + conn.executemany( + "UPDATE catalog_databases SET fingerprint = ? " + "WHERE database_name = ? AND schema_version = ?", + [(fp_json, name, version) for name, version, fp_json in rows], + ) + conn.execute("COMMIT") + finally: + conn.close() + + +class SchemaWatcher: + def __init__(self, ds): + self.ds = ds + self.states = {} + self._scan_resource_lock = threading.Lock() + # States whose catalog must be (re)built: registered after startup, + # changed by a write whose notification may never run, or whose last + # scan failed for a transient reason. Flushed by the next request or + # sweep. + self._pending = set() + self._pending_removals = set() # names whose catalog rows must go + # Guards claiming scans (WatchState.needs_scan / scan_future) and + # _pending: refresh() runs on several event loops in different + # threads, and write threads mark states from their own thread + self._lock = threading.Lock() + # The current polling task. At most one is live: a poller whose loop + # has closed or stopped is replaced, and a replaced poller exits at + # its next wake-up (it is no longer self._task) + self._task = None + self._pollers = set() + self._spawned = set() + self._last_sweep = 0.0 + # Bumped after every write to the catalog tables, so readers can + # cache things derived from the catalog (see datasette.utils.catalog). + # Locked: several event loops (threads) can store scans at once, and + # a lost increment would leave a stale cache in place + self.catalog_generation = 0 + self._generation_lock = threading.Lock() + self.counters = { + "sweeps": 0, + "sweep_seconds": 0.0, + "stats": 0, + "pragma_checks": 0, + "header_checks": 0, + "scans": 0, + "write_detections": 0, + "replaced": 0, + "missing": 0, + "busy": 0, + "restored_from_persisted": 0, + } + + def _bump_catalog_generation(self): + with self._generation_lock: + self.catalog_generation += 1 + + # ------------------------------------------------------------------ + # configuration + # ------------------------------------------------------------------ + @property + def interval_s(self): + ms = self.ds.setting("schema_watch_interval_ms") + return (ms or 0) / 1000.0 + + def _config_mode(self, name): + config = getattr(self.ds, "config", None) or {} + db_config = (config.get("databases") or {}).get(name) or {} + return db_config.get("schema_watch") + + def _default_mode_for_files(self): + from .utils import StartupError + + mode = self.ds.setting("default_schema_watch") or "external" + if mode not in DEFAULT_MODES: + raise StartupError( + f"Invalid default_schema_watch setting {mode!r}, expected one of " + f"{', '.join(DEFAULT_MODES)}" + ) + return mode + + def _resolve_mode(self, state): + db = state.db + mode = self._config_mode(state.name) or state.requested_mode + if mode is None: + if state.is_file or db.memory_name: + # Other code may change it: a plugin's own connection, a + # backup() into a named in-memory database, another process + mode = self._default_mode_for_files() + else: + # A private ":memory:" database: nothing else can see it + mode = "owned" + elif mode == FILES_DEFAULT: + mode = self._default_mode_for_files() + if mode not in MODES: + from .utils import StartupError + + raise StartupError( + f"Invalid schema_watch mode {mode!r} for database {state.name!r}, " + f"expected one of {', '.join(MODES)}" + ) + if not db.is_mutable: + # Nothing can change an immutable database + mode = "immutable" + elif mode == "external" and not state.is_file and not db.memory_name: + # A private ":memory:" database: there is nothing to poll. + # Named in-memory databases are polled with PRAGMA + # schema_version on their write connection. + mode = "owned" + return mode + + def configure(self): + """Resolve modes once ds.config is available (end of __init__).""" + # Validate the setting even when no database uses it + self._default_mode_for_files() + for state in self.states.values(): + state.mode = self._resolve_mode(state) + + # ------------------------------------------------------------------ + # registration + # ------------------------------------------------------------------ + def register(self, db, requested_mode=None): + if requested_mode not in (None, FILES_DEFAULT) and requested_mode not in MODES: + raise ValueError( + f"schema_watch must be one of {', '.join(MODES)}, not {requested_mode!r}" + ) + state = WatchState(db, db.name, requested_mode) + if hasattr(self.ds, "config"): + state.mode = self._resolve_mode(state) + old = self.states.get(db.name) + if old is not None: + old.removed = True + self.states[db.name] = state + db._watch_state = state + self._pending_removals.discard(db.name) + if self.ds.internal_db_created: + with self._lock: + self._pending.add(state) + if not db.memory_name: + self._spawn(self.refresh([state])) + # A shared-cache in-memory database is scanned by the next + # request or sweep instead: plugins typically fill one right + # after adding it (VACUUM INTO, backup()) from their own + # connection, and shared-cache table locks do not wait - a scan + # reading sqlite_master at that moment makes their writes fail + # with "database table is locked" + return state + + def unregister(self, name): + state = self.states.pop(name, None) + if state is None: + return + state.removed = True + state.db._watch_state = None + with self._lock: + self._pending.discard(state) + if self.ds.internal_db_created: + self._pending_removals.add(name) + self._spawn(self._delete_catalog([name])) + + # ------------------------------------------------------------------ + # task helpers + # ------------------------------------------------------------------ + async def _off_loop(self, fn, *args): + """Run blocking stat()/scan work in a worker thread - or inline with + num_sql_threads=0, which must work where threads cannot be started + at all (Pyodide).""" + if self.ds.executor is None: + return fn(*args) + return await asyncio.to_thread(fn, *args) + + def _spawn(self, coro): + try: + loop = asyncio.get_running_loop() + except RuntimeError: + # No loop (sync caller): the work stays pending and is picked + # up by the next flush/sweep + coro.close() + return None + task = loop.create_task(coro) + self._spawned.add(task) + + def _done(t): + self._spawned.discard(t) + if not t.cancelled() and t.exception() is not None: + logger.error("schema watcher task failed", exc_info=t.exception()) + + task.add_done_callback(_done) + return task + + def _poller_alive(self): + task = self._task + if task is None or task.done(): + return False + loop = task.get_loop() + return not loop.is_closed() and loop.is_running() + + def loop_running(self): + """True while a polling task is alive, on any event loop.""" + return self._poller_alive() + + async def start(self): + """Start the polling loop unless one is already alive on any event + loop. Runs one sweep first so the catalog is current when the first + request is served.""" + if self.interval_s <= 0 or self.ds._closed: + return + loop = asyncio.get_running_loop() + with self._lock: + if self._poller_alive(): + return + # Claim the slot before awaiting so concurrent start() calls, + # on this loop or another, no-op + task = loop.create_task(self._run(), name="datasette-schema-watcher") + self._task = task + self._pollers.add(task) + task.add_done_callback(self._poller_done) + if self.ds.internal_db_created: + try: + await self.ds._refresh_schemas(background=True) + except Exception: + logger.exception("schema watcher initial sweep failed") + + def _poller_done(self, task): + with self._lock: + self._pollers.discard(task) + + async def _run(self): + me = asyncio.current_task() + # Wait one interval first: start() has just swept. A poller that + # has been replaced (its loop stopped for a while and another loop + # started a new one) exits instead of sweeping alongside it. + while not self.ds._closed and self._task is me: + await asyncio.sleep(self.interval_s) + if self.ds._closed or self._task is not me: + return + try: + await self.ds._refresh_schemas(background=True) + except asyncio.CancelledError: + raise + except Exception: + logger.exception("schema watcher sweep failed") + + def _all_tasks(self): + with self._lock: + return [t for t in [self._task, *self._pollers, *self._spawned] if t] + + def stop(self): + tasks = self._all_tasks() + self._task = None + for task in tasks: + if task.done(): + continue + loop = task.get_loop() + try: + running = asyncio.get_running_loop() + except RuntimeError: + running = None + if running is loop: + task.cancel() + elif not loop.is_closed(): + try: + loop.call_soon_threadsafe(task.cancel) + except RuntimeError: + pass + + async def astop(self): + tasks = self._all_tasks() + self.stop() + loop = asyncio.get_running_loop() + mine = [t for t in tasks if t.get_loop() is loop and not t.done()] + if mine: + await asyncio.wait(mine, timeout=5) + + # ------------------------------------------------------------------ + # request path: O(1) unless add/remove_database or a write left work + # pending + # ------------------------------------------------------------------ + async def on_request(self): + if (self._pending or self._pending_removals) and self.ds.internal_db_created: + try: + await self.flush_pending() + except Exception: + # Never fail the request: what is still pending is retried + # (with back-off) by later requests and sweeps + logger.warning("Could not update the catalog", exc_info=True) + if ( + self.interval_s > 0 + and self.ds.internal_db_created + and not self._poller_alive() + ): + await self.start() + + def _ready_pending(self): + now = time.monotonic() + with self._lock: + return [s for s in self._pending if not s.removed and s.retry_at <= now] + + async def flush_pending(self): + if self._pending_removals: + await self._delete_catalog(list(self._pending_removals)) + ready = self._ready_pending() + if ready: + await self.refresh(ready) + + # ------------------------------------------------------------------ + # write path + # ------------------------------------------------------------------ + def check_after_write(self, state, conn, loop=None): + """Called on the thread that owns ``conn`` right after a write task. + + ``loop`` is the event loop to notify (None when already on it).""" + try: + version = conn.execute("PRAGMA schema_version").fetchone()[0] + except Exception: # noqa: BLE001 + return + with self._lock: + if ( + state.removed + or version == state.catalog_version + or version == state.notified_version + ): + return + state.notified_version = version + # Marked here, on the write thread, before the write's result is + # delivered: a blocking writer always finds needs_scan set in + # Database._after_write(). The notification below only starts a + # scan sooner; if it never runs (a block=False write from an + # event loop that has since closed or stopped) the next request + # or sweep picks the state up from _pending. + state.needs_scan = True + state.retry_at = 0.0 + self._pending.add(state) + state.stats["write_detections"] += 1 + self.counters["write_detections"] += 1 + if loop is None: + self._on_write_schema_change(state) + else: + try: + loop.call_soon_threadsafe(self._on_write_schema_change, state) + except RuntimeError: + pass + + def _on_write_schema_change(self, state): + if state.removed: + return + if self.ds.internal_db_created: + self._spawn(self.refresh([state])) + + async def after_write(self, state): + """Awaited by blocking writes: gives read-your-writes on the catalog.""" + if not self.ds.internal_db_created or state.removed: + return + try: + await self.refresh([state]) + except Exception: + logger.exception("catalog refresh after write to %r failed", state.name) + + # ------------------------------------------------------------------ + # sweeps (stat prefilter) + # ------------------------------------------------------------------ + async def sweep(self, *, background): + """Check databases for changes made outside Datasette. + + background=True (the polling loop): external-mode files only. + background=False (explicit refresh_schemas()): every mutable + database - files via the stat prefilter, named in-memory databases + via a PRAGMA schema_version check. + """ + t0 = time.perf_counter() + if self._pending_removals: + await self._delete_catalog(list(self._pending_removals)) + candidates = [] + for state in list(self.states.values()): + if state.mode == "immutable" or state.mode is None: + if state.needs_scan: + candidates.append(state) + continue + if background and state.mode != "external": + continue + candidates.append(state) + outcomes = [] + file_candidates = [s for s in candidates if s.is_file] + if file_candidates: + outcomes = await self._off_loop( + self._sweep_sync, file_candidates, background + ) + for state in candidates: + if not state.is_file and state.db.memory_name and not state.removed: + outcomes.append(await self._memory_pragma_check(state)) + to_scan = self._ready_pending() + to_clear = [] + for state, kind, fp, t_ns, version in outcomes: + if state.removed: + continue + self.counters["stats"] += 1 + if fp is not None: + state.fp = fp + state.fp_racy = is_racy(fp, t_ns) + if fp[0] is not None: + state.db._file_seen = True + if kind == "clean": + pass + elif kind == "failed_unchanged": + continue + elif kind == "missing": + self.counters["missing"] += 1 + if not state.missing: + state.missing = True + state.db._invalidate_connections() + to_clear.append(state) + elif kind == "replaced": + self.counters["replaced"] += 1 + state.missing = False + state.db._invalidate_connections() + state.needs_scan = True + elif kind == "busy": + self.counters["busy"] += 1 + state.fp_racy = True # check again next sweep + elif kind in ("checked", "header"): + self.counters[ + "pragma_checks" if kind == "checked" else "header_checks" + ] += 1 + state.stats["checks"] += 1 + if version != state.catalog_version: + state.needs_scan = True + elif kind == "error": + state.fp_racy = True + if state.needs_scan: + to_scan.append(state) + if to_clear: + await self._clear_missing(to_clear) + if to_scan: + await self.refresh(to_scan) + self.counters["sweeps"] += 1 + self.counters["sweep_seconds"] += time.perf_counter() - t0 + self._last_sweep = time.monotonic() + + def _sweep_sync(self, states, background=False): + with self._scan_resource_lock: + return self._sweep_serial(states, background) + + def _sweep_serial(self, states, background=False): + """Runs in a worker thread. Reads state, never writes it.""" + by_dir = {} + for state in states: + if state.is_file: + by_dir.setdefault( + os.path.dirname(os.path.abspath(state.db.path)), [] + ).append(state) + out = [] + for directory, dir_states in by_dir.items(): + present = None + if len(dir_states) >= SCANDIR_THRESHOLD: + try: + with os.scandir(directory) as it: + present = {entry.name for entry in it} + except OSError: + present = None + for state in dir_states: + t_ns = time.time_ns() + fp = fingerprint(state.db.path, present) + prev = state.fp + if fp == prev and not state.fp_racy and not state.needs_scan: + out.append((state, "clean", fp, t_ns, None)) + continue + if ( + background + and state.error is not None + and fp == prev + # Not state.fp_racy: failed checks set that flag to + # force a recheck, which would retry forever + and not is_racy(fp, t_ns) + ): + # The last scan failed (missing extension module, not + # a database, ...) and the file has not changed since: + # retrying would fail again. An explicit + # refresh_schemas() still retries. + out.append((state, "failed_unchanged", fp, t_ns, None)) + continue + if fp[0] is None: + out.append((state, "missing", fp, t_ns, None)) + continue + if prev is None or prev[0] is None or prev[0][:2] != fp[0][:2]: + # New inode/device (atomic replace, delete+recreate) or + # first sight: the schema_version alone cannot be trusted + out.append((state, "replaced", fp, t_ns, None)) + continue + if USE_HEADER_CHECK and fp[1] is None and fp[2] is None: + # No -wal/-journal content: read the cookie from the + # header. No connection, no SHARED lock that could make + # a concurrent writer back off. + version = header_schema_version(os.fspath(state.db.path)) + if version is not None: + out.append((state, "header", fp, t_ns, version)) + continue + out.append(self._pragma_check(state, fp, t_ns)) + return out + + async def _memory_pragma_check(self, state): + def _version(conn): + return conn.execute("PRAGMA schema_version").fetchone()[0] + + try: + version = await state.db._execute_on_write_connection(_version) + except Exception: # noqa: BLE001 + return (state, "error", None, None, None) + return (state, "checked", None, None, version) + + def _pragma_check(self, state, fp, t_ns): + db = state.db + try: + conn = self._connect(db, prepare=False) + except Exception: # noqa: BLE001 + return (state, "error", fp, t_ns, None) + try: + conn.execute(f"PRAGMA busy_timeout={PRAGMA_BUSY_TIMEOUT_MS}") + version = conn.execute("PRAGMA schema_version").fetchone()[0] + return (state, "checked", fp, t_ns, version) + except sqlite3.OperationalError as e: + msg = str(e) + if "locked" in msg or "busy" in msg: + return (state, "busy", fp, t_ns, None) + return (state, "error", fp, t_ns, None) + except Exception: # noqa: BLE001 + return (state, "error", fp, t_ns, None) + finally: + self._close(db, conn) + + # ------------------------------------------------------------------ + # scanning + storing + # ------------------------------------------------------------------ + def _connect(self, db, prepare): + # Untracked: Database.close() on the event loop thread must not close + # this connection while a worker thread is using it (that segfaults). + # It is always closed by _close() in the same thread that opened it. + # Counted separately so database lifecycle operations can wait for + # scans without closing a connection while its thread is using it. + db._untracked_connection_opened() + try: + conn = db.connect(track=False) + except BaseException: + db._untracked_connection_closed() + raise + try: + if prepare: + # crossdb=False: a scan of _memory reads its own schema only, + # it must not ATTACH (open) up to ten other database files + self.ds._prepare_connection(conn, db.name, crossdb=False) + else: + conn.row_factory = None + except BaseException: + self._close(db, conn) + raise + return conn + + @staticmethod + def _close(db, conn): + try: + conn.close() + finally: + db._untracked_connection_closed() + + def _scan_sync(self, state): + """Read schema_version + full schema in one read transaction.""" + db = state.db + fp = t_ns = None + if state.is_file: + t_ns = time.time_ns() + fp = fingerprint(db.path) + if fp[0] is None: + return {"state": state, "missing": True, "fp": fp, "t": t_ns} + try: + conn = self._connect(db, prepare=True) + try: + conn.row_factory = sqlite3.Row + conn.execute("BEGIN") + try: + version = conn.execute("PRAGMA schema_version").fetchone()[0] + schema = collect_schema(conn, db.name) + finally: + conn.rollback() + finally: + self._close(db, conn) + except Exception as e: # noqa: BLE001 + # Keep the fingerprint: a background sweep retries a failed scan + # only once the file changes, instead of every interval + return {"state": state, "error": e, "fp": fp, "t": t_ns} + return { + "state": state, + "fp": fp, + "t": t_ns, + "version": version, + "schema": schema, + } + + async def _scan_on_write_connection(self, state): + name = state.db.name + + def _scan(conn): + version = conn.execute("PRAGMA schema_version").fetchone()[0] + # The write thread's post-task schema check must not re-notify + # for the version this scan is about to store + state.notified_version = version + return version, collect_schema(conn, name) + + try: + version, schema = await state.db._execute_on_write_connection(_scan) + except Exception as e: # noqa: BLE001 + return {"state": state, "error": e} + return { + "state": state, + "fp": None, + "t": None, + "version": version, + "schema": schema, + } + + def _scan_chunk_sync(self, states): + with self._scan_resource_lock: + return self._scan_chunk_serial(states) + + def _scan_chunk_serial(self, states): + results = [] + for state in states: + try: + results.append(self._scan_sync(state)) + except Exception as e: # noqa: BLE001 + results.append({"state": state, "error": e}) + return results + + async def refresh(self, states): + """Rebuild the catalog for every state that needs it, coalescing with + scans already in flight for the same database - on this event loop + or any other.""" + loop = asyncio.get_running_loop() + pending = list(states) + # States whose scan this call ran and that failed transiently: left + # in _pending for a later request or sweep, not retried here + failed = set() + for _ in range(10): + mine = [] + waits = [] + with self._lock: + for state in pending: + inflight = state.scan_future + if inflight is not None and not inflight.done(): + if not inflight.abandoned(): + waits.append(inflight.future) + continue + # Its event loop has closed or stopped, so nothing + # will finish that scan. Take over: its results are + # discarded (_store() checks scan_future) and its + # waiters are woken to wait for this one instead + state.scan_future = None + state.needs_scan = True + inflight.finish() + if state.needs_scan and not state.removed and state not in failed: + state.needs_scan = False + state.scan_future = InFlight(loop) + mine.append(state) + if mine: + records = {state: state.scan_future for state in mine} + handled = set() + try: + await self._scan_and_store(mine, records, handled) + finally: + now = time.monotonic() + with self._lock: + for state, record in records.items(): + if state.scan_future is not record: + # Taken over by another loop (see above) + continue + state.scan_future = None + if state not in handled and not state.removed: + # The scan or the catalog write failed + # transiently (internal database locked, a + # shared-cache lock) or was cancelled: the + # catalog is not current, so make sure a + # later request or sweep tries again + state.needs_scan = True + state.retry_at = now + state.retry_delay + state.retry_delay = min( + state.retry_delay * 2, RETRY_DELAY_MAX_S + ) + self._pending.add(state) + failed.add(state) + for record in records.values(): + record.finish() + if waits: + await asyncio.wait([wait_for_concurrent(f) for f in waits]) + pending = [ + s + for s in pending + if not s.removed + and s not in failed + and ( + s.needs_scan + or (s.scan_future is not None and not s.scan_future.done()) + ) + and s.error is None + ] + if not pending: + break + with self._lock: + for state in states: + # Failed scans are retried by sweeps, not by every request + if not state.needs_scan or state.error is not None: + self._pending.discard(state) + + async def wait_for_scan(self, state): + """Wait for a catalog scan of state that is in flight, if any.""" + inflight = state.scan_future + if inflight is not None and not inflight.done() and not inflight.abandoned(): + await wait_for_concurrent(inflight.future) + + async def _scan_and_store(self, states, records, handled): + memory = [s for s in states if not s.is_file and s.db.memory_name] + if memory: + results = [] + for state in memory: + results.append(await self._scan_on_write_connection(state)) + await self._store(results, records, handled) + states = [s for s in states if s not in memory] + if not states: + return + chunks = [states[i : i + SCAN_CHUNK] for i in range(0, len(states), SCAN_CHUNK)] + semaphore = asyncio.Semaphore(SCAN_CONCURRENCY) + + async def one(chunk): + async with semaphore: + results = await self._off_loop(self._scan_chunk_sync, chunk) + await self._store(results, records, handled) + + if len(chunks) == 1: + await one(chunks[0]) + else: + await asyncio.gather(*(one(c) for c in chunks)) + + async def _store(self, results, records=None, handled=None): + """Write scan results to the catalog. ``records`` maps each state to + the InFlight this scan was started as: results for a state whose + scan has since been taken over by another loop are dropped. States + whose outcome was recorded are added to ``handled``; the others + (transient failures) are retried later by refresh().""" + records = records or {} + if handled is None: + handled = set() + + def current(state): + record = records.get(state) + return record is None or state.scan_future is record + + entries = [] + missing = [] + stored = [] + for r in results: + state = r["state"] + if state.removed or not current(state): + handled.add(state) + continue + if "error" in r: + if state.db._closed: + # Closed (or deleted) while the scan was running + 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( + "Transient error reading schema of database %r, will retry: %s", + state.name, + r["error"], + ) + continue + state.error = r["error"] + state.needs_scan = True + if r.get("fp") is not None and r["fp"][0] is not None: + state.fp = r["fp"] + state.fp_racy = is_racy(r["fp"], r["t"]) + logger.warning( + "Could not read schema of database %r: %s", state.name, r["error"] + ) + handled.add(state) + continue + state.error = None + if r.get("missing"): + missing.append(state) + state.fp = r["fp"] + state.fp_racy = True + handled.add(state) + continue + db = state.db + fp_json = _fp_to_json(r["fp"], r["t"]) if r["fp"] is not None else None + entries.append( + ( + state, + ( + state.name, + str(db.path) if db.path is not None else None, + db.is_memory, + r["version"], + fp_json, + r["schema"], + ), + ) + ) + stored.append(r) + if entries: + watcher = self + + def _write(conn): + # Runs on the internal database's write thread, which + # serializes catalog writes: skip databases removed (or + # re-added) while we were scanning, and scans another loop + # has taken over - its newer rows must not be overwritten + live = [ + entry + for state, entry in entries + if watcher.states.get(entry[0]) is state and current(state) + ] + write_catalog_entries(conn, live) + + await self.ds.get_internal_database().execute_write_fn(_write) + for r in stored: + state = r["state"] + handled.add(state) + if state.removed or not current(state): + continue + state.catalog_version = r["version"] + state.missing = False + state.retry_at = 0.0 + state.retry_delay = RETRY_DELAY_S + if r["fp"] is not None: + state.fp = r["fp"] + state.fp_racy = is_racy(r["fp"], r["t"]) + if r["fp"][0] is not None: + state.db._file_seen = True + state.stats["scans"] += 1 + self.counters["scans"] += 1 + remember_derived_table_dependencies( + state, [(t[1], t[3]) for t in r["schema"]["tables"]] + ) + # After the states: a reader that sees the new generation must + # also see their new catalog versions + self._bump_catalog_generation() + if missing: + for state in missing: + if not state.missing: + state.missing = True + state.db._invalidate_connections() + await self._clear_missing(missing) + + async def _clear_missing(self, states): + names = [s.name for s in states] + + def _clear(conn): + for name in names: + for table in CATALOG_TABLES[:-1]: + conn.execute(f"DELETE FROM {table} WHERE database_name = ?", [name]) + conn.execute( + "UPDATE catalog_databases SET schema_version = NULL, fingerprint = NULL " + "WHERE database_name = ?", + [name], + ) + + await self.ds.get_internal_database().execute_write_fn(_clear) + for state in states: + state.catalog_version = None + self._bump_catalog_generation() + + async def _delete_catalog(self, names): + watcher = self + + def _delete(conn): + for name in names: + if name in watcher.states: + # Re-added under the same name; its own scan replaces rows + continue + for table in CATALOG_TABLES: + conn.execute(f"DELETE FROM {table} WHERE database_name = ?", [name]) + + await self.ds.get_internal_database().execute_write_fn(_delete) + self._bump_catalog_generation() + for name in names: + self._pending_removals.discard(name) + + # ------------------------------------------------------------------ + # startup + # ------------------------------------------------------------------ + async def initial_scan(self): + """Full catalog build at startup. + + Deletes catalog rows for databases that are not attached (a + persistent --internal database can carry rows from a previous run), + then scans every database - except those whose persisted fingerprint + proves the file is unchanged since the catalog was written. + """ + internal = self.ds.get_internal_database() + names = list(self.states.keys()) + names_json = json.dumps(names) + + def _delete_stale(conn): + for table in CATALOG_TABLES: + conn.execute( + f"DELETE FROM {table} WHERE database_name IS NULL OR database_name NOT IN " + "(SELECT value FROM json_each(?))", + [names_json], + ) + + await internal.execute_write_fn(_delete_stale) + self._bump_catalog_generation() + self._pending_removals.clear() + states = [s for s in self.states.values() if not s.removed] + persisted = {} + if not internal.is_temp_disk: + rows = await internal.execute( + "select database_name, path, schema_version, fingerprint from catalog_databases" + ) + persisted = {row["database_name"]: row for row in rows.rows} + to_scan = [] + if persisted: + reuse = await self._off_loop(self._reusable_sync, states, persisted) + for state in states: + hit = reuse.get(state.name) + if hit is None: + to_scan.append(state) + continue + fp, t_ns, version = hit + state.db._file_seen = True + state.fp = fp + state.fp_racy = is_racy(fp, t_ns) + state.catalog_version = version + state.needs_scan = False + self.counters["restored_from_persisted"] += 1 + else: + to_scan = states + for state in to_scan: + state.needs_scan = True + if to_scan: + await self.refresh(to_scan) + with self._lock: + # Keep states that a write marked during the scan, or whose scan + # failed transiently + self._pending = {s for s in self._pending if s.needs_scan and not s.removed} + + def _reusable_sync(self, states, persisted): + """Databases whose stored fingerprint matches the file right now and + was not racy when it was stored. Stat calls only - no connections.""" + out = {} + for state in states: + row = persisted.get(state.name) + if row is None or not state.is_file or row["fingerprint"] is None: + continue + if row["path"] != str(state.db.path) or row["schema_version"] is None: + continue + stored_fp, stored_t, closed = _fp_from_json(row["fingerprint"]) + if stored_fp is None or (is_racy(stored_fp, stored_t) and not closed): + continue + t_ns = time.time_ns() + fp = fingerprint(state.db.path) + if fp == stored_fp: + out[state.name] = (fp, t_ns, row["schema_version"]) + return out + + # ------------------------------------------------------------------ + # shutdown + # ------------------------------------------------------------------ + def closing_fingerprints(self): + """Called by Datasette.close() once every attached database has been + closed. For each owned file database whose catalog is current, + returns ``(name, schema_version, fingerprint_json)`` for the file as + it is now. Owned databases are only changed through Datasette, so a + restart with a persistent internal database can reuse their catalog + rows without opening them - even if data was written (or the -wal + checkpointed by the final close) after the catalog was last built. + Stat calls only.""" + out = [] + for state in list(self.states.values()): + if ( + state.mode != "owned" + or not state.is_file + or state.removed + or state.missing + or state.needs_scan + or state.scan_future is not None + or state.error is not None + or state.catalog_version is None + or state.notified_version not in (None, state.catalog_version) + ): + continue + t_ns = time.time_ns() + try: + fp = fingerprint(state.db.path) + except OSError: + continue + if fp[0] is None or fp[1] is not None or fp[2] is not None: + # Missing, or a -wal/-journal with content: the header + # check below would not describe the whole database + continue + if state.fp is None or state.fp[0] is None or state.fp[0][:2] != fp[0][:2]: + # Replaced since the catalog was built + continue + if header_schema_version(os.fspath(state.db.path)) != state.catalog_version: + # Changed by something other than Datasette while it ran: + # the catalog is stale and must not be restored next time + continue + out.append( + (state.name, state.catalog_version, _fp_to_json(fp, t_ns, closed=True)) + ) + return out + + # ------------------------------------------------------------------ + # single-database freshness checks (permission decisions) + # ------------------------------------------------------------------ + def catalog_is_current(self, state): + """True if this database's catalog rows provably describe it as it + is right now - checked without opening a SQLite connection. + + Used before trusting the catalog for a decision about one database + (the derived-table permission rule), where a catalog that lags an + out-of-band change could expose a table. Immutable databases cannot + change. A file database is current if it is the same file the + catalog was built from, has no -wal/-journal content, and the schema + cookie in its header equals the catalog's version. Anything else + (in-memory databases, live WAL content) returns False and the caller + asks the database itself. + """ + if state.mode == "immutable": + return True + if ( + not state.is_file + or state.catalog_version is None + or state.needs_scan + or state.missing + ): + return False + path = os.fspath(state.db.path) + try: + fp = fingerprint(path) + except OSError: + return False + if fp[0] is None or fp[1] is not None or fp[2] is not None: + return False + prev = state.fp + if prev is None or prev[0] is None or prev[0][:2] != fp[0][:2]: + return False + return header_schema_version(path) == state.catalog_version + + def note_live_version(self, state, version): + """A caller read ``PRAGMA schema_version`` from the database itself. + If the catalog is behind, schedule a rescan instead of waiting for + the next write (owned) or poll (external).""" + with self._lock: + if ( + state.removed + or state.error is not None + or state.needs_scan + or state.scan_future is not None + or version == state.catalog_version + ): + return + state.needs_scan = True + self._pending.add(state) + if self.ds.internal_db_created: + self._spawn(self.refresh([state])) + + # ------------------------------------------------------------------ + # introspection + # ------------------------------------------------------------------ + def status(self): + modes = {} + for state in self.states.values(): + modes[state.mode] = modes.get(state.mode, 0) + 1 + return { + "interval_s": self.interval_s, + "loop_running": self.loop_running(), + "modes": modes, + "pending": len(self._pending), + "pending_removals": len(self._pending_removals), + **self.counters, + } diff --git a/datasette/telemetry.py b/datasette/telemetry.py index 2fd5a9dc..1ef1865b 100644 --- a/datasette/telemetry.py +++ b/datasette/telemetry.py @@ -311,9 +311,7 @@ sql_operation_duration = meter.create_histogram( write_queue_wait = meter.create_histogram( M_WRITE_QUEUE_WAIT, unit=M_WRITE_QUEUE_WAIT.unit, - description=( - "Time a write spent queued behind the single write thread for its database" - ), + description=("Time a write spent queued behind the write thread for its database"), explicit_bucket_boundaries_advisory=M_WRITE_QUEUE_WAIT.buckets, ) @@ -426,7 +424,7 @@ def observe_pending_queries(options=None): def observe_write_queue_depth(options=None): - "Writes queued behind the single write thread, per database." + "Writes queued behind a database's write thread, per database." for ds in _live_instances(): for db in _databases_of(ds): write_queue = db._write_queue @@ -437,7 +435,11 @@ def observe_write_queue_depth(options=None): def observe_open_connections(options=None): - "Open SQLite connections tracked for closing, per database." + """ + Open SQLite connections tracked for closing, per database: pooled read + connections (idle and leased), the write connection while its thread + runs, and isolated connections in use. + """ for ds in _live_instances(): for db in _databases_of(ds): yield otel_metrics.Observation( @@ -470,7 +472,7 @@ write_queue_depth_gauge = meter.create_observable_gauge( M_WRITE_QUEUE_DEPTH, callbacks=[observe_write_queue_depth], unit=M_WRITE_QUEUE_DEPTH.unit, - description="Writes queued behind a database's single write thread", + description="Writes queued behind a database's write thread", ) open_connections_gauge = meter.create_observable_gauge( diff --git a/datasette/telemetry_registry.py b/datasette/telemetry_registry.py index 63d5b89b..9792718f 100644 --- a/datasette/telemetry_registry.py +++ b/datasette/telemetry_registry.py @@ -311,7 +311,9 @@ DB_WRITE_QUEUE_WAIT = SpanName( "Time a write spent waiting in its database's write queue. For " "``block=True``, this is a child of ``db.query``. For ``block=False``, " "it is a root span linked to the span that queued the write, since the " - "write can outlive that request.", + "write can outlive that request. For the first write after the write " + "thread was stopped for being idle, this includes opening and preparing " + "a new write connection.", ) DB_WRITE_EXECUTE = SpanName( @@ -429,7 +431,9 @@ M_WRITE_QUEUE_WAIT = MetricName( "datasette.write.queue_wait", HISTOGRAM, "s", - "Time each write waited in its database's write queue.", + "Time each write waited in its database's write queue, including " + "opening a new write connection for the first write after the write " + "thread was stopped for being idle.", (DB_NAMESPACE,), buckets=DURATION_BUCKETS, ) @@ -476,9 +480,10 @@ M_WRITE_QUEUE_DEPTH = MetricName( "datasette.write.queue_depth", GAUGE, "{write}", - "Writes waiting for a database's single write thread. Increasing " - "``num_sql_threads`` does not increase write concurrency. Not reported for " - "databases that have never been written to.", + "Writes waiting for a database's write thread. Each database has at " + "most one, so increasing ``num_sql_threads`` does not increase write " + "concurrency; it stops after ``connection_idle_timeout_ms`` without " + "writes. Not reported for databases that have never been written to.", (DB_NAMESPACE,), ) @@ -486,7 +491,11 @@ M_CONNECTIONS_OPEN = MetricName( "datasette.connections.open", GAUGE, "{connection}", - "Open SQLite connections managed by Datasette.", + "Open SQLite connections managed by Datasette: pooled read connections " + "(capped across databases by ``max_open_connections``), the write " + "connection while the database's write thread is running, and isolated " + "connections in use. Short-lived connections used to check for schema " + "changes are not counted.", (DB_NAMESPACE,), ) diff --git a/datasette/utils/__init__.py b/datasette/utils/__init__.py index ad39025f..50967a69 100644 --- a/datasette/utils/__init__.py +++ b/datasette/utils/__init__.py @@ -67,21 +67,22 @@ class PaginatedResources: for resource in self.resources: yield resource - # Continue fetching subsequent pages if there are more - next_token = self.next - while next_token: - page = await self._datasette.allowed_resources( + # Fetch everything after the first page in one query. Each page + # re-runs the whole permission query, so paging through N resources + # 1000 at a time cost O(N^2 / 1000) - 21 seconds for the index page + # with 87,000 tables. The caller sees the same sequence either way. + if self.next: + rest = await self._datasette._allowed_resources_page( self._action, self._actor, parent=self._parent, include_is_private=self._include_is_private, include_reasons=self._include_reasons, - limit=self._limit, - next=next_token, + limit=None, + next=self.next, ) - for resource in page.resources: + for resource in rest.resources: yield resource - next_token = page.next # From https://www.sqlite.org/lang_keywords.html @@ -851,19 +852,32 @@ def detect_fts_sql(table): ) +_detected_json1 = None + + def detect_json1(conn=None): - close_conn = False + """Does SQLite have the JSON functions? With no ``conn`` this checks the + SQLite library itself, which cannot change while the process runs, so + the answer is cached instead of opening a connection on every call (it + is called for every HTML table page and facet request).""" + global _detected_json1 if conn is None: - conn = sqlite3.connect(":memory:") - close_conn = True + if _detected_json1 is None: + probe = sqlite3.connect(":memory:") + try: + _detected_json1 = _json1_works(probe) + finally: + probe.close() + return _detected_json1 + return _json1_works(conn) + + +def _json1_works(conn): try: conn.execute("SELECT json('{}')") return True except sqlite3.Error: return False - finally: - if close_conn: - conn.close() def table_columns(conn, table): @@ -1212,12 +1226,16 @@ class SpatialiteConnectionProblem(ConnectionProblem): def check_connection(conn): - tables = [ - r[0] - for r in conn.execute( - "select name from sqlite_master where type='table'" - ).fetchall() - ] + try: + tables = [ + r[0] + for r in conn.execute( + "select name from sqlite_master where type='table'" + ).fetchall() + ] + except sqlite3.DatabaseError as e: + # e.g. "file is not a database" + raise ConnectionProblem(e) for table in tables: try: conn.execute( diff --git a/datasette/utils/catalog.py b/datasette/utils/catalog.py new file mode 100644 index 00000000..344b16b6 --- /dev/null +++ b/datasette/utils/catalog.py @@ -0,0 +1,388 @@ +"""Schema facts read from the ``_internal`` catalog instead of the databases. + +Pages that summarise many databases at once (the index page, permission +listings) must not open every database: with a capped connection pool a loop +over thousands of databases opens and evicts a connection per database at a +0% hit rate. The SchemaWatcher keeps the ``catalog_*`` tables current (owned +databases from the write path, external ones by polling), so these helpers +answer the same questions as the live introspection methods on +:class:`~datasette.database.Database` with a handful of ``_internal`` queries +that cover any number of databases. + +They describe whatever the catalog currently holds - exactly what +``allowed_resources()`` lists - so a listing and the facts used to filter or +decorate it always come from the same snapshot. +""" + +import json + +from .sqlite import ( + _VIRTUAL_TABLE_SHADOW_SUFFIXES, + SQLITE_STAT_TABLES, + _is_fts_content_virtual_table, + _virtual_table_module, + derived_table_dependencies_from_rows, +) + +# Tables hidden in databases that contain a geometry_columns table, matching +# Database.hidden_table_names() +SPATIALITE_HIDDEN_TABLES = ( + "ElementaryGeometries", + "SpatialIndex", + "geometry_columns", + "spatial_ref_sys", + "spatialite_history", + "sql_statements_log", + "sqlite_sequence", + "views_geometry_columns", + "virts_geometry_columns", + "data_licenses", + "KNN", + "KNN2", +) + + +def _names_param(names): + return json.dumps(list(names)) + + +def derived_cache_key(state): + """Identifies one version of a database's catalog rows. A replaced file + can keep its schema_version, so the scan count is part of the key.""" + return ("catalog", state.catalog_version, state.stats["scans"]) + + +def remember_derived_table_dependencies(state, table_rows): + """Called by the SchemaWatcher with the ``(name, sql)`` table rows it has + just written to the catalog, so permission checks never need to read + them back.""" + state.db._cached_derived_table_dependencies = ( + derived_cache_key(state), + derived_table_dependencies_from_rows(table_rows), + ) + + +async def catalog_derived_table_dependencies(datasette, databases): + """``{name: {implementation table: source table}}`` for the given watched + ``Database`` objects, from the rows in the catalog. + + Each database's map is cached against its catalog rows (normally filled + in by the SchemaWatcher when it scans the database); the ones missing - + databases restored from a persisted catalog - are read back in one + ``_internal`` query, never from the databases themselves. + """ + result = {} + missing = [] + for db in databases: + state = db._watch_state + if state.missing or state.catalog_version is None: + # No catalog rows, so nothing listed that could need filtering + result[db.name] = {} + continue + key = derived_cache_key(state) + cached = db._cached_derived_table_dependencies + if cached is not None and cached[0] == key: + result[db.name] = cached[1] + else: + missing.append((db, key)) + if missing: + # Only databases with a virtual table can have dependencies + rows = await datasette.get_internal_database().execute( + """ + select database_name, table_name, sql from catalog_tables + where database_name in ( + select database_name from catalog_tables + where coalesce(rootpage, 0) = 0 + and database_name in (select value from json_each(:names)) + ) + order by rowid + """, + {"names": _names_param(db.name for db, _ in missing)}, + ) + by_database = {} + for database_name, table_name, sql in rows.rows: + by_database.setdefault(database_name, []).append((table_name, sql)) + for db, key in missing: + deps = derived_table_dependencies_from_rows(by_database.get(db.name, ())) + db._cached_derived_table_dependencies = (key, deps) + result[db.name] = deps + return result + + +async def all_derived_table_dependencies(datasette): + """``{database: dependencies}`` for every watched database, cached until + the catalog next changes (``SchemaWatcher.catalog_generation``), so each + page of ``allowed_resources()`` costs a dictionary lookup instead of a + query against every attached database.""" + generation = datasette._schema_watcher.catalog_generation + cached = datasette._catalog_derived_cache + if cached is not None and cached[0] == generation: + return cached[1] + watched = [db for db in datasette.databases.values() if db._watch_state is not None] + result = await catalog_derived_table_dependencies(datasette, watched) + result = {name: deps for name, deps in result.items() if deps} + datasette._catalog_derived_cache = (generation, result) + return result + + +class CatalogSummary: + """What the index page needs to know about one database's tables.""" + + def __init__(self, name): + self.name = name + self.views = set() + # (table, sql, type) for virtual tables and tables typed as shadow + self.special_rows = [] + self.has_spatialite = False + self._shadow = None + self._content_fts = None + + def _classify(self): + virtual = [(t, sql) for t, sql, _ in self.special_rows if sql is not None] + typed_shadow = {t for t, _, type_ in self.special_rows if type_ == "shadow"} + untyped = any(type_ is None for _, _, type_ in self.special_rows) + shadow = set(typed_shadow) + if untyped: + # Rows written without PRAGMA table_list types: fall back to the + # DDL-based rule for SQLite's built-in modules. Candidate names + # that do not exist are harmless - only existing tables are + # looked up. + for virtual_table, sql in virtual: + module = _virtual_table_module(sql) + for suffix in _VIRTUAL_TABLE_SHADOW_SUFFIXES.get(module, ()): + shadow.add(virtual_table + suffix) + self._shadow = shadow + self._content_fts = { + t for t, sql in virtual if _is_fts_content_virtual_table(sql) + } + + def is_hidden(self, table, config_hidden=()): + """Matches ``table in await db.hidden_table_names()`` for an + existing table.""" + if self._shadow is None: + self._classify() + if table in config_hidden: + return True + if table in SQLITE_STAT_TABLES or table.startswith("_"): + return True + if table in self._shadow or table in self._content_fts: + return True + return self.has_spatialite and ( + table in SPATIALITE_HIDDEN_TABLES + # sqlite_master "name like 'idx_%'": case-insensitive, "_" is + # a single-character wildcard + or (len(table) >= 4 and table[:3].lower() == "idx") + ) + + +async def catalog_summaries(datasette, database_names): + """``{database: CatalogSummary}`` for the named databases, in two + ``_internal`` queries whatever the number of databases.""" + internal = datasette.get_internal_database() + params = {"names": _names_param(database_names)} + summaries = {name: CatalogSummary(name) for name in database_names} + views = await internal.execute( + """ + select database_name, view_name from catalog_views + where database_name in (select value from json_each(:names)) + """, + params, + ) + for database_name, view_name in views.rows: + summary = summaries.get(database_name) + if summary is not None: + summary.views.add(view_name) + if view_name == "geometry_columns": + summary.has_spatialite = True + special = await internal.execute( + """ + select database_name, table_name, + case when coalesce(rootpage, 0) = 0 then sql end as sql, type + from catalog_tables + where database_name in (select value from json_each(:names)) + and ( + coalesce(rootpage, 0) = 0 + or type is null + or type = 'shadow' + or table_name = 'geometry_columns' + ) + order by rowid + """, + params, + ) + for database_name, table_name, sql, type_ in special.rows: + summary = summaries.get(database_name) + if summary is None: + continue + if table_name == "geometry_columns": + summary.has_spatialite = True + if sql is not None and _virtual_table_module(sql) is None: + sql = None + if sql is not None or type_ in (None, "shadow"): + summary.special_rows.append((table_name, sql, type_)) + return summaries + + +async def catalog_table_details(datasette, pairs): + """``{(database, table): {"columns": [...], "primary_keys": [...], + "fts_table": str | None}}`` for the given pairs, matching + ``table_columns()``, ``primary_keys()`` and ``fts_table()``.""" + from . import detect_fts_sql + + pairs = list(pairs) + details = { + pair: {"columns": [], "primary_keys": [], "fts_table": None} for pair in pairs + } + if not pairs: + return details + internal = datasette.get_internal_database() + columns = await internal.execute( + """ + select database_name, table_name, name, is_pk from catalog_columns + where (database_name, table_name) in ( + select json_extract(value, '$[0]'), json_extract(value, '$[1]') + from json_each(:pairs) + ) + order by database_name, table_name, cid + """, + {"pairs": json.dumps(pairs)}, + ) + pks = {} + for database_name, table_name, name, is_pk in columns.rows: + entry = details.get((database_name, table_name)) + if entry is None: + continue + entry["columns"].append(name) + if is_pk: + pks.setdefault((database_name, table_name), []).append((is_pk, name)) + for pair, pk_columns in pks.items(): + # Stable sort by position in the primary key, like detect_primary_keys() + pk_columns.sort(key=lambda item: item[0]) + details[pair]["primary_keys"] = [name for _, name in pk_columns] + # detect_fts() against the catalog: the same LIKE patterns, so the same + # matching rules; first match in sqlite_master (= catalog rowid) order + fts_args = [] + for database_name, table in pairs: + _, fts_params = detect_fts_sql(table) + fts_args.append( + [ + database_name, + table, + fts_params["fts_double_quoted"], + fts_params["fts_bracket_quoted"], + ] + ) + fts = await internal.execute( + r""" + select p.db, p.tbl, ( + select c.table_name from catalog_tables c + where c.database_name = p.db + and coalesce(c.rootpage, 0) = 0 + and ( + c.sql like p.dq escape char(92) + or c.sql like p.bq escape char(92) + or (c.table_name = p.tbl and c.sql like '%VIRTUAL TABLE%USING FTS%') + ) + order by c.rowid limit 1 + ) + from ( + select json_extract(value, '$[0]') as db, json_extract(value, '$[1]') as tbl, + json_extract(value, '$[2]') as dq, json_extract(value, '$[3]') as bq + from json_each(:args) + ) p + """, + {"args": json.dumps(fts_args)}, + ) + for database_name, table, fts_table in fts.rows: + entry = details.get((database_name, table)) + if entry is not None: + entry["fts_table"] = fts_table + return details + + +async def catalog_relationship_counts(datasette, database_names): + """``{database: {table: incoming + outgoing foreign keys}}``, counted the + way ``get_all_foreign_keys()`` builds its lists: compound foreign keys + are skipped, and so are references to tables that do not exist.""" + internal = datasette.get_internal_database() + rows = await internal.execute( + """ + with fk as ( + select database_name, table_name, id, max("table") as other + from catalog_foreign_keys + where database_name in (select value from json_each(:names)) + group by database_name, table_name, id + having count(*) = 1 + ), + valid as ( + select fk.database_name, fk.table_name, fk.other from fk + where exists ( + select 1 from catalog_tables t + where t.database_name = fk.database_name and t.table_name = fk.other + ) + ) + select database_name, table_name, count(*) from ( + select database_name, table_name from valid + union all + select database_name, other as table_name from valid + ) + group by database_name, table_name + """, + {"names": _names_param(database_names)}, + ) + counts = {} + for database_name, table_name, n in rows.rows: + counts.setdefault(database_name, {})[table_name] = n + return counts + + +async def catalog_all_foreign_keys(datasette, database_name): + """``get_all_foreign_keys()`` for one database, from the catalog: + ``{table: {"incoming": [...], "outgoing": [...]}}`` for every table.""" + internal = datasette.get_internal_database() + tables = await internal.execute( + "select table_name from catalog_tables where database_name = ? order by table_name", + [database_name], + ) + table_to_foreign_keys = { + row[0]: {"incoming": [], "outgoing": []} for row in tables.rows + } + rows = await internal.execute( + """ + select table_name, id, "table", "from", "to" from catalog_foreign_keys + where database_name = ? + order by table_name, id, seq + """, + [database_name], + ) + by_table = {} + for table_name, fk_id, other_table, from_, to_ in rows.rows: + by_table.setdefault(table_name, {}).setdefault(fk_id, []).append( + (other_table, from_, to_) + ) + for table, fks in by_table.items(): + if table not in table_to_foreign_keys: + continue + for parts in fks.values(): + if len(parts) != 1: + # Compound foreign keys are left out, as get_outbound_foreign_keys() does + continue + other_table, from_, to_ = parts[0] + if other_table not in table_to_foreign_keys: + # Weird edge case where something refers to a table that does + # not actually exist + continue + table_to_foreign_keys[other_table]["incoming"].append( + {"other_table": table, "column": to_, "other_column": from_} + ) + table_to_foreign_keys[table]["outgoing"].append( + {"other_table": other_table, "column": from_, "other_column": to_} + ) + for foreign_keys in table_to_foreign_keys.values(): + foreign_keys["incoming"].sort( + key=lambda fk: (fk["other_table"], fk["column"], fk["other_column"]) + ) + foreign_keys["outgoing"].sort( + key=lambda fk: (fk["other_table"], fk["column"], fk["other_column"]) + ) + return table_to_foreign_keys diff --git a/datasette/utils/inflight.py b/datasette/utils/inflight.py new file mode 100644 index 00000000..f73b5310 --- /dev/null +++ b/datasette/utils/inflight.py @@ -0,0 +1,63 @@ +"""Waiting for work started on another event loop. + +One Datasette can be driven by several event loops at once, each in its own +thread (``datasette --get``, TestClient, pytest-asyncio, an embedding app). +asyncio primitives (``asyncio.Lock``, ``asyncio.Future``) belong to one loop, +so "one call does the work, concurrent callers wait for it" is built from a +``concurrent.futures.Future`` that any loop can wait for, and a +``threading.Lock`` held only for the moment it takes to claim the work. +""" + +import asyncio +import concurrent.futures + + +class InFlight: + """One piece of work in progress, started by ``task`` on ``loop``.""" + + __slots__ = ("future", "loop", "task") + + def __init__(self, loop, task=None): + self.future = concurrent.futures.Future() + self.loop = loop + self.task = task + + def done(self): + return self.future.done() + + def abandoned(self): + """True if nothing will ever finish this work: the loop running it + has closed, or has stopped (``run_until_complete()`` returned while + the task doing it was still pending).""" + loop = self.loop + return loop.is_closed() or not loop.is_running() + + def finish(self): + """Wake every waiter (they check for themselves whether the work + succeeded). Safe to call more than once, from any thread.""" + try: + self.future.set_result(None) + except concurrent.futures.InvalidStateError: + pass + + +def wait_for_concurrent(cf_future): + """An asyncio future on the running loop that completes when + ``cf_future`` does. Cancelling it does not cancel ``cf_future``, which + other waiters - maybe on other loops - share.""" + loop = asyncio.get_running_loop() + waiter = loop.create_future() + + def _set(): + if not waiter.done(): + waiter.set_result(None) + + def _done(_): + try: + loop.call_soon_threadsafe(_set) + except RuntimeError: + # That loop has closed; nobody is waiting any more + pass + + cf_future.add_done_callback(_done) + return waiter diff --git a/datasette/utils/internal_db.py b/datasette/utils/internal_db.py index 0ddeb847..a71d732c 100644 --- a/datasette/utils/internal_db.py +++ b/datasette/utils/internal_db.py @@ -3,7 +3,12 @@ import textwrap from sqlite_utils import Database as SQLiteUtilsDatabase from sqlite_utils import Migrations -from datasette.utils import escape_sqlite, table_column_details +from datasette.utils import escape_sqlite, sqlite3, table_column_details +from datasette.utils.sqlite import ( + sqlite_table_list_types, + supports_table_xinfo, + table_types_from_rows, +) INTERNAL_DB_SCHEMA_TABLES = { "catalog_databases", @@ -180,87 +185,202 @@ async def init_internal_db(db): await db.execute_write_fn(apply_migrations, transaction=False) -async def populate_schema_tables(internal_db, db, schema_version): - database_name = db.name +# Children first, catalog_databases last, so deletes work with FK enforcement +CATALOG_TABLES = ( + "catalog_columns", + "catalog_foreign_keys", + "catalog_indexes", + "catalog_views", + "catalog_tables", + "catalog_databases", +) - tables = (await db.execute("select * from sqlite_master WHERE type = 'table'")).rows - views = (await db.execute("select * from sqlite_master WHERE type = 'view'")).rows - def collect_info(conn): - tables_to_insert = [] - views_to_insert = [] - columns_to_insert = [] - foreign_keys_to_insert = [] - indexes_to_insert = [] +@internal_migrations(name="0002_catalog_fingerprint") +def catalog_fingerprint_column(db): + # stat() fingerprint of the database file when the catalog rows were + # written, so a restart with a persistent internal database can skip + # unchanged files without opening them + if "fingerprint" not in db["catalog_databases"].columns_dict: + db.execute("ALTER TABLE catalog_databases ADD COLUMN fingerprint TEXT") - for view in views: - view_name = view["name"] - views_to_insert.append( - (database_name, view_name, view["rootpage"], view["sql"]) - ) - for table in tables: - table_name = table["name"] - tables_to_insert.append( - (database_name, table_name, table["rootpage"], table["sql"]) - ) - columns = table_column_details(conn, table_name) - columns_to_insert.extend( - { - "database_name": database_name, - "table_name": table_name, - **column._asdict(), - } - for column in columns - ) - foreign_keys = conn.execute( - f"PRAGMA foreign_key_list({escape_sqlite(table_name)})" - ).fetchall() - foreign_keys_to_insert.extend( - { - "database_name": database_name, - "table_name": table_name, - **dict(foreign_key), - } - for foreign_key in foreign_keys - ) - indexes = conn.execute( - f"PRAGMA index_list({escape_sqlite(table_name)})" - ).fetchall() - indexes_to_insert.extend( - { - "database_name": database_name, - "table_name": table_name, - **dict(index), - } - for index in indexes - ) - return ( - tables_to_insert, - views_to_insert, - columns_to_insert, - foreign_keys_to_insert, - indexes_to_insert, +@internal_migrations(name="0003_catalog_table_type") +def catalog_table_type_column(db): + # PRAGMA table_list's type for each table ("table", "virtual" or + # "shadow"), so hidden tables can be worked out from the catalog + # without opening the database. NULL means unknown. + if "type" not in db["catalog_tables"].columns_dict: + db.execute("ALTER TABLE catalog_tables ADD COLUMN type TEXT") + # Rows written before this column existed have no type: forget + # their fingerprints so the next startup rescans those databases + # instead of restoring them from the persisted catalog + db.execute("UPDATE catalog_databases SET fingerprint = NULL") + + +_COLUMNS_SQL = """ +SELECT m.name, p.cid, p.name, p.type, p."notnull", p.dflt_value, p.pk, p.hidden +FROM sqlite_master m, pragma_table_xinfo(m.name) p WHERE m.type = 'table' +""" +_FOREIGN_KEYS_SQL = """ +SELECT m.name, p.id, p.seq, p."table", p."from", p."to", p.on_update, p.on_delete, p."match" +FROM sqlite_master m, pragma_foreign_key_list(m.name) p WHERE m.type = 'table' +""" +_INDEXES_SQL = """ +SELECT m.name, p.seq, p.name, p."unique", p.origin, p.partial +FROM sqlite_master m, pragma_index_list(m.name) p WHERE m.type = 'table' +""" + + +def collect_schema(conn, database_name): + """Read everything the catalog needs from one connection. + + Run it inside a read transaction so it sees a single snapshot. Uses + three joined pragma table-valued-function queries per database instead + of three PRAGMA calls per table, falling back to the per-table loop if + that fails (e.g. a virtual table whose module is not loaded). + """ + schema = None + if supports_table_xinfo(): + try: + schema = _collect_schema_joined(conn, database_name) + except sqlite3.DatabaseError: + pass + if schema is None: + schema = _collect_schema_per_table(conn, database_name) + # Record each table's type (table/virtual/shadow) as SQLite itself + # reports it - with this connection's extensions loaded, so shadow + # tables of extension modules are recognised too + types = sqlite_table_list_types(conn) + if types is None: + types = table_types_from_rows((t[1], t[3]) for t in schema["tables"]) + schema["tables"] = [(*t, types.get(t[1])) for t in schema["tables"]] + return schema + + +def _collect_schema_joined(conn, database_name): + tables = conn.execute("select * from sqlite_master WHERE type = 'table'").fetchall() + views = conn.execute("select * from sqlite_master WHERE type = 'view'").fetchall() + columns = [ + { + "database_name": database_name, + "table_name": r[0], + "cid": r[1], + "name": r[2], + "type": r[3], + "notnull": r[4], + "default_value": r[5], + "is_pk": r[6], + "hidden": r[7], + } + for r in conn.execute(_COLUMNS_SQL).fetchall() + ] + foreign_keys = [ + { + "database_name": database_name, + "table_name": r[0], + "id": r[1], + "seq": r[2], + "table": r[3], + "from": r[4], + "to": r[5], + "on_update": r[6], + "on_delete": r[7], + "match": r[8], + } + for r in conn.execute(_FOREIGN_KEYS_SQL).fetchall() + ] + indexes = [ + { + "database_name": database_name, + "table_name": r[0], + "seq": r[1], + "name": r[2], + "unique": r[3], + "origin": r[4], + "partial": r[5], + } + for r in conn.execute(_INDEXES_SQL).fetchall() + ] + return { + "tables": [(database_name, t["name"], t["rootpage"], t["sql"]) for t in tables], + "views": [(database_name, v["name"], v["rootpage"], v["sql"]) for v in views], + "columns": columns, + "foreign_keys": foreign_keys, + "indexes": indexes, + } + + +def _collect_schema_per_table(conn, database_name): + tables = conn.execute("select * from sqlite_master WHERE type = 'table'").fetchall() + views = conn.execute("select * from sqlite_master WHERE type = 'view'").fetchall() + tables_to_insert = [] + views_to_insert = [] + columns_to_insert = [] + foreign_keys_to_insert = [] + indexes_to_insert = [] + + for view in views: + views_to_insert.append( + (database_name, view["name"], view["rootpage"], view["sql"]) ) - ( - tables_to_insert, - views_to_insert, - columns_to_insert, - foreign_keys_to_insert, - indexes_to_insert, - ) = await db.execute_fn(collect_info) + for table in tables: + table_name = table["name"] + tables_to_insert.append( + (database_name, table_name, table["rootpage"], table["sql"]) + ) + columns = table_column_details(conn, table_name) + columns_to_insert.extend( + { + "database_name": database_name, + "table_name": table_name, + **column._asdict(), + } + for column in columns + ) + foreign_keys = conn.execute( + f"PRAGMA foreign_key_list({escape_sqlite(table_name)})" + ).fetchall() + foreign_keys_to_insert.extend( + { + "database_name": database_name, + "table_name": table_name, + **dict(foreign_key), + } + for foreign_key in foreign_keys + ) + indexes = conn.execute( + f"PRAGMA index_list({escape_sqlite(table_name)})" + ).fetchall() + indexes_to_insert.extend( + { + "database_name": database_name, + "table_name": table_name, + **dict(index), + } + for index in indexes + ) + return { + "tables": tables_to_insert, + "views": views_to_insert, + "columns": columns_to_insert, + "foreign_keys": foreign_keys_to_insert, + "indexes": indexes_to_insert, + } - def replace_catalog(conn): + +def write_catalog_entries(conn, entries): + """Replace the catalog rows for each database in ``entries``. + + Each entry is (database_name, path, is_memory, schema_version, + fingerprint_json, schema) where schema comes from collect_schema(). + Runs on the internal database write connection, in one transaction. + """ + for database_name, path, is_memory, schema_version, fingerprint, schema in entries: # Delete child rows before their catalog_tables parents so this also # works if a prepare_connection plugin enables foreign key enforcement. - for table in ( - "catalog_columns", - "catalog_foreign_keys", - "catalog_indexes", - "catalog_views", - "catalog_tables", - ): + for table in CATALOG_TABLES[:-1]: conn.execute( f"DELETE FROM {table} WHERE database_name = ?", [database_name], @@ -268,29 +388,24 @@ async def populate_schema_tables(internal_db, db, schema_version): conn.execute( """ INSERT OR REPLACE INTO catalog_databases ( - database_name, path, is_memory, schema_version - ) VALUES (?, ?, ?, ?) + database_name, path, is_memory, schema_version, fingerprint + ) VALUES (?, ?, ?, ?, ?) """, - [ - database_name, - str(db.path) if db.path is not None else None, - db.is_memory, - schema_version, - ], + [database_name, path, is_memory, schema_version, fingerprint], ) conn.executemany( """ - INSERT INTO catalog_tables (database_name, table_name, rootpage, sql) - values (?, ?, ?, ?) + INSERT INTO catalog_tables (database_name, table_name, rootpage, sql, type) + values (?, ?, ?, ?, ?) """, - tables_to_insert, + [t if len(t) == 5 else (*t, None) for t in schema["tables"]], ) conn.executemany( """ INSERT INTO catalog_views (database_name, view_name, rootpage, sql) values (?, ?, ?, ?) """, - views_to_insert, + schema["views"], ) conn.executemany( """ @@ -300,7 +415,7 @@ async def populate_schema_tables(internal_db, db, schema_version): :database_name, :table_name, :cid, :name, :type, :notnull, :default_value, :is_pk, :hidden ) """, - columns_to_insert, + schema["columns"], ) conn.executemany( """ @@ -310,7 +425,7 @@ async def populate_schema_tables(internal_db, db, schema_version): :database_name, :table_name, :id, :seq, :table, :from, :to, :on_update, :on_delete, :match ) """, - foreign_keys_to_insert, + schema["foreign_keys"], ) conn.executemany( """ @@ -320,7 +435,29 @@ async def populate_schema_tables(internal_db, db, schema_version): :database_name, :table_name, :seq, :name, :unique, :origin, :partial ) """, - indexes_to_insert, + schema["indexes"], ) + +async def populate_schema_tables(internal_db, db, schema_version): + """Rebuild the catalog rows for one database (kept for compatibility - + the SchemaWatcher reads the schema on its own short-lived connection).""" + database_name = db.name + + def _collect(conn): + return collect_schema(conn, database_name) + + schema = await db.execute_fn(_collect) + entry = ( + database_name, + str(db.path) if db.path is not None else None, + db.is_memory, + schema_version, + None, + schema, + ) + + def replace_catalog(conn): + write_catalog_entries(conn, [entry]) + await internal_db.execute_write_fn(replace_catalog) diff --git a/datasette/utils/sqlite.py b/datasette/utils/sqlite.py index 2ae1be9b..f4b88be9 100644 --- a/datasette/utils/sqlite.py +++ b/datasette/utils/sqlite.py @@ -139,6 +139,11 @@ def check_structured_write_table(conn, table: str, *, allow_missing=False): raise ValueError("Structured writes require an ordinary table") +SQLITE_STAT_TABLES = frozenset( + {"sqlite_stat1", "sqlite_stat2", "sqlite_stat3", "sqlite_stat4"} +) + + def sqlite_hidden_table_names(conn, *, schema: str | None = "main") -> list[str]: schema_table = _sqlite_schema_table(schema) try: @@ -147,13 +152,27 @@ def sqlite_hidden_table_names(conn, *, schema: str | None = "main") -> list[str] ).fetchall() except sqlite3.DatabaseError: return [] + return hidden_table_names_from_rows( + rows, + table_type=lambda name: sqlite_table_type(conn, name, schema=schema), + ) + + +def hidden_table_names_from_rows(rows, table_type) -> list[str]: + """Hidden tables, given ``(name, sql)`` rows for every table in one + schema (in ``sqlite_master`` order) and a ``table_type(name)`` callable. + + Shared by the live introspection above and by the ``_internal`` catalog + (see ``datasette.utils.catalog``), which stores ``PRAGMA table_list`` + types when it scans a database. + """ hidden_tables = [] content_fts_tables = [] for name, sql in rows: if ( - name in {"sqlite_stat1", "sqlite_stat2", "sqlite_stat3", "sqlite_stat4"} + name in SQLITE_STAT_TABLES or name.startswith("_") - or sqlite_table_type(conn, name, schema=schema) == "shadow" + or table_type(name) == "shadow" ): hidden_tables.append(name) elif _is_fts_content_virtual_table(sql): @@ -161,6 +180,48 @@ def sqlite_hidden_table_names(conn, *, schema: str | None = "main") -> list[str] return sorted(hidden_tables) + content_fts_tables +def table_types_from_rows(rows) -> dict[str, SQLiteTableType]: + """Classify ``(name, sql)`` table rows as table/virtual/shadow from their + DDL alone - the fallback used when ``PRAGMA table_list`` is unavailable. + Only knows the shadow tables of SQLite's built-in modules.""" + rows = list(rows) + names = {name for name, _ in rows} + types = {} + for name, sql in rows: + types[name] = "virtual" if _virtual_table_module(sql) is not None else "table" + for virtual_table, sql in rows: + module = _virtual_table_module(sql) + if module is None: + continue + for suffix in _VIRTUAL_TABLE_SHADOW_SUFFIXES.get(module, ()): + if virtual_table + suffix in names: + types[virtual_table + suffix] = "shadow" + return types + + +def sqlite_table_list_types(conn, *, schema: str = "main") -> dict | None: + """``{table name: type}`` from ``PRAGMA table_list`` for one schema, or + None if this SQLite does not support it (before 3.37) or it fails.""" + if not supports_table_list(): + return None + try: + # The statement form cannot be shadowed by a user-created relation + # called pragma_table_list - see sqlite_table_type() + cursor = conn.execute(f"PRAGMA {_quote_identifier(schema)}.table_list") + columns = [description[0] for description in cursor.description] + types = {} + for row in cursor.fetchall(): + record = dict(zip(columns, row)) + if record.get("schema") != schema: + continue + row_type = record.get("type") + if row_type in {"table", "virtual", "shadow"}: + types[record.get("name")] = row_type + return types + except sqlite3.DatabaseError: + return None + + def sqlite_derived_table_dependencies( conn, *, schema: str | None = "main" ) -> dict[str, str]: @@ -178,7 +239,16 @@ def sqlite_derived_table_dependencies( rows = conn.execute( f"select name, sql from {schema_table} where type = 'table'" ).fetchall() + return derived_table_dependencies_from_rows(rows, schema=schema) + +def derived_table_dependencies_from_rows( + rows, *, schema: str | None = "main" +) -> dict[str, str]: + """``sqlite_derived_table_dependencies()`` for ``(name, sql)`` rows of + every table in one schema - from a live connection or from the + ``_internal`` catalog.""" + rows = list(rows) table_names = {row[0] for row in rows} # SQLite identifiers fold ASCII letters only. identifier_case = str.maketrans( @@ -284,6 +354,13 @@ def _quote_identifier(value: str) -> str: def _virtual_table_module(sql: str | None) -> str | None: if not sql: return None + if "virtual" not in sql.lower(): + # Fast path: removing comments only ever replaces text with spaces, + # so it cannot produce the VIRTUAL keyword. Skips the + # character-by-character comment scan for ordinary tables, which + # cost ~1.4s of event loop time per 17,500 tables when the + # SchemaWatcher computes derived-table dependencies at startup. + return None match = _VIRTUAL_TABLE_MODULE_RE.search(_strip_sql_comments(sql)) if match is None: return None diff --git a/datasette/views/base.py b/datasette/views/base.py index 7262edb4..75a152d9 100644 --- a/datasette/views/base.py +++ b/datasette/views/base.py @@ -116,7 +116,7 @@ class BaseView: async def dispatch_request(self, request): if self.ds: - await self.ds.refresh_schemas() + await self.ds._schema_watcher.on_request() handler = getattr(self, request.method.lower(), None) response = await handler(request) if self.ds.cors: diff --git a/datasette/views/database.py b/datasette/views/database.py index da207601..3034a891 100644 --- a/datasette/views/database.py +++ b/datasette/views/database.py @@ -36,6 +36,11 @@ from datasette.utils import ( named_parameters as derive_named_parameters, ) from datasette.utils.asgi import AsgiFileDownload, Forbidden, NotFound, Response +from datasette.utils.catalog import ( + catalog_all_foreign_keys, + catalog_summaries, + catalog_table_details, +) from datasette.write_sql import QueryWriteRejected from . import Context @@ -80,7 +85,7 @@ class DatabaseView(View): async def get(self, request, datasette): format_ = request.url_vars.get("format") or "html" - await datasette.refresh_schemas() + await datasette._schema_watcher.on_request() db = await datasette.resolve_database(request) database = db.name @@ -497,25 +502,39 @@ async def get_tables(datasette, request, db, allowed_dict) -> list[DatabaseTable """ tables = [] table_counts = await db.table_counts(100) - hidden_table_names = set(await db.hidden_table_names()) - all_foreign_keys = await db.get_all_foreign_keys() + # Everything except the row counts comes from the _internal catalog - the + # same catalog allowed_dict was built from - in a few queries, instead of + # three introspection round trips per table (and an O(tables^2) + # hidden-table check) against the database itself + summary = (await catalog_summaries(datasette, [db.name]))[db.name] + db_config = datasette.config.get("databases", {}).get(db.name, {}) + config_hidden = { + t + for t, table_config in db_config.get("tables", {}).items() + if table_config.get("hidden") + } + all_foreign_keys = await catalog_all_foreign_keys(datasette, db.name) + shown = [table for table in table_counts if table in allowed_dict] + details = await catalog_table_details( + datasette, [(db.name, table) for table in shown] + ) - for table in table_counts: - if table not in allowed_dict: + for table in shown: + if table not in all_foreign_keys: + # Not in the catalog (yet), so not in allowed_dict either continue - - table_columns = await db.table_columns(table) + table_details = details[(db.name, table)] tables.append( DatabaseTable( name=table, - columns=table_columns, - primary_keys=await db.primary_keys(table), + columns=table_details["columns"], + primary_keys=table_details["primary_keys"], count=table_counts[table], count_truncated=_table_count_truncated( datasette, db, table, table_counts[table] ), - hidden=table in hidden_table_names, - fts_table=await db.fts_table(table), + hidden=summary.is_hidden(table, config_hidden), + fts_table=table_details["fts_table"], foreign_keys=all_foreign_keys[table], private=allowed_dict[table].private, ) @@ -711,7 +730,7 @@ class QueryView(View): async def get(self, request, datasette): from datasette.app import TableNotFound - await datasette.refresh_schemas() + await datasette._schema_watcher.on_request() db = await datasette.resolve_database(request) database = db.name diff --git a/datasette/views/execute_write.py b/datasette/views/execute_write.py index c4f5e3fe..4de43493 100644 --- a/datasette/views/execute_write.py +++ b/datasette/views/execute_write.py @@ -209,13 +209,6 @@ async def _create_table_template_sql(datasette, db, actor): return None -def _analysis_changes_schema(analysis): - return any( - operation.operation in {"create", "alter", "drop"} - for operation in analysis.operations - ) - - class ExecuteWriteView(BaseView): name = "execute-write" has_json_alternate = False @@ -400,8 +393,10 @@ class ExecuteWriteView(BaseView): status=400, ) - if _analysis_changes_schema(analysis): - await self.ds.refresh_schemas(force=True) + # No refresh_schemas() here: execute_write() already rebuilt this + # database's catalog rows before returning if the SQL changed its + # schema (the SchemaWatcher's post-write check), and a forced + # refresh would stat every attached database on every write if cursor.rowcount == -1: message = "Query executed" diff --git a/datasette/views/index.py b/datasette/views/index.py index f73ee38a..b41ba3f8 100644 --- a/datasette/views/index.py +++ b/datasette/views/index.py @@ -1,5 +1,6 @@ import json +from datasette.database import DatasetteClosedError from datasette.plugins import pm from datasette.utils import ( UNSTABLE_API_MESSAGE, @@ -7,8 +8,15 @@ from datasette.utils import ( add_cors_headers, await_me_maybe, make_slot_function, + sqlite3, + tilde_encode, ) from datasette.utils.asgi import Response +from datasette.utils.catalog import ( + catalog_relationship_counts, + catalog_summaries, + catalog_table_details, +) from datasette.version import __version__ from .base import BaseView @@ -19,6 +27,15 @@ TRUNCATE_AT = 5 # Only attempt counts if database less than this size in bytes: COUNT_DB_SIZE_LIMIT = 100 * 1024 * 1024 +# Row counts need a count(*) per table in every database listed, so the +# index page only runs them when it lists at most this many databases that +# would need counting. Above that it shows table counts only - plus row +# counts that are already known without opening the database (immutable +# databases with inspect data, or whose counts a page computed earlier). +# Everything else on the page comes from the _internal catalog, so the +# index page never opens a database it is not counting. +COUNT_MAX_DATABASES = 20 + class IndexView(BaseView): name = "index" @@ -32,7 +49,14 @@ class IndexView(BaseView): "view-database", request.actor, include_is_private=True ) allowed_databases = [r async for r in db_page.all()] - allowed_db_dict = {r.parent: r for r in allowed_databases} + # One snapshot for the whole request: add_database() and + # remove_database() replace ds.databases rather than changing it, so + # a database removed while this page is being built is still listed + # instead of failing the page with a KeyError + all_databases = self.ds.databases + allowed_db_dict = { + r.parent: r for r in allowed_databases if r.parent in all_databases + } # Group tables by database tables_by_db = {} @@ -44,71 +68,57 @@ class IndexView(BaseView): tables_by_db[t.parent] = {} tables_by_db[t.parent][t.child] = t - databases = [] - # Iterate over allowed databases instead of all databases - for name, allowed_db in allowed_db_dict.items(): - db = self.ds.databases[name] - database_private = allowed_db.private + names = list(allowed_db_dict) + summaries = await catalog_summaries(self.ds, names) + all_counts = await self._table_counts(names, all_databases) + sort_by_relationships = request.args.get("_sort") == "relationships" + need_relationships = { + name for name in names if sort_by_relationships or not all_counts[name] + } + relationship_counts = ( + await catalog_relationship_counts(self.ds, need_relationships) + if need_relationships + else {} + ) + prepared = [] + for name, allowed_db in allowed_db_dict.items(): + summary = summaries[name] # Get allowed tables/views for this database allowed_for_db = tables_by_db.get(name, {}) - - # Get table names from allowed set instead of db.table_names() - table_names = [child_name for child_name in allowed_for_db] - - hidden_table_names = set(await db.hidden_table_names()) - - # Determine which allowed items are views - view_names_set = set(await db.view_names()) + db_config = self.ds.config.get("databases", {}).get(name, {}) + config_hidden = { + t + for t, table_config in db_config.get("tables", {}).items() + if table_config.get("hidden") + } views = [ {"name": child_name, "private": resource.private} for child_name, resource in allowed_for_db.items() - if child_name in view_names_set + if child_name in summary.views ] - - # Filter to just tables (not views) for table processing - table_names = [name for name in table_names if name not in view_names_set] - - # Perform counts only for immutable or DBS with <= COUNT_TABLE_LIMIT tables - table_counts = {} - if not db.is_mutable or db.size < COUNT_DB_SIZE_LIMIT: - table_counts = await db.table_counts(10) - # If any of these are None it means at least one timed out - ignore them all - if any(v is None for v in table_counts.values()): - table_counts = {} - + table_counts = all_counts[name] + relationships = relationship_counts.get(name, {}) tables = {} - for table in table_names: - # Check if table is in allowed set - if table not in allowed_for_db: + for table, resource in allowed_for_db.items(): + if table in summary.views: continue - - table_columns = await db.table_columns(table) tables[table] = { "name": table, - "columns": table_columns, - "primary_keys": await db.primary_keys(table), "count": table_counts.get(table), - "hidden": table in hidden_table_names, - "fts_table": await db.fts_table(table), - "num_relationships_for_sorting": 0, - "private": allowed_for_db[table].private, + "hidden": summary.is_hidden(table, config_hidden), + "num_relationships_for_sorting": ( + relationships.get(table, 0) if name in need_relationships else 0 + ), + "private": resource.private, } - if request.args.get("_sort") == "relationships" or not table_counts: - # We will be sorting by number of relationships, so populate that field - all_foreign_keys = await db.get_all_foreign_keys() - for table, foreign_keys in all_foreign_keys.items(): - if table in tables: - count = len(foreign_keys["incoming"] + foreign_keys["outgoing"]) - tables[table]["num_relationships_for_sorting"] = count - hidden_tables = [t for t in tables.values() if t["hidden"]] visible_tables = [t for t in tables.values() if not t["hidden"]] tables_and_views_truncated = list( sorted( - (t for t in tables.values() if t not in hidden_tables), + visible_tables, key=lambda t: ( t["num_relationships_for_sorting"], t["count"] or 0, @@ -117,7 +127,50 @@ class IndexView(BaseView): reverse=True, )[:TRUNCATE_AT] ) + prepared.append( + ( + name, + allowed_db, + tables, + views, + hidden_tables, + visible_tables, + table_counts, + tables_and_views_truncated, + ) + ) + # Columns, primary keys and FTS tables, for the tables shown only + details = await catalog_table_details( + self.ds, + [(name, t["name"]) for name, *_, truncated in prepared for t in truncated], + ) + + databases = [] + for ( + name, + allowed_db, + tables, + views, + hidden_tables, + visible_tables, + table_counts, + tables_and_views_truncated, + ) in prepared: + db = all_databases[name] + tables_and_views_truncated = [ + { + "name": t["name"], + "columns": details[(name, t["name"])]["columns"], + "primary_keys": details[(name, t["name"])]["primary_keys"], + "count": t["count"], + "hidden": t["hidden"], + "fts_table": details[(name, t["name"])]["fts_table"], + "num_relationships_for_sorting": t["num_relationships_for_sorting"], + "private": t["private"], + } + for t in tables_and_views_truncated + ] # Only add views if this is less than TRUNCATE_AT if len(tables_and_views_truncated) < TRUNCATE_AT: num_views_to_add = TRUNCATE_AT - len(tables_and_views_truncated) @@ -128,7 +181,8 @@ class IndexView(BaseView): "name": name, "hash": db.hash, "color": db.color, - "path": self.ds.urls.database(name), + # From the snapshot: name may have been removed since + "path": self.ds.urls.path(tilde_encode(db.route)), "tables_and_views_truncated": tables_and_views_truncated, "tables_and_views_more": (len(visible_tables) + len(views)) > TRUNCATE_AT, @@ -140,10 +194,14 @@ class IndexView(BaseView): ), "hidden_tables_count": len(hidden_tables), "views_count": len(views), - "private": database_private, + "private": allowed_db.private, } ) + # Leave out databases removed while this page was being built: the + # HTML template links to them by name + databases = [d for d in databases if d["name"] in self.ds.databases] + if as_format: headers = {} if self.ds.cors: @@ -189,3 +247,38 @@ class IndexView(BaseView): "noindex": request.path == "/-/", }, ) + + async def _table_counts(self, names, all_databases): + """``{database: {table: count}}`` - empty for a database whose + counts are skipped or timed out (see COUNT_MAX_DATABASES).""" + counts = {} + to_count = [] + for name in names: + db = all_databases[name] + if not db.is_mutable and db.cached_table_counts is not None: + # Known without opening the database + counts[name] = db.cached_table_counts + else: + counts[name] = {} + to_count.append(db) + if len(to_count) > COUNT_MAX_DATABASES: + return counts + for db in to_count: + if db.is_mutable: + try: + if db.size >= COUNT_DB_SIZE_LIMIT: + continue + except OSError: + # The file has gone: there is nothing to count + continue + try: + table_counts = await db.table_counts(10) + except (DatasetteClosedError, sqlite3.Error): + # Closed or removed while this page was being built, or its + # file deleted or replaced: show it without counts + continue + # If any of these are None it means at least one timed out - ignore them all + if any(v is None for v in table_counts.values()): + table_counts = {} + counts[db.name] = table_counts + return counts diff --git a/datasette/views/special.py b/datasette/views/special.py index 75078e43..bb50435d 100644 --- a/datasette/views/special.py +++ b/datasette/views/special.py @@ -18,6 +18,7 @@ from datasette.utils import ( tilde_encode, ) from datasette.utils.asgi import Forbidden, Response +from datasette.utils.catalog import catalog_summaries, catalog_table_details from .base import BaseView, View @@ -110,16 +111,27 @@ class AutocompleteDebugView(BaseView): scanned = 0 reached_scan_limit = False suggestions = [] - for database_name, db in self.ds.databases.items(): + # The first 100 tables, in database order then by name, from the + # _internal catalog rather than each database's sqlite_master + names = list(self.ds.databases) + candidates = await self.ds.get_internal_database().execute( + """ + select database_name, table_name + from catalog_tables + join json_each(:names) as j on j.value = catalog_tables.database_name + order by j.key, table_name + limit 100 + """, + {"names": json.dumps(names)}, + ) + by_database = {} + for database_name, table_name in candidates.rows: + by_database.setdefault(database_name, []).append(table_name) + for database_name, table_names in by_database.items(): if scanned >= 100 or len(suggestions) >= 5: break - remaining = 100 - scanned - results = await db.execute( - "select name from sqlite_master where type = 'table' order by name limit ?", - [remaining], - ) - for row in results.rows: - table_name = row["name"] + db = self.ds.databases[database_name] + for table_name in table_names: scanned += 1 if scanned >= 100: reached_scan_limit = True @@ -316,7 +328,7 @@ class AllowedResourcesView(BaseView): async def get(self, request): await self.ds.ensure_permission(action="view-instance", actor=request.actor) - await self.ds.refresh_schemas() + await self.ds._schema_watcher.on_request() # Check if user has permissions-debug (to show sensitive fields) has_debug_permission = await self.ds.allowed( @@ -534,7 +546,7 @@ class PermissionRulesView(BaseView): union_sql, union_params, _restriction_sqls = await build_permission_rules_sql( self.ds, actor, action ) - await self.ds.refresh_schemas() + await self.ds._schema_watcher.on_request() db = self.ds.get_internal_database() count_query = f""" @@ -928,24 +940,32 @@ class ApiExplorerView(BaseView): has_json_alternate = False async def example_links(self, request): + # Which databases and tables the actor can see, and their columns, + # come from the _internal catalog: building this page must not open + # every attached database + database_page = await self.ds.allowed_resources("view-database", request.actor) + visible_databases = {r.parent async for r in database_page.all()} + table_page = await self.ds.allowed_resources("view-table", request.actor) + visible_tables = {} + async for resource in table_page.all(): + visible_tables.setdefault(resource.parent, []).append(resource.child) + names = [name for name in self.ds.databases if name in visible_databases] + summaries = await catalog_summaries(self.ds, names) + mutable_tables = [ + (name, table) + for name in names + if self.ds.databases[name].is_mutable + for table in visible_tables.get(name, []) + if table not in summaries[name].views + ] + details = await catalog_table_details(self.ds, mutable_tables) + databases = [] - for name, db in self.ds.databases.items(): - database_visible, _ = await self.ds.check_visibility( - request.actor, - action="view-database", - resource=DatabaseResource(database=name), - ) - if not database_visible: - continue + for name in names: + db = self.ds.databases[name] tables = [] - table_names = await db.table_names() - for table in table_names: - visible, _ = await self.ds.check_visibility( - request.actor, - action="view-table", - resource=TableResource(database=name, table=table), - ) - if not visible: + for table in visible_tables.get(name, []): + if table in summaries[name].views: continue table_links = [] tables.append({"name": table, "links": table_links}) @@ -965,7 +985,8 @@ class ApiExplorerView(BaseView): resource=TableResource(database=name, table=table), actor=request.actor, ): - pks = await db.primary_keys(table) + pks = details[(name, table)]["primary_keys"] + table_columns = details[(name, table)]["columns"] table_links.extend( [ { @@ -976,7 +997,7 @@ class ApiExplorerView(BaseView): "rows": [ { column: None - for column in await db.table_columns(table) + for column in table_columns if column not in pks } ] @@ -999,7 +1020,7 @@ class ApiExplorerView(BaseView): ) for column in ( (["rowid"] if not pks else []) - + await db.table_columns(table) + + table_columns ) } ] @@ -1273,15 +1294,16 @@ class SchemaBaseView(BaseView): has_json_alternate = False - async def get_database_schema(self, database_name, actor): + async def get_database_schema(self, database_name, actor, allowed_table_names=None): """Get schema SQL for a database.""" db = self.ds.databases[database_name] - allowed_tables_page = await self.ds.allowed_resources( - "view-table", actor, parent=database_name - ) - allowed_table_names = { - resource.child async for resource in allowed_tables_page.all() - } + if allowed_table_names is None: + allowed_tables_page = await self.ds.allowed_resources( + "view-table", actor, parent=database_name + ) + allowed_table_names = { + resource.child async for resource in allowed_tables_page.all() + } result = await db.execute( "select tbl_name, sql from sqlite_master where sql is not null" ) @@ -1346,10 +1368,30 @@ class InstanceSchemaView(SchemaBaseView): ) allowed_databases = [r.parent async for r in allowed_databases_page.all()] - # Get schema for each database + # One permission query for every table, not one per database + allowed_tables_page = await self.ds.allowed_resources( + "view-table", request.actor + ) + allowed_tables = {} + async for resource in allowed_tables_page.all(): + allowed_tables.setdefault(resource.parent, set()).add(resource.child) + + # Get schema for each database. This page is a dump of every + # database the actor can see, so it reads each one's sqlite_master + # (the catalog has no index or trigger SQL) - but it skips databases + # with no visible tables or views instead of opening them. schemas = [] for database_name in allowed_databases: - schema = await self.get_database_schema(database_name, request.actor) + if database_name not in self.ds.databases: + continue + visible = allowed_tables.get(database_name, set()) + schema = ( + await self.get_database_schema( + database_name, request.actor, allowed_table_names=visible + ) + if visible + else "" + ) schemas.append({"database": database_name, "schema": schema}) if format_ == "json": diff --git a/datasette/views/table.py b/datasette/views/table.py index 0a908c69..2f412e9a 100644 --- a/datasette/views/table.py +++ b/datasette/views/table.py @@ -1738,7 +1738,7 @@ async def _sort_order(table_metadata, sortable_columns, request, order_by): async def table_view(datasette, request): - await datasette.refresh_schemas() + await datasette._schema_watcher.on_request() with tracer.trace_child_tasks(): response = await table_view_traced(datasette, request) diff --git a/datasette/write_budget.py b/datasette/write_budget.py new file mode 100644 index 00000000..3ab06fc2 --- /dev/null +++ b/datasette/write_budget.py @@ -0,0 +1,376 @@ +"""Bound retained writers and queued writes without changing per-database FIFO. + +The coordinator only admits work and retires idle workers. SQLite connections +are opened, used and closed by their owning writer. A retiring thread keeps +its slot until it has actually exited. Internal catalog writes use their own +single writer, so user saturation cannot prevent catalog progress. +""" + +import collections +import contextlib +import contextvars +import errno +import os +import threading +import time + +in_database_callback = contextvars.ContextVar("in_database_callback", default=False) +RETIRE_WRITER = object() +_STARTING = object() + + +class DatabaseResourceError(RuntimeError): + code = "database_resource_unavailable" + execution_started = False + + +class DatabaseQueueFull(DatabaseResourceError): + code = "database_queue_full" + + +class DatabaseAdmissionTimeout(DatabaseResourceError): + code = "database_admission_timeout" + + +class DatabaseReentrancyError(RuntimeError): + code = "database_reentrant_operation" + + +def check_reentrancy(): + if in_database_callback.get(): + raise DatabaseReentrancyError( + "Schedule dependent database work after the current callback returns" + ) + + +@contextlib.contextmanager +def callback_scope(): + token = in_database_callback.set(True) + try: + yield + finally: + in_database_callback.reset(token) + + +def connect_with_retry(db, connect): + """Retry only opening a connection, never preparation or a callback.""" + from .utils import sqlite3 + + for attempt in range(2): + try: + return connect() + except (OSError, sqlite3.OperationalError) as error: + pressure = isinstance(error, OSError) and error.errno in ( + errno.EMFILE, + errno.ENFILE, + errno.ENOMEM, + ) + # Python 3.10 exposes neither SQLite result constants nor + # sqlite_errorcode. Keep its open-error recovery path working. + code = getattr(error, "sqlite_errorcode", None) + cantopen_code = getattr(sqlite3, "SQLITE_CANTOPEN", None) + cantopen = isinstance(error, sqlite3.OperationalError) and ( + ( + code is not None + and cantopen_code is not None + and code & 0xFF == cantopen_code + ) + or "unable to open database file" in str(error).lower() + ) + if not pressure and (not cantopen or not os.path.exists(db.path)): + raise + if attempt: + if pressure: + raise DatabaseResourceError( + "Unable to open a database connection; try again later" + ) from error + raise + pool = db.ds._read_pool_or_none + if pool is not None: + pool.evict_idle() + # Writers own their connections: ask them to retire rather + # than closing SQLite objects from this thread. + budget = db.ds.__dict__.get("_write_budget_instance") + if budget is not None: + budget.retire_idle() + + +class WriteBudget: + def __init__(self, max_writers, max_pending, timeout_ms): + self.max_writers = max_writers + self.max_pending = max_pending + self.timeout = timeout_ms / 1000 + self._condition = threading.Condition() + self._active = {} # db -> starting marker or thread (including retiring) + self._pending = {} # task -> (db, monotonic deadline) + self._ready = collections.OrderedDict() + self._busy = set() + self._retirement_requested = set() + self._thread = None + self._inline_lock = threading.RLock() + self._inline = collections.OrderedDict() + self.stats = collections.Counter() + + def submit(self, db, task): + # Caller holds db._write_thread_lock. Never take a database lock + # while holding _condition in the coordinator. + with self._condition: + if len(self._pending) >= self.max_pending: + self.stats["rejected"] += 1 + raise DatabaseQueueFull( + "Too many queued database writes; try again later" + ) + self._pending[task] = (db, time.monotonic() + self.timeout) + db._write_queue.put(task) + self._ready.setdefault(db, None) + self.stats["peak_pending"] = max( + self.stats["peak_pending"], len(self._pending) + ) + if self._thread is None: + thread = threading.Thread( + target=self._run, name="datasette-write-scheduler", daemon=True + ) + try: + thread.start() + except BaseException: + self._pending.pop(task) + self._remove_queued(db, task) + self._ready.pop(db, None) + raise + self._thread = thread + self._condition.notify_all() + + @staticmethod + def _remove_queued(db, task): + # Remove expired entries instead of retaining unbounded tombstones. + q = db._write_queue + with q.mutex: + try: + q.queue.remove(task) + except ValueError: + pass # A worker took it; take() will reject the expired task. + q.not_full.notify() + + @staticmethod + def _fail(task, error): + from .database import _deliver_write_result + + _deliver_write_result(task, None, error) + + def take(self, db, task): + with self._condition: + item = self._pending.pop(task, None) + self._condition.notify_all() + if item is None: + return False + if time.monotonic() >= item[1]: + self.stats["expired"] += 1 + expired = True + else: + self._busy.add(db) + expired = False + if expired: + self._fail( + task, DatabaseAdmissionTimeout("Write expired before execution started") + ) + return not expired + + def completed(self, db): + with self._condition: + self._busy.discard(db) + self._condition.notify_all() + + def should_yield(self, db): + # Named in-memory databases need a keeper connection. Until a + # separate keeper is introduced their writer slot is pinned. + if db.is_memory or db._closed: + return False + with self._condition: + return any( + other is not db and other not in self._active + for other, _ in self._pending.values() + ) + + def wake(self): + with self._condition: + self._condition.notify_all() + + def drain_database(self, db, timeout): + """Wait for accepted work and its owning threads, including admission. + + The caller has stopped submissions and queued the shutdown sentinel. + Normal admission deadlines still apply. If shutdown itself times out, + only callbacks that have not started are cancelled; running callbacks + keep ownership of their connections until they return. + """ + with self._condition: + drained = self._condition.wait_for( + lambda: db not in self._active + and not any(owner is db for owner, _ in self._pending.values()), + timeout=timeout, + ) + if not drained: + self.cancel_database(db) + else: + with self._inline_lock: + self._inline.pop(db, None) + return drained + + def cancel_database(self, db): + from .database import DatasetteClosedError + + cancelled = [] + with self._condition: + for task, (owner, _) in list(self._pending.items()): + if owner is db: + self._pending.pop(task) + self._remove_queued(db, task) + cancelled.append(task) + self._ready.pop(db, None) + self._condition.notify_all() + for task in cancelled: + self._fail( + task, DatasetteClosedError("Database closed before write started") + ) + with self._inline_lock: + self._inline.pop(db, None) + + def _run(self): + while True: + expired = [] + start = [] + retire = [] + with self._condition: + for db, thread in list(self._active.items()): + if thread is not _STARTING and not thread.is_alive(): + del self._active[db] + self._retirement_requested.discard(db) + self._busy.discard(db) + self._ready.pop(db, None) + self._ready[db] = None + now = time.monotonic() + for task, (db, deadline) in list(self._pending.items()): + if now >= deadline: + self._pending.pop(task) + self._remove_queued(db, task) + self.stats["expired"] += 1 + expired.append(task) + pending_dbs = {db for db, _ in self._pending.values()} + for db in list(self._ready): + if db not in pending_dbs: + del self._ready[db] + elif ( + db not in self._active and len(self._active) < self.max_writers + ): + del self._ready[db] + self._active[db] = _STARTING + start.append(db) + self.stats["peak_writers"] = max( + self.stats["peak_writers"], len(self._active) + ) + if any(db not in self._active for db in pending_dbs): + for db in self._active: + if ( + db not in self._busy + and not db.is_memory + and db not in self._retirement_requested + ): + self._retirement_requested.add(db) + retire.append(db) + if not self._pending and not self._active: + self._ready.clear() + self._thread = None + stop = True + else: + stop = False + # close() may be waiting for a starting or retiring worker, + # not just for callbacks to complete. + self._condition.notify_all() + for task in expired: + self._fail( + task, + DatabaseAdmissionTimeout("Write expired before execution started"), + ) + if stop: + return + for db in start: + try: + with db._write_thread_lock: + # close() rejects new submissions, but accepted work + # remains eligible for this normally budgeted slot. + # Its shutdown sentinel follows the accepted tasks. + db._start_write_thread() + thread = db._write_thread + with self._condition: + if thread is None: + self._active.pop(db, None) + else: + self._active[db] = thread + except Exception as error: # noqa: BLE001 + failures = [] + with self._condition: + self._active.pop(db, None) + for task, (owner, _) in list(self._pending.items()): + if owner is db: + self._pending.pop(task) + self._remove_queued(db, task) + failures.append(task) + for task in failures: + self._fail(task, error) + for db in retire: + with db._write_thread_lock: + if db._write_thread is not None and not db._closed: + db._write_queue.put(RETIRE_WRITER) + # Timed wake also observes completed thread retirement and + # expires queued work while all slots are pinned or busy. + with self._condition: + self._condition.wait(0.01) + # No completed database is retained in this long-lived frame. + del start, retire, expired, pending_dbs + db = thread = task = owner = failures = None + + @contextlib.contextmanager + def inline(self, db): + """Non-threaded mode: bounded idle connections, no helper threads.""" + check_reentrancy() + with self._inline_lock: + for other in list(self._inline): + if other._write_connection is None or other._closed: + self._inline.pop(other) + if db not in self._inline and len(self._inline) >= self.max_writers: + victim = next((d for d in self._inline if not d.is_memory), None) + if victim is None: + raise DatabaseResourceError( + "Writer capacity is held by in-memory databases" + ) + self._inline.pop(victim) + conn, victim._write_connection = victim._write_connection, None + victim._forget_connection(conn) + self._inline.pop(db, None) + self._inline[db] = None + yield + + def retire_idle(self): + with self._condition: + candidates = [ + db + for db in self._active + if db not in self._busy + and not db.is_memory + and db not in self._retirement_requested + ] + self._retirement_requested.update(candidates) + for db in candidates: + with db._write_thread_lock: + if db._write_thread is not None: + db._write_queue.put(RETIRE_WRITER) + + def snapshot(self): + with self._condition: + return { + "writers": len(self._active), + "pending": len(self._pending), + "max_writers": self.max_writers, + "max_pending": self.max_pending, + "stats": dict(self.stats), + } diff --git a/docs/changelog.rst b/docs/changelog.rst index 2775cbc5..37b46ee1 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -4,6 +4,22 @@ Changelog ========= +.. _changelog_unreleased: + +Unreleased +---------- + +Connection management +~~~~~~~~~~~~~~~~~~~~~ + +Datasette now pools read connections and bounds writer threads and queued writes across the instance, allowing it to serve many database files without keeping each one open. (:pr:`2989`) + +- New settings default to :ref:`32 pooled file read connections `, :ref:`8 user database writers `, :ref:`256 pending writes ` and a :ref:`five-second write queue timeout `. Connections and writer threads retire after :ref:`30 seconds idle ` or under capacity pressure. These limits do not impose a hard cap on total memory or file descriptors. Named in-memory databases retain writer slots to preserve their contents. +- Index and database listings use the internal schema catalog. The new :ref:`schema watch settings ` control detection of changes made outside Datasette. Immutable table row counts are now calculated when needed instead of at startup. +- Connections, cursors, blob handles and dump iterators supplied to plugin callbacks expire when the callback finishes. Using an expired resource, or returning a live resource from ``execute_fn()``, raises ``datasette.database.ConnectionLeaseError``. Fetch results inside the callback and use :ref:`plugin_hook_prepare_connection` for setup needed on every managed connection. See :ref:`database_connection_leases` for details. +- Full write queues and admission timeouts now raise explicit :ref:`resource errors `, with HTTP 503 responses for requests. The new :ref:`plugin_hook_write_task_completed` hook reports completion or failure of accepted non-blocking writes. Plugins that create background tasks should observe their results so admission failures are handled. +- Graceful shutdown drains accepted writes subject to their admission deadlines and a per-database shutdown timeout. Running callbacks retain ownership of their connections until they finish; see :ref:`database_close`. + .. _v1_0_a41: 1.0a41 (2026-09-24) diff --git a/docs/cli-reference.rst b/docs/cli-reference.rst index db903e8e..7fb39c21 100644 --- a/docs/cli-reference.rst +++ b/docs/cli-reference.rst @@ -251,6 +251,26 @@ 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 - 0 for no limit + (default=32) + max_write_connections Maximum concurrent user database writers + (default=8) + max_pending_writes Maximum queued user writes across all databases + (default=256) + write_queue_timeout_ms Maximum wait before a queued write starts + (default=5000) + connection_idle_timeout_ms Close read connections, and stop write threads, + that have been idle for this many milliseconds - + 0 to keep them open (default=30000) + schema_watch_interval_ms How often to check external database files for + schema changes - 0 to disable polling + (default=1000) + default_schema_watch Schema watch mode for database files and named + in-memory databases without a mode of their own: + external (poll for changes made outside + Datasette) or owned (only Datasette changes them) + (default=external) 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..4a17df02 100644 --- a/docs/internals.rst +++ b/docs/internals.rst @@ -1315,8 +1315,8 @@ Removes the column type assignment for the specified column. .. _datasette_add_database: -.add_database(db, name=None, route=None) ----------------------------------------- +.add_database(db, name=None, route=None, schema_watch=None) +----------------------------------------------------------- ``db`` - datasette.database.Database instance The database to be attached. @@ -1327,6 +1327,9 @@ Removes the column type assignment for the specified column. ``route`` - string, optional This will be used in the URL path. If not specified, it will default to the same thing as the ``name``. +``schema_watch`` - string, optional + How Datasette keeps its catalog of this database's tables current: ``"owned"`` (only Datasette's own writes change the schema), ``"external"`` (other code may change it too - another process, or a plugin's own connection - so it is polled) or ``"immutable"``. The default for a database file or a named in-memory database is the :ref:`setting_default_schema_watch` setting, ``"external"`` unless configured. Pass ``"owned"`` if every change to the database goes through Datasette's write methods: it is then never polled. A ``schema_watch`` setting for this database in ``datasette.yaml`` takes precedence. See :ref:`setting_schema_watch_interval_ms`. + The ``datasette.add_database(db)`` method lets you add a new database to the current Datasette instance. The ``db`` parameter should be an instance of the ``datasette.database.Database`` class. For example: @@ -1342,7 +1345,7 @@ The ``db`` parameter should be an instance of the ``datasette.database.Database` ) ) -This will add a mutable database and serve it at ``/my-new-database``. +This will add a mutable database and serve it at ``/my-new-database``. If the file does not exist yet, the first write to the database creates it. Use ``is_mutable=False`` to add an immutable database. @@ -1397,7 +1400,7 @@ This removes a database that has been previously added. ``name=`` is the unique .close() -------- -Release all resources held by this ``Datasette`` instance. This calls :ref:`database_close` on every attached database (including the internal database), shuts down the thread pool executor used to run SQL queries, and unlinks the temporary file used to back the internal database if one was created. +Release resources held by this ``Datasette`` instance. This stops new submissions to every user database before draining their accepted writes, then closes the internal database and shuts down the thread pool executor used to run SQL queries. Each database follows the deadlines and cleanup rules in :ref:`database_close`, including deferring cleanup of connections and temporary files still owned by running write callbacks. ``close()`` is synchronous, idempotent and one-way: after a call to ``close()`` any attempt to use the Datasette instance to execute SQL will raise a ``datasette.database.DatasetteClosedError`` exception. A closed ``Datasette`` cannot be reopened — callers that need a fresh instance should construct a new one. @@ -1797,10 +1800,10 @@ Application lifecycle Datasette guarantees a fixed sequence of events between the moment a ``Datasette`` instance is constructed and the moment its resources are released: 1. ``Datasette(...)`` — the constructor runs synchronously and does not run plugin hooks. -2. **Startup** — ``await datasette.invoke_startup()`` runs once: it populates the internal database's catalog of table schemas (:ref:`internals_internal`), loads canned queries and column type configuration, then calls every registered :ref:`plugin_hook_startup` hook, in plugin registration order. When Datasette is being served, table-count precomputation for immutable databases runs immediately before this, as part of the same startup sequence. +2. **Startup** — ``await datasette.invoke_startup()`` runs once, even if several event loops call it at the same time (the others wait for it to finish): it populates the internal database's catalog of table schemas (:ref:`internals_internal`), loads canned queries and column type configuration, then calls every registered :ref:`plugin_hook_startup` hook, in plugin registration order. Immutable table counts are calculated when first requested, unless supplied using :ref:`performance_inspect`. 3. **Background-task launch** — once *every* ``startup`` hook has finished (not before), every task registered with :ref:`datasette_add_background_task` — by any plugin — is launched. A task registered by one plugin's ``startup`` hook can safely depend on state set up by another plugin's ``startup`` hook, because launch only happens after the whole round of hooks completes. 4. **Serving** — the instance handles requests (or, for headless or CLI use, does whatever the embedding program does with it). -5. **Shutdown** — triggered by the ASGI ``lifespan.shutdown`` event (Ctrl-C, ``SIGTERM``) or the end of a ``datasette serve`` process: every :ref:`plugin_hook_shutdown` hook runs first, while background tasks are still alive, so a plugin can tell its own task to wind down gracefully; every still-running background task is then cancelled and given a five-second grace period to actually stop; finally every database connection is released via :ref:`datasette_close`. +5. **Shutdown** — triggered by the ASGI ``lifespan.shutdown`` event (Ctrl-C, ``SIGTERM``) or the end of a ``datasette serve`` process: every :ref:`plugin_hook_shutdown` hook runs first, while background tasks are still alive, so a plugin can tell its own task to wind down gracefully; every still-running background task is then cancelled and given a five-second grace period to actually stop; finally database resources are cleaned up via :ref:`datasette_close`. .. admonition:: Startup hooks run on the event loop that serves requests @@ -2051,6 +2054,66 @@ Database class Instances of the ``Database`` class can be used to execute queries against attached SQLite databases, and to run introspection against their schemas. +.. _database_connections: + +Connections and threads +----------------------- + +Datasette is designed to serve anything from one database to thousands of database files from one process, without holding connections open to all of them. You do not normally need to know how connections are managed, but plugins that use the connection-level APIs (:ref:`database_execute_fn`, :ref:`database_execute_write_fn` and :ref:`database_execute_isolated_fn`) should follow the rule in :ref:`database_connection_leases`. + +**Reads** run on a shared pool of :ref:`setting_num_sql_threads` threads. Each read borrows a connection from a pool for the duration of one query or one ``execute_fn()`` callback, then returns it. At most :ref:`setting_max_open_connections` pooled read connections are open at once across all databases; when a new one is needed the least recently used idle connection is closed first, and connections idle for longer than :ref:`setting_connection_idle_timeout_ms` are closed. Read connections to the database files Datasette serves are read-only. + +**Writes** to each database are run one at a time, in the order they were submitted, by a write thread dedicated to that database and holding its single write connection. At most :ref:`setting_max_write_connections` user writers exist at once. The thread stops, closing its connection, after :ref:`setting_connection_idle_timeout_ms` without writes, or after a callback when another database needs its slot; the next write waits for capacity and starts a new one. Queued writes are bounded by :ref:`setting_max_pending_writes` and :ref:`setting_write_queue_timeout_ms`. The internal catalog uses its own writer. + +**Isolated functions** (:ref:`database_execute_isolated_fn`) get a connection of their own that is opened and closed for that one call. + +**Schema changes** made by Datasette's own writes are detected as soon as the write finishes, and databases that other processes may change are polled - see :ref:`setting_schema_watch_interval_ms`. If a database file is replaced (for example by renaming a new file over it) or deleted, connections to the old file are discarded. Once Datasette has seen a database file exist it will not create it again: writes to a deleted file fail with an error rather than silently creating a new empty database. A database added with :ref:`datasette_add_database` whose file does not exist yet is created by its first write. + +**Pages that cover many databases** - the index page, ``allowed_resources()`` listings and their derived-table permission rules, ``/-/api`` and the like - read what they need about each database's tables from the :ref:`catalog tables in the internal database ` instead of opening every database. Plugins that summarise many databases should do the same: querying ``catalog_tables``, ``catalog_columns`` or ``catalog_foreign_keys`` avoids opening each database, while calling methods such as ``db.table_names()`` on every database needs a connection to each of them. Catalog queries still cost time and memory according to the data they process, and plugins must apply appropriate permission checks before exposing catalog data to an actor. Use the live methods on a single :ref:`Database ` when you need the database's own current answer. + +With :ref:`setting_num_sql_threads` set to ``0`` (for example in Pyodide) there are no threads: reads and writes run on the event loop, read connections are still pooled and closed when idle, and retained write connections are bounded by :ref:`setting_max_write_connections`. Idle file write connections can be evicted when another database needs capacity; named in-memory connections are retained to preserve their contents. + +Every new connection is set up by the :ref:`plugin_hook_prepare_connection` plugin hook, except isolated connections, the brief connections used to check a file's schema version, and connections to the internal database. Because connections are opened and closed as needed, anything set up on a connection - functions, temporary tables, ``ATTACH``, ``PRAGMA`` settings - only lasts as long as that connection, and is not seen by other connections to the same database. Use ``prepare_connection`` for anything every connection needs. + +.. _database_connection_leases: + +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, 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. + +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. + +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. + +.. _database_resource_errors: + +Resource errors and retries +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The exceptions in ``datasette.write_budget`` distinguish capacity failures before execution from errors raised after a callback has started: + +``DatabaseResourceError`` + Base class for resource admission failures. Its ``code`` is ``"database_resource_unavailable"`` and ``execution_started`` is ``False``. Reads raise this with the message ``"Timed out waiting for a read connection"`` if they cannot obtain pooled connection capacity within five seconds. Connection opening can also raise it for resource exhaustion before preparation or execution starts. +``DatabaseQueueFull`` + A subclass of ``DatabaseResourceError`` with ``code="database_queue_full"`` and message ``"Too many queued database writes; try again later"``. The pending write limit was reached, so this submission was not accepted. +``DatabaseAdmissionTimeout`` + A subclass of ``DatabaseResourceError`` with ``code="database_admission_timeout"`` and message ``"Write expired before execution started"``. A queued write was accepted but expired before execution started. +``DatabaseReentrancyError`` + A ``RuntimeError`` with ``code="database_reentrant_operation"``. A callback attempted a nested Datasette database operation that could deadlock with finite capacity. Move the dependent operation outside the callback instead of retrying it in place. + +Resource admission errors become HTTP 503 responses with ``Retry-After: 1``; JSON responses include ``code`` and ``execution_started: false``. A caller can retry the database operation after an admission error because it has not started. This guarantee applies to that database operation, not to earlier work in the same request. Plugins that create background tasks must observe their results and handle errors; an exception in an unobserved task cannot produce a useful error response for its caller. + +Other errors do not carry that guarantee. A transaction may already have committed when catalog refresh or event dispatch fails. Cancelling a calling request after execution starts does not cancel the write. Datasette never automatically retries a callback; only connection opening can be retried once, before preparation hooks or the callback run. + .. _database_constructor: Database(ds, path=None, is_mutable=True, is_memory=False, memory_name=None, is_temp_disk=False) @@ -2160,7 +2223,7 @@ The ``Results`` object also has the following properties and methods: await db.execute_fn(fn) ----------------------- -Executes a given callback function against a read-only database connection running in a thread. The function will be passed a SQLite connection, and the return value from the function will be returned by the ``await``. +Executes a given callback function against a read connection, in one of the SQL threads. The function will be passed a SQLite connection, and the return value from the function will be returned by the ``await``. Read connections to the database files Datasette serves are opened read-only (the internal database is an exception). Example usage: @@ -2174,6 +2237,10 @@ Example usage: version = await db.execute_fn(get_version) +The connection is borrowed from a pool for the duration of the call (see :ref:`database_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 that is still reading from the connection, raises the same error: fetch the rows you need before returning. See :ref:`database_connection_leases`. + +Calls to ``execute_fn()`` are not guaranteed to use the same connection, so temporary tables, ``ATTACH`` statements and functions created inside one call may not be visible in the next. + 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: @@ -2259,6 +2326,10 @@ This method works like ``.execute_write()``, but instead of a SQL statement you The function can then perform multiple actions, safe in the knowledge that it has exclusive access to the single writable connection for as long as it is executing. +The connection is only valid while your function is running: using it after the function has returned raises ``datasette.database.ConnectionLeaseError`` - see :ref:`database_connection_leases`. The write connection is closed after :ref:`setting_connection_idle_timeout_ms` without writes, so temporary tables, ``ATTACH`` statements and functions created by one write function may be gone by the next. Every new write connection has ``PRAGMA recursive_triggers`` turned on (before any :ref:`plugin_hook_prepare_connection` hook runs), which is what ``sqlite-utils`` sets and what the triggers it creates to keep full-text search indexes current rely on for ``INSERT OR REPLACE``. + +If your function raises an exception, ``execute_write_fn()`` raises it too and the transaction is rolled back. An exception that is not a subclass of ``Exception``, such as ``SystemExit``, is raised as a ``RuntimeError`` whose ``__cause__`` is the original exception. + Like ``execute_fn()``, the call is traced as a ``db.query`` OpenTelemetry span carrying ``datasette.callback`` rather than ``db.query.text``, above the write-queue spans - see :ref:`internals_telemetry`. A named function gives the span a readable identity; a lambda reports ````. .. warning:: @@ -2283,6 +2354,8 @@ For example: except Exception as e: print("An error occurred:", e) +With ``transaction=False``, finish any transaction you start before returning. An unfinished transaction is rolled back and raises ``ConnectionLeaseError``. A connection can be retired between callbacks, so transaction state and connection-local state must not span callbacks. Do not synchronously wait for another Datasette database operation inside a callback: finite connection limits can deadlock that pattern. Straightforward nested calls raise ``DatabaseReentrancyError``; schedule dependent operations after the callback has returned. + Your function can optionally accept a ``track_event`` parameter in addition to ``conn``. If it does, it will be passed a callable that can be used to queue events for dispatch after the write transaction commits successfully. Events queued this way are discarded if the write raises an exception. .. code-block:: python @@ -2337,14 +2410,14 @@ For example, archive an article and record the change in an audit log: await database.execute_write_fn(archive_article) -If you specify ``block=False`` the method becomes fire-and-forget, queueing your function to be executed and then allowing your code after the call to ``.execute_write_fn()`` to continue running while the underlying thread waits for an opportunity to run your function. A UUID representing the queued task will be returned. Any exceptions in your code will be silently swallowed. +With ``block=False``, a UUID is returned when the task is accepted into the queue; this does not mean that the write has succeeded. A full queue raises ``DatabaseQueueFull`` immediately. Later failures are logged, and accepted tasks report success or failure through :ref:`plugin_hook_write_task_completed`, including expiration before execution. Delivery requires the submitting event loop to remain running and is not durable. With :ref:`setting_num_sql_threads` set to ``0``, the write runs inline before the UUID is returned and failures can raise directly from the call. See :ref:`database_resource_errors` before implementing retries. .. _database_execute_isolated_fn: await db.execute_isolated_fn(fn) -------------------------------- -This method works is similar to :ref:`execute_write_fn() ` but executes the provided function in an entirely isolated SQLite connection, which is opened, used and then closed again in a single call to this method. +This method works is similar to :ref:`execute_write_fn() ` but executes the provided function in an entirely isolated SQLite connection, which is opened, used and then closed again in a single call to this method. As with the other methods, the connection is only valid inside the function - see :ref:`database_connection_leases`. The :ref:`prepare_connection() ` plugin hook is not executed against this connection. @@ -2354,7 +2427,7 @@ Running ``VACUUM`` using this method also ensures it won't trigger incorrect :cl Plugins can also use this method to load potentially dangerous SQLite extensions, use them to perform an operation and then have them safely unloaded at the end of the call, without risk of exposing them to other connections. -Functions run using ``execute_isolated_fn()`` share the same queue as ``execute_write_fn()``, which guarantees that no writes can be executed at the same time as the isolated function is executing. +Functions run using ``execute_isolated_fn()`` share the same queue as ``execute_write_fn()``, which guarantees that no writes can be executed at the same time as the isolated function is executing. For immutable databases, which have no write queue, the function runs on a read-only connection in one of the SQL threads. The return value of the function will be returned by this method. Any exceptions raised by the function will be raised out of the ``await`` line as well. @@ -2363,7 +2436,11 @@ The return value of the function will be returned by this method. Any exceptions db.close() ---------- -Release all resources held by this ``Database`` instance. This shuts down the background write thread (if one was started by a previous call to :ref:`database_execute_write_fn` or similar), closes the write connection, and closes any cached read connections. +Stop accepting new work and release resources held by this ``Database`` instance. Accepted writes continue in FIFO order through normal writer admission, including tasks still waiting for their first writer. Their existing :ref:`admission deadlines ` remain in effect, so accepted writes can still expire before starting. + +Shutdown waits up to ten seconds per database for writers to drain. At that deadline, queued user writes that have not started are cancelled with ``DatasetteClosedError``. A callback already running is not interrupted: its thread keeps ownership of its connection and closes it when the callback finishes. Temporary database files are retained until their owning writer exits. A shutdown timeout is reported to standard error. + +Running read queries are waited for separately before their connections are closed, so the ten-second writer deadline is not a deadline for the entire ``close()`` call. Database shutdown does not make the write queue durable; process termination can still prevent queued work and completion notifications from finishing. After ``db.close()`` has been called, any further call to :ref:`database_execute`, :ref:`database_execute_fn`, :ref:`database_execute_write`, :ref:`database_execute_write_fn`, :ref:`database_execute_write_many`, :ref:`database_execute_write_script` or :ref:`database_execute_isolated_fn` will raise a ``datasette.database.DatasetteClosedError`` exception. @@ -2625,7 +2702,7 @@ A request to a table page produces a span named, in full:: No attributes. ``db.write.queue_wait`` - Time a write spent waiting in its database's write queue. For ``block=True``, this is a child of ``db.query``. For ``block=False``, it is a root span linked to the span that queued the write, since the write can outlive that request. + Time a write spent waiting in its database's write queue. For ``block=True``, this is a child of ``db.query``. For ``block=False``, it is a root span linked to the span that queued the write, since the write can outlive that request. For the first write after the write thread was stopped for being idle, this includes opening and preparing a new write connection. No attributes. @@ -2673,7 +2750,7 @@ This reference is also generated from ``datasette/telemetry_registry.py``: - ``error.type`` *(optional)* - The exception class name for a failed operation. On HTTP spans, also set to the status code as a string for 5xx responses. A 4xx response alone does not set this attribute or an error status. ``datasette.write.queue_wait`` - Histogram, unit ``s``. Time each write waited in its database's write queue. + Histogram, unit ``s``. Time each write waited in its database's write queue, including opening a new write connection for the first write after the write thread was stopped for being idle. Bucket boundaries: ``0.0001``, ``0.0005``, ``0.001``, ``0.005``, ``0.01``, ``0.05``, ``0.1``, ``0.5``, ``1``, ``5``, ``10``. @@ -2706,14 +2783,14 @@ This reference is also generated from ``datasette/telemetry_registry.py``: - ``db.namespace`` - Name of the database being queried. ``datasette.write.queue_depth`` - Observable gauge, unit ``{write}``. Writes waiting for a database's single write thread. Increasing ``num_sql_threads`` does not increase write concurrency. Not reported for databases that have never been written to. + Observable gauge, unit ``{write}``. Writes waiting for a database's write thread. Each database has at most one, so increasing ``num_sql_threads`` does not increase write concurrency; it stops after ``connection_idle_timeout_ms`` without writes. Not reported for databases that have never been written to. Attributes: - ``db.namespace`` - Name of the database being queried. ``datasette.connections.open`` - Observable gauge, unit ``{connection}``. Open SQLite connections managed by Datasette. + Observable gauge, unit ``{connection}``. Open SQLite connections managed by Datasette: pooled read connections (capped across databases by ``max_open_connections``), the write connection while the database's write thread is running, and isolated connections in use. Short-lived connections used to check for schema changes are not counted. Attributes: @@ -2769,7 +2846,7 @@ No token, cookie, or hidden form field is needed. Any ``
`` i Datasette's internal database ============================= -Datasette maintains an "internal" SQLite database used for configuration, caching, and storage. Plugins can store configuration, settings, and other data inside this database. By default, Datasette will use a temporary in-memory SQLite database as the internal database, which is created at startup and destroyed at shutdown. +Datasette maintains an "internal" SQLite database used for configuration, caching, and storage. Plugins can store configuration, settings, and other data inside this database. By default, Datasette will use a temporary file-backed SQLite database as the internal database, which is created at startup and destroyed at shutdown. To persist internal data across Datasette instances, use the ``--internal`` option to specify the path to a SQLite database: @@ -2784,7 +2861,7 @@ You can also set the ``DATASETTE_INTERNAL`` environment variable to specify this export DATASETTE_INTERNAL=/path/to/internal.db datasette mydatabase.db -Datasette maintains tables called ``catalog_databases``, ``catalog_tables``, ``catalog_views``, ``catalog_columns``, ``catalog_indexes``, ``catalog_foreign_keys`` with details of the attached databases and their schemas. These tables should not be considered a stable API - they may change between Datasette releases. +Datasette maintains tables called ``catalog_databases``, ``catalog_tables``, ``catalog_views``, ``catalog_columns``, ``catalog_indexes``, ``catalog_foreign_keys`` with details of the attached databases and their schemas. ``catalog_tables.type`` is the type ``PRAGMA table_list`` reports for each table (``table``, ``virtual`` or ``shadow``), or ``NULL`` if unknown. These tables should not be considered a stable API - they may change between Datasette releases. Metadata is stored in tables ``metadata_instance``, ``metadata_databases``, ``metadata_resources`` and ``metadata_columns``. Plugins can interact with these tables via the :ref:`get_*_metadata() and set_*_metadata() methods `. @@ -2796,7 +2873,7 @@ Plugin authors are asked to practice good etiquette when using the internal data 1. Use a unique prefix when creating tables, indices, and triggers in the internal database. If your plugin is called ``datasette-xyz``, then prefix names with ``datasette_xyz_*``. 2. Avoid long-running write statements that may stall or block other plugins that are trying to write at the same time. -3. Use temporary tables or shared in-memory attached databases when possible. +3. Use temporary tables or shared in-memory attached databases when possible. Temporary tables only exist on the connection that created them, and connections are not kept forever (see :ref:`database_connections`), so create and use them inside a single function call. 4. Avoid implementing features that could expose private data stored in the internal database by other plugins. .. _internals_internal_schema: @@ -2826,12 +2903,12 @@ The internal database schema is as follows: path TEXT, is_memory INTEGER, schema_version INTEGER - ); + , fingerprint TEXT); CREATE TABLE catalog_tables ( database_name TEXT, table_name TEXT, rootpage INTEGER, - sql TEXT, + sql TEXT, type TEXT, PRIMARY KEY (database_name, table_name), FOREIGN KEY (database_name) REFERENCES catalog_databases(database_name) ); diff --git a/docs/pages.rst b/docs/pages.rst index 65c03a49..de2e57bd 100644 --- a/docs/pages.rst +++ b/docs/pages.rst @@ -23,6 +23,10 @@ Add ``/.json`` to the end of the URL for the JSON version of the underlying data The index page can also be accessed at ``/-/``, useful for if the default index page has been replaced using an :ref:`index.html custom template `. The ``/-/`` page will always render the default Datasette ``index.html`` template. +.. _IndexView_row_counts: + +The index page reads table names, columns, primary keys, hidden tables and foreign key relationships from the :ref:`internal database catalog `, so it does not need to open each attached database. Row counts are the exception: they need a ``count(*)`` per table. The index page shows row counts when it lists at most 20 databases that would need counting, each smaller than 100MB, and the counts for a database are dropped if any of them takes longer than 10ms. With more databases than that it shows the number of tables in each database, sorts the most interesting tables first by their number of foreign key relationships, and only includes row counts that are already known without opening a database: counts for :ref:`immutable databases ` from ``--inspect-file`` or calculated by an earlier page. Visit a database's page for its row counts. + .. _DatabaseView: Database diff --git a/docs/performance.rst b/docs/performance.rst index 4427757c..3cf2e123 100644 --- a/docs/performance.rst +++ b/docs/performance.rst @@ -22,14 +22,14 @@ To open a file in immutable mode pass it to the datasette command using the ``-i datasette -i data.db -When you open a file in immutable mode like this Datasette will also calculate and cache the row counts for each table in that database when it first starts up, further improving performance. +When you open a file in immutable mode like this Datasette will also calculate and cache the row counts for each table in that database the first time a page needs them, further improving performance. (Earlier versions counted every immutable database at startup.) .. _performance_inspect: Using "datasette inspect" ------------------------- -Counting the rows in a table can be a very expensive operation on larger databases. In immutable mode Datasette performs this count only once and caches the results, but this can still cause server startup time to increase by several seconds or more. +Counting the rows in a table can be a very expensive operation on larger databases. In immutable mode Datasette performs this count only once and caches the results, but this can still make the first page that shows those counts take several seconds or more. If you know that a database is never going to change you can precalculate the table row counts once and store then in a JSON file, then use that file when you later start the server. @@ -37,7 +37,7 @@ To create a JSON file containing the calculated row counts for a database, use t datasette inspect data.db --inspect-file=counts.json -Then later you can start Datasette against the ``counts.json`` file and use it to skip the row counting step and speed up server startup:: +Then later you can start Datasette against the ``counts.json`` file and use it to skip the row counting step entirely:: datasette -i data.db --inspect-file=counts.json diff --git a/docs/plugin_hooks.rst b/docs/plugin_hooks.rst index 3efff7a4..e7033afb 100644 --- a/docs/plugin_hooks.rst +++ b/docs/plugin_hooks.rst @@ -57,7 +57,9 @@ arguments and can be called like this:: select random_integer(1, 10); -``prepare_connection()`` hooks are not called for Datasette's :ref:`internal database `. +``prepare_connection()`` hooks are not called for Datasette's :ref:`internal database `, :ref:`isolated callbacks ` or the brief connections used only to check a file's schema version. + +For pooled read connections, managed write connections and the short-lived connections used to rebuild the schema catalog, the hook runs once when each connection is opened. Datasette closes connections that have been idle for :ref:`setting_connection_idle_timeout_ms` and caps read connections with :ref:`setting_max_open_connections` and writers with :ref:`setting_max_write_connections`, opening new ones as needed, so expect the hook to run many times over the life of a server - keep it fast. Anything a plugin needs on every managed connection (functions, collations, authorizers, ``PRAGMA`` settings) belongs here rather than in a write function, because connections are not shared and do not live forever. Examples: `datasette-jellyfish `__, `datasette-jq `__, `datasette-haversine `__, `datasette-rure `__ @@ -78,7 +80,7 @@ write_wrapper(datasette, database, request, transaction) ``transaction`` - bool ``True`` if the write will be wrapped in a database transaction. -Return a generator function that accepts a ``conn`` argument (a SQLite connection object) and optionally a ``track_event`` argument. The generator should ``yield`` exactly once. Code before the ``yield`` runs before the write function executes; code after the ``yield`` runs after it completes. +Return a generator function that accepts a ``conn`` argument (a SQLite connection object) and optionally a ``track_event`` argument. The generator should ``yield`` exactly once. Like the write function itself, ``conn`` is only valid while the write runs - see :ref:`database_connection_leases`. Code before the ``yield`` runs before the write function executes; code after the ``yield`` runs after it completes. The result of the write function is sent back through the ``yield``, so you can capture it with ``result = yield``. @@ -2602,3 +2604,15 @@ Tokens can then be created and verified using :ref:`datasette.create_token() `__. +.. _setting_max_open_connections: + +max_open_connections +~~~~~~~~~~~~~~~~~~~~ + +.. [[[cog + setting_default(cog, "max_open_connections") +.. ]]] + +Default: ``32`` + +.. [[[end]]] + +Maximum pooled read connections across file-backed databases. When all connections are busy, reads wait up to five seconds for capacity instead of exceeding the cap. Idle connections are evicted before opening a connection to another database. The configured limit is respected even when it is smaller than :ref:`setting_num_sql_threads`. Set this to 0 for no limit. + +Immutable isolated callbacks use fresh connections counted against this same cap. In-memory read connections are exempt, except for the ``_memory`` database under ``--crossdb``, whose connections attach database files and count against the cap. + +This limits connections, not exact file descriptors or RSS. WAL files, cross-database attachments, HTTP sockets, temporary files and plugins consume additional resources. Datasette logs a warning at startup if a conservative estimate of configured connection costs and headroom exceeds the OS file descriptor limit. Read queries waiting for executor threads, result sizes, callback payload sizes and plugin memory remain outside these limits. + +.. _setting_max_write_connections: + +max_write_connections +~~~~~~~~~~~~~~~~~~~~~ + +.. [[[cog + setting_default(cog, "max_write_connections") +.. ]]] + +Default: ``8`` + +.. [[[end]]] + +Maximum simultaneous user database writer threads and their retained connections. Must be greater than zero. Each database preserves FIFO write order. When other databases are waiting, a writer closes its connection and releases its slot after its current callback completes. A retiring writer counts against the limit until its thread exits. Idle writers also release their slots under pressure, even if :ref:`setting_connection_idle_timeout_ms` is 0. + +The internal catalog has a separate writer, and file schema scans are serialized. Mutable isolated callbacks occupy a writer slot without retaining an additional ordinary file connection. Named in-memory writers occupy non-evictable slots to preserve their data; waiting writes can time out if these consume all capacity. Schema catalog scans also use these writers, so serving read-only queries to named in-memory databases can consume writer slots. Allow capacity for those databases plus any file writers that need to run alongside them. In non-threaded mode, the same setting bounds retained writer connections, evicting an idle file connection before opening another. + +.. _setting_max_pending_writes: + +max_pending_writes +~~~~~~~~~~~~~~~~~~ + +.. [[[cog + setting_default(cog, "max_pending_writes") +.. ]]] + +Default: ``256`` + +.. [[[end]]] + +Maximum queued user writes across all databases, excluding callbacks already executing. Must be greater than zero. A full queue immediately raises ``datasette.write_budget.DatabaseQueueFull`` without accepting the write. Callers are not queued waiting for queue space. The internal catalog writer is separate. + +.. _setting_write_queue_timeout_ms: + +write_queue_timeout_ms +~~~~~~~~~~~~~~~~~~~~~~ + +.. [[[cog + setting_default(cog, "write_queue_timeout_ms") +.. ]]] + +Default: ``5000`` + +.. [[[end]]] + +Maximum milliseconds an accepted write can wait before execution starts. Must be greater than zero. Expired tasks are removed and raise ``datasette.write_budget.DatabaseAdmissionTimeout``. Queue-full and admission-timeout errors have ``execution_started=False`` and become HTTP 503 responses with ``Retry-After: 1`` and an error ``code`` in JSON responses. + +After execution starts, cancellation of the calling request does not cancel the write. Datasette never automatically retries a write callback. A connection-open failure can reclaim idle connections and retry opening once, before any preparation hook or callback runs. ``block=False`` calls return a task ID on acceptance; later failures are logged and reported to :ref:`plugin_hook_write_task_completed` while the calling event loop remains running. + +.. _setting_connection_idle_timeout_ms: + +connection_idle_timeout_ms +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. [[[cog + setting_default(cog, "connection_idle_timeout_ms") +.. ]]] + +Default: ``30000`` + +.. [[[end]]] + +Connections that have not been used for this many milliseconds are closed, so a database that stops receiving traffic stops holding file descriptors, memory and threads. + +- **Read connections** in the pool are closed once they have been idle this long. The next read to that database opens a new connection. +- **Write threads**: each database that receives writes gets its own write thread, holding a single write connection, and writes to that database are run one at a time by that thread. A write thread that has had no work for this long closes its connection and exits. The next write to that database starts a new thread and opens a new connection. + +Opening a new connection costs around a millisecond (more if :ref:`plugin_hook_prepare_connection` hooks or extensions such as SpatiaLite have to be loaded again), which is negligible for a database that has been idle for 30 seconds, so the default keeps busy databases warm and releases forgotten ones quickly. + +Anything a function attaches to a connection itself - temporary tables, ``ATTACH`` statements, ``PRAGMA`` settings, functions registered with ``conn.create_function()`` - is lost when that connection is closed, and is never shared between pooled connections. Set that up in the :ref:`plugin_hook_prepare_connection` plugin hook instead, which runs once for every new connection. + +In-memory databases are exempt, since closing their only write connection would discard their contents. With :ref:`setting_num_sql_threads` set to 0 there are no threads to time out, so idle write connections are closed under capacity pressure or when Datasette shuts down; read connections are still closed when idle. + +Set this to 0 to keep read connections open until they are evicted by :ref:`setting_max_open_connections`, and write threads running until capacity pressure or Datasette shutdown:: + + datasette mydatabase.db --setting connection_idle_timeout_ms 0 + +.. _setting_schema_watch_interval_ms: + +schema_watch_interval_ms +~~~~~~~~~~~~~~~~~~~~~~~~ + +.. [[[cog + setting_default(cog, "schema_watch_interval_ms") +.. ]]] + +Default: ``1000`` + +.. [[[end]]] + +Datasette keeps a catalog of every attached database's tables, columns, indexes and foreign keys in its :ref:`internal database `. Each database has a *schema watch mode* that decides how that catalog is kept current: + +``owned`` + Only Datasette changes this database's schema. After every write Datasette makes to it, Datasette checks whether the schema changed. Blocking writes normally refresh a changed schema in the catalog before returning. If a refresh fails transiently, Datasette retains the previous catalog and retries later; the write may already have committed. Owned databases are never polled. +``external`` + Other processes may change the file too - for example ``sqlite-utils`` adding a table while Datasette is serving it. On top of the checks after Datasette's own writes, the file is polled every ``schema_watch_interval_ms`` milliseconds. Each check is a ``stat()`` of the database file and its ``-wal`` and ``-journal`` files; a database is only opened if those changed, and its catalog is only rebuilt if its schema changed. Databases whose file is replaced (for example by renaming another file over it) or deleted are noticed too: open connections to the old file are discarded and the catalog is updated. +``immutable`` + The catalog is built once at startup. Immutable databases (opened with ``-i``) always use this mode. + +Set this to ``0`` to turn polling off. External databases are then only checked after Datasette's own writes to them:: + + datasette mydatabase.db --setting schema_watch_interval_ms 0 + +Database files, and named in-memory databases, use the :ref:`setting_default_schema_watch` mode - including those added by plugins with :ref:`datasette.add_database() `, unless the plugin passes ``schema_watch="owned"`` (or ``"external"``). A named in-memory database in ``external`` mode is checked with ``PRAGMA schema_version`` on its write connection every ``schema_watch_interval_ms``, which notices tables created through a plugin's own connection to it. You can set the mode for an individual database in ``datasette.yaml``, which overrides both: + +.. code-block:: yaml + + databases: + mydatabase: + schema_watch: owned # or external, or immutable + +.. _setting_default_schema_watch: + +default_schema_watch +~~~~~~~~~~~~~~~~~~~~ + +.. [[[cog + setting_default(cog, "default_schema_watch") +.. ]]] + +Default: ``external`` + +.. [[[end]]] + +The schema watch mode, ``external`` or ``owned``, for mutable database files and named in-memory databases that do not have a mode of their own: files passed on the command line or with ``Datasette(files=...)``, and databases added with :ref:`datasette.add_database() ` without a ``schema_watch`` argument - see :ref:`setting_schema_watch_interval_ms`. + +The default, ``external``, suits the common workflow of changing a database with another tool such as ``sqlite-utils`` while Datasette is serving it: new tables and columns show up within about a second. If Datasette is the only thing that ever changes your database files, set this to ``owned`` and they will not be polled. This is worth doing for instances with thousands of databases, where polling costs a few milliseconds of CPU per second:: + + datasette *.db --setting default_schema_watch owned + + .. _setting_allow_facet: allow_facet diff --git a/tests/connection_budget.py b/tests/connection_budget.py new file mode 100644 index 00000000..3cc6d906 --- /dev/null +++ b/tests/connection_budget.py @@ -0,0 +1,204 @@ +"""Record which databases a block of code touches, and assert a budget. + +With a capped, idle-reaped connection pool, anything that loops over every +attached database opens (and evicts) a connection per database at a 0% hit +rate. These helpers make "this request only touches the databases it is +about" a testable property:: + + async with connection_budget() as budget: + await ds.client.get("/db1/t1.json") + budget.assert_only("db1") # plus _internal, always allowed + +What counts as touching a database: + +* ``opened``: ``Database.connect()`` calls - pooled read connections, write + connections, isolated connections and the SchemaWatcher's short-lived + scan/pragma connections (``track=False``) +* ``reads``: ``Database._execute_fn()`` calls (every ``execute()``, + ``execute_fn()`` and introspection helper), whether or not the pool had an + idle connection - so a warm pool does not hide a loop over every database +* ``writes``: ``execute_write*()``, ``execute_isolated_fn()`` and work run on + the write connection +* ``attached``: databases ATTACHed to a connection while it was prepared + (``--crossdb``) +* ``probes``: ``sqlite3.connect()`` calls that bypass ``Database.connect()`` + (``:memory:`` feature probes and the like), keyed by filename + +The patches are installed process-wide for the duration of the block, so +background work (the SchemaWatcher's polling task, the read pool's reaper) +that runs during the block is counted too. Tests that want to measure one +request should turn polling off (``schema_watch_interval_ms=0``) or measure +it separately. +""" + +import collections +import contextlib +import threading + +from datasette import app as app_module +from datasette.database import Database +from datasette.utils import sqlite3 + +INTERNAL = "__INTERNAL__" + +_local = threading.local() + + +class ConnectionBudget: + def __init__(self): + self._lock = threading.Lock() + self.opened = collections.Counter() + self.reads = collections.Counter() + self.writes = collections.Counter() + self.attached = collections.Counter() + self.probes = collections.Counter() + + def _add(self, counter, key): + with self._lock: + counter[key] += 1 + + def touched(self): + """Names of every database opened, read, written or ATTACHed.""" + names = set() + for counter in (self.opened, self.reads, self.writes, self.attached): + names.update(counter) + return names + + def summary(self): + return { + "opened": dict(self.opened), + "reads": dict(self.reads), + "writes": dict(self.writes), + "attached": dict(self.attached), + "probes": dict(self.probes), + } + + def assert_only(self, *names, internal=True, probes=0): + """Assert that nothing outside ``names`` was touched. + + ``_internal`` is allowed unless ``internal=False``. ``probes`` is the + number of direct ``sqlite3.connect()`` calls allowed. + """ + allowed = set(names) + if internal: + allowed.add(INTERNAL) + extra = self.touched() - allowed + assert not extra, ( + f"touched {len(extra)} database(s) outside the budget " + f"{sorted(allowed)}: {sorted(extra)[:10]}" + f"{' ...' if len(extra) > 10 else ''}\n{self.summary()}" + ) + total_probes = sum(self.probes.values()) + assert ( + total_probes <= probes + ), f"{total_probes} direct sqlite3.connect() call(s): {dict(self.probes)}" + + def assert_at_most(self, n, *, internal=True): + """Assert that at most ``n`` user databases were touched.""" + touched = self.touched() + if internal: + touched.discard(INTERNAL) + assert len(touched) <= n, ( + f"touched {len(touched)} databases, budget {n}: " + f"{sorted(touched)[:10]}\n{self.summary()}" + ) + + def assert_no_connections(self, internal=True): + """Assert that no connection was opened at all (``_internal`` aside).""" + opened = { + k: v for k, v in self.opened.items() if not (internal and k == INTERNAL) + } + assert not opened, f"opened connections: {opened}" + assert not self.attached, f"attached: {dict(self.attached)}" + assert not self.probes, f"probes: {dict(self.probes)}" + + +@contextlib.contextmanager +def _patched(budget): + original_connect = Database.connect + original_execute_fn = Database._execute_fn + original_execute_write_fn = Database._execute_write_fn + original_isolated = Database.execute_isolated_fn + original_on_write_conn = Database._execute_on_write_connection + original_prepare = app_module.Datasette._prepare_connection + original_sqlite_connect = sqlite3.connect + + def connect(self, *args, **kwargs): + budget._add(budget.opened, self.name) + _local.depth = getattr(_local, "depth", 0) + 1 + try: + return original_connect(self, *args, **kwargs) + finally: + _local.depth -= 1 + + async def _execute_fn(self, fn): + budget._add(budget.reads, self.name) + return await original_execute_fn(self, fn) + + async def _execute_write_fn(self, fn, *args, **kwargs): + budget._add(budget.writes, self.name) + return await original_execute_write_fn(self, fn, *args, **kwargs) + + async def execute_isolated_fn(self, fn): + budget._add(budget.writes, self.name) + return await original_isolated(self, fn) + + async def _execute_on_write_connection(self, fn): + budget._add(budget.writes, self.name) + return await original_on_write_conn(self, fn) + + def _prepare_connection(self, conn, database, **kwargs): + result = original_prepare(self, conn, database, **kwargs) + try: + rows = conn.execute("PRAGMA database_list").fetchall() + except Exception: # noqa: BLE001 + rows = [] + for row in rows: + name = row[1] + if name not in ("main", "temp"): + budget._add(budget.attached, name) + return result + + def sqlite_connect(database, *args, **kwargs): + if not getattr(_local, "depth", 0): + budget._add(budget.probes, str(database)) + return original_sqlite_connect(database, *args, **kwargs) + + Database.connect = connect + Database._execute_fn = _execute_fn + Database._execute_write_fn = _execute_write_fn + Database.execute_isolated_fn = execute_isolated_fn + Database._execute_on_write_connection = _execute_on_write_connection + app_module.Datasette._prepare_connection = _prepare_connection + sqlite3.connect = sqlite_connect + try: + yield budget + finally: + Database.connect = original_connect + Database._execute_fn = original_execute_fn + Database._execute_write_fn = original_execute_write_fn + Database.execute_isolated_fn = original_isolated + Database._execute_on_write_connection = original_on_write_conn + app_module.Datasette._prepare_connection = original_prepare + sqlite3.connect = original_sqlite_connect + + +class connection_budget: + """Context manager (sync or async) yielding a ConnectionBudget.""" + + def __init__(self): + self.budget = ConnectionBudget() + self._cm = None + + def __enter__(self): + self._cm = _patched(self.budget) + return self._cm.__enter__() + + def __exit__(self, *exc): + return self._cm.__exit__(*exc) + + async def __aenter__(self): + return self.__enter__() + + async def __aexit__(self, *exc): + return self.__exit__(*exc) diff --git a/tests/connection_torture.py b/tests/connection_torture.py new file mode 100644 index 00000000..f33f8d48 --- /dev/null +++ b/tests/connection_torture.py @@ -0,0 +1,1251 @@ +""" +Randomized torture harness for Datasette's connection management. + +Run one scenario per process, so that a segfault or a hang is a result rather +than the end of the test run:: + + python tests/connection_torture.py SCENARIO --seed N [--duration S] + +Several event loops, each in its own thread (and each running several +tasks), drive one Datasette with extreme settings while a separate process +changes some of the database files. The last line of output is a JSON +summary. Exit status: 0 every check passed, 1 a check failed, 3 something +hung (an operation or a thread did not finish in time); a crash shows up as +death by a signal (faulthandler prints the stacks). + +Checks: + +* no operation hangs, no crash; +* only documented exceptions (see ``Torture.expected()``); +* per database, each writer's rows are in the order it wrote them, and + every acknowledged blocking write (and every non-blocking one, once the + queue has drained) is in the file; +* reads see the reader's own earlier blocking writes; +* the catalog lists a table as soon as a blocking DDL write returns (when + the checkout has a SchemaWatcher), and matches the files after a forced + refresh; +* after ``Datasette.close()`` and ``gc.collect()``: no file descriptor open + on any database file, no write or SQL thread left, the instance itself + collected. + +Works against older checkouts too: features they lack (settings and schema watch modes) are detected and skipped. +""" + +import argparse +import asyncio +import collections +import faulthandler +import gc +import inspect +import json +import os +import random +import shutil +import sqlite3 +import subprocess +import sys +import tempfile +import threading +import time +import traceback +import weakref + +from datasette.app import Datasette +from datasette.database import Database, QueryInterrupted + +try: + from datasette.app import DEFAULT_SETTINGS +except ImportError: # pragma: no cover + DEFAULT_SETTINGS = {} +try: + from datasette.database import DatasetteClosedError +except ImportError: # pragma: no cover + DatasetteClosedError = RuntimeError +try: + from datasette.connection_pool import ConnectionLeaseError +except ImportError: + ConnectionLeaseError = None + +OP_TIMEOUT = 30.0 +SLOW_SQL = ( + "with recursive c(x) as (select 1 union all select x + 1 from c " + "where x < 5000000) select count(*) from c" +) + +SCENARIOS = { + # Reads, writes of every kind, DDL, isolated functions, analyze_sql, + # cancellation, external schema changes - on databases that stay put + "mixed": { + "loops": 3, + "tasks": 4, + "threads": 3, + "duration": 6.0, + "stable": 4, + "ext": 2, + "external_process": True, + "ops": { + "read": 6, + "execute_fn": 3, + "http": 3, + "write": 6, + "write_nonblocking": 3, + "write_many": 2, + "write_script": 1, + "ddl": 2, + "isolated": 1, + "analyze": 1, + "cancel_read": 1, + "cancel_write": 1, + }, + }, + # Databases come and go while they are being used + "lifecycle": { + "loops": 3, + "tasks": 3, + "threads": 3, + "duration": 6.0, + "stable": 3, + "ext": 1, + "churn": 3, + "external_process": True, + "ops": { + "read": 4, + "write": 4, + "write_nonblocking": 1, + "http": 2, + "isolated": 1, + "churn_read": 4, + "churn_write": 3, + "churn_slow": 1, + "add_remove": 2, + "replace_delete": 1, + "db_close": 1, + "cancel_read": 1, + }, + }, + # Datasette.close() from one loop while the others keep going + "close": { + "loops": 3, + "tasks": 4, + "threads": 3, + "duration": 4.0, + "stable": 3, + "ext": 1, + "churn": 2, + "external_process": True, + "close_midflight": True, + "ops": { + "read": 4, + "write": 4, + "write_nonblocking": 2, + "http": 2, + "ddl": 1, + "isolated": 1, + "churn_slow": 1, + "cancel_read": 1, + }, + }, + # Slow queries with a progress handler, cancelled or with their database + # closed underneath them (the close-while-stepping segfault) + "cancel": { + "loops": 2, + "tasks": 6, + "threads": 2, + "duration": 5.0, + "stable": 2, + "churn": 3, + "ops": { + "cancel_read": 4, + "cancel_write": 2, + "churn_slow": 4, + "churn_isolated_slow": 3, + "add_remove": 3, + "write": 2, + "read": 2, + }, + }, + # Databases removed (closed) while isolated functions with a progress + # handler step them - on main, immutable isolated functions were not + # tracked as pending work, and close() could free a connection mid-step + # (segfault) + "closerace": { + "loops": 3, + "tasks": 6, + "threads": 3, + "duration": 6.0, + "stable": 1, + "churn": 4, + "churn_mutable_every": 1, + "isolated_rows": 20000, + "ops": { + "churn_isolated_slow": 8, + "churn_slow": 2, + "churn_read": 2, + "add_remove": 6, + }, + }, + # Write-heavy: execute_write_many() and execute_write() from several + # loops at once (a returned cursor freed on the caller's thread used to + # reset the write thread's statement: "bad parameter or other API + # misuse" on Python < 3.12) + "writes": { + "loops": 3, + "tasks": 6, + "threads": 3, + "duration": 6.0, + "stable": 2, + "ext": 1, + "external_process": True, + "ops": { + "write_many": 6, + "write": 4, + "write_nonblocking": 3, + "cancel_write": 2, + "write_script": 1, + "ddl": 1, + "read": 3, + "isolated": 1, + }, + }, + # Many databases, the smallest pool, 1ms idle timeout, crossdb + "pool": { + "loops": 3, + "tasks": 6, + "threads": 3, + "duration": 5.0, + "stable": 12, + "crossdb": True, + "private_cap": 1, + "ops": { + "read": 6, + "execute_fn": 3, + "http": 3, + "crossdb_read": 2, + "write": 3, + }, + }, + # num_sql_threads=0: one thread, a new event loop every round (the + # TestClient pattern), every kind of operation + "nothreads": { + "loops": 1, + "rounds": True, + "tasks": 6, + "threads": 0, + "duration": 5.0, + "stable": 3, + "ext": 1, + "churn": 2, + "external_process": True, + "ops": { + "read": 4, + "execute_fn": 2, + "http": 2, + "write": 4, + "write_nonblocking": 2, + "write_many": 1, + "write_script": 1, + "ddl": 2, + "isolated": 1, + "analyze": 1, + "churn_read": 2, + "add_remove": 1, + "replace_delete": 1, + }, + }, +} + +EXTERNAL_WRITER = r""" +import random, sqlite3, sys, time +paths = sys.argv[2:] +rng = random.Random(int(sys.argv[1])) +n = 0 +while True: + path = rng.choice(paths) + try: + conn = sqlite3.connect(path, timeout=1, isolation_level=None) + n += 1 + kind = rng.random() + if kind < 0.4: + conn.execute(f"create table if not exists ext_{n % 7} (id integer, v text)") + elif kind < 0.6: + conn.execute(f"drop table if exists ext_{rng.randrange(7)}") + elif kind < 0.8: + try: + conn.execute(f"alter table ext_{rng.randrange(7)} add column c{n} text") + except sqlite3.OperationalError: + pass + else: + conn.execute("insert into data (v) values ('external')") + conn.close() + except sqlite3.Error: + pass + time.sleep(rng.uniform(0.005, 0.05)) +""" + + +class Hang(Exception): + pass + + +def open_file_paths(): + # Paths of every file this process has open, including files that have + # since been deleted (psutil's open_files() leaves those out) + paths = [] + if os.path.isdir("/proc/self/fd"): + for fd in os.listdir("/proc/self/fd"): + try: + paths.append(os.readlink(f"/proc/self/fd/{fd}")) + except OSError: + continue + return paths + # macOS: no /proc, but F_GETPATH works on any open descriptor + import fcntl + + for fd in os.listdir("/dev/fd"): + try: + path = fcntl.fcntl(int(fd), fcntl.F_GETPATH, bytes(1024)) + except OSError: + continue + paths.append(os.fsdecode(path.rstrip(b"\0"))) + return paths + + +def make_file(path, wal=False, extra_sql=""): + conn = sqlite3.connect(path) + if wal: + conn.execute("PRAGMA journal_mode=wal") + conn.executescript( + "create table if not exists log (id integer primary key, writer text, " + "seq integer, kind text);" + "create index if not exists log_writer on log (writer, seq);" + "create table if not exists data (id integer primary key, v text);" + "insert into data (v) values ('one'), ('two');" + extra_sql + ) + conn.commit() + conn.close() + + +def describe(exc): + return f"{type(exc).__module__}.{type(exc).__name__}: {exc}"[:300] + + +class Writer: + """One task's identity: its own sequence numbers per database.""" + + def __init__(self, name, rng): + self.name = name + self.rng = rng + self.seq = collections.Counter() + self.ddl = collections.Counter() + + +class Torture: + def __init__(self, scenario, seed, duration=None): + self.name = scenario + self.spec = dict(SCENARIOS[scenario]) + if duration is not None: + self.spec["duration"] = duration + self.seed = seed + self.rng = random.Random(seed) + self.tmp = tempfile.mkdtemp(prefix="datasette_torture_") + self.lock = threading.Lock() + self.problems = [] + self.counts = collections.Counter() + self.expected_errors = collections.Counter() + # (db name, writer) -> list of seqs acknowledged by blocking writes + self.acked = collections.defaultdict(list) + # (db name, writer) -> seqs issued as non-blocking writes, and seqs + # whose write was cancelled (may or may not have committed) + self.nonblocking = collections.defaultdict(list) + self.maybe = collections.defaultdict(list) + self.closing = False + self.closed_at = None + self.stop_at = None + self.churn_seq = 0 + self.features = {} + self.stable_names = [] + self.ext_names = [] + self.churn_paths = {} + self.ds = None + self.external = None + self.writers = {} + + # -- bookkeeping ----------------------------------------------------- + + def problem(self, kind, message): + with self.lock: + self.problems.append({"kind": kind, "message": message[:2000]}) + + def dump_stacks_once(self, why): + """Every thread's stack on stderr, the first time something hangs.""" + with self.lock: + if getattr(self, "_dumped", False): + return + self._dumped = True + sys.stderr.write(f"\n=== {why}: stacks of all threads ===\n") + faulthandler.dump_traceback(all_threads=True) + + def count(self, key): + with self.lock: + self.counts[key] += 1 + + def expected(self, op, group, exc): + """Is exc a documented outcome of op on a database in group? + Returns a label for the counts, or None for a problem.""" + message = str(exc) + if isinstance(exc, QueryInterrupted): + return "QueryInterrupted" + if isinstance(exc, DatasetteClosedError) and (group == "churn" or self.closing): + return type(exc).__name__ + if self.closing and isinstance(exc, RuntimeError): + # Work submitted to the executor after close() shut it down + return "RuntimeError after close" + if isinstance(exc, KeyError) and group == "churn": + return "KeyError (database removed)" + if isinstance(exc, sqlite3.Error): + if group == "churn": + # Files replaced, deleted and closed underneath + return f"sqlite3 on {group}: {message[:40]}" + if group == "mem" and "is locked" in message: + # Shared-cache in-memory databases do not wait for table or + # schema locks (documented) + return "mem: shared-cache lock" + if group == "ext" and ( + "no such table" in message + or "no such column" in message + or "database is locked" in message + or "schema has changed" in message + ): + return f"ext: {message[:30]}" + if self.closing and "closed database" in message: + return "closed database after close" + return None + + async def call(self, op, group, awaitable, timeout=OP_TIMEOUT): + """Await with a timeout; classify exceptions. Returns (ok, value).""" + try: + value = await asyncio.wait_for(awaitable, timeout) + except asyncio.TimeoutError: + self.problem("hang", f"{op} on {group} did not finish in {timeout}s") + self.dump_stacks_once(f"{op} on {group} hung") + return False, None + except asyncio.CancelledError: + raise + except Exception as e: # noqa: BLE001 + label = self.expected(op, group, e) + if label is None: + if ConnectionLeaseError is not None and isinstance( + e, ConnectionLeaseError + ): + kind = "lease" + else: + kind = "unexpected" + self.problem( + kind, + f"{op} on {group}: {describe(e)}\n" + + "".join(traceback.format_exception(e)[-6:]), + ) + else: + with self.lock: + self.expected_errors[label] += 1 + return False, None + self.count(op) + return True, value + + # -- setup ----------------------------------------------------------- + + def settings(self): + wanted = { + "num_sql_threads": self.spec["threads"], + "max_open_connections": 1, + "connection_idle_timeout_ms": 1, + "schema_watch_interval_ms": 10, + } + return {k: v for k, v in wanted.items() if k in DEFAULT_SETTINGS} + + def build(self): + spec = self.spec + files = [] + for i in range(spec.get("stable", 0)): + path = os.path.join(self.tmp, f"stable_{i}.db") + make_file(path, wal=i % 2 == 0) + files.append(path) + ext_paths = [] + for i in range(spec.get("ext", 0)): + path = os.path.join(self.tmp, f"ext_{i}.db") + make_file(path, wal=i % 2 == 1) + ext_paths.append(path) + imm = os.path.join(self.tmp, "imm.db") + make_file(imm) + self.template = os.path.join(self.tmp, "template.db") + make_file(self.template) + kwargs = {} + if spec.get("crossdb"): + kwargs["crossdb"] = True + self.features["schema_watch"] = ( + "schema_watch" in inspect.signature(Datasette.add_database).parameters + ) + self.ds = Datasette( + files + ext_paths, immutables=[imm], settings=self.settings(), **kwargs + ) + ds = self.ds + self.features["watcher"] = hasattr(ds, "_schema_watcher") + self.stable_names = [os.path.basename(p)[:-3] for p in files] + self.ext_names = [os.path.basename(p)[:-3] for p in ext_paths] + # Half the stable databases owned (when the checkout has modes): + # added with add_database() rather than on the "command line" + for i in range(spec.get("stable", 0)): + if i % 2 == 1 and self.features["schema_watch"]: + name = f"stable_{i}" + ds.remove_database(name) + ds.add_database( + Database(ds, path=files[i]), name=name, schema_watch="owned" + ) + ds.add_memory_database(f"torture_mem_{os.getpid()}_{self.seed}", name="mem") + self.groups = {name: "stable" for name in self.stable_names} + self.groups.update({name: "ext" for name in self.ext_names}) + self.groups["mem"] = "mem" + self.groups["imm"] = "imm" + for i in range(spec.get("churn", 0)): + self.add_churn() + if spec.get("external_process") and ext_paths: + self.external = subprocess.Popen( + [sys.executable, "-c", EXTERNAL_WRITER, str(self.seed), *ext_paths], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + + def add_churn(self): + with self.lock: + self.churn_seq += 1 + n = self.churn_seq + name = f"churn_{n}" + path = os.path.join(self.tmp, f"{name}.db") + shutil.copy(self.template, path) + # Every nth churn database is immutable (all of them with 1) + mutable = n % self.spec.get("churn_mutable_every", 4) != 0 + db = Database(self.ds, path=path, is_mutable=mutable) + kwargs = {"schema_watch": "external"} if self.features["schema_watch"] else {} + self.ds.add_database(db, name=name, **kwargs) + with self.lock: + self.churn_paths[name] = path + self.groups[name] = "churn" + return name + + # -- operations ------------------------------------------------------ + + def pick(self, rng, groups): + names = [n for n, g in list(self.groups.items()) if g in groups] + if not names: + return None, None + name = rng.choice(names) + return name, self.groups.get(name) + + def db(self, name): + return self.ds.databases.get(name) + + async def op_read(self, w): + name, group = self.pick(w.rng, ("stable", "ext", "mem", "imm")) + db = self.db(name) + if db is None: + return + if group == "stable" and w.seq[name] and w.rng.random() < 0.5: + # Read-your-writes: this writer's last acknowledged seq + last = self.acked.get((name, w.name)) + ok, results = await self.call( + "read", + group, + db.execute("select max(seq) from log where writer = ?", [w.name]), + ) + if ok and last: + got = results.first()[0] + if got is None or got < last[-1]: + self.problem( + "read-your-writes", + f"{name}: writer {w.name} read max(seq)={got} after " + f"its write {last[-1]} was acknowledged", + ) + return + sql = w.rng.choice( + [ + "select count(*) from data", + "select * from data order by id desc limit 3", + "select name from sqlite_master", + "PRAGMA schema_version", + ] + ) + await self.call("read", group, db.execute(sql)) + + async def op_execute_fn(self, w): + name, group = self.pick(w.rng, ("stable", "ext", "mem", "imm")) + db = self.db(name) + if db is None: + return + + def fn(conn): + n = conn.execute("select count(*) from data").fetchone()[0] + tables = [r[0] for r in conn.execute("select name from sqlite_master")] + return n, tables + + await self.call("execute_fn", group, db.execute_fn(fn)) + + async def op_http(self, w): + name, group = self.pick(w.rng, ("stable", "ext", "imm", "mem")) + path = w.rng.choice( + [ + f"/{name}.json", + f"/{name}/data.json?_size=2", + "/.json", + "/-/databases.json", + ] + ) + ok, response = await self.call("http", group, self.ds.client.get(path)) + if ok and response.status_code != 200: + if self.closing or (group == "ext" and response.status_code == 404): + # 404: data or the database page during an external change + self.count(f"http {response.status_code} (expected)") + return + if response.status_code in (400, 500) and ( + "schema is locked" in response.text + or "table is locked" in response.text + ): + # SQLITE_LOCKED from the shared-cache in-memory database: it + # does not wait for locks (documented, the same on main). The + # table page reports it as a 400 "Invalid SQL" (P45) + with self.lock: + self.expected_errors[ + f"http {response.status_code}: shared-cache lock" + ] += 1 + return + self.problem( + "http", f"GET {path} -> {response.status_code}: {response.text[:300]}" + ) + + def _next_seq(self, w, name): + w.seq[name] += 1 + return w.seq[name] + + async def _write(self, w, name, group, block): + db = self.db(name) + if db is None: + return + seq = self._next_seq(w, name) + ok, _ = await self.call( + "write" if block else "write_nonblocking", + group, + db.execute_write( + "insert into log (writer, seq, kind) values (?, ?, ?)", + [w.name, seq, "write"], + block=block, + ), + ) + # Recorded once the call has returned: a block=False write is in the + # queue by then, so a later drain write is behind it + if ok and block: + self.acked[(name, w.name)].append(seq) + elif ok: + self.nonblocking[(name, w.name)].append(seq) + elif not block: + # Failed to queue, or maybe queued then failed: either way it + # may or may not be in the file + self.maybe[(name, w.name)].append(seq) + + async def op_write(self, w): + name, group = self.pick(w.rng, ("stable",)) + if name: + await self._write(w, name, group, True) + + async def op_write_nonblocking(self, w): + name, group = self.pick(w.rng, ("stable",)) + if name: + await self._write(w, name, group, False) + + async def op_write_many(self, w): + name, group = self.pick(w.rng, ("stable",)) + db = self.db(name) + if db is None: + return + seqs = [self._next_seq(w, name) for _ in range(w.rng.randint(1, 4))] + ok, _ = await self.call( + "write_many", + group, + db.execute_write_many( + "insert into log (writer, seq, kind) values (?, ?, 'many')", + [(w.name, s) for s in seqs], + ), + ) + if ok: + self.acked[(name, w.name)].extend(seqs) + + async def op_write_script(self, w): + name, group = self.pick(w.rng, ("stable",)) + db = self.db(name) + if db is None: + return + seqs = [self._next_seq(w, name) for _ in range(2)] + script = "".join( + f"insert into log (writer, seq, kind) values ('{w.name}', {s}, 'script');" + for s in seqs + ) + ok, _ = await self.call("write_script", group, db.execute_write_script(script)) + if ok: + self.acked[(name, w.name)].extend(seqs) + + async def op_ddl(self, w): + name, group = self.pick(w.rng, ("stable", "mem")) + db = self.db(name) + if db is None: + return + w.ddl[name] += 1 + n = w.ddl[name] + if n % 3 == 0: + sql = f"drop table if exists ddl_{w.name}_{n - 1}" + table = None + else: + table = f"ddl_{w.name}_{n}" + sql = f"create table if not exists {table} (id integer primary key)" + ok, _ = await self.call("ddl", group, db.execute_write(sql)) + if ok and table and self.features["watcher"] and not self.closing: + # The catalog lists a table once a blocking DDL write returns + ok2, rows = await self.call( + "ddl_catalog", + group, + self.ds.get_internal_database().execute( + "select 1 from catalog_tables where database_name = ? " + "and table_name = ?", + [name, table], + ), + ) + if ok2 and not rows.rows and not self.closing: + self.problem( + "catalog", f"{name}: {table} not in the catalog after its DDL" + ) + + async def op_isolated(self, w): + name, group = self.pick(w.rng, ("stable", "imm")) + db = self.db(name) + if db is None: + return + + def fn(conn): + return conn.execute("select count(*) from data").fetchone()[0] + + await self.call("isolated", group, db.execute_isolated_fn(fn)) + + async def op_analyze(self, w): + name, group = self.pick(w.rng, ("stable",)) + db = self.db(name) + if db is None or not hasattr(db, "analyze_sql"): + return + await self.call( + "analyze", group, db.analyze_sql("select * from log where writer = 'x'") + ) + + async def op_cancel_read(self, w): + name, group = self.pick(w.rng, ("stable", "imm")) + db = self.db(name) + if db is None: + return + task = asyncio.ensure_future(db.execute(SLOW_SQL, custom_time_limit=300)) + await asyncio.sleep(w.rng.uniform(0, 0.05)) + task.cancel() + try: + await asyncio.wait_for(task, OP_TIMEOUT) + self.count("cancel_read finished") + except asyncio.CancelledError: + self.count("cancel_read cancelled") + except asyncio.TimeoutError: + self.problem("hang", f"cancelled read on {name} did not finish") + except Exception as e: # noqa: BLE001 + if self.expected("cancel_read", group, e) is None: + self.problem("unexpected", f"cancel_read on {name}: {describe(e)}") + + async def op_cancel_write(self, w): + name, group = self.pick(w.rng, ("stable",)) + db = self.db(name) + if db is None: + return + seq = self._next_seq(w, name) + self.maybe[(name, w.name)].append(seq) + task = asyncio.ensure_future( + db.execute_write( + "insert into log (writer, seq, kind) values (?, ?, 'cancelled')", + [w.name, seq], + ) + ) + await asyncio.sleep(w.rng.uniform(0, 0.005)) + task.cancel() + try: + await asyncio.wait_for(task, OP_TIMEOUT) + self.count("cancel_write finished") + except asyncio.CancelledError: + self.count("cancel_write cancelled") + except asyncio.TimeoutError: + self.problem("hang", f"cancelled write on {name} did not finish") + except Exception as e: # noqa: BLE001 + if self.expected("cancel_write", group, e) is None: + self.problem("unexpected", f"cancel_write on {name}: {describe(e)}") + + async def op_crossdb_read(self, w): + db = self.db("_memory") + name, group = self.pick(w.rng, ("stable",)) + if db is None or name is None: + return + attached = {r.name for r in await db.attached_databases()} + if name not in attached: + return + await self.call( + "crossdb_read", group, db.execute(f"select count(*) from [{name}].data") + ) + + async def op_churn_read(self, w): + name, group = self.pick(w.rng, ("churn",)) + db = self.db(name) + if db is None: + return + if w.rng.random() < 0.5: + await self.call( + "churn_read", group, db.execute("select count(*) from data") + ) + else: + ok, response = await self.call( + "churn_http", group, self.ds.client.get(f"/{name}/data.json?_size=1") + ) + # 400: SQL errors on a replaced or deleted file are reported as + # "Invalid SQL" by the table page + if ok and response.status_code not in (200, 400, 404, 500): + self.problem("http", f"GET /{name}/data.json -> {response.status_code}") + + async def op_churn_write(self, w): + name, group = self.pick(w.rng, ("churn",)) + db = self.db(name) + if db is None or not db.is_mutable: + return + await self.call( + "churn_write", + group, + db.execute_write( + "insert into data (v) values ('churn')", block=w.rng.random() < 0.7 + ), + ) + + async def op_churn_slow(self, w): + # A slow query with Datasette's progress-handler time limit, on a + # database that another task may close or remove meanwhile + name, group = self.pick(w.rng, ("churn",)) + db = self.db(name) + if db is None: + return + await self.call( + "churn_slow", group, db.execute(SLOW_SQL, custom_time_limit=200) + ) + + async def op_churn_isolated_slow(self, w): + # The close-while-stepping crash on main: an isolated function on an + # immutable database (not tracked as pending work there) with a + # progress handler installed, while the database is removed + name, group = self.pick(w.rng, ("churn",)) + db = self.db(name) + if db is None: + return + + rows = self.spec.get("isolated_rows", 300000) + + def fn(conn): + conn.set_progress_handler(lambda: 0, 100) + try: + return conn.execute( + "with recursive c(x) as (select 1 union all select x + 1 " + f"from c where x < {rows}) select count(*) from c" + ).fetchone()[0] + finally: + conn.set_progress_handler(None, 0) + + await self.call("churn_isolated_slow", group, db.execute_isolated_fn(fn)) + + async def op_add_remove(self, w): + churn = [n for n, g in list(self.groups.items()) if g == "churn"] + if len(churn) < 2 or (len(churn) < 6 and w.rng.random() < 0.5): + if not self.closing: + self.add_churn() + self.count("add_database") + return + name = w.rng.choice(churn) + with self.lock: + if self.groups.get(name) != "churn": + return + del self.groups[name] + try: + self.ds.remove_database(name) + self.count("remove_database") + except KeyError: + pass + except Exception as e: # noqa: BLE001 + if self.expected("remove_database", "churn", e) is None: + self.problem("unexpected", f"remove_database({name}): {describe(e)}") + + async def op_replace_delete(self, w): + name, _group = self.pick(w.rng, ("churn",)) + path = self.churn_paths.get(name) + db = self.db(name) + if path is None or db is None or not db.is_mutable: + return + if w.rng.random() < 0.6: + new = f"{path}.{w.name}.{w.rng.randrange(10**9)}.new" + make_file(new, extra_sql="create table replaced (id integer);") + os.replace(new, path) + self.count("file replaced") + else: + for suffix in ("", "-wal", "-shm", "-journal"): + try: + os.unlink(path + suffix) + except OSError: + pass + self.count("file deleted") + + async def op_db_close(self, w): + name, _group = self.pick(w.rng, ("churn",)) + db = self.db(name) + if db is None or not hasattr(db, "close"): + return + with self.lock: + if self.groups.get(name) != "churn": + return + self.groups.pop(name, None) + db.close() + self.count("Database.close") + + # -- driving --------------------------------------------------------- + + async def worker(self, w, deadline): + ops = self.spec["ops"] + names = list(ops) + weights = [ops[n] for n in names] + while time.monotonic() < deadline: + op = w.rng.choices(names, weights)[0] + if self.closing and op in ("add_remove", "replace_delete"): + op = "read" + try: + await getattr(self, f"op_{op}")(w) + except asyncio.CancelledError: + raise + except Exception as e: # noqa: BLE001 + self.problem( + "harness", + f"{op}: {describe(e)}\n" + + "".join(traceback.format_exception(e)[-5:]), + ) + await asyncio.sleep(0) + + def writer(self, loop_index, t): + # One Writer per task slot for the whole run: in "rounds" mode each + # round's event loop reuses the same writers (and their sequences) + key = (loop_index, t) + with self.lock: + w = self.writers.get(key) + if w is None: + rng = random.Random(f"{self.seed}-{loop_index}-{t}") + w = self.writers[key] = Writer(f"w{loop_index}_{t}", rng) + return w + + async def loop_main(self, loop_index, deadline): + tasks = [] + for t in range(self.spec["tasks"]): + tasks.append(self.worker(self.writer(loop_index, t), deadline)) + if self.spec.get("close_midflight") and loop_index == 0: + tasks.append(self.close_midflight(deadline)) + await asyncio.gather(*tasks) + + async def close_midflight(self, deadline): + await asyncio.sleep( + self.rng.uniform(0.3, max(0.4, deadline - time.monotonic() - 1)) + ) + await self.drain_and_verify() + self.closing = True + self.closed_at = time.monotonic() + t0 = time.monotonic() + self.ds.close() + self.counts["close seconds x100"] = int((time.monotonic() - t0) * 100) + + def run_loops(self): + duration = self.spec["duration"] + deadline = time.monotonic() + duration + if self.spec.get("rounds"): + # One thread, a fresh event loop per round + while time.monotonic() < deadline: + round_end = min(deadline, time.monotonic() + 0.4) + asyncio.run(self.loop_main(0, round_end)) + self.count("rounds") + return + errors = [] + + def runner(i): + try: + asyncio.run(self.loop_main(i, deadline)) + except BaseException as e: # noqa: BLE001 + errors.append(e) + + threads = [ + threading.Thread( + target=runner, args=(i,), name=f"torture-loop-{i}", daemon=True + ) + for i in range(self.spec["loops"]) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(duration + OP_TIMEOUT + 15) + if thread.is_alive(): + raise Hang(f"{thread.name} still running") + for e in errors: + self.problem("harness", f"loop died: {describe(e)}") + + # -- checks ---------------------------------------------------------- + + async def drain_and_verify(self): + """Wait for every write queue to drain, then check ordering and that + nothing acknowledged is missing. Runs before close().""" + ds = self.ds + # Other loops may still be writing (close_midflight): only writes + # acknowledged or queued before the drain must be visible to it + with self.lock: + acked = {k: list(v) for k, v in self.acked.items()} + nonblocking = {k: list(v) for k, v in self.nonblocking.items()} + for name in self.stable_names: + db = ds.databases.get(name) + await self.call("drain", "stable", db.execute_write("select 1")) + for name in self.stable_names: + db = ds.databases.get(name) + ok, results = await self.call( + "verify", + "stable", + db.execute( + "select writer, seq from log order by id", custom_time_limit=None + ), + ) + if ok: + self.check_log( + name, + [(r[0], r[1]) for r in results.rows], + "before close", + acked, + nonblocking, + ) + + def check_log(self, name, rows, phase, acked=None, nonblocking=None): + if acked is None: + acked = self.acked + if nonblocking is None: + nonblocking = self.nonblocking + by_writer = collections.defaultdict(list) + for writer, seq in rows: + by_writer[writer].append(seq) + for (db_name, writer), seqs in list(acked.items()): + if db_name != name: + continue + present = set(by_writer.get(writer, [])) + missing = [s for s in seqs if s not in present] + if missing: + self.problem( + "lost-write", + f"{name} ({phase}): {len(missing)} acknowledged writes by " + f"{writer} missing, e.g. seq {missing[:5]}", + ) + for (db_name, writer), seqs in list(nonblocking.items()): + if db_name != name: + continue + present = set(by_writer.get(writer, [])) + missing = [s for s in seqs if s not in present] + if missing: + self.problem( + "lost-write", + f"{name} ({phase}): {len(missing)} block=False writes by " + f"{writer} missing after the queue drained, e.g. seq {missing[:5]}", + ) + for writer, seqs in by_writer.items(): + if seqs != sorted(seqs) or len(seqs) != len(set(seqs)): + self.problem( + "write-order", + f"{name} ({phase}): writer {writer} rows out of order: {seqs[:20]}", + ) + self.count("verified databases") + + async def check_catalog(self): + """After a forced refresh, the catalog matches each stable and + external database file.""" + ds = self.ds + ok, _ = await self.call("refresh", "stable", ds.refresh_schemas(force=True)) + if not ok: + return + for name in self.stable_names + self.ext_names: + db = ds.databases.get(name) + internal = ds.get_internal_database() + ok1, catalog = await self.call( + "catalog", + "stable", + internal.execute( + "select table_name from catalog_tables where database_name = ?", + [name], + ), + ) + ok2, live = await self.call( + "catalog", + self.groups.get(name, "stable"), + db.execute("select name from sqlite_master where type = 'table'"), + ) + if ok1 and ok2: + catalog_set = {r[0] for r in catalog.rows} + live_set = {r[0] for r in live.rows} + if catalog_set != live_set: + if name in self.ext_names: + # The external process may have changed it since; + # one more forced refresh must converge + await ds.refresh_schemas(force=True) + catalog = await internal.execute( + "select table_name from catalog_tables where database_name = ?", + [name], + ) + live = await db.execute( + "select name from sqlite_master where type = 'table'" + ) + catalog_set = {r[0] for r in catalog.rows} + live_set = {r[0] for r in live.rows} + if catalog_set == live_set or self.external is not None: + continue + self.problem( + "catalog", + f"{name}: catalog {sorted(catalog_set ^ live_set)} differ " + "from the file after refresh_schemas(force=True)", + ) + + def check_files_after_close(self): + for name in self.stable_names: + path = os.path.join(self.tmp, f"{name}.db") + conn = sqlite3.connect(path) + try: + rows = conn.execute( + "select writer, seq from log order by id" + ).fetchall() + finally: + conn.close() + self.check_log(name, rows, "file after close") + + def check_resources(self, baseline_threads): + deadline = time.monotonic() + 15 + # On macOS the temp directory is reported via its /private symlink target + tmp = os.path.realpath(self.tmp) + while True: + gc.collect() + leaked_fds = [ + target + for target in open_file_paths() + if target.startswith(tmp) or "datasette_temp_" in target + ] + threads = [ + t.name + for t in threading.enumerate() + if t.name.startswith("_execute_writes") + or t.name.startswith("ThreadPoolExecutor") + or t.name.startswith("datasette-") + and t.name != "datasette-read-pool-reaper" + ] + if (not leaked_fds and not threads) or time.monotonic() > deadline: + break + time.sleep(0.1) + if leaked_fds: + self.problem( + "fd-leak", + f"{len(leaked_fds)} fds still open after close: " + + ", ".join(sorted({os.path.basename(p) for p in leaked_fds}))[:500], + ) + if threads: + self.problem("thread-leak", f"threads after close: {sorted(set(threads))}") + extra = threading.active_count() - baseline_threads + self.counts["threads above baseline after close"] = max(extra, 0) + + def run(self): + baseline_threads = threading.active_count() + self.build() + ds = self.ds + + async def start(): + await ds.invoke_startup() + await ds.get_database("mem").execute_write_script( + "create table data (id integer primary key, v text);" + "insert into data (v) values ('one'), ('two');" + ) + if self.spec.get("private_cap") and hasattr(ds, "_read_pool"): + # Below the public minimum: the pool's soft-cap path + ds._read_pool.max_open = self.spec["private_cap"] + + asyncio.run(start()) + hung = None + try: + self.run_loops() + except Hang as e: + hung = e + self.problem("hang", str(e)) + if hung is None and not self.closing: + + async def finish(): + await self.drain_and_verify() + await self.check_catalog() + + asyncio.run(finish()) + if self.external is not None: + self.external.kill() + self.external.wait() + if hung is None: + if not self.closing: + ds.close() + ref = weakref.ref(ds) + self.ds = ds = None + self.check_files_after_close() + self.check_resources(baseline_threads) + gc.collect() + if ref() is not None: + self.problem( + "instance-leak", + "Datasette still alive after close() + gc.collect()", + ) + shutil.rmtree(self.tmp, ignore_errors=True) + return hung is None + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("scenario", choices=sorted(SCENARIOS)) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--duration", type=float, default=None) + args = parser.parse_args() + faulthandler.enable() + spec = SCENARIOS[args.scenario] + duration = args.duration if args.duration is not None else spec["duration"] + # A last resort: dump every thread's stack and exit if the whole run + # takes far longer than it should + faulthandler.dump_traceback_later(duration + OP_TIMEOUT + 60, exit=True) + torture = Torture(args.scenario, args.seed, args.duration) + t0 = time.monotonic() + finished = torture.run() + summary = { + "scenario": args.scenario, + "seed": args.seed, + "python": sys.version.split()[0], + "gil": getattr(sys, "_is_gil_enabled", lambda: True)(), + "seconds": round(time.monotonic() - t0, 2), + "features": torture.features, + "ok": finished and not torture.problems, + "problems": torture.problems[:50], + "problem_kinds": dict(collections.Counter(p["kind"] for p in torture.problems)), + "counts": dict(torture.counts), + "expected_errors": dict(torture.expected_errors), + } + print(json.dumps(summary), flush=True) + faulthandler.cancel_dump_traceback_later() + if not finished or any(p["kind"] == "hang" for p in torture.problems): + os._exit(3) + os._exit(0 if summary["ok"] else 1) + + +if __name__ == "__main__": + main() diff --git a/tests/test_api.py b/tests/test_api.py index 690cb080..ec862015 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -759,6 +759,13 @@ async def test_settings_json(ds_client): "suggest_facets": True, "default_cache_ttl": 5, "num_sql_threads": 1, + "max_open_connections": 32, + "max_write_connections": 8, + "max_pending_writes": 256, + "write_queue_timeout_ms": 5000, + "connection_idle_timeout_ms": 30000, + "schema_watch_interval_ms": 1000, + "default_schema_watch": "external", "cache_size_kb": 0, "allow_csv_stream": True, "max_csv_mb": 100, diff --git a/tests/test_connection_budget.py b/tests/test_connection_budget.py new file mode 100644 index 00000000..63853816 --- /dev/null +++ b/tests/test_connection_budget.py @@ -0,0 +1,652 @@ +"""Requests, startup and background work only touch the databases they need. + +With a capped, idle-reaped connection pool, a loop over every attached +database opens (and evicts) one connection per database. These tests attach +50 databases and use tests/connection_budget.py to assert which databases +each endpoint opens, reads or writes. Every test uses a fresh instance, so +caches that are only warm after a first request do not hide a sweep. +""" + +import json +import os +import shutil +import sqlite3 +import time + +import pytest + +from datasette.app import Datasette +from datasette.utils.catalog import ( + catalog_relationship_counts, + catalog_summaries, + catalog_table_details, +) +from datasette.utils.sqlite import ( + sqlite_derived_table_dependencies, + supports_table_list, +) + +from .connection_budget import INTERNAL, connection_budget + +N = 50 + + +@pytest.fixture(autouse=True) +def only_default_plugins(): + """Plugins other tests leave registered (tests/plugins/my_plugin.py and + friends query datasette.get_database() from their hooks) are not part + of what these tests measure: unregister them for the duration.""" + from datasette.plugins import DEFAULT_PLUGINS, pm + + removed = [ + (name, plugin) + for name, plugin in pm.list_name_plugin() + if name not in DEFAULT_PLUGINS + ] + for name, _ in removed: + pm.unregister(name=name) + try: + yield + finally: + for name, plugin in removed: + if pm.get_plugin(name) is None: + pm.register(plugin, name=name) + + +def _make_db(path, i): + conn = sqlite3.connect(path) + conn.executescript(""" + create table t2 (id integer primary key, title text); + create table t1 ( + id integer primary key, name text, other_id integer references t2(id) + ); + """) + conn.executemany( + "insert into t1 (name, other_id) values (?, ?)", + [(f"n{j}", j % 3) for j in range(20)], + ) + conn.executemany( + "insert into t2 (title) values (?)", [(f"t{j}",) for j in range(3)] + ) + if i % 5 == 0: + conn.executescript(""" + create virtual table t1_fts using fts5( + name, content="t1", content_rowid='id' + ); + insert into t1_fts(t1_fts) values('rebuild'); + """) + if i % 3 == 0: + conn.execute("create view v1 as select * from t1") + if i % 7 == 0: + conn.execute("create table _hidden (x)") + conn.commit() + conn.close() + + +@pytest.fixture(scope="module") +def db_files(tmp_path_factory): + directory = tmp_path_factory.mktemp("budget") + paths = [] + for i in range(N): + path = str(directory / f"db{i:03d}.db") + _make_db(path, i) + paths.append(path) + # Older than the SchemaWatcher's racy window, so a sweep is stat-only + old = time.time() - 60 + for path in paths: + os.utime(path, (old, old)) + return paths + + +def _ds(paths, **kwargs): + settings = {"schema_watch_interval_ms": 0} + settings.update(kwargs.pop("settings", {})) + ds = Datasette(paths, settings=settings, **kwargs) + ds.root_enabled = True + return ds + + +def _root(ds): + return {"ds_actor": ds.client.actor_cookie({"id": "root"})} + + +async def _get(ds, path, cookies=None): + async with connection_budget() as budget: + response = await ds.client.get(path, cookies=cookies or {}) + assert response.status_code == 200, (path, response.status_code, response.text) + return response, budget + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/", "/.json", "/-/", "/-/.json"]) +@pytest.mark.parametrize("as_root", [False, True]) +async def test_index_page_opens_no_database(db_files, path, as_root): + ds = _ds(db_files) + await ds.invoke_startup() + response, budget = await _get(ds, path, _root(ds) if as_root else None) + budget.assert_only() + budget.assert_no_connections() + if path.endswith(".json"): + data = response.json() + assert len(data["databases"]) == N + # More databases than COUNT_MAX_DATABASES: no row counts + assert not any(d["show_table_row_counts"] for d in data["databases"]) + db000 = next(d for d in data["databases"] if d["name"] == "db000") + assert db000["tables_count"] == 2 # t1, t2 + # _hidden, and t1_fts (an external content FTS table is hidden). + # Its fts5 shadow tables are not listed at all: they derive from + # t1_fts, which is itself derived (the one-hop permission policy) + assert db000["hidden_tables_count"] == 2 + assert db000["views_count"] == 1 + shown = {t["name"]: t for t in db000["tables_and_views_truncated"]} + assert shown["t1"]["columns"] == ["id", "name", "other_id"] + assert shown["t1"]["primary_keys"] == ["id"] + assert shown["t1"]["fts_table"] == "t1_fts" + # Sorted by foreign key relationships when there are no counts + assert shown["t1"]["num_relationships_for_sorting"] == 1 + ds.close() + + +@pytest.mark.asyncio +async def test_index_page_counts_rows_for_few_databases(db_files): + # At most COUNT_MAX_DATABASES databases: rows are counted, and only the + # databases listed are touched + ds = _ds(db_files[:3]) + await ds.invoke_startup() + response, budget = await _get(ds, "/.json") + budget.assert_only("db000", "db001", "db002") + data = response.json() + assert all(d["show_table_row_counts"] for d in data["databases"]) + assert data["databases"][0]["table_rows_sum"] == 23 + ds.close() + + +@pytest.mark.asyncio +async def test_index_page_uses_cached_immutable_counts(db_files): + # Immutable databases' counts are computed once, when first needed, and + # reused even when there are too many databases to count + ds = Datasette( + db_files[3:], + immutables=db_files[:3], + settings={"schema_watch_interval_ms": 0}, + ) + await ds.invoke_startup() + for name in ("db000", "db001", "db002"): + await ds.get_database(name).table_counts() + response, budget = await _get(ds, "/.json") + budget.assert_no_connections() + budget.assert_only() + by_name = {d["name"]: d for d in response.json()["databases"]} + assert by_name["db000"]["show_table_row_counts"] + assert not by_name["db010"]["show_table_row_counts"] + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/-/databases", "/-/databases.json"]) +async def test_databases_endpoint(db_files, path): + ds = _ds(db_files) + await ds.invoke_startup() + _, budget = await _get(ds, path) + budget.assert_only() + budget.assert_no_connections() + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path", + [ + "/db001", + "/db001.json", + "/db000.json?_search=n", + "/db001/t1", + "/db001/t1.json", + "/db001/t1?_facet=name", + "/db001/t1.json?_facet=name", + "/db000/t1_fts.json?_search=n1", + "/db001/-/query?sql=select+1", + "/db001/-/query.json?sql=select+1", + "/db001/-/query", + "/db001/t1/1", + "/db001/t1/1.json", + "/db001/-/schema.json", + "/db001/t1/-/autocomplete?q=n", + ], +) +@pytest.mark.parametrize("as_root", [False, True]) +async def test_database_scoped_pages_touch_only_their_database(db_files, path, as_root): + ds = _ds(db_files) + await ds.invoke_startup() + database = path.split("/")[1].split(".")[0] + _, budget = await _get(ds, path, _root(ds) if as_root else None) + budget.assert_only(database) + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "path,as_root", + [ + ("/-/versions", False), + ("/-/versions.json", False), + ("/-/plugins.json", False), + ("/-/settings.json", False), + ("/-/config.json", False), + ("/-/actor.json", False), + ("/-/threads", True), + ("/-/threads.json", True), + ("/-/tasks.json", True), + ("/-/actions.json", True), + ("/-/jump?q=t1", False), + ("/-/jump.json?q=t1", False), + ("/-/jump.json", False), + ("/-/allowed.json?action=view-table", False), + ("/-/allowed.json?action=view-table", True), + ("/-/allowed.json?action=view-database", False), + ("/-/allowed?action=view-table", True), + ("/-/rules.json?action=view-table", True), + ("/-/check.json?action=view-table&parent=db001&child=t1", True), + ("/-/permissions", True), + ("/-/allow-debug", False), + ("/-/api", False), + ("/-/api", True), + ("/-/create-token", True), + ("/-/queries.json", False), + ("/-/patterns", False), + ], +) +async def test_instance_pages_open_no_user_database(db_files, path, as_root): + ds = _ds(db_files) + await ds.invoke_startup() + # /-/versions probes the SQLite library once per instance + await ds.client.get("/-/versions.json") + _, budget = await _get(ds, path, _root(ds) if as_root else None) + budget.assert_only() + budget.assert_no_connections() + ds.close() + + +@pytest.mark.asyncio +async def test_versions_probes_sqlite_once(db_files): + ds = _ds(db_files[:12], crossdb=True) + await ds.invoke_startup() + async with connection_budget() as budget: + for _ in range(3): + assert (await ds.client.get("/-/versions.json")).status_code == 200 + # One :memory: connection for the first request, none after; it is not + # given the --crossdb ATTACHes + assert sum(budget.probes.values()) == 1 + assert not budget.attached + budget.assert_only(probes=1) + ds.close() + + +@pytest.mark.asyncio +async def test_detect_json1_is_cached(db_files): + ds = _ds(db_files[:2]) + await ds.invoke_startup() + await ds.client.get("/db001/t1") + _, budget = await _get(ds, "/db001/t1?_facet=name") + assert not budget.probes + ds.close() + + +@pytest.mark.asyncio +async def test_instance_schema_reads_each_visible_database_once(db_files): + # /-/schema dumps every visible database's sqlite_master: one read each, + # nothing for databases the actor cannot see + ds = _ds( + db_files[:10], + config={"databases": {"db003": {"allow": False}}}, + ) + await ds.invoke_startup() + _, budget = await _get(ds, "/-/schema.json") + expected = {f"db{i:03d}" for i in range(10)} - {"db003"} + assert budget.touched() - {INTERNAL} == expected + assert all(budget.reads[name] == 1 for name in expected) + ds.close() + + +@pytest.mark.asyncio +async def test_autocomplete_debug_is_bounded(db_files): + ds = _ds(db_files) + await ds.invoke_startup() + _, budget = await _get(ds, "/-/debug/autocomplete") + # Only databases holding one of the (at most 5) suggestions + budget.assert_at_most(5) + ds.close() + + +@pytest.mark.asyncio +async def test_execute_write_touches_only_its_database(db_files, tmp_path): + paths = [] + for path in db_files[:20]: + target = tmp_path / os.path.basename(path) + shutil.copy(path, target) + paths.append(str(target)) + ds = _ds(paths) + await ds.invoke_startup() + async with connection_budget() as budget: + response = await ds.client.post( + "/db004/-/execute-write", + actor={"id": "root"}, + data={"sql": "create table new_table (id integer primary key)"}, + ) + assert response.status_code == 200, response.text + assert "Query executed" in response.text + budget.assert_only("db004") + # The catalog already has the new table + rows = await ds.get_internal_database().execute( + "select 1 from catalog_tables where database_name = 'db004' and table_name = 'new_table'" + ) + assert rows.rows + ds.close() + + +@pytest.mark.asyncio +async def test_telemetry_metric_callbacks_open_nothing(db_files): + from datasette import telemetry + + ds = _ds(db_files) + await ds.invoke_startup() + telemetry.register_datasette(ds) + try: + with connection_budget() as budget: + observations = 0 + for callback in ( + telemetry.observe_sql_thread_limit, + telemetry.observe_sql_thread_queue_depth, + telemetry.observe_pending_queries, + telemetry.observe_write_queue_depth, + telemetry.observe_open_connections, + ): + observations += len(list(callback())) + assert observations + budget.assert_only(internal=False) + budget.assert_no_connections(internal=False) + finally: + telemetry.unregister_datasette(ds) + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", [0, 3]) +async def test_startup_opens_each_database_at_most_once(db_files, num_sql_threads): + imm = db_files[:5] + with connection_budget() as budget: + ds = Datasette( + db_files[5:], + immutables=imm, + settings={ + "schema_watch_interval_ms": 0, + "num_sql_threads": num_sql_threads, + }, + ) + await ds._startup_sequence() + # One catalog scan connection per database, no pooled reads or writes + # (immutable table counts are no longer precomputed at startup) + assert set(budget.opened) - {INTERNAL, "_memory"} == { + f"db{i:03d}" for i in range(N) + } + assert all( + count == 1 for name, count in budget.opened.items() if name != INTERNAL + ), budget.opened + assert set(budget.reads) <= {INTERNAL}, budget.reads + assert set(budget.writes) <= {INTERNAL}, budget.writes + assert not budget.attached + ds.close() + + +@pytest.mark.asyncio +async def test_startup_with_crossdb_attaches_nothing(db_files): + with connection_budget() as budget: + ds = Datasette( + db_files[:12], + settings={"schema_watch_interval_ms": 0}, + crossdb=True, + ) + await ds._startup_sequence() + assert not budget.attached + # A crossdb query still sees the attached databases + response = await ds.client.get( + "/_memory/-/query.json?sql=select+count(*)+from+db001.t1&_shape=array" + ) + assert response.status_code == 200 + ds.close() + + +@pytest.mark.asyncio +async def test_restart_with_persisted_catalog_opens_nothing( + db_files, tmp_path, monkeypatch +): + from datasette import schema_watcher + + # The files' ctimes are recent (os.utime() in the fixture): without this + # their stored fingerprints would be racy, and so rescanned + monkeypatch.setattr(schema_watcher, "RACY_WINDOW_NS", 0) + internal = str(tmp_path / "internal.db") + ds = Datasette(db_files, internal=internal) + await ds.invoke_startup() + ds.close() + with connection_budget() as budget: + ds = Datasette( + db_files, internal=internal, settings={"schema_watch_interval_ms": 0} + ) + await ds._startup_sequence() + response = await ds.client.get("/.json") + assert response.status_code == 200 + budget.assert_only() + ds.close() + + +@pytest.mark.asyncio +async def test_idle_sweeps_open_nothing(db_files): + ds = _ds(db_files) + await ds.invoke_startup() + with connection_budget() as budget: + for _ in range(3): + await ds._refresh_schemas(background=True) + await ds.refresh_schemas(force=True) + budget.assert_only(internal=False) + budget.assert_no_connections(internal=False) + ds.close() + + +@pytest.mark.asyncio +async def test_failed_scan_is_not_retried_by_idle_sweeps( + db_files, tmp_path, monkeypatch +): + from datasette import schema_watcher + + # The file was just written: do not wait out the racy window + monkeypatch.setattr(schema_watcher, "RACY_WINDOW_NS", 0) + bad = tmp_path / "bad.db" + bad.write_bytes(b"this is not a SQLite database" * 10) + old = time.time() - 60 + os.utime(bad, (old, old)) + ds = _ds(db_files[:3] + [str(bad)]) + await ds.invoke_startup() + assert ds.get_database("bad")._watch_state.error is not None + with connection_budget() as budget: + for _ in range(3): + await ds._refresh_schemas(background=True) + assert "bad" not in budget.touched() + # ...until the file changes + good = sqlite3.connect(tmp_path / "good.db") + good.execute("create table fixed (id integer primary key)") + good.commit() + good.close() + os.replace(tmp_path / "good.db", bad) + await ds._refresh_schemas(background=True) + rows = await ds.get_internal_database().execute( + "select table_name from catalog_tables where database_name = 'bad'" + ) + assert [r[0] for r in rows.rows] == ["fixed"] + ds.close() + + +@pytest.mark.asyncio +async def test_check_databases_reuses_catalog_scan(db_files, tmp_path): + from datasette.cli import check_databases + from datasette.utils import ConnectionProblem + + ds = _ds(db_files) + with connection_budget() as budget: + await check_databases(ds) + # The catalog scan opened each database once; no second pass + assert all(count == 1 for name, count in budget.opened.items() if name != INTERNAL) + assert set(budget.reads) <= {INTERNAL, "_memory"}, budget.reads + ds.close() + + bad = tmp_path / "bad.db" + bad.write_bytes(b"this is not a SQLite database" * 10) + ds = _ds(db_files[:3] + [str(bad)]) + with pytest.raises(Exception) as excinfo: + await check_databases(ds) + assert "bad.db" in str(excinfo.value) + assert "file is not a database" in str(excinfo.value) + assert not isinstance(excinfo.value, ConnectionProblem) + ds.close() + + +# Parity: the catalog helpers answer exactly what the live helpers do + + +MISC_SQL = """ + create table docs (id integer primary key, title text, body text); + create virtual table docs_fts using fts4(title, body, content="docs"); + create virtual table docs_vocab using fts4aux(docs_fts); + create virtual table plain using fts5(x); + create virtual table geo using rtree(id, minx, maxx); + create table _private (x); + create table geometry_columns (f_table_name text); + create table idx_foo_bar (x); + create table IDXA (x); + create table parent (id integer primary key); + create table child ( + id integer primary key, + p integer references parent(id), + q integer references missing(id) + ); + create table comp (a, b, foreign key (a, b) references parent(id, id)); + create table self_ref (id integer primary key, up integer references self_ref(id)); + create table wide (b text, a integer, c, primary key (a, b)); + create view v_docs as select * from docs; +""" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("name", ["misc", "fixtures"]) +async def test_catalog_helpers_match_live_introspection(tmp_path, name): + from datasette.fixtures import write_fixture_database + from datasette.utils.catalog import ( + all_derived_table_dependencies, + catalog_all_foreign_keys, + ) + + path = str(tmp_path / f"{name}.db") + if name == "fixtures": + write_fixture_database(path) + config_hidden = set() + config = {} + else: + conn = sqlite3.connect(path) + conn.executescript(MISC_SQL) + conn.close() + config_hidden = {"docs"} + config = {"databases": {"misc": {"tables": {"docs": {"hidden": True}}}}} + ds = Datasette([path], config=config, settings={"schema_watch_interval_ms": 0}) + await ds.invoke_startup() + db = ds.get_database(name) + tables = await db.table_names() + summary = (await catalog_summaries(ds, [name]))[name] + # hidden_table_names() also lists SpatiaLite names that do not exist + hidden = set(await db.hidden_table_names()) & set(tables) + assert {t for t in tables if summary.is_hidden(t, config_hidden)} == hidden + assert summary.views == set(await db.view_names()) + + details = await catalog_table_details(ds, [(name, t) for t in tables]) + for table in tables: + assert details[(name, table)]["columns"] == await db.table_columns(table) + assert details[(name, table)]["primary_keys"] == await db.primary_keys(table) + assert details[(name, table)]["fts_table"] == await db.fts_table(table) + + all_foreign_keys = await db.get_all_foreign_keys() + assert await catalog_all_foreign_keys(ds, name) == all_foreign_keys + relationships = (await catalog_relationship_counts(ds, [name])).get(name, {}) + for table, fks in all_foreign_keys.items(): + assert relationships.get(table, 0) == len(fks["incoming"] + fks["outgoing"]) + + derived = await all_derived_table_dependencies(ds) + live = await db.execute_fn(sqlite_derived_table_dependencies) + assert derived.get(name, {}) == live + # Restored from the catalog rather than remembered from the scan + db._cached_derived_table_dependencies = None + assert await db.derived_table_dependencies() == live + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.skipif(not supports_table_list(), reason="needs PRAGMA table_list") +async def test_catalog_records_table_types(tmp_path): + path = str(tmp_path / "types.db") + conn = sqlite3.connect(path) + conn.executescript(""" + create table plain (id integer primary key); + create virtual table f using fts5(x); + """) + conn.close() + ds = Datasette([path], settings={"schema_watch_interval_ms": 0}) + await ds.invoke_startup() + rows = await ds.get_internal_database().execute( + "select table_name, type from catalog_tables where database_name = 'types'" + ) + types = dict(rows.rows) + assert types["plain"] == "table" + assert types["f"] == "virtual" + assert types["f_data"] == "shadow" + ds.close() + + +@pytest.mark.asyncio +async def test_catalog_derived_dependencies_follow_catalog_changes(tmp_path): + path = str(tmp_path / "live.db") + sqlite3.connect(path).execute("create table docs (body text)").connection.close() + ds = Datasette([path], settings={"schema_watch_interval_ms": 0}) + await ds.invoke_startup() + from datasette.utils.catalog import all_derived_table_dependencies + + assert "live" not in await all_derived_table_dependencies(ds) + db = ds.get_database("live") + await db.execute_write( + "create virtual table docs_fts using fts5(body, content='docs')" + ) + derived = await all_derived_table_dependencies(ds) + assert derived["live"]["docs_fts"] == "docs" + assert derived["live"]["docs_fts_data"] == "docs_fts" + ds.close() + + +def test_connection_budget_helper_counts(db_files): + # The helper itself: a direct execute_fn is a read, connect() an open + import asyncio + + async def run(): + ds = _ds(db_files[:2]) + await ds.invoke_startup() + db = ds.get_database("db001") + async with connection_budget() as budget: + await db.execute("select 1") + conn = sqlite3.connect(":memory:") + conn.close() + assert budget.reads["db001"] == 1 + assert budget.touched() == {"db001"} + with pytest.raises(AssertionError, match="db001"): + budget.assert_only("db000") + ds.close() + return budget + + budget = asyncio.run(run()) + assert sum(budget.probes.values()) == 1 + assert json.dumps(budget.summary()) diff --git a/tests/test_connection_leases.py b/tests/test_connection_leases.py new file mode 100644 index 00000000..550ea280 --- /dev/null +++ b/tests/test_connection_leases.py @@ -0,0 +1,283 @@ +""" +Connections handed to execute_write_fn(), the SQL write methods and +execute_isolated_fn() callbacks are lease proxies, like read connections: +they stop working once the callback returns. +""" + +import asyncio +import sqlite3 +import threading +import time + +import pytest +import sqlite_utils + +from datasette.app import Datasette +from datasette.database import ConnectionLeaseError + + +def _make_db(tmp_path, name="leases"): + path = str(tmp_path / f"{name}.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() + return path + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", [0, 2]) +@pytest.mark.parametrize("transaction", [True, False]) +async def test_write_callback_connection_expires( + tmp_path, num_sql_threads, transaction +): + ds = Datasette([_make_db(tmp_path)], settings={"num_sql_threads": num_sql_threads}) + db = ds.get_database("leases") + stash = {} + + def write(conn): + stash["conn"] = conn + assert isinstance(conn, sqlite3.Connection) + conn.execute("insert into t (v) values ('y')") + if not transaction: + # transaction=False DML is not committed for you (unchanged) + conn.commit() + return conn.execute("select count(*) from t").fetchone()[0] + + assert await db.execute_write_fn(write, transaction=transaction) == 2 + with pytest.raises(ConnectionLeaseError, match="execute_write_fn"): + stash["conn"].execute("select 1") + with pytest.raises(ConnectionLeaseError): + stash["conn"].row_factory = None + # The write connection itself is fine for the next write + await db.execute_write("insert into t (v) values ('z')") + assert (await db.execute("select count(*) from t")).single_value() == 3 + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", [0, 2]) +async def test_write_callback_may_return_cursor(tmp_path, num_sql_threads): + # A common plugin pattern: lambda conn: conn.execute(...). The returned + # cursor's attributes stay readable; only the connection is leased. + ds = Datasette([_make_db(tmp_path)], settings={"num_sql_threads": num_sql_threads}) + db = ds.get_database("leases") + cursor = await db.execute_write_fn( + lambda conn: conn.execute("insert into t (v) values ('y')") + ) + assert cursor.lastrowid == 2 + assert cursor.rowcount == 1 + many = await db.execute_write_many("insert into t (v) values (?)", [["a"], ["b"]]) + assert many.rowcount == 2 + await db.execute_write_script("insert into t (v) values ('c');") + assert (await db.execute("select count(*) from t")).single_value() == 5 + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", [0, 2]) +async def test_with_conn_inside_write_callback_yields_proxy(tmp_path, num_sql_threads): + ds = Datasette([_make_db(tmp_path)], settings={"num_sql_threads": num_sql_threads}) + db = ds.get_database("leases") + stash = {} + + def write(conn): + with conn as inner: + stash["inner"] = inner + inner.execute("insert into t (v) values ('y')") + + await db.execute_write_fn(write, transaction=False) + with pytest.raises(ConnectionLeaseError): + stash["inner"].execute("select 1") + assert (await db.execute("select count(*) from t")).single_value() == 2 + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", [0, 2]) +async def test_sqlite_utils_on_write_lease(tmp_path, num_sql_threads): + ds = Datasette([_make_db(tmp_path)], settings={"num_sql_threads": num_sql_threads}) + db = ds.get_database("leases") + + def write(conn): + sdb = sqlite_utils.Database(conn) + sdb["dogs"].insert_all([{"id": 1, "name": "Cleo"}], pk="id") + sdb["dogs"].transform(rename={"name": "dog_name"}) + return sdb["dogs"].columns_dict + + assert await db.execute_write_fn(write) == {"id": int, "dog_name": str} + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", [0, 2]) +@pytest.mark.parametrize("immutable", [False, True]) +async def test_isolated_callback_connection_expires( + tmp_path, num_sql_threads, immutable +): + path = _make_db(tmp_path) + if immutable: + ds = Datasette(immutables=[path], settings={"num_sql_threads": num_sql_threads}) + else: + ds = Datasette([path], settings={"num_sql_threads": num_sql_threads}) + db = ds.get_database("leases") + stash = {} + + def isolated(conn): + stash["conn"] = conn + return conn.execute("select count(*) from t").fetchone()[0] + + assert await db.execute_isolated_fn(isolated) == 1 + with pytest.raises(ConnectionLeaseError, match="execute_isolated_fn"): + stash["conn"].execute("select 1") + ds.close() + + +@pytest.mark.asyncio +async def test_write_wrapper_receives_lease(tmp_path): + from datasette import hookimpl + from datasette.plugins import pm + + stash = {} + + class WrapperPlugin: + __name__ = "WrapperPlugin" + + @hookimpl + def write_wrapper(self, datasette, database, request, transaction): + if database != "leases": + return None + + def wrapper(conn): + stash["conn"] = conn + yield + + return wrapper + + pm.register(WrapperPlugin(), name="test_write_wrapper_receives_lease") + try: + ds = Datasette([_make_db(tmp_path)]) + db = ds.get_database("leases") + await db.execute_write("insert into t (v) values ('y')") + with pytest.raises(ConnectionLeaseError): + stash["conn"].execute("select 1") + ds.close() + finally: + pm.unregister(name="test_write_wrapper_receives_lease") + + +@pytest.mark.asyncio +async def test_close_waits_for_immutable_isolated_fn(tmp_path): + # On main this ran on the executor untracked, so close() could close + # the connection while the callback was still stepping it (a segfault + # with a progress handler installed) + ds = Datasette(immutables=[_make_db(tmp_path)], settings={"num_sql_threads": 2}) + db = ds.get_database("leases") + started = threading.Event() + + def slow(conn): + started.set() + time.sleep(0.3) + return conn.execute("select count(*) from t").fetchone()[0] + + task = asyncio.ensure_future(db.execute_isolated_fn(slow)) + assert await asyncio.to_thread(started.wait, 5) + assert len(db._pending_execute_futures) == 1 + db.close() + assert await task == 1 + ds.close() + + +NO_THREADS_SCRIPT = r""" +import asyncio, os, sqlite3, sys, threading + +def no_threads(self): + raise RuntimeError("can't start new thread") + +threading.Thread.start = no_threads + +from datasette.app import Datasette +from datasette.database import Database + +path, scratch = sys.argv[1], sys.argv[2] +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() + + +async def main(): + ds = Datasette( + [path], + settings={ + "num_sql_threads": 0, + "connection_idle_timeout_ms": 50, + "schema_watch_interval_ms": 50, + }, + ) + await ds.invoke_startup() + db = ds.get_database("nothreads") + assert (await db.execute("select count(*) from t")).single_value() == 1 + await db.execute_write("insert into t (v) values ('y')") + await db.execute_write("create table t2 (id)") + tables = await ds.get_internal_database().execute( + "select table_name from catalog_tables where database_name = 'nothreads' " + "order by table_name" + ) + assert [r[0] for r in tables.rows] == ["t", "t2"], tables.rows + assert await db.execute_isolated_fn( + lambda conn: conn.execute("select count(*) from t").fetchone()[0] + ) == 2 + await db.execute_write("insert into t (v) values ('z')", block=False) + response = await ds.client.get("/nothreads/t.json?_shape=array") + assert response.status_code == 200, response.text + assert len(response.json()) == 3 + # External change detected by the polling task, inline on the loop + other = sqlite3.connect(path) + other.execute("create table external_t (id)") + other.commit() + other.close() + for _ in range(100): + await asyncio.sleep(0.02) + tables = await ds.get_internal_database().execute( + "select 1 from catalog_tables where table_name = 'external_t'" + ) + if tables.rows: + break + else: + raise AssertionError("external change not seen") + # Scratch database added later, created by its first write + sdb = ds.add_database(Database(ds, path=scratch), name="scratch") + await sdb.execute_write("create table s (id)") + assert os.path.exists(scratch) + ds.close() + print("ok", threading.active_count()) + + +asyncio.run(main()) +""" + + +def test_num_sql_threads_zero_starts_no_threads(tmp_path): + # Simulates Pyodide, where no thread can be started: nothing in the + # read pool, write path or schema watcher may need one + import subprocess + import sys + + result = subprocess.run( + [ + sys.executable, + "-c", + NO_THREADS_SCRIPT, + str(tmp_path / "nothreads.db"), + str(tmp_path / "scratch.db"), + ], + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + assert result.stdout.strip() == "ok 1" diff --git a/tests/test_connection_limits.py b/tests/test_connection_limits.py new file mode 100644 index 00000000..5294bc7c --- /dev/null +++ b/tests/test_connection_limits.py @@ -0,0 +1,92 @@ +"""Exercise real HTTP sockets and WAL files under a small OS descriptor limit.""" + +import json +import subprocess +import sys + +import pytest + + +@pytest.mark.skipif(sys.platform == "win32", reason="Requires RLIMIT_NOFILE") +def test_500_wal_databases_with_http_and_writes_under_256_fds(): + script = r""" +import asyncio, concurrent.futures, json, resource, socket, sqlite3, tempfile, threading, time, urllib.request +from pathlib import Path +import psutil +import uvicorn +from datasette.app import Datasette + +async def main(): + with tempfile.TemporaryDirectory() as directory: + paths = [] + for i in range(500): + path = str(Path(directory) / f"db{i}.db") + conn = sqlite3.connect(path) + conn.execute("pragma journal_mode=wal") + conn.execute("create table t(id integer primary key)") + conn.close() + paths.append(path) + _, hard = resource.getrlimit(resource.RLIMIT_NOFILE) + resource.setrlimit(resource.RLIMIT_NOFILE, (256, hard)) + ds = Datasette(paths) + await ds.invoke_startup() + process = psutil.Process() + baseline_rss = process.memory_info().rss + peak = {"fds": 0, "rss": 0, "writers": 0} + stopped = threading.Event() + def monitor(): + while not stopped.wait(0.002): + peak["fds"] = max(peak["fds"], process.num_fds()) + peak["rss"] = max(peak["rss"], process.memory_info().rss) + peak["writers"] = max(peak["writers"], ds._write_budget.snapshot()["writers"]) + watcher = threading.Thread(target=monitor) + watcher.start() + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + server = uvicorn.Server(uvicorn.Config(ds.app(), log_level="error")) + thread = threading.Thread(target=lambda: server.run(sockets=[sock])) + thread.start() + try: + deadline = time.monotonic() + 10 + while not server.started: + assert time.monotonic() < deadline + await asyncio.sleep(0.01) + def request(i): + with urllib.request.urlopen(f"http://127.0.0.1:{port}/db{i}/-/query.json?sql=select+count(*)+from+t", timeout=20) as response: + assert response.status == 200 + assert json.load(response)["ok"] + def reads(): + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as clients: + list(clients.map(request, range(500))) + async def writes(): + for start in range(0, 500, 16): + await asyncio.gather(*(ds.get_database(f"db{i}").execute_write("insert into t values (1)") for i in range(start, min(start + 16, 500)))) + await asyncio.gather(asyncio.to_thread(reads), writes()) + for i in range(500): + assert (await ds.get_database(f"db{i}").execute("select count(*) from t")).single_value() == 1 + assert ds._read_pool.stats["peak_open"] <= 32 + assert ds._write_budget.snapshot()["stats"]["peak_writers"] <= 8 + print(json.dumps({"peak": peak, "baseline_rss": baseline_rss, "writes": 500, "http_reads": 500}), flush=True) + finally: + server.should_exit = True + await asyncio.to_thread(thread.join, 15) + assert not thread.is_alive() + stopped.set() + watcher.join() + sock.close() + ds.close() +asyncio.run(main()) +""" + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=90, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + measurements = json.loads(result.stdout.strip().splitlines()[-1]) + assert measurements["peak"]["fds"] < 256 + assert measurements["peak"]["writers"] <= 8 + print(measurements) diff --git a/tests/test_connection_recovery.py b/tests/test_connection_recovery.py new file mode 100644 index 00000000..373f747f --- /dev/null +++ b/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() diff --git a/tests/test_connection_resource_leases.py b/tests/test_connection_resource_leases.py new file mode 100644 index 00000000..55267f97 --- /dev/null +++ b/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() diff --git a/tests/test_connection_review_fixes.py b/tests/test_connection_review_fixes.py new file mode 100644 index 00000000..3a3bad91 --- /dev/null +++ b/tests/test_connection_review_fixes.py @@ -0,0 +1,862 @@ +""" +Regression tests for problems found by the second review of the connection +management redesign (concurrency, lifecycle and public-contract findings). +Each test names the problem it pins down. +""" + +import asyncio +import gc +import os +import sqlite3 +import threading +import time +import weakref + +import pytest + +from datasette import schema_watcher as schema_watcher_module +from datasette.app import Datasette +from datasette.connection_pool import ConnectionLeaseError +from datasette.database import Database + + +def make_db(path, sql="create table t (id integer primary key, v text)"): + conn = sqlite3.connect(path) + conn.executescript(sql) + conn.commit() + conn.close() + + +async def catalog_tables(ds, name): + rows = await ds.get_internal_database().execute( + "select table_name from catalog_tables where database_name = ? " + "order by table_name", + [name], + ) + return [r[0] for r in rows.rows] + + +def live_tables(path): + conn = sqlite3.connect(path) + try: + return sorted( + r[0] + for r in conn.execute( + "select name from sqlite_master where type = 'table' " + "and name not like 'sqlite_%'" + ) + ) + finally: + conn.close() + + +def run_in_threads(*coroutine_functions, timeout=20): + """Run each coroutine function with asyncio.run() in its own thread: + several event loops driving one Datasette at the same time.""" + results = [None] * len(coroutine_functions) + + def runner(i, fn): + try: + results[i] = ("ok", asyncio.run(fn())) + except BaseException as e: # noqa: BLE001 + results[i] = ("error", e) + + threads = [ + threading.Thread(target=runner, args=(i, fn), daemon=True) + for i, fn in enumerate(coroutine_functions) + ] + for thread in threads: + thread.start() + deadline = time.monotonic() + timeout + for thread in threads: + thread.join(max(0.0, deadline - time.monotonic())) + hung = [i for i, thread in enumerate(threads) if thread.is_alive()] + assert not hung, f"event loop threads {hung} did not finish within {timeout}s" + for status, value in results: + if status == "error": + raise value + return [value for _, value in results] + + +@pytest.fixture +def owned_db(tmp_path): + path = str(tmp_path / "data.db") + make_db(path) + ds = Datasette(settings={"schema_watch_interval_ms": 0}) + db = ds.add_database(Database(ds, path=path), name="data", schema_watch="owned") + asyncio.run(ds.invoke_startup()) + yield ds, db, path + ds.close() + + +# ---------------------------------------------------------------------- +# Concurrency and lifecycle +# ---------------------------------------------------------------------- + + +def test_cross_loop_ddl_writes_do_not_hang(owned_db, monkeypatch): + # C1: refresh() treated a scan started on another *live* event loop as + # abandoned, stole it, and each loop then resolved or cleared the + # other's future: blocking writes on one loop never returned + ds, db, path = owned_db + original = schema_watcher_module.SchemaWatcher._scan_chunk_sync + + def slow_scan(self, states): + time.sleep(0.2) + return original(self, states) + + monkeypatch.setattr( + schema_watcher_module.SchemaWatcher, "_scan_chunk_sync", slow_scan + ) + + def writer(prefix, count): + async def run(): + for i in range(count): + await asyncio.wait_for( + db.execute_write(f"create table {prefix}{i} (id integer)"), 10 + ) + + return run + + def inserter(): + async def run(): + async def worker(): + for _ in range(20): + await asyncio.wait_for( + db.execute_write("insert into t (v) values ('x')"), 10 + ) + + await asyncio.gather(*(worker() for _ in range(3))) + + return run + + run_in_threads(writer("a", 3), writer("b", 3), inserter()) + + async def check(): + return await catalog_tables(ds, "data") + + assert asyncio.run(check()) == live_tables(path) + state = db._watch_state + assert state.scan_future is None + assert not state.needs_scan + + +def test_scan_from_stopped_loop_is_taken_over(owned_db): + # The rule for taking over a scan: its loop has closed or stopped, so + # nothing will finish it. A loop stopped by run_until_complete() + # returning while a scan task is pending must not hang other loops. + ds, db, _path = owned_db + loop = asyncio.new_event_loop() + try: + gate = threading.Event() + original = schema_watcher_module.SchemaWatcher._scan_chunk_sync + + def gated_scan(self, states): + gate.wait(5) + return original(self, states) + + schema_watcher_module.SchemaWatcher._scan_chunk_sync = gated_scan + try: + + async def start_scan(): + # block=False: the scan task is spawned on this loop and + # left pending when run_until_complete() returns + await db.execute_write( + "create table stopped1 (id integer)", block=False + ) + await asyncio.sleep(0.2) + + loop.run_until_complete(start_scan()) + finally: + schema_watcher_module.SchemaWatcher._scan_chunk_sync = original + gate.set() + + async def other_loop(): + await asyncio.wait_for( + db.execute_write("create table other1 (id integer)"), 10 + ) + return await catalog_tables(ds, "data") + + tables = run_in_threads(other_loop)[0] + assert {"stopped1", "other1"} <= set(tables) + finally: + loop.close() + + +def test_one_poller_across_live_event_loops(tmp_path): + # C2: every switch between two live event loops started another + # poller task, each sweeping every interval + path = str(tmp_path / "ext.db") + make_db(path) + ds = Datasette([path], settings={"schema_watch_interval_ms": 100}) + asyncio.run(ds.invoke_startup()) + watcher = ds._schema_watcher + turns = [threading.Event(), threading.Event()] + done = threading.Event() + counts = {} + + def client(me, other, rounds): + async def run(): + for _ in range(rounds): + turns[me].wait(5) + turns[me].clear() + assert (await ds.client.get("/ext.json")).status_code == 200 + turns[other].set() + if me == 0: + sweeps = watcher.counters["sweeps"] + await asyncio.sleep(1.0) + counts["sweeps_per_s"] = watcher.counters["sweeps"] - sweeps + counts["pollers"] = len([t for t in watcher._pollers if not t.done()]) + done.set() + else: + done.wait(10) + + return run + + turns[0].set() + run_in_threads(client(0, 1, 10), client(1, 0, 10)) + assert counts["pollers"] == 1 + # One poller at a 100ms interval: about 10 sweeps a second, not 10 per + # loop switch + assert counts["sweeps_per_s"] <= 15 + ds.close() + + +@pytest.mark.asyncio +async def test_failed_catalog_write_is_retried(owned_db, monkeypatch): + # C3: refresh() cleared needs_scan before scanning and never restored + # it, so an owned database lost a catalog update for good if the + # catalog write failed (internal database locked) + ds, db, _path = owned_db + original = schema_watcher_module.write_catalog_entries + locked = threading.Event() + locked.set() + + def flaky(conn, entries): + if locked.is_set(): + raise sqlite3.OperationalError("database is locked") + return original(conn, entries) + + monkeypatch.setattr(schema_watcher_module, "write_catalog_entries", flaky) + await db.execute_write("create table added_while_locked (id integer)") + # Let the refresh task spawned by the write's notification fail too + for _ in range(5): + await asyncio.sleep(0.01) + state = db._watch_state + assert state.needs_scan + assert state in ds._schema_watcher._pending + assert await catalog_tables(ds, "data") == ["t"] + # While still locked, a request is not failed by the catalog update + state.retry_at = 0.0 + assert (await ds.client.get("/data.json")).status_code == 200 + assert state.retry_at > time.monotonic() + # Once unlocked, the next request retries (after the back-off) + locked.clear() + state.retry_at = 0.0 + await ds.client.get("/data.json") + assert await catalog_tables(ds, "data") == ["added_while_locked", "t"] + response = await ds.client.get("/data.json") + assert {t["name"] for t in response.json()["tables"]} == { + "added_while_locked", + "t", + } + + +def test_block_false_write_from_closed_loop_reaches_catalog(owned_db): + # C3: the schema-change notification went to the caller's loop, which + # had closed: notified_version was recorded anyway and the change was + # never picked up + ds, db, path = owned_db + + async def write(): + await db.execute_write( + "create table added_nonblocking (id integer)", block=False + ) + + asyncio.run(write()) + # Let the write thread finish (it delivers to a closed loop) + deadline = time.monotonic() + 5 + while "added_nonblocking" not in live_tables(path): + assert time.monotonic() < deadline + time.sleep(0.01) + time.sleep(0.05) + + async def later(): + await ds.client.get("/data.json") + return await catalog_tables(ds, "data") + + assert asyncio.run(later()) == ["added_nonblocking", "t"] + + +def test_unclosed_datasette_is_collected(): + # C4: the temp internal database's atexit.register(bound method) pinned + # every Datasette even after the caller released its last reference. + ds = Datasette() + ref = weakref.ref(ds) + del ds + gc.collect() + assert ref() is None + + +def test_closed_datasette_is_collected(tmp_path): + # C4 + C5: atexit and the read-pool reaper's loop variable both kept a + # closed Datasette (and every Database, cache and plugin state) alive + path = str(tmp_path / "data.db") + make_db(path) + ds = Datasette([path], settings={"connection_idle_timeout_ms": 20}) + + async def use(ds): + await ds.invoke_startup() + await ds.get_database("data").execute("select 1") + await ds.get_database("data").execute_write("insert into t (v) values (1)") + + asyncio.run(use(ds)) + # Let the reaper scan (and close) the idle connection + time.sleep(0.3) + ds.close() + ref = weakref.ref(ds) + del ds + gc.collect() + assert ref() is None + + +@pytest.mark.asyncio +async def test_write_callback_cannot_close_write_connection(owned_db): + # C6: closing the write lease closed the write thread's connection; + # every later write failed until the thread idled out + _ds, db, _path = owned_db + + def close_it(conn): + conn.close() + + with pytest.raises(ConnectionLeaseError): + await db.execute_write_fn(close_it) + await db.execute_write("insert into t (v) values ('after')") + + # cursor.connection has the same lease protections as the connection. + def close_raw(conn): + conn.execute("select 1").connection.close() + + with pytest.raises(ConnectionLeaseError): + await db.execute_write_fn(close_raw) + for _ in range(3): + await db.execute_write("insert into t (v) values ('again')") + rows = await db.execute("select count(*) from t") + assert rows.first()[0] == 4 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", (3, 0)) +async def test_isolated_lease_close_is_allowed(tmp_path, num_sql_threads): + path = str(tmp_path / "data.db") + make_db(path) + ds = Datasette([path], settings={"num_sql_threads": num_sql_threads}) + db = ds.get_database("data") + await db.execute_isolated_fn(lambda conn: conn.close()) + await db.execute_write("insert into t (v) values (1)") + ds.close() + + +class Abort(BaseException): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("exception_class", (Abort, SystemExit)) +async def test_write_callback_base_exception_is_delivered(owned_db, exception_class): + # C7: a BaseException from a write callback killed the write thread + # without answering the caller, who waited forever + _ds, db, _path = owned_db + + def boom(conn): + raise exception_class("stop") + + with pytest.raises(RuntimeError) as excinfo: + await asyncio.wait_for(db.execute_write_fn(boom), 5) + assert isinstance(excinfo.value.__cause__, exception_class) + # Same for an isolated function + with pytest.raises(RuntimeError): + await asyncio.wait_for(db.execute_isolated_fn(boom), 5) + await asyncio.wait_for(db.execute_write("insert into t (v) values (1)"), 5) + + +# ---------------------------------------------------------------------- +# Public contract +# ---------------------------------------------------------------------- + + +def _fts_setup(path): + make_db(path, "create table secret (id integer primary key, body text)") + conn = sqlite3.connect(path) + conn.execute("insert into secret values (1, 'TOP SECRET PAYLOAD')") + conn.commit() + conn.close() + + +def _add_fts(path): + conn = sqlite3.connect(path) + conn.executescript(""" + create virtual table secret_fts using fts5(body, content='secret', content_rowid='id'); + insert into secret_fts(secret_fts) values ('rebuild'); + """) + conn.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mode", ("external", "owned")) +async def test_derived_table_rule_uses_current_schema(tmp_path, mode): + # Contract C1: the "external-content FTS tables require access to their + # content table" rule was decided from a catalog that had not seen an + # FTS table created outside Datasette: its rows were served to + # anonymous users (until the next poll, or for good for owned DBs) + path = str(tmp_path / "data.db") + _fts_setup(path) + ds = Datasette( + config={ + "databases": {"data": {"tables": {"secret": {"allow": {"id": "root"}}}}} + }, + settings={"default_allow_sql": False, "schema_watch_interval_ms": 0}, + ) + ds.add_database(Database(ds, path=path), name="data", schema_watch=mode) + await ds.invoke_startup() + assert (await ds.client.get("/data/secret.json")).status_code == 403 + _add_fts(path) + response = await ds.client.get("/data/secret_fts.json") + assert response.status_code == 403 + assert "TOP SECRET" not in response.text + # Finding the catalog behind schedules a rescan + await ds.client.get("/data.json") + assert "secret_fts" in await catalog_tables(ds, "data") + ds.close() + + +@pytest.mark.asyncio +async def test_add_database_file_changes_are_picked_up(tmp_path): + # Contract C2: add_database() defaulted to "owned" (never polled), so a + # table created by another process never appeared in listings + path = str(tmp_path / "added.db") + make_db(path) + ds = Datasette(settings={"schema_watch_interval_ms": 50}) + await ds.invoke_startup() + ds.add_database(Database(ds, path=path), name="added") + await ds.client.get("/added.json") + make_db(path, "create table t_ext (id integer)") + + async def listed(): + response = await ds.client.get("/added.json") + return {t["name"] for t in response.json()["tables"]} + + deadline = time.monotonic() + 5 + while "t_ext" not in await listed(): + assert time.monotonic() < deadline, "t_ext never appeared" + await asyncio.sleep(0.05) + ds.close() + + +@pytest.mark.asyncio +async def test_named_memory_database_out_of_band_changes(tmp_path): + # Contract C2: named in-memory databases could not be watched at all; + # datasette-app-support restores one with backup() into + # db.connect(write=True) and then shows it + ds = Datasette(settings={"schema_watch_interval_ms": 50}) + await ds.invoke_startup() + memory = ds.add_memory_database("review_fixes_temporary", name="temporary") + await memory.execute_write("create table m0 (id integer)") + source_path = str(tmp_path / "backup.db") + make_db(source_path, "create table restored_table (id integer)") + source = sqlite3.connect(source_path) + target = memory.connect(write=True) + source.backup(target) + source.close() + target.close() + + async def listed(): + response = await ds.client.get("/temporary.json") + return {t["name"] for t in response.json()["tables"]} + + deadline = time.monotonic() + 5 + while "restored_table" not in await listed(): + assert time.monotonic() < deadline, "restored_table never appeared" + await asyncio.sleep(0.05) + ds.close() + + +FTS_SETUP = """ +create table docs (id integer primary key, body text); +create virtual table docs_fts using fts5(body, content='docs', content_rowid='id'); +create trigger docs_ai after insert on docs begin + insert into docs_fts (rowid, body) values (new.id, new.body); +end; +create trigger docs_ad after delete on docs begin + insert into docs_fts (docs_fts, rowid, body) values ('delete', old.id, old.body); +end; +insert into docs (id, body) values (1, 'oldword'); +""" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", (3, 0)) +async def test_write_connection_pins_recursive_triggers(tmp_path, num_sql_threads): + # Contract C4: INSERT OR REPLACE on a sqlite-utils style FTS table only + # cleaned up the old index entry once some code had wrapped the write + # connection in sqlite_utils.Database (recursive_triggers drift) - and + # the idle write thread's fresh connections lost the drift again + path = str(tmp_path / "fts.db") + make_db(path, FTS_SETUP) + ds = Datasette( + [path], + settings={ + "num_sql_threads": num_sql_threads, + "connection_idle_timeout_ms": 50, + }, + ) + db = ds.get_database("fts") + for word in ("newword", "newerword"): + flag = await db.execute_write_fn( + lambda conn: conn.execute("PRAGMA recursive_triggers").fetchone()[0] + ) + assert flag == 1 + await db.execute_write( + "insert or replace into docs (id, body) values (1, ?)", [word] + ) + # Let the idle write thread exit: the next write opens a new + # connection + await asyncio.sleep(0.3) + stale = await db.execute( + "select rowid from docs_fts where docs_fts match 'oldword OR newword'" + ) + assert stale.rows == [] + ds.close() + + +@pytest.mark.asyncio +async def test_new_memory_database_is_not_scanned_at_registration(): + # Contract C5: add_memory_database() after startup scanned the + # shared-cache database straight away, racing plugins that fill it from + # their own connection (datasette-copy-to-memory's VACUUM INTO failed + # with "database table is locked") + ds = Datasette(settings={"schema_watch_interval_ms": 0}) + await ds.invoke_startup() + watcher = ds._schema_watcher + scans = watcher.counters["scans"] + memory = ds.add_memory_database("review_fixes_copy", name="copy") + await memory.execute("select 1 + 1") + await asyncio.sleep(0.05) + assert watcher.counters["scans"] == scans + state = memory._watch_state + assert state.scan_future is None + assert state in watcher._pending + # Fill it from a connection of its own, as the plugin does. (The plugin + # uses VACUUM INTO from a connection opened with uri=True; a plain + # connection is simpler and does not depend on whether the SQLite build + # treats URI file names as URIs without that flag) + conn = sqlite3.connect("file:review_fixes_copy?mode=memory&cache=shared", uri=True) + conn.execute("create table copied (id integer)") + conn.commit() + conn.close() + # The next request scans it + response = await ds.client.get("/copy.json") + assert {t["name"] for t in response.json()["tables"]} == {"copied"} + ds.close() + + +@pytest.mark.asyncio +async def test_locked_memory_scan_is_retried(monkeypatch): + ds = Datasette(settings={"schema_watch_interval_ms": 0}) + await ds.invoke_startup() + memory = ds.add_memory_database("review_fixes_locked", name="locked") + original = schema_watcher_module.collect_schema + locked = threading.Event() + locked.set() + + def locked_scan(conn, name): + if locked.is_set(): + raise sqlite3.OperationalError("database table is locked: sqlite_master") + return original(conn, name) + + monkeypatch.setattr(schema_watcher_module, "collect_schema", locked_scan) + await memory.execute_write("create table m (id integer)") + for _ in range(5): + await asyncio.sleep(0.01) + state = memory._watch_state + # Transient: not recorded as a broken database + assert state.error is None + assert state.needs_scan + locked.clear() + state.retry_at = 0.0 + await ds.client.get("/locked.json") + assert await catalog_tables(ds, "locked") == ["m"] + ds.close() + + +@pytest.mark.asyncio +async def test_stale_owned_catalog_is_not_persisted_at_close(tmp_path): + # Contract C6: close() stored a fresh fingerprint for an owned database + # whose catalog was stale, so a restart restored the stale catalog + path = str(tmp_path / "data.db") + internal = str(tmp_path / "internal.db") + make_db(path) + config = {"databases": {"data": {"schema_watch": "owned"}}} + + async def tables_after_start(change_before_close=None): + ds = Datasette([path], internal=internal, config=config) + await ds.invoke_startup() + response = await ds.client.get("/data.json") + if change_before_close: + # Behind Datasette's back, after the last request: an owned + # database is not polled, so the catalog is now stale + change_before_close() + ds.close() + return [t["name"] for t in response.json()["tables"]] + + assert await tables_after_start() == ["t"] + assert await tables_after_start( + lambda: make_db(path, "create table t2 (id integer)") + ) == ["t"] + # The stale catalog was not persisted as current: the restart rescans + assert await tables_after_start() == ["t", "t2"] + + +@pytest.mark.asyncio +async def test_crossdb_memory_connections_are_reaped_and_refreshed(tmp_path): + # Contract C7: --crossdb _memory connections ATTACH up to ten files but + # were never counted, reaped or refreshed: the fds stayed open and a + # removed database stayed queryable + paths = [] + for i in range(3): + path = str(tmp_path / f"db{i}.db") + make_db(path) + paths.append(path) + ds = Datasette(paths[:2], crossdb=True, settings={"connection_idle_timeout_ms": 50}) + await ds.invoke_startup() + memory = ds.get_database("_memory") + assert (await memory.execute("select count(*) from db0.t")).first()[0] == 0 + pool = ds._read_pool + assert pool.snapshot()["open"] >= 1 + deadline = time.monotonic() + 5 + while pool.snapshot()["open"]: + assert time.monotonic() < deadline, "crossdb connections were not reaped" + await asyncio.sleep(0.05) + # Added and removed databases are reflected immediately + await memory.execute("select 1") + ds.add_database(Database(ds, path=paths[2]), name="db2") + assert (await memory.execute("select count(*) from db2.t")).first()[0] == 0 + ds.remove_database("db0") + with pytest.raises(sqlite3.OperationalError): + await memory.execute("select count(*) from db0.t") + ds.close() + + +@pytest.mark.asyncio +async def test_execute_fn_may_return_generator_over_fetched_rows(): + # Contract C8: every generator was rejected, including harmless ones + ds = Datasette(memory=True) + db = ds.get_database("_memory") + + def fetched(conn): + rows = conn.execute("select 1 union all select 2").fetchall() + return (row[0] for row in rows) + + assert list(await db.execute_fn(fetched)) == [1, 2] + + def live_cursor(conn): + return (row[0] for row in conn.execute("select 1")) + + with pytest.raises(ConnectionLeaseError): + await db.execute_fn(live_cursor) + + def closure(conn): + def rows(): + yield from conn.execute("select 1") + + return rows() + + with pytest.raises(ConnectionLeaseError): + await db.execute_fn(closure) + ds.close() + + +# ---------------------------------------------------------------------- +# Found by the torture harness (tests/connection_torture.py) +# ---------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_index_survives_database_removed_mid_request(tmp_path, monkeypatch): + # The index page looked databases up by name again after listing them: + # one removed meanwhile (remove_database()) + # failed the whole page with a KeyError + from datasette.views import index as index_module + + paths = [] + for name in ("keep", "gone"): + path = str(tmp_path / f"{name}.db") + make_db(path) + paths.append(path) + ds = Datasette(paths) + await ds.invoke_startup() + original = index_module.catalog_summaries + + async def remove_during(datasette, names): + result = await original(datasette, names) + if "gone" in datasette.databases: + datasette.remove_database("gone") + return result + + monkeypatch.setattr(index_module, "catalog_summaries", remove_during) + response = await ds.client.get("/.json") + assert response.status_code == 200 + assert [d["name"] for d in response.json()["databases"]] == ["keep"] + ds.close() + + +@pytest.mark.asyncio +async def test_index_and_databases_with_closed_or_deleted_database(tmp_path): + paths = [] + for name in ("keep", "closed", "deleted"): + path = str(tmp_path / f"{name}.db") + make_db(path) + paths.append(path) + ds = Datasette(paths, settings={"schema_watch_interval_ms": 0}) + await ds.invoke_startup() + # Closed by a plugin but still attached + ds.get_database("closed").close() + # Deleted from under Datasette before any sweep noticed + os.unlink(paths[2]) + for path in ("/.json", "/", "/-/databases.json"): + response = await ds.client.get(path) + assert response.status_code == 200, (path, response.text[:300]) + databases = (await ds.client.get("/-/databases.json")).json() + if isinstance(databases, dict): + databases = databases["databases"] + sizes = {d["name"]: d["size"] for d in databases} + assert sizes["deleted"] == 0 + ds.close() + + +@pytest.mark.asyncio +async def test_cursors_are_closed_before_leaving_their_thread(owned_db): + # A cursor still references its connection's cached statement, and on + # Python < 3.12 freeing it resets that statement. execute_write_many() + # handed its cursor to the caller's event loop thread; freed there while + # the write thread ran the same SQL again, the write failed with + # "bad parameter or other API misuse" + _ds, db, _path = owned_db + cursor = await db.execute_write_many( + "insert into t (v) values (?)", [("a",), ("b",)] + ) + assert cursor.rowcount == 2 + with pytest.raises(ConnectionLeaseError): + cursor.fetchall() + cursor = await db.execute_write_fn( + lambda conn: conn.execute("insert into t (v) values ('c')") + ) + assert cursor.lastrowid == 3 + with pytest.raises(ConnectionLeaseError): + cursor.fetchone() + + # Cursors held by the frames of an exception raised by a callback are + # closed before the connection goes back to the pool (or write thread) + captured = [] + + def read_then_fail(conn): + cursor = conn.execute("select * from t") + captured.append(cursor) + raise ValueError("boom") + + with pytest.raises(ValueError): + await db.execute_fn(read_then_fail) + with pytest.raises(ValueError): + await db.execute_write_fn(read_then_fail) + for cursor in captured: + with pytest.raises(ConnectionLeaseError): + cursor.fetchone() + + +def test_startup_from_several_event_loops_runs_once(monkeypatch): + # _startup_sequence() serialized startup with an asyncio.Lock: callers + # on a second event loop got "is bound to a different event loop" (or, + # free-threaded, waited forever), and invoke_startup() itself let two + # loops run every startup hook twice + ds = Datasette(memory=True) + calls = [] + original = Datasette._apply_column_types_config + + async def slow(self): + calls.append(threading.current_thread().name) + await asyncio.sleep(0.3) + return await original(self) + + monkeypatch.setattr(Datasette, "_apply_column_types_config", slow) + + def starter(use_sequence): + async def run(): + if use_sequence: + await ds._startup_sequence() + else: + await ds.invoke_startup() + assert ds._startup_invoked + return (await ds.client.get("/_memory.json")).status_code + + return run + + statuses = run_in_threads(starter(True), starter(True), starter(False)) + assert statuses == [200, 200, 200] + assert len(calls) == 1 + ds.close() + + +@pytest.mark.asyncio +async def test_startup_hook_can_make_requests(): + # A request made from inside startup (same task) must not wait for + # startup to finish - that would wait forever + from datasette import hookimpl + from datasette.plugins import pm + + statuses = [] + + class Plugin: + __name__ = "StartupRequestPlugin" + + @hookimpl + def startup(self, datasette): + async def inner(): + statuses.append( + (await datasette.client.get("/-/versions.json")).status_code + ) + + return inner + + pm.register(Plugin(), name="startup_request_plugin") + try: + ds = Datasette(memory=True) + await asyncio.wait_for(ds.invoke_startup(), 10) + assert statuses == [200] + ds.close() + finally: + pm.unregister(name="startup_request_plugin") + + +def test_execute_write_many_from_several_event_loops(owned_db): + # T1 end to end: on Python 3.11 before the fix, dozens of these calls + # failed with "bad parameter or other API misuse" (3 of 3 runs) + _ds, db, _path = owned_db + errors = [] + + def loop(i): + async def run(): + async def worker(t): + for seq in range(60): + try: + await db.execute_write_many( + "insert into t (v) values (?)", [(f"{i}-{t}-{seq}",)] + ) + except Exception as e: # noqa: BLE001 + errors.append(repr(e)) + + await asyncio.gather(*(worker(t) for t in range(5))) + + return run + + run_in_threads(loop(0), loop(1), loop(2)) + assert errors == [] diff --git a/tests/test_connection_torture.py b/tests/test_connection_torture.py new file mode 100644 index 00000000..4c318c14 --- /dev/null +++ b/tests/test_connection_torture.py @@ -0,0 +1,75 @@ +""" +A fast, fixed-seed subset of the connection torture harness +(tests/connection_torture.py). Each scenario runs in its own process, so a +segfault or a hang fails one test instead of killing the test run. + +Several event loops (threads) drive one Datasette with extreme settings +(smallest pool, 1ms idle timeout, 10ms schema polling) through reads, +writes of every kind, DDL, isolated functions, analyze_sql, cancellation, +add/remove_database, external schema changes, files +replaced and deleted, Database.close() and Datasette.close() mid-flight, +and non-threaded mode. See the harness docstring for what is checked. +""" + +import json +import os +import signal +import subprocess +import sys + +import pytest + +HARNESS = os.path.join(os.path.dirname(__file__), "connection_torture.py") + + +@pytest.mark.parametrize( + "scenario,seed", + [ + ("mixed", 101), + ("lifecycle", 102), + ("close", 103), + ("cancel", 104), + ("pool", 105), + ("nothreads", 106), + ("writes", 107), + ], +) +def test_connection_torture(scenario, seed): + try: + proc = subprocess.run( + [ + sys.executable, + HARNESS, + scenario, + "--seed", + str(seed), + "--duration", + "1.5", + ], + capture_output=True, + text=True, + timeout=150, + check=False, + ) + except subprocess.TimeoutExpired as e: # pragma: no cover + pytest.fail(f"{scenario} hung:\n{(e.stderr or '')[-4000:]}") + if proc.returncode < 0: # pragma: no cover + pytest.fail( + f"{scenario} crashed with {signal.Signals(-proc.returncode).name}:\n" + f"{proc.stderr[-4000:]}" + ) + lines = proc.stdout.strip().splitlines() + if not lines: + pytest.fail( + f"{scenario} exited {proc.returncode} without a summary:\n" + f"{proc.stderr[-4000:]}" + ) + summary = json.loads(lines[-1]) + problems = "\n".join( + f"[{p['kind']}] {p['message'][:600]}" for p in summary["problems"] + ) + assert proc.returncode == 0, f"{scenario} seed {seed}:\n{problems}" + assert summary["ok"] + # It did something (the slow-query scenarios do a few dozen operations + # in 1.5s on a busy machine, the others hundreds) + assert sum(summary["counts"].values()) > 10 diff --git a/tests/test_fts_permissions.py b/tests/test_fts_permissions.py index 24a40cef..99044c3a 100644 --- a/tests/test_fts_permissions.py +++ b/tests/test_fts_permissions.py @@ -95,6 +95,11 @@ async def test_derived_permission_discovery_error_is_retried(monkeypatch, listin action="view-table", resource=TableResource("data", "documents") ) + # While the catalog is current, dependencies come from its rows and + # nothing is discovered live. Live discovery - the path that can fail - + # happens while a rescan of the database is pending. + db._watch_state.needs_scan = True + db._cached_derived_table_dependencies = None token = _permission_check_cache.set({}) try: with monkeypatch.context() as patch: diff --git a/tests/test_internal_db.py b/tests/test_internal_db.py index b4bb964d..c3c08716 100644 --- a/tests/test_internal_db.py +++ b/tests/test_internal_db.py @@ -29,7 +29,9 @@ async def test_internal_migrations_recorded(ds_client): order by id """) assert [tuple(row) for row in migrations.rows] == [ - ("datasette_internal", "0001_initial") + ("datasette_internal", "0001_initial"), + ("datasette_internal", "0002_catalog_fingerprint"), + ("datasette_internal", "0003_catalog_table_type"), ] @@ -61,7 +63,9 @@ async def test_internal_migrations_adopt_existing_internal_db(tmp_path): order by id """) assert [tuple(row) for row in migrations.rows] == [ - ("datasette_internal", "0001_initial") + ("datasette_internal", "0001_initial"), + ("datasette_internal", "0002_catalog_fingerprint"), + ("datasette_internal", "0003_catalog_table_type"), ] ds.close() @@ -73,7 +77,13 @@ async def test_internal_tables(ds_client): tables = await internal_db.execute("select * from catalog_tables") assert len(tables) > 5 table = tables.rows[0] - assert set(table.keys()) == {"rootpage", "table_name", "database_name", "sql"} + assert set(table.keys()) == { + "rootpage", + "table_name", + "database_name", + "sql", + "type", + } @pytest.mark.asyncio diff --git a/tests/test_internals_database.py b/tests/test_internals_database.py index 9398afb9..e399a323 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, @@ -1263,10 +1264,11 @@ async def test_database_close_shuts_down_write_thread(tmpdir): await db.execute_write("insert into t (id) values (1)") assert db._write_thread is not None assert db._write_thread.is_alive() + thread = db._write_thread db.close() - # Wait briefly for the thread to exit — the sentinel should cause it to return. - db._write_thread.join(timeout=5) - assert not db._write_thread.is_alive() + thread.join(timeout=5) + assert not thread.is_alive() + assert db._write_thread is None ds._internal_database.close() @@ -1312,10 +1314,18 @@ 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) - write_connection = await db.execute_write_fn(lambda conn: conn) + leased = await db.execute_fn(lambda conn: conn) + leased_write = await db.execute_write_fn(lambda conn: conn) + # Inspect the owned raw connections to verify physical closure as well + # as lease expiry. cursor.connection intentionally no longer leaks them. + connections = list(db._all_connections) + assert len(connections) >= 2 + with pytest.raises(ConnectionLeaseError): + leased.execute("select 1") + with pytest.raises(ConnectionLeaseError): + leased_write.execute("select 1") ds.close() - for conn in (read_connection, write_connection): + for conn in connections: with pytest.raises(sqlite3.ProgrammingError, match="closed"): conn.execute("select 1") 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() diff --git a/tests/test_lifespan.py b/tests/test_lifespan.py index 3a25ccbc..606c20db 100644 --- a/tests/test_lifespan.py +++ b/tests/test_lifespan.py @@ -215,18 +215,15 @@ async def test_concurrent_first_requests_all_wait_for_slow_startup(): @pytest.mark.asyncio -async def test_setup_db_still_runs_when_invoke_startup_ran_first(tmp_path, monkeypatch): - # Regression test: `datasette serve` (cli.py _serve_async) calls - # ds.invoke_startup() directly, before uvicorn ever sends a - # lifespan.startup event that drives _startup_sequence(). If - # _startup_sequence()'s fast path only checked `_startup_invoked`, it - # would see startup already done and skip the immutable-database - # table-count precompute (setup_db) entirely - a silent regression - # versus main, where AsgiRunOnFirstRequest ran setup_db unconditionally - # on request #1. +async def test_immutable_table_counts_computed_when_first_needed(tmp_path, monkeypatch): + # Immutable databases used to have their table counts precomputed by + # _startup_sequence(), opening every immutable database at startup. + # They are now computed the first time a page needs them, with the same + # generous time limit, and cached. db_path = tmp_path / "immutable.db" conn = sqlite3.connect(str(db_path)) conn.execute("create table t (id integer primary key)") + conn.execute("insert into t default values") conn.commit() conn.close() @@ -241,22 +238,22 @@ async def test_setup_db_still_runs_when_invoke_startup_ran_first(tmp_path, monke monkeypatch.setattr(Database, "table_counts", counting_table_counts) - # Simulate the CLI path: invoke_startup() runs directly and completes - # BEFORE _startup_sequence() ever gets a chance to run setup_db. await ds.invoke_startup() + await ds._startup_sequence() assert ds._startup_invoked is True - assert call_count["n"] == 0 - - # The lifespan/first-request path (or the CLI itself, per the fix) - # calling the shared entry point afterwards must still precompute - # table counts for immutable databases. - await ds._startup_sequence() - assert call_count["n"] == 1 assert ds._setup_db_done is True + assert call_count["n"] == 0 + assert ds.get_database("immutable").cached_table_counts is None - # Idempotency: a second call must not recompute. - await ds._startup_sequence() + response = await ds.client.get("/.json") + assert response.json()["databases"][0]["table_rows_sum"] == 1 assert call_count["n"] == 1 + assert ds.get_database("immutable").cached_table_counts == {"t": 1} + + # Cached from then on + await ds.client.get("/immutable.json") + assert ds.get_database("immutable").cached_table_counts == {"t": 1} + ds.close() @pytest.mark.asyncio diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 58e154a2..9a0bdbf8 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -48,6 +48,11 @@ def test_plugin_hooks_have_tests(plugin_hook): assert ok, f"Plugin hook is missing tests: {plugin_hook}" +def test_hook_write_task_completed(): + # Success and expiry delivery are exercised in tests/test_write_budget.py. + assert "write_task_completed" in dir(pm.hook) + + def test_hook_jump_items_sql(): # Detailed behavior is covered in tests/test_jump.py. assert "jump_items_sql" in dir(pm.hook) diff --git a/tests/test_read_pool.py b/tests/test_read_pool.py new file mode 100644 index 00000000..a4c23233 --- /dev/null +++ b/tests/test_read_pool.py @@ -0,0 +1,378 @@ +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}) + pool = ds._read_pool + # Below the 4 x num_sql_threads clamp, so set the private attribute + pool.max_open = 2 + 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) + ds._read_pool.max_open = 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.parametrize( + "num_sql_threads,configured,expected", + [ + (3, 128, 128), + (3, 5, 5), + (8, 20, 20), + (0, 1, 1), + (3, 0, 0), + ], +) +def test_max_open_connections_respects_configured_limit( + num_sql_threads, configured, expected +): + ds = Datasette( + settings={ + "num_sql_threads": num_sql_threads, + "max_open_connections": configured, + } + ) + assert ds._read_pool.max_open == expected + # /-/settings.json reports the configured value + assert ds.setting("max_open_connections") == configured + ds.close() + + +def test_connection_idle_timeout_ms_setting(): + ds = Datasette(settings={"connection_idle_timeout_ms": 1500}) + assert ds._read_pool.idle_timeout == 1.5 + ds.close() + ds = Datasette(settings={"connection_idle_timeout_ms": 0}) + assert ds._read_pool.idle_timeout == 0 + ds.close() + + +@pytest.mark.asyncio +async def test_hard_cap_when_every_connection_is_leased(tmp_path): + paths = _make_dbs(tmp_path, 3) + ds = Datasette(paths, settings={"num_sql_threads": 3, "max_open_connections": 1}) + entered = threading.Event() + release = threading.Event() + + def slow(conn): + entered.set() + assert release.wait(5) + return conn.execute("select 1").fetchone()[0] + + tasks = [ + asyncio.create_task(ds.get_database(f"db{i}").execute_fn(slow)) + for i in range(3) + ] + try: + assert await asyncio.to_thread(entered.wait, 5) + await asyncio.sleep(0.05) + assert ds._read_pool.snapshot()["open"] == 1 + assert not any(task.done() for task in tasks) + release.set() + assert await asyncio.gather(*tasks) == [1, 1, 1] + assert ds._read_pool.stats["peak_open"] == 1 + finally: + release.set() + await asyncio.gather(*tasks, return_exceptions=True) + 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() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", [0, 2]) +async def test_invalidate_database_discards_stale_connections( + tmp_path, num_sql_threads +): + paths = _make_dbs(tmp_path, 2) + ds = Datasette(paths, settings={"num_sql_threads": num_sql_threads}) + db0, db1 = ds.get_database("db0"), ds.get_database("db1") + pool = ds._read_pool + await db0.execute("select 1") + await db1.execute("select 1") + old = db0._read_pool_state.idle[-1].conn + assert db0._read_pool_state.open == 1 + # What the SchemaWatcher does when it sees db0's file replaced + db0._invalidate_connections() + # Idle connections for db0 are closed at once; db1 is untouched + assert db0._read_pool_state.open == 0 + assert db1._read_pool_state.open == 1 + assert old not in db0._all_connections + with pytest.raises(sqlite3.ProgrammingError): + old.execute("select 1") + # The next read opens a new connection + await db0.execute("select 1") + assert db0._read_pool_state.idle[-1].conn is not old + assert pool.stats["discarded_stale"] == 1 + ds.close() + + +@pytest.mark.asyncio +async def test_leased_connection_discarded_on_release_after_invalidation(tmp_path): + (path,) = _make_dbs(tmp_path, 1) + ds = Datasette([path], settings={"num_sql_threads": 2}) + db = ds.get_database("db0") + pool = ds._read_pool + seen = {} + + def callback(conn): + seen["raw"] = conn.execute("select 1").connection + # The file is replaced while this callback holds the connection: + # it must not be closed under us, only discarded on release + db._invalidate_connections() + return conn.execute("select count(*) from t").fetchone()[0] + + assert await db.execute_fn(callback) == 1 + assert db._read_pool_state.open == 0 + assert pool.stats["discarded_stale"] == 1 + with pytest.raises(ConnectionLeaseError): + seen["raw"].execute("select 1") + assert (await db.execute("select count(*) from t")).single_value() == 1 + ds.close() + + +def test_reaper_unavailable_falls_back_to_checkout_reaping(tmp_path, monkeypatch): + # Pyodide cannot start threads: idle connections must then be reaped by + # the next checkout instead, and reads must keep working + from datasette import connection_pool + + reaper = connection_pool._Reaper() + monkeypatch.setattr(connection_pool, "_reaper", reaper) + + def no_threads(self): + raise RuntimeError("can't start new thread") + + monkeypatch.setattr(threading.Thread, "start", no_threads) + paths = _make_dbs(tmp_path, 2) + ds = Datasette(paths, settings={"num_sql_threads": 0}) + pool = ds._read_pool + pool.idle_timeout = 0.05 + + async def read(name): + return (await ds.get_database(name).execute("select 1")).single_value() + + assert asyncio.run(read("db0")) == 1 + assert reaper.unavailable + assert pool.snapshot()["open"] == 1 + time.sleep(0.1) + assert asyncio.run(read("db1")) == 1 + # db0's expired connection was closed during db1's checkout + assert pool.stats["expired"] == 1 + assert ds.get_database("db0")._read_pool_state.open == 0 + ds.close() diff --git a/tests/test_schema_watcher.py b/tests/test_schema_watcher.py new file mode 100644 index 00000000..fe9b6613 --- /dev/null +++ b/tests/test_schema_watcher.py @@ -0,0 +1,718 @@ +""" +SchemaWatcher tests. Changes to "external" databases are made from a +separate OS process (a child Python interpreter using plain sqlite3, or the +sqlite-utils CLI), the way other tools change files Datasette is serving. +""" + +import asyncio +import os +import sqlite3 +import subprocess +import sys +import textwrap +import time + +import pytest + +from datasette import schema_watcher as schema_watcher_module +from datasette.app import Datasette +from datasette.database import Database + +INTERVAL_MS = 50 + + +def run_external(path, sql, journal_mode=None, keep_wal=False): + """Run sql against path in a separate Python process; returns the + child's time.monotonic() right after COMMIT.""" + script = textwrap.dedent(""" + import sqlite3, sys, time + path, sql, journal_mode, keep_wal = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4] + conn = sqlite3.connect(path, isolation_level=None, timeout=10) + if journal_mode: + conn.execute("PRAGMA journal_mode=" + journal_mode) + if keep_wal == "1": + # Never checkpoint: changes stay in the -wal file + conn.execute("PRAGMA wal_autocheckpoint=0") + conn.executescript(sql) + print(time.monotonic(), flush=True) + if keep_wal == "1": + # Keep the connection open so close() does not checkpoint + sys.stdin.read() + conn.close() + """) + if keep_wal: + proc = subprocess.Popen( + [sys.executable, "-c", script, path, sql, journal_mode or "", "1"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + text=True, + ) + line = proc.stdout.readline() + return float(line), proc + out = subprocess.run( + [sys.executable, "-c", script, path, sql, journal_mode or "", "0"], + check=True, + capture_output=True, + text=True, + ) + return float(out.stdout.strip()) + + +def run_process(*args): + subprocess.run([sys.executable, *args], check=True) + + +def make_db(path, sql="create table t (id integer primary key, a text)", wal=False): + conn = sqlite3.connect(path) + if wal: + conn.execute("PRAGMA journal_mode=wal") + conn.executescript(sql) + conn.execute("insert into t (a) values ('one')") + conn.commit() + conn.close() + + +async def catalog_columns(ds, db_name, table): + rows = await ds.get_internal_database().execute( + "select name from catalog_columns where database_name = ? and table_name = ? order by cid", + [db_name, table], + ) + return [r[0] for r in rows.rows] + + +async def catalog_tables(ds, db_name): + rows = await ds.get_internal_database().execute( + "select table_name from catalog_tables where database_name = ? order by table_name", + [db_name], + ) + return [r[0] for r in rows.rows] + + +async def wait_for(check, timeout=5.0): + t0 = time.monotonic() + while True: + result = await check() + if result: + return time.monotonic() - t0 + if time.monotonic() - t0 > timeout: + raise AssertionError(f"timed out after {timeout}s") + await asyncio.sleep(0.01) + + +@pytest.fixture +def db_path(tmp_path): + path = str(tmp_path / "ext.db") + make_db(path) + return path + + +async def start_ds(paths, interval_ms=INTERVAL_MS, settings=None, **kwargs): + ds = Datasette( + files=paths, + settings={"schema_watch_interval_ms": interval_ms, **(settings or {})}, + **kwargs, + ) + await ds.invoke_startup() + await ds.start_background_tasks() + return ds + + +@pytest.mark.asyncio +async def test_modes_and_defaults(tmp_path, db_path): + immutable_path = str(tmp_path / "imm.db") + make_db(immutable_path) + ds = Datasette(files=[db_path], immutables=[immutable_path]) + await ds.invoke_startup() + watcher = ds._schema_watcher + assert watcher.states["ext"].mode == "external" + assert watcher.states["imm"].mode == "immutable" + # add_database(): files and named in-memory databases can be changed + # by code other than Datasette's write methods (a plugin's own + # connection, a backup() into a named in-memory database), so they get + # the default_schema_watch mode too unless the caller says otherwise + added = ds.add_database(Database(ds, memory_name="sw_owned_mem"), name="added") + assert watcher.states["added"].mode == "external" + assert added._watch_state is watcher.states["added"] + added_file = str(tmp_path / "added_file.db") + make_db(added_file) + ds.add_database(Database(ds, path=added_file), name="added_file") + assert watcher.states["added_file"].mode == "external" + ds.add_database( + Database(ds, memory_name="sw_owned_mem2"), name="owned", schema_watch="owned" + ) + assert watcher.states["owned"].mode == "owned" + # A private ":memory:" database: nothing else can change it + ds.add_database(Database(ds, is_memory=True), name="private") + assert watcher.states["private"].mode == "owned" + ds.close() + # default_schema_watch=owned applies to add_database() too + ds = Datasette(files=[db_path], settings={"default_schema_watch": "owned"}) + ds.add_database(Database(ds, memory_name="sw_owned_mem3"), name="added") + assert ds._schema_watcher.states["added"].mode == "owned" + assert ds._schema_watcher.states["ext"].mode == "owned" + ds.close() + + +@pytest.mark.asyncio +async def test_mode_from_config(tmp_path, db_path): + ds = Datasette( + files=[db_path], config={"databases": {"ext": {"schema_watch": "owned"}}} + ) + assert ds._schema_watcher.states["ext"].mode == "owned" + ds.close() + + +@pytest.mark.asyncio +async def test_no_request_path_sweep(db_path): + ds = await start_ds([db_path], interval_ms=0) + watcher = ds._schema_watcher + sweeps = watcher.counters["sweeps"] + for _ in range(5): + response = await ds.client.get("/ext/t.json") + assert response.status_code == 200 + assert watcher.counters["sweeps"] == sweeps + ds.close() + + +@pytest.mark.asyncio +async def test_owned_write_path_detects_schema_change(tmp_path): + ds = Datasette(settings={"schema_watch_interval_ms": 0}) + await ds.invoke_startup() + db = ds.add_database(Database(ds, memory_name="sw_owned"), name="owned") + await ds.refresh_schemas(force=True) + watcher = ds._schema_watcher + scans = watcher.counters["scans"] + # Data-only writes: no catalog rebuild + await db.execute_write("create table if not exists t (id integer primary key)") + assert await catalog_tables(ds, "owned") == ["t"] + scans = watcher.counters["scans"] + for i in range(5): + await db.execute_write("insert into t default values") + assert watcher.counters["scans"] == scans + # Each write API: catalog is current as soon as the call returns + await db.execute_write("create table t2 (id integer primary key)") + assert "t2" in await catalog_tables(ds, "owned") + await db.execute_write_script("create table t3 (id integer); create table t4 (id);") + assert {"t3", "t4"} <= set(await catalog_tables(ds, "owned")) + scans = watcher.counters["scans"] + await db.execute_write_many("insert into t2 (id) values (?)", [[1], [2]]) + assert watcher.counters["scans"] == scans + + def plugin_write(conn): + conn.execute("alter table t add column extra text") + + await db.execute_write_fn(plugin_write) + assert await catalog_columns(ds, "owned", "t") == ["id", "extra"] + + def isolated(conn): + conn.execute("create table t6 (id)") + + await db.execute_isolated_fn(isolated) + assert "t6" in await catalog_tables(ds, "owned") + # block=False: refreshed in the background + await db.execute_write("create table t7 (id)", block=False) + await wait_for(lambda: _contains(ds, "owned", "t7")) + ds.close() + + +async def _contains(ds, db_name, table): + return table in await catalog_tables(ds, db_name) + + +@pytest.mark.asyncio +async def test_owned_write_path_non_threaded(tmp_path): + path = str(tmp_path / "nt.db") + make_db(path) + ds = Datasette(settings={"num_sql_threads": 0, "schema_watch_interval_ms": 0}) + await ds.invoke_startup() + db = ds.add_database(Database(ds, path=path), name="nt") + await ds._schema_watcher.flush_pending() + await db.execute_write("create table t2 (id)") + assert await catalog_tables(ds, "nt") == ["t", "t2"] + ds.close() + + +@pytest.mark.asyncio +async def test_external_data_only_insert_does_not_rescan(db_path): + ds = await start_ds([db_path]) + watcher = ds._schema_watcher + state = watcher.states["ext"] + scans = state.stats["scans"] + checks = state.stats["checks"] + for i in range(3): + run_external(db_path, f"insert into t (a) values ('x{i}')") + await asyncio.sleep(INTERVAL_MS * 3 / 1000) + # The fingerprint changed, so schema_version was checked... + assert state.stats["checks"] > checks + # ... but the catalog was not rebuilt + assert state.stats["scans"] == scans + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_header", (True, False)) +async def test_external_add_column(db_path, monkeypatch, use_header): + monkeypatch.setattr(schema_watcher_module, "USE_HEADER_CHECK", use_header) + ds = await start_ds([db_path]) + run_external(db_path, "alter table t add column b text") + + async def has_b(): + return await catalog_columns(ds, "ext", "t") == ["id", "a", "b"] + + await wait_for(has_b) + ds.close() + + +@pytest.mark.asyncio +async def test_external_sqlite_utils_transform(db_path): + ds = await start_ds([db_path]) + run_process( + "-m", "sqlite_utils", "transform", db_path, "t", "--rename", "a", "renamed" + ) + + async def renamed(): + return await catalog_columns(ds, "ext", "t") == ["id", "renamed"] + + await wait_for(renamed) + # And the table page uses the new column too + response = await ds.client.get("/ext/t.json?_shape=array") + assert response.json() == [{"id": 1, "renamed": "one"}] + ds.close() + + +@pytest.mark.asyncio +async def test_external_create_table_in_wal_before_checkpoint(tmp_path): + path = str(tmp_path / "walled.db") + make_db(path, wal=True) + ds = await start_ds([path]) + _, proc = run_external( + path, "create table in_wal (id integer)", journal_mode="wal", keep_wal=True + ) + try: + assert os.path.getsize(path + "-wal") > 0 + + async def seen(): + return "in_wal" in await catalog_tables(ds, "walled") + + await wait_for(seen) + # Main file still does not contain it: it was only in the WAL + assert os.path.exists(path + "-wal") + finally: + proc.stdin.close() + proc.wait() + ds.close() + + +@pytest.mark.asyncio +async def test_external_atomic_replace_by_rename(tmp_path, db_path): + ds = await start_ds([db_path]) + # Warm a cached read connection on the old inode + response = await ds.client.get("/ext/t.json?_shape=array") + assert response.json() == [{"id": 1, "a": "one"}] + replacement = str(tmp_path / "replacement.db") + conn = sqlite3.connect(replacement) + # Same schema_version as the original file (1 CREATE TABLE) but a + # different table: only the inode change gives it away + conn.execute("create table t (id integer primary key, different text)") + conn.execute("insert into t (different) values ('new file')") + conn.commit() + assert conn.execute("pragma schema_version").fetchone()[0] == 1 + conn.close() + old_inode = os.stat(db_path).st_ino + run_process( + "-c", + "import os, sys; os.replace(sys.argv[1], sys.argv[2])", + replacement, + db_path, + ) + assert os.stat(db_path).st_ino != old_inode + + async def replaced(): + return await catalog_columns(ds, "ext", "t") == ["id", "different"] + + await wait_for(replaced) + # Cached connections were invalidated, so queries see the new file + response = await ds.client.get("/ext/t.json?_shape=array") + assert response.json() == [{"id": 1, "different": "new file"}] + ds.close() + + +@pytest.mark.asyncio +async def test_old_inode_connections_without_watcher(tmp_path, db_path): + # Documents the hazard the watcher fixes: with polling off, a cached + # read connection keeps reading the replaced (unlinked) file + # One SQL thread, so every request reuses the same cached connection + ds = await start_ds([db_path], interval_ms=0, settings={"num_sql_threads": 1}) + response = await ds.client.get("/ext/t.json?_shape=array") + assert response.json() == [{"id": 1, "a": "one"}] + replacement = str(tmp_path / "replacement.db") + conn = sqlite3.connect(replacement) + conn.execute("create table t (id integer primary key, a text)") + conn.execute("insert into t (a) values ('new file')") + conn.commit() + conn.close() + os.replace(replacement, db_path) + response = await ds.client.get("/ext/t.json?_shape=array") + assert response.json() == [{"id": 1, "a": "one"}] # stale! + # An explicit refresh notices the inode change and fixes it + await ds.refresh_schemas(force=True) + response = await ds.client.get("/ext/t.json?_shape=array") + assert response.json() == [{"id": 1, "a": "new file"}] + ds.close() + + +@pytest.mark.asyncio +async def test_external_vacuum(db_path): + run_external( + db_path, + "create table big (x); insert into big select zeroblob(1000) from " + "(with recursive c(i) as (select 1 union all select i+1 from c where i < 200) select i from c);" + "drop table big;", + ) + ds = await start_ds([db_path]) + before = await ds.get_internal_database().execute( + "select schema_version from catalog_databases where database_name = 'ext'" + ) + run_external(db_path, "vacuum") + + async def bumped(): + after = await ds.get_internal_database().execute( + "select schema_version from catalog_databases where database_name = 'ext'" + ) + return after.first()[0] == before.first()[0] + 1 + + await wait_for(bumped) + ds.close() + + +@pytest.mark.asyncio +async def test_external_rapid_changes(db_path): + ds = await start_ds([db_path]) + for i in range(10): + run_external(db_path, f"create table rapid_{i} (id)") + + async def all_seen(): + tables = await catalog_tables(ds, "ext") + return all(f"rapid_{i}" in tables for i in range(10)) + + await wait_for(all_seen) + ds.close() + + +def _coarse_fingerprint(real): + # Simulate a filesystem with 1 second timestamp granularity + def fingerprint(path, present=None): + parts = real(path, present) + out = [] + for i, part in enumerate(parts): + if part is None: + out.append(None) + continue + part = list(part) + # main: (dev, ino, size, mtime, ctime); side files: (ino, size, mtime) + for idx in (3, 4) if i == 0 else (2,): + part[idx] = part[idx] // 10**9 * 10**9 + out.append(tuple(part)) + return tuple(out) + + return fingerprint + + +@pytest.mark.parametrize("journal_mode", ("delete", "wal", "truncate")) +def test_header_schema_version_matches_pragma(tmp_path, journal_mode): + path = str(tmp_path / "h.db") + conn = sqlite3.connect(path, isolation_level=None) + conn.execute(f"PRAGMA journal_mode={journal_mode}") + for i in range(300): + conn.execute(f"create table t{i} (id integer)") + conn.execute("alter table t1 add column b") + version = conn.execute("PRAGMA schema_version").fetchone()[0] + if journal_mode == "wal": + # Pages still in the -wal: the main file header is stale, which is + # why the watcher only reads it when the -wal is empty + assert schema_watcher_module.fingerprint(path)[1] is not None + conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") + assert schema_watcher_module.fingerprint(path)[1:] == (None, None) + assert schema_watcher_module.header_schema_version(path) == version == 301 + conn.close() + assert ( + schema_watcher_module.header_schema_version(str(tmp_path / "nope.db")) is None + ) + (tmp_path / "junk.db").write_bytes(b"not a database" * 20) + assert ( + schema_watcher_module.header_schema_version(str(tmp_path / "junk.db")) is None + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_header", (True, False)) +@pytest.mark.parametrize("racy_rule", (True, False)) +async def test_racily_clean_same_tick_change( + db_path, monkeypatch, racy_rule, use_header +): + """Two schema changes inside one timestamp tick, same file size and + inode: only the racy-clean rule catches the second one.""" + monkeypatch.setattr(schema_watcher_module, "USE_HEADER_CHECK", use_header) + monkeypatch.setattr( + schema_watcher_module, + "fingerprint", + _coarse_fingerprint(schema_watcher_module.fingerprint), + ) + if not racy_rule: + monkeypatch.setattr(schema_watcher_module, "RACY_WINDOW_NS", -(10**18)) + ds = Datasette(files=[db_path], settings={"schema_watch_interval_ms": 0}) + await ds.invoke_startup() + # Line the first change up with the start of a fresh second + await asyncio.sleep(1 - (time.time() % 1) + 0.05) + size = os.path.getsize(db_path) + run_external(db_path, "alter table t add column b text") + await ds._refresh_schemas() # sweep: sees change 1 + assert await catalog_columns(ds, "ext", "t") == ["id", "a", "b"] + run_external(db_path, "alter table t add column c text") + assert os.path.getsize(db_path) == size + assert time.time() % 1 < 0.9, "test machine too slow to stay in one tick" + await ds._refresh_schemas() # sweep: fingerprint identical to last one + columns = await catalog_columns(ds, "ext", "t") + if racy_rule: + assert columns == ["id", "a", "b", "c"] + else: + # Without the rule the change is missed + assert columns == ["id", "a", "b"] + ds.close() + + +@pytest.mark.asyncio +async def test_external_file_deleted_and_recreated(tmp_path, db_path): + ds = await start_ds([db_path]) + assert await catalog_tables(ds, "ext") == ["t"] + run_process("-c", "import os, sys; os.unlink(sys.argv[1])", db_path) + + async def gone(): + return await catalog_tables(ds, "ext") == [] + + await wait_for(gone) + response = await ds.client.get("/ext/t.json") + assert response.status_code in (404, 500) + # Writes fail instead of silently creating a new empty file + with pytest.raises(sqlite3.OperationalError): + await ds.get_database("ext").execute_write("insert into t (a) values ('x')") + assert not os.path.exists(db_path) + make_db(db_path, sql="create table t (id integer primary key, a text, back text)") + + async def back(): + return await catalog_columns(ds, "ext", "t") == ["id", "a", "back"] + + await wait_for(back) + response = await ds.client.get("/ext/t.json?_shape=array") + assert response.json() == [{"id": 1, "a": "one", "back": None}] + ds.close() + + +@pytest.mark.asyncio +async def test_remove_database_deletes_catalog_rows(tmp_path, db_path): + other = str(tmp_path / "other.db") + make_db(other) + ds = Datasette(files=[db_path, other]) + await ds.invoke_startup() + assert await catalog_tables(ds, "other") == ["t"] + ds.remove_database("other") + await ds._schema_watcher.on_request() + internal = ds.get_internal_database() + for table in ( + "catalog_databases", + "catalog_tables", + "catalog_columns", + "catalog_indexes", + ): + rows = await internal.execute( + f"select count(*) from {table} where database_name = 'other'" + ) + assert rows.first()[0] == 0 + assert await catalog_tables(ds, "ext") == ["t"] + ds.close() + + +@pytest.mark.asyncio +async def test_add_database_visible_on_next_request(tmp_path): + ds = Datasette(settings={"schema_watch_interval_ms": 0}) + await ds.invoke_startup() + await ds.client.get("/-/versions.json") + path = str(tmp_path / "late.db") + make_db(path) + ds.add_database(Database(ds, path=path), name="late") + response = await ds.client.get("/late.json") + assert response.status_code == 200 + assert await catalog_tables(ds, "late") == ["t"] + ds.close() + + +@pytest.mark.asyncio +async def test_persisted_fingerprints_skip_unchanged_files(tmp_path, monkeypatch): + # Shrink the racy window so freshly written files count as stable + monkeypatch.setattr(schema_watcher_module, "RACY_WINDOW_NS", 1_000_000) + internal = str(tmp_path / "internal.db") + paths = [] + for i in range(3): + path = str(tmp_path / f"p{i}.db") + make_db(path) + paths.append(path) + await asyncio.sleep(0.05) + ds = Datasette(files=paths, internal=internal) + await ds.invoke_startup() + assert ds._schema_watcher.counters["scans"] == 3 + ds.close() + # Change one file while Datasette is "down" + run_external(paths[1], "create table added_offline (id)") + ds2 = Datasette(files=paths, internal=internal) + await ds2.invoke_startup() + counters = ds2._schema_watcher.counters + assert counters["restored_from_persisted"] == 2 + assert counters["scans"] == 1 + assert await catalog_tables(ds2, "p1") == ["added_offline", "t"] + assert await catalog_tables(ds2, "p0") == ["t"] + ds2.close() + + +@pytest.mark.asyncio +async def test_remove_database_during_catalog_scan(tmp_path): + # remove_database() closes the Database while the watcher may still be + # scanning it in a worker thread. The scan connection is not tracked by + # the Database, so close() must not pull it out from under the scan + # (that used to segfault). + path = str(tmp_path / "big.db") + conn = sqlite3.connect(path) + for i in range(150): + conn.execute(f"create table t{i} (a, b, c integer references t0(a))") + conn.execute(f"create index i{i} on t{i}(a)") + conn.commit() + conn.close() + ds = Datasette(settings={"schema_watch_interval_ms": 0}) + await ds.invoke_startup() + for i in range(60): + ds.add_database(Database(ds, path=path, is_mutable=i % 2 == 0), name=f"x{i}") + await asyncio.sleep(0.0005 * (i % 5)) + ds.remove_database(f"x{i}") + await asyncio.sleep(0.2) + rows = await ds.get_internal_database().execute( + "select count(*) from catalog_tables where database_name like 'x%'" + ) + assert rows.first()[0] == 0 + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "setting,expected", (("external", "external"), ("owned", "owned")) +) +async def test_default_schema_watch_setting(tmp_path, db_path, setting, expected): + immutable_path = str(tmp_path / "imm.db") + make_db(immutable_path) + ds = Datasette( + files=[db_path], + immutables=[immutable_path], + settings={"default_schema_watch": setting}, + ) + watcher = ds._schema_watcher + assert watcher.states["ext"].mode == expected + # Immutable files are always immutable + assert watcher.states["imm"].mode == "immutable" + # add_database() follows the setting too, and can ask for either mode + other = str(tmp_path / "other.db") + make_db(other) + added = ds.add_database(Database(ds, path=other), name="added") + assert added._watch_state.mode == expected + polled = ds.add_database( + Database(ds, path=other), name="polled", schema_watch="external" + ) + assert polled._watch_state.mode == "external" + owned = ds.add_database( + Database(ds, path=other), name="owned", schema_watch="owned" + ) + assert owned._watch_state.mode == "owned" + ds.close() + + +def test_default_schema_watch_config_overrides(db_path): + ds = Datasette( + files=[db_path], + settings={"default_schema_watch": "owned"}, + config={"databases": {"ext": {"schema_watch": "external"}}}, + ) + assert ds._schema_watcher.states["ext"].mode == "external" + ds.close() + + +def test_default_schema_watch_invalid(db_path): + from datasette.utils import StartupError + + with pytest.raises(StartupError, match="default_schema_watch"): + Datasette(files=[db_path], settings={"default_schema_watch": "immutable"}) + + +@pytest.mark.asyncio +async def test_owned_default_is_not_polled(db_path): + ds = await start_ds([db_path], settings={"default_schema_watch": "owned"}) + run_external(db_path, "create table not_polled (id)") + await asyncio.sleep(INTERVAL_MS * 4 / 1000) + assert await catalog_tables(ds, "ext") == ["t"] + # Datasette's own writes are still detected immediately + await ds.get_database("ext").execute_write("create table mine (id)") + assert await catalog_tables(ds, "ext") == ["mine", "not_polled", "t"] + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", (0, 2)) +async def test_scratch_database_file_created_by_first_write(tmp_path, num_sql_threads): + # A database added before its file exists (a plugin's scratch + # database): the first write creates it. Once it has existed, deleting + # it makes writes fail rather than silently creating a new empty file. + ds = Datasette( + settings={"num_sql_threads": num_sql_threads, "schema_watch_interval_ms": 0} + ) + await ds.invoke_startup() + path = str(tmp_path / "scratch.db") + db = ds.add_database(Database(ds, path=path), name="scratch") + await ds._schema_watcher.flush_pending() + assert not os.path.exists(path) + await db.execute_write("create table t (id integer primary key)") + assert os.path.exists(path) + assert await catalog_tables(ds, "scratch") == ["t"] + await db.execute_write("insert into t default values") + os.unlink(path) + await ds.refresh_schemas(force=True) + assert await catalog_tables(ds, "scratch") == [] + with pytest.raises(sqlite3.OperationalError): + await db.execute_write("insert into t default values") + assert not os.path.exists(path) + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_sql_threads", (0, 2)) +async def test_write_connection_follows_replaced_file( + tmp_path, db_path, num_sql_threads +): + ds = await start_ds( + [db_path], interval_ms=0, settings={"num_sql_threads": num_sql_threads} + ) + db = ds.get_database("ext") + # Warm the write connection on the old inode + await db.execute_write("insert into t (a) values ('old')") + replacement = str(tmp_path / "replacement.db") + make_db(replacement) + os.replace(replacement, db_path) + await ds.refresh_schemas(force=True) + await db.execute_write("insert into t (a) values ('new')") + conn = sqlite3.connect(db_path) + try: + rows = conn.execute("select a from t order by id").fetchall() + finally: + conn.close() + # The write went to the new file, not the unlinked one + assert rows == [("one",), ("new",)] + assert [r["a"] for r in (await db.execute("select a from t order by id")).rows] == [ + "one", + "new", + ] + ds.close() diff --git a/tests/test_sqlite_error_compatibility.py b/tests/test_sqlite_error_compatibility.py new file mode 100644 index 00000000..b7da5153 --- /dev/null +++ b/tests/test_sqlite_error_compatibility.py @@ -0,0 +1,70 @@ +"""Error recovery must also work without Python 3.11's SQLite error metadata.""" + +import errno +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from datasette.schema_watcher import _is_transient +from datasette.utils import sqlite3 +from datasette.write_budget import DatabaseResourceError, connect_with_retry + + +@pytest.fixture +def old_sqlite_errors(monkeypatch): + for name in ("SQLITE_CANTOPEN", "SQLITE_NOMEM"): + monkeypatch.delattr(sqlite3, name, raising=False) + + +@pytest.mark.parametrize( + "error,expected", + [ + (sqlite3.OperationalError("unable to open database file"), True), + (sqlite3.OperationalError("out of memory"), True), + (sqlite3.OperationalError("database is locked"), True), + (sqlite3.OperationalError("no such table: missing"), False), + (OSError(errno.EMFILE, "too many open files"), True), + ], +) +def test_transient_errors_without_sqlite_constants(old_sqlite_errors, error, expected): + assert _is_transient(error) is expected + + +@pytest.mark.parametrize("os_error", [False, True]) +def test_open_recovers_without_sqlite_constants(old_sqlite_errors, tmp_path, os_error): + path = tmp_path / "db.sqlite" + path.touch() + pool = Mock() + db = SimpleNamespace(path=str(path), ds=SimpleNamespace(_read_pool_or_none=pool)) + error = ( + OSError(errno.EMFILE, "too many open files") + if os_error + else sqlite3.OperationalError("unable to open database file") + ) + connection = object() + connect = Mock(side_effect=[error, connection]) + assert connect_with_retry(db, connect) is connection + assert connect.call_count == 2 + pool.evict_idle.assert_called_once_with() + + +@pytest.mark.parametrize("os_error", [False, True]) +def test_open_failure_preserved_without_sqlite_constants( + old_sqlite_errors, tmp_path, os_error +): + path = tmp_path / "db.sqlite" + path.touch() + db = SimpleNamespace(path=str(path), ds=SimpleNamespace(_read_pool_or_none=None)) + error = ( + OSError(errno.EMFILE, "too many open files") + if os_error + else sqlite3.OperationalError("unable to open database file") + ) + connect = Mock(side_effect=error) + with pytest.raises( + DatabaseResourceError if os_error else sqlite3.OperationalError + ) as exc: + connect_with_retry(db, connect) + assert (exc.value.__cause__ if os_error else exc.value) is error + assert connect.call_count == 2 diff --git a/tests/test_write_budget.py b/tests/test_write_budget.py new file mode 100644 index 00000000..ba8deb99 --- /dev/null +++ b/tests/test_write_budget.py @@ -0,0 +1,389 @@ +import asyncio +import sqlite3 +import threading +import time + +import pytest + +from datasette.app import Datasette +from datasette.write_budget import ( + DatabaseAdmissionTimeout, + DatabaseQueueFull, + DatabaseReentrancyError, + DatabaseResourceError, +) + + +def make_databases(tmp_path, count=3, **settings): + paths = [] + for i in range(count): + path = tmp_path / f"budget{i}.db" + with sqlite3.connect(path) as conn: + conn.execute("create table t(id)") + conn.close() + paths.append(str(path)) + ds = Datasette(paths, settings={"schema_watch_interval_ms": 0, **settings}) + return ds, [ds.get_database(f"budget{i}") for i in range(count)] + + +async def wait_for(predicate): + deadline = time.monotonic() + 5 + while not predicate(): + assert time.monotonic() < deadline + await asyncio.sleep(0.005) + + +@pytest.mark.asyncio +async def test_many_writers_bounded_fifo_and_fairness(tmp_path): + ds, dbs = make_databases( + tmp_path, 30, max_write_connections=2, max_pending_writes=200 + ) + entered, release = threading.Event(), threading.Event() + order = [] + + def slow(conn): + entered.set() + assert release.wait(5) + + first = asyncio.create_task(dbs[0].execute_write_fn(slow)) + tasks = [] + try: + assert await asyncio.to_thread(entered.wait, 5) + for seq in range(4): + for db in dbs: + + def write(conn, name=db.name, seq=seq): + order.append((name, seq)) + conn.execute("insert into t values (?)", (seq,)) + + tasks.append(asyncio.create_task(db.execute_write_fn(write))) + await asyncio.sleep(0.03) + release.set() + await asyncio.gather(first, *tasks) + assert ds._write_budget.snapshot()["stats"]["peak_writers"] == 2 + for db in dbs: + assert [seq for name, seq in order if name == db.name] == list(range(4)) + assert [ + r[0] for r in (await db.execute("select id from t order by rowid")).rows + ] == list(range(4)) + assert order.index((dbs[-1].name, 0)) < order.index((dbs[0].name, 3)) + finally: + release.set() + await asyncio.gather(first, *tasks, return_exceptions=True) + ds.close() + await wait_for(lambda: ds._write_budget.snapshot()["writers"] == 0) + + +@pytest.mark.asyncio +async def test_queue_full_expiration_and_internal_progress(tmp_path): + ds, dbs = make_databases( + tmp_path, + max_write_connections=1, + max_pending_writes=1, + write_queue_timeout_ms=100, + ) + entered, release = threading.Event(), threading.Event() + ran = [] + + def slow(conn): + entered.set() + assert release.wait(5) + conn.execute("insert into t values (1)") + + first = asyncio.create_task(dbs[0].execute_write_fn(slow)) + queued = None + try: + assert await asyncio.to_thread(entered.wait, 5) + queued = asyncio.create_task( + dbs[1].execute_write_fn(lambda conn: ran.append(True)) + ) + await wait_for(lambda: ds._write_budget.snapshot()["pending"] == 1) + with pytest.raises(DatabaseQueueFull) as exc: + await dbs[2].execute_write_fn(lambda conn: ran.append(True)) + assert exc.value.execution_started is False + # Internal catalog work must not compete for the single user slot. + assert ( + await asyncio.wait_for( + ds._internal_database.execute_write_fn(lambda conn: 42), 2 + ) + == 42 + ) + with pytest.raises(DatabaseAdmissionTimeout): + await queued + assert not ran + assert dbs[1]._write_queue.qsize() == 0 + release.set() + await first + await dbs[2].execute_write("insert into t values (2)") + assert ds._write_budget.snapshot()["stats"]["peak_pending"] == 1 + finally: + release.set() + await asyncio.gather( + first, *([queued] if queued else []), return_exceptions=True + ) + ds.close() + + +@pytest.mark.asyncio +async def test_close_waiting_database_preserves_admission_deadline(tmp_path): + ds, dbs = make_databases( + tmp_path, max_write_connections=1, write_queue_timeout_ms=50 + ) + entered, release = threading.Event(), threading.Event() + + def slow(conn): + entered.set() + assert release.wait(5) + + first = asyncio.create_task(dbs[0].execute_write_fn(slow)) + try: + assert await asyncio.to_thread(entered.wait, 5) + pending = asyncio.create_task( + dbs[1].execute_write_fn(lambda conn: pytest.fail("closed callback ran")) + ) + await wait_for(lambda: ds._write_budget.snapshot()["pending"] == 1) + dbs[1].close() + with pytest.raises(DatabaseAdmissionTimeout): + await pending + assert ds._write_budget.snapshot()["pending"] == 0 + finally: + release.set() + await first + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("threads", [0, 3]) +async def test_memory_pins_capacity_and_recovers_after_close(tmp_path, threads): + ds, dbs = make_databases( + tmp_path, + max_write_connections=1, + write_queue_timeout_ms=50, + num_sql_threads=threads, + ) + memory = ds.add_memory_database("budget_memory") + try: + await memory.execute_write("create table t(id)") + with pytest.raises(DatabaseResourceError): + await dbs[0].execute_write("insert into t values (1)") + assert ( + await memory.execute_write_fn( + lambda conn: conn.execute("select count(*) from t").fetchone()[0] + ) + == 0 + ) + memory.close() + await dbs[0].execute_write("insert into t values (1)") + finally: + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("threads", [0, 3]) +async def test_isolated_writer_does_not_hold_extra_file_connection(tmp_path, threads): + ds, dbs = make_databases(tmp_path, max_write_connections=1, num_sql_threads=threads) + try: + for db in dbs: + await db.execute_write("insert into t values (1)") + assert ( + await db.execute_isolated_fn( + lambda conn, db=db: len(db._all_connections) + ) + == 1 + ) + if threads == 0: + assert sum(db._write_connection is not None for db in dbs) <= 1 + finally: + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["read", "write"]) +async def test_reentrant_database_call_fails_without_deadlock(tmp_path, kind): + ds, dbs = make_databases(tmp_path, max_write_connections=1, max_open_connections=1) + try: + + def callback(conn): + return asyncio.run(dbs[1].execute_write_fn(lambda conn: 1)) + + method = dbs[0].execute_fn if kind == "read" else dbs[0].execute_write_fn + with pytest.raises(DatabaseReentrancyError): + await asyncio.wait_for(method(callback), 2) + finally: + ds.close() + + +@pytest.mark.asyncio +async def test_admission_errors_are_http_503(tmp_path, monkeypatch): + ds, dbs = make_databases(tmp_path) + await ds.invoke_startup() + + async def fail(*args, **kwargs): + raise DatabaseQueueFull("Too many queued database writes") + + monkeypatch.setattr(dbs[0], "execute", fail) + try: + response = await ds.client.get("/budget0/-/query.json?sql=select+1") + assert response.status_code == 503, response.text + assert response.headers["retry-after"] == "1" + assert response.json()["code"] == "database_queue_full" + assert response.json()["execution_started"] is False + finally: + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("threads", [0, 3]) +async def test_unfinished_explicit_transaction_is_rolled_back(tmp_path, threads): + from datasette.connection_pool import ConnectionLeaseError + + ds, dbs = make_databases(tmp_path, num_sql_threads=threads) + try: + with pytest.raises(ConnectionLeaseError, match="rolled back"): + await dbs[0].execute_write_fn( + lambda conn: conn.execute("insert into t values (1)"), transaction=False + ) + await dbs[0].execute_write("insert into t values (2)") + assert [row[0] for row in (await dbs[0].execute("select id from t")).rows] == [ + 2 + ] + finally: + ds.close() + + +@pytest.mark.asyncio +async def test_nonblocking_expiration_notifies_plugin(tmp_path): + from datasette import hookimpl + from datasette.plugins import pm + + notifications = [] + + class Plugin: + @hookimpl + async def write_task_completed(self, database, task_id, exception): + notifications.append((database, task_id, exception)) + + plugin = Plugin() + pm.register(plugin) + ds, dbs = make_databases( + tmp_path, max_write_connections=1, write_queue_timeout_ms=50 + ) + entered, release = threading.Event(), threading.Event() + + def slow(conn): + entered.set() + assert release.wait(5) + + first = asyncio.create_task(dbs[0].execute_write_fn(slow)) + try: + assert await asyncio.to_thread(entered.wait, 5) + task_id = await dbs[1].execute_write_fn( + lambda conn: pytest.fail("expired callback ran"), block=False + ) + await wait_for(lambda: notifications) + assert notifications[0][:2] == (dbs[1].name, task_id) + assert isinstance(notifications[0][2], DatabaseAdmissionTimeout) + release.set() + await first + successful_id = await dbs[1].execute_write_fn(lambda conn: 42, block=False) + await wait_for(lambda: len(notifications) == 2) + assert notifications[1] == (dbs[1].name, successful_id, None) + finally: + release.set() + await first + ds.close() + pm.unregister(plugin) + + +@pytest.mark.asyncio +async def test_connection_open_pressure_retries_before_callback(tmp_path, monkeypatch): + import errno + + from datasette.utils import sqlite3 as sqlite_module + + ds, dbs = make_databases(tmp_path) + await ds.invoke_startup() + await dbs[1].execute("select 1") + original = sqlite_module.connect + calls = [] + + def connect(*args, **kwargs): + if "budget0.db" in str(args[0]): + calls.append(1) + if len(calls) == 1: + raise OSError(errno.EMFILE, "too many files") + return original(*args, **kwargs) + + monkeypatch.setattr(sqlite_module, "connect", connect) + ran = [] + try: + await dbs[0].execute_write_fn(lambda conn: ran.append(1)) + assert calls == [1, 1] + assert ran == [1] + assert not dbs[1]._read_pool_state.idle + finally: + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("threads", [0, 3]) +async def test_resource_error_after_execution_is_not_an_admission_error( + tmp_path, threads, monkeypatch +): + ds, dbs = make_databases(tmp_path, num_sql_threads=threads) + db = dbs[0] + await ds.invoke_startup() + try: + + def committed_then_failed(conn): + conn.execute("insert into t values (1)") + conn.commit() + raise DatabaseQueueFull("A different operation could not be admitted") + + with pytest.raises(RuntimeError, match="callback started") as exc: + await db.execute_write_fn(committed_then_failed, transaction=False) + assert not isinstance(exc.value, DatabaseResourceError) + assert isinstance(exc.value.__cause__, DatabaseQueueFull) + assert (await db.execute("select count(*) from t")).single_value() == 1 + + async def refresh_failed(*args): + raise DatabaseResourceError("Catalog connection unavailable") + + monkeypatch.setattr(ds._schema_watcher, "after_write", refresh_failed) + db._watch_state.needs_scan = True + with pytest.raises(RuntimeError, match="Write completed") as exc: + await db.execute_write("insert into t values (2)") + assert not isinstance(exc.value, DatabaseResourceError) + assert (await db.execute("select count(*) from t")).single_value() == 2 + finally: + ds.close() + + +@pytest.mark.asyncio +async def test_retiring_writer_keeps_its_slot_until_connection_closes( + tmp_path, monkeypatch +): + ds, dbs = make_databases(tmp_path, max_write_connections=1) + closing, release = threading.Event(), threading.Event() + original = dbs[0]._forget_connection + + def slow_close(conn): + closing.set() + assert release.wait(5) + original(conn) + + await dbs[0].execute_write("insert into t values (1)") + monkeypatch.setattr(dbs[0], "_forget_connection", slow_close) + pending = asyncio.create_task(dbs[1].execute_write("insert into t values (2)")) + try: + assert await asyncio.to_thread(closing.wait, 5) + assert not pending.done() + assert dbs[1]._write_thread is None + assert ds._write_budget.snapshot()["writers"] == 1 + release.set() + await pending + assert ds._write_budget.snapshot()["stats"]["peak_writers"] == 1 + finally: + release.set() + await asyncio.gather(pending, return_exceptions=True) + ds.close() diff --git a/tests/test_write_budget_shutdown.py b/tests/test_write_budget_shutdown.py new file mode 100644 index 00000000..86297286 --- /dev/null +++ b/tests/test_write_budget_shutdown.py @@ -0,0 +1,250 @@ +import asyncio +import sqlite3 +import threading +import time +from pathlib import Path + +import pytest + +from datasette.app import Datasette +from datasette.database import Database, DatasetteClosedError +from datasette.write_budget import RETIRE_WRITER, DatabaseAdmissionTimeout + + +def make_databases(tmp_path, count=2, **settings): + paths = [] + for i in range(count): + path = tmp_path / f"shutdown{i}.db" + with sqlite3.connect(path) as conn: + conn.execute("create table t(id)") + conn.close() + paths.append(str(path)) + ds = Datasette(paths, settings={"schema_watch_interval_ms": 0, **settings}) + return ds, [ds.get_database(f"shutdown{i}") for i in range(count)] + + +async def wait_for(predicate): + deadline = time.monotonic() + 5 + while not predicate(): + assert time.monotonic() < deadline + await asyncio.sleep(0.005) + + +def rows(db): + with sqlite3.connect(db.path) as conn: + result = conn.execute("select id from t order by rowid").fetchall() + conn.close() + return [row[0] for row in result] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("retirement_requested", [False, True]) +async def test_close_drains_writes_before_worker_admission( + tmp_path, monkeypatch, retirement_requested +): + ds, (db, _) = make_databases(tmp_path, max_write_connections=1) + admitted = threading.Event() + release = threading.Event() + run = ds._write_budget._run + + def paused_scheduler(): + admitted.set() + assert release.wait(5) + run() + + monkeypatch.setattr(ds._write_budget, "_run", paused_scheduler) + replies = [] + closing = None + try: + for i in range(20): + _, reply = await db._send_to_write_thread( + lambda conn, i=i: conn.execute("insert into t values (?)", [i]), + block=False, + ) + replies.append(reply) + assert admitted.wait(1) + assert db._write_thread is None + if retirement_requested: + # A previous retirement request can precede tasks accepted later. + with db._write_queue.mutex: + db._write_queue.queue.appendleft(RETIRE_WRITER) + closing = asyncio.create_task(asyncio.to_thread(db.close)) + await wait_for(lambda: db._closed) + release.set() + await asyncio.wait_for(closing, 5) + await asyncio.gather(*replies) + assert rows(db) == list(range(20)) + assert ds._write_budget.snapshot()["writers"] == 0 + assert ds._write_budget.snapshot()["stats"]["peak_writers"] == 1 + with pytest.raises(DatasetteClosedError): + await db.execute_write("insert into t values (21)") + finally: + release.set() + if closing is not None: + await closing + await asyncio.gather(*replies, return_exceptions=True) + ds.close() + + +@pytest.mark.asyncio +async def test_close_waits_for_budget_capacity_then_drains(tmp_path): + ds, (first, waiting) = make_databases(tmp_path, max_write_connections=1) + entered, release = threading.Event(), threading.Event() + + def slow(conn): + entered.set() + assert release.wait(5) + + _, first_reply = await first._send_to_write_thread(slow, block=False) + closing = None + waiting_reply = None + try: + assert await asyncio.to_thread(entered.wait, 5) + _, waiting_reply = await waiting._send_to_write_thread( + lambda conn: conn.execute("insert into t values (1)"), block=False + ) + closing = asyncio.create_task(asyncio.to_thread(waiting.close)) + await wait_for(lambda: waiting._closed) + assert ds._write_budget.snapshot()["pending"] == 1 + release.set() + await asyncio.wait_for(closing, 5) + await asyncio.gather(first_reply, waiting_reply) + assert rows(waiting) == [1] + assert ds._write_budget.snapshot()["stats"]["peak_writers"] == 1 + finally: + release.set() + if closing is not None: + await closing + await asyncio.gather( + first_reply, + *([waiting_reply] if waiting_reply else []), + return_exceptions=True, + ) + ds.close() + + +@pytest.mark.asyncio +async def test_close_preserves_admission_deadline_while_capacity_pinned(tmp_path): + ds, (waiting, _) = make_databases( + tmp_path, max_write_connections=1, write_queue_timeout_ms=50 + ) + memory = ds.add_memory_database("pinned") + ran = [] + try: + await memory.execute_write("create table t(id)") + _, reply = await waiting._send_to_write_thread( + lambda conn: ran.append(True), block=False + ) + await asyncio.wait_for(asyncio.to_thread(waiting.close), 2) + with pytest.raises(DatabaseAdmissionTimeout): + await reply + assert not ran + assert ds._write_budget.snapshot()["pending"] == 0 + assert ds._write_budget.snapshot()["stats"]["peak_writers"] == 1 + finally: + ds.close() + + +@pytest.mark.asyncio +async def test_datasette_shutdown_releases_later_memory_writer_before_drain(tmp_path): + ds, (waiting, _) = make_databases( + tmp_path, max_write_connections=1, write_queue_timeout_ms=1000 + ) + memory = ds.add_memory_database("pinned") + try: + await memory.execute_write("create table t(id)") + _, reply = await waiting._send_to_write_thread( + lambda conn: conn.execute("insert into t values (1)"), block=False + ) + await asyncio.wait_for(asyncio.to_thread(ds.close), 5) + await reply + assert rows(waiting) == [1] + assert ds._write_budget.snapshot()["writers"] == 0 + assert ds._write_budget.snapshot()["stats"]["peak_writers"] == 1 + finally: + ds.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("database_kind", ["file", "temporary", "internal"]) +async def test_close_timeout_preserves_running_connection( + tmp_path, monkeypatch, database_kind +): + import datasette.database + + monkeypatch.setattr(datasette.database, "WRITE_SHUTDOWN_TIMEOUT", 0.05) + ds, (db, _) = make_databases(tmp_path, max_write_connections=1) + if database_kind == "temporary": + db = Database(ds, is_temp_disk=True) + ds.add_database(db, name="temporary") + await db.execute_write("create table t(id)") + elif database_kind == "internal": + db = ds._internal_database + await db.execute_write("create table t(id)") + entered, release = threading.Event(), threading.Event() + + def slow(conn): + entered.set() + assert release.wait(5) + conn.execute("insert into t values (1)") + return conn.execute("select id from t").fetchall()[0][0] + + _, running = await db._send_to_write_thread(slow, block=False) + queued = None + try: + assert await asyncio.to_thread(entered.wait, 5) + if database_kind != "internal": + _, queued = await db._send_to_write_thread( + lambda conn: pytest.fail("cancelled callback ran"), block=False + ) + await asyncio.wait_for(asyncio.to_thread(db.close), 2) + if queued is not None: + with pytest.raises(DatasetteClosedError): + await queued + assert Path(db.path).exists() + assert not running.done() + release.set() + assert await running == 1 + if database_kind == "file": + assert rows(db) == [1] + else: + await wait_for(lambda: not Path(db.path).exists()) + await wait_for(lambda: db._write_thread is None) + finally: + release.set() + await asyncio.gather( + running, *([queued] if queued else []), return_exceptions=True + ) + ds.close() + + +@pytest.mark.asyncio +async def test_close_timeout_during_write_connection_open(tmp_path, monkeypatch): + import datasette.database + + monkeypatch.setattr(datasette.database, "WRITE_SHUTDOWN_TIMEOUT", 0.05) + ds, (db, _) = make_databases(tmp_path) + opened, release = threading.Event(), threading.Event() + connect = db.connect + + def paused_connect(**kwargs): + conn = connect(**kwargs) + opened.set() + assert release.wait(5) + return conn + + monkeypatch.setattr(db, "connect", paused_connect) + _, reply = await db._send_to_write_thread( + lambda conn: conn.execute("insert into t values (1)"), block=False + ) + try: + assert await asyncio.to_thread(opened.wait, 5) + await asyncio.wait_for(asyncio.to_thread(db.close), 2) + release.set() + await reply + assert rows(db) == [1] + await wait_for(lambda: db._write_thread is None) + finally: + release.set() + await asyncio.gather(reply, return_exceptions=True) + ds.close() diff --git a/tests/test_write_thread_idle.py b/tests/test_write_thread_idle.py new file mode 100644 index 00000000..807fd1b9 --- /dev/null +++ b/tests/test_write_thread_idle.py @@ -0,0 +1,363 @@ +""" +Tests for the connection_idle_timeout_ms setting: a database's write thread +closes its connection and exits after sitting idle, and the next write +starts a new one. +""" + +import asyncio +import queue +import random +import sqlite3 +import threading +import time + +import pytest + +from datasette.app import Datasette +from datasette.database import Database, DatasetteClosedError + + +def _make_db(tmp_path, name="idle", **settings): + path = tmp_path / f"{name}.db" + conn = sqlite3.connect(path) + conn.execute("create table t (id integer primary key, writer integer, seq integer)") + conn.close() + ds = Datasette([str(path)], settings=settings) + return ds, ds.get_database(name) + + +async def _wait_for(predicate, timeout=5): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + await asyncio.sleep(0.005) + return predicate() + + +def _write_threads(name): + thread_name = f"_execute_writes for database {name}" + return [t for t in threading.enumerate() if t.name == thread_name] + + +@pytest.mark.asyncio +async def test_write_thread_exits_when_idle_and_restarts(tmp_path): + ds, db = _make_db(tmp_path, connection_idle_timeout_ms=50) + try: + await db.execute_write("insert into t (writer, seq) values (1, 1)") + thread = db._write_thread + assert thread is not None and thread.is_alive() + connections_with_writer = len(db._all_connections) + assert await _wait_for(lambda: db._write_thread is None) + thread.join(timeout=5) + assert not thread.is_alive() + # Its write connection was closed and is no longer tracked + assert len(db._all_connections) == connections_with_writer - 1 + assert not _write_threads("idle") + # The next write starts a new thread + await db.execute_write("insert into t (writer, seq) values (1, 2)") + assert db._write_thread is not None and db._write_thread is not thread + assert db._write_threads_started == 2 + count = (await db.execute("select count(*) from t")).single_value() + assert count == 2 + finally: + ds.close() + + +@pytest.mark.asyncio +async def test_write_thread_idle_timeout_zero_keeps_thread(tmp_path): + ds, db = _make_db(tmp_path, connection_idle_timeout_ms=0) + try: + await db.execute_write("insert into t (writer, seq) values (1, 1)") + thread = db._write_thread + await asyncio.sleep(0.2) + assert db._write_thread is thread + assert thread.is_alive() + finally: + ds.close() + assert not thread.is_alive() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("named", (True, False)) +async def test_memory_database_write_thread_never_idles(named): + # Closing the write connection of a memory database would lose its data + ds = Datasette(settings={"connection_idle_timeout_ms": 10}) + if named: + db = ds.add_memory_database("idle_mem") + else: + db = ds.add_database(Database(ds, is_memory=True), name="idle_unnamed") + try: + await db.execute_write("create table t (id integer primary key)") + await db.execute_write("insert into t (id) values (1)") + thread = db._write_thread + await asyncio.sleep(0.1) + assert db._write_thread is thread and thread.is_alive() + count = await db.execute_write_fn( + lambda conn: conn.execute("select count(*) from t").fetchone()[0] + ) + assert count == 1 + finally: + ds.close() + + +@pytest.mark.asyncio +async def test_close_after_idle_exit(tmp_path): + ds, db = _make_db(tmp_path, connection_idle_timeout_ms=10) + await db.execute_write("insert into t (writer, seq) values (1, 1)") + thread = db._write_thread + assert await _wait_for(lambda: db._write_thread is None) + ds.close() + thread.join(timeout=5) + assert not thread.is_alive() + with pytest.raises(DatasetteClosedError): + await db.execute_write("insert into t (writer, seq) values (1, 2)") + + +@pytest.mark.asyncio +async def test_close_drains_queued_writes_with_idle_timeout(tmp_path): + ds, db = _make_db(tmp_path, connection_idle_timeout_ms=10) + release = threading.Event() + + def slow(conn): + release.wait(5) + conn.execute("insert into t (writer, seq) values (0, 0)") + + await db.execute_write_fn(slow, block=False) + futures = [] + for i in range(20): + _, future = await db._send_to_write_thread( + lambda conn, i=i: conn.execute( + "insert into t (writer, seq) values (1, ?)", [i] + ), + block=False, + ) + futures.append(future) + release.set() + # close() is synchronous, so the queued writes complete before it returns + db.close() + await asyncio.wait_for(asyncio.gather(*futures), timeout=5) + conn = sqlite3.connect(db.path) + try: + assert conn.execute("select count(*) from t").fetchone()[0] == 21 + finally: + conn.close() + ds.close() + + +@pytest.mark.asyncio +async def test_connect_failure_retried_after_queue_drains(tmp_path): + # With an idle timeout, a failed connect is reported to the writes that + # were queued, then the thread exits and the next write tries again + ds, db = _make_db(tmp_path, connection_idle_timeout_ms=30000) + + class ConnectError(Exception): + pass + + original_connect = db.connect + + def broken_connect(write=False): + if write: + raise ConnectError("no") + return original_connect(write=write) + + db.connect = broken_connect + try: + with pytest.raises(ConnectError): + await db.execute_write("insert into t (writer, seq) values (1, 1)") + finally: + db.connect = original_connect + assert await _wait_for(lambda: db._write_thread is None) + await db.execute_write("insert into t (writer, seq) values (1, 2)") + assert db._write_threads_started == 2 + ds.close() + + +class _RacyQueue(queue.Queue): + """ + The first time the write thread asks for a task while the queue is + empty, behave as if the idle timeout fired - but first pause, so the + test can enqueue a task in the window between the timeout and the + thread's exit decision. + """ + + def __init__(self): + super().__init__() + self.armed = False + self.window_open = threading.Event() + self.proceed = threading.Event() + + def get(self, block=True, timeout=None): + if self.armed and self.empty(): + self.armed = False + self.window_open.set() + self.proceed.wait(5) + raise queue.Empty + return super().get(block, timeout) + + +@pytest.mark.asyncio +async def test_task_enqueued_between_timeout_and_exit_is_not_stranded(tmp_path): + ds, db = _make_db(tmp_path, connection_idle_timeout_ms=60000) + racy = _RacyQueue() + db._write_queue = racy + try: + racy.armed = True + await db.execute_write("insert into t (writer, seq) values (1, 1)") + # The thread has finished that write and is now inside the window + assert await asyncio.to_thread(racy.window_open.wait, 5) + thread = db._write_thread + assert thread is not None + _, future = await db._send_to_write_thread( + lambda conn: conn.execute("insert into t (writer, seq) values (1, 2)"), + block=False, + ) + # The thread is still registered, so no new thread was started + assert db._write_threads_started == 1 + racy.proceed.set() + await asyncio.wait_for(future, timeout=5) + assert db._write_thread is thread + count = (await db.execute("select count(*) from t")).single_value() + assert count == 2 + finally: + racy.proceed.set() + ds.close() + + +def _recorder(log, writer, seq): + # A cheap task (no disk write), so tasks finish in microseconds and the + # idle timeout keeps firing between bursts + def record(conn): + conn.execute("select 1").fetchone() + log.append((writer, seq, threading.get_ident())) + + return record + + +@pytest.mark.asyncio +@pytest.mark.parametrize("idle_timeout_ms", (1, 5)) +async def test_idle_exit_race_stress(tmp_path, idle_timeout_ms): + """ + Many concurrent writers working in rounds. Each round starts once the + previous round's tasks are done, so the write thread's idle clock starts + at about the same moment, and the writers then wait about 0.8x-1.3x of + the idle timeout before queueing a burst. Tasks therefore keep landing + on both sides of, and right on top of, the thread's exit decision. + Every task must complete, and each writer's tasks must run in the order + they were queued. + """ + ds, db = _make_db(tmp_path, connection_idle_timeout_ms=idle_timeout_ms) + idle_timeout = idle_timeout_ms / 1000 + rng = random.Random(42) + writers = 12 + rounds = 150 if idle_timeout_ms == 1 else 60 + seqs = [0] * writers + log = [] + + async def writer_round(w, delay, burst): + await asyncio.sleep(delay) + pending = [] + for _ in range(burst): + seq = seqs[w] + seqs[w] += 1 + fn = _recorder(log, w, seq) + if seq % 3 == 0: + # Fire-and-forget write + _, future = await db._send_to_write_thread(fn, block=False) + pending.append(future) + else: + await db.execute_write_fn(fn, transaction=bool(seq % 2)) + await asyncio.gather(*pending) + + async def all_rounds(): + for _ in range(rounds): + # One shared delay per round close to the timeout, plus a little + # jitter per writer, so arrivals cluster around the exit decision + base = idle_timeout * rng.uniform(0.8, 1.2) + await asyncio.gather( + *( + writer_round( + w, base + idle_timeout * rng.uniform(0, 0.1), rng.randint(1, 3) + ) + for w in range(writers) + ) + ) + + try: + await asyncio.wait_for(all_rounds(), timeout=120) + assert len(log) == sum(seqs) + for w in range(writers): + assert [seq for writer_id, seq, _ in log if writer_id == w] == list( + range(seqs[w]) + ) + print( + f"idle_timeout_ms={idle_timeout_ms} tasks={len(log)} " + f"threads_started={db._write_threads_started} " + f"distinct_thread_idents={len({ident for _, _, ident in log})} " + f"exits_averted={db._write_thread_exits_averted}" + ) + # Only meaningful if the thread really did churn + assert db._write_threads_started > rounds // 5 + finally: + ds.close() + assert not _write_threads("idle") + + +def test_idle_exit_race_stress_many_event_loops(tmp_path): + """ + Producers on several OS threads, each with its own event loop, race + with the write thread's exit decision for real (no single-loop + ordering between producers). Rounds are synchronised with a barrier, + then each producer waits a random 0.5x-1.5x of the idle timeout. + """ + ds, db = _make_db(tmp_path, connection_idle_timeout_ms=1) + loops = 6 + rounds = 200 + barrier = threading.Barrier(loops) + errors = [] + log = [] + + def run_loop(w): + rng = random.Random(w) + + async def go(): + for seq in range(rounds): + barrier.wait(timeout=30) + # Force some completed retirements as well as racing the + # boundary. Scheduler latency must not decide coverage. + if seq % 4 == 0: + deadline = time.monotonic() + 5 + while db._write_thread is not None: + assert time.monotonic() < deadline + await asyncio.sleep(0.001) + barrier.wait(timeout=30) + await asyncio.sleep(0.001 * rng.uniform(0.5, 1.5)) + await asyncio.wait_for( + db.execute_write_fn(_recorder(log, w, seq)), timeout=10 + ) + + try: + asyncio.run(go()) + except BaseException as e: # noqa: BLE001 + errors.append(e) + barrier.abort() + + threads = [threading.Thread(target=run_loop, args=(w,)) for w in range(loops)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=120) + try: + assert not errors, errors + assert len(log) == loops * rounds + for w in range(loops): + assert [seq for writer_id, seq, _ in log if writer_id == w] == list( + range(rounds) + ) + print( + f"tasks={len(log)} threads_started={db._write_threads_started} " + f"exits_averted={db._write_thread_exits_averted}" + ) + assert db._write_threads_started > rounds // 5 + finally: + ds.close()