Repository: PrefectHQ/prefect

Issue:
ResultHandlers should be serializable to JSON
From PIN2 (https://github.com/PrefectHQ/prefect/pull/570):

In order to attach information to `state.metadata` about how each attribute should be handled, all `ResultHandler`s need to be serializable with an associated versioned schema.  Currently there are only three result handlers - Local, Cloud, and GCS but soon there will be more.


Discussion:

Resolving patch:
diff --git a/src/prefect/engine/cloud/result_handler.py b/src/prefect/engine/cloud/result_handler.py
--- a/src/prefect/engine/cloud/result_handler.py
+++ b/src/prefect/engine/cloud/result_handler.py
@@ -17,11 +17,16 @@
 class CloudResultHandler(ResultHandler):
     """
     Hook for storing and retrieving task results from Prefect cloud storage.
+
+    Args:
+        - result_handler_service (str, optional): the location of the service
+            which will further process and store the results; if not provided, will default to
+            the value of `cloud.result_handler` in your config file
     """
 
-    def __init__(self) -> None:
+    def __init__(self, result_handler_service: str = None) -> None:
         self.client = None
-        self.result_handler_service = None
+        self.result_handler_service = result_handler_service
         super().__init__()
 
     def _initialize_client(self) -> None:
diff --git a/src/prefect/serialization/result_handlers.py b/src/prefect/serialization/result_handlers.py
new file mode 100644
--- /dev/null
+++ b/src/prefect/serialization/result_handlers.py
@@ -0,0 +1,49 @@
+import json
+from typing import Any, Dict
+
+from marshmallow import fields, post_load, ValidationError
+
+from prefect.engine.cloud.result_handler import CloudResultHandler
+from prefect.engine.result_handlers import ResultHandler, LocalResultHandler
+from prefect.utilities.serialization import (
+    JSONCompatible,
+    OneOfSchema,
+    VersionedSchema,
+    to_qualified_name,
+    version,
+)
+
+
+@version("0.4.0")
+class BaseResultHandlerSchema(VersionedSchema):
+    class Meta:
+        object_class = ResultHandler
+
+
+@version("0.4.0")
+class CloudResultHandlerSchema(BaseResultHandlerSchema):
+    class Meta:
+        object_class = CloudResultHandler
+
+    result_handler_service = fields.String(allow_none=True)
+
+
+@version("0.4.0")
+class LocalResultHandlerSchema(BaseResultHandlerSchema):
+    class Meta:
+        object_class = LocalResultHandler
+
+    dir = fields.String(allow_none=True)
+
+
+class ResultHandlerSchema(OneOfSchema):
+    """
+    Field that chooses between several nested schemas
+    """
+
+    # map class name to schema
+    type_schemas = {
+        "ResultHandler": BaseResultHandlerSchema,
+        "CloudResultHandler": CloudResultHandlerSchema,
+        "LocalResultHandler": LocalResultHandlerSchema,
+    }


Test patch:
diff --git a/tests/serialization/test_result_handlers.py b/tests/serialization/test_result_handlers.py
new file mode 100644
--- /dev/null
+++ b/tests/serialization/test_result_handlers.py
@@ -0,0 +1,73 @@
+import pytest
+
+from prefect.client import Client
+from prefect.engine.cloud.result_handler import CloudResultHandler
+from prefect.engine.result_handlers import ResultHandler, LocalResultHandler
+from prefect.serialization.result_handlers import ResultHandlerSchema
+
+
+class TestBaseResultHandler:
+    def test_serialize_base_result_handler(self):
+        serialized = ResultHandlerSchema().dump(ResultHandler())
+        assert isinstance(serialized, dict)
+        assert serialized["type"] == "ResultHandler"
+
+    def test_deserialize_base_result_handler(self):
+        schema = ResultHandlerSchema()
+        obj = schema.load(schema.dump(ResultHandler()))
+        assert isinstance(obj, ResultHandler)
+        assert hasattr(obj, "logger")
+        assert obj.logger.name == "prefect.ResultHandler"
+
+
+class TestLocalResultHandler:
+    def test_serialize_local_result_handler_with_no_dir(self):
+        serialized = ResultHandlerSchema().dump(LocalResultHandler())
+        assert isinstance(serialized, dict)
+        assert serialized["type"] == "LocalResultHandler"
+        assert serialized["dir"] is None
+
+    def test_serialize_local_result_handler_with_dir(self):
+        serialized = ResultHandlerSchema().dump(LocalResultHandler(dir="/root/prefect"))
+        assert isinstance(serialized, dict)
+        assert serialized["type"] == "LocalResultHandler"
+        assert serialized["dir"] == "/root/prefect"
+
+    @pytest.mark.parametrize("dir", [None, "/root/prefect"])
+    def test_deserialize_local_result_handler(self, dir):
+        schema = ResultHandlerSchema()
+        obj = schema.load(schema.dump(LocalResultHandler(dir=dir)))
+        assert isinstance(obj, LocalResultHandler)
+        assert hasattr(obj, "logger")
+        assert obj.logger.name == "prefect.LocalResultHandler"
+        assert obj.dir == dir
+
+
+class TestCloudResultHandler:
+    def test_serialize_with_no_attributes(self):
+        serialized = ResultHandlerSchema().dump(CloudResultHandler())
+        assert isinstance(serialized, dict)
+        assert serialized["type"] == "CloudResultHandler"
+        assert serialized["result_handler_service"] is None
+        assert "client" not in serialized
+
+    def test_serialize_with_attributes(self):
+        handler = CloudResultHandler(result_handler_service="http://foo.bar")
+        handler.client = Client()
+        serialized = ResultHandlerSchema().dump(handler)
+        assert isinstance(serialized, dict)
+        assert serialized["type"] == "CloudResultHandler"
+        assert serialized["result_handler_service"] == "http://foo.bar"
+        assert "client" not in serialized
+
+    @pytest.mark.parametrize("result_handler_service", [None, "http://foo.bar"])
+    def test_deserialize_local_result_handler(self, result_handler_service):
+        schema = ResultHandlerSchema()
+        handler = CloudResultHandler(result_handler_service=result_handler_service)
+        handler.client = Client()
+        obj = schema.load(schema.dump(handler))
+        assert isinstance(obj, CloudResultHandler)
+        assert hasattr(obj, "logger")
+        assert obj.logger.name == "prefect.CloudResultHandler"
+        assert obj.result_handler_service == result_handler_service
+        assert obj.client is None