(demo) add a "demo" app to facilitate testing/working on the project

We designed it to allow creating a huge number of objects fast using
bulk creation.
This commit is contained in:
Samuel Paccoud - DINUM
2024-01-15 09:18:57 +01:00
committed by Anthony LC
parent cfc35ac23e
commit 0c550ebd1c
9 changed files with 280 additions and 0 deletions

View File

View File

@@ -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,
}

View File

View File

@@ -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)

View File

@@ -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))

View File

View File

@@ -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

View File

@@ -185,6 +185,7 @@ class Base(Configuration):
INSTALLED_APPS = [
# People
"core",
"demo",
"drf_spectacular",
# Third party apps
"corsheaders",