Repository: PrefectHQ/prefect

Issue:
Add ability to serialize / deserialize user-written ResultHandlers
Right now, only ResultHandlers which are a part of the Prefect library can be serialized / deserialized (and consequently, used).  Possibly using a `FunctionReference`, we should allow for more custom / arbitrary result handlers to be serialized.


Discussion:
Another point to add here: it is _very_ likely that user-created Result Handlers will contain state that needs to be dealt with.  For example, if I implement my own "S3" handler, I might initialize it with a bucket name.  I would expect that attribute to survive the serialization / deserialization process.
Excellent point. A few different strategies come to mind, depending on much we want/need to integrate with Secrets
I tried to take a stab at this tonight, but I'm not sure how to create a `Schema` which doesn't have a well-defined `Meta.object_class`.  User-created `ResultHandler`s will subclass the base `ResultHandler` but ultimately be called something else.  I know how to do this for _fields_ (e.g., using our custom `FunctionReference` field) but not for `Schema`s.

Any insights into how to do this @jlowin ?
Oh I'm an idiot, our `Task` schema does this exact thing!  Awesome.

Resolving patch:
diff --git a/src/prefect/serialization/result_handlers.py b/src/prefect/serialization/result_handlers.py
--- a/src/prefect/serialization/result_handlers.py
+++ b/src/prefect/serialization/result_handlers.py
@@ -1,5 +1,5 @@
 import json
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from marshmallow import ValidationError, fields, post_load
 
@@ -24,6 +24,21 @@ class Meta:
         object_class = ResultHandler
 
 
+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) -> None:
+        """Because we cannot deserialize a custom class, just return `None`"""
+        return None
+
+
 class CloudResultHandlerSchema(BaseResultHandlerSchema):
     class Meta:
         object_class = CloudResultHandler
@@ -71,3 +86,8 @@ class ResultHandlerSchema(OneOfSchema):
         "JSONResultHandler": JSONResultHandlerSchema,
         "LocalResultHandler": LocalResultHandlerSchema,
     }
+
+    def get_obj_type(self, obj: Any) -> str:
+        name = obj.__class__.__name__
+        self.type_schemas.setdefault(name, CustomResultHandlerSchema)
+        return name


Test patch:
diff --git a/tests/serialization/test_result_handlers.py b/tests/serialization/test_result_handlers.py
--- a/tests/serialization/test_result_handlers.py
+++ b/tests/serialization/test_result_handlers.py
@@ -12,10 +12,88 @@
     ResultHandler,
     S3ResultHandler,
 )
-from prefect.serialization.result_handlers import ResultHandlerSchema
+from prefect.serialization.result_handlers import (
+    ResultHandlerSchema,
+    CustomResultHandlerSchema,
+)
 from prefect.utilities.configuration import set_temporary_config
 
 
+class TestCustomSchema:
+    def test_custom_schema_dump_on_dummy_class(self):
+        class Dummy(ResultHandler):
+            def read(self, *args, **kwargs):
+                pass
+
+            def write(self, *args, **kwargs):
+                pass
+
+        serialized = CustomResultHandlerSchema().dump(Dummy())
+        assert serialized["type"].endswith("Dummy")
+        assert serialized["__version__"] == prefect.__version__
+
+    def test_custom_schema_dump_on_stateful_class(self):
+        class Stateful(ResultHandler):
+            def __init__(self, x):
+                self.x = x
+
+            def read(self, *args, **kwargs):
+                pass
+
+            def write(self, *args, **kwargs):
+                pass
+
+        serialized = CustomResultHandlerSchema().dump(Stateful(42))
+        assert serialized["type"].endswith("Stateful")
+        assert serialized["__version__"] == prefect.__version__
+
+    def test_custom_schema_roundtrip_on_base_class(self):
+        class Dummy(ResultHandler):
+            def read(self, *args, **kwargs):
+                pass
+
+            def write(self, *args, **kwargs):
+                pass
+
+        schema = CustomResultHandlerSchema()
+        obj = schema.load(schema.dump(Dummy()))
+        assert obj is None
+
+    def test_custom_schema_roundtrip_on_stateful_class(self):
+        class Stateful(ResultHandler):
+            def __init__(self, x):
+                self.x = x
+
+            def read(self, *args, **kwargs):
+                pass
+
+            def write(self, *args, **kwargs):
+                pass
+
+        schema = CustomResultHandlerSchema()
+        obj = schema.load(schema.dump(Stateful(42)))
+        assert obj is None
+
+    def test_result_handler_schema_defaults_to_custom(self):
+        class Weird(ResultHandler):
+            def __init__(self, y):
+                self.y = y
+
+            def read(self, *args, **kwargs):
+                return 99
+
+            def write(self, *args, **kwargs):
+                return type(None)
+
+        schema = ResultHandlerSchema()
+        serialized = schema.dump(Weird(dict(y="test")))
+        assert serialized["type"].endswith("Weird")
+        assert serialized["__version__"] == prefect.__version__
+
+        obj = schema.load(serialized)
+        assert obj is None
+
+
 class TestLocalResultHandler:
     def test_serialize_local_result_handler_with_no_dir(self):
         serialized = ResultHandlerSchema().dump(LocalResultHandler())