notification/notice support

Author: sfackler

tokio-postgres/src/lib.rs

@@ -34,7 +34,7 @@ pub use postgres_shared::{error, params, types};
 #[doc(inline)]
 pub use postgres_shared::{CancelData, Notification};
 
-use error::Error;
+use error::{DbError, Error};
 use params::ConnectParams;
 use tls::TlsConnect;
 use types::{FromSql, ToSql, Type};
@@ -104,6 +104,10 @@ impl Connection {
     pub fn parameter(&self, name: &str) -> Option<&str> {
         self.0.parameter(name)
     }
+
+    pub fn poll_message(&mut self) -> Poll<Option<AsyncMessage>, Error> {
+        self.0.poll_message()
+    }
 }
 
 impl Future for Connection {
@@ -115,6 +119,13 @@ impl Future for Connection {
     }
 }
 
+pub enum AsyncMessage {
+    Notice(DbError),
+    Notification(Notification),
+    #[doc(hidden)]
+    __NonExhaustive,
+}
+
 #[must_use = "futures do nothing unless polled"]
 pub struct CancelQuery(proto::CancelFuture);
 

tokio-postgres/src/proto/connection.rs

@@ -6,11 +6,10 @@ use std::collections::{HashMap, VecDeque};
 use std::io;
 use tokio_codec::Framed;
 
-use disconnected;
-use error::{self, Error};
+use error::{self, DbError, Error};
 use proto::codec::PostgresCodec;
 use tls::TlsStream;
-use {bad_response, CancelData};
+use {bad_response, disconnected, AsyncMessage, CancelData, Notification};
 
 pub struct Request {
     pub messages: Vec<u8>,
@@ -71,10 +70,10 @@ impl Connection {
         self.stream.poll()
     }
 
-    fn poll_read(&mut self) -> Result<(), Error> {
+    fn poll_read(&mut self) -> Result<Option<AsyncMessage>, Error> {
         if self.state != State::Active {
             trace!("poll_read: done");
-            return Ok(());
+            return Ok(None);
         }
 
         loop {
@@ -85,14 +84,22 @@ impl Connection {
                 }
                 Async::NotReady => {
                     trace!("poll_read: waiting on response");
-                    return Ok(());
+                    return Ok(None);
                 }
             };
 
             let message = match message {
-                Message::NoticeResponse(_) | Message::NotificationResponse(_) => {
-                    // FIXME handle these
-                    continue;
+                Message::NoticeResponse(body) => {
+                    let error = DbError::new(&mut body.fields())?;
+                    return Ok(Some(AsyncMessage::Notice(error)));
+                }
+                Message::NotificationResponse(body) => {
+                    let notification = Notification {
+                        process_id: body.process_id(),
+                        channel: body.channel()?.to_string(),
+                        payload: body.message()?.to_string(),
+                    };
+                    return Ok(Some(AsyncMessage::Notification(notification)));
                 }
                 Message::ParameterStatus(body) => {
                     self.parameters
@@ -127,7 +134,7 @@ impl Connection {
                     self.responses.push_front(sender);
                     self.pending_response = Some(message);
                     trace!("poll_read: waiting on socket");
-                    return Ok(());
+                    return Ok(None);
                 }
             }
         }
@@ -225,18 +232,26 @@ impl Connection {
             Err(e) => Err(Error::from(e)),
         }
     }
+
+    pub fn poll_message(&mut self) -> Poll<Option<AsyncMessage>, Error> {
+        let message = self.poll_read()?;
+        let want_flush = self.poll_write()?;
+        if want_flush {
+            self.poll_flush()?;
+        }
+        match message {
+            Some(message) => Ok(Async::Ready(Some(message))),
+            None => self.poll_shutdown().map(|r| r.map(|()| None)),
+        }
+    }
 }
 
 impl Future for Connection {
     type Item = ();
     type Error = Error;
 
     fn poll(&mut self) -> Poll<(), Error> {
-        self.poll_read()?;
-        let want_flush = self.poll_write()?;
-        if want_flush {
-            self.poll_flush()?;
-        }
-        self.poll_shutdown()
+        while let Some(_) = try_ready!(self.poll_message()) {}
+        Ok(Async::Ready(()))
     }
 }

tokio-postgres/tests/test.rs

@@ -1,15 +1,21 @@
 extern crate env_logger;
-extern crate futures;
 extern crate tokio;
 extern crate tokio_postgres;
 
+#[macro_use]
+extern crate futures;
+#[macro_use]
+extern crate log;
+
+use futures::future;
+use futures::sync::mpsc;
 use std::time::{Duration, Instant};
 use tokio::prelude::*;
 use tokio::runtime::current_thread::Runtime;
 use tokio::timer::Delay;
 use tokio_postgres::error::SqlState;
 use tokio_postgres::types::{Kind, Type};
-use tokio_postgres::TlsMode;
+use tokio_postgres::{AsyncMessage, TlsMode};
 
 fn smoke_test(url: &str) {
     let _ = env_logger::try_init();
@@ -447,3 +453,54 @@ fn custom_simple() {
     assert_eq!("hstore", ty.name());
     assert_eq!(&Kind::Simple, ty.kind());
 }
+
+#[test]
+fn notifications() {
+    let _ = env_logger::try_init();
+    let mut runtime = Runtime::new().unwrap();
+
+    let handshake = tokio_postgres::connect(
+        "postgres://postgres@localhost:5433".parse().unwrap(),
+        TlsMode::None,
+    );
+    let (mut client, mut connection) = runtime.block_on(handshake).unwrap();
+
+    let (tx, rx) = mpsc::unbounded();
+    let connection = future::poll_fn(move || {
+        while let Some(message) = try_ready!(connection.poll_message().map_err(|e| panic!("{}", e)))
+        {
+            if let AsyncMessage::Notification(notification) = message {
+                debug!("received {}", notification.payload);
+                tx.unbounded_send(notification).unwrap();
+            }
+        }
+
+        Ok(Async::Ready(()))
+    });
+    runtime.handle().spawn(connection).unwrap();
+
+    let listen = client.prepare("LISTEN test_notifications");
+    let listen = runtime.block_on(listen).unwrap();
+    runtime.block_on(client.execute(&listen, &[])).unwrap();
+    drop(listen); // FIXME
+
+    let notify = client.prepare("NOTIFY test_notifications, 'hello'");
+    let notify = runtime.block_on(notify).unwrap();
+    runtime.block_on(client.execute(&notify, &[])).unwrap();
+    drop(notify); // FIXME
+
+    let notify = client.prepare("NOTIFY test_notifications, 'world'");
+    let notify = runtime.block_on(notify).unwrap();
+    runtime.block_on(client.execute(&notify, &[])).unwrap();
+    drop(notify); // FIXME
+
+    drop(client);
+    runtime.run().unwrap();
+
+    let notifications = rx.collect().wait().unwrap();
+    assert_eq!(notifications.len(), 2);
+    assert_eq!(notifications[0].channel, "test_notifications");
+    assert_eq!(notifications[0].payload, "hello");
+    assert_eq!(notifications[1].channel, "test_notifications");
+    assert_eq!(notifications[1].payload, "world");
+}