feat(python-bindings): add session deletion API (#176)

# Description

Expose session deletion API to python bindings, and add tests.

Signed-off-by: Mauro Sardara <[email protected]>

Author: msardara

data-plane/gateway/service/src/lib.rs

@@ -687,7 +687,7 @@ impl Service {
             true => Ok(()),
             false => {
                 error!("error deleting session");
-                Err(ServiceError::SessionError("unknown".to_string()))
+                Err(ServiceError::SessionError("session not found".to_string()))
             }
         }
     }

data-plane/python-bindings/agp_bindings/__init__.py

@@ -5,7 +5,7 @@
 from typing import Optional
 
 from ._agp_bindings import (
-        __version__,
+    __version__,
     build_profile,
     build_info,
     SESSION_UNSPECIFIED,
@@ -21,6 +21,7 @@
     create_pyservice,
     create_rr_session,
     create_streaming_session,
+    delete_session,
     disconnect,
     init_tracing as init_tracing,
     publish,
@@ -229,6 +230,27 @@ async def create_streaming_session(
         self.sessions[session.id] = (session, asyncio.Queue(queue_size))
         return session
 
+    async def delete_session(self, session_id: int):
+        """
+        Delete a session.
+
+        Args:
+            session_id (int): The ID of the session to delete.
+
+        Returns:
+            None
+        """
+
+        # Check if the session ID is in the sessions map
+        if session_id not in self.sessions:
+            raise Exception("session not found", session_id)
+
+        # Remove the session from the map
+        del self.sessions[session_id]
+
+        # Remove the session from the gateway
+        await delete_session(self.svc, session_id)
+
     async def run_server(self, config: dict):
         """
         Start the server part of the Gateway service. The server will be started only
@@ -404,6 +426,10 @@ async def publish(
             None
         """
 
+        # Make sure the sessions exists
+        if session.id not in self.sessions:
+            raise Exception("session not found", session.id)
+
         dest = PyAgentType(organization, namespace, agent)
         await publish(self.svc, session, 1, msg, dest, agent_id)
 

data-plane/python-bindings/agp_bindings/_agp_bindings.pyi

@@ -72,6 +72,9 @@ def create_rr_session(svc:PyService, config:PyRequestResponseConfiguration=...)
 def create_streaming_session(svc:PyService, config:PyStreamingConfiguration) -> typing.Any:
     ...
 
+def delete_session(svc:PyService, session_id:builtins.int) -> typing.Any:
+    ...
+
 def disconnect(svc:PyService, conn:builtins.int) -> typing.Any:
     ...
 

data-plane/python-bindings/src/lib.rs

@@ -16,8 +16,8 @@ mod _agp_bindings {
     #[pymodule_export]
     use pyservice::{
         PyService, connect, create_ff_session, create_pyservice, create_rr_session,
-        create_streaming_session, disconnect, publish, receive, remove_route, run_server,
-        set_route, stop_server, subscribe, unsubscribe,
+        create_streaming_session, delete_session, disconnect, publish, receive, remove_route,
+        run_server, set_route, stop_server, subscribe, unsubscribe,
     };
 
     #[pymodule_export]

data-plane/python-bindings/src/pyservice.rs

@@ -95,6 +95,13 @@ impl PyService {
         ))
     }
 
+    async fn delete_session(&self, session_id: session::Id) -> Result<(), ServiceError> {
+        self.sdk
+            .service
+            .delete_session(&self.sdk.agent, session_id)
+            .await
+    }
+
     async fn run_server(&self, config: PyGrpcServerConfig) -> Result<(), ServiceError> {
         self.sdk.service.run_server(&config)
     }
@@ -250,7 +257,7 @@ pub fn create_ff_session(
             config.fire_and_forget_configuration,
         ))
         .await
-        .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+        .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }
 
@@ -267,7 +274,7 @@ pub fn create_rr_session(
             config.request_response_configuration,
         ))
         .await
-        .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+        .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }
 
@@ -284,7 +291,18 @@ pub fn create_streaming_session(
             config.streaming_configuration,
         ))
         .await
-        .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+        .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
+    })
+}
+
+#[gen_stub_pyfunction]
+#[pyfunction]
+#[pyo3(signature = (svc, session_id))]
+pub fn delete_session(py: Python, svc: PyService, session_id: u32) -> PyResult<Bound<PyAny>> {
+    pyo3_async_runtimes::tokio::future_into_py(py, async move {
+        svc.delete_session(session_id)
+            .await
+            .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }
 
@@ -299,7 +317,7 @@ pub fn run_server(py: Python, svc: PyService, config: Py<PyDict>) -> PyResult<Bo
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         svc.run_server(config)
             .await
-            .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+            .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }
 
@@ -313,7 +331,7 @@ pub fn stop_server(py: Python, svc: PyService, endpoint: String) -> PyResult<Bou
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         svc.stop_server(&endpoint)
             .await
-            .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+            .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }
 
@@ -329,7 +347,7 @@ pub fn connect(py: Python, svc: PyService, config: Py<PyDict>) -> PyResult<Bound
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         svc.connect(config)
             .await
-            .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+            .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }
 
@@ -339,7 +357,7 @@ pub fn disconnect(py: Python, svc: PyService, conn: u64) -> PyResult<Bound<PyAny
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         svc.disconnect(conn)
             .await
-            .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+            .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }
 
@@ -356,7 +374,7 @@ pub fn subscribe(
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         svc.subscribe(conn, name, id)
             .await
-            .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+            .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }
 
@@ -373,7 +391,7 @@ pub fn unsubscribe(
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         svc.unsubscribe(conn, name, id)
             .await
-            .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+            .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }
 
@@ -390,7 +408,7 @@ pub fn set_route(
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         svc.set_route(conn, name, id)
             .await
-            .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+            .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }
 
@@ -407,7 +425,7 @@ pub fn remove_route(
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         svc.remove_route(conn, name, id)
             .await
-            .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+            .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }
 
@@ -426,7 +444,7 @@ pub fn publish(
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         svc.publish(session_info.session_info, fanout, blob, name, id)
             .await
-            .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+            .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }
 
@@ -440,7 +458,7 @@ pub fn receive(py: Python, svc: PyService) -> PyResult<Bound<PyAny>> {
         async move {
             svc.receive()
                 .await
-                .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+                .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
         },
     )
 }
@@ -457,6 +475,6 @@ pub fn create_pyservice(
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
         PyService::create_pyservice(organization, namespace, agent_type, id)
             .await
-            .map_err(|e| PyErr::new::<PyException, _>(format!("{}", e.to_string())))
+            .map_err(|e| PyErr::new::<PyException, _>(e.to_string()))
     })
 }

data-plane/python-bindings/tests/test_bindings.py

@@ -63,12 +63,28 @@ async def test_end_to_end(server):
     # check if the message is correct
     assert msg_rcv == bytes(msg)
 
+    # delete sessions
+    await agp_bindings.delete_session(svc_alice, session_info.id)
+    await agp_bindings.delete_session(svc_bob, session_info.id)
+
+    # try to send a message after deleting the session - this should raise an exception
+    try:
+        await agp_bindings.publish(svc_alice, session_info, 1, msg, bob_class, None)
+    except Exception as e:
+        assert "session not found" in str(e), f"Unexpected error message: {str(e)}"
+
     # disconnect alice
     await agp_bindings.disconnect(svc_alice, conn_id_alice)
 
     # disconnect bob
     await agp_bindings.disconnect(svc_bob, conn_id_bob)
 
+    # try to delete a random session, we should get an exception
+    try:
+        await agp_bindings.delete_session(svc_alice, 123456789)
+    except Exception as e:
+        assert "session not found" in str(e)
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("server", ["127.0.0.1:12345"], indirect=True)
@@ -131,6 +147,21 @@ async def test_gateway_wrapper(server):
         # check if the message is correct
         assert msg_rcv == bytes(msg)
 
+    # delete sessions
+    await gateway1.delete_session(session_info.id)
+    await gateway2.delete_session(session_info.id)
+
+    # try to send a message after deleting the session - this should raise an exception
+    try:
+        await gateway1.publish(session_info, msg, org, ns, agent1)
+    except Exception as e:
+        assert "session not found" in str(e), f"Unexpected error message: {str(e)}"
+
+    # try to delete a random session, we should get an exception
+    try:
+        await gateway1.delete_session(123456789)
+    except Exception as e:
+        assert "session not found" in str(e), f"Unexpected error message: {str(e)}"
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("server", ["127.0.0.1:12346"], indirect=True)