class TaskInstanceSchema(Schema):
"""Task instance schema"""
task_id = fields.Str()
dag_id = fields.Str()
execution_date = fields.DateTime()
start_date = fields.DateTime()
end_date = fields.DateTime()
duration = fields.Float()
state = TaskInstanceStateField()
_try_number = fields.Int(data_key="try_number")
max_tries = fields.Int()
hostname = fields.Str()
unixname = fields.Str()
pool = fields.Str()
pool_slots = fields.Int()
queue = fields.Str()
priority_weight = fields.Int()
operator = fields.Str()
queued_dttm = fields.DateTime(data_key="queued_when")
pid = fields.Int()
executor_config = fields.Str()
sla_miss = fields.Nested(SlaMissSchema, default=None)
def get_attribute(self, obj, attr, default):
if attr == "sla_miss":
# Object is a tuple of task_instance and slamiss
# and the get_value expects a dict with key, value
# corresponding to the attr.
slamiss_instance = {"sla_miss": obj[1]}
return get_value(slamiss_instance, attr, default)
return get_value(obj[0], attr, default)
class SetTaskInstanceStateFormSchema(Schema):
"""Schema for handling the request of setting state of task instance of a DAG"""
dry_run = fields.Boolean(default=True)
task_id = fields.Str(required=True)
execution_date = fields.DateTime(required=True)
include_upstream = fields.Boolean(required=True)
include_downstream = fields.Boolean(required=True)
include_future = fields.Boolean(required=True)
include_past = fields.Boolean(required=True)
new_state = TaskInstanceStateField(required=True, validate=validate.OneOf([State.SUCCESS, State.FAILED]))
class TaskInstanceSchema(SQLAlchemySchema):
"""Task instance schema"""
class Meta:
"""Meta"""
model = TaskInstance
task_id = auto_field()
dag_id = auto_field()
run_id = auto_field(data_key="dag_run_id")
map_index = auto_field()
execution_date = auto_field()
start_date = auto_field()
end_date = auto_field()
duration = auto_field()
state = TaskInstanceStateField()
_try_number = auto_field(data_key="try_number")
max_tries = auto_field()
hostname = auto_field()
unixname = auto_field()
pool = auto_field()
pool_slots = auto_field()
queue = auto_field()
priority_weight = auto_field()
operator = auto_field()
queued_dttm = auto_field(data_key="queued_when")
pid = auto_field()
executor_config = auto_field()
sla_miss = fields.Nested(SlaMissSchema, dump_default=None)
rendered_fields = JsonObjectField(dump_default={})
def get_attribute(self, obj, attr, default):
if attr == "sla_miss":
# Object is a tuple of task_instance and slamiss
# and the get_value expects a dict with key, value
# corresponding to the attr.
slamiss_instance = {"sla_miss": obj[1]}
return get_value(slamiss_instance, attr, default)
elif attr == "rendered_fields":
return get_value(obj[0],
"rendered_task_instance_fields.rendered_fields",
default)
return get_value(obj[0], attr, default)
class SetTaskInstanceStateFormSchema(Schema):
"""Schema for handling the request of setting state of task instance of a DAG"""
dry_run = fields.Boolean(dump_default=True)
task_id = fields.Str(required=True)
execution_date = fields.DateTime(validate=validate_istimezone)
dag_run_id = fields.Str()
include_upstream = fields.Boolean(required=True)
include_downstream = fields.Boolean(required=True)
include_future = fields.Boolean(required=True)
include_past = fields.Boolean(required=True)
new_state = TaskInstanceStateField(required=True,
validate=validate.OneOf(
[State.SUCCESS, State.FAILED]))
@validates_schema
def validate_form(self, data, **kwargs):
"""Validates set task instance state form"""
if not exactly_one(data.get("execution_date"), data.get("dag_run_id")):
raise ValidationError(
"Exactly one of execution_date or dag_run_id must be provided")