alembic/env.py: an upgrade's exception ends it as itself, on one connection

alembic/env.py

@@ -36,6 +36,11 @@
   SQLite ending one by itself (``INSERT OR ROLLBACK``, ``RAISE(ROLLBACK)``). SQLite authorizes a
   statement when it prepares it, so the connection keeps no prepared statements: each execute is
   prepared, and checked, again.
+- An upgrade runs on one driver connection, which holds its transaction and the state the
+  authorizer reads. Once it is open, a connect listener closes and refuses any other before
+  anything runs on it, so a statement after SQLAlchemy invalidated the connection cannot
+  reconnect outside both. SQLAlchemy raises that statement's refusal wrapped in
+  ``StatementError``.
 - Nothing runs on the connection before ``context.configure()``: alembic would treat the
   transaction as the caller's and never commit it.
 - ``legacy_alter_table`` stays off, so batch-altering a table a view reads fails unless the
@@ -231,6 +236,29 @@ def configure_connection(dbapi_connection: sqlite3.Connection, record: Connectio
         raise
 
 
+ONE_CONNECTION = (
+    "the upgrade runs on one driver connection: no other is opened, after SQLAlchemy "
+    "invalidated it or for a second connection to the same engine"
+)
+
+
+def refuse_another_connection(
+    refused: list[str],
+) -> Callable[[sqlite3.Connection, ConnectionPoolEntry], None]:
+    """The connect listener for once the upgrade's connection is open: any other is closed and
+    refused before anything runs on it, and the refusal is recorded in ``refused``, so a revision
+    that catches it still fails the upgrade before the commit. Without it, a statement after
+    SQLAlchemy invalidated the upgrade's connection and a ``rollback()`` reconnects, outside the
+    transaction and the authorizer's state."""
+
+    def refuse(dbapi_connection: sqlite3.Connection, _: ConnectionPoolEntry) -> None:
+        dbapi_connection.close()
+        refused.append(ONE_CONNECTION)
+        raise MigrationRefused(ONE_CONNECTION)
+
+    return refuse
+
+
 def begin_immediate(connection: Connection) -> None:
     """Open the upgrade's transaction, taking the write lock at once."""
     connection.exec_driver_sql("BEGIN IMMEDIATE")
@@ -265,18 +293,30 @@ def run_migrations_online() -> None:
         event.listen(engine, "connect", configure_connection, insert=True)
         event.listen(engine, "begin", begin_immediate)
         with engine.connect() as connection:
+            refused: list[str] = []
+            event.listen(engine, "connect", refuse_another_connection(refused), insert=True)
             context.configure(
                 connection=connection,
                 target_metadata=Base.metadata,
                 transactional_ddl=True,
                 render_as_batch=True,  # autogenerate: SQLite alters tables through batch mode
             )
+            # The mapping the authorizer reads, held before the upgrade runs: SQLAlchemy
+            # invalidates the connection when an interrupt, or an error it reads as a disconnect,
+            # ends a statement, and Connection.info then raises in place of that exception. No
+            # other connection is opened, so this is the live connection's mapping.
+            state = connection.info
             with context.begin_transaction():
                 try:
                     context.run_migrations()
+                    if refused:
+                        raise MigrationRefused(
+                            f"{refused[0]}; a revision went on after the refusal, so nothing was "
+                            "committed"
+                        )
                     check_foreign_keys(connection)
                 except BaseException:
-                    connection.info[TRANSACTION_STATE] = "failed"
+                    state[TRANSACTION_STATE] = "failed"
                     raise
     finally:
         engine.dispose()

docs/database.md

@@ -67,6 +67,11 @@ converted to UTC, and returns aware UTC values. Defaults are `utc_now()`, never
   `user_version` and the pragmas reflection reads with an argument (`table_info`, `table_xinfo`, `index_list`, `index_info`, `foreign_key_list`),
   so `journal_mode`, `legacy_alter_table`, and
   `ignore_check_constraints` keep the values the connection opened with.
+- An upgrade runs on one driver connection. Once it is open, env.py refuses and closes any
+  other the engine would open: a second connection a revision asks for, or the one SQLAlchemy
+  would open after it invalidated the first for an interrupt or an error it reads as a
+  disconnect. The upgrade then fails, even when a revision catches the refusal and goes on, and
+  nothing is committed. An interrupt, a `SystemExit`, or that error ends the upgrade as itself.
 - Refused: offline `--sql` mode, databases other than SQLite, SQLite URI filenames
   (`?uri=true`), a caller-supplied connection or `sqlalchemy.url`, and the tracked
   `github_activity.db`. The connect listener asks SQLite which file it opened, so the tracked

docs/redesign-plan.md

@@ -241,6 +241,7 @@ Every gap found in the audit, and where it is resolved.
 | 108 | The data dictionary presented head-repository, milestone, and team names as GitHub's current values, though they are as of `hydrated_at` | Step 3b-2a-3: described as of `hydrated_at` (D38) |
 | 109 | Edits that may leave `updatedAt` unchanged (body, base ref, review dismissal, issue label events) were accepted as missed | Step 5 measures at least 20 cases of each that no later change the index sees follows; a kind not shown to change `updatedAt` gets an index-side check or a periodic re-hydration in pull request 5b, before Step 6a (section 8, D38) |
 | 110 | A GitHub-derived column a later migration adds to a table without `hydration_version` (accounts, compares), or a new table of such data, is never filled for existing rows | 5.1: such a migration adds a per-row read version, which the refresh or the compare reader reads, or, before Step 6b only, refuses while the file holds the rows it would leave unread (D38) |
+| 111 | An interrupt, a `SystemExit`, or an error SQLAlchemy reads as a disconnect, inside a statement of an upgrade, surfaced as `PendingRollbackError` in its place; after SQLAlchemy invalidated the upgrade's connection, a revision's `rollback()` let the next statement reconnect outside the upgrade's transaction and the state the authorizer reads | `alembic/env.py` records a failure on the connection's mapping it holds before the upgrade, so the exception ends the upgrade as itself; the upgrade runs on one driver connection, refusing and closing any other, and fails before the commit when a revision caught that refusal (5.1; `tests/db/test_migrations.py` `TestOneTransaction`) |
 
 ## 5. Design
 

tests/db/test_migrations.py

@@ -14,15 +14,18 @@
 import sqlite3
 import subprocess
 import sys
-from collections.abc import Iterator
-from contextlib import closing
+from collections.abc import Callable, Iterator
+from contextlib import closing, contextmanager
 from pathlib import Path
 from typing import Any
 
 import pytest
 from alembic import command
 from alembic.script import ScriptDirectory
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, event
+from sqlalchemy.engine import Engine
+from sqlalchemy.exc import ProgrammingError, StatementError
+from sqlalchemy.pool import Pool
 
 from github_activity_db.config import get_settings
 from github_activity_db.db.models import Base
@@ -43,6 +46,9 @@
 
 TRACKED = REPO_ROOT / "github_activity.db"
 TEST_REVISION = "f1e57000a001"
+AFTER_THE_FIRST_SCHEMA = MIGRATION_REVISIONS[MIGRATION_REVISIONS.index(FIRST_SCHEMA_HEAD) + 1 :]
+# A test-only revision's first lines: close its driver connection, as a dropped file would.
+CLOSED_CONNECTION = "bind = op.get_bind()\nbind.connection.dbapi_connection.close()\n"
 # A pr_labels row whose pull request does not exist.
 PLANTED_LABEL = "INSERT INTO pr_labels (pr_id, name, label_node_id) VALUES (999, {!r}, 'x')"
 ORPHAN = f"op.execute({PLANTED_LABEL.format('orphan')!r})"
@@ -58,6 +64,27 @@
 )
 
 
+@contextmanager
+def listening(
+    name: str, listener: Callable[..., None], target: type[Engine | Pool] = Engine
+) -> Iterator[None]:
+    """``listener`` on the ``name`` event of every engine, or every pool, env.py's own included."""
+    event.listen(target, name, listener)
+    try:
+        yield
+    finally:
+        event.remove(target, name, listener)
+
+
+def is_closed(connection: sqlite3.Connection) -> bool:
+    """Whether the driver connection is closed: a closed one runs nothing."""
+    try:
+        connection.execute("SELECT 1")
+    except sqlite3.ProgrammingError:
+        return True
+    return False
+
+
 def snapshot(database: Path) -> tuple[Any, ...]:
     """Everything an upgrade could change: the schema, every table's rows, and user_version."""
     with closing(sqlite3.connect(database)) as connection:
@@ -517,6 +544,183 @@ def table_definitions(database: Path) -> dict[str, str]:
 
 
 class TestOneTransaction:
+    @pytest.mark.parametrize(
+        "raised",
+        [
+            pytest.param(KeyboardInterrupt(), id="KeyboardInterrupt"),
+            pytest.param(SystemExit(3), id="SystemExit"),
+        ],
+    )
+    def test_an_interrupt_after_the_second_revision_applies_none(
+        self, first_schema: Path, raised: BaseException
+    ) -> None:
+        """An interrupt, or a SystemExit, at a set point inside a statement, once the second of
+        the revisions after the first schema has run: that exception itself ends the upgrade,
+        and the file stays at the revision the upgrade started from."""
+        second = AFTER_THE_FIRST_SCHEMA[1]
+        invalidated: list[BaseException | None] = []
+
+        def interrupt(_connection: Any, _cursor: Any, statement: str, *_rest: Any) -> None:
+            if statement.startswith(f"UPDATE alembic_version SET version_num='{second}'"):
+                raise raised
+
+        def record(_dbapi_connection: Any, _record: Any, exception: BaseException | None) -> None:
+            invalidated.append(exception)
+
+        before = snapshot(first_schema)
+        # Caught here, so a test that fails another way still does not end the run.
+        with (
+            listening("after_cursor_execute", interrupt),
+            listening("invalidate", record, Pool),
+            pytest.raises(type(raised)) as error,
+        ):
+            upgrade(first_schema)
+        assert error.value is raised
+        # Inside a statement, so SQLAlchemy invalidated the connection for it, once.
+        assert invalidated == [raised]
+        assert snapshot(first_schema) == before
+        assert revision_of(first_schema) == FIRST_SCHEMA_HEAD
+
+    def test_an_error_read_as_a_disconnect_applies_none(
+        self, migrated_database: Path, tmp_path: Path
+    ) -> None:
+        """A test-only revision closes its driver connection and runs a statement: SQLAlchemy
+        reads the driver's error as a disconnect and invalidates the connection, once, and that
+        error ends the upgrade."""
+        versions = tmp_path / "versions"
+        write_revision(
+            versions, TEST_REVISION, HEAD_REVISION, f"{CLOSED_CONNECTION}op.execute('SELECT 1')"
+        )
+        invalidated: list[BaseException | None] = []
+
+        def record(_dbapi_connection: Any, _record: Any, exception: BaseException | None) -> None:
+            invalidated.append(exception)
+
+        before = snapshot(migrated_database)
+        with listening("invalidate", record, Pool), pytest.raises(ProgrammingError) as error:
+            upgrade(migrated_database, test_versions=versions)
+        assert type(error.value) is ProgrammingError
+        assert type(error.value.orig) is sqlite3.ProgrammingError
+        assert error.value.statement == "SELECT 1"
+        assert error.value.connection_invalidated is True
+        assert invalidated == [error.value.orig]
+        assert snapshot(migrated_database) == before
+        assert revision_of(migrated_database) == HEAD_REVISION
+
+    @pytest.mark.parametrize(
+        "after",
+        [
+            pytest.param("", id="the revision goes on"),
+            pytest.param("raise RuntimeError('after the reconnect')", id="the revision raises"),
+        ],
+    )
+    def test_no_statement_reconnects_after_the_connection_is_invalidated(
+        self, migrated_database: Path, tmp_path: Path, after: str
+    ) -> None:
+        """After SQLAlchemy invalidated the upgrade's connection, a rollback() would let the next
+        statement reconnect on a second connection, outside the upgrade's transaction and the
+        state its authorizer reads. env.py refuses that connection and closes it: the statement
+        fails with the refusal, which SQLAlchemy wraps as it wraps any error raised while it
+        reconnects, and nothing is applied."""
+        versions = tmp_path / "versions"
+        write_revision(
+            versions,
+            TEST_REVISION,
+            HEAD_REVISION,
+            f"{CLOSED_CONNECTION}"
+            "try:\n    op.execute('SELECT 1')\nexcept Exception:\n    pass\n"
+            "bind.rollback()\n"
+            f"op.execute('CREATE TABLE late (x INTEGER)')\n{after}",
+        )
+        connections: list[sqlite3.Connection] = []
+        invalidated: list[BaseException | None] = []
+
+        def connected(dbapi_connection: sqlite3.Connection, _record: Any) -> None:
+            connections.append(dbapi_connection)
+
+        def record(_dbapi_connection: Any, _record: Any, exception: BaseException | None) -> None:
+            invalidated.append(exception)
+
+        before = snapshot(migrated_database)
+        with (
+            listening("connect", connected, Pool),
+            listening("invalidate", record, Pool),
+            pytest.raises(StatementError) as error,
+        ):
+            upgrade(migrated_database, test_versions=versions)
+        assert type(error.value) is StatementError
+        assert error.value.statement == "CREATE TABLE late (x INTEGER)"
+        assert error.value.orig is error.value.__cause__
+        assert type(error.value.orig).__name__ == "MigrationRefused"
+        assert "runs on one driver connection" in str(error.value.orig)
+        assert [type(exception) for exception in invalidated] == [sqlite3.ProgrammingError]
+        # The upgrade's connection, closed by the revision, and the refused one, closed by env.py.
+        assert [is_closed(connection) for connection in connections] == [True, True]
+        assert snapshot(migrated_database) == before
+        assert revision_of(migrated_database) == HEAD_REVISION
+
+    def test_an_upgrade_runs_on_one_driver_connection(self, first_schema: Path) -> None:
+        """The chain from the first schema to head opens one driver connection, and closes it."""
+        connections: list[sqlite3.Connection] = []
+
+        def connected(dbapi_connection: sqlite3.Connection, _record: Any) -> None:
+            connections.append(dbapi_connection)
+
+        with listening("connect", connected, Pool):
+            upgrade(first_schema)
+        assert [is_closed(connection) for connection in connections] == [True]
+        assert revision_of(first_schema) == HEAD_REVISION
+
+    def test_a_revision_gets_no_second_connection(
+        self, migrated_database: Path, tmp_path: Path
+    ) -> None:
+        """A second connection a revision asks the engine for would hold no part of the
+        upgrade's transaction: env.py refuses it before anything runs on it, and closes it."""
+        versions = tmp_path / "versions"
+        write_revision(versions, TEST_REVISION, HEAD_REVISION, "op.get_bind().engine.connect()")
+        connections: list[sqlite3.Connection] = []
+
+        def connected(dbapi_connection: sqlite3.Connection, _record: Any) -> None:
+            connections.append(dbapi_connection)
+
+        before = snapshot(migrated_database)
+        with listening("connect", connected, Pool), pytest.raises(RuntimeError) as error:
+            upgrade(migrated_database, test_versions=versions)
+        assert "runs on one driver connection" in refusal(error)
+        assert [is_closed(connection) for connection in connections] == [True, True]
+        assert snapshot(migrated_database) == before
+        assert revision_of(migrated_database) == HEAD_REVISION
+
+    def test_a_revision_that_catches_the_refusal_still_fails_the_upgrade(
+        self, migrated_database: Path, tmp_path: Path
+    ) -> None:
+        """A revision that catches the refusal of a second connection and goes on: the upgrade
+        fails before the commit, and the file keeps its revision and contents."""
+        versions = tmp_path / "versions"
+        write_revision(
+            versions,
+            TEST_REVISION,
+            HEAD_REVISION,
+            """
+            try:
+                op.get_bind().engine.connect()
+            except RuntimeError:
+                pass
+            op.execute("CREATE TABLE after_the_refusal (x INTEGER)")
+            """,
+        )
+        before = snapshot(migrated_database)
+        with pytest.raises(RuntimeError) as error:
+            upgrade(migrated_database, test_versions=versions)
+        assert type(error.value).__name__ == "MigrationRefused"
+        assert str(error.value) == (
+            "the upgrade runs on one driver connection: no other is opened, after SQLAlchemy "
+            "invalidated it or for a second connection to the same engine; a revision went on "
+            "after the refusal, so nothing was committed"
+        )
+        assert snapshot(migrated_database) == before
+        assert revision_of(migrated_database) == HEAD_REVISION
+
     def test_a_failure_in_the_second_revision_applies_neither(
         self, first_schema: Path, tmp_path: Path
     ) -> None: