From 0c550ebd1cf26a859fd20bddc26fb99fa9171ceb Mon Sep 17 00:00:00 2001 From: Samuel Paccoud - DINUM Date: Mon, 15 Jan 2024 09:18:57 +0100 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8(demo)=20add=20a=20"demo"=20app=20to?= =?UTF-8?q?=20facilitate=20testing/working=20on=20the=20project?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit We designed it to allow creating a huge number of objects fast using bulk creation. --- src/backend/demo/__init__.py | 0 src/backend/demo/defaults.py | 8 + src/backend/demo/management/__init__.py | 0 .../demo/management/commands/__init__.py | 0 .../demo/management/commands/create_demo.py | 190 ++++++++++++++++++ .../management/commands/createsuperuser.py | 49 +++++ src/backend/demo/tests/__init__.py | 0 .../demo/tests/test_commands_create_demo.py | 32 +++ src/backend/people/settings.py | 1 + 9 files changed, 280 insertions(+) create mode 100644 src/backend/demo/__init__.py create mode 100644 src/backend/demo/defaults.py create mode 100644 src/backend/demo/management/__init__.py create mode 100644 src/backend/demo/management/commands/__init__.py create mode 100755 src/backend/demo/management/commands/create_demo.py create mode 100644 src/backend/demo/management/commands/createsuperuser.py create mode 100644 src/backend/demo/tests/__init__.py create mode 100644 src/backend/demo/tests/test_commands_create_demo.py diff --git a/src/backend/demo/__init__.py b/src/backend/demo/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/backend/demo/defaults.py b/src/backend/demo/defaults.py new file mode 100644 index 0000000..a43e253 --- /dev/null +++ b/src/backend/demo/defaults.py @@ -0,0 +1,8 @@ +"""Parameters that define how the demo site will be built.""" + +NB_OBJECTS = { + "users": 1000, + "teams": 100, + "max_identities_per_user": 3, + "max_users_per_team": 100, +} diff --git a/src/backend/demo/management/__init__.py b/src/backend/demo/management/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/backend/demo/management/commands/__init__.py b/src/backend/demo/management/commands/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/backend/demo/management/commands/create_demo.py b/src/backend/demo/management/commands/create_demo.py new file mode 100755 index 0000000..a6d5cab --- /dev/null +++ b/src/backend/demo/management/commands/create_demo.py @@ -0,0 +1,190 @@ +# ruff: noqa: S311, S106 +"""create_demo management command""" + +import logging +import random +import time +from collections import defaultdict +from uuid import uuid4 + +from django import db +from django.conf import settings +from django.core.management.base import BaseCommand, CommandError + +from core import models + +from demo import defaults + +logger = logging.getLogger("people.commands.demo.create_demo") + + +def random_true_with_probability(probability): + """return True with the requested probability, False otherwise.""" + return random.random() < probability + + +class BulkQueue: + """A utility class to create Django model instances in bulk by just pushing to a queue.""" + + BATCH_SIZE = 20000 + + def __init__(self, stdout, *args, **kwargs): + """Define the queue as a dict of lists.""" + self.queue = defaultdict(list) + self.stdout = stdout + + def _bulk_create(self, objects): + """Actually create instances in bulk in the database.""" + if not objects: + return + + objects[0]._meta.model.objects.bulk_create(objects, ignore_conflicts=True) # noqa: SLF001 + # In debug mode, Django keeps query cache which creates a memory leak in this case + db.reset_queries() + self.queue[objects[0]._meta.model.__name__] = [] # noqa: SLF001 + + def push(self, obj): + """Add a model instance to queue to that it gets created in bulk.""" + objects = self.queue[obj._meta.model.__name__] # noqa: SLF001 + objects.append(obj) + if len(objects) > self.BATCH_SIZE: + self._bulk_create(objects) + self.stdout.write(".", ending="") + + def flush(self): + """Flush the queue after creating the remaining model instances.""" + for objects in self.queue.values(): + self._bulk_create(objects) + + +class Timeit: + """A utility context manager/method decorator to time execution.""" + + total_time = 0 + + def __init__(self, stdout, sentence=None): + """Set the sentence to be displayed for timing information.""" + self.sentence = sentence + self.start = None + self.stdout = stdout + + def __call__(self, func): + """Behavior on call for use as a method decorator.""" + + def timeit_wrapper(*args, **kwargs): + """wrapper to trigger/stop the timer before/after function call.""" + self.__enter__() + result = func(*args, **kwargs) + self.__exit__(None, None, None) + return result + + return timeit_wrapper + + def __enter__(self): + """Start timer upon entering context manager.""" + self.start = time.perf_counter() + if self.sentence: + self.stdout.write(self.sentence, ending=".") + + def __exit__(self, exc_type, exc_value, exc_tb): + """Stop timer and display result upon leaving context manager.""" + if exc_type is not None: + raise exc_type(exc_value) + end = time.perf_counter() + elapsed_time = end - self.start + if self.sentence: + self.stdout.write(f" Took {elapsed_time:g} seconds") + + self.__class__.total_time += elapsed_time + return elapsed_time + + +def create_demo(stdout): + """ + Create a database with demo data for developers to work in a realistic environment. + The code is engineered to create a huge number of objects fast. + """ + + queue = BulkQueue(stdout) + + with Timeit(stdout, "Creating users"): + for i in range(defaults.NB_OBJECTS["users"]): + queue.push( + models.User( + email=f"user{i:d}@example.com", + password="!", + is_superuser=False, + is_active=True, + is_staff=False, + language=random.choice(settings.LANGUAGES)[0], + ) + ) + queue.flush() + + with Timeit(stdout, "Creating identities"): + users_values = list(models.User.objects.values("id", "email")) + for user_dict in users_values: + for i in range( + random.randint(0, defaults.NB_OBJECTS["max_identities_per_user"]) + ): + user_email = user_dict["email"] + queue.push( + models.Identity( + user_id=user_dict["id"], + sub=uuid4(), + email=f"identity{i:d}{user_email:s}", + is_main=(i == 0), + ) + ) + queue.flush() + + with Timeit(stdout, "Creating teams"): + for i in range(defaults.NB_OBJECTS["teams"]): + queue.push( + models.Team( + name=f"Team {i:d}", + ) + ) + queue.flush() + + with Timeit(stdout, "Creating team accesses"): + teams_ids = list(models.Team.objects.values_list("id", flat=True)) + users_ids = list(models.User.objects.values_list("id", flat=True)) + for team_id in teams_ids: + for user_id in random.sample( + users_ids, + random.randint(1, defaults.NB_OBJECTS["max_users_per_team"]), + ): + role = random.choice(models.RoleChoices.choices) + queue.push( + models.TeamAccess(team_id=team_id, user_id=user_id, role=role[0]) + ) + queue.flush() + + +class Command(BaseCommand): + """A management command to create a demo database.""" + + help = __doc__ + + def add_arguments(self, parser): + """Add argument to require forcing execution when not in debug mode.""" + parser.add_argument( + "-f", + "--force", + action="store_true", + default=False, + help="Force command execution despite DEBUG is set to False", + ) + + def handle(self, *args, **options): + """Handling of the management command.""" + if not settings.DEBUG and not options["force"]: + raise CommandError( + ( + "This command is not meant to be used in production environment " + "except you know what you are doing, if so use --force parameter" + ) + ) + + create_demo(self.stdout) diff --git a/src/backend/demo/management/commands/createsuperuser.py b/src/backend/demo/management/commands/createsuperuser.py new file mode 100644 index 0000000..72ae10f --- /dev/null +++ b/src/backend/demo/management/commands/createsuperuser.py @@ -0,0 +1,49 @@ +""" +Management command overrring the "createsuperuser" command to allow creating users +with their email and no username. +""" +from django.contrib.auth import get_user_model +from django.core.management.base import BaseCommand + +User = get_user_model() + + +class Command(BaseCommand): + """Management command to create super users from an email and a password""" + + help = "Create a superuser with an email and a password" + + def add_arguments(self, parser): + """Define required arguments "email" and "password".""" + parser.add_argument( + "--email", + help=("Email for the user."), + ) + parser.add_argument( + "--password", + help="Password for the user.", + ) + + def handle(self, *args, **options): + """ + Given an email and a password, create a superuser or upgrade the existing + user to superuser status. + """ + email = options.get("email") + try: + user = User.objects.get(email=email) + except User.DoesNotExist: + user = User(email=email) + message = "Superuser created successfully." + else: + if user.is_superuser and user.is_staff: + message = "Superuser already exists." + else: + message = "User already existed and was upgraded to superuser." + + user.is_superuser = True + user.is_staff = True + user.set_password(options["password"]) + user.save() + + self.stdout.write(self.style.SUCCESS(message)) diff --git a/src/backend/demo/tests/__init__.py b/src/backend/demo/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/backend/demo/tests/test_commands_create_demo.py b/src/backend/demo/tests/test_commands_create_demo.py new file mode 100644 index 0000000..1bd6bd9 --- /dev/null +++ b/src/backend/demo/tests/test_commands_create_demo.py @@ -0,0 +1,32 @@ +"""Test the `create_demo` management command""" +from unittest import mock + +from django.core.management import call_command +from django.test import override_settings + +import pytest + +from core import models + +from demo import defaults + +TEST_NB_OBJECTS = { + "users": 5, + "teams": 3, + "max_identities_per_user": 3, + "max_users_per_team": 5, +} + +pytestmark = pytest.mark.django_db + + +@override_settings(DEBUG=True) +@mock.patch.dict(defaults.NB_OBJECTS, TEST_NB_OBJECTS) +def test_commands_create_demo(): + """The create_demo management command should create objects as expected.""" + call_command("create_demo") + + assert models.User.objects.count() == 5 + assert models.Identity.objects.exists() + assert models.Team.objects.count() == 3 + assert models.TeamAccess.objects.count() >= 3 diff --git a/src/backend/people/settings.py b/src/backend/people/settings.py index 7dfb4a6..2a44cdc 100755 --- a/src/backend/people/settings.py +++ b/src/backend/people/settings.py @@ -185,6 +185,7 @@ class Base(Configuration): INSTALLED_APPS = [ # People "core", + "demo", "drf_spectacular", # Third party apps "corsheaders",