[feature] Run large batch user creation operations asynchronously #608

Closes #608

---------

Co-authored-by: Federico Capoano <f.capoano@openwisp.io>

Author: Dhanus3133

docs/user/settings.rst

@@ -111,6 +111,18 @@ The default encryption format for storing radius check values.
 A list of disabled encryption formats, by default all formats are enabled
 in order to keep backward compatibility with legacy systems.
 
+``OPENWISP_RADIUS_BATCH_ASYNC_THRESHOLD``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+**Default**: ``15``
+
+When the number of users to be generated in batch user creation is greater
+than or equal to this value, the operation will be executed as a
+background task (asynchronously) using Celery. This prevents timeouts and
+keeps the user interface responsive when creating a large number of users.
+For batches smaller than the threshold, users will be created immediately
+(synchronously).
+
 ``OPENWISP_RADIUS_BATCH_DEFAULT_PASSWORD_LENGTH``
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

openwisp_radius/admin.py

@@ -7,6 +7,7 @@
 from django.contrib.admin.utils import model_ngettext
 from django.contrib.auth import get_user_model
 from django.core.exceptions import PermissionDenied
+from django.http import HttpResponseRedirect
 from django.templatetags.static import static
 from django.urls import reverse
 from django.utils.safestring import mark_safe
@@ -338,13 +339,15 @@ class RadiusBatchAdmin(MultitenantAdminMixin, TimeStampedEditableAdmin):
         "name",
         "organization",
         "strategy",
+        "status",
         "expiration_date",
         "created",
         "modified",
     ]
     fields = [
         "strategy",
         "organization",
+        "status",
         "name",
         "csvfile",
         "prefix",
@@ -393,25 +396,21 @@ def number_of_users(self, obj):
     number_of_users.short_description = _("number of users")
 
     def get_fields(self, request, obj=None):
-        fields = super().get_fields(request, obj)[:]
+        fields = super().get_fields(request, obj)
         if not obj:
+            fields = fields[:]
             fields.remove("users")
+            fields.remove("status")
         return fields
 
     def save_model(self, request, obj, form, change):
-        data = form.cleaned_data
-        strategy = data.get("strategy")
-        if not change:
-            if strategy == "csv":
-                if data.get("csvfile", False):
-                    csvfile = data.get("csvfile")
-                    obj.csvfile_upload(csvfile)
-            elif strategy == "prefix":
-                prefix = data.get("prefix")
-                n = data.get("number_of_users")
-                obj.prefix_add(prefix, n)
-        else:
-            obj.save()
+        if change:
+            super().save_model(request, obj, form, change)
+            return
+        # Save the object initially to get a PK
+        super().save_model(request, obj, form, change)
+        num_users = form.cleaned_data.get("number_of_users", 0)
+        obj.schedule_processing(number_of_users=num_users)
 
     def delete_model(self, request, obj):
         obj.users.all().delete()
@@ -466,16 +465,46 @@ def get_readonly_fields(self, request, obj=None):
         readonly_fields = super(RadiusBatchAdmin, self).get_readonly_fields(
             request, obj
         )
-        if obj:
+        if obj and obj.status != "pending":
             return (
                 "strategy",
                 "prefix",
                 "csvfile",
                 "number_of_users",
                 "users",
                 "expiration_date",
+                "name",
+                "organization",
+                "status",
             ) + readonly_fields
-        return readonly_fields
+        elif obj:
+            return ("status",) + readonly_fields
+        return ("status",) + readonly_fields
+
+    def has_delete_permission(self, request, obj=None):
+        if obj and obj.status == "processing":
+            return False
+        return super().has_delete_permission(request, obj)
+
+    def response_add(self, request, obj, post_url_continue=None):
+        if obj.status != "pending":
+            return super().response_add(request, obj, post_url_continue)
+        opts = self.model._meta
+        msg = _(
+            'The batch user creation "{obj}" was added successfully '
+            "and is now being processed in the background."
+        ).format(obj=obj)
+        if "_continue" in request.POST:
+            self.message_user(request, msg, messages.SUCCESS)
+            post_url_continue = reverse(
+                f"admin:{opts.app_label}_{opts.model_name}_change",
+                args=(obj.pk,),
+                current_app=self.admin_site.name,
+            )
+            return HttpResponseRedirect(post_url_continue)
+        else:
+            self.message_user(request, msg, messages.SUCCESS)
+            return self.response_post_save_add(request, obj)
 
 
 # Inlines for UserAdmin & OrganizationAdmin

openwisp_radius/api/serializers.py

@@ -390,7 +390,7 @@ class RadiusBatchSerializer(serializers.ModelSerializer):
     )
     pdf_link = serializers.SerializerMethodField(
         help_text=(
-            "Downlaod link to PDF file containing user credentials. "
+            "Download link to the PDF file containing user credentials. "
             "Provided only for `prefix` strategy.`"
         ),
         required=False,
@@ -405,9 +405,19 @@ class RadiusBatchSerializer(serializers.ModelSerializer):
         write_only=True,
         min_value=1,
     )
+    status = serializers.CharField(read_only=True)
+
+    def create(self, validated_data):
+        validated_data.pop("organization_slug", None)
+        validated_data.pop("number_of_users", None)
+        return super().create(validated_data)
 
     def get_pdf_link(self, obj):
-        if isinstance(obj, RadiusBatch) and obj.strategy == "prefix":
+        if (
+            isinstance(obj, RadiusBatch)
+            and obj.strategy == "prefix"
+            and obj.status == RadiusBatch.COMPLETED
+        ):
             request = self.context.get("request")
             return request.build_absolute_uri(
                 reverse(
@@ -434,7 +444,7 @@ def validate(self, data):
     class Meta:
         model = RadiusBatch
         fields = "__all__"
-        read_only_fields = ("created", "modified", "user_credentials")
+        read_only_fields = ("status", "user_credentials", "created", "modified")
 
 
 class PasswordResetSerializer(BasePasswordResetSerializer):

openwisp_radius/api/views.py

@@ -108,17 +108,17 @@ def post(self, request, *args, **kwargs):
         serializer = self.get_serializer(data=request.data)
         if serializer.is_valid():
             valid_data = serializer.validated_data.copy()
-            num_of_users = valid_data.pop("number_of_users", None)
-            valid_data["organization"] = valid_data.pop("organization_slug", None)
-            batch = serializer.create(valid_data)
-            strategy = valid_data.get("strategy")
-            if strategy == "csv":
-                batch.csvfile_upload()
-                response = RadiusBatchSerializer(batch, context={"request": request})
-            else:
-                batch.prefix_add(valid_data.get("prefix"), num_of_users)
-                response = RadiusBatchSerializer(batch, context={"request": request})
-            return Response(response.data, status=status.HTTP_201_CREATED)
+            num_of_users = valid_data.get("number_of_users", 0)
+            organization = valid_data.pop("organization_slug", None)
+            valid_data.pop("number_of_users", None)
+            batch = serializer.save(organization=organization)
+            is_async = batch.schedule_processing(number_of_users=num_of_users)
+            batch.refresh_from_db()
+            response_serializer = self.get_serializer(batch)
+            status_code = (
+                status.HTTP_202_ACCEPTED if is_async else status.HTTP_201_CREATED
+            )
+            return Response(response_serializer.data, status=status_code)
         return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
 
 

openwisp_radius/base/models.py

@@ -10,6 +10,8 @@
 import django
 import phonenumbers
 import swapper
+from asgiref.sync import async_to_sync
+from channels.layers import get_channel_layer
 from django.conf import settings
 from django.contrib.auth import get_user_model
 from django.core.cache import cache
@@ -23,13 +25,15 @@
 from django.utils.translation import gettext_lazy as _
 from jsonfield import JSONField
 from model_utils.fields import AutoLastModifiedField
+from openwisp_notifications.signals import notify
 from phonenumber_field.modelfields import PhoneNumberField
 from private_storage.fields import PrivateFileField
 
 from openwisp_radius.registration import (
     REGISTRATION_METHOD_CHOICES,
     get_registration_choices,
 )
+from openwisp_radius.tasks import process_radius_batch
 from openwisp_users.mixins import OrgMixin
 from openwisp_utils.base import KeyField, TimeStampedEditableModel, UUIDModel
 from openwisp_utils.fields import (
@@ -865,13 +869,31 @@ def _get_csv_file_location(instance, filename):
 
 
 class AbstractRadiusBatch(OrgMixin, TimeStampedEditableModel):
+    PENDING = "pending"
+    PROCESSING = "processing"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+    BATCH_STATUS_CHOICES = (
+        (PENDING, _("Pending")),
+        (PROCESSING, _("Processing")),
+        (COMPLETED, _("Completed")),
+        (FAILED, _("Failed")),
+    )
+
     strategy = models.CharField(
         _("strategy"),
         max_length=16,
         choices=_STRATEGIES,
         db_index=True,
         help_text=_("Import users from a CSV or generate using a prefix"),
     )
+    status = models.CharField(
+        max_length=16,
+        choices=BATCH_STATUS_CHOICES,
+        default=PENDING,
+        db_index=True,
+    )
     name = models.CharField(
         verbose_name=_("name"),
         max_length=128,
@@ -1064,6 +1086,71 @@ def _remove_files(self):
         if self.csvfile:
             self.csvfile.storage.delete(self.csvfile.name)
 
+    def schedule_processing(self, number_of_users=0):
+        items_to_process = 0
+        if self.strategy == "prefix":
+            items_to_process = number_of_users
+        elif self.strategy == "csv" and self.csvfile:
+            try:
+                csv_data = self.csvfile.read()
+                decoded_data = decode_byte_data(csv_data)
+                items_to_process = sum(1 for row in csv.reader(StringIO(decoded_data)))
+                self.csvfile.seek(0)
+            except Exception as e:
+                logger.error(f"Could not count rows in CSV for batch {self.pk}: {e}")
+                items_to_process = app_settings.BATCH_ASYNC_THRESHOLD
+        is_async = items_to_process >= app_settings.BATCH_ASYNC_THRESHOLD
+        if is_async:
+            process_radius_batch.delay(self.pk, number_of_users=number_of_users)
+        else:
+            self.process(number_of_users=number_of_users)
+        return is_async
+
+    def process(self, number_of_users=0, is_async=False):
+        channel_layer = get_channel_layer()
+        group_name = f"radius_batch_{self.pk}"
+        try:
+            self.status = self.PROCESSING
+            self.save(update_fields=["status"])
+            if self.strategy == "prefix":
+                self.prefix_add(self.prefix, number_of_users)
+            elif self.strategy == "csv":
+                self.csvfile_upload()
+            self.status = self.COMPLETED
+            self.save(update_fields=["status"])
+            if is_async:
+                notify.send(
+                    type="generic_message",
+                    level="success",
+                    message=_(
+                        f'The batch creation operation for "{self.name}" '
+                        "has completed successfully."
+                    ),
+                    sender=self.organization,
+                    target=self,
+                    description=_(f"Number of users processed: {self.users.count()}."),
+                )
+        except Exception as e:
+            logger.error(
+                "RadiusBatch %s failed during processing:", self.pk, exc_info=True
+            )
+            self.status = self.FAILED
+            self.save(update_fields=["status"])
+            notify.send(
+                type="generic_message",
+                level="error",
+                message=_(f'The batch creation operation for "{self.name}" failed.'),
+                sender=self.organization,
+                target=self,
+                description=_(
+                    f"An error occurred while processing the batch.\n\n" f"Error: {e}"
+                ),
+            )
+        finally:
+            async_to_sync(channel_layer.group_send)(
+                group_name, {"type": "batch_status_update", "status": self.status}
+            )
+
 
 class AbstractRadiusToken(OrgMixin, TimeStampedEditableModel, models.Model):
     # key field is a primary key so additional id field will be redundant

openwisp_radius/consumers.py

@@ -0,0 +1,47 @@
+from asgiref.sync import sync_to_async
+from channels.generic.websocket import AsyncJsonWebsocketConsumer
+from django.core.exceptions import ObjectDoesNotExist
+
+from .utils import load_model
+
+
+class RadiusBatchConsumer(AsyncJsonWebsocketConsumer):
+    def _user_can_access_batch(self, user, batch_id):
+        RadiusBatch = load_model("RadiusBatch")
+        # Superusers have access to everything,
+        if user.is_superuser:
+            return RadiusBatch.objects.filter(pk=batch_id).exists()
+        # For non-superusers, check their managed organizations
+        try:
+            RadiusBatch.objects.filter(
+                pk=batch_id, organization__in=user.organizations_managed
+            ).exists()
+            return True
+        except ObjectDoesNotExist:
+            return False
+
+    async def connect(self):
+        self.batch_id = self.scope["url_route"]["kwargs"]["batch_id"]
+        self.user = self.scope["user"]
+        self.group_name = f"radius_batch_{self.batch_id}"
+
+        if not self.user.is_authenticated or not self.user.is_staff:
+            await self.close()
+            return
+
+        has_permission = await sync_to_async(self._user_can_access_batch)(
+            self.user, self.batch_id
+        )
+
+        if not has_permission:
+            await self.close()
+            return
+
+        await self.channel_layer.group_add(self.group_name, self.channel_name)
+        await self.accept()
+
+    async def disconnect(self, close_code):
+        await self.channel_layer.group_discard(self.group_name, self.channel_name)
+
+    async def batch_status_update(self, event):
+        await self.send_json({"status": event["status"]})