prefect-client 3.0.1__py3-none-any.whl → 3.0.2__py3-none-any.whl

prefect/states.py

@@ -25,7 +25,13 @@ from prefect.exceptions import (
     UnfinishedRun,
 )
 from prefect.logging.loggers import get_logger, get_run_logger
-from prefect.results import BaseResult, R, ResultStore
+from prefect.results import (
+    BaseResult,
+    R,
+    ResultRecord,
+    ResultRecordMetadata,
+    ResultStore,
+)
 from prefect.settings import PREFECT_ASYNC_FETCH_STATE_RESULT
 from prefect.utilities.annotations import BaseAnnotation
 from prefect.utilities.asyncutils import in_async_main_thread, sync_compatible
@@ -92,7 +98,11 @@ async def _get_state_result_data_with_retries(
 
     for i in range(1, max_attempts + 1):
         try:
-            return await state.data.get()
+            if isinstance(state.data, ResultRecordMetadata):
+                record = await ResultRecord._from_metadata(state.data)
+                return record.result
+            else:
+                return await state.data.get()
         except Exception as e:
             if i == max_attempts:
                 raise
@@ -127,10 +137,12 @@ async def _get_state_result(
     ):
         raise await get_state_exception(state)
 
-    if isinstance(state.data, BaseResult):
+    if isinstance(state.data, (BaseResult, ResultRecordMetadata)):
         result = await _get_state_result_data_with_retries(
             state, retry_result_failure=retry_result_failure
         )
+    elif isinstance(state.data, ResultRecord):
+        result = state.data.result
 
     elif state.data is None:
         if state.is_failed() or state.is_crashed() or state.is_cancelled():
@@ -207,7 +219,7 @@ async def exception_to_crashed_state(
         )
 
     if result_store:
-        data = await result_store.create_result(exc)
+        data = result_store.create_result_record(exc)
     else:
         # Attach the exception for local usage, will not be available when retrieved
         # from the API
@@ -240,10 +252,10 @@ async def exception_to_failed_state(
         pass
 
     if result_store:
-        data = await result_store.create_result(exc)
+        data = result_store.create_result_record(exc)
         if write_result:
             try:
-                await data.write()
+                await result_store.apersist_result_record(data)
             except Exception as exc:
                 local_logger.warning(
                     "Failed to write result: %s Execution will continue, but the result has not been written",
@@ -309,21 +321,21 @@ async def return_value_to_state(
         state = retval
         # Unless the user has already constructed a result explicitly, use the store
         # to update the data to the correct type
-        if not isinstance(state.data, BaseResult):
-            result = await result_store.create_result(
+        if not isinstance(state.data, (BaseResult, ResultRecord, ResultRecordMetadata)):
+            result_record = result_store.create_result_record(
                 state.data,
                 key=key,
                 expiration=expiration,
             )
             if write_result:
                 try:
-                    await result.write()
+                    await result_store.apersist_result_record(result_record)
                 except Exception as exc:
                     local_logger.warning(
                         "Encountered an error while persisting result: %s Execution will continue, but the result has not been persisted",
                         exc,
                     )
-            state.data = result
+            state.data = result_record
         return state
 
     # Determine a new state from the aggregate of contained states
@@ -359,14 +371,14 @@ async def return_value_to_state(
         # TODO: We may actually want to set the data to a `StateGroup` object and just
         #       allow it to be unpacked into a tuple and such so users can interact with
         #       it
-        result = await result_store.create_result(
+        result_record = result_store.create_result_record(
             retval,
             key=key,
             expiration=expiration,
         )
         if write_result:
             try:
-                await result.write()
+                await result_store.apersist_result_record(result_record)
             except Exception as exc:
                 local_logger.warning(
                     "Encountered an error while persisting result: %s Execution will continue, but the result has not been persisted",
@@ -375,7 +387,7 @@ async def return_value_to_state(
         return State(
             type=new_state_type,
             message=message,
-            data=result,
+            data=result_record,
         )
 
     # Generators aren't portable, implicitly convert them to a list.
@@ -385,23 +397,23 @@ async def return_value_to_state(
         data = retval
 
     # Otherwise, they just gave data and this is a completed retval
-    if isinstance(data, BaseResult):
+    if isinstance(data, (BaseResult, ResultRecord)):
         return Completed(data=data)
     else:
-        result = await result_store.create_result(
+        result_record = result_store.create_result_record(
             data,
             key=key,
             expiration=expiration,
         )
         if write_result:
             try:
-                await result.write()
+                await result_store.apersist_result_record(result_record)
             except Exception as exc:
                 local_logger.warning(
                     "Encountered an error while persisting result: %s Execution will continue, but the result has not been persisted",
                     exc,
                 )
-        return Completed(data=result)
+        return Completed(data=result_record)
 
 
 @sync_compatible
@@ -442,6 +454,11 @@ async def get_state_exception(state: State) -> BaseException:
 
     if isinstance(state.data, BaseResult):
         result = await _get_state_result_data_with_retries(state)
+    elif isinstance(state.data, ResultRecord):
+        result = state.data.result
+    elif isinstance(state.data, ResultRecordMetadata):
+        record = await ResultRecord._from_metadata(state.data)
+        result = record.result
     elif state.data is None:
         result = None
     else:

prefect/task_engine.py

@@ -55,11 +55,12 @@ from prefect.exceptions import (
 )
 from prefect.futures import PrefectFuture
 from prefect.logging.loggers import get_logger, patch_print, task_run_logger
-from prefect.records.result_store import ResultRecordStore
 from prefect.results import (
     BaseResult,
+    ResultRecord,
     _format_user_supplied_storage_key,
-    get_current_result_store,
+    get_result_store,
+    should_persist_result,
 )
 from prefect.settings import (
     PREFECT_DEBUG_MODE,
@@ -418,6 +419,8 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 result = state.result(raise_on_failure=False, fetch=True)
                 if inspect.isawaitable(result):
                     result = run_coro_as_sync(result)
+            elif isinstance(state.data, ResultRecord):
+                result = state.data.result
             else:
                 result = state.data
 
@@ -441,7 +444,8 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 if inspect.isawaitable(_result):
                     _result = run_coro_as_sync(_result)
                 return _result
-
+            elif isinstance(self._return_value, ResultRecord):
+                return self._return_value.result
             # otherwise, return the value as is
             return self._return_value
 
@@ -454,10 +458,6 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             return self._raised
 
     def handle_success(self, result: R, transaction: Transaction) -> R:
-        result_store = getattr(TaskRunContext.get(), "result_store", None)
-        if result_store is None:
-            raise ValueError("Result store is not set")
-
         if self.task.cache_expiration is not None:
             expiration = pendulum.now("utc") + self.task.cache_expiration
         else:
@@ -466,7 +466,7 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         terminal_state = run_coro_as_sync(
             return_value_to_state(
                 result,
-                result_store=result_store,
+                result_store=get_result_store(),
                 key=transaction.key,
                 expiration=expiration,
             )
@@ -538,12 +538,11 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         # If the task fails, and we have retries left, set the task to retrying.
         if not self.handle_retry(exc):
             # If the task has no retries left, or the retry condition is not met, set the task to failed.
-            context = TaskRunContext.get()
             state = run_coro_as_sync(
                 exception_to_failed_state(
                     exc,
                     message="Task run encountered an exception",
-                    result_store=getattr(context, "result_store", None),
+                    result_store=get_result_store(),
                     write_result=True,
                 )
             )
@@ -595,10 +594,13 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                     log_prints=log_prints,
                     task_run=self.task_run,
                     parameters=self.parameters,
-                    result_store=get_current_result_store().update_for_task(
+                    result_store=get_result_store().update_for_task(
                         self.task, _sync=True
                     ),
                     client=client,
+                    persist_result=self.task.persist_result
+                    if self.task.persist_result is not None
+                    else should_persist_result(),
                 )
             )
             stack.enter_context(ConcurrencyContextV1())
@@ -723,17 +725,12 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else PREFECT_TASKS_REFRESH_CACHE.value()
         )
 
-        result_store = getattr(TaskRunContext.get(), "result_store", None)
-        if result_store and result_store.persist_result:
-            store = ResultRecordStore(result_store=result_store)
-        else:
-            store = None
-
         with transaction(
             key=self.compute_transaction_key(),
-            store=store,
+            store=get_result_store(),
             overwrite=overwrite,
             logger=self.logger,
+            write_on_commit=should_persist_result(),
         ) as txn:
             yield txn
 
@@ -769,10 +766,10 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         if transaction.is_committed():
             result = transaction.read()
         else:
-            if self.task.tags:
+            if self.task_run.tags:
                 # Acquire a concurrency slot for each tag, but only if a limit
                 # matching the tag already exists.
-                with concurrency(list(self.task.tags), self.task_run.id):
+                with concurrency(list(self.task_run.tags), self.task_run.id):
                     result = call_with_parameters(self.task.fn, parameters)
             else:
                 result = call_with_parameters(self.task.fn, parameters)
@@ -933,6 +930,8 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                 # Avoid fetching the result unless it is cached, otherwise we defeat
                 # the purpose of disabling `cache_result_in_memory`
                 result = await new_state.result(raise_on_failure=False, fetch=True)
+            elif isinstance(new_state.data, ResultRecord):
+                result = new_state.data.result
             else:
                 result = new_state.data
 
@@ -953,7 +952,8 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             # if the return value is a BaseResult, we need to fetch it
             if isinstance(self._return_value, BaseResult):
                 return await self._return_value.get()
-
+            elif isinstance(self._return_value, ResultRecord):
+                return self._return_value.result
             # otherwise, return the value as is
             return self._return_value
 
@@ -966,10 +966,6 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             return self._raised
 
     async def handle_success(self, result: R, transaction: Transaction) -> R:
-        result_store = getattr(TaskRunContext.get(), "result_store", None)
-        if result_store is None:
-            raise ValueError("Result store is not set")
-
         if self.task.cache_expiration is not None:
             expiration = pendulum.now("utc") + self.task.cache_expiration
         else:
@@ -977,7 +973,7 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
 
         terminal_state = await return_value_to_state(
             result,
-            result_store=result_store,
+            result_store=get_result_store(),
             key=transaction.key,
             expiration=expiration,
         )
@@ -1048,11 +1044,10 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         # If the task fails, and we have retries left, set the task to retrying.
         if not await self.handle_retry(exc):
             # If the task has no retries left, or the retry condition is not met, set the task to failed.
-            context = TaskRunContext.get()
             state = await exception_to_failed_state(
                 exc,
                 message="Task run encountered an exception",
-                result_store=getattr(context, "result_store", None),
+                result_store=get_result_store(),
             )
             self.record_terminal_state_timing(state)
             await self.set_state(state)
@@ -1102,10 +1097,13 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
                     log_prints=log_prints,
                     task_run=self.task_run,
                     parameters=self.parameters,
-                    result_store=await get_current_result_store().update_for_task(
+                    result_store=await get_result_store().update_for_task(
                         self.task, _sync=False
                     ),
                     client=client,
+                    persist_result=self.task.persist_result
+                    if self.task.persist_result is not None
+                    else should_persist_result(),
                 )
             )
             stack.enter_context(ConcurrencyContext())
@@ -1226,17 +1224,13 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             if self.task.refresh_cache is not None
             else PREFECT_TASKS_REFRESH_CACHE.value()
         )
-        result_store = getattr(TaskRunContext.get(), "result_store", None)
-        if result_store and result_store.persist_result:
-            store = ResultRecordStore(result_store=result_store)
-        else:
-            store = None
 
         with transaction(
             key=self.compute_transaction_key(),
-            store=store,
+            store=get_result_store(),
             overwrite=overwrite,
             logger=self.logger,
+            write_on_commit=should_persist_result(),
         ) as txn:
             yield txn
 
@@ -1272,10 +1266,10 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         if transaction.is_committed():
             result = transaction.read()
         else:
-            if self.task.tags:
+            if self.task_run.tags:
                 # Acquire a concurrency slot for each tag, but only if a limit
                 # matching the tag already exists.
-                async with aconcurrency(list(self.task.tags), self.task_run.id):
+                async with aconcurrency(list(self.task_run.tags), self.task_run.id):
                     result = await call_with_parameters(self.task.fn, parameters)
             else:
                 result = await call_with_parameters(self.task.fn, parameters)

prefect/transactions.py

@@ -22,7 +22,13 @@ from prefect.exceptions import MissingContextError, SerializationError
 from prefect.logging.loggers import get_logger, get_run_logger
 from prefect.records import RecordStore
 from prefect.records.base import TransactionRecord
-from prefect.results import BaseResult, ResultRecord, ResultStore
+from prefect.results import (
+    BaseResult,
+    ResultRecord,
+    ResultStore,
+    get_result_store,
+    should_persist_result,
+)
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.collections import AutoEnum
 from prefect.utilities.engine import _get_hook_name
@@ -66,19 +72,91 @@ class Transaction(ContextModel):
     logger: Union[logging.Logger, logging.LoggerAdapter] = Field(
         default_factory=partial(get_logger, "transactions")
     )
+    write_on_commit: bool = True
     _stored_values: Dict[str, Any] = PrivateAttr(default_factory=dict)
     _staged_value: Any = None
     __var__: ContextVar = ContextVar("transaction")
 
     def set(self, name: str, value: Any) -> None:
+        """
+        Set a stored value in the transaction.
+
+        Args:
+            name: The name of the value to set
+            value: The value to set
+
+        Examples:
+            Set a value for use later in the transaction:
+            ```python
+            with transaction() as txn:
+                txn.set("key", "value")
+                ...
+                assert txn.get("key") == "value"
+            ```
+        """
         self._stored_values[name] = value
 
     def get(self, name: str, default: Any = NotSet) -> Any:
-        if name not in self._stored_values:
-            if default is not NotSet:
-                return default
-            raise ValueError(f"Could not retrieve value for unknown key: {name}")
-        return self._stored_values.get(name)
+        """
+        Get a stored value from the transaction.
+
+        Child transactions will return values from their parents unless a value with
+        the same name is set in the child transaction.
+
+        Direct changes to returned values will not update the stored value. To update the
+        stored value, use the `set` method.
+
+        Args:
+            name: The name of the value to get
+            default: The default value to return if the value is not found
+
+        Returns:
+            The value from the transaction
+
+        Examples:
+            Get a value from the transaction:
+            ```python
+            with transaction() as txn:
+                txn.set("key", "value")
+                ...
+                assert txn.get("key") == "value"
+            ```
+
+            Get a value from a parent transaction:
+            ```python
+            with transaction() as parent:
+                parent.set("key", "parent_value")
+                with transaction() as child:
+                    assert child.get("key") == "parent_value"
+            ```
+
+            Update a stored value:
+            ```python
+            with transaction() as txn:
+                txn.set("key", [1, 2, 3])
+                value = txn.get("key")
+                value.append(4)
+                # Stored value is not updated until `.set` is called
+                assert value == [1, 2, 3, 4]
+                assert txn.get("key") == [1, 2, 3]
+
+                txn.set("key", value)
+                assert txn.get("key") == [1, 2, 3, 4]
+            ```
+        """
+        # deepcopy to prevent mutation of stored values
+        value = copy.deepcopy(self._stored_values.get(name, NotSet))
+        if value is NotSet:
+            # if there's a parent transaction, get the value from the parent
+            parent = self.get_parent()
+            if parent is not None:
+                value = parent.get(name, default)
+            # if there's no parent transaction, use the default
+            elif default is not NotSet:
+                value = default
+            else:
+                raise ValueError(f"Could not retrieve value for unknown key: {name}")
+        return value
 
     def is_committed(self) -> bool:
         return self.state == TransactionState.COMMITTED
@@ -101,8 +179,6 @@ class Transaction(ContextModel):
                 "Context already entered. Context enter calls cannot be nested."
             )
         parent = get_transaction()
-        if parent:
-            self._stored_values = copy.deepcopy(parent._stored_values)
         # set default commit behavior; either inherit from parent or set a default of eager
         if self.commit_mode is None:
             self.commit_mode = parent.commit_mode if parent else CommitMode.LAZY
@@ -119,7 +195,7 @@ class Transaction(ContextModel):
             and not self.store.supports_isolation_level(self.isolation_level)
         ):
             raise ValueError(
-                f"Isolation level {self.isolation_level.name} is not supported by record store type {self.store.__class__.__name__}"
+                f"Isolation level {self.isolation_level.name} is not supported by provided result store."
             )
 
         # this needs to go before begin, which could set the state to committed
@@ -229,14 +305,21 @@ class Transaction(ContextModel):
             for hook in self.on_commit_hooks:
                 self.run_hook(hook, "commit")
 
-            if self.store and self.key:
+            if self.store and self.key and self.write_on_commit:
                 if isinstance(self.store, ResultStore):
                     if isinstance(self._staged_value, BaseResult):
-                        self.store.write(self.key, self._staged_value.get(_sync=True))
+                        self.store.write(
+                            key=self.key, obj=self._staged_value.get(_sync=True)
+                        )
+                    elif isinstance(self._staged_value, ResultRecord):
+                        self.store.persist_result_record(
+                            result_record=self._staged_value
+                        )
                     else:
-                        self.store.write(self.key, self._staged_value)
+                        self.store.write(key=self.key, obj=self._staged_value)
                 else:
-                    self.store.write(self.key, self._staged_value)
+                    self.store.write(key=self.key, result=self._staged_value)
+
             self.state = TransactionState.COMMITTED
             if (
                 self.store
@@ -287,7 +370,7 @@ class Transaction(ContextModel):
 
     def stage(
         self,
-        value: Union["BaseResult", Any],
+        value: Any,
         on_rollback_hooks: Optional[List] = None,
         on_commit_hooks: Optional[List] = None,
     ) -> None:
@@ -349,6 +432,7 @@ def transaction(
     commit_mode: Optional[CommitMode] = None,
     isolation_level: Optional[IsolationLevel] = None,
     overwrite: bool = False,
+    write_on_commit: Optional[bool] = None,
     logger: Union[logging.Logger, logging.LoggerAdapter, None] = None,
 ) -> Generator[Transaction, None, None]:
     """
@@ -361,48 +445,16 @@ def transaction(
         - commit_mode: The commit mode controlling when the transaction and
             child transactions are committed
         - overwrite: Whether to overwrite an existing transaction record in the store
+        - write_on_commit: Whether to write the result to the store on commit. If not provided,
+            will default will be determined by the current run context. If no run context is
+            available, the value of `PREFECT_RESULTS_PERSIST_BY_DEFAULT` will be used.
 
     Yields:
         - Transaction: An object representing the transaction state
     """
     # if there is no key, we won't persist a record
     if key and not store:
-        from prefect.context import FlowRunContext, TaskRunContext
-        from prefect.results import ResultStore, get_default_result_storage
-
-        flow_run_context = FlowRunContext.get()
-        task_run_context = TaskRunContext.get()
-        existing_store = getattr(task_run_context, "result_store", None) or getattr(
-            flow_run_context, "result_store", None
-        )
-
-        new_store: ResultStore
-        if existing_store and existing_store.result_storage_block_id:
-            new_store = existing_store.model_copy(
-                update={
-                    "persist_result": True,
-                }
-            )
-        else:
-            default_storage = get_default_result_storage(_sync=True)
-            if existing_store:
-                new_store = existing_store.model_copy(
-                    update={
-                        "persist_result": True,
-                        "storage_block": default_storage,
-                        "storage_block_id": default_storage._block_document_id,
-                    }
-                )
-            else:
-                new_store = ResultStore(
-                    persist_result=True,
-                    result_storage=default_storage,
-                )
-        from prefect.records.result_store import ResultRecordStore
-
-        store = ResultRecordStore(
-            result_store=new_store,
-        )
+        store = get_result_store()
 
     try:
         logger = logger or get_run_logger()
@@ -415,6 +467,9 @@ def transaction(
         commit_mode=commit_mode,
         isolation_level=isolation_level,
         overwrite=overwrite,
+        write_on_commit=write_on_commit
+        if write_on_commit is not None
+        else should_persist_result(),
         logger=logger,
     ) as txn:
         yield txn

prefect/utilities/engine.py

@@ -44,12 +44,11 @@ from prefect.exceptions import (
 )
 from prefect.flows import Flow
 from prefect.futures import PrefectFuture
-from prefect.futures import PrefectFuture as NewPrefectFuture
 from prefect.logging.loggers import (
     get_logger,
     task_run_logger,
 )
-from prefect.results import BaseResult
+from prefect.results import BaseResult, ResultRecord, should_persist_result
 from prefect.settings import (
     PREFECT_LOGGING_LOG_PRINTS,
 )
@@ -122,7 +121,7 @@ async def collect_task_run_inputs(expr: Any, max_depth: int = -1) -> Set[TaskRun
 
 
 def collect_task_run_inputs_sync(
-    expr: Any, future_cls: Any = NewPrefectFuture, max_depth: int = -1
+    expr: Any, future_cls: Any = PrefectFuture, max_depth: int = -1
 ) -> Set[TaskRunInput]:
     """
     This function recurses through an expression to generate a set of any discernible
@@ -131,7 +130,7 @@ def collect_task_run_inputs_sync(
 
     Examples:
         >>> task_inputs = {
-        >>>    k: collect_task_run_inputs(v) for k, v in parameters.items()
+        >>>    k: collect_task_run_inputs_sync(v) for k, v in parameters.items()
         >>> }
     """
     # TODO: This function needs to be updated to detect parameters and constants
@@ -401,6 +400,8 @@ async def propose_state(
             # Avoid fetching the result unless it is cached, otherwise we defeat
             # the purpose of disabling `cache_result_in_memory`
             result = await state.result(raise_on_failure=False, fetch=True)
+        elif isinstance(state.data, ResultRecord):
+            result = state.data.result
         else:
             result = state.data
 
@@ -504,6 +505,8 @@ def propose_state_sync(
             result = state.result(raise_on_failure=False, fetch=True)
             if inspect.isawaitable(result):
                 result = run_coro_as_sync(result)
+        elif isinstance(state.data, ResultRecord):
+            result = state.data.result
         else:
             result = state.data
 
@@ -732,6 +735,13 @@ def emit_task_run_state_change_event(
 ) -> Event:
     state_message_truncation_length = 100_000
 
+    if isinstance(validated_state.data, ResultRecord) and should_persist_result():
+        data = validated_state.data.metadata.model_dump(mode="json")
+    elif isinstance(validated_state.data, BaseResult):
+        data = validated_state.data.model_dump(mode="json")
+    else:
+        data = None
+
     return emit_event(
         id=validated_state.id,
         occurred=validated_state.timestamp,
@@ -770,9 +780,7 @@ def emit_task_run_state_change_event(
                     exclude_unset=True,
                     exclude={"flow_run_id", "task_run_id"},
                 ),
-                "data": validated_state.data.model_dump(mode="json")
-                if isinstance(validated_state.data, BaseResult)
-                else None,
+                "data": data,
             },
             "task_run": task_run.model_dump(
                 mode="json",
@@ -822,7 +830,7 @@ def resolve_to_final_result(expr, context):
     if isinstance(context.get("annotation"), quote):
         raise StopVisiting()
 
-    if isinstance(expr, NewPrefectFuture):
+    if isinstance(expr, PrefectFuture):
         upstream_task_run = context.get("current_task_run")
         upstream_task = context.get("current_task")
         if (