fix: response headers and headers not working in const requests (#331)

* fix: allow response headers

* fix: rename get_headers to get_request_headers

* fix: header in const request bug and improve code architecture (#333)

* test: add integration tests to verify headers in non const requests

* docs: add documentation of per route headers

Co-authored-by: Antoine Romero-Romero <[email protected]>

docs/features.md

@@ -105,7 +105,22 @@ async def json(request):
 You can also add global headers for every request.
 
 ```python
-app.add_header("server", "robyn")
+app.add_request_header("server", "robyn")
+
+```
+
+## Per route headers
+You can also add headers for every route.
+
+```python
+@app.get("/request_headers")
+async def request_headers():
+    return {
+        "status_code": "200",
+        "body": "",
+        "type": "text",
+        "headers": jsonify({"Header": "header_value"}),
+    }
 
 ```
 

integration_tests/base_routes.py

@@ -63,6 +63,26 @@ async def const_request_json():
     return jsonify({"hello": "world"})
 
 
+@app.get("/const_request_headers", const=True)
+async def const_request_headers():
+    return {
+        "status_code": "200",
+        "body": "",
+        "type": "text",
+        "headers": jsonify({"Header": "header_value"}),
+    }
+
+
+@app.get("/request_headers")
+async def request_headers():
+    return {
+        "status_code": "200",
+        "body": "This is a regular response",
+        "type": "text",
+        "headers": jsonify({"Header": "header_value"}),
+    }
+
+
 @app.get("/404")
 def return_404():
     return {"status_code": "404", "body": "hello", "type": "text"}
@@ -213,7 +233,7 @@ async def redirect_route(request):
 
 
 if __name__ == "__main__":
-    app.add_header("server", "robyn")
+    app.add_request_header("server", "robyn")
     current_file_path = pathlib.Path(__file__).parent.resolve()
     app.add_directory(
         route="/test_dir",

integration_tests/test_get_requests.py

@@ -33,6 +33,14 @@ def test_queries(session):
     assert r.json() == {}
 
 
+def test_request_headers(session):
+    r = requests.get(f"{BASE_URL}/request_headers")
+    assert r.status_code == 200
+    assert r.text == "This is a regular response"
+    assert "Header" in r.headers
+    assert r.headers["Header"] == "header_value"
+
+
 def test_const_request(session):
     r = requests.get(f"{BASE_URL}/const_request")
     assert "Hello world" in r.text
@@ -43,3 +51,10 @@ def test_const_request_json(session):
     r = requests.get(f"{BASE_URL}/const_request_json")
     assert r.status_code == 200
     assert r.json() == {"hello": "world"}
+
+
+def test_const_request_headers(session):
+    r = requests.get(f"{BASE_URL}/const_request_headers")
+    assert r.status_code == 200
+    assert "Header" in r.headers
+    assert r.headers["Header"] == "header_value"

robyn/__init__.py

@@ -37,7 +37,7 @@ def __init__(self, file_object: str) -> None:
         self.router = Router()
         self.middleware_router = MiddlewareRouter()
         self.web_socket_router = WebSocketRouter()
-        self.headers = []
+        self.request_headers = []  # This needs a better type
         self.directories = []
         self.event_handlers = {}
         load_vars(project_root=directory_path)
@@ -83,8 +83,8 @@ def add_directory(
     ):
         self.directories.append((route, directory_path, index_file, show_files_listing))
 
-    def add_header(self, key: str, value: str) -> None:
-        self.headers.append((key, value))
+    def add_request_header(self, key: str, value: str) -> None:
+        self.request_headers.append((key, value))
 
     def add_web_socket(self, endpoint: str, ws: WS) -> None:
         self.web_socket_router.add_route(endpoint, ws)
@@ -140,7 +140,7 @@ def init_processpool(socket):
                     target=spawn_process,
                     args=(
                         self.directories,
-                        self.headers,
+                        self.request_headers,
                         self.router.get_routes(),
                         self.middleware_router.get_routes(),
                         self.web_socket_router.get_routes(),

robyn/processpool.py

@@ -30,7 +30,7 @@ def initialize_event_loop():
 
 def spawn_process(
     directories: Tuple[Directory, ...],
-    headers: Tuple[Header, ...],
+    request_headers: Tuple[Header, ...],
     routes: Tuple[Route, ...],
     middlewares: Tuple[MiddlewareRoute, ...],
     web_sockets: Dict[str, WS],
@@ -63,8 +63,8 @@ def spawn_process(
         route, directory_path, index_file, show_files_listing = directory
         server.add_directory(route, directory_path, index_file, show_files_listing)
 
-    for key, val in headers:
-        server.add_header(key, val)
+    for key, val in request_headers:
+        server.add_request_header(key, val)
 
     for route in routes:
         route_type, endpoint, function, is_const = route

robyn/robyn.pyi

@@ -26,7 +26,7 @@ class Server:
         show_files_listing: bool,
     ) -> None:
         pass
-    def add_header(self, key: str, value: str) -> None:
+    def add_request_header(self, key: str, value: str) -> None:
         pass
     def add_route(
         self,

robyn/router.py

@@ -5,6 +5,7 @@
 from typing import Callable, Dict, List, Tuple, Union
 from types import CoroutineType
 from robyn.robyn import FunctionInfo
+from robyn.responses import jsonify
 
 from robyn.ws import WS
 
@@ -36,7 +37,7 @@ def _format_response(self, res):
 
                 response = {"status_code": "200", "body": res["body"], **res}
         else:
-            response = {"status_code": "200", "body": res, "type": "text"}
+            response = {"status_code": "200", "body": res, "type": "text", "headers": jsonify({"Content-Type": "text/plain"})}
 
         return response
 

src/executors/mod.rs

@@ -1,7 +1,7 @@
 /// This is the module that has all the executor functions
 /// i.e. the functions that have the responsibility of parsing and executing functions.
 use crate::io_helpers::read_file;
-use crate::types::Request;
+use crate::types::{FunctionInfo, Request, Response};
 
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -10,7 +10,6 @@ use anyhow::{Context, Result};
 use log::debug;
 use pyo3_asyncio::TaskLocals;
 // pyO3 module
-use crate::types::FunctionInfo;
 use pyo3::prelude::*;
 
 fn get_function_output<'a>(
@@ -56,14 +55,8 @@ pub async fn execute_middleware_function<'a>(
     }
 }
 
-// Change this!
 #[inline]
-pub async fn execute_http_function(
-    request: &Request,
-    function: FunctionInfo,
-    // need to change this to return a response struct
-    // create a custom struct for this
-) -> Result<HashMap<String, String>> {
+pub async fn execute_http_function(request: &Request, function: FunctionInfo) -> Result<Response> {
     if function.is_async {
         let output = Python::with_gil(|py| {
             let function_output = get_function_output(&function, py, request);
@@ -79,10 +72,11 @@ pub async fn execute_http_function(
             let contents = read_file(file_path).unwrap();
             res.insert("body".to_owned(), contents);
         }
-        Ok(res)
+        Response::from_hashmap(res)
     } else {
-        Python::with_gil(|py| get_function_output(&function, py, request)?.extract())
-            .context("Failed to execute handler function")
+        Response::from_hashmap(Python::with_gil(|py| {
+            get_function_output(&function, py, request)?.extract()
+        })?)
     }
 }
 

src/io_helpers/mod.rs

@@ -4,16 +4,30 @@ use std::io::Read;
 
 use actix_web::HttpResponseBuilder;
 use anyhow::Result;
+use dashmap::DashMap;
 
 // this should be something else
 // probably inside the submodule of the http router
 #[inline]
-pub fn apply_headers(response: &mut HttpResponseBuilder, headers: &HashMap<String, String>) {
-    for (key, val) in (headers).iter() {
+pub fn apply_hashmap_headers(
+    response: &mut HttpResponseBuilder,
+    headers: &HashMap<String, String>,
+) {
+    for (key, val) in headers.iter() {
         response.insert_header((key.clone(), val.clone()));
     }
 }
 
+#[inline]
+pub fn apply_dashmap_headers(
+    response: &mut HttpResponseBuilder,
+    headers: &DashMap<String, String>,
+) {
+    for h in headers.iter() {
+        response.insert_header((h.key().clone(), h.value().clone()));
+    }
+}
+
 /// A function to read lossy files and serve it as a html response
 ///
 /// # Arguments

src/lib.rs

@@ -1,6 +1,5 @@
 mod executors;
 mod io_helpers;
-mod request_handler;
 mod routers;
 mod server;
 mod shared_socket;

src/routers/const_router.rs

@@ -4,6 +4,7 @@ use std::sync::Arc;
 use std::sync::RwLock;
 
 use crate::executors::execute_http_function;
+use crate::types::Response;
 use crate::types::{FunctionInfo, Request};
 use anyhow::Context;
 use log::debug;
@@ -16,14 +17,14 @@ use anyhow::{Error, Result};
 
 use super::Router;
 
-type RouteMap = RwLock<MatchItRouter<String>>;
+type RouteMap = RwLock<MatchItRouter<Response>>;
 
 /// Contains the thread safe hashmaps of different routes
 pub struct ConstRouter {
     routes: HashMap<Method, Arc<RouteMap>>,
 }
 
-impl Router<String, Method> for ConstRouter {
+impl Router<Response, Method> for ConstRouter {
     /// Doesn't allow query params/body/etc as variables cannot be "memoized"/"const"ified
     fn add_route(
         &self,
@@ -46,18 +47,14 @@ impl Router<String, Method> for ConstRouter {
                 .await
                 .unwrap();
             debug!("This is the result of the output {:?}", output);
-            table
-                .write()
-                .unwrap()
-                .insert(route, output.get("body").unwrap().to_string())
-                .unwrap();
+            table.write().unwrap().insert(route, output).unwrap();
             Ok(())
         })?;
 
         Ok(())
     }
 
-    fn get_route(&self, route_method: &Method, route: &str) -> Option<String> {
+    fn get_route(&self, route_method: &Method, route: &str) -> Option<Response> {
         let table = self.routes.get(route_method)?;
         let route_map = table.read().ok()?;