Allow write wrappers to roll back alter-table operations

Merge pull request #2925

datasette/views/table_create_alter.py

@@ -1309,7 +1309,9 @@ class TableAlterView(BaseView):
                     elif operation.op == "set_foreign_keys":
                         foreign_keys = [fk.tuple for fk in args.foreign_keys]
 
-                with operation_conn:
+                # Use a savepoint inside execute_write_fn's transaction so
+                # write_wrapper hooks can still reject and roll back the write.
+                with db_for_write.atomic():
                     for column in add_columns:
                         not_null_default = None
                         if column.not_null:

docs/changelog.rst

@@ -26,6 +26,7 @@ Datasette plugins can now use **background tasks** to run code independent of th
 Bug fixes
 ~~~~~~~~~
 
+- The :ref:`alter-table API <TableAlterView>` now rolls back schema changes when a :ref:`write_wrapper <plugin_hook_write_wrapper>` raises after the write. (:issue:`2924`, :pr:`2925`)
 - The :ref:`extra_template_vars() <plugin_hook_extra_template_vars>` plugin hook can now return a function or awaitable that resolves to ``None`` when no extra variables are needed. (:issue:`2005`)
 - :ref:`request.headers <internals_request>` now supports case-insensitive header lookups, so ``request.headers.get("Content-Type")`` works as well as ``request.headers.get("content-type")``. (:issue:`1861`)
 - CSV endpoints now return plain-text error messages for SQL errors. (:issue:`2129`)

docs/internals.rst

@@ -2292,7 +2292,25 @@ The value returned from ``await database.execute_write_fn(...)`` will be the ret
 
 If your function raises an exception that exception will be propagated up to the ``await`` line.
 
-By default your function will be executed inside a transaction. You can pass ``transaction=False`` to disable this behavior, though if you do that you should be careful to manually apply transactions - ideally using the ``with conn:`` pattern, or you may see ``OperationalError: database table is locked`` errors.
+By default Datasette manages the transaction. For nested transactions, use `sqlite_utils.Database(conn).atomic() <https://sqlite-utils.datasette.io/en/stable/python-api.html#grouping-changes-with-db-atomic>`__. Pass ``transaction=False`` to manage transactions yourself.
+
+For example, archive an article and record the change in an audit log:
+
+.. code-block:: python
+
+    import sqlite_utils
+
+
+    def archive_article(conn):
+        db = sqlite_utils.Database(conn)
+        with db.atomic():
+            db["articles"].update(1, {"archived": True})
+            db["audit_log"].insert(
+                {"article_id": 1, "action": "archive"}
+            )
+
+
+    await database.execute_write_fn(archive_article)
 
 If you specify ``block=False`` the method becomes fire-and-forget, queueing your function to be executed and then allowing your code after the call to ``.execute_write_fn()`` to continue running while the underlying thread waits for an opportunity to run your function. A UUID representing the queued task will be returned. Any exceptions in your code will be silently swallowed.
 

tests/test_write_wrapper.py

@@ -351,6 +351,73 @@ async def test_write_wrapper_via_api(tmp_path):
         pm.unregister(name="test_api")
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("num_sql_threads", (0, 1))
+@pytest.mark.parametrize("in_memory", (True, False))
+@pytest.mark.parametrize(
+    "operations",
+    (
+        [{"op": "add_column", "args": {"name": "extra", "type": "text"}}],
+        [{"op": "rename_column", "args": {"name": "id", "to": "renamed_id"}}],
+        [{"op": "rename_table", "args": {"to": "renamed_t"}}],
+        [
+            {"op": "add_column", "args": {"name": "extra", "type": "text"}},
+            {"op": "rename_column", "args": {"name": "id", "to": "renamed_id"}},
+            {"op": "rename_table", "args": {"to": "renamed_t"}},
+        ],
+    ),
+    ids=["add-column", "transform", "rename-table", "combined"],
+)
+async def test_write_wrapper_can_reject_alter_table_after_write(
+    tmp_path, num_sql_threads, in_memory, operations
+):
+    """Raising after yield should roll back the schema change."""
+    db_path = str(tmp_path / "demo.db")
+    ds = Datasette(
+        [] if in_memory else [db_path],
+        config={"permissions": {"alter-table": True}},
+        settings={"num_sql_threads": num_sql_threads},
+    )
+    db = (
+        ds.add_memory_database(db_path, name="demo")
+        if in_memory
+        else ds.get_database("demo")
+    )
+    await db.execute_write("CREATE TABLE t (id)")
+    await db.execute_write("INSERT INTO t (id) VALUES (1)")
+    before = await db.execute_fn(lambda conn: list(conn.iterdump()))
+
+    class Plugin:
+        __name__ = "Plugin"
+
+        @staticmethod
+        @hookimpl
+        def write_wrapper(database):
+            def wrapper(conn):
+                yield
+                raise ValueError("Rejected after write")
+
+            return wrapper if database == "demo" else None
+
+    pm.register(Plugin(), name="test_reject_alter_table")
+    try:
+        response = await ds.client.post(
+            "/demo/t/-/alter",
+            json={"operations": operations},
+        )
+        assert response.status_code == 400
+        assert response.json()["errors"] == ["Rejected after write"]
+        assert await db.execute_fn(lambda conn: list(conn.iterdump())) == before
+        assert not [
+            event
+            for event in getattr(ds, "_tracked_events", [])
+            if event.name in ("alter-table", "rename-table")
+        ]
+    finally:
+        pm.unregister(name="test_reject_alter_table")
+        ds.close()
+
+
 @pytest.mark.asyncio
 async def test_write_wrapper_change_group_pattern(datasette):
     """Test the motivating use case: activating a change group around a write."""