diff options
author | Mike Bayer <mike_mp@zzzcomputing.com> | 2015-07-03 17:29:17 -0400 |
---|---|---|
committer | Mike Bayer <mike_mp@zzzcomputing.com> | 2015-07-03 17:37:44 -0400 |
commit | a294f8cc3f2e5fc2cad048bc4ce27c57554e2688 (patch) | |
tree | 335359e7a973aea5289b7597d27d1f87cf1dc822 | |
parent | ad5390c0e344008014bcbc8edfe1050ce465ede2 (diff) | |
download | alembic-a294f8cc3f2e5fc2cad048bc4ce27c57554e2688.tar.gz |
- Implemented support for :meth:`.BatchOperations.create_primary_key`
and :meth:`.BatchOperations.create_check_constraint`.
fixes #305
- table keyword arguments are copied from the original reflected table,
such as the "mysql_engine" keyword argument.
-rw-r--r-- | alembic/ddl/postgresql.py | 3 | ||||
-rw-r--r-- | alembic/operations/batch.py | 10 | ||||
-rw-r--r-- | alembic/operations/ops.py | 11 | ||||
-rw-r--r-- | alembic/operations/schemaobj.py | 12 | ||||
-rw-r--r-- | alembic/testing/requirements.py | 11 | ||||
-rw-r--r-- | docs/build/changelog.rst | 9 | ||||
-rw-r--r-- | tests/requirements.py | 11 | ||||
-rw-r--r-- | tests/test_batch.py | 93 |
8 files changed, 150 insertions, 10 deletions
diff --git a/alembic/ddl/postgresql.py b/alembic/ddl/postgresql.py index ea423d7..5109b11 100644 --- a/alembic/ddl/postgresql.py +++ b/alembic/ddl/postgresql.py @@ -23,7 +23,8 @@ class PostgresqlImpl(DefaultImpl): def prep_table_for_batch(self, table): for constraint in table.constraints: - self.drop_constraint(constraint) + if constraint.name is not None: + self.drop_constraint(constraint) def compare_server_default(self, inspector_column, metadata_column, diff --git a/alembic/operations/batch.py b/alembic/operations/batch.py index 726df78..7135e37 100644 --- a/alembic/operations/batch.py +++ b/alembic/operations/batch.py @@ -23,7 +23,7 @@ class BatchOperationsImpl(object): self.recreate = recreate self.copy_from = copy_from self.table_args = table_args - self.table_kwargs = table_kwargs + self.table_kwargs = dict(table_kwargs) self.reflect_args = reflect_args self.reflect_kwargs = reflect_kwargs self.naming_convention = naming_convention @@ -139,11 +139,15 @@ class ApplyBatchImpl(object): for idx in self.table.indexes: self.indexes[idx.name] = idx + for k in self.table.kwargs: + self.table_kwargs.setdefault(k, self.table.kwargs[k]) + def _transfer_elements_to_new_table(self): assert self.new_table is None, "Can only create new table once" m = MetaData() schema = self.table.schema + self.new_table = new_table = Table( '_alembic_batch_temp', m, *(list(self.columns.values()) + list(self.table_args)), @@ -264,6 +268,10 @@ class ApplyBatchImpl(object): def add_constraint(self, const): if not const.name: raise ValueError("Constraint must have a name") + if isinstance(const, sql_schema.PrimaryKeyConstraint): + if self.table.primary_key in self.unnamed_constraints: + self.unnamed_constraints.remove(self.table.primary_key) + self.named_constraints[const.name] = const def drop_constraint(self, const): diff --git a/alembic/operations/ops.py b/alembic/operations/ops.py index 82fdd90..16cccb6 100644 --- a/alembic/operations/ops.py +++ b/alembic/operations/ops.py @@ -209,7 +209,11 @@ class CreatePrimaryKeyOp(AddConstraintOp): :meth:`.Operations.create_primary_key` """ - raise NotImplementedError("not yet implemented") + op = cls( + constraint_name, operations.impl.table_name, columns, + schema=operations.impl.schema + ) + return operations.invoke(op) @Operations.register_operation("create_unique_constraint") @@ -590,7 +594,10 @@ class CreateCheckConstraintOp(AddConstraintOp): :meth:`.Operations.create_check_constraint` """ - raise NotImplementedError("not yet implemented") + op = cls( + constraint_name, operations.impl.table_name, + condition, schema=operations.impl.schema, **kw) + return operations.invoke(op) @Operations.register_operation("create_index") diff --git a/alembic/operations/schemaobj.py b/alembic/operations/schemaobj.py index b590aca..f0f8105 100644 --- a/alembic/operations/schemaobj.py +++ b/alembic/operations/schemaobj.py @@ -12,11 +12,13 @@ class SchemaObjects(object): def primary_key_constraint(self, name, table_name, cols, schema=None): m = self.metadata() columns = [sa_schema.Column(n, NULLTYPE) for n in cols] - t1 = sa_schema.Table(table_name, m, - *columns, - schema=schema) - p = sa_schema.PrimaryKeyConstraint(*columns, name=name) - t1.append_constraint(p) + t = sa_schema.Table( + table_name, m, + *columns, + schema=schema) + p = sa_schema.PrimaryKeyConstraint( + *[t.c[n] for n in cols], name=name) + t.append_constraint(p) return p def foreign_key_constraint( diff --git a/alembic/testing/requirements.py b/alembic/testing/requirements.py index b981951..2889ea5 100644 --- a/alembic/testing/requirements.py +++ b/alembic/testing/requirements.py @@ -32,6 +32,17 @@ class SuiteRequirements(Requirements): ) @property + def check_constraints_w_enforcement(self): + """Target database must support check constraints + and also enforce them.""" + + return exclusions.open() + + @property + def reflects_pk_names(self): + return exclusions.closed() + + @property def fail_before_sqla_079(self): return exclusions.fails_if( lambda config: not util.sqla_079, diff --git a/docs/build/changelog.rst b/docs/build/changelog.rst index 8fd6293..8232c47 100644 --- a/docs/build/changelog.rst +++ b/docs/build/changelog.rst @@ -45,6 +45,15 @@ Changelog :version: 0.7.7 .. change:: + :tags: feature, batch + :tickets: 305 + + Implemented support for :meth:`.BatchOperations.create_primary_key` + and :meth:`.BatchOperations.create_check_constraint`. Additionally, + table keyword arguments are copied from the original reflected table, + such as the "mysql_engine" keyword argument. + + .. change:: :tags: bug, environment :tickets: 300 diff --git a/tests/requirements.py b/tests/requirements.py index 0304919..c5d538d 100644 --- a/tests/requirements.py +++ b/tests/requirements.py @@ -41,6 +41,10 @@ class DefaultRequirements(SuiteRequirements): ) @property + def check_constraints_w_enforcement(self): + return exclusions.fails_on("mysql") + + @property def unnamed_constraints(self): """constraints without names are supported.""" return exclusions.only_on(['sqlite']) @@ -53,3 +57,10 @@ class DefaultRequirements(SuiteRequirements): @property def reflects_unique_constraints_unambiguously(self): return exclusions.fails_on("mysql") + + @property + def reflects_pk_names(self): + """Target driver reflects the name of primary key constraints.""" + + return exclusions.fails_on_everything_except( + 'postgresql', 'oracle', 'mssql', 'sybase') diff --git a/tests/test_batch.py b/tests/test_batch.py index 41d1957..0f0aada 100644 --- a/tests/test_batch.py +++ b/tests/test_batch.py @@ -2,6 +2,7 @@ from contextlib import contextmanager import re from alembic.testing import exclusions +from alembic.testing import assert_raises_message from alembic.testing import TestBase, eq_, config from alembic.testing.fixtures import op_fixture from alembic.testing import mock @@ -14,8 +15,9 @@ from sqlalchemy import Integer, Table, Column, String, MetaData, ForeignKey, \ UniqueConstraint, ForeignKeyConstraint, Index, Boolean, CheckConstraint, \ Enum from sqlalchemy.engine.reflection import Inspector -from sqlalchemy.sql import column +from sqlalchemy.sql import column, text from sqlalchemy.schema import CreateTable, CreateIndex +from sqlalchemy import exc class BatchApplyTest(TestBase): @@ -627,6 +629,50 @@ class BatchAPITest(TestBase): self.mock_schema.UniqueConstraint())] ) + def test_create_pk(self): + with self._fixture() as batch: + batch.create_primary_key('pk1', ['a', 'b']) + + eq_( + self.mock_schema.Table().c.__getitem__.mock_calls, + [mock.call('a'), mock.call('b')] + ) + + eq_( + self.mock_schema.PrimaryKeyConstraint.mock_calls, + [ + mock.call( + self.mock_schema.Table().c.__getitem__(), + self.mock_schema.Table().c.__getitem__(), + name='pk1' + ) + ] + ) + eq_( + batch.impl.operations.impl.mock_calls, + [mock.call.add_constraint( + self.mock_schema.PrimaryKeyConstraint())] + ) + + def test_create_check(self): + expr = text("a > b") + with self._fixture() as batch: + batch.create_check_constraint('ck1', expr) + + eq_( + self.mock_schema.CheckConstraint.mock_calls, + [ + mock.call( + expr, name="ck1" + ) + ] + ) + eq_( + batch.impl.operations.impl.mock_calls, + [mock.call.add_constraint( + self.mock_schema.CheckConstraint())] + ) + def test_drop_constraint(self): with self._fixture() as batch: batch.drop_constraint('uq1') @@ -795,6 +841,25 @@ class BatchRoundTripTest(TestBase): context = MigrationContext.configure(self.conn) self.op = Operations(context) + def _no_pk_fixture(self): + nopk = Table( + 'nopk', self.metadata, + Column('a', Integer), + Column('b', Integer), + Column('c', Integer), + mysql_engine='InnoDB' + ) + nopk.create(self.conn) + self.conn.execute( + nopk.insert(), + [ + {"a": 1, "b": 2, "c": 3}, + {"a": 2, "b": 4, "c": 5}, + ] + + ) + return nopk + def tearDown(self): self.metadata.drop_all(self.conn) self.conn.close() @@ -854,6 +919,32 @@ class BatchRoundTripTest(TestBase): {"id": 5, "x": 9} ]) + def test_add_pk_constraint(self): + self._no_pk_fixture() + with self.op.batch_alter_table("nopk", recreate="always") as batch_op: + batch_op.create_primary_key('newpk', ['a', 'b']) + + pk_const = Inspector.from_engine(self.conn).get_pk_constraint('nopk') + with config.requirements.reflects_pk_names.fail_if(): + eq_(pk_const['name'], 'newpk') + eq_(pk_const['constrained_columns'], ['a', 'b']) + + @config.requirements.check_constraints_w_enforcement + def test_add_ck_constraint(self): + with self.op.batch_alter_table("foo", recreate="always") as batch_op: + batch_op.create_check_constraint("newck", text("x > 0")) + + # we dont support reflection of CHECK constraints + # so test this by just running invalid data in + foo = self.metadata.tables['foo'] + + assert_raises_message( + exc.IntegrityError, + "newck", + self.conn.execute, + foo.insert(), {"id": 6, "data": 5, "x": -2} + ) + @config.requirements.sqlalchemy_094 @config.requirements.unnamed_constraints def test_drop_foreign_key(self): |