nk20-scripts/management/commands/import_activities.py

#!/usr/bin/env python3

import psycopg2 as pg
import psycopg2.extras as pge
import datetime
import copy

from django.utils.timezone import make_aware
from django.db import transaction

from activity.models import ActivityType, Activity, Guest, Entry
from member.models import Club
from note.models import Note
from ._import_utils import ImportCommand, BulkCreateManager, timed

MAP_ACTIVITY = dict()

class Command(ImportCommand):
    """
    Import command for Activities Base Data (Comptes, and Aliases)
    """

    @timed
    @transaction.atomic
    def import_activities(self, cur, chunk):
        cur.execute("SELECT * FROM activites ORDER by id")
        n = cur.rowcount
        bulk_mgr = BulkCreateManager(chunk_size=chunk)
        activity_type_id = ActivityType.objects.get(name="Pot").pk  # Need to be fixed manually
        kfet = Club.objects.get(name="Kfet")
        pk_activity = 1
        for idx, row in enumerate(cur):
            self.update_line(idx, n, row["titre"])
            note = self.MAP_IDBDE[row["responsable"]]
            if note == 6244:
                 # Licorne magique ne doit pas utiliser son compte club pour proposer des activités
                note = Note.objects.get(pk=self.MAP_IDBDE[6524])
                note = note.user_id
            organizer = Club.objects.filter(name=row["signature"])
            if organizer.exists():
                # Try to find the club that organizes the activity.
                # If not found, assume it's Kfet (fix manually)
                organizer = organizer.get()
            else:
                organizer = kfet
            obj_dict = {
                "pk": pk_activity,
                "name": row["titre"],
                "description": row["description"],
                "activity_type_id": activity_type_id,  # By default Pot
                "creater_id": note,
                "organizer_id": organizer.pk,
                "attendees_club_id": kfet.pk,  # Maybe fix manually
                "date_start": make_aware(row["debut"]),
                "date_end": make_aware(row["fin"]),
                "valid": row["validepar"] is not None,
                "open": row["ouvert"],  # Should always be False
            }
            # WARNING: Fields lieu, liste, listeimprimee are missing
            MAP_ACTIVITY[row["id"]] = pk_activity
            pk_activity += 1
            bulk_mgr.add(Activity(**obj_dict))
        bulk_mgr.done()

    @timed
    @transaction.atomic
    def import_guest(self, cur, chunk):
        bulk_mgr = BulkCreateManager(chunk_size=chunk)
        cur.execute("SELECT * FROM invites ORDER by id")
        n = cur.rowcount
        for idx, row in enumerate(cur):
            self.update_line(idx, n, f"{row['nom']} {row['prenom']}")
            obj_dict = {
                "pk": row["id"],
                "activity_id": MAP_ACTIVITY[row["activite"]],
                "last_name": row["nom"],
                "first_name": row["prenom"],
                "inviter_id": self.MAP_IDBDE[row["responsable"]],
            }
            bulk_mgr.add(Guest(**obj_dict))
        bulk_mgr.done()

    @timed
    @transaction.atomic
    def import_activities_entries(self, cur, chunk):
        bulk_mgr = BulkCreateManager(chunk_size=chunk)
        cur.execute("SELECT * FROM entree_activites ORDER by id")
        n = cur.rowcount
        for idx, row in enumerate(cur):
            self.update_line(idx, n, f"{row['idbde']} {row['responsable']}")
            obj_dict = {
                "activity_id": MAP_ACTIVITY[row["activite"]],
                "time": make_aware(row["heure_entree"]),
                "note_id": self.MAP_IDBDE[row["responsable"] if row['est_invite'] else row["idbde"]],
                "guest_id": self.MAP_IDBDE[row["idbde"]] if row['est_invite'] else None,
            }
            bulk_mgr.add(Entry(**obj_dict))
        bulk_mgr.done()

    def handle(self, *args, **kwargs):
        # default args, provided by ImportCommand.
        nk15db, nk15user = kwargs['nk15db'], kwargs['nk15user']
        # connecting to nk15 database
        conn = pg.connect(database=nk15db, user=nk15user)
        cur = conn.cursor(cursor_factory=pge.DictCursor)

        if kwargs["map"]:
            self.load_map(kwargs["map"])
        self.import_activities(cur, kwargs["chunk"])
        self.import_guest(cur, kwargs["chunk"])
        self.import_activities_entries(cur, kwargs["chunk"])