Skip to content

Commit 5730903

Browse files
sfacklersfackler
committed
Lazily initialize typeinfo queries
Cuts out a little bit of traffic for connections that don't use custom types and avoids weird logic for things like CockroachDB that expose a Postgres interface but not pg_catalog.
1 parent 7509d34 commit 5730903

1 file changed

Lines changed: 84 additions & 78 deletions

File tree

‎src/lib.rs‎

Lines changed: 84 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,10 @@ pub mod stmt;
8989
pub mod transaction;
9090
pub mod types;
9191

92+
const TYPEINFO_QUERY_BIT: u8 = 0b0000_0001;
93+
const TYPEINFO_ENUM_QUERY_BIT: u8 = 0b0000_0010;
94+
const TYPEINFO_COMPOSITE_QUERY_BIT: u8 = 0b0000_0100;
95+
9296
const TYPEINFO_QUERY: &'static str = "__typeinfo";
9397
const TYPEINFO_ENUM_QUERY: &'static str = "__typeinfo_enum";
9498
const TYPEINFO_COMPOSITE_QUERY: &'static str = "__typeinfo_composite";
@@ -215,6 +219,7 @@ struct InnerConnection {
215219
trans_depth: u32,
216220
desynchronized: bool,
217221
finished: bool,
222+
typeinfo_state: u8,
218223
}
219224

220225
impl Drop for InnerConnection {
@@ -256,6 +261,7 @@ impl InnerConnection {
256261
desynchronized: false,
257262
finished: false,
258263
trans_depth: 0,
264+
typeinfo_state: 0,
259265
};
260266

261267
options.push(("client_encoding".to_owned(), "UTF8".to_owned()));
@@ -287,87 +293,9 @@ impl InnerConnection {
287293
}
288294
}
289295

290-
try!(conn.setup_typeinfo_query());
291-
292296
Ok(conn)
293297
}
294298

295-
#[cfg_attr(rustfmt, rustfmt_skip)]
296-
fn setup_typeinfo_query(&mut self) -> result::Result<(), ConnectError> {
297-
match self.raw_prepare(TYPEINFO_ENUM_QUERY,
298-
"SELECT enumlabel \
299-
FROM pg_catalog.pg_enum \
300-
WHERE enumtypid = $1 \
301-
ORDER BY enumsortorder") {
302-
Ok(..) => {}
303-
Err(Error::Io(e)) => return Err(ConnectError::Io(e)),
304-
// Postgres 9.0 doesn't have enumsortorder
305-
Err(Error::Db(ref e)) if e.code == SqlState::UndefinedColumn => {
306-
match self.raw_prepare(TYPEINFO_ENUM_QUERY,
307-
"SELECT enumlabel \
308-
FROM pg_catalog.pg_enum \
309-
WHERE enumtypid = $1 \
310-
ORDER BY oid") {
311-
Ok(..) => {}
312-
Err(Error::Io(e)) => return Err(ConnectError::Io(e)),
313-
Err(Error::Db(e)) => return Err(ConnectError::Db(e)),
314-
Err(Error::Conversion(_)) => unreachable!(),
315-
}
316-
}
317-
// Old versions of Postgres and things like Redshift don't support enums
318-
Err(Error::Db(ref e)) if e.code == SqlState::UndefinedTable => {}
319-
// Some Postgres-like databases are missing a pg_catalog (e.g. Cockroach)
320-
Err(Error::Db(ref e)) if e.code == SqlState::InvalidCatalogName => return Ok(()),
321-
Err(Error::Db(e)) => return Err(ConnectError::Db(e)),
322-
Err(Error::Conversion(_)) => unreachable!(),
323-
}
324-
325-
match self.raw_prepare(TYPEINFO_COMPOSITE_QUERY,
326-
"SELECT attname, atttypid \
327-
FROM pg_catalog.pg_attribute \
328-
WHERE attrelid = $1 \
329-
AND NOT attisdropped \
330-
AND attnum > 0 \
331-
ORDER BY attnum") {
332-
Ok(..) => {}
333-
Err(Error::Io(e)) => return Err(ConnectError::Io(e)),
334-
// Old versions of Postgres and things like Redshift don't support composites
335-
Err(Error::Db(ref e)) if e.code == SqlState::UndefinedTable => {}
336-
Err(Error::Db(e)) => return Err(ConnectError::Db(e)),
337-
Err(Error::Conversion(_)) => unreachable!(),
338-
}
339-
340-
match self.raw_prepare(TYPEINFO_QUERY,
341-
"SELECT t.typname, t.typtype, t.typelem, r.rngsubtype, \
342-
t.typbasetype, n.nspname, t.typrelid \
343-
FROM pg_catalog.pg_type t \
344-
LEFT OUTER JOIN pg_catalog.pg_range r ON \
345-
r.rngtypid = t.oid \
346-
INNER JOIN pg_catalog.pg_namespace n ON \
347-
t.typnamespace = n.oid \
348-
WHERE t.oid = $1") {
349-
Ok(..) => return Ok(()),
350-
Err(Error::Io(e)) => return Err(ConnectError::Io(e)),
351-
// Range types weren't added until Postgres 9.2, so pg_range may not exist
352-
Err(Error::Db(ref e)) if e.code == SqlState::UndefinedTable => {}
353-
Err(Error::Db(e)) => return Err(ConnectError::Db(e)),
354-
Err(Error::Conversion(_)) => unreachable!(),
355-
}
356-
357-
match self.raw_prepare(TYPEINFO_QUERY,
358-
"SELECT t.typname, t.typtype, t.typelem, NULL::OID, t.typbasetype, \
359-
n.nspname, t.typrelid \
360-
FROM pg_catalog.pg_type t \
361-
INNER JOIN pg_catalog.pg_namespace n \
362-
ON t.typnamespace = n.oid \
363-
WHERE t.oid = $1") {
364-
Ok(..) => Ok(()),
365-
Err(Error::Io(e)) => Err(ConnectError::Io(e)),
366-
Err(Error::Db(e)) => Err(ConnectError::Db(e)),
367-
Err(Error::Conversion(_)) => unreachable!(),
368-
}
369-
}
370-
371299
fn write_messages(&mut self, messages: &[Frontend]) -> std_io::Result<()> {
372300
debug_assert!(!self.desynchronized);
373301
for message in messages {
@@ -699,8 +627,41 @@ impl InnerConnection {
699627
Ok(Type::Other(ty))
700628
}
701629

630+
fn setup_typeinfo_query(&mut self) -> Result<()> {
631+
if self.typeinfo_state & TYPEINFO_QUERY_BIT != 0 {
632+
return Ok(());
633+
}
634+
635+
match self.raw_prepare(TYPEINFO_QUERY,
636+
"SELECT t.typname, t.typtype, t.typelem, r.rngsubtype, \
637+
t.typbasetype, n.nspname, t.typrelid \
638+
FROM pg_catalog.pg_type t \
639+
LEFT OUTER JOIN pg_catalog.pg_range r ON \
640+
r.rngtypid = t.oid \
641+
INNER JOIN pg_catalog.pg_namespace n ON \
642+
t.typnamespace = n.oid \
643+
WHERE t.oid = $1") {
644+
Ok(..) => {}
645+
// Range types weren't added until Postgres 9.2, so pg_range may not exist
646+
Err(Error::Db(ref e)) if e.code == SqlState::UndefinedTable => {
647+
try!(self.raw_prepare(TYPEINFO_QUERY,
648+
"SELECT t.typname, t.typtype, t.typelem, NULL::OID, \
649+
t.typbasetype, n.nspname, t.typrelid \
650+
FROM pg_catalog.pg_type t \
651+
INNER JOIN pg_catalog.pg_namespace n \
652+
ON t.typnamespace = n.oid \
653+
WHERE t.oid = $1"));
654+
}
655+
Err(e) => return Err(e),
656+
}
657+
658+
self.typeinfo_state |= TYPEINFO_QUERY_BIT;
659+
Ok(())
660+
}
661+
702662
#[allow(if_not_else)]
703663
fn read_type(&mut self, oid: Oid) -> Result<Other> {
664+
try!(self.setup_typeinfo_query());
704665
try!(self.raw_execute(TYPEINFO_QUERY, "", 0, &[Type::Oid], &[&oid]));
705666
let mut rows = VecDeque::new();
706667
try!(self.read_rows(&mut rows));
@@ -743,7 +704,34 @@ impl InnerConnection {
743704
Ok(Other::new(name, oid, kind, schema))
744705
}
745706

707+
fn setup_typeinfo_enum_query(&mut self) -> Result<()> {
708+
if self.typeinfo_state & TYPEINFO_ENUM_QUERY_BIT != 0 {
709+
return Ok(());
710+
}
711+
712+
match self.raw_prepare(TYPEINFO_ENUM_QUERY,
713+
"SELECT enumlabel \
714+
FROM pg_catalog.pg_enum \
715+
WHERE enumtypid = $1 \
716+
ORDER BY enumsortorder") {
717+
Ok(..) => {}
718+
// Postgres 9.0 doesn't have enumsortorder
719+
Err(Error::Db(ref e)) if e.code == SqlState::UndefinedColumn => {
720+
try!(self.raw_prepare(TYPEINFO_ENUM_QUERY,
721+
"SELECT enumlabel \
722+
FROM pg_catalog.pg_enum \
723+
WHERE enumtypid = $1 \
724+
ORDER BY oid"));
725+
}
726+
Err(e) => return Err(e),
727+
}
728+
729+
self.typeinfo_state |= TYPEINFO_ENUM_QUERY_BIT;
730+
Ok(())
731+
}
732+
746733
fn read_enum_variants(&mut self, oid: Oid) -> Result<Vec<String>> {
734+
try!(self.setup_typeinfo_enum_query());
747735
try!(self.raw_execute(TYPEINFO_ENUM_QUERY, "", 0, &[Type::Oid], &[&oid]));
748736
let mut rows = VecDeque::new();
749737
try!(self.read_rows(&mut rows));
@@ -759,7 +747,25 @@ impl InnerConnection {
759747
Ok(variants)
760748
}
761749

750+
fn setup_typeinfo_composite_query(&mut self) -> Result<()> {
751+
if self.typeinfo_state & TYPEINFO_COMPOSITE_QUERY_BIT != 0 {
752+
return Ok(());
753+
}
754+
755+
try!(self.raw_prepare(TYPEINFO_COMPOSITE_QUERY,
756+
"SELECT attname, atttypid \
757+
FROM pg_catalog.pg_attribute \
758+
WHERE attrelid = $1 \
759+
AND NOT attisdropped \
760+
AND attnum > 0 \
761+
ORDER BY attnum"));
762+
763+
self.typeinfo_state |= TYPEINFO_COMPOSITE_QUERY_BIT;
764+
Ok(())
765+
}
766+
762767
fn read_composite_fields(&mut self, relid: Oid) -> Result<Vec<Field>> {
768+
try!(self.setup_typeinfo_composite_query());
763769
try!(self.raw_execute(TYPEINFO_COMPOSITE_QUERY, "", 0, &[Type::Oid], &[&relid]));
764770
let mut rows = VecDeque::new();
765771
try!(self.read_rows(&mut rows));

0 commit comments

Comments
 (0)