|
| 1 | +# Copyright (C) 2021 Magenta ApS, https://magenta.dk. |
| 2 | +# Contact: info@magenta.dk. |
| 3 | + |
| 4 | +from django.core.management.base import ( |
| 5 | + BaseCommand, |
| 6 | + CommandError, |
| 7 | +) |
| 8 | +from django.db import transaction |
| 9 | +from django.contrib.auth import get_user_model |
| 10 | +from system.models import ( |
| 11 | + Site, |
| 12 | + Script, |
| 13 | + Batch, |
| 14 | + Job, |
| 15 | + PCGroup, |
| 16 | + PC, |
| 17 | +) |
| 18 | + |
| 19 | +User = get_user_model() |
| 20 | + |
| 21 | + |
| 22 | +class Command(BaseCommand): |
| 23 | + """ |
| 24 | + Run a script on pcs, group or site as maintenance jobs. |
| 25 | +
|
| 26 | + Notes: |
| 27 | + Only scripts without arguments are supported for now. |
| 28 | + Sites and Groups arguments take UID's while the PC argument takes ID's. |
| 29 | + The script does not validate that a given group, site or PC exists. In case one doesn't match, it's silently ignored. |
| 30 | +
|
| 31 | + Form: |
| 32 | + $ python manage.py run_maintenance_script <user_username> <script_uid_to_run> <batch_site_uid> {--pcs <target_pc_ids...> | --groups <target_group_uids...>| --sites <target_site_uids...>} |
| 33 | + Examples: |
| 34 | +
|
| 35 | + $ python manage.py run_maintenance_script shg 1 magenta --pcs 1 4 |
| 36 | + $ python manage.py run_maintenance_script jabbi 3 magenta-test --groups group1 |
| 37 | + $ python manage.py run_maintenance_script gitte 2 magenta --sites magenta mag test |
| 38 | +
|
| 39 | +
|
| 40 | + """ |
| 41 | + |
| 42 | + help = "Run a maintenance job script on pcs, groups or sites" |
| 43 | + |
| 44 | + def add_arguments(self, parser): |
| 45 | + parser.add_argument( |
| 46 | + "username", nargs="?", type=str, help="username of the user" |
| 47 | + ) |
| 48 | + parser.add_argument("script", nargs="?", type=int, help="the id of the script") |
| 49 | + parser.add_argument( |
| 50 | + "script_site", nargs="?", type=str, help="the uid of the site" |
| 51 | + ) |
| 52 | + parser.add_argument("--pcs", nargs="*", type=int, help="a list of pc uids") |
| 53 | + parser.add_argument("--sites", nargs="*", type=str, help="the uid of a site") |
| 54 | + parser.add_argument("--groups", nargs="*", type=str, help="the uid of a group") |
| 55 | + |
| 56 | + @transaction.atomic |
| 57 | + def handle(self, *args, **options): |
| 58 | + script_id = options["script"] |
| 59 | + script_site_uid = options["script_site"] |
| 60 | + username = options["username"] |
| 61 | + |
| 62 | + user = User.objects.filter(is_superuser=True, username=username).first() |
| 63 | + if not user: |
| 64 | + raise CommandError(f"User with username: {username} does not exist") |
| 65 | + |
| 66 | + script = Script.objects.filter(id=script_id).first() |
| 67 | + if not script: |
| 68 | + raise CommandError(f"Script with ID: {script_id} does not exist") |
| 69 | + |
| 70 | + script_site = Site.objects.filter(uid=script_site_uid).first() |
| 71 | + if not script_site: |
| 72 | + raise CommandError(f"Site with UID: {script_site_uid} does not exist") |
| 73 | + |
| 74 | + if options["pcs"]: |
| 75 | + pcs = PC.objects.filter(id__in=options["pcs"]) |
| 76 | + elif options["sites"]: |
| 77 | + site_uids = options["sites"] |
| 78 | + sites = Site.objects.filter(uid__in=site_uids) |
| 79 | + pcs = PC.objects.filter(site__in=sites) |
| 80 | + elif options["groups"]: |
| 81 | + group_ids = options["groups"] |
| 82 | + pc_groups = PCGroup.objects.filter(uid__in=group_ids) |
| 83 | + pcs = PC.objects.filter(pc_groups__in=pc_groups) |
| 84 | + else: |
| 85 | + raise CommandError("--pcs, --site or --group needs to be given") |
| 86 | + |
| 87 | + self.stdout.write( |
| 88 | + self.style.SUCCESS( |
| 89 | + f"Do you want to run {script} on site {script_site}" |
| 90 | + f" as user {username} for {pcs.count()} PCs? (Y/y for yes)" |
| 91 | + ) |
| 92 | + ) |
| 93 | + confirmation = input() |
| 94 | + if confirmation in ["y", "Y"]: |
| 95 | + batch = Batch.objects.create(site=script_site, script=script, name="") |
| 96 | + for pc in pcs: |
| 97 | + Job.objects.create(user=user, batch=batch, pc=pc) |
| 98 | + self.stdout.write( |
| 99 | + self.style.SUCCESS( |
| 100 | + f"Maintenance job was created for pc: {pc}" |
| 101 | + f" for site: {script_site}" |
| 102 | + ) |
| 103 | + ) |
| 104 | + else: |
| 105 | + self.stdout.write(self.style.WARNING("Aborting")) |
0 commit comments