[feature] Proof of concept need some code refinement

Author: dee077

openwisp_controller/connection/api/serializers.py

@@ -12,6 +12,7 @@
 DeviceConnection = load_model("connection", "DeviceConnection")
 Credentials = load_model("connection", "Credentials")
 Device = load_model("config", "Device")
+BatchCommand = load_model("connection", "BatchCommand")
 
 
 class ValidatedDeviceFieldSerializer(ValidatedModelSerializer):
@@ -43,6 +44,10 @@ class CommandSerializer(ValidatedDeviceFieldSerializer):
         required=False,
         pk_field=serializers.UUIDField(format="hex_verbose"),
     )
+    batch_command = serializers.PrimaryKeyRelatedField(
+        read_only=True,
+        pk_field=serializers.UUIDField(format="hex_verbose"),
+    )
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -115,3 +120,54 @@ class Meta:
             "is_working": {"read_only": True},
         }
         read_only_fields = ("created", "modified")
+
+
+class BatchCommandExecuteSerializer(
+    FilterSerializerByOrgManaged, serializers.ModelSerializer
+):
+    type = serializers.CharField(source="command_type")
+    input = serializers.JSONField(
+        source="command_input", allow_null=True, required=False
+    )
+    devices = serializers.PrimaryKeyRelatedField(
+        many=True,
+        queryset=Device.objects.all(),
+        required=False,
+        allow_empty=True,
+        pk_field=serializers.UUIDField(format="hex_verbose"),
+    )
+
+    class Meta:
+        model = BatchCommand
+        fields = (
+            "organization",
+            "type",
+            "input",
+            "devices",
+            "group",
+            "location",
+        )
+        extra_kwargs = {
+            "organization": {"required": False, "allow_null": True},
+        }
+
+    def validate(self, data):
+        if (
+            not data.get("organization")
+            and not self.context["request"].user.is_superuser
+        ):
+            raise serializers.ValidationError(
+                _("Only superusers can execute batch commands without an organization.")
+            )
+        if devices := data.get("devices"):
+            org = data.get("organization")
+            for device in devices:
+                if org and device.organization_id != org.id:
+                    raise serializers.ValidationError(
+                        {
+                            "devices": _(
+                                "All devices must belong to the same organization."
+                            )
+                        }
+                    )
+        return data

openwisp_controller/connection/api/urls.py

@@ -40,6 +40,11 @@ def get_api_urls(api_views):
             api_views.deviceconnection_detail_view,
             name="deviceconnection_detail",
         ),
+        path(
+            "api/v1/controller/batch-command/execute/",
+            api_views.batch_command_execute_view,
+            name="batch_command_execute",
+        ),
     ]
 
 

openwisp_controller/connection/api/views.py

@@ -1,12 +1,15 @@
 from django.utils.translation import gettext_lazy as _
 from drf_yasg import openapi
 from drf_yasg.utils import swagger_auto_schema
+from rest_framework import status
 from rest_framework.generics import (
+    GenericAPIView,
     ListCreateAPIView,
     RetrieveAPIView,
     RetrieveUpdateDestroyAPIView,
     get_object_or_404,
 )
+from rest_framework.response import Response
 from swapper import load_model
 
 from openwisp_utils.api.pagination import OpenWispPagination
@@ -17,6 +20,7 @@
     RelatedDeviceProtectedAPIMixin,
 )
 from .serializers import (
+    BatchCommandExecuteSerializer,
     CommandSerializer,
     CredentialSerializer,
     DeviceConnectionSerializer,
@@ -26,6 +30,7 @@
 Device = load_model("config", "Device")
 Credentials = load_model("connection", "Credentials")
 DeviceConnection = load_model("connection", "DeviceConnection")
+BatchCommand = load_model("connection", "BatchCommand")
 
 
 class BaseCommandView(RelatedDeviceProtectedAPIMixin):
@@ -138,6 +143,33 @@ class DeviceConnectionListCreateView(BaseDeviceConnection, ListCreateAPIView):
 DeviceConnenctionListCreateView = DeviceConnectionListCreateView
 
 
+class BatchCommandExecuteView(ProtectedAPIMixin, GenericAPIView):
+    model = BatchCommand
+    queryset = BatchCommand.objects.all()
+    serializer_class = BatchCommandExecuteSerializer
+
+    def post(self, request):
+        serializer = self.get_serializer(data=request.data)
+        serializer.is_valid(raise_exception=True)
+        batch = serializer.save()
+        batch.launch_async()
+        return Response({"batch": str(batch.pk)}, status=status.HTTP_201_CREATED)
+
+    def get(self, request):
+        serializer = self.get_serializer(data=request.query_params)
+        serializer.is_valid(raise_exception=True)
+        data = serializer.validated_data
+        device_pks = []
+        devices_list = data.pop("devices", None)
+        if devices_list:
+            device_pks = [str(d.pk) for d in devices_list]
+        batch = BatchCommand(**data)
+        if not device_pks:
+            resolved = batch.resolve_devices()
+            device_pks = [str(d.pk) for d in resolved]
+        return Response({"devices": device_pks})
+
+
 class DeviceConnectionDetailView(BaseDeviceConnection, RetrieveUpdateDestroyAPIView):
     def get_object(self):
         queryset = self.filter_queryset(self.get_queryset())
@@ -158,3 +190,5 @@ def get_object(self):
 
 # TODO: remove in version 1.4
 deviceconnection_details_view = deviceconnection_detail_view
+
+batch_command_execute_view = BatchCommandExecuteView.as_view()

openwisp_controller/connection/base/models.py

@@ -563,6 +563,8 @@ def save(self, *args, **kwargs):
         output = super().save(*args, **kwargs)
         if adding:
             self._schedule_command()
+        if self.batch_command_id and self.status != "in-progress":
+            self.batch_command.calculate_and_update_status()
         return output
 
     def _save_without_resurrecting(self):
@@ -738,6 +740,8 @@ class AbstractBatchCommand(TimeStampedEditableModel):
     organization = models.ForeignKey(
         get_model_name("openwisp_users", "Organization"),
         on_delete=models.CASCADE,
+        blank=True,
+        null=True,
     )
     status = models.CharField(
         max_length=12, choices=STATUS_CHOICES, default=STATUS_CHOICES[0][0]
@@ -761,37 +765,55 @@ class AbstractBatchCommand(TimeStampedEditableModel):
         null=True,
         verbose_name=_("location"),
     )
-    include_all_devices = models.BooleanField(default=False)
+    devices = models.ManyToManyField(
+        get_model_name("config", "Device"),
+        blank=True,
+        verbose_name=_("devices"),
+    )
     total_devices = models.PositiveIntegerField(default=0)
     successful = models.PositiveIntegerField(default=0)
     failed = models.PositiveIntegerField(default=0)
     cancelled = models.PositiveIntegerField(default=0)
 
     class Meta:
         abstract = True
-        verbose_name = _("Batch command operation")
-        verbose_name_plural = _("Batch command operations")
+        verbose_name = _("Batch command")
+        verbose_name_plural = _("Batch commands")
 
     def clean(self):
         super().clean()
-        if self.group and self.group.organization != self.organization:
-            raise ValidationError(
-                {
-                    "group": _(
-                        "The organization of the group doesn't match "
-                        "the organization of the batch command operation"
-                    )
-                }
-            )
-        if self.location and self.location.organization != self.organization:
-            raise ValidationError(
-                {
-                    "location": _(
-                        "The organization of the location doesn't match "
-                        "the organization of the batch command operation"
+        if self.organization_id:
+            if self.group and self.group.organization != self.organization:
+                raise ValidationError(
+                    {
+                        "group": _(
+                            "The organization of the group doesn't match "
+                            "the organization of the batch command operation"
+                        )
+                    }
+                )
+            if self.location and self.location.organization != self.organization:
+                raise ValidationError(
+                    {
+                        "location": _(
+                            "The organization of the location doesn't match "
+                            "the organization of the batch command operation"
+                        )
+                    }
+                )
+            if self.pk and self.devices.exists():
+                org_mismatch = self.devices.exclude(
+                    organization=self.organization
+                ).exists()
+                if org_mismatch:
+                    raise ValidationError(
+                        {
+                            "devices": _(
+                                "All devices must belong to the same "
+                                "organization as the batch command."
+                            )
+                        }
                     )
-                }
-            )
         allowed = dict(
             AbstractCommand.get_org_allowed_commands(
                 organization_id=self.organization_id
@@ -813,14 +835,16 @@ def clean(self):
             raise ValidationError({"command_input": e.message})
 
     def resolve_devices(self):
+        if self.pk and self.devices.exists():
+            return self.devices.all()
         Device = load_model("config", "Device")
-        qs = Device.objects.filter(organization=self.organization)
+        qs = Device.objects.all()
+        if self.organization_id:
+            qs = qs.filter(organization=self.organization)
         if self.group:
             qs = qs.filter(group=self.group)
         if self.location:
             qs = qs.filter(location=self.location)
-        if not self.include_all_devices and not self.group and not self.location:
-            qs = qs.none()
         return qs
 
     def launch(self):