Adjust Type::Enum to carry the variants

Author: sfackler

src/lib.rs

@@ -98,6 +98,7 @@ pub mod types;
 pub mod notification;
 
 const TYPEINFO_QUERY: &'static str = "__typeinfo";
+const TYPEINFO_ENUM_QUERY: &'static str = "__typeinfo_enum";
 const TYPEINFO_ARRAY_QUERY: &'static str = "__typeinfo_array";
 
 /// A type alias of the result returned by many methods.
@@ -464,12 +465,23 @@ impl InnerConnection {
 
     #[cfg_attr(rustfmt, rustfmt_skip)]
     fn setup_typeinfo_query(&mut self) -> result::Result<(), ConnectError> {
+        match self.raw_prepare(TYPEINFO_ENUM_QUERY,
+                               "SELECT enumlabel \
+                                FROM pg_catalog.pg_enum \
+                                WHERE enumtypid = $1 \
+                                ORDER BY enumsortorder") {
+            Ok(..) => {}
+            Err(Error::Io(e)) => return Err(ConnectError::Io(e)),
+            Err(Error::Db(e)) => return Err(ConnectError::Db(e)),
+            Err(Error::Conversion(_)) => unreachable!(),
+        }
+
         match self.raw_prepare(TYPEINFO_ARRAY_QUERY,
                                "SELECT attname, atttypid \
                                 FROM pg_catalog.pg_attribute \
                                 WHERE attrelid = $1 \
-                                AND NOT attisdropped \
-                                AND attnum > 0 \
+                                    AND NOT attisdropped \
+                                    AND attnum > 0 \
                                 ORDER BY attnum") {
             Ok(..) => {}
             Err(Error::Io(e)) => return Err(ConnectError::Io(e)),
@@ -865,7 +877,19 @@ impl InnerConnection {
         };
 
         let kind = if type_ == b'e' as i8 {
-            Kind::Enum
+            try!(self.raw_execute(TYPEINFO_ENUM_QUERY, "", 0, &[Type::Oid], &[&oid]));
+            let mut rows = VecDeque::new();
+            try!(self.read_rows(&mut rows));
+
+            let ctx = SessionInfo::new(self);
+            let mut variants = vec![];
+            for row in rows {
+                variants.push(try!(String::from_sql(&Type::Name,
+                                                    &mut &**row[0].as_ref().unwrap(),
+                                                    &ctx)));
+            }
+
+            Kind::Enum(variants)
         } else if type_ == b'p' as i8 {
             Kind::Pseudo
         } else if basetype != 0 {

src/types/mod.rs

@@ -105,8 +105,8 @@ pub type Oid = u32;
 pub enum Kind {
     /// A simple type like `VARCHAR` or `INTEGER`.
     Simple,
-    /// An enumerated type.
-    Enum,
+    /// An enumerated type along with its variants.
+    Enum(Vec<String>),
     /// A pseudo-type.
     Pseudo,
     /// An array type along with the type of its elements.

tests/types/mod.rs

@@ -307,3 +307,19 @@ fn composite() {
         t => panic!("bad type {:?}", t),
     }
 }
+
+#[test]
+fn enum_() {
+    let conn = Connection::connect("postgres://postgres@localhost", SslMode::None).unwrap();
+    conn.batch_execute("CREATE TYPE pg_temp.mood AS ENUM ('sad', 'ok', 'happy');").unwrap();
+
+    let stmt = conn.prepare("SELECT $1::mood").unwrap();
+    let type_ = &stmt.param_types()[0];
+    assert_eq!(type_.name(), "mood");
+    match type_.kind() {
+        &Kind::Enum(ref variants) => {
+            assert_eq!(variants, &["sad".to_owned(), "ok".to_owned(), "happy".to_owned()]);
+        }
+        _ => panic!("bad type"),
+    }
+}