Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/utilities/server.py: 52%
36 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-12-05 11:21 +0000
1"""
2Utilities for the Prefect REST API server.
3"""
5from collections.abc import Coroutine, Sequence 1a
6from contextlib import AsyncExitStack 1a
7from typing import TYPE_CHECKING, Any, Callable, get_type_hints 1a
9from fastapi import APIRouter, Request, Response, status 1a
10from fastapi.routing import APIRoute 1a
11from starlette.routing import BaseRoute 1a
12from starlette.routing import Route as StarletteRoute 1a
15def method_paths_from_routes(routes: Sequence[BaseRoute]) -> set[str]: 1a
16 """
17 Generate a set of strings describing the given routes in the format: <method> <path>
19 For example, "GET /logs/"
20 """
21 method_paths: set[str] = set()
22 for route in routes:
23 if isinstance(route, (APIRoute, StarletteRoute)):
24 for method in route.methods or ():
25 method_paths.add(f"{method} {route.path}")
27 return method_paths
30class PrefectAPIRoute(APIRoute): 1a
31 """
32 A FastAPIRoute class which attaches an async stack to requests that exits before
33 a response is returned.
35 Requests already have `request.scope['fastapi_astack']` which is an async stack for
36 the full scope of the request. This stack is used for managing contexts of FastAPI
37 dependencies. If we want to close a dependency before the request is complete
38 (i.e. before returning a response to the user), we need a stack with a different
39 scope. This extension adds this stack at `request.state.response_scoped_stack`.
40 """
42 def get_route_handler(self) -> Callable[[Request], Coroutine[Any, Any, Response]]: 1a
43 default_handler = super().get_route_handler() 1a
45 async def handle_response_scoped_depends(request: Request) -> Response: 1a
46 # Create a new stack scoped to exit before the response is returned
47 response = None
48 async with AsyncExitStack() as stack:
49 request.state.response_scoped_stack = stack
50 response = await default_handler(request)
52 if TYPE_CHECKING:
53 assert response is not None
54 return response
56 return handle_response_scoped_depends 1a
59class PrefectRouter(APIRouter): 1a
60 """
61 A base class for Prefect REST API routers.
62 """
64 def __init__(self, **kwargs: Any) -> None: 1a
65 kwargs.setdefault("route_class", PrefectAPIRoute) 1a
66 super().__init__(**kwargs) 1a
68 def add_api_route( 1a
69 self, path: str, endpoint: Callable[..., Any], **kwargs: Any
70 ) -> None:
71 """
72 Add an API route.
74 For routes that return content and have not specified a `response_model`,
75 use return type annotation to infer the response model.
77 For routes that return No-Content status codes, explicitly set
78 a `response_class` to ensure nothing is returned in the response body.
79 """
80 if kwargs.get("status_code") == status.HTTP_204_NO_CONTENT: 1a
81 # any routes that return No-Content status codes must
82 # explicitly set a response_class that will handle status codes
83 # and not return anything in the body
84 kwargs["response_class"] = Response 1a
85 if kwargs.get("response_model") is None: 85 ↛ 86line 85 didn't jump to line 86 because the condition on line 85 was never true1a
86 kwargs["response_model"] = get_type_hints(endpoint).get("return")
87 return super().add_api_route(path, endpoint, **kwargs) 1a