Fix bug with column dropping involving foreign keys.

Bonus: remove_from_table now understands foreign keys

docs/changelog.rst

@@ -12,6 +12,7 @@ Fixed bugs
 - cleared up test output and improved testing of deprecation warnings.
 - some documentation fixes
 - fixed bug with column dropping in sqlite (issue 96)
+- fixed bug with column dropping involving foreign keys
 
 0.6 (11.07.2010)
 ---------------------------

migrate/changeset/schema.py

@@ -2,10 +2,13 @@
    Schema module providing common schema operations.
 """
 import warnings
+
 from UserDict import DictMixin
 
 import sqlalchemy
 
+from sqlalchemy.schema import ForeignKeyConstraint
+
 from migrate.exceptions import *
 from migrate.changeset import SQLA_06
 from migrate.changeset.databases.visitor import (get_engine_visitor,
@@ -571,6 +574,7 @@ populated with defaults
         # TODO: remove primary keys, constraints, etc
         if unset_table:
             self.table = None
+            
         to_drop = set()
         for index in table.indexes:
             columns = []
@@ -582,6 +586,20 @@ populated with defaults
             else:
                 to_drop.add(index)
         table.indexes = table.indexes - to_drop
+        
+        to_drop = set()
+        for cons in table.constraints:
+            # TODO: deal with other types of constraint
+            if isinstance(cons,ForeignKeyConstraint):
+                col_names = []
+                for col_name in cons.columns:
+                    if not isinstance(col_name,basestring):
+                        col_name = col_name.name
+                    col_names.append(col_name)
+                if self.name in col_names:
+                    to_drop.add(cons)
+        table.constraints = table.constraints - to_drop
+        
         if table.c.contains_column(self):
             table.c.remove(self)
 

migrate/tests/changeset/test_changeset.py

@@ -334,6 +334,96 @@ class TestAddDropColumn(fixture.DB):
             # a crude test for 0.5.x
             Index('ix_tmp_adddropcol_d1',self.table.c.d1).drop()
             
+    def _actual_foreign_keys(self):
+        from sqlalchemy.schema import ForeignKeyConstraint
+        result = []
+        for cons in self.table.constraints:
+            if isinstance(cons,ForeignKeyConstraint):
+                col_names = []
+                for col_name in cons.columns:
+                    if not isinstance(col_name,basestring):
+                        col_name = col_name.name
+                    col_names.append(col_name)
+                result.append(col_names)
+        result.sort()
+        return result
+        
+    @fixture.usedb()
+    def test_drop_with_foreign_keys(self):
+        self.table.drop()
+        self.meta.clear()
+        
+        # create FK's target
+        reftable = Table('tmp_ref', self.meta,
+            Column('id', Integer, primary_key=True),
+        )
+        if self.engine.has_table(reftable.name):
+            reftable.drop()
+        reftable.create()
+        
+        # add a table with two foreign key columns
+        self.table = Table(
+            self.table_name, self.meta,
+            Column('id', Integer, primary_key=True),
+            Column('r1', Integer, ForeignKey('tmp_ref.id')),
+            Column('r2', Integer, ForeignKey('tmp_ref.id')),
+            )
+        self.table.create()
+
+        # paranoid check
+        self.assertEqual([['r1'],['r2']],
+                         self._actual_foreign_keys())
+        
+        # delete one
+        self.table.c.r2.drop()
+        
+        # check remaining foreign key is there
+        self.assertEqual([['r1']],
+                         self._actual_foreign_keys())
+        
+    @fixture.usedb()
+    def test_drop_with_complex_foreign_keys(self):
+        from sqlalchemy.schema import ForeignKeyConstraint
+        from sqlalchemy.schema import UniqueConstraint
+        
+        self.table.drop()
+        self.meta.clear()
+        
+        # create FK's target
+        reftable = Table('tmp_ref', self.meta,
+            Column('id', Integer, primary_key=True),
+            Column('jd', Integer),
+            UniqueConstraint('id','jd')
+            )
+        if self.engine.has_table(reftable.name):
+            reftable.drop()
+        reftable.create()
+        
+        # add a table with a complex foreign key constraint
+        self.table = Table(
+            self.table_name, self.meta,
+            Column('id', Integer, primary_key=True),
+            Column('r1', Integer),
+            Column('r2', Integer),
+            ForeignKeyConstraint(['r1','r2'],
+                                 [reftable.c.id,reftable.c.jd])
+            )
+        self.table.create()
+
+        # paranoid check
+        self.assertEqual([['r1','r2']],
+                         self._actual_foreign_keys())
+        
+        # delete one
+        self.table.c.r2.drop()
+        
+        # check the constraint is gone, since part of it
+        # is no longer there - if people hit this,
+        # they may be confused, maybe we should raise an error
+        # and insist that the constraint is deleted first, separately?
+        self.assertEqual([],
+                         self._actual_foreign_keys())
+        
 class TestRename(fixture.DB):
     """Tests for table and index rename methods"""
     level = fixture.DB.CONNECT