PyRIGS/training/management/commands/import_old_db.py

import os
import datetime
import re
import xml.etree.ElementTree as ET
from django.core.management.base import BaseCommand
from django.conf import settings
from django.db.utils import IntegrityError
from django.utils.timezone import make_aware

from training import models
from RIGS.models import Profile


class Command(BaseCommand):
    epoch = datetime.date(1970, 1, 1)
    id_map = {}

    def handle(self, *args, **options):
        self.import_Trainees()
        self.import_TrainingCatagory()
        self.import_TrainingItem()
        self.import_TrainingItemQualification()
        self.import_TrainingLevel()
        self.import_TrainingLevelQualification()
        self.import_TrainingLevelRequirements()

    @staticmethod
    def xml_path(file):
        return os.path.join(settings.BASE_DIR, "data/{}".format(file))

    @staticmethod
    def parse_xml(file):
        tree = ET.parse(file)

        return tree.getroot()

    def import_Trainees(self):
        tally = [0, 0]

        root = self.parse_xml(self.xml_path("Members.xml"))

        for child in root:
            try:
                name = child.find("Member_x0020_Name").text
                first_name = name.split()[0]
                last_name = " ".join(name.split()[1:])
                profile = Profile.objects.filter(first_name=first_name, last_name=last_name).first()

                if profile:
                    self.id_map[child.find("ID").text] = profile.pk
                    print(f"Found existing user {profile}, matching data")
                    tally[0] += 1
                else:
                    # PYTHONIC, BABY
                    initials = first_name[0] + "".join(
                        [name_section[0] for name_section in re.split("\\s*-", last_name.replace("(", ""))]
                    )
                    # print(initials)
                    new_profile = Profile.objects.create(
                        username=name.replace(" ", ""), first_name=first_name, last_name=last_name, initials=initials
                    )
                    self.id_map[child.find("ID").text] = new_profile.pk
                    tally[1] += 1
                    print(f"No match found, creating new user {new_profile}")
            except AttributeError:  # W.T.F
                print("Trainee #{} is FUBAR".format(child.find("ID").text))

        print("Trainees - Updated: {}, Created: {}".format(tally[0], tally[1]))

    def import_TrainingCatagory(self):
        tally = [0, 0]

        root = self.parse_xml(self.xml_path("Categories.xml"))

        for child in root:
            obj, created = models.TrainingCategory.objects.update_or_create(
                pk=int(child.find("ID").text),
                reference_number=int(child.find("Category_x0020_Number").text),
                name=child.find("Category_x0020_Name").text,
            )

            if created:
                tally[1] += 1
            else:
                tally[0] += 1

        print("Categories - Updated: {}, Created: {}".format(tally[0], tally[1]))

    def import_TrainingItem(self):
        tally = [0, 0]

        root = self.parse_xml(self.xml_path("Training Items.xml"))

        for child in root:
            if child.find("active").text == "0":
                active = False
            else:
                active = True

            number = int(child.find("Item_x0020_Number").text)
            name = child.find("Item_x0020_Name").text
            category = models.TrainingCategory.objects.get(pk=int(child.find("Category_x0020_ID").text))

            try:
                obj, created = models.TrainingItem.objects.update_or_create(
                    pk=int(child.find("ID").text),
                    reference_number=number,
                    description=name,
                    category=category,
                    active=active,
                )
            except IntegrityError:
                print(f"Training Item {category.reference_number}.{number} {name} has a duplicate reference number")

            if created:
                tally[1] += 1
            else:
                tally[0] += 1

        print(f"Training Items - Updated: {tally[0]}, Created: {tally[1]}")

    def import_TrainingItemQualification(self):
        tally = [0, 0, 0]

        root = self.parse_xml(self.xml_path("Training Records.xml"))

        for child in root:
            depths = [
                ("Training_Started", models.TrainingItemQualification.STARTED),
                ("Training_Complete", models.TrainingItemQualification.COMPLETE),
                ("Competency_Assessed", models.TrainingItemQualification.PASSED_OUT),
            ]

            for depth, depth_index in depths:
                if child.find(f"{depth}_Date") is not None:
                    if child.find(f"{depth}_Assessor_ID") is None:
                        print(f"Training Record #{child.find('ID').text} had no supervisor. Assigning System User.")
                        supervisor = Profile.objects.get(first_name="God")
                        continue
                    supervisor = Profile.objects.get(pk=self.id_map[child.find("{}_Assessor_ID".format(depth)).text])
                    if child.find("Member_ID") is None:
                        print(
                            "Training Record #{} didn't train anybody and has been ignored. Dammit {}".format(
                                child.find("ID").text, supervisor.name
                            )
                        )
                        tally[2] += 1
                        continue
                    try:
                        obj, created = models.TrainingItemQualification.objects.update_or_create(
                            item=models.TrainingItem.objects.get(pk=int(child.find("Training_Item_ID").text)),
                            trainee=Profile.objects.get(pk=self.id_map[child.find("Member_ID").text]),
                            depth=depth_index,
                            date=child.find("{}_Date".format(depth)).text[
                                :-9
                            ],  # Stored as datetime with time as midnight because fuck you I guess
                            supervisor=supervisor,
                        )
                        notes = child.find("{}_Notes".format(depth))
                        if notes is not None:
                            obj.notes = notes.text
                            obj.save()
                        if created:
                            tally[1] += 1
                        else:
                            tally[0] += 1
                    except IntegrityError:  # Eh?
                        print("Training Record #{} is probably duplicate. ಠ_ಠ".format(child.find("ID").text))
                    except AttributeError:
                        print(child.find("ID").text)

        print(
            "Training Item Qualifications - Updated: {}, Created: {}, Broken: {}".format(tally[0], tally[1], tally[2])
        )

    def import_TrainingLevel(self):
        tally = [0, 0]

        root = self.parse_xml(self.xml_path("Training Levels.xml"))

        for child in root:
            name = child.find("Level_x0020_Name").text
            if name == "Technical Assistant":
                level = models.TrainingLevel.TA
                depString = None
            elif "Common" in name:
                levelString = name.split()[0]
                if levelString == "Technician":
                    level = models.TrainingLevel.TECHNICIAN
                elif levelString == "Supervisor":
                    level = models.TrainingLevel.SUPERVISOR
                depString = None
            else:
                depString = name.split()[-1]
                levelString = name.split()[0]
                if levelString == "Technician":
                    level = models.TrainingLevel.TECHNICIAN
                elif levelString == "Supervisor":
                    level = models.TrainingLevel.SUPERVISOR
                else:
                    print(levelString)
                    continue
                for dep in models.TrainingLevel.DEPARTMENTS:
                    if dep[1] == depString:
                        department = dep[0]

            desc = ""
            if child.find("Desc") is not None:
                desc = child.find("Desc").text

            obj, created = models.TrainingLevel.objects.update_or_create(
                pk=int(child.find("ID").text), description=desc, level=level
            )
            if depString is not None:
                obj.department = department
                obj.save()

            if created:
                tally[1] += 1
            else:
                tally[0] += 1

        for level in models.TrainingLevel.objects.all():
            if level.department is not None:
                if level.level == models.TrainingLevel.TECHNICIAN:
                    level.prerequisite_levels.add(
                        models.TrainingLevel.objects.get(level=models.TrainingLevel.TA),
                        models.TrainingLevel.objects.get(level=models.TrainingLevel.TECHNICIAN, department=None),
                    )
                elif level.level == models.TrainingLevel.SUPERVISOR:
                    level.prerequisite_levels.add(
                        models.TrainingLevel.objects.get(
                            level=models.TrainingLevel.TECHNICIAN, department=level.department
                        ),
                        models.TrainingLevel.objects.get(level=models.TrainingLevel.SUPERVISOR, department=None),
                    )

        print("Training Levels - Updated: {}, Created: {}".format(tally[0], tally[1]))

    def import_TrainingLevelQualification(self):
        tally = [0, 0]

        root = self.parse_xml(self.xml_path("Training Level Records.xml"))

        for child in root:
            try:
                trainee = (
                    Profile.objects.get(pk=self.id_map[child.find("Member_x0020_ID").text])
                    if child.find("Member_x0020_ID") is not None
                    else False
                )
                level = (
                    models.TrainingLevel.objects.get(pk=int(child.find("Training_x0020_Level_x0020_ID").text))
                    if child.find("Training_x0020_Level_x0020_ID") is not None
                    else False
                )

                if trainee and level:
                    obj, created = models.TrainingLevelQualification.objects.update_or_create(
                        pk=int(child.find("ID").text), trainee=trainee, level=level
                    )
                else:
                    print(
                        "Training Level Qualification #{} failed to import. Trainee: {} and Level: {}".format(
                            child.find("ID").text, trainee, level
                        )
                    )
                    continue

                if child.find("Date_x0020_Level_x0020_Awarded") is not None:
                    obj.confirmed_on = make_aware(
                        datetime.datetime.strptime(
                            child.find("Date_x0020_Level_x0020_Awarded").text.split("T")[0], "%Y-%m-%d"
                        )
                    )
                    obj.save()
                    # confirmed by?

                if created:
                    tally[1] += 1
                else:
                    tally[0] += 1
            except IntegrityError:  # Eh?
                print("Training Level Qualification #{} is duplicate. ಠ_ಠ".format(child.find("ID").text))

        print("TrainingLevelQualifications - Updated: {}, Created: {}".format(tally[0], tally[1]))

    def import_TrainingLevelRequirements(self):
        tally = [0, 0]

        root = self.parse_xml(self.xml_path("Training Level Requirements.xml"))

        for child in root:
            items = child.find("Items").text.split(",")
            for item in items:
                try:
                    item = item.split(".")
                    obj, created = models.TrainingLevelRequirement.objects.update_or_create(
                        level=models.TrainingLevel.objects.get(pk=int(child.find("Level").text)),
                        item=models.TrainingItem.objects.get(
                            active=True,
                            reference_number=item[1],
                            category=models.TrainingCategory.objects.get(reference_number=item[0]),
                        ),
                        depth=int(child.find("Depth").text),
                    )

                    if created:
                        tally[1] += 1
                    else:
                        tally[0] += 1
                except models.TrainingItem.DoesNotExist:
                    print("Item with number {} does not exist".format(item))
                except models.TrainingItem.MultipleObjectsReturned:
                    print(
                        models.TrainingItem.objects.filter(
                            reference_number=item[1],
                            category=models.TrainingCategory.objects.get(reference_number=item[0]),
                        )
                    )

        print("TrainingLevelRequirements - Updated: {}, Created: {}".format(tally[0], tally[1]))