[fix] Made batch import operations atomic #551

Fixes #551

Author: stktyagi

openwisp_radius/base/models.py

@@ -17,7 +17,7 @@
 from django.core.cache import cache
 from django.core.exceptions import ObjectDoesNotExist, ValidationError
 from django.core.mail import send_mail
-from django.db import models
+from django.db import models, transaction
 from django.db.models import ProtectedError, Q
 from django.utils import timezone
 from django.utils.crypto import get_random_string
@@ -985,16 +985,17 @@ def clean(self):
     def add(self, reader, password_length=BATCH_DEFAULT_PASSWORD_LENGTH):
         users_list = []
         generated_passwords = []
-        for row in reader:
-            if len(row) == 5:
-                user, password = self.get_or_create_user(
-                    row, users_list, password_length
-                )
-                users_list.append(user)
-                if password:
-                    generated_passwords.append(password)
-        for user in users_list:
-            self.save_user(user)
+        with transaction.atomic():
+            for row in reader:
+                if len(row) == 5:
+                    user, password = self.get_or_create_user(
+                        row, users_list, password_length
+                    )
+                    users_list.append(user)
+                    if password:
+                        generated_passwords.append(password)
+            for user in users_list:
+                self.save_user(user)
         for element in generated_passwords:
             username, password, user_email = element
             send_mail(
@@ -1012,9 +1013,10 @@ def csvfile_upload(
         csv_data = csvfile.read()
         csv_data = decode_byte_data(csv_data)
         reader = csv.reader(StringIO(csv_data), delimiter=",")
-        self.full_clean()
-        self.save()
-        self.add(reader, password_length)
+        with transaction.atomic():
+            self.full_clean()
+            self.save()
+            self.add(reader, password_length)
 
     def prefix_add(self, prefix, n, password_length=BATCH_DEFAULT_PASSWORD_LENGTH):
         users_list, user_credentials = prefix_generate_users(prefix, n, password_length)

openwisp_radius/tests/test_batch_add_users.py

@@ -1,6 +1,8 @@
 from unittest.mock import patch
 
+from django.contrib.auth import get_user_model
 from django.core.exceptions import ValidationError
+from django.db import IntegrityError
 from django.urls import reverse
 
 from ..utils import load_model
@@ -145,5 +147,46 @@ def test_verified_batch_user_creation(self):
         self.assertEqual(user.registered_user.method, "manual")
 
 
+class TestBatchAtomicity(FileMixin, BaseTransactionTestCase):
+    def test_csv_upload_total_rollback(self):
+        User = get_user_model()
+        org = self._get_org()
+        data = [
+            ["user_one", "pass123", "total@example.com", "John", "Doe"],
+            ["user_two", "pass123", "total@example.com", "Jane", "Doe"],
+        ]
+        batch = RadiusBatch(
+            name="total-rollback",
+            strategy="csv",
+            organization=org,
+            csvfile=self._get_csvfile(data),
+        )
+        with self.assertRaises(IntegrityError):
+            batch.csvfile_upload()
+        self.assertFalse(RadiusBatch.objects.filter(name="total-rollback").exists())
+        self.assertFalse(User.objects.filter(username="user_one").exists())
+
+    def test_add_method_internal_atomicity(self):
+        User = get_user_model()
+        org = self._get_org()
+        data = [
+            ["user_one", "pass123", "duplicate@example.com", "John", "Doe"],
+            ["user_two", "pass123", "duplicate@example.com", "Jane", "Doe"],
+        ]
+        batch = self._create_radius_batch(
+            name="atomic-integrity-test",
+            strategy="csv",
+            organization=org,
+            csvfile=self._get_csvfile(data),
+        )
+        with self.assertRaises(IntegrityError):
+            batch.add(data)
+        self.assertFalse(User.objects.filter(username="user_one").exists())
+        self.assertTrue(
+            RadiusBatch.objects.filter(name="atomic-integrity-test").exists()
+        )
+        self.assertEqual(batch.users.count(), 0)
+
+
 del BaseTestCase
 del BaseTransactionTestCase