Initial commit

env/lib/python3.10/site-packages/wagtail/management/commands/fixtree.py

@@ -0,0 +1,165 @@
+import functools
+import operator
+
+from django.core.management.base import BaseCommand
+from django.db import models
+from django.db.models import Q
+
+from wagtail.models import Collection, Page
+
+
+class Command(BaseCommand):
+    help = "Checks for data integrity errors on the page tree, and fixes them where possible."
+    stealth_options = ("delete_orphans",)
+
+    def add_arguments(self, parser):
+        parser.add_argument(
+            "--noinput",
+            action="store_false",
+            dest="interactive",
+            default=True,
+            help="If provided, any fixes requiring user interaction will be skipped.",
+        )
+        parser.add_argument(
+            "--full",
+            action="store_true",
+            dest="full",
+            default=False,
+            help="If provided, uses a more thorough but slower method that also fixes path ordering issues.",
+        )
+
+    def numberlist_to_string(self, numberlist):
+        # Converts a list of numbers into a string
+        # Doesn't put "L" after longs
+        return "[" + ", ".join(map(str, numberlist)) + "]"
+
+    def handle(self, **options):
+        any_page_problems_fixed = False
+        for page in Page.objects.all():
+            try:
+                page.specific
+            except page.specific_class.DoesNotExist:
+                self.stdout.write(
+                    "Page %d (%s) is missing a subclass record; deleting."
+                    % (page.id, page.title)
+                )
+                any_page_problems_fixed = True
+                page.delete()
+
+        self.handle_model(Page, "page", "pages", any_page_problems_fixed, options)
+        self.handle_model(Collection, "collection", "collections", False, options)
+
+    def handle_model(
+        self, model, model_name, model_name_plural, any_problems_fixed, options
+    ):
+        fix_paths = options.get("full", False)
+
+        self.stdout.write("Checking %s tree for problems..." % model_name)
+        (bad_alpha, bad_path, orphans, bad_depth, bad_numchild) = model.find_problems()
+
+        if bad_depth:
+            self.stdout.write(
+                "Incorrect depth value found for %s: %s"
+                % (model_name_plural, self.numberlist_to_string(bad_depth))
+            )
+        if bad_numchild:
+            self.stdout.write(
+                "Incorrect numchild value found for %s: %s"
+                % (model_name_plural, self.numberlist_to_string(bad_numchild))
+            )
+
+        if orphans:
+            # The 'orphans' list as returned by treebeard only includes nodes that are
+            # missing an immediate parent; descendants of orphans are not included.
+            # Deleting only the *actual* orphans is a bit silly (since it'll just create
+            # more orphans), so generate a queryset that contains descendants as well.
+            orphan_paths = model.objects.filter(id__in=orphans).values_list(
+                "path", flat=True
+            )
+            filter_conditions = []
+            for path in orphan_paths:
+                filter_conditions.append(Q(path__startswith=path))
+
+            # combine filter_conditions into a single ORed condition
+            final_filter = functools.reduce(operator.or_, filter_conditions)
+
+            # build a queryset of all nodes to be removed; this must be a vanilla Django
+            # queryset rather than a treebeard MP_NodeQuerySet, so that we bypass treebeard's
+            # custom delete() logic that would trip up on the very same corruption that we're
+            # trying to fix here.
+            nodes_to_delete = models.query.QuerySet(model).filter(final_filter)
+
+            self.stdout.write("Orphaned %s found:" % model_name_plural)
+            for node in nodes_to_delete:
+                self.stdout.write("ID %d: %s" % (node.id, node))
+            self.stdout.write("")
+
+            if options.get("interactive", True):
+                yes_or_no = input("Delete these %s? [y/N] " % model_name_plural)
+                delete_orphans = yes_or_no.lower().startswith("y")
+                self.stdout.write("")
+            else:
+                # Running tests, check for the "delete_orphans" option
+                delete_orphans = options.get("delete_orphans", False)
+
+            if delete_orphans:
+                deletion_count = len(nodes_to_delete)
+                nodes_to_delete.delete()
+                self.stdout.write(
+                    "%d orphaned %s deleted."
+                    % (
+                        deletion_count,
+                        model_name_plural if deletion_count != 1 else model_name,
+                    )
+                )
+                any_problems_fixed = True
+
+        # fix_paths will fix problems not identified by find_problems, so if that option has been
+        # passed, run it regardless (and set any_problems_fixed=True, since we don't have a way to
+        # test whether anything was actually fixed in that process)
+        if bad_depth or bad_numchild or fix_paths:
+            model.fix_tree(destructive=False, fix_paths=fix_paths)
+            any_problems_fixed = True
+
+        if any_problems_fixed:
+            # re-run find_problems to see if any new ones have surfaced
+            (
+                bad_alpha,
+                bad_path,
+                orphans,
+                bad_depth,
+                bad_numchild,
+            ) = model.find_problems()
+
+        if any((bad_alpha, bad_path, orphans, bad_depth, bad_numchild)):
+            self.stdout.write("Remaining problems (cannot fix automatically):")
+            if bad_alpha:
+                self.stdout.write(
+                    "Invalid characters found in path for %s: %s"
+                    % (model_name_plural, self.numberlist_to_string(bad_alpha))
+                )
+            if bad_path:
+                self.stdout.write(
+                    "Invalid path length found for %s: %s"
+                    % (model_name_plural, self.numberlist_to_string(bad_path))
+                )
+            if orphans:
+                self.stdout.write(
+                    "Orphaned %s found: %s"
+                    % (model_name_plural, self.numberlist_to_string(orphans))
+                )
+            if bad_depth:
+                self.stdout.write(
+                    "Incorrect depth value found for %s: %s"
+                    % (model_name_plural, self.numberlist_to_string(bad_depth))
+                )
+            if bad_numchild:
+                self.stdout.write(
+                    "Incorrect numchild value found for %s: %s"
+                    % (model_name_plural, self.numberlist_to_string(bad_numchild))
+                )
+
+        elif any_problems_fixed:
+            self.stdout.write("All problems fixed.\n\n")
+        else:
+            self.stdout.write("No problems found.\n\n")