Implement simple query API

jwilm · jwilm · commit 99ebda229695 · 2018-04-27T17:34:24.000-07:00
The simple query API is a more robust version of `batch_execute`. Like
that method, `simple_query` allows passing a `&amp;str` of semicolon
delimited queries. Divergence from `batch_execute` is in the return
type; instead of nothing, a `Vec&lt;TextRows&gt;` is returned. Each entry in
this `Vec` is the result set of one query in the query string. Thus if
there are two semicolon delimited queries, there will be two entries in
this `Vec`.

The `TextRows` and `TextRow` types returned from `simple_query` closely
mirror existing `Rows` and `Row` types with one major difference: only
string values can be retrieved from them.

There are a few TODOs in the code:

* Are text values in this case guaranteed to be utf-8 encoded?
* unwrap call in simple_query which assumes RowDescription is always
  sent
* documentation (denoted with either STUB or TODO)
diff --git a/postgres/src/lib.rs b/postgres/src/lib.rs
@@ -103,6 +103,7 @@ use stmt::{Column, Statement};
 use tls::TlsHandshake;
 use transaction::{IsolationLevel, Transaction};
 use types::{Field, FromSql, IsNull, Kind, Oid, ToSql, Type};
+use text_rows::TextRows;
 
 #[doc(inline)]
 pub use error::Error;
@@ -118,6 +119,7 @@ pub mod notification;
 pub mod params;
 mod priv_io;
 pub mod rows;
+pub mod text_rows;
 pub mod stmt;
 pub mod tls;
 pub mod transaction;
@@ -534,18 +536,7 @@ impl InnerConnection {
             .and_then(|oid| self.get_type(oid))
             .collect()?;
 
-        let columns = match raw_columns {
-            Some(body) => body.fields()
-                .and_then(|field| {
-                    Ok(Column::new(
-                        field.name().to_owned(),
-                        self.get_type(field.type_oid())?,
-                    ))
-                })
-                .collect()?,
-            None => vec![],
-        };
-
+        let columns = self.parse_cols(raw_columns)?;
         Ok((param_types, columns))
     }
 
@@ -735,6 +726,22 @@ impl InnerConnection {
         Ok(ty)
     }
 
+
+    fn parse_cols(&mut self, raw: Option<backend::RowDescriptionBody>) -> Result<Vec<Column>> {
+        match raw {
+            Some(body) => body.fields()
+                .and_then(|field| {
+                    Ok(Column::new(
+                        field.name().to_owned(),
+                        self.get_type(field.type_oid())?,
+                    ))
+                })
+                .collect()
+                .map_err(From::from),
+            None => Ok(vec![]),
+        }
+    }
+
     fn setup_typeinfo_query(&mut self) -> Result<()> {
         if self.has_typeinfo_query {
             return Ok(());
@@ -919,6 +926,53 @@ impl InnerConnection {
         }
     }
 
+    fn simple_query_(&mut self, query: &str) -> Result<Vec<TextRows>> {
+        check_desync!(self);
+        debug!("executing query: {}", query);
+        self.stream.write_message(|buf| frontend::query(query, buf))?;
+        self.stream.flush()?;
+
+        let mut result = vec![];
+        let mut rows = vec![];
+        let mut columns = None;
+
+        loop {
+            match self.read_message()? {
+                backend::Message::ReadyForQuery(_) => break,
+                backend::Message::DataRow(body) => {
+                    rows.push(RowData::new(body)?);
+                }
+                backend::Message::CopyInResponse(_) => {
+                    self.stream.write_message(|buf| {
+                        frontend::copy_fail("COPY queries cannot be directly executed", buf)
+                    })?;
+                    self.stream.write_message(
+                        |buf| Ok::<(), io::Error>(frontend::sync(buf)),
+                    )?;
+                    self.stream.flush()?;
+                }
+                backend::Message::ErrorResponse(body) => {
+                    self.wait_for_ready()?;
+                    return Err(err(&mut body.fields()));
+                }
+                backend::Message::RowDescription(body) => {
+                    columns = Some(self.parse_cols(Some(body))?);
+                }
+                backend::Message::CommandComplete(_) => {
+                    result.push(
+                        TextRows::new(
+                            // TODO is it safe to unwrap here?
+                            columns.take().unwrap(),
+                            mem::replace(&mut rows, Vec::new())
+                        )
+                    );
+                }
+                _ => {}
+            }
+        }
+        Ok(result)
+    }
+
     fn quick_query(&mut self, query: &str) -> Result<Vec<Vec<Option<String>>>> {
         check_desync!(self);
         debug!("executing query: {}", query);
@@ -1295,6 +1349,14 @@ impl Connection {
         self.0.borrow_mut().quick_query(query).map(|_| ())
     }
 
+
+    /// Send a simple, non-prepared query
+    ///
+    /// TODO docs. Should this replace the batch_execute API?
+    pub fn simple_query(&self, query: &str) -> Result<Vec<TextRows>> {
+        self.0.borrow_mut().simple_query_(query)
+    }
+
     /// Returns a structure providing access to asynchronous notifications.
     ///
     /// Use the `LISTEN` command to register this connection for notifications.
@@ -1372,6 +1434,9 @@ pub trait GenericConnection {
 
     /// Like `Connection::is_active`.
     fn is_active(&self) -> bool;
+
+    /// Like `Connection::simple_query`.
+    fn simple_query(&self, query: &str) -> Result<Vec<TextRows>>;
 }
 
 impl GenericConnection for Connection {
@@ -1396,12 +1461,16 @@ impl GenericConnection for Connection {
     }
 
     fn batch_execute(&self, query: &str) -> Result<()> {
-        self.batch_execute(query)
+        self.batch_execute(&query)
     }
 
     fn is_active(&self) -> bool {
         self.is_active()
     }
+
+    fn simple_query(&self, query: &str) -> Result<Vec<TextRows>> {
+        self.simple_query(query)
+    }
 }
 
 impl<'a> GenericConnection for Transaction<'a> {
@@ -1426,7 +1495,11 @@ impl<'a> GenericConnection for Transaction<'a> {
     }
 
     fn batch_execute(&self, query: &str) -> Result<()> {
-        self.batch_execute(query)
+        self.batch_execute(&query)
+    }
+
+    fn simple_query(&self, query: &str) -> Result<Vec<TextRows>> {
+        self.simple_query(query)
     }
 
     fn is_active(&self) -> bool {
diff --git a/postgres/src/text_rows.rs b/postgres/src/text_rows.rs
@@ -0,0 +1,177 @@
+//! Query result rows.
+
+use postgres_shared::rows::RowData;
+use std::fmt;
+use std::slice;
+
+#[doc(inline)]
+pub use postgres_shared::rows::RowIndex;
+
+use stmt::{Column};
+
+/// The resulting rows of a query.
+pub struct TextRows {
+    columns: Vec<Column>,
+    data: Vec<RowData>,
+}
+
+impl fmt::Debug for TextRows {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        fmt.debug_struct("TextRows")
+            .field("columns", &self.columns())
+            .field("rows", &self.data.len())
+            .finish()
+    }
+}
+
+impl TextRows {
+    pub(crate) fn new(columns: Vec<Column>, data: Vec<RowData>) -> TextRows {
+        TextRows {
+            columns: columns,
+            data: data,
+        }
+    }
+
+    /// Returns a slice describing the columns of the `TextRows`.
+    pub fn columns(&self) -> &[Column] {
+        &self.columns[..]
+    }
+
+    /// Returns the number of rows present.
+    pub fn len(&self) -> usize {
+        self.data.len()
+    }
+
+    /// Determines if there are any rows present.
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    /// Returns a specific `TextRow`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `idx` is out of bounds.
+    pub fn get<'a>(&'a self, idx: usize) -> TextRow<'a> {
+        TextRow {
+            columns: &self.columns,
+            data: &self.data[idx],
+        }
+    }
+
+    /// Returns an iterator over the `TextRow`s.
+    pub fn iter<'a>(&'a self) -> Iter<'a> {
+        Iter {
+            columns: self.columns(),
+            iter: self.data.iter(),
+        }
+    }
+}
+
+impl<'a> IntoIterator for &'a TextRows {
+    type Item = TextRow<'a>;
+    type IntoIter = Iter<'a>;
+
+    fn into_iter(self) -> Iter<'a> {
+        self.iter()
+    }
+}
+
+/// An iterator over `TextRow`s.
+pub struct Iter<'a> {
+    columns: &'a [Column],
+    iter: slice::Iter<'a, RowData>,
+}
+
+impl<'a> Iterator for Iter<'a> {
+    type Item = TextRow<'a>;
+
+    fn next(&mut self) -> Option<TextRow<'a>> {
+        self.iter.next().map(|row| {
+            TextRow {
+                columns: self.columns,
+                data: row,
+            }
+        })
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.iter.size_hint()
+    }
+}
+
+impl<'a> DoubleEndedIterator for Iter<'a> {
+    fn next_back(&mut self) -> Option<TextRow<'a>> {
+        self.iter.next_back().map(|row| {
+            TextRow {
+                columns: self.columns,
+                data: row,
+            }
+        })
+    }
+}
+
+impl<'a> ExactSizeIterator for Iter<'a> {}
+
+/// A single result row of a query.
+pub struct TextRow<'a> {
+    columns: &'a [Column],
+    data: &'a RowData,
+}
+
+impl<'a> fmt::Debug for TextRow<'a> {
+    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
+        fmt.debug_struct("TextRow")
+            .field("columns", &self.columns)
+            .finish()
+    }
+}
+
+impl<'a> TextRow<'a> {
+    /// Returns the number of values in the row.
+    pub fn len(&self) -> usize {
+        self.data.len()
+    }
+
+    /// Determines if there are any values in the row.
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    /// Returns a slice describing the columns of the `TextRow`.
+    pub fn columns(&self) -> &[Column] {
+        self.columns
+    }
+
+    /// stub
+    pub fn get<I>(&self, idx: I) -> Option<&str>
+    where
+        I: RowIndex + fmt::Debug,
+    {
+        match self.get_inner(&idx) {
+            Some(value) => value,
+            None => panic!("no such column {:?}", idx),
+        }
+    }
+
+    /// stub
+    pub fn get_opt<I>(&self, idx: I) -> Option<Option<&str>>
+    where
+        I: RowIndex,
+    {
+        self.get_inner(&idx)
+    }
+
+    fn get_inner<I>(&self, idx: &I) -> Option<Option<&str>>
+    where
+        I: RowIndex,
+    {
+        let idx = match idx.__idx(self.columns) {
+            Some(idx) => idx,
+            None => return None,
+        };
+
+        // TODO can we assume these values will always be utf8?
+        Some(self.data.get(idx).map(|s| ::std::str::from_utf8(s).expect("utf8 encoded")))
+    }
+}
diff --git a/postgres/src/transaction.rs b/postgres/src/transaction.rs
@@ -4,6 +4,7 @@ use std::cell::Cell;
 use std::fmt;
 
 use rows::Rows;
+use text_rows::TextRows;
 use stmt::Statement;
 use types::ToSql;
 use {bad_response, Connection, Result};
@@ -229,6 +230,11 @@ impl<'conn> Transaction<'conn> {
         self.conn.batch_execute(query)
     }
 
+    /// Like `Connection::simple_query`.
+    pub fn simple_query(&self, query: &str) -> Result<Vec<TextRows>> {
+        self.conn.simple_query(query)
+    }
+
     /// Like `Connection::transaction`, but creates a nested transaction via
     /// a savepoint.
     ///
@@ -277,7 +283,7 @@ impl<'conn> Transaction<'conn> {
     pub fn set_config(&self, config: &Config) -> Result<()> {
         let mut command = "SET TRANSACTION".to_owned();
         config.build_command(&mut command);
-        self.batch_execute(&command)
+        self.batch_execute(&command).map(|_| ())
     }
 
     /// Determines if the transaction is currently set to commit or roll back.
