173 lines
4.6 KiB
Python
173 lines
4.6 KiB
Python
import copy
|
|
import json
|
|
import uuid
|
|
import logging
|
|
|
|
from django.db import models
|
|
|
|
from jeeves.core.registry import TaskRegistry
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BaseModel(models.Model):
|
|
"""
|
|
Abstract model used when creating models in the project.
|
|
|
|
Uses an UUID4 as primary key field and adds a created_at, modified_at
|
|
fields.
|
|
"""
|
|
|
|
id = models.UUIDField(default=uuid.uuid4, primary_key=True, editable=False)
|
|
created_at = models.DateTimeField(auto_now_add=True)
|
|
modified_at = models.DateTimeField(auto_now=True)
|
|
|
|
class Meta:
|
|
abstract = True
|
|
|
|
|
|
class DefinitionMixin:
|
|
@property
|
|
def definition(self):
|
|
return json.loads(self._definition)
|
|
|
|
@definition.setter
|
|
def definition(self, value: dict):
|
|
self._definition = json.dumps(value)
|
|
|
|
|
|
class Flow(BaseModel, DefinitionMixin):
|
|
name = models.CharField(max_length=64)
|
|
_definition = models.TextField(default=json.dumps({"tasks": []}))
|
|
|
|
def serialize(self):
|
|
definition = copy.deepcopy(self.definition)
|
|
definition["tasks"] = [task.serialize() for task in self.tasks]
|
|
return {"name": self.name, "definition": definition}
|
|
|
|
@property
|
|
def tasks(self):
|
|
for task_ref in self.definition["tasks"]:
|
|
try:
|
|
task = Task.objects.from_reference(task_ref)
|
|
yield task
|
|
except Task.DoesNotExist:
|
|
logger.warning(
|
|
f"Task {task_ref} does not exist on db, removing from definition"
|
|
)
|
|
# TODO: Dirty cleanup of dangling tasks in definitions
|
|
definition = self.definition
|
|
definition["tasks"].remove(task_ref)
|
|
self.definition = definition
|
|
self.save()
|
|
|
|
def execute(self, foreground=False):
|
|
from jeeves.db.tasks import start_execution
|
|
|
|
execution = Run.from_flow(self)
|
|
if not foreground:
|
|
start_execution.send(execution_id=str(execution.pk))
|
|
else:
|
|
start_execution(execution_id=str(execution.pk))
|
|
return Run.objects.get(pk=str(execution.pk)) # Reload model instance
|
|
|
|
|
|
class TaskManager(models.Manager):
|
|
def from_reference(self, reference):
|
|
return self.get_queryset().get(pk=reference)
|
|
|
|
|
|
class Task(DefinitionMixin, BaseModel):
|
|
name = models.CharField(max_length=128)
|
|
type = models.CharField(max_length=128)
|
|
_definition = models.TextField()
|
|
|
|
objects = TaskManager()
|
|
|
|
def serialize(self):
|
|
return {"type": self.type, "definition": self.definition}
|
|
|
|
@property
|
|
def instance(self):
|
|
return TaskRegistry.get_task(self.type, **self.definition)
|
|
|
|
def run(self):
|
|
runner = self.instance
|
|
return FlowStep(str(self.pk), *runner.execute())
|
|
|
|
|
|
class Run(DefinitionMixin, BaseModel):
|
|
PENDING = "pending"
|
|
FINISHED = "finished"
|
|
RUN_STATUS = ((PENDING, PENDING.capitalize()), (FINISHED, FINISHED.capitalize()))
|
|
status = models.CharField(max_length=32, choices=RUN_STATUS, default=PENDING)
|
|
success = models.BooleanField(default=True)
|
|
output = models.TextField(blank=True, default="")
|
|
flow = models.ForeignKey(Flow, on_delete=models.CASCADE)
|
|
_definition = models.TextField()
|
|
_result = models.TextField(default="{}")
|
|
|
|
@property
|
|
def duration(self):
|
|
# TODO: Proper ended_at
|
|
return self.modified_at - self.created_at
|
|
|
|
@property
|
|
def result(self):
|
|
return json.loads(self._result)
|
|
|
|
@result.setter
|
|
def result(self, value: dict):
|
|
self._result = json.dumps(value)
|
|
|
|
@classmethod
|
|
def from_flow(cls, flow: Flow):
|
|
execution = cls()
|
|
execution.definition = flow.serialize()
|
|
execution.flow = flow
|
|
execution.save()
|
|
return execution
|
|
|
|
|
|
class FlowStep:
|
|
task: str
|
|
executed: bool = True
|
|
error: bool = False
|
|
output: str
|
|
|
|
def __init__(self, task_id, error, output):
|
|
self.task = task_id
|
|
self.error = error
|
|
self.output = output
|
|
|
|
def __repr__(self):
|
|
return f"<FlowStep error={self.error}>"
|
|
|
|
def serialize(self):
|
|
return {
|
|
"task": self.task,
|
|
"executed": self.executed,
|
|
"error": self.error,
|
|
"output": self.output,
|
|
}
|
|
|
|
|
|
class Result:
|
|
def __init__(self):
|
|
self.steps = []
|
|
|
|
def __repr__(self):
|
|
return f"""<Result
|
|
last={self.last_flow_step}>
|
|
"""
|
|
|
|
def serialize(self):
|
|
return {"steps": [step.serialize() for step in self.steps]}
|
|
|
|
def add_step(self, flow_step):
|
|
self.steps.append(flow_step)
|
|
|
|
@property
|
|
def last_flow_step(self):
|
|
return self.steps[-1]
|