Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/utilities/server.py: 60%

1""" 

2Utilities for the Prefect REST API server. 

3""" 

4 

5from collections.abc import Coroutine, Sequence 

6from contextlib import AsyncExitStack 

7from typing import TYPE_CHECKING, Any, Callable, get_type_hints 

8 

9from fastapi import APIRouter, Request, Response, status 

10from fastapi.routing import APIRoute 

11from starlette.routing import BaseRoute 

12 

13 

14def method_paths_from_routes(routes: Sequence[BaseRoute]) -> set[str]: 

15 """ 

16 Generate a set of strings describing the given routes in the format: <method> <path> 

17 

18 For example, "GET /logs/" 

19 """ 

20 method_paths: set[str] = set() 

21 for route in routes: 

22 effective_route_contexts = getattr(route, "effective_route_contexts", None) 

23 route_candidates = ( 

24 effective_route_contexts() 

25 if callable(effective_route_contexts) 

26 else (route,) 

27 ) 

28 

29 for route_candidate in route_candidates: 

30 path = getattr(route_candidate, "path", None) 

31 methods = getattr(route_candidate, "methods", None) 

32 if path is None or methods is None: 

33 continue 

34 

35 for method in methods: 

36 method_paths.add(f"{method} {path}") 

37 

38 return method_paths 

39 

40 

41class PrefectAPIRoute(APIRoute): 

42 """ 

43 A FastAPIRoute class which attaches an async stack to requests that exits before 

44 a response is returned. 

45 

46 Requests already have `request.scope['fastapi_astack']` which is an async stack for 

47 the full scope of the request. This stack is used for managing contexts of FastAPI 

48 dependencies. If we want to close a dependency before the request is complete 

49 (i.e. before returning a response to the user), we need a stack with a different 

50 scope. This extension adds this stack at `request.state.response_scoped_stack`. 

51 """ 

52 

53 def get_route_handler(self) -> Callable[[Request], Coroutine[Any, Any, Response]]: 

54 default_handler = super().get_route_handler() 

55 

56 async def handle_response_scoped_depends(request: Request) -> Response: 

57 # Create a new stack scoped to exit before the response is returned 

58 response = None 

59 async with AsyncExitStack() as stack: 

60 request.state.response_scoped_stack = stack 

61 response = await default_handler(request) 

62 

63 if TYPE_CHECKING: 63 ↛ 64line 63 didn't jump to line 64 because the condition on line 63 was never true

64 assert response is not None 

65 return response 

66 

67 return handle_response_scoped_depends 

68 

69 

70class PrefectRouter(APIRouter): 

71 """ 

72 A base class for Prefect REST API routers. 

73 """ 

74 

75 def __init__(self, **kwargs: Any) -> None: 

76 kwargs.setdefault("route_class", PrefectAPIRoute) 

77 super().__init__(**kwargs) 

78 

79 def add_api_route( 

80 self, path: str, endpoint: Callable[..., Any], **kwargs: Any 

81 ) -> None: 

82 """ 

83 Add an API route. 

84 

85 For routes that return content and have not specified a `response_model`, 

86 use return type annotation to infer the response model. 

87 

88 For routes that return No-Content status codes, explicitly set 

89 a `response_class` to ensure nothing is returned in the response body. 

90 """ 

91 if kwargs.get("status_code") == status.HTTP_204_NO_CONTENT: 

92 # any routes that return No-Content status codes must 

93 # explicitly set a response_class that will handle status codes 

94 # and not return anything in the body 

95 kwargs["response_class"] = Response 

96 if kwargs.get("response_model") is None: 

97 kwargs["response_model"] = get_type_hints(endpoint).get("return") 

98 return super().add_api_route(path, endpoint, **kwargs)