timestrap-django

 1from django.core.management import call_command
 2from django.core.management.base import BaseCommand, CommandError
 3from django.db import DEFAULT_DB_ALIAS
 4from django.utils.six.moves import input
 5
 6
 7class Command(BaseCommand):
 8    help = (
 9        "Deletes all data from this instance and recreates the original "
10        "site and default admin account."
11    )
12
13    def add_arguments(self, parser):
14        parser.add_argument(
15            "--noinput",
16            "--no-input",
17            action="store_false",
18            dest="interactive",
19            default=True,
20            help="Tells the command to NOT prompt the user for input of any " "kind.",
21        )
22        parser.add_argument(
23            "--database",
24            action="store",
25            dest="database",
26            default=DEFAULT_DB_ALIAS,
27            help='Nominates a database to flush. Defaults to the "default" '
28            "database.",
29        )
30        parser.add_argument(
31            "--fake",
32            dest="iterations",
33            default=5,
34            help=(
35                "Fill the database with fake data after the reset. Provide a "
36                "number a number of iterations to perform (higher number = "
37                "more fake data)."
38            ),
39        )
40
41    def handle(self, *args, **kwargs):
42        interactive = kwargs["interactive"]
43        database = kwargs["database"]
44        verbosity = kwargs["verbosity"]
45        iterations = int(kwargs["iterations"])
46
47        if interactive:
48            confirm = input(
49                """You have requested a reset of the application.
50
51This will IRREVERSIBLY DESTROY all data currently in the
52database, establish the initial site, and create the initial
53admin user.
54
55Are you sure you want to do this?
56
57Type 'yes' to continue, or 'no' to cancel: """
58            )
59        else:
60            confirm = "yes"
61
62        if confirm == "yes":
63            try:
64                call_command(
65                    "flush", database=database, interactive=False, verbosity=verbosity
66                )
67                self.stdout.write(self.style.SUCCESS("Successfully flushed database."))
68            except CommandError:
69                self.stdout.write(
70                    self.style.ERROR("Database flush failed. Reset aborted.")
71                )
72                return
73            call_command("migrate", verbosity=verbosity)
74
75            if iterations > 0:
76                call_command("fake", verbosity=verbosity, iterations=iterations)
77
78            if verbosity > 0:
79                self.stdout.write(self.style.SUCCESS("Reset operation complete."))
80        else:
81            if verbosity > 0:
82                self.stdout.write(self.style.ERROR("Reset operation cancelled."))