From a8b292edb1de26b199c2d4709bd7a871dead55cb Mon Sep 17 00:00:00 2001
From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Mon, 19 Dec 2016 11:04:43 -0500
Subject: [PATCH] Ensure primary_key flag unset for PK batch drop

Fixed bug where doing ``batch_op.drop_constraint()`` against the
primary key constraint would fail to remove the "primary_key" flag
from the column, resulting in the constraint being recreated.

Change-Id: I20c04860b151ac86466337f0522018be06c6feec
Fixes: #402
---
 alembic/operations/batch.py |  8 ++++++--
 docs/build/changelog.rst    |  8 ++++++++
 tests/test_batch.py         | 21 ++++++++++++++++++++-
 3 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/alembic/operations/batch.py b/alembic/operations/batch.py
index 24c1a816..fa7899e0 100644
--- a/alembic/operations/batch.py
+++ b/alembic/operations/batch.py
@@ -1,5 +1,5 @@
 from sqlalchemy import Table, MetaData, Index, select, Column, \
-    ForeignKeyConstraint, cast, CheckConstraint
+    ForeignKeyConstraint, PrimaryKeyConstraint, cast, CheckConstraint
 from sqlalchemy import types as sqltypes
 from sqlalchemy import schema as sql_schema
 from sqlalchemy.util import OrderedDict
@@ -345,7 +345,7 @@ class ApplyBatchImpl(object):
         if not const.name:
             raise ValueError("Constraint must have a name")
         try:
-            del self.named_constraints[const.name]
+            const = self.named_constraints.pop(const.name)
         except KeyError:
             if _is_type_bound(const):
                 # type-bound constraints are only included in the new
@@ -354,6 +354,10 @@ class ApplyBatchImpl(object):
                 # Operations.implementation_for(alter_column)
                 return
             raise ValueError("No such constraint: '%s'" % const.name)
+        else:
+            if isinstance(const, PrimaryKeyConstraint):
+                for col in const.columns:
+                    self.columns[col.name].primary_key = False
 
     def create_index(self, idx):
         self.new_indexes[idx.name] = idx
diff --git a/docs/build/changelog.rst b/docs/build/changelog.rst
index b2ac0843..0baaa538 100644
--- a/docs/build/changelog.rst
+++ b/docs/build/changelog.rst
@@ -6,6 +6,14 @@ Changelog
 .. changelog::
     :version: 0.8.10
 
+    .. change:: 402
+      :tags: bug, batch
+      :tickets: 402
+
+      Fixed bug where doing ``batch_op.drop_constraint()`` against the
+      primary key constraint would fail to remove the "primary_key" flag
+      from the column, resulting in the constraint being recreated.
+
     .. change:: 356
       :tags: bug, versioning
       :tickets: 356
diff --git a/tests/test_batch.py b/tests/test_batch.py
index 3124baa7..99d79560 100644
--- a/tests/test_batch.py
+++ b/tests/test_batch.py
@@ -13,7 +13,7 @@ from alembic.runtime.migration import MigrationContext
 
 from sqlalchemy import Integer, Table, Column, String, MetaData, ForeignKey, \
     UniqueConstraint, ForeignKeyConstraint, Index, Boolean, CheckConstraint, \
-    Enum, DateTime
+    Enum, DateTime, PrimaryKeyConstraint
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.sql import column, text, select
 from sqlalchemy.schema import CreateTable, CreateIndex
@@ -58,6 +58,17 @@ class BatchApplyTest(TestBase):
         )
         return ApplyBatchImpl(t, table_args, table_kwargs, False)
 
+    def _pk_fixture(self):
+        m = MetaData()
+        t = Table(
+            'tname', m,
+            Column('id', Integer),
+            Column('x', String()),
+            Column('y', Integer),
+            PrimaryKeyConstraint('id', name="mypk")
+        )
+        return ApplyBatchImpl(t, (), {}, False)
+
     def _literal_ck_fixture(
             self, copy_from=None, table_args=(), table_kwargs={}):
         m = MetaData()
@@ -534,6 +545,14 @@ class BatchApplyTest(TestBase):
             dialect='mysql'
         )
 
+    def test_drop_pk(self):
+        impl = self._pk_fixture()
+        pk = self.op.schema_obj.primary_key_constraint("mypk", "tname", ["id"])
+        impl.drop_constraint(pk)
+        new_table = self._assert_impl(impl)
+        assert not new_table.c.id.primary_key
+        assert not len(new_table.primary_key)
+
 
 class BatchAPITest(TestBase):
     __requires__ = ('sqlalchemy_08', )
-- 
2.47.2