BITMAKER-2542 Allow setting DataPersistence field on Project/Spider m…

…odels (#150) --------- Co-authored-by: Raymond Negron <[email protected]> Co-authored-by: emegona <[email protected]>
bitmakerla · Apr 11, 2023 · d14d7b3 · d14d7b3
1 parent c6bcf35
commit d14d7b3
Show file tree

Hide file tree

Showing 31 changed files with 1,064 additions and 394 deletions.
diff --git a/estela-api/api/serializers/cronjob.py b/estela-api/api/serializers/cronjob.py
@@ -1,15 +1,20 @@
 from croniter import croniter
 from rest_framework import serializers
-from api import errors
-
-from core.models import SpiderJobArg, SpiderJobEnvVar, SpiderCronJob, SpiderJobTag
 
+from api import errors
 from api.serializers.job_specific import (
     SpiderJobArgSerializer,
     SpiderJobEnvVarSerializer,
     SpiderJobTagSerializer,
 )
-from core.cronjob import enable_cronjob, disable_cronjob, update_schedule
+from core.cronjob import disable_cronjob, enable_cronjob, update_schedule
+from core.models import (
+    DataStatus,
+    SpiderCronJob,
+    SpiderJobArg,
+    SpiderJobEnvVar,
+    SpiderJobTag,
+)
 
 
 class SpiderCronJobSerializer(serializers.ModelSerializer):
@@ -144,10 +149,10 @@ def update(self, instance, validated_data):
         if "unique_collection" in validated_data:
             instance.unique_collection = unique_collection
         if "data_status" in validated_data:
-            if data_status == SpiderCronJob.PERSISTENT_STATUS:
-                instance.data_status = SpiderCronJob.PERSISTENT_STATUS
-            elif data_status == SpiderCronJob.PENDING_STATUS:
-                instance.data_status = SpiderCronJob.PENDING_STATUS
+            if data_status == DataStatus.PERSISTENT_STATUS:
+                instance.data_status = DataStatus.PERSISTENT_STATUS
+            elif data_status == DataStatus.PENDING_STATUS:
+                instance.data_status = DataStatus.PENDING_STATUS
                 if data_expiry_days < 1:
                     raise serializers.ValidationError(
                         {"error": errors.POSITIVE_SMALL_INTEGER_FIELD}

diff --git a/estela-api/api/serializers/deploy.py b/estela-api/api/serializers/deploy.py
@@ -59,7 +59,12 @@ def update(self, instance, validated_data):
                 )
                 project.spiders.exclude(name__in=spiders_names).update(deleted=True)
                 new_spiders = [
-                    Spider(name=spider_name, project=project)
+                    Spider(
+                        name=spider_name,
+                        project=project,
+                        data_status=project.data_status,
+                        data_expiry_days=project.data_expiry_days,
+                    )
                     for spider_name in spiders_names
                     if not project.spiders.filter(name=spider_name).exists()
                 ]

diff --git a/estela-api/api/serializers/job.py b/estela-api/api/serializers/job.py
@@ -7,7 +7,13 @@
     SpiderJobTagSerializer,
 )
 from config.job_manager import job_manager
-from core.models import SpiderJob, SpiderJobArg, SpiderJobEnvVar, SpiderJobTag
+from core.models import (
+    DataStatus,
+    SpiderJob,
+    SpiderJobArg,
+    SpiderJobEnvVar,
+    SpiderJobTag,
+)
 
 
 class SpiderJobSerializer(serializers.ModelSerializer):
@@ -149,12 +155,12 @@ def update(self, instance, validated_data):
 
         if (
             "data_status" in validated_data
-            and instance.data_status != SpiderJob.DELETED_STATUS
+            and instance.data_status != DataStatus.DELETED_STATUS
         ):
-            if data_status == SpiderJob.PERSISTENT_STATUS:
-                instance.data_status = SpiderJob.PERSISTENT_STATUS
-            elif data_status == SpiderJob.PENDING_STATUS:
-                instance.data_status = SpiderJob.PENDING_STATUS
+            if data_status == DataStatus.PERSISTENT_STATUS:
+                instance.data_status = DataStatus.PERSISTENT_STATUS
+            elif data_status == DataStatus.PENDING_STATUS:
+                instance.data_status = DataStatus.PENDING_STATUS
                 if data_expiry_days < 1:
                     raise serializers.ValidationError(
                         {"error": errors.POSITIVE_SMALL_INTEGER_FIELD}

diff --git a/estela-api/api/serializers/project.py b/estela-api/api/serializers/project.py
@@ -1,10 +1,7 @@
 from django.contrib.auth.models import User
-from django.db.models import Sum
 from rest_framework import serializers
-from rest_framework.exceptions import APIException
 
-from config.job_manager import spiderdata_db_client
-from core.models import Permission, Project, SpiderJob, UsageRecord
+from core.models import DataStatus, Permission, Project, UsageRecord
 
 
 class UserDetailSerializer(serializers.ModelSerializer):
@@ -39,7 +36,15 @@ class ProjectSerializer(serializers.ModelSerializer):
 
     class Meta:
         model = Project
-        fields = ("pid", "name", "category", "container_image", "users")
+        fields = (
+            "pid",
+            "name",
+            "category",
+            "container_image",
+            "users",
+            "data_status",
+            "data_expiry_days",
+        )
 
 
 class UsageRecordSerializer(serializers.ModelSerializer):
@@ -97,7 +102,6 @@ class ProjectUpdateSerializer(serializers.ModelSerializer):
         ("DEVELOPER", "Developer"),
         ("VIEWER", "Viewer"),
     ]
-
     pid = serializers.UUIDField(
         read_only=True, help_text="A UUID identifying this project."
     )
@@ -117,7 +121,27 @@ class ProjectUpdateSerializer(serializers.ModelSerializer):
         required=False,
         help_text="New permission.",
     )
+    data_status = serializers.ChoiceField(
+        write_only=True,
+        choices=DataStatus.HIGH_LEVEL_OPTIONS,
+        required=False,
+        help_text="New data status.",
+    )
+    data_expiry_days = serializers.IntegerField(
+        write_only=True,
+        required=False,
+        help_text="New data expiry days.",
+    )
 
     class Meta:
         model = Project
-        fields = ("pid", "name", "users", "email", "action", "permission")
+        fields = (
+            "pid",
+            "name",
+            "users",
+            "email",
+            "action",
+            "permission",
+            "data_status",
+            "data_expiry_days",
+        )
diff --git a/estela-api/api/serializers/spider.py b/estela-api/api/serializers/spider.py
@@ -1,9 +1,36 @@
 from rest_framework import serializers
 
-from core.models import Spider
+from core.models import DataStatus, Spider
 
 
 class SpiderSerializer(serializers.ModelSerializer):
     class Meta:
         model = Spider
-        fields = ("sid", "name", "project")
+        fields = ("sid", "name", "project", "data_status", "data_expiry_days")
+
+
+class SpiderUpdateSerializer(serializers.ModelSerializer):
+    sid = serializers.UUIDField(
+        read_only=True, help_text="A UUID identifying this spider."
+    )
+    data_status = serializers.ChoiceField(
+        choices=DataStatus.HIGH_LEVEL_OPTIONS,
+        required=False,
+        help_text="New data status.",
+    )
+    data_expiry_days = serializers.IntegerField(
+        required=False,
+        help_text="New data expiry days.",
+    )
+
+    class Meta:
+        model = Spider
+        fields = ("sid", "name", "data_status", "data_expiry_days")
+
+    def update(self, instance, validated_data):
+        instance.data_status = validated_data.get("data_status", instance.data_status)
+        instance.data_expiry_days = validated_data.get(
+            "data_expiry_days", instance.data_expiry_days
+        )
+        instance.save()
+        return instance
diff --git a/estela-api/api/views/cronjob.py b/estela-api/api/views/cronjob.py
@@ -15,7 +15,7 @@
     SpiderCronJobUpdateSerializer,
 )
 from core.cronjob import create_cronjob, disable_cronjob, run_cronjob_once
-from core.models import Spider, SpiderCronJob
+from core.models import DataStatus, Spider, SpiderCronJob
 
 
 class SpiderCronJobViewSet(
@@ -68,9 +68,9 @@ def create(self, request, *args, **kwargs):
         spider = get_object_or_404(Spider, sid=self.kwargs["sid"], deleted=False)
         serializer = SpiderCronJobCreateSerializer(data=request.data)
         serializer.is_valid(raise_exception=True)
-        data_status = request.data.pop("data_status", SpiderCronJob.PERSISTENT_STATUS)
+        data_status = request.data.pop("data_status", DataStatus.PERSISTENT_STATUS)
 
-        if data_status == SpiderCronJob.PENDING_STATUS:
+        if data_status == DataStatus.PENDING_STATUS:
             date_str = request.data.pop("data_expiry_days", "0/1")
             data_expiry_days = self.get_days(date_str)
             if data_expiry_days < 1:
@@ -138,7 +138,7 @@ def run_once(self, request, *args, **kwargs):
         partial = kwargs.pop("partial", False)
         instance = self.get_object()
 
-        cronjob = SpiderCronJobSerializer(instance)
+        cronjob = SpiderCronJobSerializer(instance, partial=partial)
 
         run_cronjob_once(cronjob.data)
         return Response(cronjob.data, status=status.HTTP_200_OK)
diff --git a/estela-api/api/views/job.py b/estela-api/api/views/job.py
@@ -15,7 +15,7 @@
     SpiderJobUpdateSerializer,
 )
 from config.job_manager import job_manager
-from core.models import Spider, SpiderJob
+from core.models import DataStatus, Spider, SpiderJob
 
 
 class SpiderJobViewSet(
@@ -101,8 +101,8 @@ def create(self, request, *args, **kwargs):
         async_param = request.query_params.get("async", False)
         serializer = SpiderJobCreateSerializer(data=request.data)
         serializer.is_valid(raise_exception=True)
-        data_status = request.data.pop("data_status", SpiderJob.PERSISTENT_STATUS)
-        if data_status == SpiderJob.PENDING_STATUS:
+        data_status = request.data.pop("data_status", DataStatus.PERSISTENT_STATUS)
+        if data_status == DataStatus.PENDING_STATUS:
             data_expiry_days = request.data.pop("data_expiry_days", 1)
             if data_expiry_days < 1:
                 raise ParseError({"error": "Invalid data expiry days value."})

diff --git a/estela-api/api/views/project.py b/estela-api/api/views/project.py
@@ -19,6 +19,7 @@
     UsageRecordSerializer,
 )
 from core.models import (
+    DataStatus,
     Permission,
     Project,
     Spider,
@@ -86,8 +87,12 @@ def update(self, request, *args, **kwargs):
         user_email = serializer.validated_data.pop("email", "")
         action = serializer.validated_data.pop("action", "")
         permission = serializer.validated_data.pop("permission", "")
+        data_status = serializer.validated_data.pop("data_status", "")
+        data_expiry_days = serializer.validated_data.pop("data_expiry_days", 0)
+
         if name:
             instance.name = name
+
         if user_email and user_email != request.user.email:
             if not (
                 request.user.permission_set.get(project=instance).permission
@@ -120,8 +125,13 @@ def update(self, request, *args, **kwargs):
                 instance.users.add(user, through_defaults={"permission": permission})
             else:
                 raise ParseError({"error": "Action not supported."})
-        serializer.save()
 
+        if data_status:
+            instance.data_status = data_status
+            if data_status == DataStatus.PENDING_STATUS and data_expiry_days > 0:
+                instance.data_expiry_days = data_expiry_days
+
+        serializer.save()
         headers = self.get_success_headers(serializer.data)
         return Response(serializer.data, status=status.HTTP_200_OK, headers=headers)
 

diff --git a/estela-api/api/views/spider.py b/estela-api/api/views/spider.py
@@ -1,11 +1,18 @@
-from rest_framework import viewsets
+from drf_yasg.utils import swagger_auto_schema
+from rest_framework import mixins, status
+from rest_framework.response import Response
 
 from api.mixins import BaseViewSet
-from api.serializers.spider import SpiderSerializer
+from api.serializers.spider import SpiderSerializer, SpiderUpdateSerializer
 from core.models import Spider
 
 
-class SpiderViewSet(BaseViewSet, viewsets.ReadOnlyModelViewSet):
+class SpiderViewSet(
+    mixins.ListModelMixin,
+    mixins.RetrieveModelMixin,
+    mixins.UpdateModelMixin,
+    BaseViewSet,
+):
     model_class = Spider
     serializer_class = SpiderSerializer
     lookup_field = "sid"
@@ -15,3 +22,21 @@ def get_queryset(self):
         return self.model_class.objects.filter(
             project__pid=self.kwargs["pid"], deleted=False
         )
+
+    @swagger_auto_schema(
+        request_body=SpiderUpdateSerializer,
+        responses={status.HTTP_200_OK: SpiderUpdateSerializer()},
+    )
+    def update(self, request, *args, **kwargs):
+        partial = kwargs.pop("partial", False)
+        instance = self.get_object()
+        serializer = SpiderUpdateSerializer(
+            instance, data=request.data, partial=partial
+        )
+        serializer.is_valid(raise_exception=True)
+        self.perform_update(serializer)
+
+        if getattr(instance, "_prefetched_objects_cache", None):
+            instance._prefetched_objects_cache = {}
+
+        return Response(serializer.data, status=status.HTTP_200_OK)