oxisql_postgres/
notify.rs

1//! PostgreSQL LISTEN/NOTIFY async notification support.
2//!
3//! Provides [`PgNotification`], [`NotificationStream`], and the methods
4//! [`notify`] and [`listen`] for [`PgConnection`].
5//!
6//! # Architecture
7//!
8//! `tokio-postgres` delivers async messages (including `NOTIFY` payloads) via
9//! the `Connection` object's `poll_message` method, **not** via the `Client`.
10//! When [`PgConnection::connect`] is called, the `Connection` driver is spawned
11//! as a background task.  To intercept notifications, that task forwards
12//! `AsyncMessage::Notification` items into a
13//! `tokio::sync::broadcast::Sender<tokio_postgres::Notification>`.
14//!
15//! `PgConnection` stores an `Option<broadcast::Sender<_>>`.  When the
16//! connection was created via [`PgConnection::from_client`] (where no
17//! `Connection` object is available), this field is `None` and `listen()`
18//! returns [`PgError::Notify`] with a descriptive message.
19//!
20//! Multiple [`NotificationStream`] values can coexist — each holds its own
21//! `broadcast::Receiver` — and all receive every notification (filtered to
22//! the relevant channel by `recv_timeout`).
23//!
24//! # Channel identifier safety
25//!
26//! Channel names are validated to contain only ASCII alphanumeric characters
27//! and underscores, preventing SQL injection.
28
29use std::sync::Arc;
30use std::time::Duration;
31
32use tokio::sync::{broadcast, Mutex};
33
34use crate::error::PgError;
35
36/// Capacity of the broadcast channel used to buffer notifications.
37///
38/// If more than this many notifications arrive without being consumed, the
39/// oldest are silently dropped (broadcast semantics).
40const NOTIFICATION_CAPACITY: usize = 64;
41
42// ── PgNotification ────────────────────────────────────────────────────────────
43
44/// A notification received via the PostgreSQL `LISTEN/NOTIFY` mechanism.
45#[derive(Debug, Clone)]
46pub struct PgNotification {
47    /// The channel name on which the notification was raised.
48    pub channel: String,
49    /// The payload string passed by the notifying backend.
50    pub payload: String,
51    /// The process ID of the notifying backend process.
52    pub process_id: i32,
53}
54
55impl From<tokio_postgres::Notification> for PgNotification {
56    fn from(n: tokio_postgres::Notification) -> Self {
57        PgNotification {
58            process_id: n.process_id(),
59            channel: n.channel().to_string(),
60            payload: n.payload().to_string(),
61        }
62    }
63}
64
65// ── NotificationStream ────────────────────────────────────────────────────────
66
67/// A live subscription to a PostgreSQL notification channel.
68///
69/// Obtained from `PgConnection::listen`.  Drive the stream by calling
70/// [`recv_timeout`] in a loop.  Call [`unlisten`] when done to deregister the
71/// channel on the server.
72///
73/// [`recv_timeout`]: NotificationStream::recv_timeout
74/// [`unlisten`]: NotificationStream::unlisten
75pub struct NotificationStream {
76    /// Broadcast receiver from the connection's notification forwarder.
77    pub(crate) rx: broadcast::Receiver<PgNotification>,
78    /// The channel name this stream is subscribed to.
79    pub(crate) channel: String,
80    /// The client needed for `UNLISTEN`.
81    pub(crate) inner: Arc<Mutex<tokio_postgres::Client>>,
82}
83
84impl NotificationStream {
85    /// Wait up to `timeout` for the next notification on this channel.
86    ///
87    /// Returns `Some(notification)` if a matching notification arrives within
88    /// the timeout, or `None` if the timeout elapses.  Notifications for
89    /// other channels are silently discarded.
90    pub async fn recv_timeout(&mut self, timeout: Duration) -> Option<PgNotification> {
91        let deadline = tokio::time::Instant::now() + timeout;
92        loop {
93            let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
94            if remaining.is_zero() {
95                return None;
96            }
97            match tokio::time::timeout(remaining, self.rx.recv()).await {
98                // Timed out.
99                Err(_elapsed) => return None,
100                // Channel lagged — the oldest notification(s) were dropped.
101                // Continue waiting for the next one.
102                Ok(Err(broadcast::error::RecvError::Lagged(_))) => continue,
103                // Channel closed — the connection was dropped.
104                Ok(Err(broadcast::error::RecvError::Closed)) => return None,
105                Ok(Ok(n)) if n.channel == self.channel => return Some(n),
106                // Notification for a different channel — skip and keep waiting.
107                Ok(Ok(_other)) => continue,
108            }
109        }
110    }
111
112    /// Send `UNLISTEN` for this channel and consume the stream.
113    ///
114    /// After this call the connection will no longer receive notifications
115    /// for this channel (unless another `PgConnection::listen` call
116    /// re-registers it).
117    pub async fn unlisten(self) -> Result<(), PgError> {
118        validate_channel_name(&self.channel)?;
119        let client = self.inner.lock().await;
120        client
121            .batch_execute(&format!("UNLISTEN {}", self.channel))
122            .await
123            .map_err(|e| PgError::Notify(e.to_string()))
124    }
125}
126
127// ── Connection driver with notification forwarding ────────────────────────────
128
129/// Spawn the tokio-postgres `Connection` driver and forward
130/// `AsyncMessage::Notification` items to a broadcast channel.
131///
132/// Returns the `broadcast::Sender` that the `PgConnection` should store.
133/// The spawned task exits when the connection closes.
134pub(crate) fn spawn_connection_driver<S, T>(
135    mut connection: tokio_postgres::Connection<S, T>,
136) -> broadcast::Sender<PgNotification>
137where
138    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
139    T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
140{
141    let (tx, _rx) = broadcast::channel::<PgNotification>(NOTIFICATION_CAPACITY);
142    let tx_clone = tx.clone();
143
144    tokio::spawn(async move {
145        loop {
146            match futures::future::poll_fn(|cx| connection.poll_message(cx)).await {
147                None => break,
148                Some(Err(_e)) => break,
149                Some(Ok(tokio_postgres::AsyncMessage::Notification(n))) => {
150                    // If no receivers are listening, the send is a no-op.
151                    let _ = tx_clone.send(PgNotification::from(n));
152                }
153                Some(Ok(_)) => {
154                    // Notices and other async messages are ignored.
155                }
156            }
157        }
158    });
159
160    tx
161}
162
163// ── Notify (NOTIFY channel, payload) ─────────────────────────────────────────
164
165/// Send a `NOTIFY` to the given channel with an optional payload.
166///
167/// The channel name must consist only of ASCII alphanumeric characters and
168/// underscores.  The payload is single-quote–escaped before inclusion in the
169/// SQL command string.
170pub(crate) async fn notify(
171    inner: &Arc<Mutex<tokio_postgres::Client>>,
172    channel: &str,
173    payload: &str,
174) -> Result<(), PgError> {
175    validate_channel_name(channel)?;
176    let escaped_payload = payload.replace('\'', "''");
177    let sql = format!("NOTIFY {channel}, '{escaped_payload}'");
178    let client = inner.lock().await;
179    client
180        .batch_execute(&sql)
181        .await
182        .map_err(|e| PgError::Notify(e.to_string()))
183}
184
185// ── Listen (LISTEN channel) ───────────────────────────────────────────────────
186
187/// Register `LISTEN` on the given channel and return a [`NotificationStream`].
188///
189/// Requires that the `PgConnection` was created via [`PgConnection::connect`]
190/// (not `from_client`); otherwise returns [`PgError::Notify`].
191pub(crate) async fn listen(
192    inner: &Arc<Mutex<tokio_postgres::Client>>,
193    notif_tx: &Option<broadcast::Sender<PgNotification>>,
194    channel: &str,
195) -> Result<NotificationStream, PgError> {
196    validate_channel_name(channel)?;
197
198    let tx = notif_tx.as_ref().ok_or_else(|| {
199        PgError::Notify(
200            "LISTEN requires a connection created via PgConnection::connect; \
201             connections created via from_client cannot receive notifications"
202                .to_string(),
203        )
204    })?;
205
206    let sql = format!("LISTEN {channel}");
207    {
208        let client = inner.lock().await;
209        client
210            .batch_execute(&sql)
211            .await
212            .map_err(|e| PgError::Notify(e.to_string()))?;
213    }
214
215    Ok(NotificationStream {
216        rx: tx.subscribe(),
217        channel: channel.to_string(),
218        inner: Arc::clone(inner),
219    })
220}
221
222// ── Identifier validation ─────────────────────────────────────────────────────
223
224/// Validate a PostgreSQL channel (or identifier) name.
225///
226/// Accepts only ASCII alphanumeric characters and underscores to prevent SQL
227/// injection.  Returns [`PgError::Notify`] if the name is invalid.
228pub(crate) fn validate_channel_name(name: &str) -> Result<(), PgError> {
229    if name.is_empty() {
230        return Err(PgError::Notify(
231            "channel name must not be empty".to_string(),
232        ));
233    }
234    if !name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
235        return Err(PgError::Notify(format!(
236            "invalid channel name {name:?}: only ASCII alphanumeric characters and underscores are allowed"
237        )));
238    }
239    Ok(())
240}
oxisql_postgres/notify.rs

oxisql_postgres/
notify.rs
