Add middleware

Author: saviorand

lightbug_api/__init__.mojo

@@ -3,40 +3,75 @@ from lightbug_api.context import Context
 from lightbug_api.response import Response
 from lightbug_api.routing import (
     BaseRequest,
+    ErrorHandler,
     FromReq,
     Handler,
     HandlerResponse,
+    Middleware,
+    MiddlewareEntry,
+    MiddlewareResult,
     PathPattern,
     RouteMatch,
     RootRouter,
     Router,
+    abort,
+    next,
 )
 
 
+# ---------------------------------------------------------- startup hook types
+
+comptime StartupHook = fn () raises
+
+struct StartupHookEntry(Copyable):
+    var hook: StartupHook
+
+    fn __init__(out self, hook: StartupHook):
+        self.hook = hook
+
+    fn __init__(out self, *, copy: Self):
+        self.hook = copy.hook
+
+
+# ----------------------------------------------------------------------- App
+
 struct App:
     """The top-level application — register routes then call ``run()``.
 
-    Example::
+    **Quick-start**::
 
         fn main() raises:
             var app = App()
 
+            # Routes
             app.get("/",              index)
             app.get("/users/{id}",    get_user)
             app.post("/users",        create_user)
             app.delete("/users/{id}", delete_user)
 
+            # Sub-router mounted at /v1
             var api = Router("v1")
             api.get("status", health)
-            app.add_router(api^)     # mounts at /v1/status
+            app.add_router(api^)
+
+            # Middleware — runs before every handler
+            app.use(require_auth)
+
+            # Error handler — catches unhandled exceptions from handlers
+            app.on_error(my_error_handler)
+
+            # Startup hook — runs once before the server starts
+            app.on_startup(connect_db)
 
             app.run()
     """
 
     var router: RootRouter
+    var startup_hooks: List[StartupHookEntry]
 
     def __init__(out self) raises:
         self.router = RootRouter()
+        self.startup_hooks = List[StartupHookEntry]()
 
     # ------------------------------------------ route registration
 
@@ -72,22 +107,78 @@ struct App:
         """Mount a sub-router under its path fragment."""
         self.router.add_router(router^)
 
+    # ------------------------------------------ middleware
+
+    def use(mut self, middleware: Middleware) -> None:
+        """Add a middleware function that runs before every handler.
+
+        Middleware runs in registration order. Use ``next()`` to continue to
+        the next middleware / handler, or ``abort(response)`` to short-circuit.
+
+        Example::
+
+            fn log_requests(ctx: Context) raises -> MiddlewareResult:
+                print(ctx.method(), ctx.path())
+                return next()
+
+            fn require_token(ctx: Context) raises -> MiddlewareResult:
+                if not ctx.header("X-API-Key"):
+                    return abort(Response.unauthorized("missing X-API-Key"))
+                return next()
+
+            app.use(log_requests)
+            app.use(require_token)
+        """
+        self.router.use(middleware)
+
+    # ------------------------------------------ lifecycle
+
+    def on_startup(mut self, hook: StartupHook) -> None:
+        """Register a function to run once before the server starts.
+
+        Useful for opening database connections, loading config, etc.
+
+        Example::
+
+            fn connect_db() raises:
+                print("DB connected")
+
+            app.on_startup(connect_db)
+        """
+        self.startup_hooks.append(StartupHookEntry(hook))
+
+    def on_error(mut self, handler: ErrorHandler) -> None:
+        """Register a custom error handler for unhandled exceptions from handlers.
+
+        The default handler logs the error and returns 500 Internal Server Error.
+
+        Example::
+
+            fn my_errors(ctx: Context, e: Error) raises -> HTTPResponse:
+                print("Oops:", String(e))
+                return Response.internal_error(String(e))
+
+            app.on_error(my_errors)
+        """
+        self.router.error_handler = handler
+
     # ------------------------------------------ start server
 
     def run(mut self, host: String = "0.0.0.0", port: Int = 8080) raises:
         """Start the HTTP server.
 
+        Runs all startup hooks, then begins listening for connections.
+
         Args:
             host: Bind address (default ``0.0.0.0``).
             port: TCP port (default ``8080``).
         """
+        for i in range(len(self.startup_hooks)):
+            self.startup_hooks[i].hook()
         var server = Server()
         server.listen_and_serve(String(host, ":", port), self.router)
 
     def start_server(mut self, address: String = "0.0.0.0:8080") raises:
-        """Start the HTTP server.
-
-        Deprecated: use ``run(host, port)`` instead.
-        """
+        """Deprecated: use ``run(host, port)`` instead."""
         var server = Server()
         server.listen_and_serve(address, self.router)

lightbug_api/context.mojo

@@ -145,6 +145,56 @@ struct Context(Copyable):
             return result.value()
         return default
 
+    # ------------------------------------------------------- typed path params
+
+    fn path_int(self, name: String) -> Optional[Int]:
+        """Parse a path parameter as ``Int``.
+
+        Returns ``None`` if the parameter is absent or not a valid integer.
+        """
+        var s = self.path_param(name)
+        if s:
+            try:
+                return Optional(atol(s.value()))
+            except:
+                pass
+        return Optional[Int]()
+
+    fn path_int(self, name: String, default: Int) -> Int:
+        """Parse a path parameter as ``Int``, falling back to *default*."""
+        var result = self.path_int(name)
+        if result:
+            return result.value()
+        return default
+
+    # ------------------------------------------------------ typed query params
+
+    fn query_int(self, name: String, default: Int = 0) -> Int:
+        """Parse a query parameter as ``Int``, falling back to *default*.
+
+        Example: ``ctx.query_int("page", 1)``
+        """
+        var s = self.query(name)
+        if s:
+            try:
+                return atol(s.value())
+            except:
+                pass
+        return default
+
+    fn query_bool(self, name: String, default: Bool = False) -> Bool:
+        """Parse a query parameter as ``Bool``, falling back to *default*.
+
+        Truthy string values: ``"true"``, ``"1"``, ``"yes"``.
+
+        Example: ``ctx.query_bool("verbose")``
+        """
+        var s = self.query(name)
+        if s:
+            var v = s.value()
+            return v == "true" or v == "1" or v == "yes"
+        return default
+
     # ---------------------------------------------------------------- headers
 
     fn header(self, name: String) -> Optional[String]:

lightbug_api/routing.mojo

@@ -3,6 +3,7 @@ from std.utils import Variant
 
 from lightbug_http import HTTPRequest, HTTPResponse, HTTPService, NotFound, OK
 from lightbug_http.http import RequestMethod
+from lightbug_http.http.common_response import InternalError
 from lightbug_http.uri import URIDelimiters
 
 from lightbug_api.context import Context
@@ -30,6 +31,51 @@ comptime HandlerResponse = Variant[HTTPResponse, String]
 # Use Context to access the request, path params, query params, headers, body.
 comptime Handler = fn (Context) raises -> HandlerResponse
 
+# ------------------------------------------------------------ middleware types
+
+# Middleware return type:
+#   HTTPResponse  — short-circuit: send this response immediately, skip handler
+#   Bool          — continue to the next middleware / handler (value is ignored)
+#
+# Use the helpers ``next()`` and ``abort(response)`` to return these cleanly.
+comptime MiddlewareResult = Variant[HTTPResponse, Bool]
+
+# Every middleware shares this non-capturing function-pointer signature.
+comptime Middleware = fn (Context) raises -> MiddlewareResult
+
+
+fn next() -> MiddlewareResult:
+    """Signal that processing should continue to the next middleware / handler."""
+    return MiddlewareResult(True)
+
+
+fn abort(var response: HTTPResponse) -> MiddlewareResult:
+    """Short-circuit the request with *response*, skipping all further processing."""
+    return MiddlewareResult(response^)
+
+
+# Wrapper so Middleware function pointers can live in a List[MiddlewareEntry].
+struct MiddlewareEntry(Copyable):
+    var handler: Middleware
+
+    fn __init__(out self, handler: Middleware):
+        self.handler = handler
+
+    fn __init__(out self, *, copy: Self):
+        self.handler = copy.handler
+
+
+# ------------------------------------------------------------ error handler
+
+# Called when a route handler raises an unhandled error.
+# Return an appropriate HTTPResponse; the exception is consumed.
+comptime ErrorHandler = fn (Context, Error) raises -> HTTPResponse
+
+
+fn _default_error_handler(ctx: Context, e: Error) raises -> HTTPResponse:
+    print("lightbug_api error:", String(e))
+    return InternalError()
+
 
 # --------------------------------------------------------------- path matching
 
@@ -229,6 +275,8 @@ struct RouterBase[is_main_app: Bool = False](HTTPService, Copyable):
     var path_fragment: String
     var sub_routers: List[RouterBase[False]]
     var routes: List[RouteEntry]
+    var middleware: List[MiddlewareEntry]
+    var error_handler: ErrorHandler
 
     # ------------------------------------------------------------------ init
 
@@ -238,11 +286,15 @@ struct RouterBase[is_main_app: Bool = False](HTTPService, Copyable):
         self.path_fragment = "/"
         self.sub_routers = List[RouterBase[False]]()
         self.routes = List[RouteEntry]()
+        self.middleware = List[MiddlewareEntry]()
+        self.error_handler = _default_error_handler
 
     def __init__(out self: Self, path_fragment: String) raises:
         self.path_fragment = path_fragment
         self.sub_routers = List[RouterBase[False]]()
         self.routes = List[RouteEntry]()
+        self.middleware = List[MiddlewareEntry]()
+        self.error_handler = _default_error_handler
 
         if not self._validate_path_fragment(path_fragment):
             raise Error(RouterErrors.INVALID_PATH_FRAGMENT_ERROR)
@@ -251,6 +303,8 @@ struct RouterBase[is_main_app: Bool = False](HTTPService, Copyable):
         self.path_fragment = copy.path_fragment
         self.sub_routers = copy.sub_routers.copy()
         self.routes = copy.routes.copy()
+        self.middleware = copy.middleware.copy()
+        self.error_handler = copy.error_handler
 
     # -------------------------------------------------------- route registration
 
@@ -296,6 +350,23 @@ struct RouterBase[is_main_app: Bool = False](HTTPService, Copyable):
         """Mount a sub-router, nesting all its routes under its path fragment."""
         self.sub_routers.append(router^)
 
+    def use(mut self, middleware: Middleware) -> None:
+        """Add a middleware function that runs before every handler on this router.
+
+        Middleware runs in registration order. Return ``next()`` to continue,
+        or ``abort(response)`` to short-circuit.
+
+        Example::
+
+            fn require_auth(ctx: Context) raises -> MiddlewareResult:
+                if not ctx.header("Authorization"):
+                    return abort(Response.unauthorized())
+                return next()
+
+            app.use(require_auth)
+        """
+        self.middleware.append(MiddlewareEntry(middleware))
+
     # ---------------------------------------------------------------- dispatch
 
     def func(mut self, req: HTTPRequest) raises -> HTTPResponse:
@@ -310,7 +381,21 @@ struct RouterBase[is_main_app: Bool = False](HTTPService, Copyable):
             raise e^
 
         var ctx = Context(req.copy(), route_match.path_params.copy())
-        var res = route_match.handler(ctx)
+
+        # Run middleware chain — any middleware may short-circuit with a response.
+        for i in range(len(self.middleware)):
+            var mw_result = self.middleware[i].handler(ctx)
+            if mw_result.isa[HTTPResponse]():
+                return mw_result.unsafe_take[HTTPResponse]()
+            # else Bool → continue to next middleware / handler
+
+        # Dispatch to the matched handler; convert unhandled errors to responses.
+        var res: HandlerResponse
+        try:
+            res = route_match.handler(ctx)
+        except e:
+            return self.error_handler(ctx, e)
+
         return self._encode_response(res^)
 
     # --------------------------------------------------------------- internals