From: Mike Bayer Date: Thu, 24 Nov 2011 00:27:58 +0000 (-0500) Subject: - support the constraints generated from SchemaType, honoring conditional X-Git-Tag: rel_0_1_0~42 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=e7e063cc60828607d355f67790055222983be90d;p=thirdparty%2Fsqlalchemy%2Falembic.git - support the constraints generated from SchemaType, honoring conditional rule - add mssql_drop_default, mssql_drop_check flags to drop_column(), will emit special MSSQL crap to drop DEFAULT and CHECK constraints based on inline system table lookup --- diff --git a/alembic/ddl/impl.py b/alembic/ddl/impl.py index c2db03c9..2cfaca7e 100644 --- a/alembic/ddl/impl.py +++ b/alembic/ddl/impl.py @@ -90,11 +90,13 @@ class DefaultImpl(object): def add_column(self, table_name, column): self._exec(base.AddColumn(table_name, column)) - def drop_column(self, table_name, column): + def drop_column(self, table_name, column, **kw): self._exec(base.DropColumn(table_name, column)) def add_constraint(self, const): - self._exec(schema.AddConstraint(const)) + if const._create_rule is None or \ + const._create_rule(self): + self._exec(schema.AddConstraint(const)) def drop_constraint(self, const): self._exec(schema.DropConstraint(const)) diff --git a/alembic/ddl/mssql.py b/alembic/ddl/mssql.py index d9d0fac2..42c2952a 100644 --- a/alembic/ddl/mssql.py +++ b/alembic/ddl/mssql.py @@ -21,6 +21,34 @@ class MSSQLImpl(DefaultImpl): super(MSSQLImpl, self).bulk_insert(table, rows) + def drop_column(self, table_name, column, **kw): + drop_default = kw.pop('mssql_drop_default', False) + if drop_default: + self._exec( + _exec_drop_col_constraint(table_name, column, 'sys.default_constraints') + ) + drop_check = kw.pop('mssql_drop_check', False) + if drop_check: + self._exec( + _exec_drop_col_constraint(table_name, column, 'sys.check_constraints') + ) + super(MSSQLImpl, self).drop_column(table_name, column) + + +def _exec_drop_col_constraint(tname, colname, type_): + # from http://www.mssqltips.com/sqlservertip/1425/working-with-default-constraints-in-sql-server/ + # TODO: needs table formatting, etc. + return """declare @const_name varchar(256) +select @const_name = [name] from %(type)s +where parent_object_id = object_id('%(tname)s') +and col_name(parent_object_id, parent_column_id) = '%(colname)s' +exec('alter table %(tname)s drop constraint ' + @const_name)""" % { + 'type':type_, + 'tname':tname, + 'colname':colname + } + + @compiles(AddColumn, 'mssql') def visit_add_column(element, compiler, **kw): return "%s %s" % ( @@ -39,3 +67,5 @@ def visit_rename_column(element, compiler, **kw): format_column_name(compiler, element.column_name), format_column_name(compiler, element.newname) ) + + diff --git a/alembic/op.py b/alembic/op.py index d4a76ce3..e1c8a937 100644 --- a/alembic/op.py +++ b/alembic/op.py @@ -137,21 +137,36 @@ def add_column(table_name, column): table_name, column ) - for constraint in [f.constraint for f in t.foreign_keys]: - get_impl().add_constraint(constraint) + for constraint in t.constraints: + if not isinstance(constraint, schema.PrimaryKeyConstraint): + get_impl().add_constraint(constraint) -def drop_column(table_name, column_name): +def drop_column(table_name, column_name, **kw): """Issue a "drop column" instruction using the current change context. e.g.:: drop_column('organization', 'account_id') + :param table_name: name of table + :param column_name: name of column + :param mssql_drop_check: Optional boolean. When ``True``, on + Microsoft SQL Server only, first + drop the CHECK constraint on the column using a SQL-script-compatible + block that selects into a @variable from sys.check_constraints, + then exec's a separate DROP CONSTRAINT for that constraint. + :param mssql_drop_default: Optional boolean. When ``True``, on + Microsoft SQL Server only, first + drop the DEFAULT constraint on the column using a SQL-script-compatible + block that selects into a @variable from sys.default_constraints, + then exec's a separate DROP CONSTRAINT for that default. + """ get_impl().drop_column( table_name, - _column(column_name, NULLTYPE) + _column(column_name, NULLTYPE), + **kw ) diff --git a/tests/__init__.py b/tests/__init__.py index b2ecb944..1157aafe 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -101,6 +101,15 @@ def _op_fixture(dialect='default', as_sql=False): # whitespace and such eq_(self.impl.assertion, list(sql)) + def assert_contains(self, sql): + for stmt in self.impl.assertion: + if sql in stmt: + return + else: + assert False, "Could not locate fragment %r in %r" % ( + sql, + self.impl.assertion + ) return ctx(dialect, as_sql) def _sqlite_testing_config(): diff --git a/tests/test_mssql.py b/tests/test_mssql.py index f229c0e3..c84f8491 100644 --- a/tests/test_mssql.py +++ b/tests/test_mssql.py @@ -23,6 +23,19 @@ def test_alter_column_rename_mssql(): "EXEC sp_rename 't.c', 'x', 'COLUMN'" ) +def test_drop_column_w_default(): + context = _op_fixture('mssql') + op.drop_column('t1', 'c1', mssql_drop_default=True) + context.assert_contains("exec('alter table t1 drop constraint ' + @const_name)") + context.assert_contains("ALTER TABLE t1 DROP COLUMN c1") + + +def test_drop_column_w_check(): + context = _op_fixture('mssql') + op.drop_column('t1', 'c1', mssql_drop_check=True) + context.assert_contains("exec('alter table t1 drop constraint ' + @const_name)") + context.assert_contains("ALTER TABLE t1 DROP COLUMN c1") + # TODO: when we add schema support #def test_alter_column_rename_mssql_schema(): # context = _op_fixture('mssql') diff --git a/tests/test_op.py b/tests/test_op.py index 7d63e541..c9f5e77c 100644 --- a/tests/test_op.py +++ b/tests/test_op.py @@ -3,7 +3,8 @@ from tests import _op_fixture from alembic import op from sqlalchemy import Integer, Column, ForeignKey, \ - UniqueConstraint, Table, MetaData, String + UniqueConstraint, Table, MetaData, String,\ + Boolean from sqlalchemy.sql import table def test_add_column(): @@ -24,6 +25,24 @@ def test_add_column_fk(): "ALTER TABLE t1 ADD FOREIGN KEY(c1) REFERENCES c2 (id)" ) +def test_add_column_schema_type(): + """Test that a schema type generates its constraints....""" + context = _op_fixture() + op.add_column('t1', Column('c1', Boolean, nullable=False)) + context.assert_( + 'ALTER TABLE t1 ADD COLUMN c1 BOOLEAN NOT NULL', + 'ALTER TABLE t1 ADD CHECK (c1 IN (0, 1))' + ) + +def test_add_column_schema_type_checks_rule(): + """Test that a schema type doesn't generate a + constraint based on check rule.""" + context = _op_fixture('postgresql') + op.add_column('t1', Column('c1', Boolean, nullable=False)) + context.assert_( + 'ALTER TABLE t1 ADD COLUMN c1 BOOLEAN NOT NULL', + ) + def test_add_column_fk_self_referential(): context = _op_fixture() op.add_column('t1', Column('c1', Integer, ForeignKey('t1.c2'), nullable=False))