Merge pull request #1684 from SchrodingersGat/rebuild-command

Adds a management command to rebuild MPTT models

InvenTree/InvenTree/management/commands/rebuild_models.py

@@ -0,0 +1,60 @@
+"""
+Custom management command to rebuild all MPTT models
+
+- This is crucial after importing any fixtures, etc
+"""
+
+from django.core.management.base import BaseCommand
+
+
+class Command(BaseCommand):
+    """
+    Rebuild all database models which leverage the MPTT structure.
+    """
+
+    def handle(self, *args, **kwargs):
+
+        # Part model
+        try:
+            print("Rebuilding Part objects")
+
+            from part.models import Part
+            Part.objects.rebuild()
+        except:
+            print("Error rebuilding Part objects")
+
+        # Part category
+        try:
+            print("Rebuilding PartCategory objects")
+
+            from part.models import PartCategory
+            PartCategory.objects.rebuild()
+        except:
+            print("Error rebuilding PartCategory objects")
+
+        # StockItem model
+        try:
+            print("Rebuilding StockItem objects")
+
+            from stock.models import StockItem
+            StockItem.objects.rebuild()
+        except:
+            print("Error rebuilding StockItem objects")
+
+        # StockLocation model
+        try:
+            print("Rebuilding StockLocation objects")
+
+            from stock.models import StockLocation
+            StockLocation.objects.rebuild()
+        except:
+            print("Error rebuilding StockLocation objects")
+
+        # Build model
+        try:
+            print("Rebuilding Build objects")
+
+            from build.models import Build
+            Build.objects.rebuild()
+        except:
+            print("Error rebuilding Build objects")

tasks.py

@@ -129,6 +129,14 @@ def wait(c):
 
     manage(c, "wait_for_db")
 
+@task
+def rebuild(c):
+    """
+    Rebuild database models with MPTT structures
+    """
+
+    manage(c, "rebuild_models")
+
 @task
 def migrate(c):
     """
@@ -311,7 +319,7 @@ def export_records(c, filename='data.json'):
     print("Data export completed")
 
 
-@task(help={'filename': 'Input filename'})
+@task(help={'filename': 'Input filename'}, post=[rebuild])
 def import_records(c, filename='data.json'):
     """
     Import database records from a file
@@ -354,7 +362,7 @@ def import_records(c, filename='data.json'):
 
     print("Data import completed")
 
-@task
+@task(post=[rebuild])
 def import_fixtures(c):
     """
     Import fixture data into the database.