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

prefect/telemetry/run_telemetry.py

@@ -1,22 +1,33 @@
 import time
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Dict, Optional
+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.schemas import TaskRun
+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]):
     """
@@ -36,67 +47,148 @@ class RunTelemetry:
     _tracer: "Tracer" = field(
         default_factory=lambda: get_tracer("prefect", prefect.__version__)
     )
-    _span = None
+    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,
-        task_run: TaskRun,
-        parameters: Optional[Dict[str, Any]] = None,
-        labels: Optional[Dict[str, Any]] = None,
+        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 task run.
+        Start a span for a run.
         """
         if parameters is None:
             parameters = {}
-        if labels is None:
-            labels = {}
+
         parameter_attributes = {
             f"prefect.run.parameter.{k}": type(v).__name__
             for k, v in parameters.items()
         }
-        self._span = self._tracer.start_span(
-            name=task_run.name,
+
+        # 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.type": "task",
-                "prefect.run.id": str(task_run.id),
-                "prefect.tags": task_run.tags,
+                "prefect.run.name": name or run.name,
+                "prefect.run.type": run_type,
+                "prefect.run.id": str(run.id),
+                "prefect.tags": run.tags,
                 **parameter_attributes,
-                **labels,
+                **{
+                    key: value
+                    for key, value in run.labels.items()
+                    if not key.startswith("__")  # exclude internal labels
+                },
             },
         )
 
-    def end_span_on_success(self, terminal_message: str) -> None:
+        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 task run on success.
+        End a span for a run on success.
         """
-        if self._span:
-            self._span.set_status(Status(StatusCode.OK), terminal_message)
-            self._span.end(time.time_ns())
-            self._span = None
+        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: str) -> None:
+    def end_span_on_failure(self, terminal_message: Optional[str] = None) -> None:
         """
-        End a span for a task run on failure.
+        End a span for a run on failure.
         """
-        if self._span:
-            self._span.set_status(Status(StatusCode.ERROR, terminal_message))
-            self._span.end(time.time_ns())
-            self._span = None
+        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: Exception) -> None:
+    def record_exception(self, exc: BaseException) -> None:
         """
         Record an exception on a span.
         """
-        if self._span:
-            self._span.record_exception(exc)
+        if self.span:
+            self.span.record_exception(exc)
 
     def update_state(self, new_state: State) -> None:
         """
-        Update a span with the state of a task run.
+        Update a span with the state of a run.
         """
-        if self._span:
-            self._span.add_event(
+        if self.span:
+            self.span.add_event(
                 new_state.name or new_state.type,
                 {
                     "prefect.state.message": new_state.message or "",
@@ -105,3 +197,35 @@ class RunTelemetry:
                     "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/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 = ["/", "%", "&", ">", "<"]
 

prefect/variables.py

@@ -1,13 +1,14 @@
-from typing import List, Optional
+from typing import Optional
 
 from pydantic import BaseModel, Field
 
+from prefect._internal.compatibility.async_dispatch import async_dispatch
 from prefect._internal.compatibility.migration import getattr_migration
+from prefect.client.orchestration import get_client
 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 MAX_VARIABLE_NAME_LENGTH, StrictVariableValue
-from prefect.utilities.asyncutils import sync_compatible
 
 
 class Variable(BaseModel):
@@ -31,19 +32,18 @@ class Variable(BaseModel):
         description="The value of the variable",
         examples=["my-value"],
     )
-    tags: Optional[List[str]] = Field(default=None)
+    tags: Optional[list[str]] = Field(default=None)
 
     @classmethod
-    @sync_compatible
-    async def set(
+    async def aset(
         cls,
         name: str,
         value: StrictVariableValue,
-        tags: Optional[List[str]] = None,
+        tags: Optional[list[str]] = None,
         overwrite: bool = False,
     ) -> "Variable":
         """
-        Sets a new variable. If one exists with the same name, must pass `overwrite=True`
+        Asynchronously sets a new variable. If one exists with the same name, must pass `overwrite=True`
 
         Returns the newly set variable object.
 
@@ -60,8 +60,8 @@ class Variable(BaseModel):
             from prefect.variables import Variable
 
             @flow
-            def my_flow():
-                Variable.set(name="my_var",value="test_value", tags=["hi", "there"], overwrite=True)
+            async def my_flow():
+                await Variable.aset(name="my_var",value="test_value", tags=["hi", "there"], overwrite=True)
             ```
         """
         client, _ = get_or_create_client()
@@ -87,14 +87,62 @@ class Variable(BaseModel):
         return cls.model_validate(var_dict)
 
     @classmethod
-    @sync_compatible
-    async def get(
+    @async_dispatch(aset)
+    def set(
+        cls,
+        name: str,
+        value: StrictVariableValue,
+        tags: Optional[list[str]] = None,
+        overwrite: bool = False,
+    ) -> "Variable":
+        """
+        Sets a new variable. If one exists with the same name, must pass `overwrite=True`
+
+        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.
+
+        Example:
+            Set a new variable and overwrite it if it already exists.
+
+            ```
+            from prefect.variables import Variable
+
+            @flow
+            def my_flow():
+                Variable.set(name="my_var",value="test_value", tags=["hi", "there"], overwrite=True)
+            ```
+        """
+        with get_client(sync_client=True) as client:
+            variable_exists = client.read_variable_by_name(name)
+            var_dict = {"name": name, "value": value, "tags": tags or []}
+
+            if variable_exists:
+                if not overwrite:
+                    raise ValueError(
+                        f"Variable {name!r} already exists. Use `overwrite=True` to update it."
+                    )
+                client.update_variable(variable=VariableUpdate.model_validate(var_dict))
+                variable = client.read_variable_by_name(name)
+                for key in var_dict.keys():
+                    var_dict.update({key: getattr(variable, key)})
+            else:
+                client.create_variable(variable=VariableCreate.model_validate(var_dict))
+
+            return cls.model_validate(var_dict)
+
+    @classmethod
+    async def aget(
         cls,
         name: str,
         default: StrictVariableValue = None,
     ) -> StrictVariableValue:
         """
-        Get a variable's value by name.
+        Asynchronously get a variable's value by name.
 
         If the variable does not exist, return the default value.
 
@@ -109,8 +157,8 @@ class Variable(BaseModel):
             from prefect.variables import Variable
 
             @flow
-            def my_flow():
-                var = Variable.get("my_var")
+            async def my_flow():
+                var = await Variable.aget("my_var")
             ```
         """
         client, _ = get_or_create_client()
@@ -119,10 +167,41 @@ class Variable(BaseModel):
         return variable.value if variable else default
 
     @classmethod
-    @sync_compatible
-    async def unset(cls, name: str) -> bool:
+    @async_dispatch(aget)
+    def get(
+        cls,
+        name: str,
+        default: StrictVariableValue = None,
+    ) -> StrictVariableValue:
         """
-        Unset a variable by name.
+        Get a variable's value by name.
+
+        If the variable does not exist, return the default value.
+
+        Args:
+            - name: The name of the variable value to get.
+            - default: The default value to return if the variable does not exist.
+
+        Example:
+            Get a variable's value by name.
+            ```python
+            from prefect import flow
+            from prefect.variables import Variable
+
+            @flow
+            def my_flow():
+                var = Variable.get("my_var")
+            ```
+        """
+        with get_client(sync_client=True) as client:
+            variable = client.read_variable_by_name(name)
+
+            return variable.value if variable else default
+
+    @classmethod
+    async def aunset(cls, name: str) -> bool:
+        """
+        Asynchronously unset a variable by name.
 
         Args:
             - name: The name of the variable to unset.
@@ -136,8 +215,8 @@ class Variable(BaseModel):
             from prefect.variables import Variable
 
             @flow
-            def my_flow():
-                Variable.unset("my_var")
+            async def my_flow():
+                await Variable.aunset("my_var")
             ```
         """
         client, _ = get_or_create_client()
@@ -147,5 +226,34 @@ class Variable(BaseModel):
         except ObjectNotFound:
             return False
 
+    @classmethod
+    @async_dispatch(aunset)
+    def unset(cls, name: str) -> bool:
+        """
+        Unset a variable by name.
+
+        Args:
+            - name: The name of the variable to unset.
+
+        Returns `True` if the variable was deleted, `False` if the variable did not exist.
+
+        Example:
+            Unset a variable by name.
+            ```python
+            from prefect import flow
+            from prefect.variables import Variable
+
+            @flow
+            def my_flow():
+                Variable.unset("my_var")
+            ```
+        """
+        with get_client(sync_client=True) as client:
+            try:
+                client.delete_variable_by_name(name=name)
+                return True
+            except ObjectNotFound:
+                return False
+
 
 __getattr__ = getattr_migration(__name__)

{prefect_client-3.1.6.dist-info → prefect_client-3.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.1.6
+Version: 3.1.7
 Summary: Workflow orchestration and management.
 Home-page: https://www.prefect.io
 Author: Prefect Technologies, Inc.
@@ -51,6 +51,7 @@ Requires-Dist: pydantic-extra-types<3.0.0,>=2.8.2
 Requires-Dist: pydantic-settings>2.2.1
 Requires-Dist: python-dateutil<3.0.0,>=2.8.2
 Requires-Dist: python-slugify<9.0,>=5.0
+Requires-Dist: python-socks[asyncio]<3.0,>=2.5.3
 Requires-Dist: pyyaml<7.0.0,>=5.4.1
 Requires-Dist: rfc3339-validator<0.2.0,>=0.1.4
 Requires-Dist: rich<14.0,>=11.0