prefect-client 3.0.0rc17__py3-none-any.whl → 3.0.0rc19__py3-none-any.whl

prefect/records/base.py

@@ -0,0 +1,223 @@
+import abc
+import os
+import socket
+import threading
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from prefect.results import BaseResult
+    from prefect.transactions import IsolationLevel
+
+
+@dataclass
+class TransactionRecord:
+    """
+    A dataclass representation of a transaction record.
+    """
+
+    key: str
+    result: "BaseResult"
+
+
+class RecordStore(abc.ABC):
+    @abc.abstractmethod
+    def read(
+        self, key: str, holder: Optional[str] = None
+    ) -> Optional[TransactionRecord]:
+        """
+        Read the transaction record with the given key.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock. If a lock exists on
+                the record being written, the read will be blocked until the lock is
+                released if the provided holder does not match the holder of the lock.
+                If not provided, a default holder based on the current host, process,
+                and thread will be used.
+
+        Returns:
+            TransactionRecord: The transaction record with the given key.
+        """
+        ...
+
+    @abc.abstractmethod
+    def write(self, key: str, result: "BaseResult", holder: Optional[str] = None):
+        """
+        Write the transaction record with the given key.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            record: The transaction record to write.
+            holder: Unique identifier for the holder of the lock. If a lock exists on
+                the record being written, the write will be rejected if the provided
+                holder does not match the holder of the lock. If not provided,
+                a default holder based on the current host, process, and thread will
+                be used.
+        """
+        ...
+
+    @abc.abstractmethod
+    def exists(self, key: str) -> bool:
+        """
+        Check if the transaction record with the given key exists.
+
+        Args:
+            key: Unique identifier for the transaction record.
+
+        Returns:
+            bool: True if the record exists; False otherwise.
+        """
+        ...
+
+    @abc.abstractmethod
+    def supports_isolation_level(self, isolation_level: "IsolationLevel") -> bool:
+        """
+        Check if the record store supports the given isolation level.
+
+        Args:
+            isolation_level: The isolation level to check.
+
+        Returns:
+            bool: True if the record store supports the isolation level; False otherwise.
+        """
+        ...
+
+    def acquire_lock(
+        self,
+        key: str,
+        holder: Optional[str] = None,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        """
+        Acquire a lock for a transaction record with the given key. Will block other
+        actors from updating this transaction record until the lock is
+        released.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock. If not provided,
+                a default holder based on the current host, process, and thread will
+                be used.
+            acquire_timeout: Max number of seconds to wait for the record to become
+                available if it is locked while attempting to acquire a lock. Pass 0
+                to attempt to acquire a lock without waiting. Blocks indefinitely by
+                default.
+            hold_timeout: Max number of seconds to hold the lock for. Holds the lock
+                indefinitely by default.
+
+        Returns:
+            bool: True if the lock was successfully acquired; False otherwise.
+        """
+        raise NotImplementedError
+
+    def release_lock(self, key: str, holder: Optional[str] = None):
+        """
+        Releases the lock on the corresponding transaction record.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock. Must match the
+                holder provided when acquiring the lock.
+        """
+        raise NotImplementedError
+
+    def is_locked(self, key: str) -> bool:
+        """
+        Simple check to see if the corresponding record is currently locked.
+
+        Args:
+            key: Unique identifier for the transaction record.
+
+        Returns:
+            True is the record is locked; False otherwise.
+        """
+        raise NotImplementedError
+
+    def is_lock_holder(self, key: str, holder: Optional[str] = None) -> bool:
+        """
+        Check if the current holder is the lock holder for the transaction record.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock. If not provided,
+                a default holder based on the current host, process, and thread will
+                be used.
+
+        Returns:
+            bool: True if the current holder is the lock holder; False otherwise.
+        """
+        raise NotImplementedError
+
+    def wait_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        """
+        Wait for the corresponding transaction record to become free.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            timeout: Maximum time to wait. None means to wait indefinitely.
+
+        Returns:
+            bool: True if the lock becomes free within the timeout; False
+                otherwise.
+        """
+        ...
+
+    @staticmethod
+    def generate_default_holder() -> str:
+        """
+        Generate a default holder string using hostname, PID, and thread ID.
+
+        Returns:
+            str: A unique identifier string.
+        """
+        hostname = socket.gethostname()
+        pid = os.getpid()
+        thread_name = threading.current_thread().name
+        thread_id = threading.get_ident()
+        return f"{hostname}:{pid}:{thread_id}:{thread_name}"
+
+    @contextmanager
+    def lock(
+        self,
+        key: str,
+        holder: Optional[str] = None,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ):
+        """
+        Context manager to lock the transaction record during the execution
+        of the nested code block.
+
+        Args:
+            key: Unique identifier for the transaction record.
+            holder: Unique identifier for the holder of the lock. If not provided,
+                a default holder based on the current host, process, and thread will
+                be used.
+            acquire_timeout: Max number of seconds to wait for the record to become
+                available if it is locked while attempting to acquire a lock. Pass 0
+                to attempt to acquire a lock without waiting. Blocks indefinitely by
+                default.
+            hold_timeout: Max number of seconds to hold the lock for. Holds the lock
+                indefinitely by default.
+
+        Example:
+            Hold a lock while during an operation:
+                ```python
+                    with TransactionRecord(key="my-transaction-record-key").lock():
+                        do_stuff()
+                ```
+        """
+        self.acquire_lock(
+            key=key,
+            holder=holder,
+            acquire_timeout=acquire_timeout,
+            hold_timeout=hold_timeout,
+        )
+
+        try:
+            yield
+        finally:
+            self.release_lock(key=key, holder=holder)

prefect/records/filesystem.py

@@ -0,0 +1,207 @@
+import json
+import time
+from pathlib import Path
+from typing import Dict, Optional
+
+import pendulum
+from typing_extensions import TypedDict
+
+from prefect.logging.loggers import get_logger
+from prefect.records.base import RecordStore, TransactionRecord
+from prefect.results import BaseResult
+from prefect.transactions import IsolationLevel
+
+logger = get_logger(__name__)
+
+
+class _LockInfo(TypedDict):
+    """
+    A dictionary containing information about a lock.
+
+    Attributes:
+        holder: The holder of the lock.
+        expiration: Datetime when the lock expires.
+        path: Path to the lock file.
+    """
+
+    holder: str
+    expiration: Optional[pendulum.DateTime]
+    path: Path
+
+
+class FileSystemRecordStore(RecordStore):
+    """
+    A record store that stores data on the local filesystem.
+
+    Locking is implemented using a lock file with the same name as the record file,
+    but with a `.lock` extension.
+
+    Attributes:
+        records_directory: the directory where records are stored; defaults to
+            `{PREFECT_HOME}/records`
+    """
+
+    def __init__(self, records_directory: Path):
+        self.records_directory = records_directory
+        self._locks: Dict[str, _LockInfo] = {}
+
+    def _ensure_records_directory_exists(self):
+        self.records_directory.mkdir(parents=True, exist_ok=True)
+
+    def _lock_path_for_key(self, key: str) -> Path:
+        if (lock_info := self._locks.get(key)) is not None:
+            return lock_info["path"]
+        return self.records_directory.joinpath(key).with_suffix(".lock")
+
+    def _get_lock_info(self, key: str, use_cache=True) -> Optional[_LockInfo]:
+        if use_cache:
+            if (lock_info := self._locks.get(key)) is not None:
+                print("Got lock info from cache")
+                return lock_info
+
+        lock_path = self._lock_path_for_key(key)
+
+        try:
+            with open(lock_path, "r") as lock_file:
+                lock_info = json.load(lock_file)
+                lock_info["path"] = lock_path
+                expiration = lock_info.get("expiration")
+                lock_info["expiration"] = (
+                    pendulum.parse(expiration) if expiration is not None else None
+                )
+            self._locks[key] = lock_info
+            print("Got lock info from file")
+            return lock_info
+        except FileNotFoundError:
+            return None
+
+    def read(
+        self, key: str, holder: Optional[str] = None
+    ) -> Optional[TransactionRecord]:
+        if not self.exists(key):
+            return None
+
+        holder = holder or self.generate_default_holder()
+
+        if self.is_locked(key) and not self.is_lock_holder(key, holder):
+            self.wait_for_lock(key)
+        record_data = self.records_directory.joinpath(key).read_text()
+        return TransactionRecord(
+            key=key, result=BaseResult.model_validate_json(record_data)
+        )
+
+    def write(self, key: str, result: BaseResult, holder: Optional[str] = None) -> None:
+        self._ensure_records_directory_exists()
+
+        if self.is_locked(key) and not self.is_lock_holder(key, holder):
+            raise ValueError(
+                f"Cannot write to transaction with key {key} because it is locked by another holder."
+            )
+
+        record_path = self.records_directory.joinpath(key)
+        record_path.touch(exist_ok=True)
+        record_data = result.model_dump_json()
+        record_path.write_text(record_data)
+
+    def exists(self, key: str) -> bool:
+        return self.records_directory.joinpath(key).exists()
+
+    def supports_isolation_level(self, isolation_level: IsolationLevel) -> bool:
+        return isolation_level in {
+            IsolationLevel.READ_COMMITTED,
+            IsolationLevel.SERIALIZABLE,
+        }
+
+    def acquire_lock(
+        self,
+        key: str,
+        holder: Optional[str] = None,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        holder = holder or self.generate_default_holder()
+
+        self._ensure_records_directory_exists()
+        lock_path = self._lock_path_for_key(key)
+
+        if self.is_locked(key) and not self.is_lock_holder(key, holder):
+            lock_free = self.wait_for_lock(key, acquire_timeout)
+            if not lock_free:
+                return False
+
+        try:
+            Path(lock_path).touch(exist_ok=False)
+        except FileExistsError:
+            if not self.is_lock_holder(key, holder):
+                logger.debug(
+                    f"Another actor acquired the lock for record with key {key}. Trying again."
+                )
+                return self.acquire_lock(key, holder, acquire_timeout, hold_timeout)
+        expiration = (
+            pendulum.now("utc") + pendulum.duration(seconds=hold_timeout)
+            if hold_timeout is not None
+            else None
+        )
+
+        with open(Path(lock_path), "w") as lock_file:
+            json.dump(
+                {
+                    "holder": holder,
+                    "expiration": str(expiration) if expiration is not None else None,
+                },
+                lock_file,
+            )
+
+        self._locks[key] = {
+            "holder": holder,
+            "expiration": expiration,
+            "path": lock_path,
+        }
+
+        return True
+
+    def release_lock(self, key: str, holder: Optional[str] = None) -> None:
+        holder = holder or self.generate_default_holder()
+        lock_path = self._lock_path_for_key(key)
+        if not self.is_locked(key):
+            ValueError(f"No lock for transaction with key {key}")
+        if self.is_lock_holder(key, holder):
+            Path(lock_path).unlink(missing_ok=True)
+            self._locks.pop(key, None)
+        else:
+            raise ValueError(f"No lock held by {holder} for transaction with key {key}")
+
+    def is_locked(self, key: str, use_cache: bool = False) -> bool:
+        if (lock_info := self._get_lock_info(key, use_cache=use_cache)) is None:
+            return False
+
+        if (expiration := lock_info.get("expiration")) is None:
+            return True
+
+        expired = expiration < pendulum.now("utc")
+        if expired:
+            Path(lock_info["path"]).unlink()
+            self._locks.pop(key, None)
+            return False
+        else:
+            return True
+
+    def is_lock_holder(self, key: str, holder: Optional[str] = None) -> bool:
+        if not self.is_locked(key):
+            return False
+
+        holder = holder or self.generate_default_holder()
+        if not self.is_locked(key):
+            return False
+        if (lock_info := self._get_lock_info(key)) is None:
+            return False
+        return lock_info["holder"] == holder
+
+    def wait_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        seconds_waited = 0
+        while self.is_locked(key, use_cache=False):
+            if timeout and seconds_waited >= timeout:
+                return False
+            seconds_waited += 0.1
+            time.sleep(0.1)
+        return True

prefect/records/memory.py

@@ -0,0 +1,178 @@
+import threading
+from typing import Dict, Optional, TypedDict
+
+from prefect.results import BaseResult
+from prefect.transactions import IsolationLevel
+
+from .base import RecordStore, TransactionRecord
+
+
+class _LockInfo(TypedDict):
+    """
+    A dictionary containing information about a lock.
+
+    Attributes:
+        holder: The holder of the lock.
+        lock: The lock object.
+        expiration_timer: The timer for the lock expiration
+    """
+
+    holder: str
+    lock: threading.Lock
+    expiration_timer: Optional[threading.Timer]
+
+
+class MemoryRecordStore(RecordStore):
+    """
+    A record store that stores data in memory.
+    """
+
+    _instance = None
+
+    def __new__(cls, *args, **kwargs):
+        if cls._instance is None:
+            cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        self._locks_dict_lock = threading.Lock()
+        self._locks: Dict[str, _LockInfo] = {}
+        self._records: Dict[str, TransactionRecord] = {}
+
+    def read(
+        self, key: str, holder: Optional[str] = None
+    ) -> Optional[TransactionRecord]:
+        holder = holder or self.generate_default_holder()
+
+        if self.is_locked(key) and not self.is_lock_holder(key, holder):
+            self.wait_for_lock(key)
+        return self._records.get(key)
+
+    def write(self, key: str, result: BaseResult, holder: Optional[str] = None) -> None:
+        holder = holder or self.generate_default_holder()
+
+        with self._locks_dict_lock:
+            if self.is_locked(key) and not self.is_lock_holder(key, holder):
+                raise ValueError(
+                    f"Cannot write to transaction with key {key} because it is locked by another holder."
+                )
+        self._records[key] = TransactionRecord(key=key, result=result)
+
+    def exists(self, key: str) -> bool:
+        return key in self._records
+
+    def supports_isolation_level(self, isolation_level: IsolationLevel) -> bool:
+        return isolation_level in {
+            IsolationLevel.READ_COMMITTED,
+            IsolationLevel.SERIALIZABLE,
+        }
+
+    def _expire_lock(self, key: str):
+        """
+        Expire the lock for the given key.
+
+        Used as a callback for the expiration timer of a lock.
+
+        Args:
+            key: The key of the lock to expire.
+        """
+        with self._locks_dict_lock:
+            if key in self._locks:
+                lock_info = self._locks[key]
+                if lock_info["lock"].locked():
+                    lock_info["lock"].release()
+                if lock_info["expiration_timer"]:
+                    lock_info["expiration_timer"].cancel()
+                del self._locks[key]
+
+    def acquire_lock(
+        self,
+        key: str,
+        holder: Optional[str] = None,
+        acquire_timeout: Optional[float] = None,
+        hold_timeout: Optional[float] = None,
+    ) -> bool:
+        holder = holder or self.generate_default_holder()
+        with self._locks_dict_lock:
+            if key not in self._locks:
+                lock = threading.Lock()
+                lock.acquire()
+                expiration_timer = None
+                if hold_timeout is not None:
+                    expiration_timer = threading.Timer(
+                        hold_timeout, self._expire_lock, args=(key,)
+                    )
+                    expiration_timer.start()
+                self._locks[key] = _LockInfo(
+                    holder=holder, lock=lock, expiration_timer=expiration_timer
+                )
+                return True
+            elif self._locks[key]["holder"] == holder:
+                return True
+            else:
+                existing_lock_info = self._locks[key]
+
+        if acquire_timeout is not None:
+            existing_lock_acquired = existing_lock_info["lock"].acquire(
+                timeout=acquire_timeout
+            )
+        else:
+            existing_lock_acquired = existing_lock_info["lock"].acquire()
+
+        if existing_lock_acquired:
+            with self._locks_dict_lock:
+                if (
+                    expiration_timer := existing_lock_info["expiration_timer"]
+                ) is not None:
+                    expiration_timer.cancel()
+                expiration_timer = None
+                if hold_timeout is not None:
+                    expiration_timer = threading.Timer(
+                        hold_timeout, self._expire_lock, args=(key,)
+                    )
+                    expiration_timer.start()
+                self._locks[key] = _LockInfo(
+                    holder=holder,
+                    lock=existing_lock_info["lock"],
+                    expiration_timer=expiration_timer,
+                )
+            return True
+        return False
+
+    def release_lock(self, key: str, holder: Optional[str] = None) -> None:
+        holder = holder or self.generate_default_holder()
+        with self._locks_dict_lock:
+            if key in self._locks and self._locks[key]["holder"] == holder:
+                if (
+                    expiration_timer := self._locks[key]["expiration_timer"]
+                ) is not None:
+                    expiration_timer.cancel()
+                self._locks[key]["lock"].release()
+                del self._locks[key]
+            else:
+                raise ValueError(
+                    f"No lock held by {holder} for transaction with key {key}"
+                )
+
+    def is_locked(self, key: str) -> bool:
+        return key in self._locks and self._locks[key]["lock"].locked()
+
+    def is_lock_holder(self, key: str, holder: Optional[str] = None) -> bool:
+        holder = holder or self.generate_default_holder()
+        lock_info = self._locks.get(key)
+        return (
+            lock_info is not None
+            and lock_info["lock"].locked()
+            and lock_info["holder"] == holder
+        )
+
+    def wait_for_lock(self, key: str, timeout: Optional[float] = None) -> bool:
+        if lock := self._locks.get(key, {}).get("lock"):
+            if timeout is not None:
+                lock_acquired = lock.acquire(timeout=timeout)
+            else:
+                lock_acquired = lock.acquire()
+            if lock_acquired:
+                lock.release()
+            return lock_acquired
+        return True

prefect/records/result_store.py

@@ -1,22 +1,26 @@
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Optional
 
 import pendulum
 
 from prefect.results import BaseResult, PersistedResult, ResultFactory
+from prefect.transactions import IsolationLevel
 from prefect.utilities.asyncutils import run_coro_as_sync
 
-from .store import RecordStore
+from .base import RecordStore, TransactionRecord
 
 
 @dataclass
 class ResultFactoryStore(RecordStore):
     result_factory: ResultFactory
-    cache: PersistedResult = None
+    cache: Optional[PersistedResult] = None
 
     def exists(self, key: str) -> bool:
         try:
-            result = self.read(key)
+            record = self.read(key)
+            if not record:
+                return False
+            result = record.result
             result.get(_sync=True)
             if result.expiration:
                 # if the result has an expiration,
@@ -29,25 +33,26 @@ class ResultFactoryStore(RecordStore):
         except Exception:
             return False
 
-    def read(self, key: str) -> BaseResult:
+    def read(self, key: str, holder: Optional[str] = None) -> TransactionRecord:
         if self.cache:
-            return self.cache
+            return TransactionRecord(key=key, result=self.cache)
         try:
             result = PersistedResult(
                 serializer_type=self.result_factory.serializer.type,
                 storage_block_id=self.result_factory.storage_block_id,
                 storage_key=key,
             )
-            return result
+            return TransactionRecord(key=key, result=result)
         except Exception:
             # this is a bit of a bandaid for functionality
             raise ValueError("Result could not be read")
 
-    def write(self, key: str, value: Any) -> BaseResult:
-        if isinstance(value, PersistedResult):
+    def write(self, key: str, result: Any, holder: Optional[str] = None) -> None:
+        if isinstance(result, PersistedResult):
             # if the value is already a persisted result, write it
-            value.write(_sync=True)
-            return value
-        elif isinstance(value, BaseResult):
-            return value
-        return run_coro_as_sync(self.result_factory.create_result(obj=value, key=key))
+            result.write(_sync=True)
+        elif not isinstance(result, BaseResult):
+            run_coro_as_sync(self.result_factory.create_result(obj=result, key=key))
+
+    def supports_isolation_level(self, isolation_level: IsolationLevel) -> bool:
+        return isolation_level == IsolationLevel.READ_COMMITTED

prefect/results.py

@@ -672,6 +672,17 @@ class PersistedResult(BaseResult):
 
         return result
 
+    def __eq__(self, other):
+        if not isinstance(other, PersistedResult):
+            return False
+        return (
+            self.type == other.type
+            and self.serializer_type == other.serializer_type
+            and self.storage_key == other.storage_key
+            and self.storage_block_id == other.storage_block_id
+            and self.expiration == other.expiration
+        )
+
 
 class PersistedResultBlob(BaseModel):
     """