From 0b68996cc511b3a801f0cd0157bd66332d75f46f Mon Sep 17 00:00:00 2001 From: Simon Willison Date: Thu, 15 Dec 2022 13:06:45 -0800 Subject: [PATCH] Revert "Replace AsgiLifespan with AsgiRunOnFirstRequest, refs #1955" This reverts commit dc18f62089e5672d03176f217d7840cdafa5c447. --- datasette/app.py | 20 +++++++++-- datasette/utils/asgi.py | 44 ++++++++++++++++-------- docs/plugin_hooks.rst | 5 ++- docs/testing_plugins.rst | 2 +- tests/test_internals_datasette_client.py | 1 + 5 files changed, 50 insertions(+), 22 deletions(-) diff --git a/datasette/app.py b/datasette/app.py index 7e68249808..f3cb887657 100644 --- a/datasette/app.py +++ b/datasette/app.py @@ -69,6 +69,8 @@ row_sql_params_pks, ) from .utils.asgi import ( + AsgiLifespan, + Base400, Forbidden, NotFound, DatabaseNotFound, @@ -76,10 +78,11 @@ RowNotFound, Request, Response, - AsgiRunOnFirstRequest, asgi_static, asgi_send, asgi_send_file, + asgi_send_html, + asgi_send_json, asgi_send_redirect, ) from .utils.internal_db import init_internal_db, populate_schema_tables @@ -1417,7 +1420,7 @@ def app(self): async def setup_db(): # First time server starts up, calculate table counts for immutable databases - for database in self.databases.values(): + for dbname, database in self.databases.items(): if not database.is_mutable: await database.table_counts(limit=60 * 60 * 1000) @@ -1431,7 +1434,10 @@ async def setup_db(): ) if self.setting("trace_debug"): asgi = AsgiTracer(asgi) - asgi = AsgiRunOnFirstRequest(asgi, on_startup=[setup_db, self.invoke_startup]) + asgi = AsgiLifespan( + asgi, + on_startup=setup_db, + ) for wrapper in pm.hook.asgi_wrapper(datasette=self): asgi = wrapper(asgi) return asgi @@ -1720,34 +1726,42 @@ def _fix(self, path, avoid_path_rewrites=False): return path async def get(self, path, **kwargs): + await self.ds.invoke_startup() async with httpx.AsyncClient(app=self.app) as client: return await client.get(self._fix(path), **kwargs) async def options(self, path, **kwargs): + await self.ds.invoke_startup() async with httpx.AsyncClient(app=self.app) as client: return await client.options(self._fix(path), **kwargs) async def head(self, path, **kwargs): + await self.ds.invoke_startup() async with httpx.AsyncClient(app=self.app) as client: return await client.head(self._fix(path), **kwargs) async def post(self, path, **kwargs): + await self.ds.invoke_startup() async with httpx.AsyncClient(app=self.app) as client: return await client.post(self._fix(path), **kwargs) async def put(self, path, **kwargs): + await self.ds.invoke_startup() async with httpx.AsyncClient(app=self.app) as client: return await client.put(self._fix(path), **kwargs) async def patch(self, path, **kwargs): + await self.ds.invoke_startup() async with httpx.AsyncClient(app=self.app) as client: return await client.patch(self._fix(path), **kwargs) async def delete(self, path, **kwargs): + await self.ds.invoke_startup() async with httpx.AsyncClient(app=self.app) as client: return await client.delete(self._fix(path), **kwargs) async def request(self, method, path, **kwargs): + await self.ds.invoke_startup() avoid_path_rewrites = kwargs.pop("avoid_path_rewrites", None) async with httpx.AsyncClient(app=self.app) as client: return await client.request( diff --git a/datasette/utils/asgi.py b/datasette/utils/asgi.py index 566902518e..f080df91b5 100644 --- a/datasette/utils/asgi.py +++ b/datasette/utils/asgi.py @@ -156,6 +156,35 @@ def fake(cls, path_with_query_string, method="GET", scheme="http", url_vars=None return cls(scope, None) +class AsgiLifespan: + def __init__(self, app, on_startup=None, on_shutdown=None): + self.app = app + on_startup = on_startup or [] + on_shutdown = on_shutdown or [] + if not isinstance(on_startup or [], list): + on_startup = [on_startup] + if not isinstance(on_shutdown or [], list): + on_shutdown = [on_shutdown] + self.on_startup = on_startup + self.on_shutdown = on_shutdown + + async def __call__(self, scope, receive, send): + if scope["type"] == "lifespan": + while True: + message = await receive() + if message["type"] == "lifespan.startup": + for fn in self.on_startup: + await fn() + await send({"type": "lifespan.startup.complete"}) + elif message["type"] == "lifespan.shutdown": + for fn in self.on_shutdown: + await fn() + await send({"type": "lifespan.shutdown.complete"}) + return + else: + await self.app(scope, receive, send) + + class AsgiStream: def __init__(self, stream_fn, status=200, headers=None, content_type="text/plain"): self.stream_fn = stream_fn @@ -420,18 +449,3 @@ async def asgi_send(self, send): content_type=self.content_type, headers=self.headers, ) - - -class AsgiRunOnFirstRequest: - def __init__(self, asgi, on_startup): - assert isinstance(on_startup, list) - self.asgi = asgi - self.on_startup = on_startup - self._started = False - - async def __call__(self, scope, receive, send): - if not self._started: - self._started = True - for hook in self.on_startup: - await hook() - return await self.asgi(scope, receive, send) diff --git a/docs/plugin_hooks.rst b/docs/plugin_hooks.rst index cdc73f00f4..f41ca876ed 100644 --- a/docs/plugin_hooks.rst +++ b/docs/plugin_hooks.rst @@ -902,14 +902,13 @@ Potential use-cases: .. note:: - If you are writing :ref:`unit tests ` for a plugin that uses this hook and doesn't exercise Datasette by sending - any simulated requests through it you will need to explicitly call ``await ds.invoke_startup()`` in your tests. An example: + If you are writing :ref:`unit tests ` for a plugin that uses this hook you will need to explicitly call ``await ds.invoke_startup()`` in your tests. An example: .. code-block:: python @pytest.mark.asyncio async def test_my_plugin(): - ds = Datasette() + ds = Datasette([], metadata={}) await ds.invoke_startup() # Rest of test goes here diff --git a/docs/testing_plugins.rst b/docs/testing_plugins.rst index 6d2097ad08..41f50e567a 100644 --- a/docs/testing_plugins.rst +++ b/docs/testing_plugins.rst @@ -80,7 +80,7 @@ Creating a ``Datasette()`` instance like this as useful shortcut in tests, but t This method registers any :ref:`plugin_hook_startup` or :ref:`plugin_hook_prepare_jinja2_environment` plugins that might themselves need to make async calls. -If you are using ``await datasette.client.get()`` and similar methods then you don't need to worry about this - Datasette automatically calls ``invoke_startup()`` the first time it handles a request. +If you are using ``await datasette.client.get()`` and similar methods then you don't need to worry about this - those method calls ensure that ``.invoke_startup()`` has been called for you. .. _testing_plugins_pdb: diff --git a/tests/test_internals_datasette_client.py b/tests/test_internals_datasette_client.py index 7a95ed6e1c..cbbfa3c357 100644 --- a/tests/test_internals_datasette_client.py +++ b/tests/test_internals_datasette_client.py @@ -6,6 +6,7 @@ @pytest_asyncio.fixture async def datasette(app_client): + await app_client.ds.invoke_startup() return app_client.ds