Workaround for when FK rel field is itself an FK.

Author: coleifer

playhouse/migrations.py

@@ -332,9 +332,13 @@ def _build_stub(model, fields, imports):
     meta = model._meta
     lines = ['class %s(Model):' % model.__name__]
     for field in fields:
-        lines.append('    %s = %s' % (
-            field.name,
-            _build_field(field, imports, [])))
+        if isinstance(field, ForeignKeyField):
+            # A pk that is itself a fk stores a plain integer.
+            source = 'IntegerField(primary_key=True, column_name=%r)' % (
+                field.column_name)
+        else:
+            source = _build_field(field, imports, [])
+        lines.append('    %s = %s' % (field.name, source))
     lines.append('    class Meta:')
     lines.append('        database = db')
     lines.append('        table_name = %r' % meta.table_name)

tests/schema_diff.py

@@ -102,6 +102,22 @@ class Meta:
         table_name = 'sd_badge'
 
 
+class SdAccount(TestModel):
+    label = CharField()
+
+    class Meta:
+        table_name = 'sd_account'
+
+
+class SdSettings(TestModel):
+    # One-to-one extension: the pk is itself an fk.
+    account = ForeignKeyField(SdAccount, primary_key=True)
+    theme = CharField()
+
+    class Meta:
+        table_name = 'sd_settings'
+
+
 @skip_if(IS_CRDB, 'crdb introspection differs')
 class TestSchemaDiff(ModelTestCase):
     requires = SD_MODELS
@@ -571,7 +587,8 @@ def tearDown(self):
         try:
             shutil.rmtree(self.dir, ignore_errors=True)
             self.database.drop_tables([self.runner.History], safe=True)
-            for table in ('sd_badge', 'sd_slug', 'sd_tag', 'sd_sku',
+            for table in ('sd_theme', 'sd_settings', 'sd_account',
+                          'sd_badge', 'sd_slug', 'sd_tag', 'sd_sku',
                           'sd_profile', 'sd_alias', 'sd_kv'):
                 self.database.execute_sql('DROP TABLE IF EXISTS %s' % table)
         finally:
@@ -715,6 +732,28 @@ def test_fk_stub(self):
         diff = diff_models(self.database, models)
         self.assertEqual(diff.create_tables, [SdBadge])
 
+    @requires_models(SdAccount, SdSettings)
+    def test_fk_stub_pk_fk(self):
+        # sd_settings' pk is itself an fk. The stub renders it as a
+        # plain integer rather than recursing into sd_account.
+        class SdTheme(TestModel):
+            settings = ForeignKeyField(SdSettings)
+            name = CharField()
+
+            class Meta:
+                table_name = 'sd_theme'
+
+        models = [SdTheme, SdSettings, SdAccount] + SD_MODELS
+        body = template(diff_models(self.database, models))
+        self.assertIn("account = IntegerField(primary_key=True, "
+                      "column_name='account_id')", body)
+        self.apply(body, 'theme')
+        self.assertFalse(diff_models(self.database, models))
+
+        self.assertEqual(self.runner.down(), ['0001_theme'])
+        self.assertEqual(diff_models(self.database, models).create_tables,
+                         [SdTheme])
+
     def test_autofield_alias(self):
         # The aliased pk must render or the table never converges.
         class SdAlias(TestModel):