prefect-client 2.19.4__py3-none-any.whl → 3.0.0rc2__py3-none-any.whl

prefect/transactions.py

@@ -0,0 +1,224 @@
+from contextlib import contextmanager
+from contextvars import ContextVar, Token
+from typing import (
+    Any,
+    Callable,
+    Generator,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+)
+
+from pydantic import Field
+
+from prefect.context import ContextModel
+from prefect.records import RecordStore
+from prefect.utilities.collections import AutoEnum
+
+T = TypeVar("T")
+
+
+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
+    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
+    _staged_value: Any = None
+    __var__ = ContextVar("transaction")
+
+    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."
+            )
+        # set default commit behavior
+        if self.commit_mode is None:
+            parent = get_transaction()
+
+            # either inherit from parent or set a default of eager
+            if parent:
+                self.commit_mode = parent.commit_mode
+            else:
+                self.commit_mode = CommitMode.EAGER
+
+        # 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_type, exc_val, exc_tb):
+        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):
+        # currently we only support READ_COMMITTED isolation
+        # i.e., no locking behavior
+        if not self.overwrite and self.store and self.store.exists(key=self.key):
+            self.state = TransactionState.COMMITTED
+
+    def read(self) -> dict:
+        return self.store.read(key=self.key)
+
+    def reset(self) -> None:
+        parent = self.get_parent()
+
+        if parent:
+            # parent takes responsibility
+            parent.add_child(self)
+
+        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]:
+            return False
+
+        try:
+            for child in self.children:
+                child.commit()
+
+            for hook in self.on_commit_hooks:
+                hook(self)
+
+            if self.store:
+                self.store.write(key=self.key, value=self._staged_value)
+            self.state = TransactionState.COMMITTED
+            return True
+        except Exception:
+            self.rollback()
+            return False
+
+    def stage(
+        self, value: dict, on_rollback_hooks: list, on_commit_hooks: list
+    ) -> None:
+        """
+        Stage a value to be committed later.
+        """
+        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):
+                hook(self)
+
+            self.state = TransactionState.ROLLED_BACK
+
+            for child in reversed(self.children):
+                child.rollback()
+
+            return True
+        except Exception:
+            return False
+
+    @classmethod
+    def get_active(cls: Type[T]) -> Optional[T]:
+        return cls.__var__.get(None)
+
+
+def get_transaction() -> Transaction:
+    return Transaction.get_active()
+
+
+@contextmanager
+def transaction(
+    key: Optional[str] = None,
+    store: Optional[RecordStore] = None,
+    commit_mode: CommitMode = CommitMode.LAZY,
+    overwrite: bool = False,
+) -> Generator[Transaction, None, None]:
+    with Transaction(
+        key=key, store=store, commit_mode=commit_mode, overwrite=overwrite
+    ) as txn:
+        yield txn

prefect/types/__init__.py

@@ -1,112 +1,91 @@
-from typing import Any, Callable, ClassVar, Generator
+from typing import Annotated, Any, Dict, List, Union
+import orjson
+import pydantic
 
-from pydantic_core import core_schema, CoreSchema, SchemaValidator
-from typing_extensions import Self
-from datetime import timedelta
+from pydantic import (
+    BeforeValidator,
+    Field,
+    StrictBool,
+    StrictFloat,
+    StrictInt,
+    StrictStr,
+)
+from zoneinfo import available_timezones
 
+MAX_VARIABLE_NAME_LENGTH = 255
+MAX_VARIABLE_VALUE_LENGTH = 5000
 
-class NonNegativeInteger(int):
-    """An integer that must be greater than or equal to 0."""
+timezone_set = available_timezones()
 
-    schema: ClassVar[CoreSchema] = core_schema.int_schema(ge=0)
+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(timezone_set))]
 
-    @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
+BANNED_CHARACTERS = ["/", "%", "&", ">", "<"]
 
-    @classmethod
-    def validate(cls, v: Any) -> Self:
-        return SchemaValidator(schema=cls.schema).validate_python(v)
+WITHOUT_BANNED_CHARACTERS = r"^[^" + "".join(BANNED_CHARACTERS) + "]+$"
+Name = Annotated[str, Field(pattern=WITHOUT_BANNED_CHARACTERS)]
 
+WITHOUT_BANNED_CHARACTERS_EMPTY_OK = r"^[^" + "".join(BANNED_CHARACTERS) + "]*$"
+NameOrEmpty = Annotated[str, Field(pattern=WITHOUT_BANNED_CHARACTERS_EMPTY_OK)]
 
-class PositiveInteger(int):
-    """An integer that must be greater than 0."""
 
-    schema: ClassVar[CoreSchema] = core_schema.int_schema(gt=0)
+def non_emptyish(value: str) -> str:
+    if not value.strip("' \""):
+        raise ValueError("name cannot be an empty string")
 
-    @classmethod
-    def __get_validators__(cls) -> Generator[Callable[..., Any], None, None]:
-        yield cls.validate
+    return value
 
-    @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 NonNegativeFloat(float):
-    schema: ClassVar[CoreSchema] = core_schema.float_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 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",
 ]