prefect-client 3.1.5__py3-none-any.whl → 3.1.7__py3-none-any.whl

prefect/telemetry/processors.py

@@ -1,14 +1,17 @@
 import time
 from threading import Event, Lock, Thread
-from typing import Dict, Optional
+from typing import TYPE_CHECKING, Dict, Optional
 
 from opentelemetry.context import Context
-from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
-from opentelemetry.sdk.trace.export import SpanExporter
+from opentelemetry.sdk.trace import Span, SpanProcessor
+
+if TYPE_CHECKING:
+    from opentelemetry.sdk.trace import ReadableSpan, Span
+    from opentelemetry.sdk.trace.export import SpanExporter
 
 
 class InFlightSpanProcessor(SpanProcessor):
-    def __init__(self, span_exporter: SpanExporter):
+    def __init__(self, span_exporter: "SpanExporter"):
         self.span_exporter = span_exporter
         self._in_flight: Dict[int, Span] = {}
         self._lock = Lock()
@@ -26,7 +29,7 @@ class InFlightSpanProcessor(SpanProcessor):
                 if to_export:
                     self.span_exporter.export(to_export)
 
-    def _readable_span(self, span: Span) -> ReadableSpan:
+    def _readable_span(self, span: "Span") -> "ReadableSpan":
         readable = span._readable_span()
         readable._end_time = time.time_ns()
         readable._attributes = {
@@ -35,13 +38,13 @@ class InFlightSpanProcessor(SpanProcessor):
         }
         return readable
 
-    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
+    def on_start(self, span: "Span", parent_context: Optional[Context] = None) -> None:
         if not span.context or not span.context.trace_flags.sampled:
             return
         with self._lock:
             self._in_flight[span.context.span_id] = span
 
-    def on_end(self, span: ReadableSpan) -> None:
+    def on_end(self, span: "ReadableSpan") -> None:
         if not span.context or not span.context.trace_flags.sampled:
             return
         with self._lock:

prefect/telemetry/run_telemetry.py

@@ -0,0 +1,231 @@
+import time
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+from opentelemetry import propagate, trace
+from opentelemetry.context import Context
+from opentelemetry.propagators.textmap import Setter
+from opentelemetry.trace import (
+    Span,
+    Status,
+    StatusCode,
+    get_tracer,
+)
+from typing_extensions import TypeAlias
+
+import prefect
+from prefect.client.orchestration import PrefectClient, SyncPrefectClient
+from prefect.client.schemas import FlowRun, TaskRun
+from prefect.client.schemas.objects import State
+from prefect.context import FlowRunContext, TaskRunContext
+from prefect.types import KeyValueLabels
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import Tracer
+
+LABELS_TRACEPARENT_KEY = "__OTEL_TRACEPARENT"
+TRACEPARENT_KEY = "traceparent"
+
+FlowOrTaskRun: TypeAlias = Union[FlowRun, TaskRun]
+
+
+class OTELSetter(Setter[KeyValueLabels]):
+    """
+    A setter for OpenTelemetry that supports Prefect's custom labels.
+    """
+
+    def set(self, carrier: KeyValueLabels, key: str, value: str) -> None:
+        carrier[key] = value
+
+
+@dataclass
+class RunTelemetry:
+    """
+    A class for managing the telemetry of runs.
+    """
+
+    _tracer: "Tracer" = field(
+        default_factory=lambda: get_tracer("prefect", prefect.__version__)
+    )
+    span: Optional[Span] = None
+
+    async def async_start_span(
+        self,
+        run: FlowOrTaskRun,
+        client: PrefectClient,
+        name: Optional[str] = None,
+        parameters: Optional[dict[str, Any]] = None,
+    ):
+        traceparent, span = self._start_span(run, name, parameters)
+
+        if self._run_type(run) == "flow" and traceparent:
+            # Only explicitly update labels if the run is a flow as task runs
+            # are updated via events.
+            await client.update_flow_run_labels(
+                run.id, {LABELS_TRACEPARENT_KEY: traceparent}
+            )
+
+        return span
+
+    def start_span(
+        self,
+        run: FlowOrTaskRun,
+        client: SyncPrefectClient,
+        name: Optional[str] = None,
+        parameters: Optional[dict[str, Any]] = None,
+    ):
+        traceparent, span = self._start_span(run, name, parameters)
+
+        if self._run_type(run) == "flow" and traceparent:
+            # Only explicitly update labels if the run is a flow as task runs
+            # are updated via events.
+            client.update_flow_run_labels(run.id, {LABELS_TRACEPARENT_KEY: traceparent})
+
+        return span
+
+    def _start_span(
+        self,
+        run: FlowOrTaskRun,
+        name: Optional[str] = None,
+        parameters: Optional[dict[str, Any]] = None,
+    ) -> tuple[Optional[str], Span]:
+        """
+        Start a span for a run.
+        """
+        if parameters is None:
+            parameters = {}
+
+        parameter_attributes = {
+            f"prefect.run.parameter.{k}": type(v).__name__
+            for k, v in parameters.items()
+        }
+
+        # Use existing trace context if this run already has one (e.g., from
+        # server operations like Late), otherwise use parent's trace context if
+        # available (e.g., nested flow / task runs). If neither exists, this
+        # will be a root span (e.g., a top-level flow run).
+        if LABELS_TRACEPARENT_KEY in run.labels:
+            context = self._trace_context_from_labels(run.labels)
+        else:
+            parent_run = self._parent_run()
+            parent_labels = parent_run.labels if parent_run else {}
+            if LABELS_TRACEPARENT_KEY in parent_labels:
+                context = self._trace_context_from_labels(parent_labels)
+            else:
+                context = None
+
+        run_type = self._run_type(run)
+
+        self.span = self._tracer.start_span(
+            name=name or run.name,
+            context=context,
+            attributes={
+                "prefect.run.name": name or run.name,
+                "prefect.run.type": run_type,
+                "prefect.run.id": str(run.id),
+                "prefect.tags": run.tags,
+                **parameter_attributes,
+                **{
+                    key: value
+                    for key, value in run.labels.items()
+                    if not key.startswith("__")  # exclude internal labels
+                },
+            },
+        )
+
+        if traceparent := self._traceparent_from_span(self.span):
+            run.labels[LABELS_TRACEPARENT_KEY] = traceparent
+
+        return traceparent, self.span
+
+    def _run_type(self, run: FlowOrTaskRun) -> str:
+        return "task" if isinstance(run, TaskRun) else "flow"
+
+    def _trace_context_from_labels(
+        self, labels: Optional[KeyValueLabels]
+    ) -> Optional[Context]:
+        """Get trace context from run labels if it exists."""
+        if not labels or LABELS_TRACEPARENT_KEY not in labels:
+            return None
+        traceparent = labels[LABELS_TRACEPARENT_KEY]
+        carrier = {TRACEPARENT_KEY: traceparent}
+        return propagate.extract(carrier)
+
+    def _traceparent_from_span(self, span: Span) -> Optional[str]:
+        carrier = {}
+        propagate.inject(carrier, context=trace.set_span_in_context(span))
+        return carrier.get(TRACEPARENT_KEY)
+
+    def end_span_on_success(self) -> None:
+        """
+        End a span for a run on success.
+        """
+        if self.span:
+            self.span.set_status(Status(StatusCode.OK))
+            self.span.end(time.time_ns())
+            self.span = None
+
+    def end_span_on_failure(self, terminal_message: Optional[str] = None) -> None:
+        """
+        End a span for a run on failure.
+        """
+        if self.span:
+            self.span.set_status(
+                Status(StatusCode.ERROR, terminal_message or "Run failed")
+            )
+            self.span.end(time.time_ns())
+            self.span = None
+
+    def record_exception(self, exc: BaseException) -> None:
+        """
+        Record an exception on a span.
+        """
+        if self.span:
+            self.span.record_exception(exc)
+
+    def update_state(self, new_state: State) -> None:
+        """
+        Update a span with the state of a run.
+        """
+        if self.span:
+            self.span.add_event(
+                new_state.name or new_state.type,
+                {
+                    "prefect.state.message": new_state.message or "",
+                    "prefect.state.type": new_state.type,
+                    "prefect.state.name": new_state.name or new_state.type,
+                    "prefect.state.id": str(new_state.id),
+                },
+            )
+
+    def _parent_run(self) -> Union[FlowOrTaskRun, None]:
+        """
+        Identify the "parent run" for the current execution context.
+
+        Both flows and tasks can be nested "infinitely," and each creates a
+        corresponding context when executed. This method determines the most
+        appropriate parent context (either a task run or a flow run) based on
+        their relationship in the current hierarchy.
+
+        Returns:
+            FlowOrTaskRun: The parent run object (task or flow) if applicable.
+            None: If there is no parent context, implying the current run is the top-level parent.
+        """
+        parent_flow_run_context = FlowRunContext.get()
+        parent_task_run_context = TaskRunContext.get()
+
+        if parent_task_run_context and parent_flow_run_context:
+            # If both contexts exist, which is common for nested flows or tasks,
+            # check if the task's flow_run_id matches the current flow_run.
+            # If they match, the task is a child of the flow and is the parent of the current run.
+            flow_run_id = getattr(parent_flow_run_context.flow_run, "id", None)
+            if parent_task_run_context.task_run.flow_run_id == flow_run_id:
+                return parent_task_run_context.task_run
+            # Otherwise, assume the flow run is the entry point and is the parent.
+            return parent_flow_run_context.flow_run
+        elif parent_flow_run_context:
+            return parent_flow_run_context.flow_run
+        elif parent_task_run_context:
+            return parent_task_run_context.task_run
+
+        return None

prefect/transactions.py

@@ -23,7 +23,7 @@ from prefect.exceptions import (
     MissingContextError,
     SerializationError,
 )
-from prefect.logging.loggers import get_logger, get_run_logger
+from prefect.logging.loggers import LoggingAdapter, get_logger, get_run_logger
 from prefect.records import RecordStore
 from prefect.records.base import TransactionRecord
 from prefect.results import (
@@ -32,9 +32,9 @@ from prefect.results import (
     ResultStore,
     get_result_store,
 )
+from prefect.utilities._engine import get_hook_name
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.collections import AutoEnum
-from prefect.utilities.engine import _get_hook_name
 
 
 class IsolationLevel(AutoEnum):
@@ -72,13 +72,13 @@ class Transaction(ContextModel):
         default_factory=list
     )
     overwrite: bool = False
-    logger: Union[logging.Logger, logging.LoggerAdapter] = Field(
+    logger: Union[logging.Logger, LoggingAdapter] = 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")
+    __var__: ContextVar[Self] = ContextVar("transaction")
 
     def set(self, name: str, value: Any) -> None:
         """
@@ -209,7 +209,7 @@ class Transaction(ContextModel):
         self._token = self.__var__.set(self)
         return self
 
-    def __exit__(self, *exc_info):
+    def __exit__(self, *exc_info: Any):
         exc_type, exc_val, _ = exc_info
         if not self._token:
             raise RuntimeError(
@@ -254,7 +254,7 @@ class Transaction(ContextModel):
         ):
             self.state = TransactionState.COMMITTED
 
-    def read(self) -> Union["BaseResult", ResultRecord, None]:
+    def read(self) -> Union["BaseResult[Any]", ResultRecord[Any], None]:
         if self.store and self.key:
             record = self.store.read(key=self.key)
             if isinstance(record, ResultRecord):
@@ -354,8 +354,8 @@ class Transaction(ContextModel):
             self.rollback()
             return False
 
-    def run_hook(self, hook, hook_type: str) -> None:
-        hook_name = _get_hook_name(hook)
+    def run_hook(self, hook: Callable[..., Any], 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)
 
@@ -379,8 +379,8 @@ class Transaction(ContextModel):
     def stage(
         self,
         value: Any,
-        on_rollback_hooks: Optional[List] = None,
-        on_commit_hooks: Optional[List] = None,
+        on_rollback_hooks: Optional[list[Callable[..., Any]]] = None,
+        on_commit_hooks: Optional[list[Callable[..., Any]]] = None,
     ) -> None:
         """
         Stage a value to be committed later.
@@ -441,7 +441,7 @@ def transaction(
     isolation_level: Optional[IsolationLevel] = None,
     overwrite: bool = False,
     write_on_commit: bool = True,
-    logger: Union[logging.Logger, logging.LoggerAdapter, None] = None,
+    logger: Optional[Union[logging.Logger, LoggingAdapter]] = None,
 ) -> Generator[Transaction, None, None]:
     """
     A context manager for opening and managing a transaction.
@@ -465,9 +465,9 @@ def transaction(
         store = get_result_store()
 
     try:
-        logger = logger or get_run_logger()
+        _logger: Union[logging.Logger, LoggingAdapter] = logger or get_run_logger()
     except MissingContextError:
-        logger = get_logger("transactions")
+        _logger = get_logger("transactions")
 
     with Transaction(
         key=key,
@@ -476,6 +476,6 @@ def transaction(
         isolation_level=isolation_level,
         overwrite=overwrite,
         write_on_commit=write_on_commit,
-        logger=logger,
+        logger=_logger,
     ) as txn:
         yield txn

prefect/types/__init__.py

@@ -3,7 +3,8 @@ from typing import Annotated, Any, Dict, List, Optional, Set, TypeVar, Union
 from typing_extensions import Literal, TypeAlias
 import orjson
 import pydantic
-
+from pydantic_extra_types.pendulum_dt import DateTime as PydanticDateTime
+from pydantic_extra_types.pendulum_dt import Date as PydanticDate
 from pydantic import (
     BeforeValidator,
     Field,
@@ -34,6 +35,8 @@ TimeZone = Annotated[
     ),
 ]
 
+DateTime: TypeAlias = PydanticDateTime
+Date: TypeAlias = PydanticDate
 
 BANNED_CHARACTERS = ["/", "%", "&", ">", "<"]
 
@@ -96,7 +99,7 @@ def cast_none_to_empty_dict(value: Any) -> dict[str, Any]:
 
 
 KeyValueLabels = Annotated[
-    dict[str, Union[StrictBool, StrictInt, StrictFloat, str]],
+    Dict[str, Union[StrictBool, StrictInt, StrictFloat, str]],
     BeforeValidator(cast_none_to_empty_dict),
 ]
 
@@ -149,9 +152,6 @@ LogLevel = Annotated[
 ]
 
 
-KeyValueLabels: TypeAlias = dict[str, Union[StrictBool, StrictInt, StrictFloat, str]]
-
-
 def convert_none_to_empty_dict(v: Optional[KeyValueLabels]) -> KeyValueLabels:
     return v or {}
 

prefect/utilities/_engine.py

@@ -0,0 +1,96 @@
+"""Internal engine utilities"""
+
+
+from collections.abc import Callable
+from functools import partial
+from typing import TYPE_CHECKING, Any, Union
+from uuid import uuid4
+
+from prefect.context import FlowRunContext
+from prefect.flows import Flow
+from prefect.tasks import Task, TaskRunNameCallbackWithParameters
+
+
+def dynamic_key_for_task_run(
+    context: FlowRunContext, task: "Task[..., Any]", stable: bool = True
+) -> Union[int, str]:
+    if (
+        stable is False or context.detached
+    ):  # this task is running on remote infrastructure
+        return str(uuid4())
+    elif context.flow_run is None:  # this is an autonomous task run
+        context.task_run_dynamic_keys[task.task_key] = getattr(
+            task, "dynamic_key", str(uuid4())
+        )
+
+    elif task.task_key not in context.task_run_dynamic_keys:
+        context.task_run_dynamic_keys[task.task_key] = 0
+    else:
+        dynamic_key = context.task_run_dynamic_keys[task.task_key]
+        if TYPE_CHECKING:
+            assert isinstance(dynamic_key, int)
+        context.task_run_dynamic_keys[task.task_key] = dynamic_key + 1
+
+    return context.task_run_dynamic_keys[task.task_key]
+
+
+def resolve_custom_flow_run_name(
+    flow: "Flow[..., Any]", parameters: dict[str, Any]
+) -> str:
+    if callable(flow.flow_run_name):
+        flow_run_name = flow.flow_run_name()
+        if not TYPE_CHECKING:
+            if not isinstance(flow_run_name, str):
+                raise TypeError(
+                    f"Callable {flow.flow_run_name} for 'flow_run_name' returned type"
+                    f" {type(flow_run_name).__name__} but a string is required."
+                )
+    elif isinstance(flow.flow_run_name, str):
+        flow_run_name = flow.flow_run_name.format(**parameters)
+    else:
+        raise TypeError(
+            "Expected string or callable for 'flow_run_name'; got"
+            f" {type(flow.flow_run_name).__name__} instead."
+        )
+
+    return flow_run_name
+
+
+def resolve_custom_task_run_name(
+    task: "Task[..., Any]", parameters: dict[str, Any]
+) -> str:
+    if callable(task.task_run_name):
+        # If the callable accepts a 'parameters' kwarg, pass the entire parameters dict
+        if TaskRunNameCallbackWithParameters.is_callback_with_parameters(
+            task.task_run_name
+        ):
+            task_run_name = task.task_run_name(parameters=parameters)
+        else:
+            # If it doesn't expect parameters, call it without arguments
+            task_run_name = task.task_run_name()
+
+        if not TYPE_CHECKING:
+            if not isinstance(task_run_name, str):
+                raise TypeError(
+                    f"Callable {task.task_run_name} for 'task_run_name' returned type"
+                    f" {type(task_run_name).__name__} but a string is required."
+                )
+    elif isinstance(task.task_run_name, str):
+        task_run_name = task.task_run_name.format(**parameters)
+    else:
+        raise TypeError(
+            "Expected string or callable for 'task_run_name'; got"
+            f" {type(task.task_run_name).__name__} instead."
+        )
+
+    return task_run_name
+
+
+def get_hook_name(hook: Callable[..., Any]) -> str:
+    return (
+        hook.__name__
+        if hasattr(hook, "__name__")
+        else (
+            hook.func.__name__ if isinstance(hook, partial) else hook.__class__.__name__
+        )
+    )

prefect/utilities/annotations.py

@@ -1,33 +1,40 @@
 import warnings
-from abc import ABC
-from collections import namedtuple
-from typing import Generic, TypeVar
+from operator import itemgetter
+from typing import Any, cast
 
-T = TypeVar("T")
+from typing_extensions import Self, TypeVar
 
+T = TypeVar("T", infer_variance=True)
 
-class BaseAnnotation(
-    namedtuple("BaseAnnotation", field_names="value"), ABC, Generic[T]
-):
+
+class BaseAnnotation(tuple[T]):
     """
     Base class for Prefect annotation types.
 
-    Inherits from `namedtuple` for unpacking support in another tools.
+    Inherits from `tuple` for unpacking support in other tools.
     """
 
+    __slots__ = ()
+
+    def __new__(cls, value: T) -> Self:
+        return super().__new__(cls, (value,))
+
+    # use itemgetter to minimise overhead, just like namedtuple generated code would
+    value: T = cast(T, property(itemgetter(0)))
+
     def unwrap(self) -> T:
-        return self.value
+        return self[0]
 
-    def rewrap(self, value: T) -> "BaseAnnotation[T]":
+    def rewrap(self, value: T) -> Self:
         return type(self)(value)
 
-    def __eq__(self, other: "BaseAnnotation[T]") -> bool:
+    def __eq__(self, other: Any) -> bool:
         if type(self) is not type(other):
             return False
-        return self.unwrap() == other.unwrap()
+        return super().__eq__(other)
 
     def __repr__(self) -> str:
-        return f"{type(self).__name__}({self.value!r})"
+        return f"{type(self).__name__}({self[0]!r})"
 
 
 class unmapped(BaseAnnotation[T]):
@@ -38,9 +45,9 @@ class unmapped(BaseAnnotation[T]):
     operation instead of being split.
     """
 
-    def __getitem__(self, _) -> T:
+    def __getitem__(self, _: object) -> T:  # type: ignore[override]  # pyright: ignore[reportIncompatibleMethodOverride]
         # Internally, this acts as an infinite array where all items are the same value
-        return self.unwrap()
+        return super().__getitem__(0)
 
 
 class allow_failure(BaseAnnotation[T]):
@@ -87,14 +94,14 @@ class quote(BaseAnnotation[T]):
 
 
 # Backwards compatibility stub for `Quote` class
-class Quote(quote):
-    def __init__(self, expr):
+class Quote(quote[T]):
+    def __new__(cls, expr: T) -> Self:
         warnings.warn(
             "Use of `Quote` is deprecated. Use `quote` instead.",
             DeprecationWarning,
             stacklevel=2,
         )
-        super().__init__(expr)
+        return super().__new__(cls, expr)
 
 
 class NotSet: