implement column adding with foreign keys on sqlite

2 files changed, 29 insertions, 21 deletions

diff --git a/migrate/changeset/databases/sqlite.py b/migrate/changeset/databases/sqlite.py
index 4136af3..2c49b03 100644
--- a/migrate/changeset/databases/sqlite.py
+++ b/migrate/changeset/databases/sqlite.py
@@ -51,14 +51,27 @@ class SQLiteHelper(SQLiteCommon):
         self.append('DROP TABLE migration_tmp')
         self.execute()
 
-class SQLiteColumnGenerator(SQLiteSchemaGenerator, SQLiteCommon, 
-                            ansisql.ANSIColumnGenerator):
+class SQLiteColumnGenerator(SQLiteSchemaGenerator, 
+                            ansisql.ANSIColumnGenerator,
+                            # at the end so we get the normal
+                            # visit_column by default
+                            SQLiteHelper,
+                            SQLiteCommon
+                            ):
     """SQLite ColumnGenerator"""
 
-    def add_foreignkey(self, constraint):
-        """Does not support ALTER TABLE ADD FOREIGN KEY"""
-        self._not_supported("ALTER TABLE ADD CONSTRAINT")
-
+    def _modify_table(self, table, column, delta):
+        columns = ' ,'.join(map(
+                self.preparer.format_column,
+                [c for c in table.columns if c.name!=column.name]))
+        return ('INSERT INTO %%(table_name)s (%(cols)s) '
+                'SELECT %(cols)s from migration_tmp')%{'cols':columns}
+
+    def visit_column(self,column):
+        if column.foreign_keys:
+            SQLiteHelper.visit_column(self,column)
+        else:
+            super(SQLiteColumnGenerator,self).visit_column(column)
 
 class SQLiteColumnDropper(SQLiteHelper, ansisql.ANSIColumnDropper):
     """SQLite ColumnDropper"""
@@ -73,7 +86,6 @@ class SQLiteSchemaChanger(SQLiteHelper, ansisql.ANSISchemaChanger):
     """SQLite SchemaChanger"""
 
     def _modify_table(self, table, column, delta):
-        column = table.columns[delta.current_name]
         return 'INSERT INTO %(table_name)s SELECT * from migration_tmp'
 
     def visit_index(self, index):
diff --git a/migrate/tests/changeset/test_changeset.py b/migrate/tests/changeset/test_changeset.py
index b1a4aba..8eff822 100644
--- a/migrate/tests/changeset/test_changeset.py
+++ b/migrate/tests/changeset/test_changeset.py
@@ -166,23 +166,19 @@ class TestAddDropColumn(fixture.DB):
 
         # create column with fk
         col = Column('data', Integer, ForeignKey(reftable.c.id))
-        if self.url.startswith('sqlite'):
-            self.assertRaises(exceptions.NotSupportedError,
-                col.create, self.table)
-        else:
-            col.create(self.table)
+        col.create(self.table)
 
-            # check if constraint is added
-            for cons in self.table.constraints:
-                if isinstance(cons, sqlalchemy.schema.ForeignKeyConstraint):
-                    break
-            else:
-                self.fail('No constraint found')
+        # check if constraint is added
+        for cons in self.table.constraints:
+            if isinstance(cons, sqlalchemy.schema.ForeignKeyConstraint):
+                break
+        else:
+            self.fail('No constraint found')
 
-            # TODO: test on db level if constraints work
+        # TODO: test on db level if constraints work
 
-            self.assertEqual(reftable.c.id.name, col.foreign_keys[0].column.name)
-            col.drop(self.table)
+        self.assertEqual(reftable.c.id.name, col.foreign_keys[0].column.name)
+        col.drop(self.table)
 
         if self.engine.has_table(reftable.name):
             reftable.drop()