Repository: sparckles/Robyn (language: python)

Issue:
Add query params in Websockets
<!-- 
Thank you for considering improving Robyn!

Please describe your idea in depth. If you're not sure what to write, imagine the following:
  - How is this important to you? How would you use it?
  - Can you think of any alternatives?
  - Do you have any ideas about how it can be implemented? Are you willing/able to implement it? Do you need mentoring?
-->


Discussion / hints:

Fix patch:
diff --git a/docs_src/src/pages/documentation/api_reference/websockets.mdx b/docs_src/src/pages/documentation/api_reference/websockets.mdx
index 370f69895..c111d154a 100644
--- a/docs_src/src/pages/documentation/api_reference/websockets.mdx
+++ b/docs_src/src/pages/documentation/api_reference/websockets.mdx
@@ -182,6 +182,37 @@ To handle real-time communication, Batman learned how to work with WebSockets. H
   </Col>
 </Row>
 
+<Row>
+  <Col>
+    Robyn also showed Batman to work with query params.
+  </Col>
+  <Col sticky>
+    <CodeGroup title="Request" tag="GET" label="/hello_world?name=gordon&desg=commissioner">
+
+      ```python {{ title: 'untyped' }}
+
+        @websocket.on("message")
+        async def message(ws, msg, global_dependencies) -> str:
+            websocket_id = ws.id
+            if (ws.query_params.get("name") == "gordon" and ws.query_params.get("desg") == "commissioner"):
+              ws.sync_broadcast("Gordon authorized to login!")
+            return ""
+
+      ```
+
+      ```python {{title: 'typed'}}
+
+        @websocket.on("message")
+        async def message(ws: WebSocketConnector, msg: str, global_dependencies) -> str:
+            websocket_id = ws.id
+            if (ws.query_params.get("name") == "gordon" and ws.query_params.get("desg") == "commissioner"):
+              ws.sync_broadcast("Gordon authorized to login!")
+            return ""
+
+      ```
+    </CodeGroup>
+  </Col>
+</Row>
 
 ---
 
diff --git a/robyn/robyn.pyi b/robyn/robyn.pyi
index 948671e67..4d4ff0329 100644
--- a/robyn/robyn.pyi
+++ b/robyn/robyn.pyi
@@ -347,6 +347,7 @@ class WebSocketConnector:
 
     Attributes:
         id (str): The id of the client
+        query_params (QueryParams): The query parameters object
 
         async_broadcast (Callable): The function to broadcast a message to all clients
         async_send_to (Callable): The function to send a message to the client
@@ -355,6 +356,7 @@ class WebSocketConnector:
     """
 
     id: str
+    query_params: QueryParams
 
     async def async_broadcast(self, message: str) -> None:
         pass
diff --git a/src/websockets/mod.rs b/src/websockets/mod.rs
index 483fd7db6..83ec47993 100644
--- a/src/websockets/mod.rs
+++ b/src/websockets/mod.rs
@@ -2,6 +2,7 @@ pub mod registry;
 
 use crate::executors::web_socket_executors::execute_ws_function;
 use crate::types::function_info::FunctionInfo;
+use crate::types::multimap::QueryParams;
 use registry::{SendMessageToAll, SendText};
 
 use actix::prelude::*;
@@ -26,6 +27,7 @@ pub struct WebSocketConnector {
     pub router: HashMap<String, FunctionInfo>,
     pub task_locals: TaskLocals,
     pub registry_addr: Addr<WebSocketRegistry>,
+    pub query_params: QueryParams,
 }
 
 // By default mailbox capacity is 16 messages.
@@ -173,6 +175,11 @@ impl WebSocketConnector {
     pub fn get_id(&self) -> String {
         self.id.to_string()
     }
+
+    #[getter]
+    pub fn get_query_params(&self) -> QueryParams {
+        self.query_params.clone()
+    }
 }
 
 static REGISTRY_ADDRESSES: OnceCell<RwLock<HashMap<String, Addr<WebSocketRegistry>>>> =
@@ -207,12 +214,26 @@ pub async fn start_web_socket(
 ) -> Result<HttpResponse, Error> {
     let registry_addr = get_or_init_registry_for_endpoint(endpoint);
 
+    let mut query_params = QueryParams::new();
+
+    if !req.query_string().is_empty() {
+        let split = req.query_string().split('&');
+        for s in split {
+            let path_params = s.split_once('=').unwrap_or((s, ""));
+            let key = path_params.0.to_string();
+            let value = path_params.1.to_string();
+
+            query_params.set(key, value);
+        }
+    }
+
     ws::start(
         WebSocketConnector {
             router,
             task_locals,
             id: Uuid::new_v4(),
             registry_addr,
+            query_params,
         },
         &req,
         stream,


Test patch:
diff --git a/integration_tests/base_routes.py b/integration_tests/base_routes.py
index 7ebb87b64..1513472fd 100644
--- a/integration_tests/base_routes.py
+++ b/integration_tests/base_routes.py
@@ -74,6 +74,8 @@ async def message(ws: WebSocketConnector, msg: str, global_dependencies) -> str:
     elif state == 1:
         resp = "Whooo??"
     elif state == 2:
+        await ws.async_broadcast(ws.query_params.get("one"))
+        ws.sync_send_to(websocket_id, ws.query_params.get("two"))
         resp = "*chika* *chika* Slim Shady."
     websocket_state[websocket_id] = (state + 1) % 3
     return resp
diff --git a/integration_tests/test_web_sockets.py b/integration_tests/test_web_sockets.py
index e1e1589d0..722590dd1 100644
--- a/integration_tests/test_web_sockets.py
+++ b/integration_tests/test_web_sockets.py
@@ -8,7 +8,7 @@
 
 @pytest.mark.benchmark
 def test_web_socket_raw_benchmark(session):
-    ws = create_connection(f"{BASE_URL}/web_socket")
+    ws = create_connection(f"{BASE_URL}/web_socket?one=hi&two=hello")
     assert ws.recv() == "Hello world, from ws"
 
     ws.send("My name is?")
@@ -20,6 +20,8 @@ def test_web_socket_raw_benchmark(session):
     assert ws.recv() == "Whooo??"
 
     ws.send("My name is?")
+    ws.recv() == "hi"
+    ws.recv() == "hello"
     assert ws.recv() == "*chika* *chika* Slim Shady."