Replace uvicorn access log with our logger (#453)

mlrun/api/main.py

@@ -1,5 +1,8 @@
+import uuid
+
 import fastapi
 import uvicorn
+import uvicorn.protocols.utils
 from fastapi.exception_handlers import http_exception_handler
 
 import mlrun.errors
@@ -51,6 +54,63 @@ async def http_status_error_handler(
     )
 
 
+def get_client_address(scope):
+    # uvicorn expects this to be a tuple while starlette test client sets it to be a list
+    if isinstance(scope.get("client"), list):
+        scope["client"] = tuple(scope.get("client"))
+    return uvicorn.protocols.utils.get_client_addr(scope)
+
+
+@app.middleware("http")
+async def log_request_response(request: fastapi.Request, call_next):
+    request_id = str(uuid.uuid4())
+    silent_logging_paths = [
+        "healthz",
+    ]
+    path_with_query_string = uvicorn.protocols.utils.get_path_with_query_string(
+        request.scope
+    )
+    if not any(
+        silent_logging_path in path_with_query_string
+        for silent_logging_path in silent_logging_paths
+    ):
+        logger.debug(
+            "Received request",
+            method=request.method,
+            client_address=get_client_address(request.scope),
+            http_version=request.scope["http_version"],
+            request_id=request_id,
+            uri=path_with_query_string,
+        )
+    try:
+        response = await call_next(request)
+    except Exception:
+        logger.warning(
+            "Request handling failed. Sending response",
+            # User middleware (like this one) runs after the exception handling middleware, the only thing running after
+            # it is Starletter's ServerErrorMiddleware which is responsible for catching any un-handled exception
+            # and transforming it to 500 response. therefore we can statically assign status code to 500
+            status_code=500,
+            request_id=request_id,
+            uri=path_with_query_string,
+            method=request.method,
+        )
+        raise
+    else:
+        if not any(
+            silent_logging_path in path_with_query_string
+            for silent_logging_path in silent_logging_paths
+        ):
+            logger.debug(
+                "Sending response",
+                status_code=response.status_code,
+                request_id=request_id,
+                uri=path_with_query_string,
+                method=request.method,
+            )
+        return response
+
+
 @app.on_event("startup")
 async def startup_event():
     logger.info("configuration dump", dumped_config=config.dump_yaml())
@@ -98,6 +158,7 @@ def main():
         host="0.0.0.0",
         port=config.httpdb.port,
         debug=config.httpdb.debug,
+        access_log=False,
     )
 
 

tests/api/test_logging_middleware.py

@@ -0,0 +1,118 @@
+import typing
+from http import HTTPStatus
+
+import fastapi
+import pydantic
+import pytest
+from fastapi.exception_handlers import http_exception_handler
+from fastapi.testclient import TestClient
+from sqlalchemy.orm import Session
+
+from mlrun.api.main import app
+from mlrun.utils import logger
+
+
+class HandledException1(Exception):
+    pass
+
+
+class HandledException2(Exception):
+    pass
+
+
+class UnhandledException(Exception):
+    pass
+
+
+@app.exception_handler(HandledException1)
+async def handler_returning_response(request: fastapi.Request, exc: HandledException1):
+    logger.warning("Handler caught HandledException1 exception, returning 204 response")
+    return fastapi.Response(status_code=HTTPStatus.NO_CONTENT.value)
+
+
+@app.exception_handler(HandledException2)
+async def handler_returning_http_exception(
+    request: fastapi.Request, exc: HandledException2
+):
+    logger.warning(
+        "Handler caught HandledException2 exception, returning HTTPException with 401"
+    )
+    return await http_exception_handler(
+        request, fastapi.HTTPException(status_code=HTTPStatus.UNAUTHORIZED.value)
+    )
+
+
+test_router = fastapi.APIRouter()
+
+
+@test_router.get("/success")
+def success():
+    logger.info("Success endpoint received request, returning 202")
+    return fastapi.Response(status_code=202)
+
+
+@test_router.get("/handled_exception_1")
+def handled_exception_1():
+    logger.info(
+        "handled_exception_1 endpoint received request, raising handled exception 1"
+    )
+    raise HandledException1("handled exception 1")
+
+
+@test_router.get("/handled_exception_2")
+def handled_exception_2():
+    logger.info(
+        "handled_exception_2 endpoint received request, raising handled exception 2"
+    )
+    raise HandledException2("handled exception 2")
+
+
+@test_router.get("/unhandled_exception")
+def unhandled_exception():
+    logger.info("unhandled endpoint received request, raising unhandled exception")
+    raise UnhandledException("Unhandled exception")
+
+
+class SomeScheme(pydantic.BaseModel):
+    id: str
+
+
+@test_router.post("/fastapi_handled_exception")
+def fastapi_handled_exception(model: SomeScheme):
+    logger.info("Should not get here, will fail on body validation")
+
+
+middleware_modes = [
+    "with_middleware",
+    "without_middleware",
+]
+
+
+# must add it here since we're adding routes
+@pytest.fixture(params=middleware_modes)
+def client(request) -> typing.Generator:
+    if request.param == "without_middleware":
+        # HACKY
+        app.user_middleware = []
+        app.middleware_stack = app.build_middleware_stack()
+    app.include_router(test_router, prefix="/test")
+    with TestClient(app) as c:
+        yield c
+
+
+def test_logging_middleware(db: Session, client: TestClient) -> None:
+    resp = client.get("/test/success")
+    assert resp.status_code == HTTPStatus.ACCEPTED.value
+
+    resp = client.get("/test/handled_exception_1")
+    assert resp.status_code == HTTPStatus.NO_CONTENT.value
+
+    resp = client.get("/test/handled_exception_2")
+    assert resp.status_code == HTTPStatus.UNAUTHORIZED.value
+
+    resp = client.post("/test/fastapi_handled_exception")
+    assert resp.status_code == HTTPStatus.UNPROCESSABLE_ENTITY.value
+
+    with pytest.raises(UnhandledException):
+        # In a real fastapi (and not test) unhandled exception returns 500
+        client.get("/test/unhandled_exception")

tests/conftest.py

@@ -41,6 +41,7 @@
 environ["PYTHONPATH"] = root_path
 environ["MLRUN_DBPATH"] = rundb_path
 environ["MLRUN_httpdb__dirpath"] = rundb_path
+environ["MLRUN_log_level"] = "DEBUG"
 pytest_plugins = ["tests.common_fixtures"]
 
 # import package stuff after setting env vars so it will take effect