From e845f7e8042e0d34263ed759e97f97ae6d632bd6 Mon Sep 17 00:00:00 2001 From: nicholas Date: Tue, 28 Feb 2023 12:36:31 -0500 Subject: [PATCH] Add flow run id to task run result artifact, add tests for both --- src/prefect/server/orchestration/rules.py | 4 ++++ tests/server/orchestration/test_rules.py | 6 ++++++ 2 files changed, 10 insertions(+) diff --git a/src/prefect/server/orchestration/rules.py b/src/prefect/server/orchestration/rules.py index 35b54b7705a1..5f2e03fcdb0e 100644 --- a/src/prefect/server/orchestration/rules.py +++ b/src/prefect/server/orchestration/rules.py @@ -392,6 +392,10 @@ async def _validate_proposed_state( if state_data is not None: state_result_artifact = core.Artifact.from_result(state_data) state_result_artifact.task_run_id = self.run.id + + flow_run = await self.flow_run() + state_result_artifact.flow_run_id = flow_run.id + await artifacts.create_artifact(self.session, state_result_artifact) state_payload["result_artifact_id"] = state_result_artifact.id diff --git a/tests/server/orchestration/test_rules.py b/tests/server/orchestration/test_rules.py index daf48d97c13a..0232cb8f4c09 100644 --- a/tests/server/orchestration/test_rules.py +++ b/tests/server/orchestration/test_rules.py @@ -1461,6 +1461,12 @@ async def test_context_validation_writes_result_artifact( orm_artifact = await models.artifacts.read_artifact(ctx.session, artifact_id) assert orm_artifact.data == {"value": "some special data"} + if run_type == "task": + assert orm_artifact.task_run_id == ctx.run.id + assert orm_artifact.flow_run_id == ctx.run.flow_run_id + else: + assert orm_artifact.flow_run_id == ctx.run.id + async def test_context_validation_writes_result_artifact_with_metadata( self, session, run_type, initialize_orchestration ):