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

prefect/transactions.py

@@ -18,11 +18,21 @@ from pydantic import Field, PrivateAttr
 from typing_extensions import Self
 
 from prefect.context import ContextModel
-from prefect.exceptions import MissingContextError, SerializationError
+from prefect.exceptions import (
+    ConfigurationError,
+    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 +76,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 +183,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
@@ -118,8 +198,10 @@ class Transaction(ContextModel):
             and self.key
             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__}"
+            raise ConfigurationError(
+                f"Isolation level {self.isolation_level.name} is not supported by provided "
+                "configuration. Please ensure you've provided a lock file directory or lock "
+                "manager when using the SERIALIZABLE isolation level."
             )
 
         # this needs to go before begin, which could set the state to committed
@@ -229,14 +311,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 +376,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 +438,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 +451,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 +473,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/asyncutils.py

@@ -341,6 +341,13 @@ def sync_compatible(
         will submit the async method to the event loop.
     - If we cannot find an event loop, we will create a new one and run the async method
         then tear down the loop.
+
+    Note: Type checkers will infer functions decorated with `@sync_compatible` are synchronous. If
+    you want to use the decorated function in an async context, you will need to ignore the types
+    and "cast" the return type to a coroutine. For example:
+    ```
+    python result: Coroutine = sync_compatible(my_async_function)(arg1, arg2) # type: ignore
+    ```
     """
 
     @wraps(async_fn)

prefect/utilities/collections.py

@@ -18,7 +18,6 @@ from typing import (
     Generator,
     Hashable,
     Iterable,
-    Iterator,
     List,
     Optional,
     Set,
@@ -192,7 +191,9 @@ def extract_instances(
     return ret
 
 
-def batched_iterable(iterable: Iterable[T], size: int) -> Iterator[Tuple[T, ...]]:
+def batched_iterable(
+    iterable: Iterable[T], size: int
+) -> Generator[Tuple[T, ...], None, None]:
     """
     Yield batches of a certain size from an iterable
 

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
 
@@ -624,6 +627,9 @@ def link_state_to_result(state: State, result: Any) -> None:
     """
 
     flow_run_context = FlowRunContext.get()
+    # Drop the data field to avoid holding a strong reference to the result
+    # Holding large user objects in memory can cause memory bloat
+    linked_state = state.model_copy(update={"data": None})
 
     def link_if_trackable(obj: Any) -> None:
         """Track connection between a task run result and its associated state if it has a unique ID.
@@ -640,7 +646,7 @@ def link_state_to_result(state: State, result: Any) -> None:
         ):
             state.state_details.untrackable_result = True
             return
-        flow_run_context.task_run_results[id(obj)] = state
+        flow_run_context.task_run_results[id(obj)] = linked_state
 
     if flow_run_context:
         visit_collection(expr=result, visit_fn=link_if_trackable, max_depth=1)
@@ -732,6 +738,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 +783,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 +833,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 (

prefect/utilities/importtools.py

@@ -398,6 +398,7 @@ def safe_load_namespace(
         # Save original sys.path and modify it
         original_sys_path = sys.path.copy()
         sys.path.insert(0, parent_dir)
+        sys.path.insert(0, file_dir)
 
         # Create a temporary module for import context
         temp_module = ModuleType(package_name)

prefect/utilities/urls.py

@@ -1,17 +1,22 @@
 import inspect
+import ipaddress
+import socket
 import urllib.parse
-from typing import Any, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
+from urllib.parse import urlparse
 from uuid import UUID
 
 from pydantic import BaseModel
 
 from prefect import settings
-from prefect.blocks.core import Block
-from prefect.events.schemas.automations import Automation
-from prefect.events.schemas.events import ReceivedEvent, Resource
-from prefect.futures import PrefectFuture
 from prefect.logging.loggers import get_logger
-from prefect.variables import Variable
+
+if TYPE_CHECKING:
+    from prefect.blocks.core import Block
+    from prefect.events.schemas.automations import Automation
+    from prefect.events.schemas.events import ReceivedEvent, Resource
+    from prefect.futures import PrefectFuture
+    from prefect.variables import Variable
 
 logger = get_logger("utilities.urls")
 
@@ -58,6 +63,54 @@ URLType = Literal["ui", "api"]
 RUN_TYPES = {"flow-run", "task-run"}
 
 
+def validate_restricted_url(url: str):
+    """
+    Validate that the provided URL is safe for outbound requests.  This prevents
+    attacks like SSRF (Server Side Request Forgery), where an attacker can make
+    requests to internal services (like the GCP metadata service, localhost addresses,
+    or in-cluster Kubernetes services)
+
+    Args:
+        url: The URL to validate.
+
+    Raises:
+        ValueError: If the URL is a restricted URL.
+    """
+
+    try:
+        parsed_url = urlparse(url)
+    except ValueError:
+        raise ValueError(f"{url!r} is not a valid URL.")
+
+    if parsed_url.scheme not in ("http", "https"):
+        raise ValueError(
+            f"{url!r} is not a valid URL.  Only HTTP and HTTPS URLs are allowed."
+        )
+
+    hostname = parsed_url.hostname or ""
+
+    # Remove IPv6 brackets if present
+    if hostname.startswith("[") and hostname.endswith("]"):
+        hostname = hostname[1:-1]
+
+    if not hostname:
+        raise ValueError(f"{url!r} is not a valid URL.")
+
+    try:
+        ip_address = socket.gethostbyname(hostname)
+        ip = ipaddress.ip_address(ip_address)
+    except socket.gaierror:
+        try:
+            ip = ipaddress.ip_address(hostname)
+        except ValueError:
+            raise ValueError(f"{url!r} is not a valid URL.  It could not be resolved.")
+
+    if ip.is_private:
+        raise ValueError(
+            f"{url!r} is not a valid URL.  It resolves to the private address {ip}."
+        )
+
+
 def convert_class_to_name(obj: Any) -> str:
     """
     Convert CamelCase class name to dash-separated lowercase name
@@ -69,12 +122,12 @@ def convert_class_to_name(obj: Any) -> str:
 
 def url_for(
     obj: Union[
-        PrefectFuture,
-        Block,
-        Variable,
-        Automation,
-        Resource,
-        ReceivedEvent,
+        "PrefectFuture",
+        "Block",
+        "Variable",
+        "Automation",
+        "Resource",
+        "ReceivedEvent",
         BaseModel,
         str,
     ],
@@ -105,6 +158,11 @@ def url_for(
         url_for(obj=my_flow_run)
         url_for("flow-run", obj_id="123e4567-e89b-12d3-a456-426614174000")
     """
+    from prefect.blocks.core import Block
+    from prefect.events.schemas.automations import Automation
+    from prefect.events.schemas.events import ReceivedEvent, Resource
+    from prefect.futures import PrefectFuture
+
     if isinstance(obj, PrefectFuture):
         name = "task-run"
     elif isinstance(obj, Block):

prefect/workers/base.py

@@ -146,6 +146,12 @@ class BaseJobConfiguration(BaseModel):
         )
         variables.update(values)
 
+        # deep merge `env`
+        if isinstance(job_config.get("env"), dict) and (
+            hardcoded_env := variables.get("env")
+        ):
+            job_config["env"] = hardcoded_env | job_config.get("env")
+
         populated_configuration = apply_values(template=job_config, values=variables)
         populated_configuration = await resolve_block_document_references(
             template=populated_configuration, client=client
@@ -863,19 +869,15 @@ class BaseWorker(abc.ABC):
 
         if flow_run.deployment_id:
             deployment = await self._client.read_deployment(flow_run.deployment_id)
-        if deployment and deployment.concurrency_limit:
-            limit_name = f"deployment:{deployment.id}"
-            concurrency_limit = deployment.concurrency_limit
+        if deployment and deployment.global_concurrency_limit:
+            limit_name = deployment.global_concurrency_limit.name
             concurrency_ctx = concurrency
         else:
-            limit_name = None
-            concurrency_limit = None
+            limit_name = ""
             concurrency_ctx = asyncnullcontext
 
         try:
-            async with concurrency_ctx(
-                limit_name, occupy=concurrency_limit, max_retries=0
-            ):
+            async with concurrency_ctx(limit_name, max_retries=0, strict=True):
                 configuration = await self._get_configuration(flow_run, deployment)
                 submitted_event = self._emit_flow_run_submitted_event(configuration)
                 result = await self.run(

{prefect_client-3.0.1.dist-info → prefect_client-3.0.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.0.1
+Version: 3.0.3
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.