Avoid result factories if persistence is turned off for a task (#15055)

PrefectHQ · Aug 23, 2024 · 0e4c2d9 · 0e4c2d9
1 parent 2f000f9
commit 0e4c2d9
Showing 1 changed file with 15 additions and 6 deletions.
diff --git a/src/prefect/task_engine.py b/src/prefect/task_engine.py
@@ -705,17 +705,22 @@ def start(
 
     @contextmanager
     def transaction_context(self) -> Generator[Transaction, None, None]:
-        result_factory = getattr(TaskRunContext.get(), "result_factory", None)
-
         # refresh cache setting is now repurposes as overwrite transaction record
         overwrite = (
             self.task.refresh_cache
             if self.task.refresh_cache is not None
             else PREFECT_TASKS_REFRESH_CACHE.value()
         )
+
+        result_factory = getattr(TaskRunContext.get(), "result_factory", None)
+        if result_factory and result_factory.persist_result:
+            store = ResultFactoryStore(result_factory=result_factory)
+        else:
+            store = None
+
         with transaction(
             key=self.compute_transaction_key(),
-            store=ResultFactoryStore(result_factory=result_factory),
+            store=store,
             overwrite=overwrite,
             logger=self.logger,
         ) as txn:
@@ -1199,17 +1204,21 @@ async def start(
 
     @asynccontextmanager
     async def transaction_context(self) -> AsyncGenerator[Transaction, None]:
-        result_factory = getattr(TaskRunContext.get(), "result_factory", None)
-
         # refresh cache setting is now repurposes as overwrite transaction record
         overwrite = (
             self.task.refresh_cache
             if self.task.refresh_cache is not None
             else PREFECT_TASKS_REFRESH_CACHE.value()
         )
+        result_factory = getattr(TaskRunContext.get(), "result_factory", None)
+        if result_factory and result_factory.persist_result:
+            store = ResultFactoryStore(result_factory=result_factory)
+        else:
+            store = None
+
         with transaction(
             key=self.compute_transaction_key(),
-            store=ResultFactoryStore(result_factory=result_factory),
+            store=store,
             overwrite=overwrite,
             logger=self.logger,
         ) as txn: