prefect-client 2.20.4__py3-none-any.whl → 3.0.0__py3-none-any.whl

prefect/transactions.py

@@ -0,0 +1,410 @@
+import copy
+import logging
+from contextlib import contextmanager
+from contextvars import ContextVar, Token
+from functools import partial
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    Type,
+    Union,
+)
+
+from pydantic import Field, PrivateAttr
+from typing_extensions import Self
+
+from prefect.context import ContextModel, FlowRunContext, TaskRunContext
+from prefect.exceptions import MissingContextError, SerializationError
+from prefect.logging.loggers import get_logger, get_run_logger
+from prefect.records import RecordStore
+from prefect.results import (
+    BaseResult,
+    ResultStore,
+    get_default_result_storage,
+)
+from prefect.utilities.annotations import NotSet
+from prefect.utilities.collections import AutoEnum
+from prefect.utilities.engine import _get_hook_name
+
+
+class IsolationLevel(AutoEnum):
+    READ_COMMITTED = AutoEnum.auto()
+    SERIALIZABLE = AutoEnum.auto()
+
+
+class CommitMode(AutoEnum):
+    EAGER = AutoEnum.auto()
+    LAZY = AutoEnum.auto()
+    OFF = AutoEnum.auto()
+
+
+class TransactionState(AutoEnum):
+    PENDING = AutoEnum.auto()
+    ACTIVE = AutoEnum.auto()
+    STAGED = AutoEnum.auto()
+    COMMITTED = AutoEnum.auto()
+    ROLLED_BACK = AutoEnum.auto()
+
+
+class Transaction(ContextModel):
+    """
+    A base model for transaction state.
+    """
+
+    store: Optional[RecordStore] = None
+    key: Optional[str] = None
+    children: List["Transaction"] = Field(default_factory=list)
+    commit_mode: Optional[CommitMode] = None
+    isolation_level: Optional[IsolationLevel] = IsolationLevel.READ_COMMITTED
+    state: TransactionState = TransactionState.PENDING
+    on_commit_hooks: List[Callable[["Transaction"], None]] = Field(default_factory=list)
+    on_rollback_hooks: List[Callable[["Transaction"], None]] = Field(
+        default_factory=list
+    )
+    overwrite: bool = False
+    logger: Union[logging.Logger, logging.LoggerAdapter] = Field(
+        default_factory=partial(get_logger, "transactions")
+    )
+    _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:
+        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)
+
+    def is_committed(self) -> bool:
+        return self.state == TransactionState.COMMITTED
+
+    def is_rolled_back(self) -> bool:
+        return self.state == TransactionState.ROLLED_BACK
+
+    def is_staged(self) -> bool:
+        return self.state == TransactionState.STAGED
+
+    def is_pending(self) -> bool:
+        return self.state == TransactionState.PENDING
+
+    def is_active(self) -> bool:
+        return self.state == TransactionState.ACTIVE
+
+    def __enter__(self):
+        if self._token is not None:
+            raise RuntimeError(
+                "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
+        # set default isolation level; either inherit from parent or set a default of read committed
+        if self.isolation_level is None:
+            self.isolation_level = (
+                parent.isolation_level if parent else IsolationLevel.READ_COMMITTED
+            )
+
+        assert self.isolation_level is not None, "Isolation level was not set correctly"
+        if (
+            self.store
+            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__}"
+            )
+
+        # this needs to go before begin, which could set the state to committed
+        self.state = TransactionState.ACTIVE
+        self.begin()
+        self._token = self.__var__.set(self)
+        return self
+
+    def __exit__(self, *exc_info):
+        exc_type, exc_val, _ = exc_info
+        if not self._token:
+            raise RuntimeError(
+                "Asymmetric use of context. Context exit called without an enter."
+            )
+        if exc_type:
+            self.rollback()
+            self.reset()
+            raise exc_val
+
+        if self.commit_mode == CommitMode.EAGER:
+            self.commit()
+
+        # if parent, let them take responsibility
+        if self.get_parent():
+            self.reset()
+            return
+
+        if self.commit_mode == CommitMode.OFF:
+            # if no one took responsibility to commit, rolling back
+            # note that rollback returns if already committed
+            self.rollback()
+        elif self.commit_mode == CommitMode.LAZY:
+            # no one left to take responsibility for committing
+            self.commit()
+
+        self.reset()
+
+    def begin(self):
+        if (
+            self.store
+            and self.key
+            and self.isolation_level == IsolationLevel.SERIALIZABLE
+        ):
+            self.logger.debug(f"Acquiring lock for transaction {self.key!r}")
+            self.store.acquire_lock(self.key)
+        if (
+            not self.overwrite
+            and self.store
+            and self.key
+            and self.store.exists(key=self.key)
+        ):
+            self.state = TransactionState.COMMITTED
+
+    def read(self) -> Optional[BaseResult]:
+        if self.store and self.key:
+            record = self.store.read(key=self.key)
+            if record is not None:
+                return record.result
+        return None
+
+    def reset(self) -> None:
+        parent = self.get_parent()
+
+        if parent:
+            # parent takes responsibility
+            parent.add_child(self)
+
+        if self._token:
+            self.__var__.reset(self._token)
+            self._token = None
+
+        # do this below reset so that get_transaction() returns the relevant txn
+        if parent and self.state == TransactionState.ROLLED_BACK:
+            parent.rollback()
+
+    def add_child(self, transaction: "Transaction") -> None:
+        self.children.append(transaction)
+
+    def get_parent(self) -> Optional["Transaction"]:
+        prev_var = getattr(self._token, "old_value")
+        if prev_var != Token.MISSING:
+            parent = prev_var
+        else:
+            parent = None
+        return parent
+
+    def commit(self) -> bool:
+        if self.state in [TransactionState.ROLLED_BACK, TransactionState.COMMITTED]:
+            if (
+                self.store
+                and self.key
+                and self.isolation_level == IsolationLevel.SERIALIZABLE
+            ):
+                self.logger.debug(f"Releasing lock for transaction {self.key!r}")
+                self.store.release_lock(self.key)
+
+            return False
+
+        try:
+            for child in self.children:
+                child.commit()
+
+            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)
+            self.state = TransactionState.COMMITTED
+            if (
+                self.store
+                and self.key
+                and self.isolation_level == IsolationLevel.SERIALIZABLE
+            ):
+                self.logger.debug(f"Releasing lock for transaction {self.key!r}")
+                self.store.release_lock(self.key)
+            return True
+        except SerializationError as exc:
+            if self.logger:
+                self.logger.warning(
+                    f"Encountered an error while serializing result for transaction {self.key!r}: {exc}"
+                    " Code execution will continue, but the transaction will not be committed.",
+                )
+            self.rollback()
+            return False
+        except Exception:
+            if self.logger:
+                self.logger.exception(
+                    f"An error was encountered while committing transaction {self.key!r}",
+                    exc_info=True,
+                )
+            self.rollback()
+            return False
+
+    def run_hook(self, hook, hook_type: str) -> None:
+        hook_name = _get_hook_name(hook)
+        # Undocumented way to disable logging for a hook. Subject to change.
+        should_log = getattr(hook, "log_on_run", True)
+
+        if should_log:
+            self.logger.info(f"Running {hook_type} hook {hook_name!r}")
+
+        try:
+            hook(self)
+        except Exception as exc:
+            if should_log:
+                self.logger.error(
+                    f"An error was encountered while running {hook_type} hook {hook_name!r}",
+                )
+            raise exc
+        else:
+            if should_log:
+                self.logger.info(
+                    f"{hook_type.capitalize()} hook {hook_name!r} finished running successfully"
+                )
+
+    def stage(
+        self,
+        value: BaseResult,
+        on_rollback_hooks: Optional[List] = None,
+        on_commit_hooks: Optional[List] = None,
+    ) -> None:
+        """
+        Stage a value to be committed later.
+        """
+        on_commit_hooks = on_commit_hooks or []
+        on_rollback_hooks = on_rollback_hooks or []
+
+        if self.state != TransactionState.COMMITTED:
+            self._staged_value = value
+            self.on_rollback_hooks += on_rollback_hooks
+            self.on_commit_hooks += on_commit_hooks
+            self.state = TransactionState.STAGED
+
+    def rollback(self) -> bool:
+        if self.state in [TransactionState.ROLLED_BACK, TransactionState.COMMITTED]:
+            return False
+
+        try:
+            for hook in reversed(self.on_rollback_hooks):
+                self.run_hook(hook, "rollback")
+
+            self.state = TransactionState.ROLLED_BACK
+
+            for child in reversed(self.children):
+                child.rollback()
+
+            return True
+        except Exception:
+            if self.logger:
+                self.logger.exception(
+                    f"An error was encountered while rolling back transaction {self.key!r}",
+                    exc_info=True,
+                )
+            return False
+        finally:
+            if (
+                self.store
+                and self.key
+                and self.isolation_level == IsolationLevel.SERIALIZABLE
+            ):
+                self.logger.debug(f"Releasing lock for transaction {self.key!r}")
+                self.store.release_lock(self.key)
+
+    @classmethod
+    def get_active(cls: Type[Self]) -> Optional[Self]:
+        return cls.__var__.get(None)
+
+
+def get_transaction() -> Optional[Transaction]:
+    return Transaction.get_active()
+
+
+@contextmanager
+def transaction(
+    key: Optional[str] = None,
+    store: Optional[RecordStore] = None,
+    commit_mode: Optional[CommitMode] = None,
+    isolation_level: Optional[IsolationLevel] = None,
+    overwrite: bool = False,
+    logger: Union[logging.Logger, logging.LoggerAdapter, None] = None,
+) -> Generator[Transaction, None, None]:
+    """
+    A context manager for opening and managing a transaction.
+
+    Args:
+        - key: An identifier to use for the transaction
+        - store: The store to use for persisting the transaction result. If not provided,
+            a default store will be used based on the current run context.
+        - 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
+
+    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_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,
+        )
+
+    try:
+        logger = logger or get_run_logger()
+    except MissingContextError:
+        logger = get_logger("transactions")
+
+    with Transaction(
+        key=key,
+        store=store,
+        commit_mode=commit_mode,
+        isolation_level=isolation_level,
+        overwrite=overwrite,
+        logger=logger,
+    ) as txn:
+        yield txn

prefect/types/__init__.py

@@ -1,112 +1,98 @@
-from typing import Any, Callable, ClassVar, Generator
-
-from pydantic_core import core_schema, CoreSchema, SchemaValidator
-from typing_extensions import Self
-from datetime import timedelta
-
-
-class NonNegativeInteger(int):
-    """An integer that must be greater than or equal to 0."""
-
-    schema: ClassVar[CoreSchema] = core_schema.int_schema(ge=0)
-
-    @classmethod
-    def __get_validators__(cls) -> Generator[Callable[..., Any], None, None]:
-        yield cls.validate
-
-    @classmethod
-    def __get_pydantic_core_schema__(
-        cls, source_type: Any, handler: Callable[..., Any]
-    ) -> CoreSchema:
-        return cls.schema
-
-    @classmethod
-    def validate(cls, v: Any) -> Self:
-        return SchemaValidator(schema=cls.schema).validate_python(v)
-
-
-class PositiveInteger(int):
-    """An integer that must be greater than 0."""
-
-    schema: ClassVar[CoreSchema] = core_schema.int_schema(gt=0)
-
-    @classmethod
-    def __get_validators__(cls) -> Generator[Callable[..., Any], None, None]:
-        yield cls.validate
+from typing import Annotated, Any, Dict, List, Union
+import orjson
+import pydantic
+
+from pydantic import (
+    BeforeValidator,
+    Field,
+    StrictBool,
+    StrictFloat,
+    StrictInt,
+    StrictStr,
+)
+from zoneinfo import available_timezones
+
+MAX_VARIABLE_NAME_LENGTH = 255
+MAX_VARIABLE_VALUE_LENGTH = 5000
+
+NonNegativeInteger = Annotated[int, Field(ge=0)]
+PositiveInteger = Annotated[int, Field(gt=0)]
+NonNegativeFloat = Annotated[float, Field(ge=0.0)]
+
+TimeZone = Annotated[
+    str,
+    Field(
+        default="UTC",
+        pattern="|".join(
+            [z for z in sorted(available_timezones()) if "localtime" not in z]
+        ),
+    ),
+]
 
-    @classmethod
-    def __get_pydantic_core_schema__(
-        cls, source_type: Any, handler: Callable[..., Any]
-    ) -> CoreSchema:
-        return cls.schema
 
-    @classmethod
-    def validate(cls, v: Any) -> Self:
-        return SchemaValidator(schema=cls.schema).validate_python(v)
+BANNED_CHARACTERS = ["/", "%", "&", ">", "<"]
 
+WITHOUT_BANNED_CHARACTERS = r"^[^" + "".join(BANNED_CHARACTERS) + "]+$"
+Name = Annotated[str, Field(pattern=WITHOUT_BANNED_CHARACTERS)]
 
-class NonNegativeFloat(float):
-    schema: ClassVar[CoreSchema] = core_schema.float_schema(ge=0)
+WITHOUT_BANNED_CHARACTERS_EMPTY_OK = r"^[^" + "".join(BANNED_CHARACTERS) + "]*$"
+NameOrEmpty = Annotated[str, Field(pattern=WITHOUT_BANNED_CHARACTERS_EMPTY_OK)]
 
-    @classmethod
-    def __get_validators__(cls) -> Generator[Callable[..., Any], None, None]:
-        yield cls.validate
 
-    @classmethod
-    def __get_pydantic_core_schema__(
-        cls, source_type: Any, handler: Callable[..., Any]
-    ) -> CoreSchema:
-        return cls.schema
+def non_emptyish(value: str) -> str:
+    if not value.strip("' \""):
+        raise ValueError("name cannot be an empty string")
 
-    @classmethod
-    def validate(cls, v: Any) -> Self:
-        return SchemaValidator(schema=cls.schema).validate_python(v)
+    return value
 
 
-class NonNegativeDuration(timedelta):
-    """A timedelta that must be greater than or equal to 0."""
+NonEmptyishName = Annotated[
+    str,
+    Field(pattern=WITHOUT_BANNED_CHARACTERS),
+    BeforeValidator(non_emptyish),
+]
 
-    schema: ClassVar = core_schema.timedelta_schema(ge=timedelta(seconds=0))
 
-    @classmethod
-    def __get_validators__(cls) -> Generator[Callable[..., Any], None, None]:
-        yield cls.validate
+VariableValue = Union[
+    StrictStr,
+    StrictInt,
+    StrictBool,
+    StrictFloat,
+    None,
+    Dict[str, Any],
+    List[Any],
+]
 
-    @classmethod
-    def __get_pydantic_core_schema__(
-        cls, source_type: Any, handler: Callable[..., Any]
-    ) -> CoreSchema:
-        return cls.schema
 
-    @classmethod
-    def validate(cls, v: Any) -> Self:
-        return SchemaValidator(schema=cls.schema).validate_python(v)
+def check_variable_value(value: object) -> object:
+    try:
+        json_string = orjson.dumps(value)
+    except orjson.JSONEncodeError:
+        raise ValueError("Variable value must be serializable to JSON")
 
+    if value is not None and len(json_string) > MAX_VARIABLE_VALUE_LENGTH:
+        raise ValueError(
+            f"Variable value must be less than {MAX_VARIABLE_VALUE_LENGTH} characters"
+        )
+    return value
 
-class PositiveDuration(timedelta):
-    """A timedelta that must be greater than 0."""
 
-    schema: ClassVar = core_schema.timedelta_schema(gt=timedelta(seconds=0))
+StrictVariableValue = Annotated[VariableValue, BeforeValidator(check_variable_value)]
 
-    @classmethod
-    def __get_validators__(cls) -> Generator[Callable[..., Any], None, None]:
-        yield cls.validate
+LaxUrl = Annotated[str, BeforeValidator(lambda x: str(x).strip())]
 
-    @classmethod
-    def __get_pydantic_core_schema__(
-        cls, source_type: Any, handler: Callable[..., Any]
-    ) -> CoreSchema:
-        return cls.schema
 
-    @classmethod
-    def validate(cls, v: Any) -> Self:
-        return SchemaValidator(schema=cls.schema).validate_python(v)
+class SecretDict(pydantic.Secret[Dict[str, Any]]):
+    pass
 
 
 __all__ = [
     "NonNegativeInteger",
     "PositiveInteger",
     "NonNegativeFloat",
-    "NonNegativeDuration",
-    "PositiveDuration",
+    "Name",
+    "NameOrEmpty",
+    "NonEmptyishName",
+    "SecretDict",
+    "StrictVariableValue",
 ]

prefect/types/entrypoint.py

@@ -0,0 +1,13 @@
+from enum import Enum
+
+
+class EntrypointType(Enum):
+    """
+    Enum representing a entrypoint type.
+
+    File path entrypoints are in the format: `path/to/file.py:function_name`.
+    Module path entrypoints are in the format: `path.to.module.function_name`.
+    """
+
+    FILE_PATH = "file_path"
+    MODULE_PATH = "module_path"

prefect/utilities/annotations.py

@@ -21,8 +21,8 @@ class BaseAnnotation(
     def rewrap(self, value: T) -> "BaseAnnotation[T]":
         return type(self)(value)
 
-    def __eq__(self, other: object) -> bool:
-        if not type(self) == type(other):
+    def __eq__(self, other: "BaseAnnotation[T]") -> bool:
+        if type(self) is not type(other):
             return False
         return self.unwrap() == other.unwrap()
 
@@ -90,10 +90,11 @@ class quote(BaseAnnotation[T]):
 class Quote(quote):
     def __init__(self, expr):
         warnings.warn(
-            DeprecationWarning,
             "Use of `Quote` is deprecated. Use `quote` instead.",
+            DeprecationWarning,
             stacklevel=2,
         )
+        super().__init__(expr)
 
 
 class NotSet: