Run startup via ASGI lifespan instead of waiting for the first request (#2887)

* Run startup via ASGI lifespan instead of waiting for the first request
* Ensure immutable table counts still precompute when startup ran first

tests/test_lifespan.py

@@ -0,0 +1,259 @@
+"""
+Tests for wiring Datasette startup (setup_db table counts + invoke_startup)
+into the ASGI lifespan protocol.
+
+These exercise Datasette._startup_sequence() via three different callers:
+- AsgiLifespan, by hand-driving lifespan.startup messages (no HTTP request)
+- AsgiRunOnFirstRequest, the fallback for hosts that never send lifespan
+  events (this is what DatasetteClient / plain httpx.ASGITransport uses)
+- Both at once, to prove startup hooks run at most once
+"""
+
+import asyncio
+import contextlib
+import sqlite3
+
+import httpx
+import pytest
+
+from datasette import hookimpl
+from datasette.app import Datasette
+from datasette.database import Database
+from datasette.plugins import pm
+
+
+async def _drive_lifespan_startup(app):
+    """Send a single lifespan.startup message into app's ASGI lifespan loop
+    and return the list of messages sent back - without ever sending
+    lifespan.shutdown. Mirrors what a real server does: after startup
+    completes it parks waiting for the next event. We cancel that wait
+    once we've observed the startup response, rather than closing the
+    Datasette instance down with a shutdown message.
+    """
+    messages_sent = []
+    startup_responded = asyncio.Event()
+    delivered = False
+
+    async def receive():
+        nonlocal delivered
+        if not delivered:
+            delivered = True
+            return {"type": "lifespan.startup"}
+        # No further messages: block until the task is cancelled below,
+        # same as a real server parked waiting for lifespan.shutdown.
+        await asyncio.Event().wait()
+
+    async def send(message):
+        messages_sent.append(message)
+        startup_responded.set()
+
+    task = asyncio.create_task(app({"type": "lifespan"}, receive, send))
+    try:
+        await asyncio.wait_for(startup_responded.wait(), timeout=5)
+    finally:
+        task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await task
+    return messages_sent
+
+
+@pytest.mark.asyncio
+async def test_lifespan_startup_runs_before_any_request():
+    ds = Datasette(memory=True)
+    assert ds._startup_invoked is False
+    app = ds.app()
+
+    messages = await _drive_lifespan_startup(app)
+
+    assert {"type": "lifespan.startup.complete"} in messages
+    assert ds._startup_invoked is True
+    # Internal catalog tables should be populated too, entirely without an
+    # HTTP request having been made.
+    internal_db = ds.get_internal_database()
+    databases = await internal_db.execute("select * from catalog_databases")
+    assert len(databases.rows) >= 1
+
+
+@pytest.mark.asyncio
+async def test_lifespan_startup_failure_reports_lifespan_startup_failed():
+    class RaisingStartupPlugin:
+        __name__ = "RaisingStartupPlugin"
+
+        @hookimpl
+        def startup(self, datasette):
+            async def inner():
+                raise RuntimeError("boom from startup hook")
+
+            return inner
+
+    ds = Datasette(memory=True)
+    pm.register(RaisingStartupPlugin(), name="raising_startup_plugin")
+    try:
+        app = ds.app()
+        messages = await _drive_lifespan_startup(app)
+    finally:
+        pm.unregister(name="raising_startup_plugin")
+
+    assert messages == [
+        {"type": "lifespan.startup.failed", "message": "boom from startup hook"}
+    ]
+    # The exception happened before invoke_startup() got to the end of its
+    # body, so startup is not considered to have completed.
+    assert ds._startup_invoked is False
+
+
+@pytest.mark.asyncio
+async def test_startup_runs_exactly_once_across_lifespan_and_first_request():
+    call_count = {"n": 0}
+
+    class CountingStartupPlugin:
+        __name__ = "CountingStartupPlugin"
+
+        @hookimpl
+        def startup(self, datasette):
+            async def inner():
+                call_count["n"] += 1
+
+            return inner
+
+    ds = Datasette(memory=True)
+    pm.register(CountingStartupPlugin(), name="counting_startup_plugin")
+    try:
+        # Build the ASGI app once, the way a real deployment does - and
+        # reuse the SAME app instance for both the lifespan drive and the
+        # HTTP requests below, since a fresh ds.app() call would reset the
+        # AsgiRunOnFirstRequest fallback's state.
+        app = ds.app()
+
+        messages = await _drive_lifespan_startup(app)
+        assert {"type": "lifespan.startup.complete"} in messages
+        assert call_count["n"] == 1
+
+        # A first HTTP request (as if the host never sent lifespan events,
+        # or lifespan already ran) should not run the hook again.
+        transport = httpx.ASGITransport(app=app)
+        async with httpx.AsyncClient(
+            transport=transport, base_url="http://localhost"
+        ) as client:
+            response1 = await client.get("/-/versions.json")
+            assert response1.status_code == 200
+            # ... nor should a second, repeat request.
+            response2 = await client.get("/-/versions.json")
+            assert response2.status_code == 200
+    finally:
+        pm.unregister(name="counting_startup_plugin")
+
+    assert call_count["n"] == 1
+
+
+@pytest.mark.asyncio
+async def test_no_lifespan_first_request_still_triggers_startup():
+    # Pin today's behavior: a client that never drives ASGI lifespan events
+    # at all (like httpx.ASGITransport, which DatasetteClient uses) still
+    # gets startup armed by the AsgiRunOnFirstRequest fallback.
+    ds = Datasette(memory=True)
+    assert ds._startup_invoked is False
+    app = ds.app()
+    transport = httpx.ASGITransport(app=app)
+    async with httpx.AsyncClient(
+        transport=transport, base_url="http://localhost"
+    ) as client:
+        response = await client.get("/-/versions.json")
+        assert response.status_code == 200
+
+    assert ds._startup_invoked is True
+    internal_db = ds.get_internal_database()
+    databases = await internal_db.execute("select * from catalog_databases")
+    assert len(databases.rows) >= 1
+
+
+@pytest.mark.asyncio
+async def test_datasette_client_first_request_triggers_startup():
+    # Same as above, but through the real DatasetteClient (ds.client) that
+    # plugins and tests actually use, to confirm nothing regressed there.
+    ds = Datasette(memory=True)
+    assert ds._startup_invoked is False
+    response = await ds.client.get("/-/versions.json")
+    assert response.status_code == 200
+    assert ds._startup_invoked is True
+
+
+@pytest.mark.asyncio
+async def test_concurrent_first_requests_all_wait_for_slow_startup():
+    call_count = {"n": 0}
+
+    class SlowStartupPlugin:
+        __name__ = "SlowStartupPlugin"
+
+        @hookimpl
+        def startup(self, datasette):
+            async def inner():
+                call_count["n"] += 1
+                await asyncio.sleep(0.2)
+
+            return inner
+
+    ds = Datasette(memory=True)
+    pm.register(SlowStartupPlugin(), name="slow_startup_plugin")
+    try:
+        app = ds.app()
+        transport = httpx.ASGITransport(app=app)
+        async with httpx.AsyncClient(
+            transport=transport, base_url="http://localhost"
+        ) as client:
+            responses = await asyncio.gather(
+                *[client.get("/-/versions.json") for _ in range(10)]
+            )
+    finally:
+        pm.unregister(name="slow_startup_plugin")
+
+    # Every one of the 10 simultaneous first requests must have blocked
+    # until startup actually finished, not raced ahead of it.
+    assert all(response.status_code == 200 for response in responses)
+    assert call_count["n"] == 1
+    assert ds._startup_invoked is True
+
+
+@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.
+    db_path = tmp_path / "immutable.db"
+    conn = sqlite3.connect(str(db_path))
+    conn.execute("create table t (id integer primary key)")
+    conn.commit()
+    conn.close()
+
+    ds = Datasette([], immutables=[str(db_path)])
+
+    call_count = {"n": 0}
+    original_table_counts = Database.table_counts
+
+    async def counting_table_counts(self, *args, **kwargs):
+        call_count["n"] += 1
+        return await original_table_counts(self, *args, **kwargs)
+
+    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()
+    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
+
+    # Idempotency: a second call must not recompute.
+    await ds._startup_sequence()
+    assert call_count["n"] == 1