class WorkloadCacheData(EmbeddedDocument):
workload_system = StringField(choices=['loadleveler', 'slurm'])
user_name = StringField()
collected_time = DateTimeField()
update_time = DateTimeField(default=datetime.datetime.utcnow)
request = DictField()
content = GenericEmbeddedDocumentField(
choices=[JobListContent, QueueInfoListContent])
class Job(Document, BrewtilsJob):
meta = {
"auto_create_index":
False,
"index_background":
True,
"indexes": [{
"name": "next_run_time_index",
"fields": ["next_run_time"],
"sparse": True
}],
}
TRIGGER_MODEL_MAPPING = {
"date": DateTrigger,
"cron": CronTrigger,
"interval": IntervalTrigger,
}
name = StringField(required=True)
trigger_type = StringField(required=True,
choices=BrewtilsJob.TRIGGER_TYPES)
trigger = GenericEmbeddedDocumentField(
choices=list(TRIGGER_MODEL_MAPPING.values()))
request_template = EmbeddedDocumentField("RequestTemplate")
misfire_grace_time = IntField()
coalesce = BooleanField(default=True)
next_run_time = DateTimeField()
success_count = IntField(required=True, default=0, min_value=0)
error_count = IntField(required=True, default=0, min_value=0)
status = StringField(required=True,
choices=BrewtilsJob.STATUS_TYPES,
default="RUNNING")
max_instances = IntField(default=3, min_value=1)
def __str__(self):
return BrewtilsJob.__str__(self)
def __repr__(self):
return BrewtilsJob.__repr__(self)
def clean(self):
"""Validate before saving to the database"""
if self.trigger_type not in self.TRIGGER_MODEL_MAPPING:
raise ModelValidationError(
"Cannot save job. No matching model for trigger type: %s" %
self.trigger_type)
if not isinstance(self.trigger,
self.TRIGGER_MODEL_MAPPING[self.trigger_type]):
raise ModelValidationError(
"Cannot save job. Trigger type: %s but got trigger: %s" %
(self.trigger_type, type(self.trigger)))
class BlobData(EmbeddedDocument):
type = StringField(default=None)
name = StringField(default=None)
content = GenericEmbeddedDocumentField(default=None)
meta = {
'allow_inheritance': True,
}
def to_dict(self):
if self.content is None:
content_dict = None
else:
content_dict = self.content.to_dict()
return {'type': self.type, 'name': self.name, 'content': content_dict}
class Job(MongoModel, Document):
brewtils_model = brewtils.models.Job
meta = {
"auto_create_index": False,
"index_background": True,
"indexes": [
{"name": "next_run_time_index", "fields": ["next_run_time"], "sparse": True}
],
}
TRIGGER_MODEL_MAPPING = {
"date": DateTrigger,
"cron": CronTrigger,
"interval": IntervalTrigger,
"file": FileTrigger,
}
name = StringField(required=True)
trigger_type = StringField(required=True, choices=BrewtilsJob.TRIGGER_TYPES)
trigger = GenericEmbeddedDocumentField(choices=list(TRIGGER_MODEL_MAPPING.values()))
request_template = EmbeddedDocumentField("RequestTemplate", required=True)
misfire_grace_time = IntField()
coalesce = BooleanField(default=True)
next_run_time = DateTimeField()
success_count = IntField(required=True, default=0, min_value=0)
error_count = IntField(required=True, default=0, min_value=0)
status = StringField(
required=True, choices=BrewtilsJob.STATUS_TYPES, default="RUNNING"
)
max_instances = IntField(default=3, min_value=1)
def clean(self):
"""Validate before saving to the database"""
if self.trigger_type not in self.TRIGGER_MODEL_MAPPING:
raise ModelValidationError(
f"Cannot save job. No mongo model for trigger type {self.trigger_type}"
)
trigger_class = self.TRIGGER_MODEL_MAPPING.get(self.trigger_type)
if not isinstance(self.trigger, trigger_class):
raise ModelValidationError(
f"Cannot save job. Expected trigger type {self.trigger_type} but "
f"actual type was {type(self.trigger)}"
)
class Base(Document):
ticket_id = IntField(default=None)
owner = StringField(default=None)
repo = StringField(default=None)
timestamp = DateTimeField(required=True, default=datetime.datetime.utcnow)
data = GenericEmbeddedDocumentField(default=None)
meta = {'allow_inheritance': True, 'abstract': True}
def set_data(self, data):
if not isinstance(data, EmbeddedDocument):
raise TypeError("data must be EmbeddedDocument.")
self.data = data
def is_valid(self):
if self.ticket_id is None:
return False
if self.owner is None:
return False
if self.repo is None:
return False
return True
def to_dict(self):
if not self.is_valid():
return None
if self.data is None:
data_dict = None
elif hasattr(self.data, 'to_dict'):
data_dict = self.data.to_dict()
else:
try:
data_dict = dict(self.data)
except ValueError as e:
data_dict = self.data
result = {
'ticket_id': self.ticket_id,
'owner': self.owner,
'repo': self.repo,
'timestamp': self.timestamp,
'data': data_dict
}
return result
class Cart(EmbeddedDocument):
items = ListField(GenericEmbeddedDocumentField())
subtotal = FloatField()
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def __getitem__(self, item: int) -> BaseItem:
return self.items[item]
def __setitem__(self, key: int, value: BaseItem):
self.items[key] = value
def __delitem__(self, key: int):
if key < len(self.items):
del self.items[key]
def get_safe(self) -> dict:
res = document_to_dict(self)
for i in range(len(self.items)):
res['items'][i] = self.items[i].get_safe()
return res
class Person(Document):
settings = GenericEmbeddedDocumentField(choices=[BaseSettings])
class Person(Document):
settings = GenericEmbeddedDocumentField(choices=(AdminSettings,
NonAdminSettings))
class BlogPost(Document):
comments = ListField(
GenericEmbeddedDocumentField(choices=(Comments, )))
class Person(Document):
name = StringField()
likes = ListField(GenericEmbeddedDocumentField(choices=(Dish, )))
class Person(Document):
name = StringField()
like = GenericEmbeddedDocumentField()
