Merge pull request #17 from cuebook/notebookTests

Refactoring notebook jobs services

api/genie/services/__init__.py

@@ -2,4 +2,5 @@
 from .schemas import Schemas
 from .connections import Connections
 from .notebookTemplates import NotebookTemplateService
-from .kubernetes import KubernetesServices
+from .kubernetes import KubernetesServices
+from .schedules import ScheduleService

api/genie/services/notebookJobs.py

@@ -8,7 +8,7 @@
 from django.template import Template, Context
 # from django_celery_beat.models import CrontabSchedule
 from genie.models import NOTEBOOK_STATUS_ABORT, NOTEBOOK_STATUS_QUEUED, NOTEBOOK_STATUS_RUNNING, NOTEBOOK_STATUS_ABORTING, NotebookObject, NotebookJob, RunStatus, Connection, NotebookTemplate, CustomSchedule as Schedule
-from genie.serializers import NotebookObjectSerializer, ScheduleSerializer, RunStatusSerializer
+from genie.serializers import NotebookObjectSerializer, RunStatusSerializer
 from workflows.models import Workflow, NotebookJob as WorkflowNotebookJob
 from utils.apiResponse import ApiResponse
 from utils.zeppelinAPI import Zeppelin
@@ -307,19 +307,6 @@ def addNotebookJob(notebookId: str, scheduleId: int):
  res.update(True, "NotebookJob added successfully", None)
  return res
 
- @staticmethod
- def updateNotebookJob(notebookJobId: int, scheduleId: int):
- """
- Service to update crontab of an existing NotebookJob
- :param notebookId: ID of the NotebookJob for which to update crontab
- :param scheduleId: ID of schedule
- """
- res = ApiResponse()
- scheduleObj = Schedule.objects.get(crontabschedule_ptr_id=scheduleId)
- NotebookJob.objects.filter(id=notebookJobId).update(crontab=scheduleObj)
- res.update(True, "NotebookJob updated successfully", None)
- return res
-
  @staticmethod
  def deleteNotebookJob(notebookId: int):
  """
@@ -331,112 +318,6 @@ def deleteNotebookJob(notebookId: int):
  res.update(True, "NotebookJob deleted successfully", None)
  return res
 
- @staticmethod
- def toggleNotebookJob(notebookId: int, enabled: bool):
- """
- Service to update crontab of an existing NotebookJob
- :param notebookId: ID of the NotebookJob for which to update crontab
- :param scheduleId: ID of schedule
- """
- res = ApiResponse()
- NotebookJob.objects.filter(notebookId=notebookId).update(enabled=enabled)
- res.update(True, "NotebookJob updated successfully", None)
- return res
-
- @staticmethod
- def getSchedules():
- """
- Service to get all schedule objects
- """
- res = ApiResponse()
- schedules = Schedule.objects.exclude(id=1)
- data = ScheduleSerializer(schedules, many=True).data
- res.update(True, "Schedules fetched successfully", data)
- return res
-
- @staticmethod
- def addSchedule(cron: str, timezone: str = None, name: str = ""):
- """
- Service to add Schedule
- :param cron: Crontab in string format
- :param timezone: Timezone string for which to configure Schedule
- :param name: Name of schedule provided by user
- """
- res = ApiResponse()
- cronElements = cron.split()
- if len(cronElements) != 5:
- res.update(False, "Crontab must contain five elements")
- return res 
- timezone = timezone if timezone else "UTC"
- schedule = Schedule.objects.create(
- minute=cronElements[0],
- hour=cronElements[1],
- day_of_month=cronElements[2],
- month_of_year=cronElements[3],
- day_of_week=cronElements[4],
- timezone=timezone,
- name=name,
- )
- res.update(True, "Schedule added successfully", schedule.id)
- return res
-
- @staticmethod
- def getSingleSchedule(scheduleId: int):
- """
- Service to get singleSchedule
- :param scheduleId: int
- """
-
- res = ApiResponse()
- schedules = Schedule.objects.filter(crontabschedule_ptr_id=scheduleId)
- data = ScheduleSerializer(schedules, many=True).data
- res.update(True, "Schedules fetched successfully", data)
- return res
-
- @staticmethod
- def updateSchedule(id, crontab, timezone, name):
- """
- Service to update Schedule
- param id: int
- param cron: Crontab in string format
- param timezone: Timezone in string format
- param name: String
- """
- res = ApiResponse()
- cronElements = crontab.split(" ")
- if len(cronElements) != 5:
- res.update(False, "Crontab must contain five elements")
- return res 
- schedule = Schedule.objects.get(crontabschedule_ptr_id=id)
- schedule.minute=cronElements[0]
- schedule.hour=cronElements[1]
- schedule.day_of_month=cronElements[2]
- schedule.month_of_year=cronElements[3]
- schedule.day_of_week=cronElements[4]
- schedule.timezone = timezone
- schedule.name = name
- schedule.save()
- res.update(True, "Schedules updated successfully", [])
- return res
-
- @staticmethod
- def deleteSchedule(scheduleId: int):
- """ Service to delete schedule of given scheduleId """
- res = ApiResponse()
- Schedule.objects.filter(id=scheduleId).delete()
- res.update(True, "Schedules deleted successfully", [])
- return res
-
- @staticmethod
- def getTimezones():
- """
- Service to fetch all pytz timezones
- """
- res = ApiResponse()
- timezones = pytz.all_timezones
- res.update(True, "Timezones fetched successfully", timezones)
- return res
-
  @staticmethod
  def runNotebookJob(notebookId: str):
  """

api/genie/services/schedules.py

@@ -0,0 +1,104 @@
+import pytz
+import logging
+# from django_celery_beat.models import CrontabSchedule
+from genie.models import CustomSchedule as Schedule
+from genie.serializers import ScheduleSerializer
+from utils.apiResponse import ApiResponse
+
+# Get an instance of a logger
+logger = logging.getLogger(__name__)
+
+class ScheduleService:
+ @staticmethod
+ def getSchedules():
+ """
+ Service to get all schedule objects
+ """
+ res = ApiResponse()
+ schedules = Schedule.objects.exclude(id=1)
+ data = ScheduleSerializer(schedules, many=True).data
+ res.update(True, "Schedules fetched successfully", data)
+ return res
+
+ @staticmethod
+ def addSchedule(cron: str, timezone: str = None, name: str = ""):
+ """
+ Service to add Schedule
+ :param cron: Crontab in string format
+ :param timezone: Timezone string for which to configure Schedule
+ :param name: Name of schedule provided by user
+ """
+ res = ApiResponse()
+ cronElements = cron.split()
+ if len(cronElements) != 5:
+ res.update(False, "Crontab must contain five elements")
+ return res 
+ timezone = timezone if timezone else "UTC"
+ schedule = Schedule.objects.create(
+ minute=cronElements[0],
+ hour=cronElements[1],
+ day_of_month=cronElements[2],
+ month_of_year=cronElements[3],
+ day_of_week=cronElements[4],
+ timezone=timezone,
+ name=name,
+ )
+ res.update(True, "Schedule added successfully", schedule.id)
+ return res
+
+ @staticmethod
+ def getSingleSchedule(scheduleId: int):
+ """
+ Service to get singleSchedule
+ :param scheduleId: int
+ """
+
+ res = ApiResponse()
+ schedules = Schedule.objects.filter(crontabschedule_ptr_id=scheduleId)
+ data = ScheduleSerializer(schedules, many=True).data
+ res.update(True, "Schedules fetched successfully", data)
+ return res
+
+ @staticmethod
+ def updateSchedule(id, crontab, timezone, name):
+ """
+ Service to update Schedule
+ param id: int
+ param cron: Crontab in string format
+ param timezone: Timezone in string format
+ param name: String
+ """
+ res = ApiResponse()
+ cronElements = crontab.split(" ")
+ if len(cronElements) != 5:
+ res.update(False, "Crontab must contain five elements")
+ return res 
+ schedule = Schedule.objects.get(crontabschedule_ptr_id=id)
+ schedule.minute=cronElements[0]
+ schedule.hour=cronElements[1]
+ schedule.day_of_month=cronElements[2]
+ schedule.month_of_year=cronElements[3]
+ schedule.day_of_week=cronElements[4]
+ schedule.timezone = timezone
+ schedule.name = name
+ schedule.save()
+ res.update(True, "Schedules updated successfully", [])
+ return res
+
+ @staticmethod
+ def deleteSchedule(scheduleId: int):
+ """ Service to delete schedule of given scheduleId """
+ res = ApiResponse()
+ Schedule.objects.filter(id=scheduleId).delete()
+ res.update(True, "Schedules deleted successfully", [])
+ return res
+
+ @staticmethod
+ def getTimezones():
+ """
+ Service to fetch all pytz timezones
+ """
+ res = ApiResponse()
+ timezones = pytz.all_timezones
+ res.update(True, "Timezones fetched successfully", timezones)
+ return res

api/genie/tests/test_views_notebookJobs.py

@@ -99,26 +99,6 @@ def test_notebookJob(client, populate_seed_data, mocker):
  assert response.status_code == 200
  assert response.data['success'] == True
 
- # TODO Test and fix update notebook job
- # Test Update Notebook Schedule
- # mixer.blend("genie.customSchedule", id=2)
- # data = {"notebookId": "BX976MDDE", "scheduleId": 2}
- # response = client.put(path, data=data, content_type="application/json")
- # assert response.status_code == 200
- # assert response.data['success'] == True
-
- # Test disable notebook job
- data = {"notebookId": "BX976MDDE", "enabled": False}
- response = client.put(path, data=data, content_type="application/json")
- assert response.status_code == 200
- assert response.data['success'] == True
-
- # Test enabled notebook job
- data = {"notebookId": "BX976MDDE", "enabled": True}
- response = client.put(path, data=data, content_type="application/json")
- assert response.status_code == 200
- assert response.data['success'] == True
-
  # Test delete notebook job
  response = client.delete(path, data=data, content_type="application/json")
  assert response.status_code == 200

api/genie/views.py

@@ -2,7 +2,7 @@
 import json
 from rest_framework.views import APIView
 from rest_framework.response import Response
-from genie.services import NotebookJobServices, Connections, NotebookTemplateService, KubernetesServices, Schemas
+from genie.services import NotebookJobServices, Connections, NotebookTemplateService, KubernetesServices, Schemas, ScheduleService
 from rest_framework.decorators import api_view
 
 class NotebookOperationsView(APIView):
@@ -99,16 +99,6 @@ def post(self, request):
  scheduleId = request.data["scheduleId"]
  res = NotebookJobServices.addNotebookJob(notebookId=notebookId, scheduleId=scheduleId)
  return Response(res.json())
-
- def put(self, request):
- notebookId = request.data["notebookId"]
- if "scheduleId" in request.data:
- scheduleId = request.data["scheduleId"]
- res = NotebookJobServices.updateNotebookJob(notebookId=notebookId, scheduleId=scheduleId)
- elif "enabled" in request.data:
- enabled = request.data["enabled"]
- res = NotebookJobServices.toggleNotebookJob(notebookId=notebookId, enabled=enabled)
- return Response(res.json())
 
  def delete(self, request, notebookId=None):
  res = NotebookJobServices.deleteNotebookJob(notebookId=notebookId)
@@ -119,22 +109,22 @@ class ScheduleView(APIView):
  Class to get and add available crontab schedules
  """
  def get(self, request):
- res = NotebookJobServices.getSchedules()
+ res = ScheduleService.getSchedules()
  return Response(res.json())
 
  def post(self, request):
  name = request.data["name"]
  cron = request.data["crontab"]
  timezone = request.data["timezone"]
- res = NotebookJobServices.addSchedule(cron=cron, timezone=timezone, name=name)
+ res = ScheduleService.addSchedule(cron=cron, timezone=timezone, name=name)
  return Response(res.json())
 
  def put(self,request):
  id = request.data["id"]
  name = request.data["name"]
  crontab = request.data["crontab"]
  timezone = request.data["timezone"]
- res = NotebookJobServices.updateSchedule(id=id, crontab=crontab, timezone=timezone, name=name)
+ res = ScheduleService.updateSchedule(id=id, crontab=crontab, timezone=timezone, name=name)
  return Response(res.json())
 
 @api_view(["GET", "PUT", "DELETE"])
@@ -145,10 +135,10 @@ def schedule(request: HttpRequest, scheduleId: int) -> Response:
  :param connection_id: Connection Id
  """
  if request.method == "GET":
- res = NotebookJobServices.getSingleSchedule(scheduleId)
+ res = ScheduleService.getSingleSchedule(scheduleId)
  return Response(res.json())
  if request.method == "DELETE":
- res = NotebookJobServices.deleteSchedule(scheduleId)
+ res = ScheduleService.deleteSchedule(scheduleId)
  return Response(res.json())
 
 
@@ -157,7 +147,7 @@ class TimzoneView(APIView):
  Class to get standard pytz timezones
  """
  def get(self, request):
- res = NotebookJobServices.getTimezones()
+ res = ScheduleService.getTimezones()
  return Response(res.json())