def stopWorkflow(workflowId: int):
"""
Stops given workflow
"""
res = ApiResponse(message="Error in stopping workflow")
notebookIds = list(
NotebookJob.objects.filter(workflow_id=workflowId).values_list(
"notebookId", flat=True))
workflowRuns = WorkflowRun.objects.filter(
workflow_id=workflowId).order_by("-startTimestamp")
if workflowRuns.count():
workflowRun = workflowRuns[0]
workflowRun.status = STATUS_ABORTED
workflowRun.save()
notebookIds = Workflow.objects.get(
id=workflowId).notebookjob_set.all().values_list("notebookId",
flat=True)
responses = [
NotebookJobServices.stopNotebookJob(notebookId)
for notebookId in notebookIds
]
res.update(True, "Stopped workflow successfully")
return res
def getDatasetDetails(datasetLocation, datasourceName):
"""
Service to fetch S3 dataset details
:param datasetLocation: Location of the S3 bucket
"""
res = ApiResponse()
schema = DruidIngestionSpecGenerator._getSchemaForDatasourceInS3(
datasetLocation)
ingestionSpec = DruidIngestionSpecGenerator.getIngestionSpec(
datasetLocation=datasetLocation,
datasourceName=datasourceName,
datasetSchema=schema)
s3DatasetSchema = list(
map(
lambda x: {
"columnName": x.name,
"dataType": x.logical_type.type
}, schema))
datasetDetails = {
"dremioSchema": s3DatasetSchema,
"druidIngestionSpec": ingestionSpec
}
res.update(True, "Dataset schema retrieved successfully",
datasetDetails)
return res
def updateSchedule(workflowId: int, scheduleId: int):
"""Update given workflow's schedule"""
res = ApiResponse(message="Error in updating workflow schedule")
updateStatus = Workflow.objects.filter(id=workflowId).update(
crontab_id=scheduleId)
res.update(True, "Workflow schedule updated successfully", True)
return res
def getConnectionTypes():
res = ApiResponse()
connectionTypes = ConnectionType.objects.all()
serializer = ConnectionTypeSerializer(connectionTypes, many=True)
res.update(True, "Successfully retrieved connection types",
serializer.data)
return res
def addNotebook(payload):
res = ApiResponse(message="Error adding notebook")
notebookTemplate = NotebookTemplate.objects.get(
id=payload.get("notebookTemplateId", 0))
context = payload # Storing payload in context variable so that it can be used for rendering
# Handling connection variables
if payload.get("sourceConnection", False):
connection = Connection.objects.get(id=payload["sourceConnection"])
connectionParams = connection.cpvc.all()
for cp in connectionParams:
paramName = cp.connectionParam.name
context["sourceConnection_" + paramName] = cp.value
if payload.get("targetConnection", False):
connection = Connection.objects.get(id=payload["sourceConnection"])
connectionParams = connection.cpvc.all()
for cp in connectionParams:
paramName = cp.connectionParam.name
context["sourceConnection_" + paramName] = cp.value
# Handling S3 path - Splitting it to get the table name
if payload.get("destinationTableS3Path", False):
destinationTableName = payload["destinationTableS3Path"].rsplit(
'/', 1)[1]
warehouseLocation = payload["destinationTableS3Path"].rsplit(
'/', 1)[0]
context["destinationTableName"] = destinationTableName
context["warehouseLocation"] = warehouseLocation
# Adding a temp table name to the context
context["tempTableName"] = "tempTable_" + str(round(
time.time() * 1000))
notebook = Template(notebookTemplate.template).render(Context(context))
response = Zeppelin.addNotebook(notebook)
if response:
res.update(True, "Notebook added successfully")
return res
def getWorkflows(offset: int = 0,
limit: int = 25,
sortColumn: str = None,
sortOrder: str = None):
"""
Service to fetch and serialize Workflows
:param offset: Offset for fetching NotebookJob objects
"""
res = ApiResponse(message="Error retrieving workflows")
workflows = Workflow.objects.order_by("-id")
total = workflows.count()
if (sortColumn):
isAscending = True if sortOrder == "ascend" else False
workflows = WorkflowServices.sortingOnWorkflows(
workflows, sortColumn, isAscending)
data = WorkflowSerializer(workflows[offset:offset + limit],
many=True).data
res.update(
True,
"Workflows retrieved successfully",
{
"total": total,
"workflows": data
},
)
return res
def getTimezones():
"""
Service to fetch all pytz timezones
"""
res = ApiResponse()
timezones = pytz.all_timezones
res.update(True, "Timezones fetched successfully", timezones)
return res
def runNotebookJob(notebookId: str):
"""
Service to run notebook job
"""
res = ApiResponse("Error in running notebook")
runNotebookJobTask.delay(notebookId=notebookId, runType="Manual")
res.update(True, "Notebook triggered successfully", None)
return res
def getAllAccountSettings():
"""
Service to get keys and values of all AccountSettings
"""
res = ApiResponse()
data = list(AccountSetting.objects.all().values("key", "value", "label", "type"))
res.update(True, "Fetched account settings successfully", data)
return res
def getDriversCount():
"""
Gets Driver and executors count
"""
res = ApiResponse()
data = Kubernetes.getDriversCount()
res.update(True, "Pods status retrieved successfully", data)
return res
def unarchiveNotebook(notebookId: str, notebookName: str):
"""
Service to unarchive notebook
"""
res = ApiResponse(message="Error in archiving notebook")
response = Zeppelin.renameNotebook(notebookId, notebookName)
if response:
res.update(True, "Notebook unarchived successfully", None)
return res
def getAllAccountSettings():
"""
Service to get keys and values of all AccountSettings
"""
res = ApiResponse()
accountSettings = AccountSetting.objects.all()
data = AccountSettingSerializer(accountSettings, many=True).data
res.update(True, "Fetched account settings successfully", data)
return res
def deleteWorkflow(workflowId: int):
"""
Delete workflow
:param workflowId: id of Workflows.Workflow
"""
res = ApiResponse(message="Error in deleting workflow logs")
Workflow.objects.filter(id=workflowId).delete()
res.update(True, "Workflow deleted successfully")
return res
def cloneNotebook(notebookId: str, payload: dict):
"""
Service to run notebook job
"""
res = ApiResponse(message="Error in cloning notebook")
response = Zeppelin.cloneNotebook(notebookId, json.dumps(payload))
if response:
res.update(True, "Notebook cloned successfully", None)
return res
def deleteNotebook(notebookId: str):
"""
Service to run notebook job
"""
res = ApiResponse(message="Error in cloning notebook")
response = Zeppelin.deleteNotebook(notebookId)
if response:
res.update(True, "Notebook deleted successfully", None)
return res
def getSchedules():
"""
Service to get all CrontabSchedule objects
"""
res = ApiResponse()
crontabSchedules = CrontabSchedule.objects.all()
data = CrontabScheduleSerializer(crontabSchedules, many=True).data
res.update(True, "Schedules fetched successfully", data)
return res
def clearNotebookResults(notebookId: str):
"""
Service to run notebook job
"""
res = ApiResponse(message="Error in clearing notebook")
response = Zeppelin.clearNotebookResults(notebookId)
if response:
res.update(True, "Notebook cleared successfully", None)
return res
def deleteNotebookJob(notebookId: int):
"""
Service to update crontab of an existing NotebookJob
:param notebookId: ID of the Notebook for which to delete
"""
res = ApiResponse()
NotebookJob.objects.filter(name=notebookId).delete()
res.update(True, "NotebookJob deleted successfully", None)
return res
def updateTriggerWorkflow(workflowId: int, triggerWorkflowId: int,
triggerWorkflowStatus: int):
"""Update given workflow's trigger workflow"""
res = ApiResponse(message="Error in updating trigger workflow")
updateStatus = Workflow.objects.filter(id=workflowId).update(
triggerWorkflow_id=triggerWorkflowId,
triggerWorkflowStatus=triggerWorkflowStatus)
res.update(True, "Trigger workflow updated successfully", updateStatus)
return res
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
def archivedNotebooks():
"""
Get archived notebooks
"""
res = ApiResponse(message="Error retrieving archived notebooks")
notebooks = Zeppelin.getAllNotebooks("~Trash")
if notebooks:
res.update(True, "Archived notebooks retrieved successfully",
notebooks)
return res
def getNotebookObject(notebookObjId: int):
"""
Service to fetch specified NotebookObject
:param notebookObjId: ID of the notebook object
"""
res = ApiResponse(message="Error retrieving specified Notebook Object")
notebookObj = NotebookObject.objects.get(id=notebookObjId)
data = NotebookObjectSerializer(notebookObj).data
res.update(True, "NotebookObject retrieved successfully", data)
return res
def stopNotebookJob(notebookId: str):
"""
Service to run notebook job
"""
res = ApiResponse(message="Error in stopping notebook")
# Updating runStatus that the task was aborted
response = Zeppelin.stopNotebookJob(notebookId)
if response:
res.update(True, "Notebook stopped successfully", None)
return res
def deleteNotebook(notebookId: str):
"""
Service to run notebook job
"""
res = ApiResponse(message="Error in deleting notebook")
response = Zeppelin.deleteNotebook(notebookId)
if response:
NotebookObject.objects.filter(
notebookZeppelinId=notebookId).delete()
res.update(True, "Notebook deleted successfully", None)
return res
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 crontabScheduleId: ID of CrontabSchedule
"""
res = ApiResponse()
NotebookJob.objects.filter(notebookId=notebookId).update(
enabled=enabled)
res.update(True, "NotebookJob updated successfully", None)
return res
def updateNotebookJob(notebookJobId: int, crontabScheduleId: int):
"""
Service to update crontab of an existing NotebookJob
:param notebookId: ID of the NotebookJob for which to update crontab
:param crontabScheduleId: ID of CrontabSchedule
"""
res = ApiResponse()
crontabScheduleObj = CrontabSchedule.objects.get(id=crontabScheduleId)
NotebookJob.objects.filter(id=notebookJobId).update(
crontab=crontabScheduleObj)
res.update(True, "NotebookJob updated successfully", None)
return res
def updateWorkflow(
id: int,
name: str,
scheduleId: int,
triggerWorkflowId: int,
triggerWorkflowStatus: str,
notebookIds: List[int],
):
"""
Updates workflow
:param name: name of new workflow
:param scheduleId: crontab id
:param triggerWorkflowId: id of workflow which triggers this workflow
:param triggerWorkflowStatus: ["success", "failure", "always"] required
status of triggerWorkflow to trigger this workflow
:param notebookIds: notebookIds for workflow
"""
res = ApiResponse(message="Error in updating workflow")
workflow = Workflow.objects.get(id=id)
if not workflow:
return res
workflow.name = name
workflow.triggerWorkflow_id = triggerWorkflowId
workflow.triggerWorkflowStatus = triggerWorkflowStatus
workflow.save()
if scheduleId:
if workflow.periodictask:
workflow.periodictask.crontab_id = scheduleId
workflow.periodictask.save()
else:
periodictask = PeriodicTask.objects.create(
crontab_id=scheduleId,
name=name,
task="workflows.tasks.runWorkflowJob",
args=str([workflow.id]))
workflow.periodictask = periodictask
workflow.save()
else:
if workflow.periodictask:
PeriodicTask.objects.get(id=workflow.periodictask).delete()
workflow.periodictask = None
workflow.save()
WorkflowNotebookMap.objects.filter(workflow_id=id).delete()
notebookJobs = [
WorkflowNotebookMap(workflow_id=id, notebookId=notebookId)
for notebookId in notebookIds
]
WorkflowNotebookMap.objects.bulk_create(notebookJobs)
res.update(True, "Workflow updated successfully", None)
return res
def addConnection(payload):
res = ApiResponse()
connectionType = ConnectionType.objects.get(id=payload["connectionType_id"])
connection = Connection.objects.create(
name=payload["name"], description=payload["description"], connectionType=connectionType
)
for param in payload["params"]:
cp = ConnectionParam.objects.get(name=param, connectionType=connectionType)
ConnectionParamValue.objects.create(
connectionParam=cp, value=payload["params"][param], connection=connection
)
res.update(True, "Connection added successfully")
return res
def updateAccountSettings(settings: list):
"""
Service to set specified value of AccountSetting with specified key
:param key: Key of the Account Setting
:param value: Value to be set in the Account Setting
"""
res = ApiResponse()
for setting in settings:
accountSetting = AccountSetting.objects.get(key=setting["key"])
accountSetting.value = setting["value"]
accountSetting.save()
res.update(True, "Updated account setting successfully")
return res
def runNotebookJob(notebookId: str):
"""
Service to run notebook job
"""
res = ApiResponse("Error in running notebook")
notebookRunLogs = NotebookRunLogs.objects.create(
notebookId=notebookId,
status=NOTEBOOK_STATUS_QUEUED,
runType="Manual")
runNotebookJobTask.delay(notebookId=notebookId,
notebookRunLogsId=notebookRunLogs.id,
runType="Manual")
res.update(True, "Notebook triggered successfully", None)
return res
