mirror of
https://github.com/etesync/server
synced 2024-11-22 08:48:07 +00:00
80 lines
3.0 KiB
Python
80 lines
3.0 KiB
Python
import typing as t
|
|
|
|
from fastapi.routing import APIRoute, get_request_handler
|
|
from pydantic import BaseModel
|
|
from starlette.requests import Request
|
|
from starlette.responses import Response
|
|
|
|
from .db_hack import django_db_cleanup_decorator
|
|
from .utils import msgpack_decode, msgpack_encode
|
|
|
|
|
|
class MsgpackRequest(Request):
|
|
media_type = "application/msgpack"
|
|
|
|
async def raw_body(self) -> bytes:
|
|
return await super().body()
|
|
|
|
async def body(self) -> bytes:
|
|
if not hasattr(self, "_json"):
|
|
body = await self.raw_body()
|
|
self._json = msgpack_decode(body)
|
|
return self._json
|
|
|
|
|
|
class MsgpackResponse(Response):
|
|
media_type = "application/msgpack"
|
|
|
|
def render(self, content: t.Optional[t.Any]) -> bytes:
|
|
if content is None:
|
|
return b""
|
|
|
|
if isinstance(content, BaseModel):
|
|
content = content.model_dump()
|
|
return msgpack_encode(content)
|
|
|
|
|
|
class MsgpackRoute(APIRoute):
|
|
# keep track of content-type -> request classes
|
|
REQUESTS_CLASSES = {MsgpackRequest.media_type: MsgpackRequest}
|
|
# keep track of content-type -> response classes
|
|
ROUTES_HANDLERS_CLASSES = {MsgpackResponse.media_type: MsgpackResponse}
|
|
|
|
def __init__(self, path: str, endpoint: t.Callable[..., t.Any], *args, **kwargs):
|
|
endpoint = django_db_cleanup_decorator(endpoint)
|
|
super().__init__(path, endpoint, *args, **kwargs)
|
|
|
|
def _get_media_type_route_handler(self, media_type):
|
|
return get_request_handler(
|
|
dependant=self.dependant,
|
|
body_field=self.body_field,
|
|
status_code=self.status_code,
|
|
# use custom response class or fallback on default self.response_class
|
|
response_class=self.ROUTES_HANDLERS_CLASSES.get(media_type, self.response_class),
|
|
response_field=self.response_field,
|
|
response_model_include=self.response_model_include,
|
|
response_model_exclude=self.response_model_exclude,
|
|
response_model_by_alias=self.response_model_by_alias,
|
|
response_model_exclude_unset=self.response_model_exclude_unset,
|
|
response_model_exclude_defaults=self.response_model_exclude_defaults,
|
|
response_model_exclude_none=self.response_model_exclude_none,
|
|
dependency_overrides_provider=self.dependency_overrides_provider,
|
|
)
|
|
|
|
def get_route_handler(self) -> t.Callable:
|
|
async def custom_route_handler(request: Request) -> Response:
|
|
content_type = request.headers.get("Content-Type")
|
|
if content_type is not None:
|
|
try:
|
|
request_cls = self.REQUESTS_CLASSES[content_type]
|
|
request = request_cls(request.scope, request.receive)
|
|
except KeyError:
|
|
# nothing registered to handle content_type, process given requests as-is
|
|
pass
|
|
|
|
accept = request.headers.get("Accept")
|
|
route_handler = self._get_media_type_route_handler(accept)
|
|
return await route_handler(request)
|
|
|
|
return custom_route_handler
|