Repair as_revision_number to return a tuple for "heads"

Fixed bug where the :meth:`.Script.as_revision_number` method
did not accommodate for the 'heads' identifier, which in turn
caused the :meth:`.EnvironmentContext.get_head_revisions`
and :meth:`.EnvironmentContext.get_revision_argument` methods
to be not usable when multiple heads were present.
The :meth:.`EnvironmentContext.get_head_revisions` method returns
a tuple in all cases as documented.

Change-Id: I085d9b6c3f4ceafd6828d24983768a3d3916ce00
Fixes: #482

Author: zzzeek

alembic/script/base.py

@@ -237,6 +237,8 @@ def as_revision_number(self, id_):
         if not rev:
             # convert () to None
             return None
+        elif id_ == "heads":
+            return rev
         else:
             return rev[0]
 

alembic/testing/env.py

@@ -319,6 +319,74 @@ def downgrade():
     return a, b, c
 
 
+def multi_heads_fixture(cfg, a, b, c):
+    """Create a multiple head fixture from the three-revs fixture"""
+
+    d = util.rev_id()
+    e = util.rev_id()
+    f = util.rev_id()
+
+    script = ScriptDirectory.from_config(cfg)
+    script.generate_revision(
+        d, "revision d from b", head=b, splice=True, refresh=True)
+    write_script(script, d, """\
+"Rev D"
+revision = '%s'
+down_revision = '%s'
+
+from alembic import op
+
+
+def upgrade():
+    op.execute("CREATE STEP 4")
+
+
+def downgrade():
+    op.execute("DROP STEP 4")
+
+""" % (d, b))
+
+    script.generate_revision(
+        e, "revision e from d", head=d, splice=True, refresh=True)
+    write_script(script, e, """\
+"Rev E"
+revision = '%s'
+down_revision = '%s'
+
+from alembic import op
+
+
+def upgrade():
+    op.execute("CREATE STEP 5")
+
+
+def downgrade():
+    op.execute("DROP STEP 5")
+
+""" % (e, d))
+
+    script.generate_revision(
+        f, "revision f from b", head=b, splice=True, refresh=True)
+    write_script(script, f, """\
+"Rev F"
+revision = '%s'
+down_revision = '%s'
+
+from alembic import op
+
+
+def upgrade():
+    op.execute("CREATE STEP 6")
+
+
+def downgrade():
+    op.execute("DROP STEP 6")
+
+""" % (f, b))
+
+    return d, e, f
+
+
 def _multidb_testing_config(engines):
     """alembic.ini fixture to work exactly with the 'multidb' template"""
 

docs/build/unreleased/482.rst

@@ -0,0 +1,13 @@
+.. change::
+    :tags: bug, runtime
+    :tickets: 482
+
+    Fixed bug where the :meth:`.Script.as_revision_number` method
+    did not accommodate for the 'heads' identifier, which in turn
+    caused the :meth:`.EnvironmentContext.get_head_revisions`
+    and :meth:`.EnvironmentContext.get_revision_argument` methods
+    to be not usable when multiple heads were present.
+    The :meth:.`EnvironmentContext.get_head_revisions` method returns
+    a tuple in all cases as documented.
+
+

tests/test_offline_environment.py

@@ -4,7 +4,8 @@
 
 from alembic.testing import assert_raises_message
 from alembic.testing.env import staging_env, _no_sql_testing_config, \
-    three_rev_fixture, clear_staging_env, env_file_fixture
+    three_rev_fixture, clear_staging_env, env_file_fixture, \
+    multi_heads_fixture
 import re
 
 a = b = c = None
@@ -80,6 +81,12 @@ def test_destination_rev_pre_context(self):
         command.stamp(self.cfg, b, sql=True)
         command.downgrade(self.cfg, "%s:%s" % (c, b), sql=True)
 
+    def test_destination_rev_pre_context_multihead(self):
+        d, e, f = multi_heads_fixture(self.cfg, a, b, c)
+        env_file_fixture("""
+assert set(context.get_revision_argument()) == set(('%s', '%s', '%s', ))
+""" % (f, e, c))
+        command.upgrade(self.cfg, 'heads', sql=True)
 
     def test_destination_rev_post_context(self):
         env_file_fixture("""
@@ -90,25 +97,56 @@ def test_destination_rev_post_context(self):
         command.downgrade(self.cfg, "%s:%s" % (c, b), sql=True)
         command.stamp(self.cfg, b, sql=True)
 
+    def test_destination_rev_post_context_multihead(self):
+        d, e, f = multi_heads_fixture(self.cfg, a, b, c)
+        env_file_fixture("""
+context.configure(dialect_name='sqlite')
+assert set(context.get_revision_argument()) == set(('%s', '%s', '%s', ))
+""" % (f, e, c))
+        command.upgrade(self.cfg, 'heads', sql=True)
+
     def test_head_rev_pre_context(self):
         env_file_fixture("""
 assert context.get_head_revision() == '%s'
-""" % c)
+assert context.get_head_revisions() == ('%s', )
+""" % (c, c))
         command.upgrade(self.cfg, b, sql=True)
         command.downgrade(self.cfg, "%s:%s" % (b, a), sql=True)
         command.stamp(self.cfg, b, sql=True)
         command.current(self.cfg)
 
+    def test_head_rev_pre_context_multihead(self):
+        d, e, f = multi_heads_fixture(self.cfg, a, b, c)
+        env_file_fixture("""
+assert set(context.get_head_revisions()) == set(('%s', '%s', '%s', ))
+""" % (e, f, c))
+        command.upgrade(self.cfg, e, sql=True)
+        command.downgrade(self.cfg, "%s:%s" % (e, b), sql=True)
+        command.stamp(self.cfg, c, sql=True)
+        command.current(self.cfg)
+
     def test_head_rev_post_context(self):
         env_file_fixture("""
 context.configure(dialect_name='sqlite')
 assert context.get_head_revision() == '%s'
-""" % c)
+assert context.get_head_revisions() == ('%s', )
+""" % (c, c))
         command.upgrade(self.cfg, b, sql=True)
         command.downgrade(self.cfg, "%s:%s" % (b, a), sql=True)
         command.stamp(self.cfg, b, sql=True)
         command.current(self.cfg)
 
+    def test_head_rev_post_context_multihead(self):
+        d, e, f = multi_heads_fixture(self.cfg, a, b, c)
+        env_file_fixture("""
+context.configure(dialect_name='sqlite')
+assert set(context.get_head_revisions()) == set(('%s', '%s', '%s', ))
+""" % (e, f, c))
+        command.upgrade(self.cfg, e, sql=True)
+        command.downgrade(self.cfg, "%s:%s" % (e, b), sql=True)
+        command.stamp(self.cfg, c, sql=True)
+        command.current(self.cfg)
+
     def test_tag_pre_context(self):
         env_file_fixture("""
 assert context.get_tag_argument() == 'hi'