- 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 c2db03c9c9adb4dd6b73e31f8f553dfbf712f26e..2cfaca7e10135c598ad69c6b5666a63a3891022f 100644 (file)
--- 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 d9d0fac2cf88f4d6af30892ca8c64e4c1d31b3c7..42c2952a7aac47ed9d89d64960206a95c81860f9 100644 (file)
--- 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 d4a76ce3a0516858c4943cae1dcd3bb7c7036333..e1c8a93716bba8e97e1fe743e36aa98d77d8e0d1 100644 (file)
--- 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 b2ecb944a2e85ea7a018c74345a1d1822287ee9b..1157aafe62bce1bdf757244a5dc02aace028fac8 100644 (file)
--- 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 f229c0e31edbb503ed79a50aba5e3d8fb506bcda..c84f849147d13228900fb81d99eaa52a0ea009aa 100644 (file)
--- 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 7d63e54156d311ee499c70a3ecdbd8e1402036be..c9f5e77c3d8f9304f9a2ab23c223a5b604da2b0d 100644 (file)
--- 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))