From a335ecd6239b960cc44c6d9606c61f1d87649cf7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=BC=D0=B8=D1=82=D1=80=D0=B8=D0=B9=20=D0=A7=D0=B5?= <39742182+Dmi4er4@users.noreply.github.com> Date: Fri, 21 Jun 2024 16:18:15 +0300 Subject: [PATCH] Script fix (#850) * script fix * script fix * script fix --- .../commands/create_delete_role_group.py | 48 ++++++++++++------- 1 file changed, 32 insertions(+), 16 deletions(-) diff --git a/apps/admission/management/commands/create_delete_role_group.py b/apps/admission/management/commands/create_delete_role_group.py index 6a1033bc6..72167045d 100644 --- a/apps/admission/management/commands/create_delete_role_group.py +++ b/apps/admission/management/commands/create_delete_role_group.py @@ -3,18 +3,17 @@ from django.db import transaction from django.core.management import BaseCommand, CommandError -from admission.management.commands._utils import CurrentCampaignMixin from core.models import Branch from users.constants import Roles -from users.models import User +from users.models import User, UserGroup from django.conf import settings -class Command(CurrentCampaignMixin, BaseCommand): +class Command(BaseCommand): help = """ Give or take back interviewer role from Users in csv Example of usage: - ./manage.py create_delete_role_group --filename=interviewers.csv --default_branch=msk --role=INTERVIEWER + ./manage.py create_delete_role_group --filename=interviewers.csv --branch=msk --role=INTERVIEWER """ def add_arguments(self, parser): @@ -32,9 +31,9 @@ def add_arguments(self, parser): help="csv delimiter", ) parser.add_argument( - "--default_branch", + "--branch", type=str, - help="Default branch set if user doesn't have one", + help="Branch used for role", ) parser.add_argument( "--take_back", @@ -50,33 +49,50 @@ def add_arguments(self, parser): help="Role to give or take back", ) + def get_group_or_none(self, user, **kwargs): + try: + return user.groups.get(**kwargs) + except UserGroup.DoesNotExist: + return None + def handle(self, *args, **options): delimiter = options["delimiter"] filename = options["filename"] take_back = options["take_back"] - default_branch = options["default_branch"] - role = options["role"] + branch = options["branch"] + role = getattr(Roles, options["role"]) available = Branch.objects.filter( active=True, site_id=settings.SITE_ID ) cs = [c.code for c in available] - if not default_branch or default_branch not in cs: - msg = f"Provide the code of the branch with --default_branch. Options: {cs}" + if not branch or branch not in cs: + msg = f"Provide the code of the branch with --branch. Options: {cs}" raise CommandError(msg) - default_branch = Branch.objects.get(code=default_branch) + branch = Branch.objects.get(code=branch) with open(filename) as csvfile: reader = csv.reader(csvfile, delimiter=delimiter) with transaction.atomic(): headers = next(reader) for row in reader: user: User = User.objects.get(email__iexact=row[0]) - branch = user.branch - if not branch: - self.stdout.write(self.style.WARNING(f"{user} doesn't have branch. Using default one")) - branch = default_branch - role = getattr(Roles, role) + if take_back: + if self.get_group_or_none(user, role=role, branch=branch, site_id=settings.SITE_ID) is None: + self.stdout.write( + self.style.WARNING(f'{user} does not has group with this role and branch')) + continue + user.remove_group(role, branch=branch) else: + if self.get_group_or_none(user, role=role, branch__isnull=True, site_id=settings.SITE_ID): + self.stdout.write( + self.style.WARNING(f'{user} already has group with this role and with no branch')) + continue + + if self.get_group_or_none(user, role=role, branch=branch, site_id=settings.SITE_ID): + self.stdout.write( + self.style.WARNING(f'{user} already has group with this role and branch')) + continue + user.add_group(role, branch=branch)