Support scalar queries.

tokio-postgres/src/client.rs

@@ -22,7 +22,7 @@ use futures_channel::mpsc;
 use futures_util::{future, pin_mut, ready, StreamExt, TryStreamExt};
 use parking_lot::Mutex;
 use postgres_protocol::message::{backend::Message, frontend};
-use postgres_types::BorrowToSql;
+use postgres_types::{BorrowToSql, FromSqlOwned};
 use std::collections::HashMap;
 use std::fmt;
 #[cfg(feature = "runtime")]
@@ -256,6 +256,16 @@ impl Client {
             .await
     }
 
+    /// Like [`Client::query`] but returns a vector of scalars.
+    pub async fn query_scalar<T: FromSqlOwned>(
+        &self,
+        sql: &str,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<Vec<T>, Error> {
+        let rows = self.query(sql, params).await?;
+        rows.into_iter().map(|r| r.try_get(0)).collect()
+    }
+
     /// Executes a statement which returns a single row, returning it.
     ///
     /// Returns an error if the query does not return exactly one row.
@@ -279,6 +289,16 @@ impl Client {
             .and_then(|res| res.ok_or_else(Error::row_count))
     }
 
+    /// Like [`Client::query_one`] but returns one scalar.
+    pub async fn query_one_scalar<T: FromSqlOwned>(
+        &self,
+        sql: &str,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<T, Error> {
+        let row = self.query_one(sql, params).await?;
+        row.try_get(0)
+    }
+
     /// Executes a statements which returns zero or one rows, returning it.
     ///
     /// Returns an error if the query returns more than one row.
@@ -318,6 +338,16 @@ impl Client {
         Ok(first)
     }
 
+    /// Like [`Client::query_opt`] but returns an optional scalar.
+    pub async fn query_opt_scalar<S: FromSqlOwned>(
+        &self,
+        sql: &str,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<Option<S>, Error> {
+        let row = self.query_opt(sql, params).await?;
+        row.map(|x| (x.try_get::<_, S>(0))).transpose()
+    }
+
     /// 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

tokio-postgres/src/transaction.rs

@@ -15,6 +15,7 @@ use crate::{
 use bytes::Buf;
 use futures_util::TryStreamExt;
 use postgres_protocol::message::frontend;
+use postgres_types::FromSqlOwned;
 use tokio::io::{AsyncRead, AsyncWrite};
 
 /// A representation of a PostgreSQL database transaction.
@@ -114,7 +115,16 @@ impl<'a> Transaction<'a> {
         self.client.query(statement, params).await
     }
 
-    /// Like `Client::query_one`.
+    /// Like [`Client::query_scalar`].
+    pub async fn query_scalar<T: FromSqlOwned>(
+        &self,
+        sql: &str,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<Vec<T>, Error> {
+        self.client.query_scalar(sql, params).await
+    }
+
+    /// Like [`Client::query_one`].
     pub async fn query_one<T>(
         &self,
         statement: &T,
@@ -126,7 +136,16 @@ impl<'a> Transaction<'a> {
         self.client.query_one(statement, params).await
     }
 
-    /// Like `Client::query_opt`.
+    /// Like [`Client::query_one_scalar`].
+    pub async fn query_one_scalar<T: FromSqlOwned>(
+        &self,
+        sql: &str,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<T, Error> {
+        self.client.query_one_scalar(sql, params).await
+    }
+
+    /// Like [`Client::query_opt`].
     pub async fn query_opt<T>(
         &self,
         statement: &T,
@@ -138,7 +157,16 @@ impl<'a> Transaction<'a> {
         self.client.query_opt(statement, params).await
     }
 
-    /// Like `Client::query_raw`.
+    /// Like [`Client::query_opt_scalar`].
+    pub async fn query_opt_scalar<S: FromSqlOwned>(
+        &self,
+        sql: &str,
+        params: &[&(dyn ToSql + Sync)],
+    ) -> Result<Option<S>, Error> {
+        self.client.query_opt_scalar(sql, params).await
+    }
+
+    /// Like [`Client::query_raw`]
     pub async fn query_raw<T, P, I>(&self, statement: &T, params: I) -> Result<RowStream, Error>
     where
         T: ?Sized + ToStatement,

tokio-postgres/tests/test/main.rs

@@ -952,3 +952,49 @@ async fn deferred_constraint() {
         .await
         .unwrap_err();
 }
+
+#[tokio::test]
+async fn query_opt_scalar() {
+    let client = connect("user=postgres").await;
+    client
+        .batch_execute(
+            "CREATE TEMPORARY TABLE person (
+                id serial,
+                name text NOT NULL,
+                age integer
+            );
+            INSERT INTO person (name, age) VALUES ('steven', 18);
+            INSERT INTO person (name, age) VALUES ('fred', NULL);
+            ",
+        )
+        .await
+        .unwrap();
+
+    let age: Option<i32> = client
+        .query_opt_scalar("SELECT age FROM person WHERE name = $1", &[&"steven"])
+        .await
+        .unwrap();
+
+    assert_eq!(age, Some(18));
+
+    let age: Option<Option<i32>> = client
+        .query_opt_scalar("SELECT age FROM person WHERE name = $1", &[&"fred"])
+        .await
+        .unwrap();
+
+    assert_eq!(age, Some(None));
+
+    let age: Option<Option<i32>> = client
+        .query_opt_scalar("SELECT age FROM person WHERE name = $1", &[&"steven"])
+        .await
+        .unwrap();
+
+    assert_eq!(age, Some(Some(18)));
+
+    let age: Option<Option<i32>> = client
+        .query_opt_scalar("SELECT age FROM person WHERE name = $1", &[&"barney"])
+        .await
+        .unwrap();
+
+    assert_eq!(age, None);
+}