Allow custom notice handling

It was a tossup between a ~fn and a ~Trait *shrug*

Author: sfackler

src/lib.rs

@@ -54,9 +54,16 @@ use types::{PostgresType, ToSql, FromSql};
 mod message;
 mod types;
 
-fn handle_notice_response(fields: ~[(u8, ~str)]) {
-    let err = PostgresDbError::new(fields);
-    info2!("{}: {}", err.severity, err.message);
+pub trait PostgresNoticeHandler {
+    fn handle(&mut self, notice: PostgresDbError);
+}
+
+pub struct DefaultNoticeHandler;
+
+impl PostgresNoticeHandler for DefaultNoticeHandler {
+    fn handle(&mut self, notice: PostgresDbError) {
+        info2!("{}: {}", notice.severity, notice.message);
+    }
 }
 
 #[deriving(ToStr)]
@@ -121,7 +128,8 @@ impl PostgresDbError {
 
 pub struct PostgresConnection {
     priv stream: Cell<TcpStream>,
-    priv next_stmt_id: Cell<int>
+    priv next_stmt_id: Cell<int>,
+    priv notice_handler: Cell<~PostgresNoticeHandler>
 }
 
 impl Drop for PostgresConnection {
@@ -171,7 +179,8 @@ impl PostgresConnection {
 
         let conn = PostgresConnection {
             stream: Cell::new(stream),
-            next_stmt_id: Cell::new(0)
+            next_stmt_id: Cell::new(0),
+            notice_handler: Cell::new(~DefaultNoticeHandler as ~PostgresNoticeHandler)
         };
 
         args.push((~"client_encoding", ~"UTF8"));
@@ -250,10 +259,13 @@ impl PostgresConnection {
             };
 
             match msg {
-                NoticeResponse { fields } =>
-                    handle_notice_response(fields),
+                NoticeResponse { fields } => {
+                    let mut handler = self.notice_handler.take();
+                    handler.handle(PostgresDbError::new(fields));
+                    self.notice_handler.put_back(handler);
+                }
                 ParameterStatus { parameter, value } =>
-                    info!("Parameter %s = %s", parameter, value),
+                    debug!("Parameter %s = %s", parameter, value),
                 msg => return msg
             }
         }
@@ -298,6 +310,13 @@ impl PostgresConnection {
         }
     }
 
+    pub fn set_notice_handler(&self, handler: ~PostgresNoticeHandler)
+            -> ~PostgresNoticeHandler {
+        let old_handler = self.notice_handler.take();
+        self.notice_handler.put_back(handler);
+        old_handler
+    }
+
     pub fn prepare<'a>(&'a self, query: &str) -> NormalPostgresStatement<'a> {
         match self.try_prepare(query) {
             Ok(stmt) => stmt,

src/test.rs

@@ -8,7 +8,8 @@ use extra::uuid::Uuid;
 use std::f32;
 use std::f64;
 
-use postgres::{DbError,
+use postgres::{PostgresNoticeHandler,
+               DbError,
                DnsError,
                MissingPassword,
                Position,
@@ -338,6 +339,24 @@ fn test_get_named_fail() {
     let _: i32 = result.next().unwrap()["asdf"];
 }
 
+#[test]
+fn test_custom_notice_handler() {
+    static mut count: uint = 0;
+    struct Handler;
+
+    impl PostgresNoticeHandler for Handler {
+        fn handle(&mut self, _notice: PostgresDbError) {
+            unsafe { count += 1; }
+        }
+    }
+
+    let conn = PostgresConnection::connect("postgres://postgres@localhost");
+    conn.set_notice_handler(~Handler as ~PostgresNoticeHandler);
+    conn.update("CREATE TEMPORARY TABLE foo (id INT PRIMARY KEY)", []);
+
+    assert_eq!(unsafe { count }, 1);
+}
+
 #[test]
 fn test_plaintext_pass() {
     PostgresConnection::connect("postgres://pass_user:password@localhost");