nk20-scripts/management/commands/import_transaction.py

#!/usr/bin/env python3
import re
import psycopg2 as pg
import psycopg2.extras as pge
import pytz
import datetime
import copy

from django.utils.timezone import make_aware
from django.db import transaction
from django.contrib.contenttypes.models import ContentType

from note.models import (TemplateCategory,
                         TransactionTemplate,
                         Transaction,
                         RecurrentTransaction,
                         SpecialTransaction,
                         MembershipTransaction,
                         )
from note.models import Note, NoteClub
from activity.models import Guest, GuestTransaction

from member.models import Membership
from ._import_utils import ImportCommand, BulkCreateManager, timed

# from member/fixtures/initial
BDE_PK = 1
KFET_PK = 2

# from note/fixtures/initial
NOTE_SPECIAL_CODE = {
    "espèce": 1,
    "carte": 2,
    "chèque": 3,
    "virement": 4,
}
# from permission/fixtures/initial
BDE_ROLE_PK = 1
KFET_ROLE_PK = 2

CT = {
    "RecurrentTransaction": ContentType.objects.get(app_label="note", model="recurrenttransaction"),
    "SpecialTransaction": ContentType.objects.get(app_label="note", model="specialtransaction"),
    "MembershipTransaction": ContentType.objects.get(app_label="note", model="membershiptransaction"),
    "GuestTransaction": ContentType.objects.get(app_label="activity", model="guesttransaction"),
}

def get_date_end(date_start):
    date_end = copy.deepcopy(date_start)
    if date_start.month > 8:
        date_end = date_start.replace(year=date_start.year+1)
    date_end = date_end.replace(month=9, day=30)
    return date_end


class Command(ImportCommand):
    """
    Import command for People base data (Comptes, and Aliases)
    """

    def add_arguments(self, parser):
        parser.add_argument('-b', '--buttons', action='store_true', help="import buttons")
        parser.add_argument('-t', '--transactions', action='store', default=0, help="start id for transaction import")

    @timed
    def import_buttons(self, cur, chunk_size):
        self.categories = dict()
        self.buttons = dict()
        bulk_mgr = BulkCreateManager(chunk_size=chunk_size)
        cur.execute("SELECT * FROM boutons;")
        n = cur.rowcount
        for idx, row in enumerate(cur):
            self.update_line(idx, n, row["label"])
            if row["categorie"] not in self.categories:
                cat = TemplateCategory.objects.create(name=row["categorie"])
                cat.save()
                self.categories[row["categorie"]] = cat.pk
            obj_dict = {
                "pk": row["id"],
                "name": row["label"],
                "amount": row["montant"],
                "destination_id": self.MAP_IDBDE[row["destinataire"]],
                "category_id": self.categories[row["categorie"]],
                "display": row["affiche"],
                "description": row["description"],
            }
            if row["label"] in self.buttons:
                obj_dict["name"] = f"{obj_dict['name']}_{obj_dict['destination_id']}"
            bulk_mgr.add(TransactionTemplate(**obj_dict))
            self.buttons[obj_dict["name"]] = (row["id"], self.categories[row["categorie"]])
        bulk_mgr.done()

    def _basic_transaction(self, row, obj_dict, child_dict):
        if len(row["description"]) > 255:
            obj_dict["reason"] = obj_dict["reason"][:250] + "...)"
        return obj_dict, None, None

    def _template_transaction(self, row, obj_dict, child_dict):
        if self.buttons.get(row["description"]):
            child_dict["category_id"] = self.buttons[row["description"]][1]
            child_dict["template_id"] = self.buttons[row["description"]][0]
        # elif self.categories.get(row["categorie"]):
        #     child_dict["category_id"] = self.categories[row["categorie"]]
        elif "WEI" in row["description"]:
            return obj_dict, None, None
        else:
            return obj_dict, None, None
        obj_dict["polymorphic_ctype"] = CT["RecurrentTransaction"]
        return obj_dict, child_dict, RecurrentTransaction

    def _membership_transaction(self, row, obj_dict, child_dict, pk_membership):
        obj_dict["polymorphic_ctype"] = CT["MembershipTransaction"]
        obj_dict2 = obj_dict.copy()
        child_dict2 = child_dict.copy()
        child_dict2["membership_id"] = pk_membership

        return obj_dict2, child_dict2, MembershipTransaction

    def _special_transaction(self, row, obj_dict, child_dict):
        # Some transaction uses BDE (idbde=0) as source or destination,
        # lets fix that.
        obj_dict["polymorphic_ctype"] = CT["SpecialTransaction"]
        field_id = "source_id" if row["type"] == "crédit" else "destination_id"
        if "espèce" in row["description"]:
            obj_dict[field_id] = 1
        elif "carte" in row["description"]:
            obj_dict[field_id] = 2
        elif "cheques" in row["description"]:
            obj_dict[field_id] = 3
        elif "virement" in row["description"]:
            obj_dict[field_id] = 4
        # humans and clubs have always the biggest id
        actor_pk = max(row["destinataire"], row["emetteur"])
        actor = Note.objects.get(id=self.MAP_IDBDE[actor_pk])
        # custom fields of SpecialTransaction
        if actor.__class__.__name__ == "NoteUser":
            child_dict["first_name"] = actor.user.first_name
            child_dict["last_name"] = actor.user.last_name
        else:
            child_dict["first_name"] = actor.club.name
            child_dict["last_name"] = actor.club.name
        return obj_dict, child_dict, SpecialTransaction

    def _guest_transaction(self, row, obj_dict, child_dict):
        # Currently GuestTransaction is related to a Guest.
        # This is not ideal and should be change to the Entry of this Guest.
        obj_dict["polymorphic_ctype"] = CT["GuestTransaction"]
        m = re.search(r"Invitation (.*?)(?:\s\()(.*?)\s(.*?)\)", row["description"])
        if m:
            first_name, last_name = m.group(2), m.group(3)
            guest_id = Guest.objects.filter(first_name__iexact=first_name,
                                            last_name__iexact=last_name).first().pk
            child_dict["guest_id"] = guest_id
        else:
            raise(f"Guest not Found {row['id']} {first_name}, last_name")

        return obj_dict, child_dict, GuestTransaction

    @timed
    @transaction.atomic
    def import_transaction(self, cur, chunk_size, idmin):
        bulk_mgr = BulkCreateManager(chunk_size=chunk_size)
        cur.execute(
            f"SELECT t.date AS transac_date, t.type, t.emetteur,\
              t.destinataire,t.quantite, t.montant, t.description,\
              t.valide, t.cantinvalidate, t.categorie, \
              a.idbde, a.annee, a.wei, a.date AS adh_date, a.section\
            FROM transactions AS t \
            LEFT JOIN adhesions  AS a ON t.id = a.idtransaction \
            WHERE t.id >= {idmin} \
            ORDER BY t.id;")
        n = cur.rowcount
        pk_membership = 1
        pk_transaction = 1
        for idx, row in enumerate(cur):
            self.update_line(idx, n, row["description"])
            try:
                date = make_aware(row["transac_date"])
            except (pytz.NonExistentTimeError, pytz.AmbiguousTimeError):
                date = make_aware(row["transac_date"] + datetime.timedelta(hours=1))

            # standart transaction object
            obj_dict = {
                "pk": pk_transaction,
                "destination_id": self.MAP_IDBDE[row["destinataire"]],
                "polymorphic_ctype": None,
                "source_id": self.MAP_IDBDE[row["emetteur"]],
                "amount": row["montant"],
                "created_at": date,
                "destination_alias": "",
                "invalidity_reason": None,
                "quantity": row["quantite"],
                "reason": row["description"],
                "source_alias": "",
                "valid": row["valide"],
            }
            # for child transaction Models
            child_dict = {"pk": pk_transaction}
            ttype = row["type"]
            # Membership transaction detection and import
            if row["valide"] and (ttype == "adhésion" or row["description"].lower() == "inscription"):
                note = Note.objects.get(pk=obj_dict["source_id"])
                if isinstance(note, NoteClub):
                    child_transaction = None # don't bother register clubs
                else:
                    user_id = note.user_id
                    montant = obj_dict["amount"]
                    (obj_dict0,
                     child_dict0,
                     child_transaction) = self._membership_transaction(row, obj_dict, child_dict, pk_membership)
                    bde_dict = {
                        "pk": pk_membership,
                        "user_id": user_id,
                        "club_id": BDE_PK,
                        "date_start": date.date(),  # Only date, not time
                        "date_end": get_date_end(date.date()),
                        "fee": min(500, montant)
                    }
                    pk_membership += 1
                    pk_transaction += 1
                    obj_dict, child_dict, child_transaction = self._membership_transaction(row, obj_dict, child_dict, pk_membership)
                    # Kfet membership
                    # BDE Membership
                    obj_dict["pk"] = pk_transaction
                    child_dict["pk"] = pk_transaction
                    kfet_dict = {
                        "pk": pk_membership,
                        "user_id": user_id,
                        "club_id": KFET_PK,
                        "date_start": date.date(),  # Only date, not time
                        "date_end": get_date_end(date.date()),
                        "fee": max(montant - 500, 0),
                    }
                    obj_dict0["amount"] = bde_dict["fee"]
                    obj_dict["amount"] = kfet_dict["fee"]
                    # BDE membership Transaction is inserted before the Kfet membershipTransaction
                    pk_membership += 1
                    pk_transaction += 1
                    bulk_mgr.add(
                        Membership(**bde_dict),
                        Membership(**kfet_dict),
                        Transaction(**obj_dict0),
                        child_transaction(**child_dict0),
                        Transaction(**obj_dict),
                        child_transaction(**child_dict),
                    )
                    continue
            elif ttype == "bouton":
                obj_dict, child_dict, child_transaction = self._template_transaction(row, obj_dict, child_dict)
            elif ttype == "crédit" or ttype == "retrait":
                obj_dict, child_dict, child_transaction = self._special_transaction(row, obj_dict, child_dict)
            elif ttype == "invitation":
                obj_dict, child_dict, child_transaction = self._guest_transaction(row, obj_dict, child_dict)
            elif ttype == "don" or ttype == "transfert":
                obj_dict, child_dict, child_transaction = self._basic_transaction(row, obj_dict, child_dict)
            else:
                child_transaction = None
                # create base transaction object and typed one
            bulk_mgr.add(Transaction(**obj_dict))
            if child_transaction is not None:
                bulk_mgr.add(child_transaction(**child_dict))
            pk_transaction += 1
        bulk_mgr.done()

    def set_roles(self):
        bulk_mgr = BulkCreateManager(chunk_size=10000)
        membership_ids = Membership.objects.values_list('id',flat=True)
        for m_id in membership_ids:
            bulk_mgr.add(
                Membership.roles.trough(membership_id=m_id,role_id=BDE_ROLE_PK),
                Membership.roles.trough(membership_id=m_id,role_id=KFET_ROLE_PK),
            )
        bulk_mgr.done()

    @timed
    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_buttons(cur, kwargs["chunk"])
        self.import_transaction(cur, kwargs["chunk"], 0)