Merge 1a97c6e23d into c435a80030

docs/plugin_hooks.rst

@@ -1175,20 +1175,19 @@ Here is an example that validates required plugin configuration. The server will
                 "my-plugin requires setting required-setting"
             )
 
-You can also return an async function, which will be awaited on startup. Use this option if you need to execute any database queries, for example this function which creates the ``my_table`` database table if it does not yet exist:
+The hook can be an ``async def`` function. Datasette awaits it before startup completes, so you can execute database queries directly. For example, this hook creates the ``my_table`` database table if it does not yet exist:
 
 .. code-block:: python
 
     @hookimpl
-    def startup(datasette):
-        async def inner():
-            db = datasette.get_database()
-            if "my_table" not in await db.table_names():
-                await db.execute_write("""
-                    create table my_table (mycol text)
-                """)
+    async def startup(datasette):
+        db = datasette.get_database()
+        if "my_table" not in await db.table_names():
+            await db.execute_write("""
+                create table my_table (mycol text)
+            """)
 
-        return inner
+A regular ``def startup(datasette)`` function can also return an async function or a coroutine, which Datasette will await. Exceptions raised by an async startup hook, including ``StartupError``, prevent startup in the same way as exceptions from a synchronous hook.
 
 Potential use-cases:
 

tests/test_plugins.py

@@ -980,6 +980,65 @@ async def test_hook_startup(ds_client):
     assert 2 == ds_client.ds._startup_hook_calculation
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("return_style", ["async_def", "async_function", "coroutine"])
+async def test_hook_startup_async(return_style, tmp_path):
+    db_path = tmp_path / "startup.db"
+    sqlite3.connect(db_path).close()
+    ds = Datasette([db_path])
+
+    async def initialize(datasette):
+        await datasette.get_database().execute_write(
+            "create table startup_test (id integer)"
+        )
+
+    class AsyncStartupPlugin:
+        __name__ = "AsyncStartupPlugin"
+
+        if return_style == "async_def":
+
+            @hookimpl
+            async def startup(self, datasette):
+                await initialize(datasette)
+
+        else:
+
+            @hookimpl
+            def startup(self, datasette):
+                async def inner():
+                    await initialize(datasette)
+
+                return inner if return_style == "async_function" else inner()
+
+    ds.pm.register(AsyncStartupPlugin(), name="async_startup_plugin")
+    try:
+        await ds.invoke_startup()
+        result = await ds.get_database().execute("select count(*) from startup_test")
+        assert result.first()[0] == 0
+    finally:
+        ds.pm.unregister(name="async_startup_plugin")
+
+
+@pytest.mark.asyncio
+async def test_hook_startup_async_error():
+    ds = Datasette(memory=True)
+
+    class FailingStartupPlugin:
+        __name__ = "FailingStartupPlugin"
+
+        @hookimpl
+        async def startup(self, datasette):
+            await datasette.get_database().execute("select 1")
+            raise StartupError("Invalid plugin configuration")
+
+    ds.pm.register(FailingStartupPlugin(), name="failing_startup_plugin")
+    try:
+        with pytest.raises(StartupError, match="Invalid plugin configuration"):
+            await ds.invoke_startup()
+    finally:
+        ds.pm.unregister(name="failing_startup_plugin")
+
+
 @pytest.mark.asyncio
 async def test_hook_startup_metadata_available(ds_client):
     # Metadata from metadata.yaml should be populated before startup() fires