Limit total facet suggestion time per page (#862)

datasette/facets.py

@@ -157,6 +157,7 @@ class ColumnFacet(Facet):
         columns = await self.get_columns(self.sql, self.params)
         facet_size = self.get_facet_size()
         suggested_facets = []
+        self._suggestions_so_far = suggested_facets
         already_enabled = [c["config"]["simple"] for c in self.get_configs()]
         for column in columns:
             if column in already_enabled:
@@ -311,6 +312,7 @@ class ArrayFacet(Facet):
     async def suggest(self):
         columns = await self.get_columns(self.sql, self.params)
         suggested_facets = []
+        self._suggestions_so_far = suggested_facets
         already_enabled = [c["config"]["simple"] for c in self.get_configs()]
         for column in columns:
             if column in already_enabled:
@@ -469,6 +471,7 @@ class DateFacet(Facet):
         columns = await self.get_columns(self.sql, self.params)
         already_enabled = [c["config"]["simple"] for c in self.get_configs()]
         suggested_facets = []
+        self._suggestions_so_far = suggested_facets
         for column in columns:
             if column in already_enabled:
                 continue

datasette/views/table_extras.py

@@ -1,4 +1,6 @@
+import asyncio
 import itertools
+import time
 from dataclasses import dataclass
 from typing import ClassVar
 
@@ -290,10 +292,30 @@ class SuggestedFacetsExtra(Extra):
             and not context.nofacet
             and not context.nosuggest
         ):
-            facet_suggest_awaitables = [facet.suggest() for facet in facet_instances]
-            for suggest_result in await context.run_sequential(
-                *facet_suggest_awaitables
-            ):
+            time_limit_ms = context.datasette.setting("facet_suggest_time_limit_ms")
+            deadline = (
+                time.monotonic() + 10 * time_limit_ms / 1000
+                if time_limit_ms > 0
+                else None
+            )
+            for facet in facet_instances:
+                if deadline is None:
+                    (suggest_result,) = await context.run_sequential(facet.suggest())
+                else:
+                    remaining = deadline - time.monotonic()
+                    if remaining <= 0:
+                        break
+                    try:
+                        (suggest_result,) = await asyncio.wait_for(
+                            context.run_sequential(facet.suggest()), timeout=remaining
+                        )
+                    except asyncio.TimeoutError:
+                        if time.monotonic() < deadline:
+                            raise
+                        suggested_facets.extend(
+                            getattr(facet, "_suggestions_so_far", [])
+                        )
+                        break
                 suggested_facets.extend(suggest_result)
         return suggested_facets
 

docs/settings.rst

@@ -284,6 +284,10 @@ Default: ``50``
 
 When Datasette calculates suggested facets it needs to run a SQL query for every column in your table. This time limit, in milliseconds, applies separately to each query. If the time limit is exceeded the column will not be suggested as a facet.
 
+Facet suggestions for a page stop after ten times this limit in total. Any
+suggestions found before the total limit are still returned. A value of ``0``
+disables this total limit.
+
 You can increase this time limit like so::
 
     datasette mydatabase.db --setting facet_suggest_time_limit_ms 500

tests/test_facets.py

@@ -1,4 +1,6 @@
+import asyncio
 import json
+from types import SimpleNamespace
 from urllib.parse import parse_qsl, urlsplit
 
 import pytest
@@ -14,6 +16,8 @@ from datasette.facets import (
 )
 from datasette.utils import detect_json1
 from datasette.utils.asgi import Request
+from datasette.views.table import run_sequential
+from datasette.views.table_extras import SuggestedFacetsExtra
 
 from .fixtures import make_app_client
 
@@ -74,6 +78,132 @@ async def test_column_facet_suggest(ds_client):
     ] == suggestions
 
 
+@pytest.mark.asyncio
+async def test_suggested_facets_stop_after_page_time_limit():
+    class SuggestingFacet:
+        def __init__(self, name, wait=False):
+            self.name = name
+            self.wait = wait
+            self.called = False
+            self.cancelled = False
+
+        async def suggest(self):
+            self.called = True
+            if self.wait:
+                self._suggestions_so_far = [{"name": self.name}]
+                try:
+                    await asyncio.Event().wait()
+                except asyncio.CancelledError:
+                    self.cancelled = True
+                    raise
+            return [{"name": self.name}]
+
+    first = SuggestingFacet("first")
+    slow = SuggestingFacet("slow", wait=True)
+    last = SuggestingFacet("last")
+    settings = {
+        "suggest_facets": True,
+        "allow_facet": True,
+        "facet_suggest_time_limit_ms": 10,
+    }
+    context = SimpleNamespace(
+        datasette=SimpleNamespace(setting=settings.__getitem__),
+        next_arg=None,
+        nofacet=False,
+        nosuggest=False,
+        run_sequential=run_sequential,
+    )
+
+    suggestions = await asyncio.wait_for(
+        SuggestedFacetsExtra().resolve(context, [first, slow, last]), timeout=2
+    )
+
+    assert suggestions == [{"name": "first"}, {"name": "slow"}]
+    assert first.called and slow.called and slow.cancelled
+    assert not last.called
+
+
+@pytest.mark.asyncio
+async def test_suggested_facets_do_not_hide_facet_timeout():
+    class FailingFacet:
+        async def suggest(self):
+            raise asyncio.TimeoutError("facet failed")
+
+    settings = {
+        "suggest_facets": True,
+        "allow_facet": True,
+        "facet_suggest_time_limit_ms": 50,
+    }
+    context = SimpleNamespace(
+        datasette=SimpleNamespace(setting=settings.__getitem__),
+        next_arg=None,
+        nofacet=False,
+        nosuggest=False,
+        run_sequential=run_sequential,
+    )
+
+    with pytest.raises(asyncio.TimeoutError, match="facet failed"):
+        await SuggestedFacetsExtra().resolve(context, [FailingFacet()])
+
+
+@pytest.mark.asyncio
+async def test_suggested_facets_zero_time_limit_keeps_existing_behavior():
+    class FastFacet:
+        async def suggest(self):
+            return [{"name": "available"}]
+
+    settings = {
+        "suggest_facets": True,
+        "allow_facet": True,
+        "facet_suggest_time_limit_ms": 0,
+    }
+    context = SimpleNamespace(
+        datasette=SimpleNamespace(setting=settings.__getitem__),
+        next_arg=None,
+        nofacet=False,
+        nosuggest=False,
+        run_sequential=run_sequential,
+    )
+
+    suggestions = await SuggestedFacetsExtra().resolve(context, [FastFacet()])
+    assert suggestions == [{"name": "available"}]
+
+
+@pytest.mark.asyncio
+async def test_suggested_facets_keep_completed_columns_on_timeout(monkeypatch):
+    ds = Datasette([], memory=True, settings={"facet_suggest_time_limit_ms": 20})
+    db = ds.add_database(Database(ds, memory_name="suggest_deadline"))
+    await db.execute_write("create table items (first text, second text)")
+    for row in (("a", "x"), ("a", "y"), ("b", "y")):
+        await db.execute_write("insert into items values (?, ?)", row)
+
+    execute = ds.execute
+    suggestion_queries = 0
+
+    async def delay_second_suggestion(*args, **kwargs):
+        nonlocal suggestion_queries
+        sql = args[1].lower()
+        if "group by value" in sql and "select rowid as value" not in sql:
+            suggestion_queries += 1
+            if suggestion_queries == 2:
+                await asyncio.Event().wait()
+        return await execute(*args, **kwargs)
+
+    monkeypatch.setattr(ds, "execute", delay_second_suggestion)
+    try:
+        response = await asyncio.wait_for(
+            ds.client.get("/suggest_deadline/items.json?_extra=suggested_facets"),
+            timeout=3,
+        )
+        assert response.status_code == 200
+        assert [facet["name"] for facet in response.json()["suggested_facets"]] == [
+            "first"
+        ]
+        assert suggestion_queries == 2
+    finally:
+        ds.close()
+
+
 @pytest.mark.asyncio
 async def test_column_facet_suggest_skip_if_already_selected(ds_client):
     facet = ColumnFacet(