Use two models in generated migrations. Test column addition and removal.

3 files changed, 95 insertions, 51 deletions

diff --git a/migrate/tests/versioning/test_genmodel.py b/migrate/tests/versioning/test_genmodel.py
index 77e7a66..c12900d 100644
--- a/migrate/tests/versioning/test_genmodel.py
+++ b/migrate/tests/versioning/test_genmodel.py
@@ -42,14 +42,15 @@ class TestSchemaDiff(fixture.DB):
 
         def assertDiff(isDiff, tablesMissingInDatabase, tablesMissingInModel, tablesWithDiff):
             diff = schemadiff.getDiffOfModelAgainstDatabase(self.meta, self.engine, excludeTables=['migrate_version'])
-            eq_(bool(diff), isDiff)
             eq_(
                 (diff.tables_missing_from_B,
                  diff.tables_missing_from_A,
-                 diff.tables_different.keys()),
+                 diff.tables_different.keys(),
+                 bool(diff)),
                 (tablesMissingInDatabase,
                  tablesMissingInModel,
-                 tablesWithDiff)
+                 tablesWithDiff,
+                 isDiff)
                 )
 
         # Model is defined but database is empty.
@@ -64,8 +65,9 @@ class TestSchemaDiff(fixture.DB):
         if repr(String()) == 'String()':
             self.assertEqualsIgnoreWhitespace(decls, '''
             from migrate.changeset import schema
-            meta = MetaData()
-            tmp_schemadiff = Table('tmp_schemadiff', meta,
+            pre_meta = MetaData()
+            post_meta = MetaData()
+            tmp_schemadiff = Table('tmp_schemadiff', post_meta,
                 Column('id', Integer, primary_key=True, nullable=False),
                 Column('name', UnicodeText),
                 Column('data', UnicodeText),
@@ -74,19 +76,14 @@ class TestSchemaDiff(fixture.DB):
         else:
             self.assertEqualsIgnoreWhitespace(decls, '''
             from migrate.changeset import schema
-            meta = MetaData()
-            tmp_schemadiff = Table('tmp_schemadiff', meta,
+            pre_meta = MetaData()
+            post_meta = MetaData()
+            tmp_schemadiff = Table('tmp_schemadiff', post_meta,
                 Column('id', Integer, primary_key=True, nullable=False),
                 Column('name', UnicodeText(length=None)),
                 Column('data', UnicodeText(length=None)),
             )
             ''')
-        self.assertEqualsIgnoreWhitespace(upgradeCommands,
-            '''meta.bind = migrate_engine
-            tmp_schemadiff.create()''')
-        self.assertEqualsIgnoreWhitespace(downgradeCommands,
-            '''meta.bind = migrate_engine
-            tmp_schemadiff.drop()''')
 
         # Create table in database, now model should match database.
         self._applyLatestModel()
@@ -111,12 +108,14 @@ class TestSchemaDiff(fixture.DB):
             else:
                 dataId = result.last_inserted_ids()[0]
 
-        # Modify table in model (by removing it and adding it back to model) -- drop column data and add column data2.
+        # Modify table in model (by removing it and adding it back to model)
+        # Drop column data, add columns data2 and data3.
         self.meta.remove(self.table)
         self.table = Table(self.table_name,self.meta,
             Column('id',Integer(),primary_key=True),
             Column('name',UnicodeText(length=None)),
             Column('data2',Integer(),nullable=True),
+            Column('data3',Integer(),nullable=True),
         )
         assertDiff(True, [], [], [self.table_name])
 
@@ -124,6 +123,36 @@ class TestSchemaDiff(fixture.DB):
         self._applyLatestModel()
         assertDiff(False, [], [], [])
 
+        # Drop column data3, add data4
+        self.meta.remove(self.table)
+        self.table = Table(self.table_name,self.meta,
+            Column('id',Integer(),primary_key=True),
+            Column('name',UnicodeText(length=None)),
+            Column('data2',Integer(),nullable=True),
+            Column('data4',Float(),nullable=True),
+        )
+        assertDiff(True, [], [], [self.table_name])
+
+        diff = schemadiff.getDiffOfModelAgainstDatabase(
+            self.meta, self.engine, excludeTables=['migrate_version'])
+        decls, upgradeCommands, downgradeCommands = genmodel.ModelGenerator(diff,self.engine).genB2AMigration(indent='')
+
+        # decls have changed since genBDefinition
+        exec decls in locals()
+        # migration commands expect a namespace containing migrate_engine
+        migrate_engine = self.engine
+        # run the migration up and down
+        exec upgradeCommands in locals()
+        assertDiff(False, [], [], [])
+
+        exec decls in locals()
+        exec downgradeCommands in locals()
+        assertDiff(True, [], [], [self.table_name])
+
+        exec decls in locals()
+        exec upgradeCommands in locals()
+        assertDiff(False, [], [], [])
+
         if not self.engine.name == 'oracle':
             # Make sure data is still present.
             result = self.engine.execute(self.table.select(self.table.c.id==dataId))
diff --git a/migrate/tests/versioning/test_script.py b/migrate/tests/versioning/test_script.py
index 6c9d585..53ef929 100644
--- a/migrate/tests/versioning/test_script.py
+++ b/migrate/tests/versioning/test_script.py
@@ -158,8 +158,8 @@ def upgrade(migrate_engine):
             repository=self.repo_path,
         )
 
-        self.assertTrue('User.create()' in source_script)
-        self.assertTrue('User.drop()' in source_script)
+        self.assertTrue("['User'].create()" in source_script)
+        self.assertTrue("['User'].drop()" in source_script)
 
     @fixture.usedb()
     def test_make_update_script_for_equal_models(self):
@@ -196,9 +196,9 @@ def upgrade(migrate_engine):
 
         self.assertTrue(0
                         < source_script.find('upgrade')
-                        < source_script.find('User.create()')
+                        < source_script.find("['User'].create()")
                         < source_script.find('downgrade')
-                        < source_script.find('User.drop()'))
+                        < source_script.find("['User'].drop()"))
 
     def setup_model_params(self):
         self.script_path = self.tmp_py()
diff --git a/migrate/versioning/genmodel.py b/migrate/versioning/genmodel.py
index 06d3175..85df627 100644
--- a/migrate/versioning/genmodel.py
+++ b/migrate/versioning/genmodel.py
@@ -96,7 +96,7 @@ class ModelGenerator(object):
         else:
             return """Column(%(name)r, %(commonStuff)s)""" % data
 
-    def _getTableDefn(self, table):
+    def _getTableDefn(self, table, metaName='meta'):
         out = []
         tableName = table.name
         if self.declarative:
@@ -107,8 +107,8 @@ class ModelGenerator(object):
                 out.append("    %s" % self.column_repr(col))
             out.append('\n')
         else:
-            out.append("%(table)s = Table('%(table)s', meta," %
-                            {'table': tableName})
+            out.append("%(table)s = Table('%(table)s', %(meta)s," %
+                       {'table': tableName, 'meta': metaName})
             for col in table.columns:
                 out.append("    %s," % self.column_repr(col))
             out.append(")\n")
@@ -151,46 +151,61 @@ class ModelGenerator(object):
         '''
 
         decls = ['from migrate.changeset import schema',
-                 'meta = MetaData()']
-        for table in self._get_tables(
-            missingA=True,missingB=True,modified=True
-            ):
-            decls.extend(self._getTableDefn(table))
-
-        upgradeCommands, downgradeCommands = [], []
-        for tableName in self.diff.tables_missing_from_A:
-            upgradeCommands.append("%(table)s.drop()" % {'table': tableName})
-            downgradeCommands.append("%(table)s.create()" % \
-                                         {'table': tableName})
-        for tableName in self.diff.tables_missing_from_B:
-            upgradeCommands.append("%(table)s.create()" % {'table': tableName})
-            downgradeCommands.append("%(table)s.drop()" % {'table': tableName})
-
-        for tableName in self.diff.tables_different:
-            dbTable = self.diff.metadataB.tables[tableName]
-            td = self.diff.tables_different[tableName]
+                 'pre_meta = MetaData()',
+                 'post_meta = MetaData()',
+                ]
+        upgradeCommands = ['pre_meta.bind = migrate_engine',
+                           'post_meta.bind = migrate_engine']
+        downgradeCommands = list(upgradeCommands)
+
+        for tn in self.diff.tables_missing_from_A:
+            pre_table = self.diff.metadataB.tables[tn]
+            decls.extend(self._getTableDefn(pre_table, metaName='pre_meta'))
+            upgradeCommands.append(
+                "pre_meta.tables[%(table)r].drop()" % {'table': tn})
+            downgradeCommands.append(
+                "pre_meta.tables[%(table)r].create()" % {'table': tn})
+
+        for tn in self.diff.tables_missing_from_B:
+            post_table = self.diff.metadataA.tables[tn]
+            decls.extend(self._getTableDefn(post_table, metaName='post_meta'))
+            upgradeCommands.append(
+                "post_meta.tables[%(table)r].create()" % {'table': tn})
+            downgradeCommands.append(
+                "post_meta.tables[%(table)r].drop()" % {'table': tn})
+
+        for (tn, td) in self.diff.tables_different.iteritems():
+            if td.columns_missing_from_A or td.columns_different:
+                pre_table = self.diff.metadataB.tables[tn]
+                decls.extend(self._getTableDefn(
+                    pre_table, metaName='pre_meta'))
+            if td.columns_missing_from_B or td.columns_different:
+                post_table = self.diff.metadataA.tables[tn]
+                decls.extend(self._getTableDefn(
+                    post_table, metaName='post_meta'))
+
             for col in td.columns_missing_from_A:
-                upgradeCommands.append('%s.append_column(%s)' % (
-                   tableName,
-                   self.column_repr(
-                      self.diff.metadataB.tables[tableName].columns[col])))
-                downgradeCommands.append('%s.columns[%r].drop()' % (tableName, col))
+                upgradeCommands.append(
+                    'pre_meta.tables[%r].columns[%r].drop()' % (tn, col))
+                downgradeCommands.append(
+                    'pre_meta.tables[%r].columns[%r].create()' % (tn, col))
             for col in td.columns_missing_from_B:
-                upgradeCommands.append('%s.columns[%r].drop()' % (tableName, col))
-                downgradeCommands.append('%s.columns[%r].create()' % (tableName, col))
+                upgradeCommands.append(
+                    'post_meta.tables[%r].columns[%r].create()' % (tn, col))
+                downgradeCommands.append(
+                    'post_meta.tables[%r].columns[%r].drop()' % (tn, col))
             for modelCol, databaseCol, modelDecl, databaseDecl in td.columns_different:
                 upgradeCommands.append(
                     'assert False, "Can\'t alter columns: %s:%s=>%s"' % (
-                    tableName, modelCol.name, databaseCol.name))
+                    tn, modelCol.name, databaseCol.name))
                 downgradeCommands.append(
                     'assert False, "Can\'t alter columns: %s:%s=>%s"' % (
-                    tableName, modelCol.name, databaseCol.name))
-        pre_command = '    meta.bind = migrate_engine'
+                    tn, modelCol.name, databaseCol.name))
 
         return (
             '\n'.join(decls),
-            '\n'.join([pre_command] + ['%s%s' % (indent, line) for line in upgradeCommands]),
-            '\n'.join([pre_command] + ['%s%s' % (indent, line) for line in downgradeCommands]))
+            '\n'.join('%s%s' % (indent, line) for line in upgradeCommands),
+            '\n'.join('%s%s' % (indent, line) for line in downgradeCommands))
 
     def _db_can_handle_this_change(self,td):
         """Check if the database can handle going from B to A."""