Remove column from primary key when dropping

Fixed issue in batch where dropping a primary key column, then adding it
back under the same name but without the primary_key flag, would not remove
it from the existing PrimaryKeyConstraint.  If a new PrimaryKeyConstraint
is added, it is used as-is, as was the case before.

Change-Id: Id79c793fbde1a17393adeb75c2da39f191e676e6
Fixes: #502

4 files changed, 75 insertions, 2 deletions

diff --git a/alembic/operations/batch.py b/alembic/operations/batch.py
index 84d29d9..5b562d6 100644
--- a/alembic/operations/batch.py
+++ b/alembic/operations/batch.py
@@ -7,7 +7,7 @@ from .. import util
 if util.sqla_08:
     from sqlalchemy.events import SchemaEventTarget
 from ..util.sqla_compat import _columns_for_constraint, \
-    _is_type_bound, _fk_is_self_referential
+    _is_type_bound, _fk_is_self_referential, _remove_column_from_collection
 
 
 class BatchOperationsImpl(object):
@@ -334,6 +334,11 @@ class ApplyBatchImpl(object):
         self.column_transfers[column.name] = {}
 
     def drop_column(self, table_name, column, **kw):
+        if column.name in self.table.primary_key.columns:
+            _remove_column_from_collection(
+                self.table.primary_key.columns,
+                column
+            )
         del self.columns[column.name]
         del self.column_transfers[column.name]
 
diff --git a/alembic/util/sqla_compat.py b/alembic/util/sqla_compat.py
index 4620cda..0bb3606 100644
--- a/alembic/util/sqla_compat.py
+++ b/alembic/util/sqla_compat.py
@@ -114,6 +114,15 @@ def _find_columns(clause):
     return cols
 
 
+def _remove_column_from_collection(collection, column):
+    """remove a column from a ColumnCollection."""
+
+    # workaround for older SQLAlchemy, remove the
+    # same object that's present
+    to_remove = collection[column.key]
+    collection.remove(to_remove)
+
+
 def _textual_index_column(table, text_):
     """a workaround for the Index construct's severe lack of flexibility"""
     if isinstance(text_, compat.string_types):
diff --git a/docs/build/unreleased/502.rst b/docs/build/unreleased/502.rst
new file mode 100644
index 0000000..0155536
--- /dev/null
+++ b/docs/build/unreleased/502.rst
@@ -0,0 +1,8 @@
+.. change::
+    :tags: bug, batch
+    :tickets: 502
+
+    Fixed issue in batch where dropping a primary key column, then adding it
+    back under the same name but without the primary_key flag, would not remove
+    it from the existing PrimaryKeyConstraint.  If a new PrimaryKeyConstraint
+    is added, it is used as-is, as was the case before.
diff --git a/tests/test_batch.py b/tests/test_batch.py
index 03be9f9..07a5536 100644
--- a/tests/test_batch.py
+++ b/tests/test_batch.py
@@ -918,7 +918,7 @@ class CopyFromTest(TestBase):
 
 
 class BatchRoundTripTest(TestBase):
-    __requires__ = ('sqlalchemy_08', )
+    __requires__ = ('sqlalchemy_09', )
     __only_on__ = "sqlite"
 
     def setUp(self):
@@ -1207,6 +1207,36 @@ class BatchRoundTripTest(TestBase):
             {"id": 5, "x": 9}
         ])
 
+    def test_drop_pk_col_readd_col(self):
+        # drop a column, add it back without primary_key=True, should no
+        # longer be in the constraint
+        with self.op.batch_alter_table("foo") as batch_op:
+            batch_op.drop_column('id')
+            batch_op.add_column(Column('id', Integer))
+
+        pk_const = Inspector.from_engine(self.conn).get_pk_constraint('foo')
+        eq_(pk_const['constrained_columns'], [])
+
+    def test_drop_pk_col_readd_pk_col(self):
+        # drop a column, add it back with primary_key=True, should remain
+        with self.op.batch_alter_table("foo") as batch_op:
+            batch_op.drop_column('id')
+            batch_op.add_column(Column('id', Integer, primary_key=True))
+
+        pk_const = Inspector.from_engine(self.conn).get_pk_constraint('foo')
+        eq_(pk_const['constrained_columns'], ['id'])
+
+    def test_drop_pk_col_readd_col_also_pk_const(self):
+        # drop a column, add it back without primary_key=True, but then
+        # also make anew PK constraint that includes it, should remain
+        with self.op.batch_alter_table("foo") as batch_op:
+            batch_op.drop_column('id')
+            batch_op.add_column(Column('id', Integer))
+            batch_op.create_primary_key('newpk', ['id'])
+
+        pk_const = Inspector.from_engine(self.conn).get_pk_constraint('foo')
+        eq_(pk_const['constrained_columns'], ['id'])
+
     def test_add_pk_constraint(self):
         self._no_pk_fixture()
         with self.op.batch_alter_table("nopk", recreate="always") as batch_op:
@@ -1427,6 +1457,16 @@ class BatchRoundTripMySQLTest(BatchRoundTripTest):
     __only_on__ = "mysql"
 
     @exclusions.fails()
+    def test_drop_pk_col_readd_pk_col(self):
+        super(BatchRoundTripMySQLTest, self).test_drop_pk_col_readd_pk_col()
+
+    @exclusions.fails()
+    def test_drop_pk_col_readd_col_also_pk_const(self):
+        super(
+            BatchRoundTripMySQLTest, self
+        ).test_drop_pk_col_readd_col_also_pk_const()
+
+    @exclusions.fails()
     def test_rename_column_pk(self):
         super(BatchRoundTripMySQLTest, self).test_rename_column_pk()
 
@@ -1458,6 +1498,17 @@ class BatchRoundTripPostgresqlTest(BatchRoundTripTest):
     __only_on__ = "postgresql"
 
     @exclusions.fails()
+    def test_drop_pk_col_readd_pk_col(self):
+        super(
+            BatchRoundTripPostgresqlTest, self).test_drop_pk_col_readd_pk_col()
+
+    @exclusions.fails()
+    def test_drop_pk_col_readd_col_also_pk_const(self):
+        super(
+            BatchRoundTripPostgresqlTest, self
+        ).test_drop_pk_col_readd_col_also_pk_const()
+
+    @exclusions.fails()
     def test_change_type(self):
         super(BatchRoundTripPostgresqlTest, self).test_change_type()