From 4c46f9fa89fd75336bf75831d5c91df2c338b6b9 Mon Sep 17 00:00:00 2001
From: Steven Fackler <sfackler@gmail.com>
Date: Wed, 17 Feb 2016 22:22:31 -0800
Subject: [PATCH] Adjust Type::Enum to carry the variants

---
 src/lib.rs         | 30 +++++++++++++++++++++++++++---
 src/types/mod.rs   |  4 ++--
 tests/types/mod.rs | 16 ++++++++++++++++
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 9b24f527..5ed2c04d 100644
--- a/src/lib.rs
+++ b/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 {
diff --git a/src/types/mod.rs b/src/types/mod.rs
index cdeb0f62..1fddbff1 100644
--- a/src/types/mod.rs
+++ b/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.
diff --git a/tests/types/mod.rs b/tests/types/mod.rs
index 10c8cf22..fe71b053 100644
--- a/tests/types/mod.rs
+++ b/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"),
+    }
+}