Update server-side State.data to accept arbitrary JSON

src/prefect/orion/database/orm_models.py

@@ -132,7 +132,7 @@ def flow_run_id(cls):
         default=schemas.states.StateDetails,
         nullable=False,
     )
-    data = sa.Column(Pydantic(schemas.data.DataDocument), nullable=True)
+    data = sa.Column(sa.JSON, nullable=True)
 
     @declared_attr
     def flow_run(cls):
@@ -186,7 +186,7 @@ def task_run_id(cls):
         default=schemas.states.StateDetails,
         nullable=False,
     )
-    data = sa.Column(Pydantic(schemas.data.DataDocument), nullable=True)
+    data = sa.Column(sa.JSON, nullable=True)
 
     @declared_attr
     def task_run(cls):

src/prefect/orion/schemas/actions.py

@@ -137,7 +137,7 @@ class StateCreate(ActionBaseModel):
     type: schemas.states.StateType = FieldFrom(schemas.states.State)
     name: Optional[str] = FieldFrom(schemas.states.State)
     message: Optional[str] = FieldFrom(schemas.states.State)
-    data: Optional[schemas.data.DataDocument] = FieldFrom(schemas.states.State)
+    data: Optional[Any] = FieldFrom(schemas.states.State)
     state_details: schemas.states.StateDetails = FieldFrom(schemas.states.State)
 
     # DEPRECATED

src/prefect/orion/schemas/states.py

@@ -4,13 +4,12 @@
 
 import datetime
 import warnings
-from typing import Generic, Optional, Type, TypeVar
+from typing import Any, Generic, Optional, Type, TypeVar
 from uuid import UUID
 
 import pendulum
 from pydantic import Field, root_validator, validator
 
-from prefect.orion.schemas.data import DataDocument
 from prefect.orion.utilities.schemas import DateTimeTZ, IDBaseModel, PrefectBaseModel
 from prefect.utilities.collections import AutoEnum
 
@@ -58,8 +57,12 @@ class Config:
     name: Optional[str] = Field(default=None)
     timestamp: DateTimeTZ = Field(default_factory=lambda: pendulum.now("UTC"))
     message: Optional[str] = Field(default=None, example="Run started")
-    data: Optional[DataDocument[R]] = Field(
+    data: Optional[Any] = Field(
         default=None,
+        description=(
+            "Data associated with the state, e.g. a result. "
+            "Content must be storable as JSON."
+        ),
     )
     state_details: StateDetails = Field(default_factory=StateDetails)
 

tests/orion/api/test_flow_runs.py

@@ -664,6 +664,27 @@ async def test_set_flow_run_state_force_skips_orchestration(
         assert response2.status_code == status.HTTP_201_CREATED
         assert response2.json()["status"] == "ACCEPT"
 
+    @pytest.mark.parametrize("data", [1, "test", {"foo": "bar"}])
+    async def test_set_flow_run_state_accepts_any_jsonable_data(
+        self, flow_run, client, session, data
+    ):
+        response = await client.post(
+            f"/flow_runs/{flow_run.id}/set_state",
+            json=dict(state=dict(type="COMPLETED", data=data)),
+        )
+        assert response.status_code == 201
+
+        api_response = OrchestrationResult.parse_obj(response.json())
+        assert api_response.status == responses.SetStateStatus.ACCEPT
+
+        flow_run_id = flow_run.id
+        session.expire(flow_run)
+
+        run = await models.flow_runs.read_flow_run(
+            session=session, flow_run_id=flow_run_id
+        )
+        assert run.state.data == data
+
 
 class TestFlowRunHistory:
     async def test_history_interval_must_be_one_second_or_larger(self, client):