Merge remote-tracking branch 'origin/HH-224752' into EXP-96105

frontik/app.py

@@ -19,15 +19,15 @@
 
 import frontik.producers.json_producer
 import frontik.producers.xml_producer
-from frontik import integrations, media_types
+from frontik import integrations, media_types, request_context
 from frontik.debug import get_frontik_and_apps_versions
 from frontik.handler import PageHandler, get_current_handler
-from frontik.handler_asgi import execute_page
+from frontik.handler_asgi import serve_request
 from frontik.handler_return_values import ReturnedValueHandlers, get_default_returned_value_handlers
 from frontik.integrations.statsd import StatsDClient, StatsDClientStub, create_statsd_client
 from frontik.options import options
 from frontik.process import WorkerState
-from frontik.routing import import_all_pages, router
+from frontik.routing import import_all_pages, method_not_allowed_router, not_found_router, router
 from frontik.service_discovery import UpstreamManager
 from frontik.util import check_request_id, generate_uniq_timestamp_request_id
 
@@ -102,6 +102,8 @@ def __init__(self) -> None:
         self.returned_value_handlers: ReturnedValueHandlers = get_default_returned_value_handlers()
 
         import_all_pages(self.app_module_name)
+        assert len(not_found_router.routes) < 2
+        assert len(method_not_allowed_router.routes) < 2
 
         self.ui_methods: dict = {}
         self.ui_modules: dict = {}
@@ -114,14 +116,14 @@ def __call__(self, tornado_request: httputil.HTTPServerRequest) -> Optional[Awai
         request_id = tornado_request.headers.get('X-Request-Id') or generate_uniq_timestamp_request_id()
         if options.validate_request_id:
             check_request_id(request_id)
+        tornado_request.request_id = request_id  # type: ignore
 
         async def _serve_tornado_request(
             frontik_app: FrontikApplication,
             _tornado_request: httputil.HTTPServerRequest,
-            _request_id: str,
             asgi_app: FrontikAsgiApp,
         ) -> None:
-            status, reason, headers, data = await execute_page(frontik_app, _tornado_request, _request_id, asgi_app)
+            status, reason, headers, data = await serve_request(frontik_app, _tornado_request, asgi_app)
 
             assert _tornado_request.connection is not None
             _tornado_request.connection.set_close_callback(None)  # type: ignore
@@ -131,7 +133,8 @@ async def _serve_tornado_request(
             _tornado_request.connection.finish()
             return await future
 
-        return asyncio.create_task(_serve_tornado_request(self, tornado_request, request_id, self.asgi_app))
+        with request_context.request_context(request_id):
+            return asyncio.create_task(_serve_tornado_request(self, tornado_request, self.asgi_app))
 
     def create_upstream_manager(
         self,

frontik/handler.py

@@ -194,7 +194,9 @@ def get_path_argument(self, name: str, default: Any = _ARG_DEFAULT) -> str:
             if default is _ARG_DEFAULT:
                 raise DefaultValueError(name)
             return default
-        value = _remove_control_chars_regex.sub(' ', value)
+
+        if isinstance(value, str):
+            value = _remove_control_chars_regex.sub(' ', value)
         return value
 
     @overload
@@ -400,6 +402,9 @@ async def delete(self, *args, **kwargs):
     async def head(self, *args, **kwargs):
         await self._execute_page()
 
+    async def options(self, *args, **kwargs):
+        await self._execute_page()
+
     async def _execute_page(self) -> None:
         self.stages_logger.commit_stage('prepare')
 
@@ -604,7 +609,10 @@ def send_error(self, status_code: int = 500, **kwargs: Any) -> None:
 
         try:
             self.write_error(status_code, **kwargs)
-        except Exception:
+        except Exception as exc:
+            if isinstance(exc, FinishSignal):
+                return
+
             self.log.exception('Uncaught exception in write_error')
             if not self._finished:
                 self.finish()

frontik/handler_asgi.py

@@ -6,14 +6,14 @@
 
 from fastapi.routing import APIRoute
 from tornado import httputil
-from tornado.httputil import HTTPHeaders
+from tornado.httputil import HTTPHeaders, HTTPServerRequest
 
 from frontik import media_types, request_context
 from frontik.debug import DebugMode, DebugTransform
 from frontik.handler import PageHandler, get_default_headers, log_request
 from frontik.handler_active_limit import request_limiter
 from frontik.json_builder import JsonBuilder
-from frontik.routing import find_route, get_allowed_methods
+from frontik.routing import find_route, get_allowed_methods, method_not_allowed_router, not_found_router
 
 if TYPE_CHECKING:
     from frontik.app import FrontikApplication, FrontikAsgiApp
@@ -22,77 +22,119 @@
 log = logging.getLogger('handler')
 
 
-async def execute_page(
+async def serve_request(
     frontik_app: FrontikApplication,
-    tornado_request: httputil.HTTPServerRequest,
-    request_id: str,
+    tornado_request: HTTPServerRequest,
     asgi_app: FrontikAsgiApp,
 ) -> tuple[int, str, HTTPHeaders, bytes]:
-    with request_context.request_context(request_id), request_limiter(frontik_app.statsd_client) as accepted:
+    with request_limiter(frontik_app.statsd_client) as accepted:
         log.info('requested url: %s', tornado_request.uri)
-        tornado_request.request_id = request_id  # type: ignore
+        if not accepted:
+            log_request(tornado_request, http.client.SERVICE_UNAVAILABLE)
+            return make_not_accepted_response()
+
+        debug_mode = DebugMode(tornado_request)
+        if debug_mode.auth_failed():
+            assert debug_mode.failed_auth_header is not None
+            log_request(tornado_request, http.client.UNAUTHORIZED)
+            return make_debug_auth_failed_response(debug_mode.failed_auth_header)
+
         assert tornado_request.method is not None
+
         route, page_cls, path_params = find_route(tornado_request.path, tornado_request.method)
+        if route is None and tornado_request.method == 'HEAD':
+            route, page_cls, path_params = find_route(tornado_request.path, 'GET')
 
-        debug_mode = DebugMode(tornado_request)
         data: bytes
 
-        if not accepted:
-            status, reason, headers, data = make_not_accepted_response()
-        elif debug_mode.auth_failed():
-            assert debug_mode.failed_auth_header is not None
-            status, reason, headers, data = make_debug_auth_failed_response(debug_mode.failed_auth_header)
-        elif route is None:
-            status, reason, headers, data = await make_not_found_response(frontik_app, tornado_request)
+        if route is None:
+            status, reason, headers, data = await make_not_found_response(
+                frontik_app, asgi_app, tornado_request, debug_mode
+            )
         else:
-            request_context.set_handler_name(f'{route.endpoint.__module__}.{route.endpoint.__name__}')
-
-            if page_cls is not None:
-                status, reason, headers, data = await legacy_process_request(
-                    frontik_app, tornado_request, route, page_cls, path_params, debug_mode
-                )
-            else:
-                result = {'headers': get_default_headers()}
-                scope, receive, send = convert_tornado_request_to_asgi(
-                    frontik_app, tornado_request, route, path_params, debug_mode, result
-                )
-                await asgi_app(scope, receive, send)
-
-                status = result['status']
-                reason = httputil.responses.get(status, 'Unknown')
-                headers = HTTPHeaders(result['headers'])
-                data = result['data']
-
-                if not scope['json_builder'].is_empty():
-                    if data != b'null':
-                        raise RuntimeError('Cant have return and json.put at the same time')
-
-                    headers['Content-Type'] = media_types.APPLICATION_JSON
-                    data = scope['json_builder'].to_bytes()
-                    headers['Content-Length'] = str(len(data))
+            status, reason, headers, data = await execute_page(
+                frontik_app, asgi_app, tornado_request, route, page_cls, path_params, debug_mode
+            )
 
         if debug_mode.enabled:
             debug_transform = DebugTransform(frontik_app, debug_mode)
             status, headers, data = debug_transform.transform_chunk(tornado_request, status, headers, data)
             reason = httputil.responses.get(status, 'Unknown')
 
         log_request(tornado_request, status)
-
         return status, reason, headers, data
 
 
+async def execute_page(
+    frontik_app: FrontikApplication,
+    asgi_app: FrontikAsgiApp,
+    tornado_request: HTTPServerRequest,
+    route: APIRoute,
+    page_cls: type[PageHandler] | None,
+    path_params: dict,
+    debug_mode: DebugMode,
+) -> tuple[int, str, HTTPHeaders, bytes]:
+    request_context.set_handler_name(f'{route.endpoint.__module__}.{route.endpoint.__name__}')
+
+    if page_cls is not None:
+        return await execute_tornado_page(frontik_app, tornado_request, route, page_cls, path_params, debug_mode)
+
+    result: dict = {'headers': get_default_headers()}
+    scope, receive, send = convert_tornado_request_to_asgi(
+        frontik_app, tornado_request, route, path_params, debug_mode, result
+    )
+    await asgi_app(scope, receive, send)
+
+    status: int = result['status']
+    reason = httputil.responses.get(status, 'Unknown')
+    headers = HTTPHeaders(result['headers'])
+    data = result['data']
+
+    if not scope['json_builder'].is_empty():
+        if data != b'null':
+            raise RuntimeError('Cant have "return" and "json.put" at the same time')
+
+        headers['Content-Type'] = media_types.APPLICATION_JSON
+        data = scope['json_builder'].to_bytes()
+        headers['Content-Length'] = str(len(data))
+
+    return status, reason, headers, data
+
+
 async def make_not_found_response(
-    frontik_app: FrontikApplication, tornado_request: httputil.HTTPServerRequest
+    frontik_app: FrontikApplication,
+    asgi_app: FrontikAsgiApp,
+    tornado_request: httputil.HTTPServerRequest,
+    debug_mode: DebugMode,
 ) -> tuple[int, str, HTTPHeaders, bytes]:
-    allowed_methods = get_allowed_methods(tornado_request.path.strip('/'))
+    allowed_methods = get_allowed_methods(tornado_request.path)
     default_headers = get_default_headers()
-
-    if allowed_methods:
+    headers: Any
+
+    if allowed_methods and len(method_not_allowed_router.routes) != 0:
+        status, _, headers, data = await execute_page(
+            frontik_app,
+            asgi_app,
+            tornado_request,
+            method_not_allowed_router.routes[0],  # type: ignore
+            method_not_allowed_router._cls,
+            {'allowed_methods': allowed_methods},
+            debug_mode,
+        )
+    elif allowed_methods:
         status = 405
         headers = {'Allow': ', '.join(allowed_methods)}
         data = b''
-    elif hasattr(frontik_app, 'application_404_handler'):
-        status, headers, data = await frontik_app.application_404_handler(tornado_request)
+    elif len(not_found_router.routes) != 0:
+        status, _, headers, data = await execute_page(
+            frontik_app,
+            asgi_app,
+            tornado_request,
+            not_found_router.routes[0],  # type: ignore
+            not_found_router._cls,
+            {},
+            debug_mode,
+        )
     else:
         status, headers, data = build_error_data(404, 'Not Found')
 
@@ -126,7 +168,7 @@ def build_error_data(
     return status_code, headers, data
 
 
-async def legacy_process_request(
+async def execute_tornado_page(
     frontik_app: FrontikApplication,
     tornado_request: httputil.HTTPServerRequest,
     route: APIRoute,

frontik/routing.py

@@ -17,7 +17,7 @@
 routing_logger = logging.getLogger('frontik.routing')
 
 routers: list[APIRouter] = []
-_plain_routes: dict[tuple, tuple] = {}
+_plain_routes: dict[tuple[str, str], tuple[APIRoute, type[PageHandler] | None]] = {}
 _regex_mapping: list[tuple[re.Pattern, APIRoute, Type[PageHandler]]] = []
 
 
@@ -48,16 +48,22 @@ def head(self, path: str, cls: Optional[Type[PageHandler]] = None, **kwargs: Any
         self._cls = self._base_cls or cls
         return super().head(path, **kwargs)
 
-    def add_api_route(self, *args, **kwargs):
+    def options(self, path: str, cls: Optional[Type[PageHandler]] = None, **kwargs: Any) -> Callable:
+        self._cls = self._base_cls or cls
+        return super().options(path, **kwargs)
+
+    def add_api_route(self, *args: Any, cls: Optional[Type[PageHandler]] = None, **kwargs: Any) -> None:
         super().add_api_route(*args, **kwargs)
+        self._cls = self._base_cls or cls or self._cls
         route: APIRoute = self.routes[-1]  # type: ignore
         method = next(iter(route.methods), None)
+        assert method is not None
         path = route.path.strip('/')
 
         if _plain_routes.get((path, method), None) is not None:
             raise RuntimeError(f'route for {method} {path} already exists')
 
-        _plain_routes[(path, method)] = (route, self._cls)  # we need our routing, for get route object
+        _plain_routes[(path, method)] = (route, self._cls)  # we need our routing, for getting route object
 
 
 class FrontikRegexRouter(APIRouter):
@@ -87,6 +93,10 @@ def head(self, path: str, cls: Optional[Type[PageHandler]] = None, **kwargs: Any
         self._cls = self._base_cls or cls
         return super().head(path, **kwargs)
 
+    def options(self, path: str, cls: Optional[Type[PageHandler]] = None, **kwargs: Any) -> Callable:
+        self._cls = self._base_cls or cls
+        return super().options(path, **kwargs)
+
     def add_api_route(self, *args: Any, cls: Optional[Type[PageHandler]] = None, **kwargs: Any) -> None:
         super().add_api_route(*args, **kwargs)
         self._cls = self._base_cls or cls or self._cls
@@ -131,6 +141,8 @@ def import_all_pages(app_module: str) -> None:
 
 
 router = FrontikRouter()
+not_found_router = FrontikRouter()
+method_not_allowed_router = FrontikRouter()
 regex_router = FrontikRegexRouter()
 routers.extend((router, regex_router))
 
@@ -146,7 +158,7 @@ def _find_regex_route(
     return None, None, None
 
 
-def find_route(path: str, method: str) -> tuple[APIRoute, type, dict]:
+def find_route(path: str, method: str) -> tuple[APIRoute, type[PageHandler], dict]:
     route: APIRoute
     route, page_cls, path_params = _find_regex_route(path, method)  # type: ignore
 
@@ -164,7 +176,10 @@ def find_route(path: str, method: str) -> tuple[APIRoute, type, dict]:
 def get_allowed_methods(path: str) -> list[str]:
     allowed_methods = []
     for method in ('GET', 'POST', 'PUT', 'DELETE', 'HEAD'):
-        route, _ = _plain_routes.get((path, method), (None, None))
+        route, _ = _plain_routes.get((path.strip('/'), method), (None, None))
+        if route is None:
+            route, _, _ = _find_regex_route(path, method)
+
         if route is not None:
             allowed_methods.append(method)