prefect-client 3.0.0rc20__py3-none-any.whl → 3.0.2__py3-none-any.whl

prefect/transactions.py

@@ -17,17 +17,19 @@ from typing import (
 from pydantic import Field, PrivateAttr
 from typing_extensions import Self
 
-from prefect.context import ContextModel, FlowRunContext, TaskRunContext
+from prefect.context import ContextModel
 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,
-    ResultFactory,
-    get_default_result_storage,
+    ResultRecord,
+    ResultStore,
+    get_result_store,
+    should_persist_result,
 )
 from prefect.utilities.annotations import NotSet
-from prefect.utilities.asyncutils import run_coro_as_sync
 from prefect.utilities.collections import AutoEnum
 from prefect.utilities.engine import _get_hook_name
 
@@ -56,7 +58,7 @@ class Transaction(ContextModel):
     A base model for transaction state.
     """
 
-    store: Optional[RecordStore] = None
+    store: Union[RecordStore, ResultStore, None] = None
     key: Optional[str] = None
     children: List["Transaction"] = Field(default_factory=list)
     commit_mode: Optional[CommitMode] = None
@@ -70,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
@@ -105,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
@@ -123,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
@@ -177,10 +249,14 @@ class Transaction(ContextModel):
         ):
             self.state = TransactionState.COMMITTED
 
-    def read(self) -> Optional[BaseResult]:
+    def read(self) -> Union["BaseResult", ResultRecord, None]:
         if self.store and self.key:
             record = self.store.read(key=self.key)
-            if record is not None:
+            if isinstance(record, ResultRecord):
+                return record
+            # for backwards compatibility, if we encounter a transaction record, return the result
+            # This happens when the transaction is using a `ResultStore`
+            if isinstance(record, TransactionRecord):
                 return record.result
         return None
 
@@ -229,8 +305,21 @@ class Transaction(ContextModel):
             for hook in self.on_commit_hooks:
                 self.run_hook(hook, "commit")
 
-            if self.store and self.key:
-                self.store.write(key=self.key, result=self._staged_value)
+            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(
+                            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(key=self.key, obj=self._staged_value)
+                else:
+                    self.store.write(key=self.key, result=self._staged_value)
+
             self.state = TransactionState.COMMITTED
             if (
                 self.store
@@ -281,7 +370,7 @@ class Transaction(ContextModel):
 
     def stage(
         self,
-        value: BaseResult,
+        value: Any,
         on_rollback_hooks: Optional[List] = None,
         on_commit_hooks: Optional[List] = None,
     ) -> None:
@@ -339,10 +428,11 @@ def get_transaction() -> Optional[Transaction]:
 @contextmanager
 def transaction(
     key: Optional[str] = None,
-    store: Optional[RecordStore] = None,
+    store: Union[RecordStore, ResultStore, None] = None,
     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]:
     """
@@ -355,47 +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:
-        flow_run_context = FlowRunContext.get()
-        task_run_context = TaskRunContext.get()
-        existing_factory = getattr(task_run_context, "result_factory", None) or getattr(
-            flow_run_context, "result_factory", None
-        )
-
-        new_factory: ResultFactory
-        if existing_factory and existing_factory.storage_block_id:
-            new_factory = existing_factory.model_copy(
-                update={
-                    "persist_result": True,
-                }
-            )
-        else:
-            default_storage = get_default_result_storage(_sync=True)
-            if existing_factory:
-                new_factory = existing_factory.model_copy(
-                    update={
-                        "persist_result": True,
-                        "storage_block": default_storage,
-                        "storage_block_id": default_storage._block_document_id,
-                    }
-                )
-            else:
-                new_factory = run_coro_as_sync(
-                    ResultFactory.default_factory(
-                        persist_result=True,
-                        result_storage=default_storage,
-                    )
-                )
-        from prefect.records.result_store import ResultFactoryStore
-
-        store = ResultFactoryStore(
-            result_factory=new_factory,
-        )
+        store = get_result_store()
 
     try:
         logger = logger or get_run_logger()
@@ -408,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/callables.py

@@ -263,9 +263,7 @@ def parameter_docstrings(docstring: Optional[str]) -> Dict[str, str]:
     if not docstring:
         return param_docstrings
 
-    with disable_logger("griffe.docstrings.google"), disable_logger(
-        "griffe.agents.nodes"
-    ):
+    with disable_logger("griffe"):
         parsed = parse(Docstring(docstring), Parser.google)
         for section in parsed:
             if section.kind != DocstringSectionKind.parameters:

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 (

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/variables.py

@@ -1,19 +1,18 @@
-from typing import List, Optional, Union
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
 
 from prefect._internal.compatibility.migration import getattr_migration
-from prefect.client.schemas.actions import VariableCreate as VariableRequest
-from prefect.client.schemas.actions import VariableUpdate as VariableUpdateRequest
-from prefect.client.schemas.objects import Variable as VariableResponse
+from prefect.client.schemas.actions import VariableCreate, VariableUpdate
 from prefect.client.utilities import get_or_create_client
 from prefect.exceptions import ObjectNotFound
-from prefect.types import StrictVariableValue
+from prefect.types import MAX_VARIABLE_NAME_LENGTH, StrictVariableValue
 from prefect.utilities.asyncutils import sync_compatible
 
 
-class Variable(VariableRequest):
+class Variable(BaseModel):
     """
-    Variables are named, mutable string values, much like environment variables. Variables are scoped to a Prefect server instance or a single workspace in Prefect Cloud.
-    https://docs.prefect.io/latest/concepts/variables/
+    Variables are named, mutable JSON values that can be shared across tasks and flows.
 
     Arguments:
         name: A string identifying the variable.
@@ -21,6 +20,19 @@ class Variable(VariableRequest):
         tags: An optional list of strings to associate with the variable.
     """
 
+    name: str = Field(
+        default=...,
+        description="The name of the variable",
+        examples=["my_variable"],
+        max_length=MAX_VARIABLE_NAME_LENGTH,
+    )
+    value: StrictVariableValue = Field(
+        default=...,
+        description="The value of the variable",
+        examples=["my-value"],
+    )
+    tags: Optional[List[str]] = Field(default=None)
+
     @classmethod
     @sync_compatible
     async def set(
@@ -29,22 +41,21 @@ class Variable(VariableRequest):
         value: StrictVariableValue,
         tags: Optional[List[str]] = None,
         overwrite: bool = False,
-        as_object: bool = False,
-    ):
+    ) -> "Variable":
         """
         Sets a new variable. If one exists with the same name, must pass `overwrite=True`
 
-        Returns the newly set value. If `as_object=True`, return the full Variable object
+        Returns the newly set variable object.
 
         Args:
             - name: The name of the variable to set.
             - value: The value of the variable to set.
             - tags: An optional list of strings to associate with the variable.
             - overwrite: Whether to overwrite the variable if it already exists.
-            - as_object: Whether to return the full Variable object.
 
         Example:
             Set a new variable and overwrite it if it already exists.
+
             ```
             from prefect.variables import Variable
 
@@ -62,14 +73,17 @@ class Variable(VariableRequest):
                 raise ValueError(
                     f"Variable {name!r} already exists. Use `overwrite=True` to update it."
                 )
-            await client.update_variable(variable=VariableUpdateRequest(**var_dict))
+            await client.update_variable(variable=VariableUpdate(**var_dict))
             variable = await client.read_variable_by_name(name)
+            var_dict = {
+                "name": variable.name,
+                "value": variable.value,
+                "tags": variable.tags or [],
+            }
         else:
-            variable = await client.create_variable(
-                variable=VariableRequest(**var_dict)
-            )
+            await client.create_variable(variable=VariableCreate(**var_dict))
 
-        return variable if as_object else variable.value
+        return cls(**var_dict)
 
     @classmethod
     @sync_compatible
@@ -77,19 +91,15 @@ class Variable(VariableRequest):
         cls,
         name: str,
         default: StrictVariableValue = None,
-        as_object: bool = False,
-    ) -> Union[StrictVariableValue, VariableResponse]:
+    ) -> StrictVariableValue:
         """
         Get a variable's value by name.
 
         If the variable does not exist, return the default value.
 
-        If `as_object=True`, return the full variable object. `default` is ignored in this case.
-
         Args:
-            - name: The name of the variable to get.
+            - name: The name of the variable value to get.
             - default: The default value to return if the variable does not exist.
-            - as_object: Whether to return the full variable object.
 
         Example:
             Get a variable's value by name.
@@ -105,7 +115,7 @@ class Variable(VariableRequest):
         client, _ = get_or_create_client()
         variable = await client.read_variable_by_name(name)
 
-        return variable if as_object else (variable.value if variable else default)
+        return variable.value if variable else default
 
     @classmethod
     @sync_compatible

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
@@ -865,17 +871,19 @@ class BaseWorker(abc.ABC):
             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
             concurrency_ctx = concurrency
+
+            # ensure that the global concurrency limit is available
+            # and up-to-date before attempting to acquire a slot
+            await self._client.upsert_global_concurrency_limit_by_name(
+                limit_name, deployment.concurrency_limit
+            )
         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/workers/process.py

@@ -144,9 +144,7 @@ class ProcessWorker(BaseWorker):
         " when first getting started."
     )
     _display_name = "Process"
-    _documentation_url = (
-        "https://docs.prefect.io/latest/api-ref/prefect/workers/process/"
-    )
+    _documentation_url = "https://docs.prefect.io/latest/get-started/quickstart"
     _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/356e6766a91baf20e1d08bbe16e8b5aaef4d8643-48x48.png"
 
     async def start(

{prefect_client-3.0.0rc20.dist-info → prefect_client-3.0.2.dist-info}/METADATA

@@ -1,12 +1,12 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.0.0rc20
+Version: 3.0.2
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.
 Author-email: help@prefect.io
 License: UNKNOWN
-Project-URL: Changelog, https://github.com/PrefectHQ/prefect/blob/main/RELEASE-NOTES.md
+Project-URL: Changelog, https://github.com/PrefectHQ/prefect/releases
 Project-URL: Documentation, https://docs.prefect.io
 Project-URL: Source, https://github.com/PrefectHQ/prefect
 Project-URL: Tracker, https://github.com/PrefectHQ/prefect/issues