mirror of
https://github.com/simonw/datasette.git
synced 2026-09-02 14:44:07 +02:00
Run datasette serve startup and uvicorn on a single event loop (#2886)
* Run datasette serve startup and uvicorn on a single event loop * Move the serve-subprocess test plumbing into a conftest fixture * Fix datasette-litestream URL and trim marker-task test comments * Explain why serve_with_plugins needs a subprocess and plugin files * Apply ruff 0.16 and black fixes * Tweaked some comments
This commit is contained in:
Alex Garcia
GitHub
parent
0337fba234
commit
e78b8a2e6a
4 changed files with 253 additions and 41 deletions
|
|
@ -663,16 +663,6 @@ def serve(
|
|||
# Private utility mechanism for writing unit tests
|
||||
return ds
|
||||
|
||||
# Run async soundness checks before startup hooks, since invoke_startup
|
||||
# now populates internal tables which requires querying each database
|
||||
run_sync(lambda: check_databases(ds))
|
||||
|
||||
# Run the "startup" plugin hooks
|
||||
try:
|
||||
run_sync(ds.invoke_startup)
|
||||
except StartupError as e:
|
||||
raise click.ClickException(e.args[0])
|
||||
|
||||
if headers and not get:
|
||||
raise click.ClickException("--headers can only be used with --get")
|
||||
|
||||
|
|
@ -680,6 +670,14 @@ def serve(
|
|||
raise click.ClickException("--token can only be used with --get")
|
||||
|
||||
if get:
|
||||
# --get means we don't run Uvicorn at all
|
||||
run_sync(lambda: check_databases(ds))
|
||||
|
||||
try:
|
||||
run_sync(ds.invoke_startup)
|
||||
except StartupError as e:
|
||||
raise click.ClickException(e.args[0])
|
||||
|
||||
client = TestClient(ds)
|
||||
request_headers = {}
|
||||
if token:
|
||||
|
|
@ -704,34 +702,51 @@ def serve(
|
|||
sys.exit(exit_code)
|
||||
return
|
||||
|
||||
# Start the server
|
||||
url = None
|
||||
if root:
|
||||
ds.root_enabled = True
|
||||
url = "http://{}:{}{}?token={}".format(
|
||||
host, port, ds.urls.path("-/auth-token"), ds._root_token
|
||||
)
|
||||
click.echo(url)
|
||||
if open_browser:
|
||||
if url is None:
|
||||
# Figure out most convenient URL - to table, database or homepage
|
||||
path = run_sync(lambda: initial_path_for_datasette(ds))
|
||||
url = f"http://{host}:{port}{path}"
|
||||
webbrowser.open(url)
|
||||
uvicorn_kwargs = {
|
||||
"host": host,
|
||||
"port": port,
|
||||
"log_level": "info",
|
||||
"lifespan": "on",
|
||||
"workers": 1,
|
||||
}
|
||||
if uds:
|
||||
uvicorn_kwargs["uds"] = uds
|
||||
if ssl_keyfile:
|
||||
uvicorn_kwargs["ssl_keyfile"] = ssl_keyfile
|
||||
if ssl_certfile:
|
||||
uvicorn_kwargs["ssl_certfile"] = ssl_certfile
|
||||
uvicorn.run(ds.app(), **uvicorn_kwargs)
|
||||
# check_databases, invoke_startup() and the uvicorn server all run on a
|
||||
# single event loop, so that anything a plugin's "startup" hook schedules
|
||||
# on the loop (asyncio.create_task, Lock/Queue/Event objects, ...) is
|
||||
# still alive when the server starts handling requests.
|
||||
async def _serve_async():
|
||||
# Populate internal catalog tables before invoke_startup
|
||||
await check_databases(ds)
|
||||
|
||||
# Run the "startup" plugin hooks
|
||||
try:
|
||||
await ds.invoke_startup()
|
||||
except StartupError as e:
|
||||
raise click.ClickException(e.args[0])
|
||||
|
||||
# Start the server
|
||||
url = None
|
||||
if root:
|
||||
ds.root_enabled = True
|
||||
url = "http://{}:{}{}?token={}".format(
|
||||
host, port, ds.urls.path("-/auth-token"), ds._root_token
|
||||
)
|
||||
click.echo(url)
|
||||
if open_browser:
|
||||
if url is None:
|
||||
# Figure out most convenient URL - to table, database or homepage
|
||||
path = await initial_path_for_datasette(ds)
|
||||
url = f"http://{host}:{port}{path}"
|
||||
webbrowser.open(url)
|
||||
uvicorn_kwargs = {
|
||||
"host": host,
|
||||
"port": port,
|
||||
"log_level": "info",
|
||||
"lifespan": "on",
|
||||
"workers": 1,
|
||||
}
|
||||
if uds:
|
||||
uvicorn_kwargs["uds"] = uds
|
||||
if ssl_keyfile:
|
||||
uvicorn_kwargs["ssl_keyfile"] = ssl_keyfile
|
||||
if ssl_certfile:
|
||||
uvicorn_kwargs["ssl_certfile"] = ssl_certfile
|
||||
server = uvicorn.Server(uvicorn.Config(ds.app(), **uvicorn_kwargs))
|
||||
await server.serve()
|
||||
|
||||
asyncio.run(_serve_async())
|
||||
|
||||
|
||||
@cli.command()
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ dependencies = [
|
|||
"hupper>=1.9",
|
||||
"httpx>=0.20,<1.0",
|
||||
"pluggy>=1.0",
|
||||
"uvicorn>=0.11",
|
||||
"uvicorn>=0.29",
|
||||
"aiofiles>=0.4",
|
||||
"PyYAML>=5.3",
|
||||
"mergedeep>=1.1.1",
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@ import importlib.metadata
|
|||
import os
|
||||
import pathlib
|
||||
import re
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
|
|
@ -32,17 +33,31 @@ UNDOCUMENTED_PERMISSIONS = {
|
|||
}
|
||||
|
||||
|
||||
def wait_until_responds(url, timeout=5.0, client=httpx, **kwargs):
|
||||
def wait_until_responds(url, timeout=5.0, client=httpx, process=None, **kwargs):
|
||||
start = time.time()
|
||||
while time.time() - start < timeout:
|
||||
# If the server died there is no point waiting out the timeout - fail
|
||||
# now, with its output, instead of after `timeout` seconds of silence
|
||||
if process is not None and process.poll() is not None:
|
||||
raise AssertionError(
|
||||
"Server exited early with returncode {}\n{}".format(
|
||||
process.returncode, process.stdout.read().decode("utf-8")
|
||||
)
|
||||
)
|
||||
try:
|
||||
client.get(url, **kwargs)
|
||||
return
|
||||
except httpx.ConnectError:
|
||||
except httpx.TransportError:
|
||||
time.sleep(0.1)
|
||||
raise AssertionError(f"Timed out waiting for {url} to respond")
|
||||
|
||||
|
||||
def find_free_port():
|
||||
with socket.socket() as sock:
|
||||
sock.bind(("127.0.0.1", 0))
|
||||
return sock.getsockname()[1]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bare_ds():
|
||||
"""
|
||||
|
|
@ -301,6 +316,71 @@ def ds_unix_domain_socket_server(tmp_path_factory):
|
|||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def serve_with_plugins(tmp_path):
|
||||
"""Factory fixture for starting ``datasette serve`` in a subprocess with
|
||||
plugins written to a temporary ``--plugins-dir``.
|
||||
|
||||
For tests that need the real serve path: event-loop wiring, exit codes,
|
||||
signals. The usual in-process ``pm.register`` plugin pattern can't reach
|
||||
a subprocess, so plugin source is written out as importable files instead.
|
||||
|
||||
Unlike ``ds_localhost_http_server`` this is function-scoped and takes a
|
||||
fresh port each time, because each test needs its own plugins. Call it as::
|
||||
|
||||
proc, port = serve_with_plugins({"my_plugin": PLUGIN_SOURCE})
|
||||
|
||||
``plugins`` maps module name to Python source. Pass
|
||||
``wait_for_startup=False`` when the server is expected to fail during
|
||||
startup rather than begin serving. Extra CLI arguments are passed through.
|
||||
Every process started is terminated when the test ends.
|
||||
"""
|
||||
processes = []
|
||||
|
||||
def start(plugins, *extra_args, wait_for_startup=True):
|
||||
plugins_dir = tmp_path / "plugins"
|
||||
plugins_dir.mkdir(exist_ok=True)
|
||||
for module_name, source in plugins.items():
|
||||
(plugins_dir / f"{module_name}.py").write_text(source, "utf-8")
|
||||
port = find_free_port()
|
||||
proc = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"datasette",
|
||||
"--memory",
|
||||
"--plugins-dir",
|
||||
str(plugins_dir),
|
||||
"-h",
|
||||
"127.0.0.1",
|
||||
"-p",
|
||||
str(port),
|
||||
*extra_args,
|
||||
],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
# Avoid FileNotFoundError: [Errno 2] No such file or directory:
|
||||
cwd=tempfile.gettempdir(),
|
||||
)
|
||||
processes.append(proc)
|
||||
if wait_for_startup:
|
||||
wait_until_responds(
|
||||
f"http://127.0.0.1:{port}/-/versions.json", process=proc
|
||||
)
|
||||
return proc, port
|
||||
|
||||
yield start
|
||||
|
||||
for proc in processes:
|
||||
if proc.poll() is None:
|
||||
proc.terminate()
|
||||
try:
|
||||
proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
proc.wait()
|
||||
|
||||
|
||||
# Import fixtures from fixtures.py to make them available
|
||||
from .fixtures import ( # noqa: F401
|
||||
TEMP_PLUGIN_SECRET_FILE,
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import socket
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -28,3 +29,119 @@ def test_serve_unix_domain_socket(ds_unix_domain_socket_server):
|
|||
"path": "/_memory",
|
||||
"tables": [],
|
||||
}.items() <= response.json().items()
|
||||
|
||||
|
||||
# Shaped after datasette-litestream's startup hook, which schedules a
|
||||
# background task with asyncio.get_running_loop().create_task(...):
|
||||
# https://github.com/datasette/datasette-litestream
|
||||
MARKER_TASK_PLUGIN = """
|
||||
import asyncio
|
||||
from datasette import hookimpl
|
||||
from datasette.utils.asgi import Response
|
||||
|
||||
|
||||
@hookimpl
|
||||
def startup(datasette):
|
||||
datasette._startup_calls = getattr(datasette, "_startup_calls", 0) + 1
|
||||
|
||||
async def _mark():
|
||||
# Must await before setting the flag: a task with no internal
|
||||
# await point could finish on the throwaway loop before it
|
||||
# closed, masking the regression this test guards against.
|
||||
await asyncio.sleep(0.2)
|
||||
datasette._marker_task_ran = True
|
||||
|
||||
asyncio.get_running_loop().create_task(_mark())
|
||||
|
||||
|
||||
@hookimpl
|
||||
def register_routes():
|
||||
async def marker_status(datasette):
|
||||
return Response.json(
|
||||
{
|
||||
"marker_task_ran": getattr(datasette, "_marker_task_ran", False),
|
||||
"startup_calls": getattr(datasette, "_startup_calls", 0),
|
||||
}
|
||||
)
|
||||
|
||||
return [(r"^/-/marker-task-ran$", marker_status)]
|
||||
"""
|
||||
|
||||
|
||||
STARTUP_ERROR_PLUGIN = """
|
||||
from datasette import hookimpl
|
||||
from datasette.utils import StartupError
|
||||
|
||||
|
||||
@hookimpl
|
||||
def startup(datasette):
|
||||
raise StartupError("boom from plugin")
|
||||
"""
|
||||
|
||||
|
||||
@pytest.mark.serial
|
||||
def test_startup_hook_background_task_runs_on_serving_loop(serve_with_plugins):
|
||||
"""
|
||||
Litestream-shaped regression test: a startup hook that does
|
||||
asyncio.get_running_loop().create_task(...) must have that task
|
||||
actually execute before/while the server is handling requests. This
|
||||
only holds if invoke_startup() and uvicorn.Server.serve() share one
|
||||
event loop. This test fails against unmodified main, where
|
||||
invoke_startup() runs on a throwaway loop that is closed before
|
||||
uvicorn opens its own loop to serve.
|
||||
"""
|
||||
_, port = serve_with_plugins({"marker_task_plugin": MARKER_TASK_PLUGIN})
|
||||
# The fixture has already waited for the server to answer requests. The
|
||||
# marker task deliberately awaits before setting its flag, so poll for a
|
||||
# moment rather than assuming it landed before the first request arrived.
|
||||
deadline = time.time() + 3.0
|
||||
payload = {}
|
||||
while time.time() < deadline:
|
||||
payload = httpx.get(
|
||||
f"http://127.0.0.1:{port}/-/marker-task-ran", timeout=1.0
|
||||
).json()
|
||||
if payload["marker_task_ran"]:
|
||||
break
|
||||
time.sleep(0.05)
|
||||
assert payload.get("marker_task_ran"), (
|
||||
"The startup hook's asyncio.create_task(...) never ran - "
|
||||
"invoke_startup() and the server are not sharing an event loop"
|
||||
)
|
||||
# Polling above means this test would also pass if the startup hook were
|
||||
# re-run on the serving loop by the first-request fallback - which would
|
||||
# hide exactly the bug being tested. invoke_startup() is idempotent today
|
||||
# so that cannot happen; assert it explicitly so that if the idempotency
|
||||
# guard is ever removed this test fails loudly instead of silently
|
||||
# becoming a no-op.
|
||||
assert payload["startup_calls"] == 1, (
|
||||
"startup hook ran {} times - the marker may have been set by a "
|
||||
"re-run on the serving loop rather than by the original task".format(
|
||||
payload["startup_calls"]
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.serial
|
||||
def test_startup_error_fails_fast_before_port_binds(serve_with_plugins):
|
||||
"""
|
||||
A "startup" plugin hook that raises StartupError must fail fast: print
|
||||
the message, exit non-zero, and never accept a connection on the port -
|
||||
the failure must happen before uvicorn.Server binds the socket.
|
||||
"""
|
||||
proc, port = serve_with_plugins(
|
||||
{"startup_error_plugin": STARTUP_ERROR_PLUGIN}, wait_for_startup=False
|
||||
)
|
||||
stdout, _ = proc.communicate(timeout=15)
|
||||
output = stdout.decode("utf-8")
|
||||
assert proc.returncode not in (0, None), output
|
||||
assert "boom from plugin" in output, output
|
||||
|
||||
# Nothing is listening on the port now the process has exited. This
|
||||
# confirms the socket was not left bound; on its own it cannot prove the
|
||||
# failure preceded the bind, since a port nothing ever touched also
|
||||
# refuses connections.
|
||||
with (
|
||||
pytest.raises(OSError),
|
||||
socket.create_connection(("127.0.0.1", port), timeout=0.2),
|
||||
):
|
||||
pass
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue