prefect-client 3.0.1__py3-none-any.whl → 3.0.3__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,
@@ -76,7 +77,7 @@ from prefect.states import (
     exception_to_failed_state,
     return_value_to_state,
 )
-from prefect.transactions import Transaction, transaction
+from prefect.transactions import IsolationLevel, Transaction, transaction
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.callables import call_with_parameters, parameters_to_args_kwargs
@@ -363,7 +364,6 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         new_state = Running()
 
         self.task_run.start_time = new_state.timestamp
-        self.task_run.run_count += 1
 
         flow_run_context = FlowRunContext.get()
         if flow_run_context and flow_run_context.flow_run:
@@ -411,6 +411,9 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self.task_run.state_type = new_state.type
         self.task_run.state_name = new_state.name
 
+        if new_state.is_running():
+            self.task_run.run_count += 1
+
         if new_state.is_final():
             if isinstance(state.data, BaseResult) and state.data.has_cached_object():
                 # Avoid fetching the result unless it is cached, otherwise we defeat
@@ -418,6 +421,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 +446,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 +460,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 +468,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,
             )
@@ -511,7 +513,6 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else:
                 delay = None
                 new_state = Retrying()
-                self.task_run.run_count += 1
 
             self.logger.info(
                 "Task run failed with exception: %r - " "Retry %s/%s will start %s",
@@ -538,12 +539,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 +595,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())
@@ -690,8 +693,9 @@ class SyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         if scheduled_time := self.state.state_details.scheduled_time:
             sleep_time = (scheduled_time - pendulum.now("utc")).total_seconds()
             await anyio.sleep(sleep_time if sleep_time > 0 else 0)
+            new_state = Retrying() if self.state.name == "AwaitingRetry" else Running()
             self.set_state(
-                Retrying() if self.state.name == "AwaitingRetry" else Running(),
+                new_state,
                 force=True,
             )
 
@@ -723,17 +727,21 @@ 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
+        isolation_level = (
+            IsolationLevel(self.task.cache_policy.isolation_level)
+            if self.task.cache_policy
+            and self.task.cache_policy is not NotSet
+            and self.task.cache_policy.isolation_level is not None
+            else 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(),
+            isolation_level=isolation_level,
         ) as txn:
             yield txn
 
@@ -769,10 +777,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)
@@ -877,7 +885,6 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         new_state = Running()
 
         self.task_run.start_time = new_state.timestamp
-        self.task_run.run_count += 1
 
         flow_run_context = FlowRunContext.get()
         if flow_run_context:
@@ -925,6 +932,9 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         self.task_run.state_type = new_state.type
         self.task_run.state_name = new_state.name
 
+        if new_state.is_running():
+            self.task_run.run_count += 1
+
         if new_state.is_final():
             if (
                 isinstance(new_state.data, BaseResult)
@@ -933,6 +943,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 +965,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 +979,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 +986,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,
         )
@@ -1021,7 +1030,6 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
             else:
                 delay = None
                 new_state = Retrying()
-                self.task_run.run_count += 1
 
             self.logger.info(
                 "Task run failed with exception: %r - " "Retry %s/%s will start %s",
@@ -1048,11 +1056,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 +1109,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())
@@ -1192,8 +1202,9 @@ class AsyncTaskRunEngine(BaseTaskRunEngine[P, R]):
         if scheduled_time := self.state.state_details.scheduled_time:
             sleep_time = (scheduled_time - pendulum.now("utc")).total_seconds()
             await anyio.sleep(sleep_time if sleep_time > 0 else 0)
+            new_state = Retrying() if self.state.name == "AwaitingRetry" else Running()
             await self.set_state(
-                Retrying() if self.state.name == "AwaitingRetry" else Running(),
+                new_state,
                 force=True,
             )
 
@@ -1226,17 +1237,21 @@ 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
+        isolation_level = (
+            IsolationLevel(self.task.cache_policy.isolation_level)
+            if self.task.cache_policy
+            and self.task.cache_policy is not NotSet
+            and self.task.cache_policy.isolation_level is not None
+            else 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(),
+            isolation_level=isolation_level,
         ) as txn:
             yield txn
 
@@ -1272,10 +1287,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)