def test_stateful_trigger(trigger, bounds):
kwargs = dict(zip(("at_least", "at_most"), bounds))
t = Task(trigger=trigger(**kwargs))
serialized = TaskSchema().dump(t)
assert serialized["trigger"]["fn"] == to_qualified_name(trigger)
assert set(serialized["trigger"]["kwargs"].values()) == set(bounds)
t2 = TaskSchema().load(serialized)
assert t2.trigger is not trigger # the trigger is not the factory function
assert to_qualified_name(t2.trigger) == to_qualified_name(trigger(*bounds))
def test_stateful_validators(validator, validate_on):
with pytest.warns(UserWarning):
t = Task(cache_validator=validator(validate_on))
serialized = TaskSchema().dump(t)
assert serialized["cache_validator"]["fn"] == to_qualified_name(validator)
assert set(serialized["cache_validator"]["kwargs"]["validate_on"]) == set(
validate_on)
t2 = TaskSchema().load(serialized)
assert (t2.cache_validator
is not validator) # the validator is not the factory function
assert to_qualified_name(t2.cache_validator) == to_qualified_name(
validator(validate_on))
class CustomResultHandlerSchema(ObjectSchema):
class Meta:
object_class = lambda: ResultHandler
exclude_fields = ["type"]
type = fields.Function(lambda handler: to_qualified_name(type(handler)),
lambda x: x)
@post_load
def create_object(self, data: dict, **kwargs: Any) -> None:
"""Because we cannot deserialize a custom class, just return `None`"""
return None
class ParameterSchema(TaskMethodsMixin, ObjectSchema):
class Meta:
object_class = lambda: prefect.core.task.Parameter # type: ignore
exclude_fields = ["type", "outputs", "slug"]
type = fields.Function(lambda task: to_qualified_name(type(task)), lambda x: x)
name = fields.String(required=True)
slug = fields.String(allow_none=True)
default = JSONCompatible(allow_none=True)
required = fields.Boolean(allow_none=True)
description = fields.String(allow_none=True)
tags = fields.List(fields.String())
outputs = fields.Method("load_outputs", allow_none=True)
class CustomResultSchema(ObjectSchema):
class Meta:
object_class = lambda: result.Result
exclude_fields = ["type"]
location = fields.Str(allow_none=True)
type = fields.Function(lambda result: to_qualified_name(type(result)), lambda x: x)
@post_load
def create_object(self, data: dict, **kwargs: Any) -> result.Result:
"""Because we cannot deserialize a custom class, return a base `Result`"""
return result.Result(location=data.get("location"))
class TaskSchema(TaskMethodsMixin, ObjectSchema):
class Meta:
object_class = lambda: prefect.core.Task
exclude_fields = ["type", "inputs", "outputs"]
type = fields.Function(lambda task: to_qualified_name(type(task)),
lambda x: x)
name = fields.String(allow_none=True)
slug = fields.String(allow_none=True)
description = fields.String(allow_none=True)
tags = fields.List(fields.String())
max_retries = fields.Integer(allow_none=True)
retry_delay = fields.TimeDelta(allow_none=True)
inputs = fields.Method("load_inputs", allow_none=True)
outputs = fields.Method("load_outputs", allow_none=True)
timeout = fields.Integer(allow_none=True)
trigger = StatefulFunctionReference(
valid_functions=[
prefect.triggers.all_finished,
prefect.triggers.manual_only,
prefect.triggers.always_run,
prefect.triggers.all_successful,
prefect.triggers.all_failed,
prefect.triggers.any_successful,
prefect.triggers.any_failed,
prefect.triggers.some_failed,
prefect.triggers.some_successful,
],
# don't reject custom functions, just leave them as strings
reject_invalid=False,
allow_none=True,
)
skip_on_upstream_skip = fields.Boolean(allow_none=True)
cache_for = fields.TimeDelta(allow_none=True)
cache_key = fields.String(allow_none=True)
cache_validator = StatefulFunctionReference(
valid_functions=[
prefect.engine.cache_validators.never_use,
prefect.engine.cache_validators.duration_only,
prefect.engine.cache_validators.all_inputs,
prefect.engine.cache_validators.all_parameters,
prefect.engine.cache_validators.partial_inputs_only,
prefect.engine.cache_validators.partial_parameters_only,
],
# don't reject custom functions, just leave them as strings
reject_invalid=False,
allow_none=True,
)
auto_generated = fields.Boolean(allow_none=True)
class FlowSchema(ObjectSchema):
# All nested 'many' types that are stored as a `Set` in the `Flow` must be sorted
# using a `value_selection_fn` so `Flow.serialized_hash()` is deterministic
class Meta:
object_class = lambda: prefect.core.Flow
exclude_fields = ["type", "parameters"]
# ordered to make sure Task objects are loaded before Edge objects, due to Task caching
ordered = True
project = fields.String(allow_none=True)
name = fields.String(required=True, allow_none=True)
version = fields.String(allow_none=True)
description = fields.String(allow_none=True)
type = fields.Function(lambda flow: to_qualified_name(type(flow)),
lambda x: x)
schedule = fields.Nested(ScheduleSchema, allow_none=True)
parameters = Nested(ParameterSchema,
value_selection_fn=get_parameters,
many=True)
tasks = Nested(TaskSchema, value_selection_fn=get_tasks, many=True)
edges = Nested(EdgeSchema, value_selection_fn=get_edges, many=True)
reference_tasks = Nested(TaskSchema,
value_selection_fn=get_reference_tasks,
many=True,
only=["slug"])
environment = fields.Nested(EnvironmentSchema, allow_none=True)
run_config = fields.Nested(RunConfigSchema, allow_none=True)
storage = fields.Nested(StorageSchema, allow_none=True)
@post_load
def create_object(self, data: dict, **kwargs: Any) -> prefect.core.Flow:
"""
Flow edges are validated, for example to make sure the keys match Task inputs,
but because we are deserializing all Tasks as base Tasks, the edge validation will
fail (base Tasks have no inputs). Therefore we hold back the edges from Flow
initialization and assign them explicitly.
Args:
- data (dict): the deserialized data
Returns:
- Flow
"""
data["validate"] = False
flow = super().create_object(data)
return flow
class CustomEnvironmentSchema(ObjectSchema):
class Meta:
object_class = lambda: Environment
exclude_fields = ["type"]
labels = fields.List(fields.String())
type = fields.Function(
lambda environment: to_qualified_name(type(environment)), lambda x: x)
@post_load
def create_object(self, data: dict, **kwargs: Any) -> Environment:
"""
Because we cannot deserialize a custom class, we return an empty
Base Environment with the appropriate labels.
"""
return Environment(labels=data.get("labels"))
class FlowSchema(ObjectSchema):
class Meta:
object_class = lambda: prefect.core.Flow
exclude_fields = ["id", "type", "parameters"]
# ordered to make sure Task objects are loaded before Edge objects, due to Task caching
ordered = True
id = fields.String()
project = fields.String(allow_none=True)
name = fields.String(required=True, allow_none=True)
version = fields.String(allow_none=True)
description = fields.String(allow_none=True)
type = fields.Function(lambda flow: to_qualified_name(type(flow)),
lambda x: x)
schedule = fields.Nested(ScheduleSchema, allow_none=True)
parameters = Nested(ParameterSchema,
value_selection_fn=get_parameters,
many=True)
tasks = fields.Nested(TaskSchema, many=True)
edges = fields.Nested(EdgeSchema, many=True)
reference_tasks = Nested(TaskSchema,
value_selection_fn=get_reference_tasks,
many=True,
only=["id"])
environment = fields.Nested(EnvironmentSchema, allow_none=True)
@post_load
def create_object(self, data: dict) -> prefect.core.Flow:
"""
Flow edges are validated, for example to make sure the keys match Task inputs,
but because we are deserializing all Tasks as base Tasks, the edge validation will
fail (base Tasks have no inputs). Therefore we hold back the edges from Flow
initialization and assign them explicitly.
Args:
- data (dict): the deserialized data
Returns:
- Flow
"""
data["validate"] = False
flow = super().create_object(data)
if "id" in data:
flow.id = data["id"]
return flow
