mirror of
https://github.com/simonw/datasette.git
synced 2026-09-03 15:14:08 +02:00
Run datasette serve startup and uvicorn on a single event loop
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Alex Garcia
parent
0337fba234
commit
cd9daa266c
3 changed files with 238 additions and 39 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,16 @@ def serve(
|
|||
raise click.ClickException("--token can only be used with --get")
|
||||
|
||||
if get:
|
||||
# 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])
|
||||
|
||||
client = TestClient(ds)
|
||||
request_headers = {}
|
||||
if token:
|
||||
|
|
@ -704,34 +704,52 @@ 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():
|
||||
# Run async soundness checks before startup hooks, since invoke_startup
|
||||
# now populates internal tables which requires querying each database
|
||||
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",
|
||||
|
|
|
|||
|
|
@ -1,4 +1,8 @@
|
|||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -28,3 +32,180 @@ def test_serve_unix_domain_socket(ds_unix_domain_socket_server):
|
|||
"path": "/_memory",
|
||||
"tables": [],
|
||||
}.items() <= response.json().items()
|
||||
|
||||
|
||||
def _find_free_port():
|
||||
with socket.socket() as sock:
|
||||
sock.bind(("127.0.0.1", 0))
|
||||
return sock.getsockname()[1]
|
||||
|
||||
|
||||
# Shaped after datasette-litestream's (sync) startup hook, which schedules a
|
||||
# background task with asyncio.get_running_loop().create_task(...):
|
||||
# https://github.com/simonw/datasette-litestream/blob/main/datasette_litestream/__init__.py
|
||||
# That only has a chance to actually run if invoke_startup() executes on the
|
||||
# same event loop that goes on to serve requests - if it runs on a throwaway
|
||||
# loop that gets closed straight after (as on unmodified main), the task is
|
||||
# scheduled but never gets a turn before the loop is torn down.
|
||||
MARKER_TASK_PLUGIN = '''
|
||||
import asyncio
|
||||
from datasette import hookimpl
|
||||
from datasette.utils.asgi import Response
|
||||
|
||||
|
||||
@hookimpl
|
||||
def startup(datasette):
|
||||
async def _mark():
|
||||
# The await is essential to the regression: a task with no
|
||||
# internal await point can complete during the brief window
|
||||
# between run_until_complete()'s coroutine finishing and the
|
||||
# temporary loop actually stopping, masking the bug this test
|
||||
# guards against. Real background tasks (like
|
||||
# datasette-litestream's credential_refresh_loop) always have an
|
||||
# internal await, and never get to resume once their throwaway
|
||||
# loop is closed.
|
||||
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)}
|
||||
)
|
||||
|
||||
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(tmp_path):
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
plugins_dir = tmp_path / "plugins"
|
||||
plugins_dir.mkdir()
|
||||
(plugins_dir / "marker_task_plugin.py").write_text(MARKER_TASK_PLUGIN, "utf-8")
|
||||
|
||||
port = _find_free_port()
|
||||
ds_proc = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"datasette",
|
||||
"--memory",
|
||||
"--plugins-dir",
|
||||
str(plugins_dir),
|
||||
"-p",
|
||||
str(port),
|
||||
],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
cwd=tempfile.gettempdir(),
|
||||
)
|
||||
try:
|
||||
url = f"http://localhost:{port}/-/marker-task-ran"
|
||||
deadline = time.time() + 15.0
|
||||
marker_task_ran = False
|
||||
while time.time() < deadline:
|
||||
if ds_proc.poll() is not None:
|
||||
raise AssertionError(
|
||||
"datasette serve exited early\n"
|
||||
+ ds_proc.stdout.read().decode("utf-8")
|
||||
)
|
||||
try:
|
||||
response = httpx.get(url, timeout=1.0)
|
||||
except httpx.TransportError:
|
||||
time.sleep(0.1)
|
||||
continue
|
||||
if response.status_code == 200 and response.json().get(
|
||||
"marker_task_ran"
|
||||
):
|
||||
marker_task_ran = True
|
||||
break
|
||||
time.sleep(0.1)
|
||||
assert marker_task_ran, (
|
||||
"The startup hook's asyncio.create_task(...) never ran - "
|
||||
"invoke_startup() and the server are not sharing an event loop"
|
||||
)
|
||||
finally:
|
||||
ds_proc.terminate()
|
||||
try:
|
||||
ds_proc.wait(timeout=5)
|
||||
except subprocess.TimeoutExpired:
|
||||
ds_proc.kill()
|
||||
ds_proc.wait()
|
||||
|
||||
|
||||
@pytest.mark.serial
|
||||
def test_startup_error_fails_fast_before_port_binds(tmp_path):
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
plugins_dir = tmp_path / "plugins"
|
||||
plugins_dir.mkdir()
|
||||
(plugins_dir / "startup_error_plugin.py").write_text(
|
||||
STARTUP_ERROR_PLUGIN, "utf-8"
|
||||
)
|
||||
|
||||
port = _find_free_port()
|
||||
ds_proc = subprocess.Popen(
|
||||
[
|
||||
sys.executable,
|
||||
"-m",
|
||||
"datasette",
|
||||
"--memory",
|
||||
"--plugins-dir",
|
||||
str(plugins_dir),
|
||||
"-p",
|
||||
str(port),
|
||||
],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
cwd=tempfile.gettempdir(),
|
||||
)
|
||||
try:
|
||||
deadline = time.time() + 15.0
|
||||
# While the process is still alive (it should crash almost
|
||||
# immediately) repeatedly confirm nothing is listening on the port
|
||||
while ds_proc.poll() is None and time.time() < deadline:
|
||||
with pytest.raises(OSError):
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=0.2):
|
||||
pass
|
||||
time.sleep(0.05)
|
||||
|
||||
stdout, _ = ds_proc.communicate(timeout=5)
|
||||
output = stdout.decode("utf-8")
|
||||
assert ds_proc.returncode not in (0, None), output
|
||||
assert "boom from plugin" in output, output
|
||||
|
||||
# And confirm it never accepted a connection even now it has exited
|
||||
with pytest.raises(OSError):
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=0.2):
|
||||
pass
|
||||
finally:
|
||||
if ds_proc.poll() is None:
|
||||
ds_proc.kill()
|
||||
ds_proc.wait()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue