From 618f0d9d05f89302467b71eaf4b67f1bd943aaf3 Mon Sep 17 00:00:00 2001 From: Mike Bayer Date: Wed, 11 Jul 2018 14:05:02 -0400 Subject: [PATCH] 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 --- alembic/operations/batch.py | 7 ++++- alembic/util/sqla_compat.py | 9 ++++++ docs/build/unreleased/502.rst | 8 ++++++ tests/test_batch.py | 53 ++++++++++++++++++++++++++++++++++- 4 files changed, 75 insertions(+), 2 deletions(-) create mode 100644 docs/build/unreleased/502.rst diff --git a/alembic/operations/batch.py b/alembic/operations/batch.py index 84d29d95..5b562d62 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 4620cda3..0bb36065 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 00000000..01555363 --- /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 03be9f98..07a55363 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: @@ -1426,6 +1456,16 @@ class BatchRoundTripTest(TestBase): 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() @@ -1457,6 +1497,17 @@ class BatchRoundTripMySQLTest(BatchRoundTripTest): 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() -- 2.47.2