superworker.py

from getpass import getpass

from django.core.management.base import BaseCommand
from django.db import IntegrityError

from web.models.constants import GLOBAL_STUDY_ID
from web.models.worker_study_role import ROLE_CHOICES_TECHNICIAN
from ...models import Worker, User, Location, Language, WorkerStudyRole


class Command(BaseCommand):
    help = 'creates super worker (superuser + worker)'

    def add_arguments(self, parser):
        parser.add_argument('-u', '--username', type=str, required=True)
        parser.add_argument('-e', '--email', type=str, required=True)
        parser.add_argument('-f', '--first-name', type=str, required=True)
        parser.add_argument('-l', '--last-name', type=str, required=True)

    def handle(self, *args, **kwargs):
        first_name = kwargs['first_name']
        last_name = kwargs['last_name']
        email = kwargs['email']
        username = kwargs['username']

        password = getpass()
        user = None
        try:
            user = User.objects.create(username=username, email=email)
            user.is_superuser = True
            user.is_staff = True
            user.is_admin = True
            user.set_password(password)
            user.save()
        except IntegrityError:
            self.stderr.write('User already exists')
            return

        try:
            worker = Worker.objects.create(first_name=first_name, last_name=last_name, email=email, user=user)
            locations = Location.objects.all()
            worker.locations.set(locations)
            languages = Language.objects.all()
            worker.languages.set(languages)
            worker.save()
            WorkerStudyRole.objects.update_or_create(worker=worker,
                                                     study_id=GLOBAL_STUDY_ID,
                                                     name=ROLE_CHOICES_TECHNICIAN)
        except IntegrityError:
            self.stderr.write('Worker already exists')
            return

        self.stderr.write(f'Superworker {username} created')