Add query_one

Author: sfackler

postgres/src/client.rs

@@ -96,8 +96,6 @@ impl Client {
     /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
     /// with the `prepare` method.
     ///
-    /// The `query_iter` method can be used to avoid buffering all rows in memory at once.
-    ///
     /// # Panics
     ///
     /// Panics if the number of parameters provided does not match the number expected.
@@ -125,6 +123,41 @@ impl Client {
         executor::block_on(self.0.query(query, params))
     }
 
+    /// Executes a statement which returns a single row, returning it.
+    ///
+    /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
+    /// provided, 1-indexed.
+    ///
+    /// The `query` argument can either be a `Statement`, or a raw query string. If the same statement will be
+    /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
+    /// with the `prepare` method.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the number of parameters provided does not match the number expected.
+    ///
+    /// # Examples
+    ///
+    /// ```no_run
+    /// use postgres::{Client, NoTls};
+    ///
+    /// # fn main() -> Result<(), postgres::Error> {
+    /// let mut client = Client::connect("host=localhost user=postgres", NoTls)?;
+    ///
+    /// let baz = true;
+    /// let row = client.query_one("SELECT foo FROM bar WHERE baz = $1", &[&baz])?;
+    /// let foo: i32 = row.get("foo");
+    /// println!("foo: {}", foo);
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn query_one<T>(&mut self, query: &T, params: &[&(dyn ToSql + Sync)]) -> Result<Row, Error>
+    where
+        T: ?Sized + ToStatement,
+    {
+        executor::block_on(self.0.query_one(query, params))
+    }
+
     /// A maximally-flexible version of `query`.
     ///
     /// It takes an iterator of parameters rather than a slice, and returns an iterator of rows rather than collecting

tokio-postgres/src/client.rs

@@ -9,6 +9,7 @@ use crate::query::RowStream;
 use crate::slice_iter;
 #[cfg(feature = "runtime")]
 use crate::tls::MakeTlsConnect;
+use pin_utils::pin_mut;
 use crate::tls::TlsConnect;
 use crate::to_statement::ToStatement;
 use crate::types::{Oid, ToSql, Type};
@@ -195,6 +196,13 @@ impl Client {
 
     /// Executes a statement, returning a vector of the resulting rows.
     ///
+    /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
+    /// provided, 1-indexed.
+    ///
+    /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
+    /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
+    /// with the `prepare` method.
+    ///
     /// # Panics
     ///
     /// Panics if the number of parameters provided does not match the number expected.
@@ -212,8 +220,52 @@ impl Client {
             .await
     }
 
+    /// Executes a statement which returns a single row, returning it.
+    ///
+    /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
+    /// provided, 1-indexed.
+    ///
+    /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
+    /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
+    /// with the `prepare` method.
+    ///
+    /// Returns an error if the query does not return exactly one row.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the number of parameters provided does not match the number expected.
+    pub async fn query_one<T>(
+        &self,
+        statement: &T,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<Row, Error>
+    where
+        T: ?Sized + ToStatement
+    {
+        let stream = self.query_raw(statement, slice_iter(params)).await?;
+        pin_mut!(stream);
+
+        let row = match stream.try_next().await? {
+            Some(row) => row,
+            None => return Err(Error::row_count()),
+        };
+
+        if stream.try_next().await?.is_some() {
+            return Err(Error::row_count());
+        }
+
+        Ok(row)
+    }
+
     /// The maximally flexible version of [`query`].
     ///
+    /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
+    /// provided, 1-indexed.
+    ///
+    /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
+    /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
+    /// with the `prepare` method.
+    ///
     /// # Panics
     ///
     /// Panics if the number of parameters provided does not match the number expected.
@@ -231,6 +283,13 @@ impl Client {
 
     /// Executes a statement, returning the number of rows modified.
     ///
+    /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
+    /// provided, 1-indexed.
+    ///
+    /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
+    /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
+    /// with the `prepare` method.
+    ///
     /// If the statement does not modify any rows (e.g. `SELECT`), 0 is returned.
     ///
     /// # Panics
@@ -249,6 +308,13 @@ impl Client {
 
     /// The maximally flexible version of [`execute`].
     ///
+    /// A statement may contain parameters, specified by `$n`, where `n` is the index of the parameter of the list
+    /// provided, 1-indexed.
+    ///
+    /// The `statement` argument can either be a `Statement`, or a raw query string. If the same statement will be
+    /// repeatedly executed (perhaps with different query parameters), consider preparing the statement up front
+    /// with the `prepare` method.
+    ///
     /// # Panics
     ///
     /// Panics if the number of parameters provided does not match the number expected.

tokio-postgres/src/error/mod.rs

@@ -345,6 +345,7 @@ enum Kind {
     Authentication,
     ConfigParse,
     Config,
+    RowCount,
     #[cfg(feature = "runtime")]
     Connect,
 }
@@ -383,6 +384,7 @@ impl fmt::Display for Error {
             Kind::Authentication => fmt.write_str("authentication error")?,
             Kind::ConfigParse => fmt.write_str("invalid connection string")?,
             Kind::Config => fmt.write_str("invalid configuration")?,
+            Kind::RowCount => fmt.write_str("query returned an unexpected number of rows")?,
             #[cfg(feature = "runtime")]
             Kind::Connect => fmt.write_str("error connecting to server")?,
         };
@@ -483,6 +485,10 @@ impl Error {
         Error::new(Kind::Config, Some(e))
     }
 
+    pub(crate) fn row_count() -> Error {
+        Error::new(Kind::RowCount, None)
+    }
+
     #[cfg(feature = "runtime")]
     pub(crate) fn connect(e: io::Error) -> Error {
         Error::new(Kind::Connect, Some(Box::new(e)))

tokio-postgres/src/transaction.rs

@@ -106,6 +106,18 @@ impl<'a> Transaction<'a> {
         self.client.query(statement, params).await
     }
 
+    /// Like `Client::query_one`.
+    pub async fn query_one<T>(
+        &self,
+        statement: &T,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<Row, Error>
+        where
+            T: ?Sized + ToStatement,
+    {
+        self.client.query_one(statement, params).await
+    }
+
     /// Like `Client::query_raw`.
     pub async fn query_raw<'b, T, I>(&self, statement: &T, params: I) -> Result<RowStream, Error>
     where

tokio-postgres/tests/test/main.rs

@@ -642,3 +642,19 @@ async fn check_send() {
     is_send(&f);
     drop(f);
 }
+
+#[tokio::test]
+async fn query_one() {
+    let client = connect("user=postgres").await;
+
+    client.batch_execute("
+        CREATE TEMPORARY TABLE foo (
+            name TEXT
+        );
+        INSERT INTO foo (name) VALUES ('alice'), ('bob'), ('carol');
+    ").await.unwrap();
+
+    client.query_one("SELECT * FROM foo WHERE name = 'dave'", &[]).await.err().unwrap();
+    client.query_one("SELECT * FROM foo WHERE name = 'alice'", &[]).await.unwrap();
+    client.query_one("SELECT * FROM foo", &[]).await.err().unwrap();
+}