Foreign key fixes for pwiz, fixes #426

playhouse/tests_pwiz.py

@@ -248,6 +248,7 @@ class TestPwiz(unittest.TestCase):
         self.assertEqual(col_types_id.get_field_parameters(), {
             'db_column': "'col_types_id'",
             'rel_model': 'Coltypes',
+            'to_field': "'f11'",
         })
 
         col_types_nullable_id = rel_model['col_types_nullable_id']
@@ -255,13 +256,25 @@ class TestPwiz(unittest.TestCase):
             'db_column': "'col_types_nullable_id'",
             'null': True,
             'rel_model': 'Coltypes',
+            'to_field': "'f11'",
         })
 
         fkpk = columns['fkpk']
         self.assertEqual(fkpk['col_types_id'].get_field_parameters(), {
             'db_column': "'col_types_id'",
             'rel_model': 'Coltypes',
-            'primary_key': True})
+            'primary_key': True,
+            'to_field': "'f11'"})
+
+        category = columns['category']
+
+        parent_id = category['parent_id']
+        self.assertEqual(parent_id.get_field_parameters(), {
+            'db_column': "'parent_id'",
+            'null': True,
+            'rel_model': "'self'",
+            'to_field': "'id'",
+        })
 
     @generative_test
     def test_get_field(self, introspector):
@@ -287,14 +300,15 @@ class TestPwiz(unittest.TestCase):
             ('fkpk', (
                 ('col_types_id', 'col_types = ForeignKeyField('
                  'db_column=\'col_types_id\', primary_key=True, '
-                 'rel_model=Coltypes)'),
+                 'rel_model=Coltypes, to_field=\'f11\')'),
             )),
             ('relmodel', (
                 ('col_types_id', 'col_types = ForeignKeyField('
-                 'db_column=\'col_types_id\', rel_model=Coltypes)'),
+                 'db_column=\'col_types_id\', rel_model=Coltypes, '
+                 'to_field=\'f11\')'),
                 ('col_types_nullable_id', 'col_types_nullable = '
                  'ForeignKeyField(db_column=\'col_types_nullable_id\', '
-                 'null=True, rel_model=Coltypes)'),
+                 'null=True, rel_model=Coltypes, to_field=\'f11\')'),
             )),
             ('underscores', (
                 ('_id', '_id = PrimaryKeyField()'),
@@ -303,7 +317,8 @@ class TestPwiz(unittest.TestCase):
             ('category', (
                 ('name', 'name = CharField(max_length=10)'),
                 ('parent_id', 'parent = ForeignKeyField('
-                 'db_column=\'parent_id\', null=True, rel_model=\'self\')'),
+                 'db_column=\'parent_id\', null=True, rel_model=\'self\', '
+                 'to_field=\'id\')'),
             )),
         )
 

pwiz.py

@@ -84,19 +84,21 @@ class Column(object):
         # Handle ForeignKeyField-specific attributes.
         if self.field_class is ForeignKeyField:
             params['rel_model'] = self.rel_model
+            if self.to_field:
+                params['to_field'] = "'%s'" % self.to_field
 
         return params
 
     def is_primary_key(self):
         return self.field_class is PrimaryKeyField or self.primary_key
 
-    def set_foreign_key(self, foreign_key, model_names):
+    def set_foreign_key(self, foreign_key, model_names, dest=None):
         self.field_class = ForeignKeyField
         if foreign_key.dest_table == foreign_key.table:
             self.rel_model = "'self'"
         else:
             self.rel_model = model_names[foreign_key.dest_table]
-        self.to_field = foreign_key.dest_column
+        self.to_field = dest and dest.name or None
 
     def get_field(self):
         # Generate the field definition for this column.
@@ -442,11 +444,17 @@ class SqliteMetadata(Metadata):
         table_definition = cursor.fetchone()[0].strip()
 
         try:
-            columns = re.search('\((.+)\)', table_definition).groups()[0]
+            columns = re.search(
+                '\((.+)\)',
+                table_definition,
+                re.MULTILINE | re.DOTALL).groups()[0]
         except AttributeError:
             print_('Unable to read table definition for "%s"' % table)
             return []
 
+        # Replace any new-lines or other junk with whitespace.
+        columns = re.sub('[\s\n\r]+', ' ', columns).strip()
+
         fks = []
         for column_def in columns.split(','):
             column_def = column_def.strip()
@@ -529,13 +537,22 @@ class Introspector(object):
 
         # On the second pass convert all foreign keys.
         for table in tables:
-            for foreign_key in foreign_keys[table]:
-                src = columns[foreign_key.table][foreign_key.column]
-                src.set_foreign_key(foreign_key, model_names)
-
             for column_name, column in columns[table].items():
                 column.name = self.make_column_name(column_name)
 
+            for foreign_key in foreign_keys[table]:
+                src = columns[foreign_key.table][foreign_key.column]
+                try:
+                    dest = columns[foreign_key.dest_table][
+                        foreign_key.dest_column]
+                except KeyError:
+                    dest = None
+
+                src.set_foreign_key(
+                    foreign_key,
+                    model_names,
+                    dest)
+
         return columns, foreign_keys, model_names
 
     def print_models(self, tables=None):