Add "--check-heads" option to "current" command

Added :paramref:`.command.current.check_heads` parameter to
:func:`.command.current` command, available from the command line via the
``--check-heads`` option to ``alembic current``.  This tests if all head
revisions are applied to the database and raises :class:`.DatabaseNotAtHead`
(or from the command line, exits with a non-zero exit code) if this is not
the case.  The parameter operates equvialently to the cookbook recipe
:ref:`cookbook_check_heads`. Pull request courtesy Stefan Scherfke.

Fixes: #1705
Closes: #1739
Pull-request: #1739
Pull-request-sha: 955f9b7

Change-Id: I1173217ff6a4927226ec9e7b8c7ef67c6431d728

Author: sscherfke

alembic/command.py

@@ -681,11 +681,18 @@ def branches(config: Config, verbose: bool = False) -> None:
             )
 
 
-def current(config: Config, verbose: bool = False) -> None:
+def current(
+    config: Config, check_heads: bool = False, verbose: bool = False
+) -> None:
     """Display the current revision for a database.
 
     :param config: a :class:`.Config` instance.
 
+    :param check_heads: Check if all head revisions are applied to the
+        database.  Raises :class:`.DatabaseNotAtHead` if this is not the case.
+
+        .. versionadded:: 1.17.1
+
     :param verbose: output in verbose mode.
 
     """
@@ -698,6 +705,12 @@ def display_version(rev, context):
                 "Current revision(s) for %s:",
                 util.obfuscate_url_pw(context.connection.engine.url),
             )
+        if check_heads and (
+            set(context.get_current_heads()) != set(script.get_heads())
+        ):
+            raise util.DatabaseNotAtHead(
+                "Database is not on all head revisions"
+            )
         for rev in script.get_all_current(rev):
             config.print_stdout(rev.cmd_format(verbose))
 

alembic/config.py

@@ -808,6 +808,17 @@ def __init__(self, prog: Optional[str] = None) -> None:
                 "environment and version locations",
             ),
         ),
+        "check_heads": (
+            "-c",
+            "--check-heads",
+            dict(
+                action="store_true",
+                help=(
+                    "Check if all head revisions are applied to the database. "
+                    "Exit with an error code if this is not the case."
+                ),
+            ),
+        ),
     }
     _POSITIONAL_OPTS = {
         "directory": dict(help="location of scripts directory"),

alembic/util/__init__.py

@@ -1,6 +1,7 @@
 from .editor import open_in_editor as open_in_editor
 from .exc import AutogenerateDiffsDetected as AutogenerateDiffsDetected
 from .exc import CommandError as CommandError
+from .exc import DatabaseNotAtHead as DatabaseNotAtHead
 from .langhelpers import _with_legacy_names as _with_legacy_names
 from .langhelpers import asbool as asbool
 from .langhelpers import dedupe_tuple as dedupe_tuple

alembic/util/exc.py

@@ -10,10 +10,28 @@
 
 
 class CommandError(Exception):
-    pass
+    """Base command error for all exceptions"""
+
+
+class DatabaseNotAtHead(CommandError):
+    """Indicates the database is not at current head revisions.
+
+    Raised by the :func:`.command.current` command when the
+    :paramref:`.command.current.check_heads` parameter is used.
+
+    .. versionadded:: 1.17.1
+
+    """
 
 
 class AutogenerateDiffsDetected(CommandError):
+    """Raised when diffs were detected by the :func:`.command.check`
+    command.
+
+    .. versionadded:: 1.9.0
+
+    """
+
     def __init__(
         self,
         message: str,

docs/build/api/exceptions.rst

@@ -0,0 +1,9 @@
+.. _alembic.runtime.exceptions.toplevel:
+
+=======================
+Exception Objects
+=======================
+
+
+.. automodule:: alembic.util.exc
+    :members:

docs/build/api/index.rst

@@ -29,4 +29,5 @@ to run commands programmatically, as discussed in the section :doc:`/api/command
    autogenerate
    script
    ddl
+   exceptions
 

docs/build/cookbook.rst

@@ -1480,10 +1480,24 @@ at :func:`.autogenerate.compare_metadata`::
 
 
 
+.. _cookbook_check_heads:
 
 Test current database revision is at head(s)
 ============================================
 
+.. versionchanged:: 1.17.1  This recipe is now part of the ``alembic current``
+   command using the :paramref:`.command.current.check_heads` parameter,
+   available from the command line as ``--check-heads``:
+
+   .. sourcecode::
+
+       alembic current --check-heads
+
+       INFO  [alembic.runtime.migration] Context impl SQLiteImpl.
+       INFO  [alembic.runtime.migration] Will assume non-transactional DDL.
+       ERROR [alembic.util.messaging] Database is not on all head revisions
+       FAILED: Database is not on all head revisions
+
 A recipe to determine if a database schema is up to date in terms of applying
 Alembic migrations.   May be useful for test or installation suites to
 determine if the target database is up to date.   Makes use of the

docs/build/unreleased/1705.rst

@@ -0,0 +1,11 @@
+.. change::
+   :tags: usecase, commands
+   :tickets: 1705
+
+   Added :paramref:`.command.current.check_heads` parameter to
+   :func:`.command.current` command, available from the command line via the
+   ``--check-heads`` option to ``alembic current``.  This tests if all head
+   revisions are applied to the database and raises :class:`.DatabaseNotAtHead`
+   (or from the command line, exits with a non-zero exit code) if this is not
+   the case.  The parameter operates equvialently to the cookbook recipe
+   :ref:`cookbook_check_heads`. Pull request courtesy Stefan Scherfke.

tests/test_command.py

@@ -311,6 +311,8 @@ def _assert_lines(self, revs):
         eq_(lines, set(revs))
 
     def test_doesnt_create_alembic_version(self):
+        with self.bind.begin() as conn:
+            conn.exec_driver_sql("drop table if exists alembic_version")
         command.current(self.cfg)
         engine = self.bind
         with engine.connect() as conn:
@@ -345,6 +347,40 @@ def test_heads_upg(self):
         with self._assert_lines(["a2", "b3"]):
             command.current(self.cfg)
 
+    def test_check_heads_success(self):
+        """
+        "--check-heads" succeeds if all head revisions are applied.
+        """
+        command.stamp(self.cfg, ())
+        command.stamp(self.cfg, (self.a3.revision, self.b3.revision))
+        with self._assert_lines(["a3", "b3"]):
+            command.current(self.cfg, check_heads=True)
+
+    @testing.combinations(
+        ["a2"],
+        ["a3"],
+        ["b3"],
+        ["a2", "b3"],
+        ["a3", "b2"],
+        argnames="revs",
+    )
+    def test_check_heads_fail(self, revs):
+        """
+        "--check-heads" succeeds if all head revisions are applied.
+        """
+        command.stamp(self.cfg, ())
+        command.stamp(
+            self.cfg, tuple(getattr(self, rev).revision for rev in revs)
+        )
+        assert_raises_message(
+            util.DatabaseNotAtHead,
+            "Database is not on all head revisions",
+            command.current,
+            self.cfg,
+            check_heads=True,
+        )
+        command.stamp(self.cfg, ())
+
 
 class RevisionTest(TestBase):
     def setUp(self):
@@ -1562,7 +1598,7 @@ def existing_pyproject_fixture(self):
         yield config.Config(
             self.cfg.config_file_name, toml_file=root / "pyproject.toml"
         )
-        shutil.rmtree(root)
+        os.unlink(root / "pyproject.toml")
 
     @testing.variation("cmd", ["list_templates", "init"])
     def test_init_custom_template_location(self, cmd, custom_template_fixture):