diff options
-rw-r--r-- | alembic/batch.py | 2 | ||||
-rw-r--r-- | alembic/ddl/impl.py | 10 | ||||
-rw-r--r-- | alembic/ddl/postgresql.py | 4 | ||||
-rw-r--r-- | tests/test_batch.py | 16 |
4 files changed, 29 insertions, 3 deletions
diff --git a/alembic/batch.py b/alembic/batch.py index a64d2bc..8470404 100644 --- a/alembic/batch.py +++ b/alembic/batch.py @@ -170,6 +170,8 @@ class ApplyBatchImpl(object): def _create(self, op_impl): self._transfer_elements_to_new_table() + + op_impl.prep_table_for_batch(self.table) op_impl.create_table(self.new_table) op_impl._exec( diff --git a/alembic/ddl/impl.py b/alembic/ddl/impl.py index ca65677..da77052 100644 --- a/alembic/ddl/impl.py +++ b/alembic/ddl/impl.py @@ -74,6 +74,16 @@ class DefaultImpl(with_metaclass(ImplMeta)): """ return False + def prep_table_for_batch(self, table): + """perform any operations needed on a table before a new + one is created to replace it in batch mode. + + the PG dialect uses this to drop constraints on the table + before the new one uses those same names. + + """ + + @property def bind(self): return self.connection diff --git a/alembic/ddl/postgresql.py b/alembic/ddl/postgresql.py index 6193cda..5156cea 100644 --- a/alembic/ddl/postgresql.py +++ b/alembic/ddl/postgresql.py @@ -14,6 +14,10 @@ class PostgresqlImpl(DefaultImpl): __dialect__ = 'postgresql' transactional_ddl = True + def prep_table_for_batch(self, table): + for constraint in table.constraints: + self.drop_constraint(constraint) + def compare_server_default(self, inspector_column, metadata_column, rendered_metadata_default, diff --git a/tests/test_batch.py b/tests/test_batch.py index 2d9bb51..bbbb094 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -362,6 +362,7 @@ class BatchRoundTripTest(TestBase): Column('x', Integer) ) t1.create(self.conn) + self.conn.execute( t1.insert(), [ @@ -409,6 +410,18 @@ class BatchRoundTripTest(TestBase): {"id": 5, "x": 9} ]) + def test_drop_column_fk_recreate(self): + with self.op.batch_alter_table("foo", recreate='always') as batch_op: + batch_op.drop_column('data') + + self._assert_data([ + {"id": 1, "x": 5}, + {"id": 2, "x": 6}, + {"id": 3, "x": 7}, + {"id": 4, "x": 8}, + {"id": 5, "x": 9} + ]) + def test_rename_column(self): with self.op.batch_alter_table("foo") as batch_op: batch_op.alter_column('x', new_column_name='y') @@ -496,6 +509,3 @@ class BatchRoundTripPostgresqlTest(BatchRoundTripTest): def test_change_type(self): super(BatchRoundTripPostgresqlTest, self).test_change_type() - @exclusions.fails() - def test_add_column_recreate(self): - super(BatchRoundTripPostgresqlTest, self).test_add_column_recreate() |