refactoring in atomic function

apps/member/management/commands/import_nk15.py

@@ -11,29 +11,22 @@ from django.core.exceptions import ValidationError
 
 from django.contrib.auth.models import User
 from note.models import Note, NoteSpecial, NoteUser, NoteClub
+from note.models import Transaction, TransactionTemplate, TransactionCategory, TransactionType
 from member.models import Profile, Club
 
-class Command(BaseCommand):
-    """
-    Command for importing the database of NK15.
-    Need to be run by a user with a registered role in postgres for the database nk15. 
-    """
 
-    def handle(self, *args, **options):
-        conn = pg.connect(database="nk15",user="nk15_user")
-        cur = conn.cursor(cursor_factory = pge.DictCursor)
-
-        #  Start with Special accounts
+@transaction.atomic
+def import_special(cur):
     cur.execute("SELECT * FROM comptes WHERE idbde <0 ORDER BY idbde;")
-
+    map_idbde = dict()
     for row in cur:
-            with transaction.atomic():
         obj,created = NoteSpecial.objects.get_or_create(special_type = row["pseudo"],
                                                         balance = row["solde"],
                                                         is_active =True)
         if created:
             obj.save()
-                    # The rest
+            map_idbde[row["idbde"]] = obj.pk
+
     cur.execute("SELECT * FROM comptes WHERE idbde=0;")
     res = cur.fetchone()
     clubBde, c = Club.objects.get_or_create(pk = 1,
@@ -55,12 +48,15 @@ class Command(BaseCommand):
     clubBde.save()
     clubKfet.save()
     clubBde.note.solde=res["solde"]
+    map_idbde[0] = clubKfet.note.pk
+    return map_idbde
 
+
+@transaction.atomic
+def import_comptes(cur,map_idbde):
     cur.execute("SELECT * FROM comptes WHERE idbde > 0 ORDER BY idbde;")
     pkclub = 3
-        with transaction.atomic():
     for row in cur:
-                row["idbde"] += 7 # do not overwrite the already populated id.
         if row["type"] == "personne":
             #sanitize password
             if row["passwd"] != "*|*":
@@ -85,6 +81,10 @@ class Command(BaseCommand):
                         last_name = row["prenom"],
                         email = row["mail"],
                     )
+                else:
+                    raise(e)
+            else:
+                pass
             profile = Profile.objects.create(
                 phone_number = row["tel"],
                 address = row["adresse"],
@@ -96,7 +96,6 @@ class Command(BaseCommand):
 
             obj_list =[user, profile, note]
         else:#club
-                    print(row)
             club,c = Club.objects.get_or_create(pk=pkclub,
                                                 name = row["pseudo"],
                                                 email = row["mail"],
@@ -111,4 +110,54 @@ class Command(BaseCommand):
             obj_list = [club,note]
         for obj in obj_list:
             obj.save()
-                    #endfor
+            map_idbde[row["idbde"]] = note.pk
+            #
+    return map_idbde
+
+@transaction.atomic
+def import_boutons(cur,map_idbde):
+    cur.execute("SELECT * FROM boutons;")
+    for row in cur:
+        cat, created = TransactionCategory.objects.get_or_create(name=row["categorie"])
+
+        button = TransactionTemplate.objects.create(pk=row["id"],
+                                            name=row["label"],
+                                            amount=row["montant"],
+                                            destination_id=map_idbde[row["destinataire"]],
+                                            category = cat,
+                                            display = row["affiche"],
+                                            description = row["description"],
+        )
+        if created:
+            cat.save()
+        button.save()
+
+
+class Command(BaseCommand):
+    """
+    Command for importing the database of NK15.
+    Need to be run by a user with a registered role in postgres for the database nk15. 
+    """
+    def add_arguments(self,parser):
+        parser.add_argument('-s', '--special', action = 'store_true')
+        parser.add_argument('-c', '--comptes', action = 'store_true')
+        parser.add_argument('-b', '--boutons', action = 'store_true')
+
+
+    def handle(self, *args, **kwargs):
+        conn = pg.connect(database="nk15",user="nk15_user")
+        cur = conn.cursor(cursor_factory = pge.DictCursor)
+
+        if kwargs["special"]:
+            map_idbde = import_special(cur)
+            print("Minimal setup created")
+
+        if kwargs["comptes"]:
+            map_idbde = import_comptes(cur,map_idbde)
+            print("comptes table imported")
+
+        if kwargs["boutons"]:
+            import_boutons(cur,map_idbde)
+            print("boutons table imported")
+        if kwargs["transaction"]:
+            import_transaction(cur)