From cfe522034e02b5f164efec94556b83c43908415d Mon Sep 17 00:00:00 2001 From: Ryan Leckey Date: Sun, 30 Jun 2019 06:35:30 -0700 Subject: [PATCH] Implement RowDescription --- sqlx-postgres-protocol/src/lib.rs | 2 + sqlx-postgres-protocol/src/message.rs | 7 +- sqlx-postgres-protocol/src/row_description.rs | 172 ++++++++++++++++++ sqlx-postgres/src/connection/query.rs | 4 + 4 files changed, 184 insertions(+), 1 deletion(-) create mode 100644 sqlx-postgres-protocol/src/row_description.rs diff --git a/sqlx-postgres-protocol/src/lib.rs b/sqlx-postgres-protocol/src/lib.rs index 001c3aaf..7422959e 100644 --- a/sqlx-postgres-protocol/src/lib.rs +++ b/sqlx-postgres-protocol/src/lib.rs @@ -10,6 +10,7 @@ mod password_message; mod query; mod ready_for_query; mod response; +mod row_description; mod startup_message; mod terminate; @@ -24,6 +25,7 @@ pub use self::{ query::Query, ready_for_query::{ReadyForQuery, TransactionStatus}, response::{Response, ResponseBuilder, Severity}, + row_description::{FieldDescription, FieldDescriptions, RowDescription}, startup_message::StartupMessage, terminate::Terminate, }; diff --git a/sqlx-postgres-protocol/src/message.rs b/sqlx-postgres-protocol/src/message.rs index b2e55659..53a7f762 100644 --- a/sqlx-postgres-protocol/src/message.rs +++ b/sqlx-postgres-protocol/src/message.rs @@ -1,4 +1,7 @@ -use crate::{Authentication, BackendKeyData, Decode, ParameterStatus, ReadyForQuery, Response}; +use crate::{ + Authentication, BackendKeyData, Decode, ParameterStatus, ReadyForQuery, Response, + RowDescription, +}; use byteorder::{BigEndian, ByteOrder}; use bytes::BytesMut; use std::io; @@ -9,6 +12,7 @@ pub enum Message { ParameterStatus(ParameterStatus), BackendKeyData(BackendKeyData), ReadyForQuery(ReadyForQuery), + RowDescription(RowDescription), Response(Box), } @@ -46,6 +50,7 @@ impl Message { b'Z' => Message::ReadyForQuery(ReadyForQuery::decode(src)?), b'R' => Message::Authentication(Authentication::decode(src)?), b'K' => Message::BackendKeyData(BackendKeyData::decode(src)?), + b'T' => Message::RowDescription(RowDescription::decode(src)?), _ => unimplemented!("decode not implemented for token: {}", token as char), })) diff --git a/sqlx-postgres-protocol/src/row_description.rs b/sqlx-postgres-protocol/src/row_description.rs new file mode 100644 index 00000000..fe806012 --- /dev/null +++ b/sqlx-postgres-protocol/src/row_description.rs @@ -0,0 +1,172 @@ +use crate::Decode; +use byteorder::{BigEndian, ByteOrder}; +use bytes::Bytes; +use memchr::memchr; +use std::{ + io, + mem::size_of_val, + num::{NonZeroI16, NonZeroU32}, + str, +}; + +// TODO: Custom Debug for RowDescription and FieldDescription + +/// A descriptive record on a single field received from PostgreSQL. +#[derive(Debug)] +pub struct FieldDescription<'a> { + name: &'a str, + table_oid: Option, + column_attribute_num: Option, + type_oid: u32, + type_size: i16, + type_modifier: i32, + format: i16, +} + +impl<'a> FieldDescription<'a> { + #[inline] + pub fn name(&self) -> &'a str { + self.name + } + + #[inline] + pub fn table_oid(&self) -> Option { + self.table_oid.map(Into::into) + } + + #[inline] + pub fn column_attribute_num(&self) -> Option { + self.column_attribute_num.map(Into::into) + } + + #[inline] + pub fn type_oid(&self) -> u32 { + self.type_oid + } + + #[inline] + pub fn type_size(&self) -> i16 { + self.type_size + } + + #[inline] + pub fn type_modifier(&self) -> i32 { + self.type_modifier + } + + #[inline] + pub fn format(&self) -> i16 { + self.format + } +} + +#[derive(Debug)] +pub struct RowDescription { + // The number of fields in a row (can be zero). + len: u16, + data: Bytes, +} + +impl RowDescription { + pub fn fields(&self) -> FieldDescriptions<'_> { + FieldDescriptions { + rem: self.len, + buf: &*self.data, + } + } +} + +impl Decode for RowDescription { + fn decode(src: Bytes) -> io::Result { + let len = BigEndian::read_u16(&src[..2]); + + Ok(Self { + len, + data: src.slice_from(2), + }) + } +} + +pub struct FieldDescriptions<'a> { + rem: u16, + buf: &'a [u8], +} + +impl<'a> Iterator for FieldDescriptions<'a> { + type Item = FieldDescription<'a>; + + fn size_hint(&self) -> (usize, Option) { + (self.rem as usize, Some(self.rem as usize)) + } + + fn next(&mut self) -> Option { + if self.rem == 0 { + return None; + } + + self.rem -= 1; + + let name_end = memchr(0, &self.buf).unwrap(); + let mut idx = name_end + 1; + let name = unsafe { str::from_utf8_unchecked(&self.buf[..name_end]) }; + + let table_oid = BigEndian::read_u32(&self.buf[idx..]); + idx += size_of_val(&table_oid); + + let column_attribute_num = BigEndian::read_i16(&self.buf[idx..]); + idx += size_of_val(&column_attribute_num); + + let type_oid = BigEndian::read_u32(&self.buf[idx..]); + idx += size_of_val(&type_oid); + + let type_size = BigEndian::read_i16(&self.buf[idx..]); + idx += size_of_val(&type_size); + + let type_modifier = BigEndian::read_i32(&self.buf[idx..]); + idx += size_of_val(&type_modifier); + + let format = BigEndian::read_i16(&self.buf[idx..]); + + Some(FieldDescription { + name, + table_oid: NonZeroU32::new(table_oid), + column_attribute_num: NonZeroI16::new(column_attribute_num), + type_oid, + type_size, + type_modifier, + format, + }) + } +} + +impl<'a> ExactSizeIterator for FieldDescriptions<'a> {} + +#[cfg(test)] +mod tests { + use super::RowDescription; + use crate::Decode; + use bytes::Bytes; + use std::io; + + const ROW_DESC_1: &[u8] = b"\0\x01?column?\0\0\0\0\0\0\0\0\0\0\x17\0\x04\xff\xff\xff\xff\0\0D\0\0\0\x0b\0\x01\0\0\0\x011"; + + #[test] + fn it_decodes_row_description() -> io::Result<()> { + let src = Bytes::from_static(ROW_DESC_1); + let message = RowDescription::decode(src)?; + assert_eq!(message.fields().len(), 1); + + let mut fields = message.fields(); + + let field_1 = fields.next().unwrap(); + assert_eq!(field_1.name(), "?column?"); + assert_eq!(field_1.table_oid(), None); + assert_eq!(field_1.column_attribute_num(), None); + assert_eq!(field_1.type_oid(), 23); + assert_eq!(field_1.type_size(), 4); + assert_eq!(field_1.type_modifier(), -1); + assert_eq!(field_1.format(), 0); + + Ok(()) + } +} diff --git a/sqlx-postgres/src/connection/query.rs b/sqlx-postgres/src/connection/query.rs index 85ccba01..e0a5c58d 100644 --- a/sqlx-postgres/src/connection/query.rs +++ b/sqlx-postgres/src/connection/query.rs @@ -8,6 +8,10 @@ pub async fn query<'a: 'b, 'b>(conn: &'a mut Connection, query: &'b str) -> io:: while let Some(message) = conn.stream.next().await { match message? { + Message::RowDescription(_) => { + // Do nothing + } + Message::ReadyForQuery(_) => { break; }