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

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()
@@ -73,27 +73,76 @@ class Variable(BaseModel):
                 raise ValueError(
                     f"Variable {name!r} already exists. Use `overwrite=True` to update it."
                 )
-            await client.update_variable(variable=VariableUpdate(**var_dict))
+            await client.update_variable(
+                variable=VariableUpdate.model_validate(var_dict)
+            )
             variable = await client.read_variable_by_name(name)
-            var_dict = {
-                "name": variable.name,
-                "value": variable.value,
-                "tags": variable.tags or [],
-            }
+            for key in var_dict.keys():
+                var_dict.update({key: getattr(variable, key)})
         else:
-            await client.create_variable(variable=VariableCreate(**var_dict))
+            await client.create_variable(
+                variable=VariableCreate.model_validate(var_dict)
+            )
+
+        return cls.model_validate(var_dict)
+
+    @classmethod
+    @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 []}
 
-        return cls(**var_dict)
+            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
-    @sync_compatible
-    async def get(
+    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.
 
@@ -108,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()
@@ -118,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.
@@ -135,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()
@@ -146,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/workers/base.py

@@ -53,6 +53,7 @@ from prefect.states import (
     Pending,
     exception_to_failed_state,
 )
+from prefect.types import KeyValueLabels
 from prefect.utilities.dispatch import get_registry_for_type, register_base_type
 from prefect.utilities.engine import propose_state
 from prefect.utilities.services import critical_service_loop
@@ -988,7 +989,6 @@ class BaseWorker(abc.ABC):
         try:
             configuration = await self._get_configuration(flow_run)
             submitted_event = self._emit_flow_run_submitted_event(configuration)
-            await self._give_worker_labels_to_flow_run(flow_run.id)
             result = await self.run(
                 flow_run=flow_run,
                 task_status=task_status,
@@ -1222,13 +1222,20 @@ class BaseWorker(abc.ABC):
         Give this worker's identifying labels to the specified flow run.
         """
         if self._cloud_client:
-            await self._cloud_client.update_flow_run_labels(
-                flow_run_id,
-                {
-                    "prefect.worker.name": self.name,
-                    "prefect.worker.type": self.type,
-                },
-            )
+            labels: KeyValueLabels = {
+                "prefect.worker.name": self.name,
+                "prefect.worker.type": self.type,
+            }
+
+            if self._work_pool:
+                labels.update(
+                    {
+                        "prefect.work-pool.name": self._work_pool.name,
+                        "prefect.work-pool.id": str(self._work_pool.id),
+                    }
+                )
+
+            await self._cloud_client.update_flow_run_labels(flow_run_id, labels)
 
     async def __aenter__(self):
         self._logger.debug("Entering worker context...")

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: prefect-client
-Version: 3.1.5
+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
@@ -165,7 +166,9 @@ Start with our [friendly tutorial](https://docs.prefect.io/tutorials) or explore
 
 ## Join the community
 
-Prefect is made possible by the fastest growing community of thousands of friendly data engineers. Join us in building a new kind of workflow system. The [Prefect Slack community](https://prefect.io/slack) is a fantastic place to learn more about Prefect, ask questions, or get help with workflow design. All community forums, including code contributions, issue discussions, and slack messages are subject to our [Code of Conduct](https://discourse.prefect.io/faq).
+Prefect is made possible by the fastest growing community of thousands of friendly data engineers. Join us in building a new kind of workflow system. 
+The [Prefect Slack community](https://prefect.io/slack) is a fantastic place to learn more about Prefect, ask questions, or get help with workflow design. 
+All community forums, including code contributions, issue discussions, and Slack messages are subject to our [Code of Conduct](https://github.com/PrefectHQ/prefect/blob/main/CODE_OF_CONDUCT.md).
 
 ## Contribute