Merge pull request #50 from citusdata/dropmodel-migration

fix migration for foreign keys when model doesn't exist anymore

django_multitenant/backends/postgresql/base.py

@@ -10,6 +10,7 @@
     DatabaseIntrospection
 )
 from django_multitenant.fields import TenantForeignKey
+from django_multitenant.utils import get_model_by_db_table, get_tenant_column
 
 logger = logging.getLogger(__name__)
 
@@ -42,17 +43,24 @@ def _alter_field(self, model, old_field, new_field, old_type, new_type,
     # Override
     def _create_fk_sql(self, model, field, suffix):
         if isinstance(field, TenantForeignKey):
-            from_table = model._meta.db_table
-            from_columns = field.column, self.get_tenant_column(model)
-            to_table = field.target_field.model._meta.db_table
-            to_columns = field.target_field.column, self.get_tenant_column(field.target_field.model)
+            try:
+                # test if both models exists
+                # This case happens when we are running from scratch migrations and one model was removed from code
+                # In the previous migrations we would still be creating the foreign key
+                from_model = get_model_by_db_table(model._meta.db_table)
+                to_model = get_model_by_db_table(field.target_field.model._meta.db_table)
+            except ValueError:
+                return None
+
+            from_columns = field.column, get_tenant_column(from_model)
+            to_columns = field.target_field.column, get_tenant_column(to_model)
             suffix = suffix % {
-                "to_table": to_table,
+                "to_table": field.target_field.model._meta.db_table,
                 "to_column": '_'.join(to_columns),
             }
 
             return self.sql_create_fk % {
-                "table": self.quote_name(from_table),
+                "table": self.quote_name(model._meta.db_table),
                 "name": self.quote_name(self._create_index_name(model, from_columns, suffix=suffix)),
                 "column": ', '.join([self.quote_name(from_col) for from_col in from_columns]),
                 "to_table": self.quote_name(field.target_field.model._meta.db_table),
@@ -65,25 +73,13 @@ def _create_fk_sql(self, model, field, suffix):
     def execute(self, sql, params=()):
         # Hack: Citus will throw the following error if these statements are
         # not executed separately: "ERROR: cannot execute multiple utility events"
-        if not params:
+        if sql and not params:
             for statement in str(sql).split(';'):
                 if statement and not statement.isspace():
                     super(DatabaseSchemaEditor, self).execute(statement)
-        else:
+        elif sql:
             super(DatabaseSchemaEditor, self).execute(sql, params)
 
-    @staticmethod
-    def get_tenant_column(model):
-        app_label = model._meta.app_label
-        model_name = model._meta.model_name
-        for candidate_model in apps.get_models():
-            if (candidate_model._meta.app_label == app_label
-                and candidate_model._meta.model_name == model_name):
-                return candidate_model.tenant_id
-        else:
-            # here you can do fallback logic if no model with db_table found
-            raise ValueError('Model {}.{} not found!'.format(app_label, model_name))
-
 
     def _create_index_name(self, model, column_names, suffix=""):
         # compat with django 2.X and django 1.X

django_multitenant/tests/migrations/0010_auto_20190517_1514.py

@@ -0,0 +1,32 @@
+# -*- coding: utf-8 -*-
+# Generated by Django 1.11 on 2019-05-17 15:14
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+import django.db.models.deletion
+import django_multitenant.fields
+import django_multitenant.mixins
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('tests', '0009_auto_20190412_0839'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='TempModel',
+            fields=[
+                ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('name', models.CharField(max_length=255)),
+                ('account', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='tests.Account', db_constraint=False)),
+            ],
+            options={
+                'abstract': False,
+            },
+            bases=(django_multitenant.mixins.TenantModelMixin, models.Model),
+        ),
+
+
+    ]

django_multitenant/tests/migrations/0011_distribute_new_table.py

@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+# Generated by Django 1.11 on 2019-05-17 15:14
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+import django.db.models.deletion
+import django_multitenant.fields
+import django_multitenant.mixins
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('tests', '0010_auto_20190517_1514'),
+    ]
+
+    operations = [
+        migrations.RunSQL("ALTER TABLE tests_tempmodel DROP CONSTRAINT tests_tempmodel_pkey CASCADE;"),
+        migrations.RunSQL("ALTER TABLE tests_tempmodel ADD CONSTRAINT tests_tempmodel_pkey PRIMARY KEY (account_id, id);"),
+        migrations.RunSQL("SELECT create_distributed_table('tests_tempmodel', 'account_id');"),
+
+    ]

django_multitenant/tests/migrations/0012_auto_20190517_1606.py

@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+# Generated by Django 1.11 on 2019-05-17 16:06
+from __future__ import unicode_literals
+
+from django.db import migrations, models
+import django.db.models.deletion
+import django_multitenant.fields
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('tests', '0011_distribute_new_table'),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name='projectmanager',
+            name='project',
+            field=django_multitenant.fields.TenantForeignKey(on_delete=django.db.models.deletion.CASCADE, to='tests.Project'),
+        ),
+
+        migrations.AddField(
+            model_name='tempmodel',
+            name='project',
+            field=django_multitenant.fields.TenantForeignKey(null=True, on_delete=django.db.models.deletion.CASCADE, to='tests.Project'),
+        ),
+        migrations.AlterField(
+            model_name='tempmodel',
+            name='account',
+            field=models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, to='tests.Account'),
+        ),
+    ]

django_multitenant/tests/migrations/0013_auto_20190517_1607.py

@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+# Generated by Django 1.11 on 2019-05-17 16:07
+from __future__ import unicode_literals
+
+from django.db import migrations
+import django.db.models.deletion
+import django_multitenant.fields
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('tests', '0012_auto_20190517_1606'),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name='tempmodel',
+            name='project',
+        ),
+        migrations.DeleteModel(
+            name='TempModel',
+        ),
+    ]