LISTEN/NOTIFY funcionality

Signed-off-by: chandr-andr (Kiselev Aleksandr) <chandr@chandr.net>

Author: chandr-andr

python/psqlpy/__init__.py

@@ -7,7 +7,7 @@
     IsolationLevel,
     KeepaliveConfig,
     Listener,
-    ListenerNotification,
+    ListenerNotificationMsg,
     LoadBalanceHosts,
     QueryResult,
     ReadVariant,
@@ -28,7 +28,7 @@
     "IsolationLevel",
     "KeepaliveConfig",
     "Listener",
-    "ListenerNotification",
+    "ListenerNotificationMsg",
     "LoadBalanceHosts",
     "QueryResult",
     "ReadVariant",

python/psqlpy/_internal/__init__.pyi

@@ -2,7 +2,7 @@ import types
 from enum import Enum
 from io import BytesIO
 from ipaddress import IPv4Address, IPv6Address
-from typing import Any, Callable, Sequence, TypeVar
+from typing import Any, Awaitable, Callable, Sequence, TypeVar
 
 from typing_extensions import Buffer, Self
 
@@ -1360,6 +1360,9 @@ class ConnectionPool:
                 res = await connection.execute(...)
         ```
         """
+    def listener(self: Self) -> Listener:
+        """Create new listener."""
+
     def close(self: Self) -> None:
         """Close the connection pool."""
 
@@ -1752,6 +1755,70 @@ class ConnectionPoolBuilder:
 class Listener:
     """Result."""
 
+    connection: Connection
 
-class ListenerNotification:
-    """Result."""
+    def __aiter__(self: Self) -> Self: ...
+    async def __anext__(self: Self) -> ListenerNotificationMsg: ...
+    async def __aenter__(self: Self) -> Self: ...
+    async def __aexit__(
+        self: Self,
+        exception_type: type[BaseException] | None,
+        exception: BaseException | None,
+        traceback: types.TracebackType | None,
+    ) -> None: ...
+    async def startup(self: Self) -> None:
+        """Startup the listener.
+
+        Each listener MUST be started up.
+        """
+    async def add_callback(
+        self: Self,
+        channel: str,
+        callback: Callable[
+            [str, str, int, Connection], Awaitable[None],
+        ],
+    ) -> None:
+        """Add callback to the channel.
+        
+        Callback must be async function and have signature like this:
+        ```python
+        async def callback(
+            channel: str,
+            payload: str,
+            process_id: str,
+            connection: Connection,
+        ) -> None:
+            ...
+        ```
+        """
+    
+    async def clear_channel_callbacks(self, channel: str) -> None:
+        """Remove all callbacks for the channel.
+        
+        ### Parameters:
+        - `channel`: name of the channel.
+        """
+
+    async def listen(self: Self) -> None:
+        """Start listening.
+
+        Start actual listening.  
+        In the background it creates task in Rust event loop.  
+        You must save returned Future to the array.
+        """
+
+    async def abort_listen(self: Self) -> None:
+        """Abort listen.
+        
+        If `listen()` method was called, stop listening,  
+        else don't do anything.
+        """
+
+
+class ListenerNotificationMsg:
+    """Listener message in async iterator."""
+
+    process_id: int
+    channel: str
+    payload: str
+    connection: Connection

src/driver/connection_pool.rs

@@ -500,7 +500,7 @@ impl ConnectionPool {
         Connection::new(None, Some(self.pool.clone()))
     }
 
-    pub async fn add_listener(
+    pub fn listener(
         self_: pyo3::Py<Self>,
     ) -> RustPSQLDriverPyResult<Listener> {
         let (pg_config, ca_file, ssl_mode) = pyo3::Python::with_gil(|gil| {

src/driver/listener.rs

@@ -55,11 +55,10 @@ impl ChannelCallbacks {
 
 
 #[derive(Clone, Debug)]
-#[pyclass]
 pub struct ListenerNotification {
-    process_id: i32,
-    channel: String,
-    payload: String,
+    pub process_id: i32,
+    pub channel: String,
+    pub payload: String,
 }
 
 impl From::<Notification> for ListenerNotification {
@@ -73,13 +72,47 @@ impl From::<Notification> for ListenerNotification {
 }
 
 #[pyclass]
-struct ListenerNotificationMsg {
+pub struct ListenerNotificationMsg {
     process_id: i32,
     channel: String,
     payload: String,
     connection: Connection,
 }
 
+#[pymethods]
+impl ListenerNotificationMsg {
+    #[getter]
+    fn process_id(&self) -> i32 {
+        self.process_id
+    }
+    
+    #[getter]
+    fn channel(&self) -> String {
+        self.channel.clone()
+    }
+    
+    #[getter]
+    fn payload(&self) -> String {
+        self.payload.clone()
+    }
+    
+    #[getter]
+    fn connection(&self) -> Connection {
+        self.connection.clone()
+    }
+}
+
+impl ListenerNotificationMsg {
+    fn new(value: ListenerNotification, conn: Connection) -> Self {
+        ListenerNotificationMsg {
+            process_id: value.process_id,
+            channel: String::from(value.channel),
+            payload: String::from(value.payload),
+            connection: conn,
+        }
+    }
+}
+
 struct ListenerCallback {
     task_locals: Option<TaskLocals>,
     callback: Py<PyAny>,
@@ -111,7 +144,15 @@ impl ListenerCallback {
         if let Some(task_locals) = task_locals {
             tokio_runtime().spawn(pyo3_async_runtimes::tokio::scope(task_locals, async move {
                 let future = Python::with_gil(|py| {
-                    let awaitable = callback.call1(py, (lister_notification, connection)).unwrap();
+                    let awaitable = callback.call1(
+                        py,
+                        (
+                            lister_notification.channel,
+                            lister_notification.payload,
+                            lister_notification.process_id,
+                            connection,
+                        )
+                    ).unwrap();
                     pyo3_async_runtimes::tokio::into_future(awaitable.into_bound(py)).unwrap()
                 });
                 future.await.unwrap();
@@ -226,6 +267,41 @@ impl Listener {
         return Err(RustPSQLDriverError::ListenerClosedError)
     }
 
+    fn __anext__(&self) -> RustPSQLDriverPyResult<Option<Py<PyAny>>> {
+        let Some(client) = self.connection.db_client() else {
+            return Err(RustPSQLDriverError::ListenerStartError(
+                "Listener doesn't have underlying client, please call startup".into(),
+            ));
+        };
+        let Some(receiver) = self.receiver.clone() else {
+            return Err(RustPSQLDriverError::ListenerStartError(
+                "Listener doesn't have underlying receiver, please call startup".into(),
+            ));
+        };
+
+        let is_listened_clone = self.is_listened.clone();
+        let listen_query_clone = self.listen_query.clone();
+        let connection = self.connection.clone();
+
+        let py_future = Python::with_gil(move |gil| {
+            rustdriver_future(gil, async move {
+                {
+                    call_listen(&is_listened_clone, &listen_query_clone, &client).await?;
+                };
+                let next_element = {
+                    let mut write_receiver = receiver.write().await;
+                    write_receiver.next().await
+                };
+
+                let inner_notification = process_message(next_element)?;
+
+                Ok(ListenerNotificationMsg::new(inner_notification, connection))
+            })
+        });
+
+        Ok(Some(py_future?))
+    }
+
     #[getter]
     fn connection(&self) -> Connection {
         self.connection.clone()
@@ -280,40 +356,6 @@ impl Listener {
         Ok(())
     }
 
-    fn __anext__(&self) -> RustPSQLDriverPyResult<Option<PyObject>> {
-        let Some(client) = self.connection.db_client() else {
-            return Err(RustPSQLDriverError::ListenerStartError(
-                "Listener doesn't have underlying client, please call startup".into(),
-            ));
-        };
-        let Some(receiver) = self.receiver.clone() else {
-            return Err(RustPSQLDriverError::ListenerStartError(
-                "Listener doesn't have underlying receiver, please call startup".into(),
-            ));
-        };
-
-        let is_listened_clone = self.is_listened.clone();
-        let listen_query_clone = self.listen_query.clone();
-
-        let py_future = Python::with_gil(move |gil| {
-            rustdriver_future(gil, async move {
-                {
-                    call_listen(&is_listened_clone, &listen_query_clone, &client).await?;
-                };
-                let next_element = {
-                    let mut write_receiver = receiver.write().await;
-                    write_receiver.next().await
-                };
-
-                let inner_notification = process_message(next_element)?;
-
-                Ok(inner_notification)
-            })
-        });
-
-        Ok(Some(py_future?))
-    }
-
     #[pyo3(signature = (channel, callback))]
     async fn add_callback(
         &mut self,
@@ -337,10 +379,6 @@ impl Listener {
             callback,
         );
 
-        // let awaitable = callback.call1(()).unwrap();
-        // println!("8888888 {:?}", awaitable);
-        // let bbb = pyo3_async_runtimes::tokio::into_future(awaitable).unwrap();
-        // println!("999999");
         {
             let mut write_channel_callbacks = self.channel_callbacks.write().await;
             write_channel_callbacks.add_callback(channel, listener_callback);
@@ -351,6 +389,15 @@ impl Listener {
         Ok(())
     }
 
+    async fn clear_channel_callbacks(&mut self, channel: String) {
+        {
+            let mut write_channel_callbacks = self.channel_callbacks.write().await;
+            write_channel_callbacks.clear_channel_callbacks(channel);
+        }
+
+        self.update_listen_query().await;
+    }
+
     async fn listen(&mut self) -> RustPSQLDriverPyResult<()> {
         let Some(client) = self.connection.db_client() else {
             return Err(RustPSQLDriverError::BaseConnectionError("test".into()));
@@ -401,6 +448,14 @@ impl Listener {
 
         Ok(())
     }
+
+    async fn abort_listen(&mut self) {
+        if let Some(listen_abort_handler) = &self.listen_abort_handler {
+            listen_abort_handler.abort();
+        }
+
+        self.listen_abort_handler = None;
+    }
 }
 
 async fn dispatch_callback(

src/lib.rs

@@ -30,7 +30,7 @@ fn psqlpy(py: Python<'_>, pymod: &Bound<'_, PyModule>) -> PyResult<()> {
     pymod.add_class::<driver::transaction::Transaction>()?;
     pymod.add_class::<driver::cursor::Cursor>()?;
     pymod.add_class::<driver::listener::Listener>()?;
-    pymod.add_class::<driver::listener::ListenerNotification>()?;
+    pymod.add_class::<driver::listener::ListenerNotificationMsg>()?;
     pymod.add_class::<driver::transaction_options::IsolationLevel>()?;
     pymod.add_class::<driver::transaction_options::SynchronousCommit>()?;
     pymod.add_class::<driver::transaction_options::ReadVariant>()?;