wasi_pg_client/

notification.rs

//! PostgreSQL LISTEN/NOTIFY support.
//!
//! PostgreSQL provides an asynchronous notification system where clients can
//! subscribe to named channels (`LISTEN`) and other clients can send
//! notifications to those channels (`NOTIFY`). Notifications arrive as
//! `NotificationResponse` backend messages, which can appear at any time
//! between other messages.
//!
//! # Example
//! ```ignore
//! // Connection A: listen
//! conn_a.listen("my_channel").await?;
//!
//! // Connection B: notify
//! conn_b.notify("my_channel", "hello!").await?;
//!
//! // Connection A: receive
//! conn_a.wait_for_notification(None).await?;
//! let notifications = conn_a.notifications();
//! ```

use std::time::Duration;

use crate::protocol::{BackendMessage, TransactionStatus};

use crate::connection::{Connection, ConnectionState};
use crate::error::{PgError, Result};

#[cfg(feature = "tracing")]
use crate::tracing_ext::TARGET_NOTIFICATION;

// ---------------------------------------------------------------------------
// Notification (re-exported from connection, but documented here)
// ---------------------------------------------------------------------------

/// An asynchronous notification received from PostgreSQL.
///
/// Notifications are delivered via the `LISTEN`/`NOTIFY` protocol. A
/// notification includes the process ID of the notifying backend, the
/// channel name, and an optional payload string.
///
/// Notifications can arrive at any time — they are interleaved with other
/// backend messages. The connection buffers them in an internal queue.
/// Use [`Connection::notifications`] or [`Connection::drain_notifications`]
/// to retrieve buffered notifications, or [`Connection::wait_for_notification`]
/// to block until one arrives.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct Notification {
    /// Backend process ID that sent the notification.
    pub process_id: i32,
    /// Channel name.
    pub channel: String,
    /// Payload string (empty string if no payload was sent).
    pub payload: String,
}

// ---------------------------------------------------------------------------
// Connection methods
// ---------------------------------------------------------------------------

impl Connection {
    /// Start listening for notifications on a channel.
    ///
    /// This executes `LISTEN <channel>` on the server. After this call,
    /// any `NOTIFY` on the same channel (from any connection) will cause
    /// a `NotificationResponse` message to be sent to this connection.
    ///
    /// # Example
    /// ```ignore
    /// conn.listen("events").await?;
    /// ```
    #[must_use = "listen errors should be checked"]
    pub async fn listen(&mut self, channel: &str) -> Result<()> {
        let sql = Self::build_listen_sql(channel);
        self.execute(&sql).await?;
        self.session_state.track_listen(channel);
        #[cfg(feature = "tracing")]
        tracing::info!(target: TARGET_NOTIFICATION, channel = %channel, "LISTEN: subscribed to channel");
        Ok(())
    }

    /// Stop listening on a channel.
    ///
    /// This executes `UNLISTEN <channel>` on the server.
    #[must_use = "unlisten errors should be checked"]
    pub async fn unlisten(&mut self, channel: &str) -> Result<()> {
        let sql = Self::build_unlisten_sql(channel);
        self.execute(&sql).await?;
        self.session_state.track_unlisten(channel);
        Ok(())
    }

    /// Stop listening on all channels.
    ///
    /// This executes `UNLISTEN *` on the server.
    #[must_use = "unlisten errors should be checked"]
    pub async fn unlisten_all(&mut self) -> Result<()> {
        self.execute("UNLISTEN *").await?;
        self.session_state.clear_listen_channels();
        Ok(())
    }

    /// Send a notification on a channel.
    ///
    /// This uses `pg_notify(channel, payload)` which properly handles
    /// identifier quoting and payload escaping.
    ///
    /// # Example
    /// ```ignore
    /// conn.notify("events", "user_logged_in").await?;
    /// ```
    #[must_use = "notify errors should be checked"]
    pub async fn notify(&mut self, channel: &str, payload: &str) -> Result<()> {
        #[cfg(feature = "tracing")]
        tracing::debug!(target: TARGET_NOTIFICATION, channel = %channel, payload_len = payload.len(), "NOTIFY: sending notification");
        self.execute_params("SELECT pg_notify($1, $2)", &[&channel, &payload])
            .await?;
        Ok(())
    }

    /// Take all buffered notifications from the internal queue.
    ///
    /// Notifications can arrive at any time during other operations. They
    /// are buffered in an internal queue. This method drains the queue
    /// and returns all notifications that have arrived since the last call.
    ///
    /// This is a synchronous operation — no I/O is performed.
    pub fn notifications(&mut self) -> Vec<Notification> {
        self.notification_queue.drain(..).collect()
    }

    #[allow(dead_code)]
    async fn read_next_notification_blocking(&mut self) -> Result<Notification> {
        if !self.is_idle() {
            return Err(PgError::InvalidState(
                "connection must be idle while waiting for notifications".into(),
            ));
        }

        loop {
            let msg = self.codec.read_message(&mut self.transport).await?;
            match msg {
                BackendMessage::NotificationResponse(body) => {
                    return Ok(Notification {
                        process_id: body.process_id(),
                        channel: body.channel().unwrap_or("").to_string(),
                        payload: body.message().unwrap_or("").to_string(),
                    });
                }
                BackendMessage::NoticeResponse(body) => {
                    if let Ok(notice) = crate::query::Notice::from_fields(&body) {
                        self.handle_notice(&notice);
                    }
                }
                BackendMessage::ParameterStatus(body) => {
                    if let (Ok(name), Ok(value)) = (body.name(), body.value()) {
                        self.server_params
                            .params
                            .insert(name.to_string(), value.to_string());
                    }
                }
                BackendMessage::ReadyForQuery(body) => {
                    self.transaction_status = TransactionStatus::from_u8(body.status())
                        .unwrap_or(TransactionStatus::Idle);
                    self.state = ConnectionState::Idle;
                }
                BackendMessage::EmptyQueryResponse => {}
                _ => {
                    return Err(PgError::InvalidState(
                        "received unexpected backend message while waiting for notification".into(),
                    ));
                }
            }
        }
    }

    /// Wait for the next notification to arrive.
    ///
    /// If a notification is already buffered, it is returned immediately.
    /// Otherwise, this method blocks (async) until a notification arrives
    /// or the optional timeout expires.
    ///
    /// On native tokio builds, timeout semantics are real: the read wait is
    /// raced against the requested deadline. On other targets the timeout path
    /// remains best-effort and uses an empty query cycle to flush pending
    /// notifications.
    ///
    /// # Example
    /// ```ignore
    /// // Wait up to 5 seconds for a notification
    /// if let Some(notification) = conn.wait_for_notification(Some(Duration::from_secs(5))).await? {
    ///     println!("Got notification on {}: {}", notification.channel, notification.payload);
    /// }
    /// ```
    #[must_use = "notification errors should be checked"]
    pub async fn wait_for_notification(
        &mut self,
        timeout: Option<Duration>,
    ) -> Result<Option<Notification>> {
        // Check queue first
        if let Some(n) = self.notification_queue.pop_front() {
            return Ok(Some(n));
        }

        // If timeout is Some(0), return immediately without sending a query
        if let Some(d) = timeout {
            if d.is_zero() {
                return Ok(None);
            }
        }

        #[cfg(all(not(target_arch = "wasm32"), feature = "tokio-transport"))]
        {
            match timeout {
                Some(deadline) => {
                    match tokio::time::timeout(deadline, self.read_next_notification_blocking())
                        .await
                    {
                        Ok(result) => result.map(Some),
                        Err(_) => Ok(None),
                    }
                }
                None => self.read_next_notification_blocking().await.map(Some),
            }
        }

        #[cfg(not(all(not(target_arch = "wasm32"), feature = "tokio-transport")))]
        {
            // Best-effort fallback for targets without a runtime timeout race.
            let _ = timeout;

            // Send an empty query to trigger a ReadyForQuery cycle.
            // The server will deliver any pending notifications before
            // sending ReadyForQuery.
            self.transition(ConnectionState::ActiveSimpleQuery)?;

            self.codec
                .send(
                    &mut self.transport,
                    &crate::protocol::FrontendMessage::Query { sql: String::new() },
                )
                .await
                .map_err(crate::error::Error::from)?;

            // Read messages, collecting notifications
            loop {
                let msg = self.codec.read_message(&mut self.transport).await?;
                match msg {
                    BackendMessage::NotificationResponse(body) => {
                        let notification = Notification {
                            process_id: body.process_id(),
                            channel: body.channel().unwrap_or("").to_string(),
                            payload: body.message().unwrap_or("").to_string(),
                        };

                        // Continue reading until ReadyForQuery to drain the cycle
                        self.read_until_ready().await?;

                        return Ok(Some(notification));
                    }
                    BackendMessage::EmptyQueryResponse => {}
                    BackendMessage::ReadyForQuery(body) => {
                        self.transaction_status = TransactionStatus::from_u8(body.status())
                            .unwrap_or(TransactionStatus::Idle);
                        self.state = ConnectionState::Idle;
                        break;
                    }
                    BackendMessage::NoticeResponse(body) => {
                        if let Ok(notice) = crate::query::Notice::from_fields(&body) {
                            self.handle_notice(&notice);
                        }
                    }
                    BackendMessage::ParameterStatus(body) => {
                        if let (Ok(name), Ok(value)) = (body.name(), body.value()) {
                            self.server_params
                                .params
                                .insert(name.to_string(), value.to_string());
                        }
                    }
                    _ => {}
                }
            }

            // No notification arrived during this cycle
            Ok(self.notification_queue.pop_front())
        }
    }

    /// Wait for a notification with a timeout.
    ///
    /// If the notification queue already contains a notification, returns it
    /// immediately. Otherwise, waits up to `timeout` for a new notification.
    /// Returns `Ok(None)` if no notification arrives within the timeout.
    ///
    /// # Example
    /// ```ignore
    /// if let Some(n) = conn.wait_for_notification_with_timeout(Duration::from_secs(5)).await? {
    ///     println!("Got notification on {}: {}", n.channel, n.payload);
    /// }
    /// ```
    #[must_use = "notification errors should be checked"]
    pub async fn wait_for_notification_with_timeout(
        &mut self,
        timeout: Duration,
    ) -> Result<Option<Notification>> {
        self.wait_for_notification(Some(timeout)).await
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;
    use crate::auth::{Codec, ServerParams};
    use crate::config::Config;
    use crate::connection::ConnectionState;
    use crate::protocol::TransactionStatus;
    use crate::transport::{BufferedTransport, ClientTransport, MockTransport, PgTransport};
    use std::collections::VecDeque;

    fn make_connection(read_data: Vec<u8>) -> Connection {
        let transport = PgTransport::Plain(BufferedTransport::new(ClientTransport::Mock(
            MockTransport::new(read_data),
        )));
        Connection {
            transport,
            codec: Codec::new(),
            server_params: ServerParams::default(),
            state: ConnectionState::Idle,
            config: Config::new(),
            transaction_status: TransactionStatus::Idle,
            notification_queue: VecDeque::new(),
            notice_handler: None,
            statement_counter: 0,
            needs_recovery: false,
            health: crate::reconnect::session::ConnectionHealth::new(),
            session_state: crate::reconnect::session::SessionState::new(),
        }
    }

    fn build_command_complete_msg(tag: &str) -> Vec<u8> {
        let mut buf = vec![b'C'];
        let mut body = Vec::new();
        body.extend_from_slice(tag.as_bytes());
        body.push(0);
        let len = (body.len() + 4) as i32;
        buf.extend_from_slice(&len.to_be_bytes());
        buf.extend_from_slice(&body);
        buf
    }

    fn build_ready_for_query(status: u8) -> Vec<u8> {
        vec![b'Z', 0, 0, 0, 5, status]
    }

    fn build_notification_response(pid: i32, channel: &str, payload: &str) -> Vec<u8> {
        let mut buf = vec![b'A'];
        let mut body = Vec::new();
        body.extend_from_slice(&pid.to_be_bytes());
        body.extend_from_slice(channel.as_bytes());
        body.push(0);
        body.extend_from_slice(payload.as_bytes());
        body.push(0);
        let len = (body.len() + 4) as i32;
        buf.extend_from_slice(&len.to_be_bytes());
        buf.extend_from_slice(&body);
        buf
    }

    fn build_row_description_msg(fields: &[(&str, u32)]) -> Vec<u8> {
        let mut buf = vec![b'T'];
        let mut body = Vec::new();
        body.extend_from_slice(&(fields.len() as i16).to_be_bytes());
        for (name, type_oid) in fields {
            body.extend_from_slice(name.as_bytes());
            body.push(0);
            body.extend_from_slice(&0u32.to_be_bytes()); // table_oid
            body.extend_from_slice(&0i16.to_be_bytes()); // column_id
            body.extend_from_slice(&type_oid.to_be_bytes()); // type_oid
            body.extend_from_slice(&(-1i16).to_be_bytes()); // type_size
            body.extend_from_slice(&(-1i32).to_be_bytes()); // type_modifier
            body.extend_from_slice(&0i16.to_be_bytes()); // format
        }
        let len = (body.len() + 4) as i32;
        buf.extend_from_slice(&len.to_be_bytes());
        buf.extend_from_slice(&body);
        buf
    }

    fn build_data_row_msg(values: &[Option<&str>]) -> Vec<u8> {
        let mut buf = vec![b'D'];
        let mut body = Vec::new();
        body.extend_from_slice(&(values.len() as i16).to_be_bytes());
        for val in values {
            match val {
                Some(v) => {
                    let bytes = v.as_bytes();
                    body.extend_from_slice(&(bytes.len() as i32).to_be_bytes());
                    body.extend_from_slice(bytes);
                }
                None => {
                    body.extend_from_slice(&(-1i32).to_be_bytes());
                }
            }
        }
        let len = (body.len() + 4) as i32;
        buf.extend_from_slice(&len.to_be_bytes());
        buf.extend_from_slice(&body);
        buf
    }

    #[tokio::test]
    async fn test_listen() {
        let mut data = Vec::new();
        data.extend_from_slice(&build_command_complete_msg("LISTEN"));
        data.extend_from_slice(&build_ready_for_query(b'I'));

        let mut conn = make_connection(data);
        conn.listen("my_channel").await.unwrap();
        assert!(conn.is_idle());
        assert!(conn
            .session_state()
            .listen_channels()
            .contains("my_channel"));
    }

    #[tokio::test]
    async fn test_unlisten() {
        let mut data = Vec::new();
        data.extend_from_slice(&build_command_complete_msg("UNLISTEN"));
        data.extend_from_slice(&build_ready_for_query(b'I'));

        let mut conn = make_connection(data);
        conn.session_state.track_listen("my_channel");
        conn.unlisten("my_channel").await.unwrap();
        assert!(conn.is_idle());
        assert!(!conn
            .session_state()
            .listen_channels()
            .contains("my_channel"));
    }

    #[tokio::test]
    async fn test_unlisten_all() {
        let mut data = Vec::new();
        data.extend_from_slice(&build_command_complete_msg("UNLISTEN"));
        data.extend_from_slice(&build_ready_for_query(b'I'));

        let mut conn = make_connection(data);
        conn.session_state.track_listen("ch1");
        conn.session_state.track_listen("ch2");
        conn.unlisten_all().await.unwrap();
        assert!(conn.is_idle());
        assert!(conn.session_state().listen_channels().is_empty());
    }

    #[tokio::test]
    async fn test_notify() {
        let mut data = Vec::new();
        // pg_notify returns a row
        data.extend_from_slice(&build_row_description_msg(&[(
            "pg_notify",
            crate::types::TEXT_OID,
        )]));
        data.extend_from_slice(&build_data_row_msg(&[Some("LISTEN")]));
        data.extend_from_slice(&build_command_complete_msg("SELECT 1"));
        data.extend_from_slice(&build_ready_for_query(b'I'));

        let mut conn = make_connection(data);
        conn.notify("my_channel", "hello").await.unwrap();
        assert!(conn.is_idle());
    }

    #[tokio::test]
    async fn test_notifications_buffered() {
        let mut conn = make_connection(vec![]);

        // Manually push notifications into the queue
        conn.notification_queue.push_back(Notification {
            process_id: 1,
            channel: "ch1".to_string(),
            payload: "hello".to_string(),
        });
        conn.notification_queue.push_back(Notification {
            process_id: 2,
            channel: "ch2".to_string(),
            payload: "world".to_string(),
        });

        let notifications = conn.notifications();
        assert_eq!(notifications.len(), 2);
        assert_eq!(notifications[0].channel, "ch1");
        assert_eq!(notifications[1].channel, "ch2");

        // Queue should be empty now
        assert!(conn.notifications().is_empty());
    }

    #[tokio::test]
    async fn test_wait_for_notification_from_queue() {
        let mut conn = make_connection(vec![]);

        // Pre-buffer a notification
        conn.notification_queue.push_back(Notification {
            process_id: 42,
            channel: "test".to_string(),
            payload: "payload".to_string(),
        });

        // Should return immediately from the queue
        let n = conn.wait_for_notification(None).await.unwrap();
        assert!(n.is_some());
        let n = n.unwrap();
        assert_eq!(n.process_id, 42);
        assert_eq!(n.channel, "test");
        assert_eq!(n.payload, "payload");
    }

    #[tokio::test]
    async fn test_wait_for_notification_from_server() {
        let mut data = Vec::new();
        // EmptyQueryResponse for the empty query
        data.extend_from_slice(&[b'I', 0, 0, 0, 4]); // EmptyQueryResponse
                                                     // NotificationResponse
        data.extend_from_slice(&build_notification_response(99, "events", "user_login"));
        // ReadyForQuery
        data.extend_from_slice(&build_ready_for_query(b'I'));

        let mut conn = make_connection(data);
        let n = conn.wait_for_notification(None).await.unwrap();
        assert!(n.is_some());
        let n = n.unwrap();
        assert_eq!(n.process_id, 99);
        assert_eq!(n.channel, "events");
        assert_eq!(n.payload, "user_login");
    }

    #[tokio::test]
    async fn test_wait_for_notification_with_timeout_from_queue() {
        let mut conn = make_connection(vec![]);

        // Pre-buffer a notification
        conn.notification_queue.push_back(Notification {
            process_id: 7,
            channel: "timeout_ch".to_string(),
            payload: "timeout_payload".to_string(),
        });

        // Should return immediately from the queue, ignoring the timeout
        let n = conn
            .wait_for_notification_with_timeout(Duration::from_secs(60))
            .await
            .unwrap();
        assert!(n.is_some());
        let n = n.unwrap();
        assert_eq!(n.process_id, 7);
        assert_eq!(n.channel, "timeout_ch");
        assert_eq!(n.payload, "timeout_payload");
    }

    #[tokio::test]
    async fn test_wait_for_notification_zero_timeout() {
        let mut conn = make_connection(vec![]);

        // With zero timeout and empty queue, should return None immediately
        let n = conn
            .wait_for_notification(Some(Duration::ZERO))
            .await
            .unwrap();
        assert!(n.is_none());
    }

    #[tokio::test]
    async fn test_wait_for_notification_with_timeout_zero() {
        let mut conn = make_connection(vec![]);

        // wait_for_notification_with_timeout with zero duration
        let n = conn
            .wait_for_notification_with_timeout(Duration::ZERO)
            .await
            .unwrap();
        assert!(n.is_none());
    }
}