prefect-client 2.16.2__py3-none-any.whl → 2.16.3__py3-none-any.whl

prefect/task_engine.py

@@ -1,5 +1,6 @@
 from contextlib import AsyncExitStack
 from typing import (
+    Any,
     Dict,
     Iterable,
     Optional,
@@ -9,6 +10,7 @@ import anyio
 from typing_extensions import Literal
 
 from prefect._internal.concurrency.api import create_call, from_async, from_sync
+from prefect.client.orchestration import PrefectClient
 from prefect.client.schemas.objects import TaskRun
 from prefect.context import EngineContext
 from prefect.engine import (
@@ -30,11 +32,11 @@ async def submit_autonomous_task_run_to_engine(
     task: Task,
     task_run: TaskRun,
     task_runner: BaseTaskRunner,
-    parameters: Optional[Dict] = None,
+    parameters: Optional[Dict[str, Any]] = None,
     wait_for: Optional[Iterable[PrefectFuture]] = None,
     mapped: bool = False,
     return_type: EngineReturnType = "future",
-    client=None,
+    client: Optional[PrefectClient] = None,
 ) -> PrefectFuture:
     async with AsyncExitStack() as stack:
         parameters = parameters or {}
@@ -58,11 +60,17 @@ async def submit_autonomous_task_run_to_engine(
                 task_runner=task_runner,
             )
             if task.isasync:
-                future = await from_async.wait_for_call_in_loop_thread(begin_run)
+                future_result_or_state = await from_async.wait_for_call_in_loop_thread(
+                    begin_run
+                )
             else:
-                future = from_sync.wait_for_call_in_loop_thread(begin_run)
+                future_result_or_state = from_sync.wait_for_call_in_loop_thread(
+                    begin_run
+                )
 
-            await wait_for_task_runs_and_report_crashes(
-                task_run_futures=[future],
-                client=client,
-            )
+            if return_type == "future":
+                await wait_for_task_runs_and_report_crashes(
+                    task_run_futures=[future_result_or_state],
+                    client=client,
+                )
+            return future_result_or_state

prefect/tasks.py

@@ -555,6 +555,7 @@ class Task(Generic[P, R]):
         the result is wrapped in a Prefect State which provides error handling.
         """
         from prefect.engine import enter_task_run_engine
+        from prefect.task_engine import submit_autonomous_task_run_to_engine
         from prefect.task_runners import SequentialTaskRunner
 
         # Convert the call args/kwargs to a parameter dict
@@ -568,6 +569,21 @@ class Task(Generic[P, R]):
                 self.isasync, self.name, parameters, self.viz_return_value
             )
 
+        if (
+            PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING.value()
+            and not FlowRunContext.get()
+        ):
+            from prefect import get_client
+
+            return submit_autonomous_task_run_to_engine(
+                task=self,
+                task_run=None,
+                task_runner=SequentialTaskRunner(),
+                parameters=parameters,
+                return_type=return_type,
+                client=get_client(),
+            )
+
         return enter_task_run_engine(
             self,
             parameters=parameters,
@@ -792,16 +808,16 @@ class Task(Generic[P, R]):
             PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING.value()
             and not FlowRunContext.get()
         ):
-            create_autonomous_task_run = create_call(
+            create_autonomous_task_run_call = create_call(
                 create_autonomous_task_run, task=self, parameters=parameters
             )
             if self.isasync:
                 return from_async.wait_for_call_in_loop_thread(
-                    create_autonomous_task_run
+                    create_autonomous_task_run_call
                 )
             else:
                 return from_sync.wait_for_call_in_loop_thread(
-                    create_autonomous_task_run
+                    create_autonomous_task_run_call
                 )
 
         return enter_task_run_engine(
@@ -968,7 +984,7 @@ class Task(Generic[P, R]):
             [[11, 21], [12, 22], [13, 23]]
         """
 
-        from prefect.engine import enter_task_run_engine
+        from prefect.engine import begin_task_map, enter_task_run_engine
 
         # Convert the call args/kwargs to a parameter dict; do not apply defaults
         # since they should not be mapped over
@@ -981,6 +997,25 @@ class Task(Generic[P, R]):
                 "`task.map()` is not currently supported by `flow.visualize()`"
             )
 
+        if (
+            PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING.value()
+            and not FlowRunContext.get()
+        ):
+            map_call = create_call(
+                begin_task_map,
+                task=self,
+                parameters=parameters,
+                flow_run_context=None,
+                wait_for=wait_for,
+                return_type=return_type,
+                task_runner=None,
+                autonomous=True,
+            )
+            if self.isasync:
+                return from_async.wait_for_call_in_loop_thread(map_call)
+            else:
+                return from_sync.wait_for_call_in_loop_thread(map_call)
+
         return enter_task_run_engine(
             self,
             parameters=parameters,

prefect/utilities/schema_tools/hydration.py

@@ -0,0 +1,218 @@
+import json
+from typing import Any, Callable, Dict, Optional
+
+from prefect._internal.pydantic import HAS_PYDANTIC_V2
+
+if HAS_PYDANTIC_V2:
+    from pydantic.v1 import BaseModel, Field
+else:
+    from pydantic import BaseModel, Field
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing_extensions import TypeAlias
+
+from prefect.server.models.variables import read_variables
+
+
+class HydrationContext(BaseModel):
+    workspace_variables: Dict[str, str] = Field(default_factory=dict)
+    raise_on_error: bool = Field(default=False)
+
+    @classmethod
+    async def build(
+        cls,
+        session: AsyncSession,
+        raise_on_error: bool = False,
+    ) -> "HydrationContext":
+        variables = await read_variables(
+            session=session,
+        )
+        return cls(
+            workspace_variables={
+                variable.name: variable.value for variable in variables
+            },
+            raise_on_error=raise_on_error,
+        )
+
+
+Handler: TypeAlias = Callable[[Dict, HydrationContext], Any]
+PrefectKind: TypeAlias = Optional[str]
+
+_handlers: Dict[PrefectKind, Handler] = {}
+
+
+class Placeholder:
+    def __eq__(self, other):
+        return isinstance(other, type(self))
+
+    @property
+    def is_error(self) -> bool:
+        return False
+
+
+class RemoveValue(Placeholder):
+    pass
+
+
+def _remove_value(value) -> bool:
+    return isinstance(value, RemoveValue)
+
+
+class HydrationError(Placeholder, Exception):
+    def __init__(self, detail: Optional[str] = None):
+        self.detail = detail
+
+    @property
+    def is_error(self) -> bool:
+        return True
+
+    @property
+    def message(self):
+        raise NotImplementedError("Must be implemented by subclass")
+
+    def __eq__(self, other):
+        return isinstance(other, type(self)) and self.message == other.message
+
+    def __str__(self):
+        return self.message
+
+
+class KeyNotFound(HydrationError):
+    @property
+    def message(self):
+        return f"Missing '{self.key}' key in __prefect object"
+
+    @property
+    def key(self) -> str:
+        raise NotImplementedError("Must be implemented by subclass")
+
+
+class ValueNotFound(KeyNotFound):
+    @property
+    def key(self):
+        return "value"
+
+
+class VariableNameNotFound(KeyNotFound):
+    @property
+    def key(self):
+        return "variable_name"
+
+
+class InvalidJSON(HydrationError):
+    @property
+    def message(self):
+        message = "Invalid JSON"
+        if self.detail:
+            message += f": {self.detail}"
+        return message
+
+
+class WorkspaceVariableNotFound(HydrationError):
+    @property
+    def variable_name(self) -> str:
+        assert self.detail is not None
+        return self.detail
+
+    @property
+    def message(self):
+        return f"Variable '{self.detail}' not found."
+
+
+def handler(kind: PrefectKind) -> Callable:
+    def decorator(func: Handler) -> Handler:
+        _handlers[kind] = func
+        return func
+
+    return decorator
+
+
+def call_handler(kind: PrefectKind, obj: Dict, ctx: HydrationContext) -> Any:
+    if kind not in _handlers:
+        return (obj or {}).get("value", None)
+
+    res = _handlers[kind](obj, ctx)
+    if ctx.raise_on_error and isinstance(res, HydrationError):
+        raise res
+    return res
+
+
+@handler("none")
+def null_handler(obj: Dict, ctx: HydrationContext):
+    if "value" in obj:
+        # null handler is a pass through, so we want to continue to hydrate
+        return _hydrate(obj["value"], ctx)
+    else:
+        return ValueNotFound()
+
+
+@handler("json")
+def json_handler(obj: Dict, ctx: HydrationContext):
+    if "value" in obj:
+        try:
+            return json.loads(obj["value"])
+        except json.decoder.JSONDecodeError as e:
+            return InvalidJSON(detail=str(e))
+    else:
+        # If `value` is not in the object, we need special handling to help
+        # the UI. For now if an object looks like {"__prefect_kind": "json"}
+        # We will remove it from the parent object. e.x.
+        # {"a": {"__prefect_kind": "json"}} -> {}
+        # or
+        # [{"__prefect_kind": "json"}] -> []
+        return RemoveValue()
+
+
+@handler("workspace_variable")
+def workspace_variable_handler(obj: Dict, ctx: HydrationContext):
+    if "variable_name" in obj:
+        variable = obj["variable_name"]
+        if variable in ctx.workspace_variables:
+            return ctx.workspace_variables[variable]
+        else:
+            return WorkspaceVariableNotFound(detail=variable)
+    else:
+        # Special handling if `variable_name` is not in the object.
+        # If an object looks like {"__prefect_kind": "workspace_variable"}
+        # we will remove it from the parent object. e.x.
+        # {"a": {"__prefect_kind": "workspace_variable"}} -> {}
+        # or
+        # [{"__prefect_kind": "workspace_variable"}] -> []
+        # or
+        # {"__prefect_kind": "workspace_variable"} -> {}
+        return RemoveValue()
+
+
+def hydrate(obj: Dict, ctx: Optional[HydrationContext] = None):
+    res = _hydrate(obj, ctx)
+
+    if _remove_value(res):
+        return {}
+
+    return res
+
+
+def _hydrate(obj, ctx: Optional[HydrationContext] = None):
+    if ctx is None:
+        ctx = HydrationContext()
+
+    prefect_object = isinstance(obj, dict) and "__prefect_kind" in obj
+
+    if prefect_object:
+        prefect_kind = obj.get("__prefect_kind")
+        return call_handler(prefect_kind, obj, ctx)
+    else:
+        if isinstance(obj, dict):
+            return {
+                key: hydrated_value
+                for key, value in obj.items()
+                if not _remove_value(hydrated_value := _hydrate(value, ctx))
+            }
+        elif isinstance(obj, list):
+            return [
+                hydrated_element
+                for element in obj
+                if not _remove_value(hydrated_element := _hydrate(element, ctx))
+            ]
+        else:
+            return obj

prefect/utilities/schema_tools/validation.py

@@ -0,0 +1,240 @@
+from collections import defaultdict, deque
+from copy import deepcopy
+from typing import Dict, List
+
+import jsonschema
+from jsonschema.exceptions import ValidationError as JSONSchemaValidationError
+from jsonschema.validators import Draft202012Validator, create
+
+from prefect.utilities.collections import remove_nested_keys
+from prefect.utilities.schema_tools.hydration import HydrationError, Placeholder
+
+
+class CircularSchemaRefError(Exception):
+    pass
+
+
+class ValidationError(Exception):
+    pass
+
+
+PLACEHOLDERS_VALIDATOR_NAME = "_placeholders"
+
+
+def _build_validator():
+    def _applicable_validators(schema):
+        # the default implementation returns `schema.items()`
+        return {**schema, PLACEHOLDERS_VALIDATOR_NAME: None}.items()
+
+    def _placeholders(validator, _, instance, schema):
+        if isinstance(instance, HydrationError):
+            yield JSONSchemaValidationError(instance.message)
+
+    validators = dict(Draft202012Validator.VALIDATORS)
+    validators.update({PLACEHOLDERS_VALIDATOR_NAME: _placeholders})
+
+    # It is necessary to `create` a new validator instead of using `extend` because
+    # the `extend` method does not accept an `application_validators` parameter.
+    # We want `_placeholders` to be applicable always, without needing to modify
+    # the schema itself.
+    return create(
+        meta_schema=Draft202012Validator.META_SCHEMA,
+        validators=validators,
+        version="prefect",
+        type_checker=Draft202012Validator.TYPE_CHECKER,
+        format_checker=Draft202012Validator.FORMAT_CHECKER,
+        id_of=Draft202012Validator.ID_OF,
+        applicable_validators=_applicable_validators,
+    )
+
+
+_VALIDATOR = _build_validator()
+
+
+def is_valid_schema(schema: Dict, preprocess: bool = True):
+    if preprocess:
+        schema = preprocess_schema(schema)
+    try:
+        if schema is not None:
+            _VALIDATOR.check_schema(schema, format_checker=_VALIDATOR.FORMAT_CHECKER)
+    except jsonschema.SchemaError as exc:
+        raise ValueError(f"Invalid schema: {exc.message}") from exc
+
+
+def validate(
+    obj: Dict,
+    schema: Dict,
+    raise_on_error: bool = False,
+    preprocess: bool = True,
+    ignore_required: bool = False,
+) -> List[JSONSchemaValidationError]:
+    if preprocess:
+        schema = preprocess_schema(schema)
+
+    if ignore_required:
+        schema = remove_nested_keys(["required"], schema)
+
+    if raise_on_error:
+        try:
+            jsonschema.validate(obj, schema, _VALIDATOR)
+        except RecursionError:
+            raise CircularSchemaRefError
+        except JSONSchemaValidationError as exc:
+            if exc.json_path == "$":
+                error_message = "Validation failed."
+            else:
+                error_message = (
+                    f"Validation failed for field {exc.json_path.replace('$.', '')!r}."
+                )
+            error_message += f" Failure reason: {exc.message}"
+            raise ValidationError(error_message) from exc
+        return []
+    else:
+        try:
+            validator = _VALIDATOR(schema, format_checker=_VALIDATOR.FORMAT_CHECKER)
+            errors = list(validator.iter_errors(obj))
+        except RecursionError:
+            raise CircularSchemaRefError
+        return errors
+
+
+def is_valid(
+    obj: Dict,
+    schema: Dict,
+) -> bool:
+    errors = validate(obj, schema)
+    return len(errors) == 0
+
+
+def prioritize_placeholder_errors(errors):
+    errors_by_path = defaultdict(list)
+    for error in errors:
+        path_str = "->".join(str(p) for p in error.relative_path)
+        errors_by_path[path_str].append(error)
+
+    filtered_errors = []
+    for path, grouped_errors in errors_by_path.items():
+        placeholders_errors = [
+            error
+            for error in grouped_errors
+            if error.validator == PLACEHOLDERS_VALIDATOR_NAME
+        ]
+
+        if placeholders_errors:
+            filtered_errors.extend(placeholders_errors)
+        else:
+            filtered_errors.extend(grouped_errors)
+
+    return filtered_errors
+
+
+def build_error_obj(errors: List[JSONSchemaValidationError]) -> Dict:
+    error_response: dict = {"errors": []}
+
+    # If multiple errors are present for the same path and one of them
+    # is a placeholder error, we want only want to use the placeholder error.
+    errors = prioritize_placeholder_errors(errors)
+
+    for error in errors:
+        # If the Placeholder is not representing an error, we can skip it
+        if isinstance(error.instance, Placeholder) and not error.instance.is_error:
+            continue
+
+        path = deque(error.relative_path)
+
+        # Required errors should be moved one level down to the property
+        # they're associated with, so we add an extra level to the path.
+        if error.validator == "required":
+            required_field = error.message.split(" ")[0].strip("'")
+            path.append(required_field)
+
+        current = error_response["errors"]
+
+        # error at the root, just append the error message
+        if not path:
+            current.append(error.message)
+
+        while path:
+            part = path.popleft()
+            if isinstance(part, int):
+                if not path:
+                    current.append({"index": part, "errors": [error.message]})
+                else:
+                    for entry in current:
+                        if entry.get("index") == part:
+                            current = entry["errors"]
+                            break
+                    else:
+                        new_entry = {"index": part, "errors": []}
+                        current.append(new_entry)
+                        current = new_entry["errors"]
+            else:
+                if not path:
+                    current.append({"property": part, "errors": [error.message]})
+                else:
+                    for entry in current:
+                        if entry.get("property") == part:
+                            current = entry.get("errors", [])
+                            break
+                    else:
+                        new_entry = {"property": part, "errors": []}
+                        current.append(new_entry)
+                        current = new_entry["errors"]
+
+    valid = len(error_response["errors"]) == 0
+    error_response["valid"] = valid
+
+    return error_response
+
+
+def _fix_null_typing(key: str, schema: Dict, required_fields: List[str]):
+    """
+    Pydantic V1 does not generate a valid Draft2020-12 schema for null types.
+    """
+    if (
+        key not in required_fields
+        and "type" in schema
+        and schema.get("type") != "null"
+        and "default" not in schema
+    ):
+        schema["anyOf"] = [{"type": schema["type"]}, {"type": "null"}]
+        del schema["type"]
+
+
+def _fix_tuple_items(schema: Dict):
+    """
+    Pydantic V1 does not generate a valid Draft2020-12 schema for tuples.
+    """
+    if (
+        schema.get("items")
+        and isinstance(schema["items"], list)
+        and not schema.get("prefixItems")
+    ):
+        schema["prefixItems"] = deepcopy(schema["items"])
+        del schema["items"]
+
+
+def process_properties(properties, required_fields):
+    for key, schema in properties.items():
+        _fix_null_typing(key, schema, required_fields)
+        _fix_tuple_items(schema)
+
+        if "properties" in schema:
+            required_fields = schema.get("required", [])
+            process_properties(schema["properties"], required_fields)
+
+
+def preprocess_schema(schema):
+    schema = deepcopy(schema)
+
+    if "properties" in schema:
+        required_fields = schema.get("required", [])
+        process_properties(schema["properties"], required_fields)
+
+    if "definitions" in schema:  # Also process definitions for reused models
+        for definition in schema["definitions"].values():
+            if "properties" in definition:
+                required_fields = definition.get("required", [])
+                process_properties(definition["properties"], required_fields)
+
+    return schema