Adjust Type::Enum to carry the variants

This commit is contained in:
Steven Fackler 2016-02-17 22:22:31 -08:00
parent bdb8513f52
commit 4c46f9fa89
3 changed files with 45 additions and 5 deletions

View File

@ -98,6 +98,7 @@ pub mod types;
pub mod notification; pub mod notification;
const TYPEINFO_QUERY: &'static str = "__typeinfo"; const TYPEINFO_QUERY: &'static str = "__typeinfo";
const TYPEINFO_ENUM_QUERY: &'static str = "__typeinfo_enum";
const TYPEINFO_ARRAY_QUERY: &'static str = "__typeinfo_array"; const TYPEINFO_ARRAY_QUERY: &'static str = "__typeinfo_array";
/// A type alias of the result returned by many methods. /// A type alias of the result returned by many methods.
@ -464,12 +465,23 @@ impl InnerConnection {
#[cfg_attr(rustfmt, rustfmt_skip)] #[cfg_attr(rustfmt, rustfmt_skip)]
fn setup_typeinfo_query(&mut self) -> result::Result<(), ConnectError> { 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, match self.raw_prepare(TYPEINFO_ARRAY_QUERY,
"SELECT attname, atttypid \ "SELECT attname, atttypid \
FROM pg_catalog.pg_attribute \ FROM pg_catalog.pg_attribute \
WHERE attrelid = $1 \ WHERE attrelid = $1 \
AND NOT attisdropped \ AND NOT attisdropped \
AND attnum > 0 \ AND attnum > 0 \
ORDER BY attnum") { ORDER BY attnum") {
Ok(..) => {} Ok(..) => {}
Err(Error::Io(e)) => return Err(ConnectError::Io(e)), Err(Error::Io(e)) => return Err(ConnectError::Io(e)),
@ -865,7 +877,19 @@ impl InnerConnection {
}; };
let kind = if type_ == b'e' as i8 { 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 { } else if type_ == b'p' as i8 {
Kind::Pseudo Kind::Pseudo
} else if basetype != 0 { } else if basetype != 0 {

View File

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

View File

@ -307,3 +307,19 @@ fn composite() {
t => panic!("bad type {:?}", t), 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"),
}
}