diff --git a/oxyroute/app.py b/oxyroute/app.py index a67f5d4..ea26cc9 100644 --- a/oxyroute/app.py +++ b/oxyroute/app.py @@ -98,6 +98,14 @@ def include_router( kw: dict[str, Any] = {k: v for k, v in merged.items() if k in allowed} reg(self, full, **kw)(handler) + def add_exception_handler( + self, exc_type: type[BaseException], handler: Callable[..., Any] + ) -> None: + """ + Register a global exception handler for a specific exception type. + """ + self._app.add_exception_handler(exc_type, handler) + def set_middleware(self, handler: Callable[..., Any] | None) -> None: """ One optional pre-route callback ``(scope, protocol)`` — return ``None`` to pass through. diff --git a/src/dispatch.rs b/src/dispatch.rs index 4448c4a..5867377 100644 --- a/src/dispatch.rs +++ b/src/dispatch.rs @@ -7,6 +7,7 @@ use jsonwebtoken::errors::ErrorKind; use jsonwebtoken::{decode, Validation}; use pyo3::prelude::*; use pyo3::types::{PyBytes, PyDict, PyList, PyString, PyTuple}; +use pyo3::IntoPyObjectExt; use serde_json::Value as JsonValue; use crate::config; @@ -164,13 +165,51 @@ fn send_python_error_sync( method: &str, path: &str, err: PyErr, + scope: Option<&pyo3::Bound<'_, PyAny>>, + state: Option<&std::sync::Arc>>, ) -> PyResult<()> { + if let (Some(sc), Some(st)) = (scope, state) { + let snap = st.read().hot_snapshot(); + for (exc_type, handler, is_async) in snap.exception_handlers.iter().rev() { + if let Ok(exc_obj) = err.clone_ref(py).into_bound_py_any(py) { + if exc_obj.is_instance(exc_type.bind(py)).unwrap_or(false) { + if *is_async { + log::error!( + "Async exception handler cannot be called in sync route fallback: {}", + method + ); + continue; + } + let exc_obj_any = exc_obj.clone().into_any(); + if let Ok(res) = handler.bind(py).call1((sc.clone(), exc_obj_any)) { + match map_handler_return(py, &res.clone().unbind()) { + Ok(mapped) => { + return send_handler_map_inline( + py, + protocol, + method == "HEAD", + mapped, + ) + } + Err(e) => { + log::error!( + "Exception handler returned invalid type or map failed: {:?}", + e + ); + } + } + } + } + } + } + } if try_http_exception_sync(py, protocol, &err)? { return Ok(()); } send_internal_error_sync(py, protocol, method, path, err) } +#[allow(clippy::too_many_arguments)] fn run_trivial_sync_route( py: Python<'_>, protocol: &Py, @@ -178,22 +217,24 @@ fn run_trivial_sync_route( path: &str, is_head: bool, entry: &RouteEntry, + scope: &pyo3::Bound<'_, PyAny>, + state: &std::sync::Arc>, ) -> PyResult<()> { let handler = entry.handler.bind(py); let out = match handler.call0() { Ok(x) => x.unbind(), Err(e) => { - return send_python_error_sync(py, protocol, method, path, e); + return send_python_error_sync(py, protocol, method, path, e, Some(scope), Some(state)); } }; let mapped = match map_handler_return(py, &out) { Ok(m) => m, Err(e) => { - return send_python_error_sync(py, protocol, method, path, e); + return send_python_error_sync(py, protocol, method, path, e, Some(scope), Some(state)); } }; if let Err(e) = send_handler_map_inline(py, protocol, is_head, mapped) { - return send_python_error_sync(py, protocol, method, path, e); + return send_python_error_sync(py, protocol, method, path, e, Some(scope), Some(state)); } Ok(()) } @@ -204,7 +245,59 @@ async fn send_python_error( method: &str, path: &str, err: PyErr, + scope: Option<&Py>, + state: Option<&std::sync::Arc>>, ) -> PyResult { + if let (Some(sc), Some(st)) = (scope, state) { + let snap = st.read().hot_snapshot(); + let coro_or_res = Python::with_gil(|py| -> PyResult, bool)>> { + for (exc_type, handler, is_async) in snap.exception_handlers.iter().rev() { + if let Ok(exc_obj) = err.clone_ref(py).into_bound_py_any(py) { + if exc_obj.is_instance(exc_type.bind(py)).unwrap_or(false) { + let exc_obj_any = exc_obj.clone().into_any(); + if let Ok(res) = handler.bind(py).call1((sc.bind(py).clone(), exc_obj_any)) + { + return Ok(Some((res.unbind(), *is_async))); + } + } + } + } + Ok(None) + }); + + if let Ok(Some((res_py, is_async))) = coro_or_res { + let final_res = if is_async { + let fut = Python::with_gil(|py| { + pyo3_async_runtimes::tokio::into_future(res_py.bind(py).clone()) + }); + if let Ok(f) = fut { + match f.await { + Ok(x) => x, + Err(e) => return send_internal_error(protocol, method, path, e).await, + } + } else { + res_py + } + } else { + res_py + }; + + let mapped_res = Python::with_gil(|py| map_handler_return(py, &final_res)); + if let Ok(mapped) = mapped_res { + let res = Python::with_gil(|py| { + send_handler_map_inline(py, protocol, method == "HEAD", mapped) + }); + if res.is_ok() { + return Ok(Python::with_gil(|py| py.None())); + } + } else { + log::error!( + "Async exception handler returned invalid type or map failed: {:?}", + mapped_res.err() + ); + } + } + } if let Some(res) = try_http_exception(protocol, &err).await? { return Ok(res); } @@ -298,7 +391,16 @@ pub fn try_rsgi_sync_short_circuit( return Err(pyo3::exceptions::PyRuntimeError::new_err("route index")); }; if route_is_trivial_sync(entry) { - run_trivial_sync_route(py, &protocol_py, &method, &path, is_head, entry)?; + run_trivial_sync_route( + py, + &protocol_py, + &method, + &path, + is_head, + entry, + scope, + state, + )?; return Ok(Some(py.None())); } Ok(None) @@ -420,7 +522,8 @@ pub async fn run_rsgi( }) { Ok(x) => x, Err(e) => { - return send_python_error(&protocol, &method, &path, e).await; + return send_python_error(&protocol, &method, &path, e, Some(&scope), Some(&state)) + .await; } }; let skip = Python::with_gil(|py| out.bind(py).is_none()); @@ -493,7 +596,10 @@ pub async fn run_rsgi( send_handler_map_inline(py, &protocol, is_head, mapped) }) { Ok(()) => Ok(Python::with_gil(|py| py.None())), - Err(e) => send_python_error(&protocol, &method, &path, e).await, + Err(e) => { + send_python_error(&protocol, &method, &path, e, Some(&scope), Some(&state)) + .await + } }; } } @@ -755,7 +861,15 @@ pub async fn run_rsgi( let ct = match ct { Ok(x) => x, Err(e) => { - return send_python_error(&protocol, &method, &path, e).await; + return send_python_error( + &protocol, + &method, + &path, + e, + Some(&scope), + Some(&state), + ) + .await; } }; if ct.as_deref().map(str::is_empty) != Some(false) { @@ -820,7 +934,8 @@ pub async fn run_rsgi( }) { Ok(o) => Some(o), Err(e) => { - return send_python_error(&protocol, &method, &path, e).await; + return send_python_error(&protocol, &method, &path, e, Some(&scope), Some(&state)) + .await; } } } else { @@ -848,7 +963,15 @@ pub async fn run_rsgi( }) { Ok(x) => x, Err(e) => { - return send_python_error(&protocol, &method, &path, e).await; + return send_python_error( + &protocol, + &method, + &path, + e, + Some(&scope), + Some(&state), + ) + .await; } }; let fut = match Python::with_gil(|py| { @@ -857,13 +980,29 @@ pub async fn run_rsgi( }) { Ok(f) => f, Err(e) => { - return send_python_error(&protocol, &method, &path, e).await; + return send_python_error( + &protocol, + &method, + &path, + e, + Some(&scope), + Some(&state), + ) + .await; } }; match fut.await { Ok(x) => x, Err(e) => { - return send_python_error(&protocol, &method, &path, e).await; + return send_python_error( + &protocol, + &method, + &path, + e, + Some(&scope), + Some(&state), + ) + .await; } } } else { @@ -886,7 +1025,15 @@ pub async fn run_rsgi( }) { Ok(x) => x, Err(e) => { - return send_python_error(&protocol, &method, &path, e).await; + return send_python_error( + &protocol, + &method, + &path, + e, + Some(&scope), + Some(&state), + ) + .await; } } }; @@ -905,7 +1052,17 @@ pub async fn run_rsgi( if let Some(pool) = snapshot.db_pool.as_ref() { match crate::db::execute_query(pool, &db_query).await { Ok(res) => res, - Err(e) => return send_python_error(&protocol, &method, &path, e).await, + Err(e) => { + return send_python_error( + &protocol, + &method, + &path, + e, + Some(&scope), + Some(&state), + ) + .await + } } } else { return send_python_error( @@ -915,12 +1072,17 @@ pub async fn run_rsgi( pyo3::exceptions::PyRuntimeError::new_err( "DBQuery returned by dependency but no database pool configured", ), + Some(&scope), + Some(&state), ) .await; } } Ok(None) => o, - Err(e) => return send_python_error(&protocol, &method, &path, e).await, + Err(e) => { + return send_python_error(&protocol, &method, &path, e, Some(&scope), Some(&state)) + .await + } }; dep_out.push(resolved); @@ -1009,7 +1171,8 @@ pub async fn run_rsgi( }) { Ok(x) => x, Err(e) => { - return send_python_error(&protocol, &method, &path, e).await; + return send_python_error(&protocol, &method, &path, e, Some(&scope), Some(&state)) + .await; } }; let handler_out: PyObject = if run_async { @@ -1019,13 +1182,15 @@ pub async fn run_rsgi( }) { Ok(f) => f, Err(e) => { - return send_python_error(&protocol, &method, &path, e).await; + return send_python_error(&protocol, &method, &path, e, Some(&scope), Some(&state)) + .await; } }; match fut.await { Ok(x) => x, Err(e) => { - return send_python_error(&protocol, &method, &path, e).await; + return send_python_error(&protocol, &method, &path, e, Some(&scope), Some(&state)) + .await; } } } else { @@ -1100,7 +1265,7 @@ pub async fn run_rsgi( send_handler_map_inline(py, &protocol, is_head, mapped) }) { Ok(()) => Ok(Python::with_gil(|py| py.None())), - Err(e) => send_python_error(&protocol, &method, &path, e).await, + Err(e) => send_python_error(&protocol, &method, &path, e, Some(&scope), Some(&state)).await, } } diff --git a/src/lib.rs b/src/lib.rs index f94ae4c..836900b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -358,6 +358,26 @@ impl App { /// Single optional pre-route hook. Return ``None`` to continue; otherwise the return value /// is mapped like a route handler (e.g. :class:`oxyroute.Response`, ``dict`` with ``status`` / ``body`` / ``headers``). + + #[pyo3(signature = (exc_type, handler))] + fn add_exception_handler( + &self, + exc_type: pyo3::Bound<'_, pyo3::types::PyType>, + handler: Py, + ) -> PyResult<()> { + let mut st = self.state.write(); + let is_async = pyo3::Python::with_gil(|py| -> PyResult { + let inspect = py.import("inspect")?; + inspect + .getattr("iscoroutinefunction")? + .call1((&handler,))? + .extract::() + }) + .unwrap_or(false); + Arc::make_mut(&mut st.exception_handlers).push((exc_type.unbind(), handler, is_async)); + Ok(()) + } + fn set_middleware(&self, handler: Option>) -> PyResult<()> { let mut st = self.state.write(); if let Some(h) = handler { diff --git a/src/state.rs b/src/state.rs index 15ce5f6..ae2256e 100644 --- a/src/state.rs +++ b/src/state.rs @@ -75,6 +75,8 @@ pub fn route_is_trivial_sync(entry: &RouteEntry) -> bool { entry.trivial_sync } +pub type ExceptionHandlerList = Arc, Py, bool)>>; + pub struct AppState { /// Wrapped in `Arc>` so the hot path can clone a cheap pointer **once** per request and /// release [`AppState`]'s `RwLock` immediately. Mutation goes through [`Arc::make_mut`]. @@ -100,6 +102,7 @@ pub struct AppState { pub request_middleware: Arc>>, /// Stack of `(scope, response_dict) -> Response | dict` response hooks. Runs before CORS/Security headers. pub response_middleware: Arc>>, + pub exception_handlers: ExceptionHandlerList, /// Optional Python CORS config (e.g. :class:`oxyroute.cors.CORSConfig`) for response headers. pub cors: Option>, /// Optional :class:`oxyroute.security_headers.SecurityHeadersConfig` (or compatible @@ -132,6 +135,7 @@ impl AppState { include_openapi: true, request_middleware: Arc::new(Vec::new()), response_middleware: Arc::new(Vec::new()), + exception_handlers: Arc::new(Vec::new()), cors: None, security_headers: None, db_pool: None, @@ -153,6 +157,7 @@ impl AppState { security_headers: self.security_headers.clone(), request_middleware: Arc::clone(&self.request_middleware), response_middleware: Arc::clone(&self.response_middleware), + exception_handlers: Arc::clone(&self.exception_handlers), include_openapi: self.include_openapi, db_pool: self.db_pool.clone(), } @@ -183,6 +188,7 @@ pub struct HotSnapshot { pub security_headers: Option>, pub request_middleware: Arc>>, pub response_middleware: Arc>>, + pub exception_handlers: ExceptionHandlerList, pub include_openapi: bool, pub db_pool: Option, } diff --git a/tests/test_exception_handlers.py b/tests/test_exception_handlers.py new file mode 100644 index 0000000..3ef68d0 --- /dev/null +++ b/tests/test_exception_handlers.py @@ -0,0 +1,58 @@ +import asyncio + +import httpx +from oxyroute import App, Response +from tests._rsgi_test_transport import asgi_test_app + + +def test_exception_handlers(): + class CustomError(Exception): + def __init__(self, msg: str): + self.msg = msg + + class SubCustomError(CustomError): + pass + + app = App() + + # Note: the `add_exception_handler` method might be used as a decorator or a regular method. + # We didn't implement it as a decorator returning the function, but in our `app.py` it's just: + # def add_exception_handler(self, exc_type: type[BaseException], handler: Callable[..., Any]) -> None: + # So we call it directly. + + def handle_custom_error(scope, exc): + return Response(status=400, body=exc.msg.encode()) + + async def handle_sub_custom_error(scope, exc): + return {"status": 418, "body": "sub error"} + + app.add_exception_handler(CustomError, handle_custom_error) + app.add_exception_handler(SubCustomError, handle_sub_custom_error) + + @app.get("/error1") + def error1() -> str: + raise CustomError("test1") + + @app.get("/error2") + async def error2() -> str: + raise SubCustomError("test2") + + @app.get("/unhandled") + def unhandled() -> str: + raise ValueError("oh no") + + async def _run() -> None: + transport = httpx.ASGITransport(app=asgi_test_app(app)) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as c: + r1 = await c.get("/error1") + assert r1.status_code == 400 + assert r1.text == "test1" + + r2 = await c.get("/error2") + assert r2.status_code == 418 + assert r2.text == "sub error" + + r3 = await c.get("/unhandled") + assert r3.status_code == 500 + + asyncio.run(_run())