Fix genmodel for SQLA 0.9

Problem:
* Some python code was auto generated and exec'ed in that
package.
* The python code that was problematic had a 'Table' definition
* The generated code imports '*' from sqlalchemy
* One among the 'Column' was defined as an 'INTEGER' type, which
points to sqlalchemy.sql.sqltypes.INTEGER
* The INTEGER class was initialised with a parameter display_width
which contradicts with sqlalchemy.sql.sqltypes.INTEGER.__init__,
which does not accept any parameters
* The 'INTEGER' class should have been imported from mysql dialects'
type module
Solution:
* While generating, in the header part, I am now checking if any of
the column.type.__class__.__name__ has 'dialects' in it.
* If I find any, I am adding the import rule such that the type is
imported from the dialects' type.

This patch has been tested with SQLA 0.9.3, for which it fixes the
unit tests, and with SQLA 0.8.2, which doesn't have (new) problems
with this patch.

Change-Id: Ie0e09b45388462629100017bea3ea8a314d148d8

1 files changed, 27 insertions, 10 deletions

diff --git a/migrate/versioning/genmodel.py b/migrate/versioning/genmodel.py
index 85df627..efff67f 100644
--- a/migrate/versioning/genmodel.py
+++ b/migrate/versioning/genmodel.py
@@ -20,13 +20,11 @@ HEADER = """
 ## File autogenerated by genmodel.py
 
 from sqlalchemy import *
-meta = MetaData()
 """
 
-DECLARATIVE_HEADER = """
-## File autogenerated by genmodel.py
+META_DEFINITION = "meta = MetaData()"
 
-from sqlalchemy import *
+DECLARATIVE_DEFINITION = """
 from sqlalchemy.ext import declarative
 
 Base = declarative.declarative_base()
@@ -125,6 +123,30 @@ class ModelGenerator(object):
                 for name in names:
                     yield metadata.tables.get(name)
 
+    def _genModelHeader(self, tables):
+        out = []
+        import_index = []
+
+        out.append(HEADER)
+
+        for table in tables:
+            for col in table.columns:
+                if "dialects" in col.type.__module__ and \
+                   col.type.__class__ not in import_index:
+                    out.append("from " + col.type.__module__ +
+                               " import " + col.type.__class__.__name__)
+                    import_index.append(col.type.__class__)
+
+        out.append("")
+
+        if self.declarative:
+            out.append(DECLARATIVE_DEFINITION)
+        else:
+            out.append(META_DEFINITION)
+        out.append("")
+
+        return out
+
     def genBDefinition(self):
         """Generates the source code for a definition of B.
 
@@ -134,11 +156,7 @@ class ModelGenerator(object):
         """
 
         out = []
-        if self.declarative:
-            out.append(DECLARATIVE_HEADER)
-        else:
-            out.append(HEADER)
-        out.append("")
+        out.extend(self._genModelHeader(self._get_tables(missingA=True)))
         for table in self._get_tables(missingA=True):
             out.extend(self._getTableDefn(table))
         return '\n'.join(out)
@@ -282,4 +300,3 @@ class ModelGenerator(object):
                 except:
                     trans.rollback()
                     raise
-