Add query_one

gwenn · gwenn · commit ae0e2fda874c · 2025-05-26T07:48:04.000+02:00
diff --git a/src/error.rs b/src/error.rs
@@ -50,6 +50,10 @@ pub enum Error {
     /// for [`query_row`](crate::Connection::query_row)) did not return any.
     QueryReturnedNoRows,
 
+    /// Error when a query that was expected to return only one row (e.g.,
+    /// for [`query_one`](crate::Connection::query_one)) did return more than one.
+    QueryReturnedMoreThanOneRow,
+
     /// Error when the value of a particular column is requested, but the index
     /// is out of range for the statement.
     InvalidColumnIndex(usize),
@@ -152,6 +156,7 @@ impl PartialEq for Error {
             (Self::InvalidPath(p1), Self::InvalidPath(p2)) => p1 == p2,
             (Self::ExecuteReturnedResults, Self::ExecuteReturnedResults) => true,
             (Self::QueryReturnedNoRows, Self::QueryReturnedNoRows) => true,
+            (Self::QueryReturnedMoreThanOneRow, Self::QueryReturnedMoreThanOneRow) => true,
             (Self::InvalidColumnIndex(i1), Self::InvalidColumnIndex(i2)) => i1 == i2,
             (Self::InvalidColumnName(n1), Self::InvalidColumnName(n2)) => n1 == n2,
             (Self::InvalidColumnType(i1, n1, t1), Self::InvalidColumnType(i2, n2, t2)) => {
@@ -278,6 +283,7 @@ impl fmt::Display for Error {
                 write!(f, "Execute returned results - did you mean to call query?")
             }
             Self::QueryReturnedNoRows => write!(f, "Query returned no rows"),
+            Self::QueryReturnedMoreThanOneRow => write!(f, "Query returned more than one row"),
             Self::InvalidColumnIndex(i) => write!(f, "Invalid column index: {i}"),
             Self::InvalidColumnName(ref name) => write!(f, "Invalid column name: {name}"),
             Self::InvalidColumnType(i, ref name, ref t) => {
@@ -336,6 +342,7 @@ impl error::Error for Error {
             | Self::InvalidParameterName(_)
             | Self::ExecuteReturnedResults
             | Self::QueryReturnedNoRows
+            | Self::QueryReturnedMoreThanOneRow
             | Self::InvalidColumnIndex(_)
             | Self::InvalidColumnName(_)
             | Self::InvalidColumnType(..)
diff --git a/src/functions.rs b/src/functions.rs
@@ -1005,7 +1005,7 @@ mod test {
         })?;
 
         let res: bool =
-            db.one_column("SELECT example(0, i) FROM (SELECT 0 as i UNION SELECT 1)")?;
+            db.first_column("SELECT example(0, i) FROM (SELECT 0 as i UNION SELECT 1)")?;
         // Doesn't actually matter, we'll assert in the function if there's a problem.
         assert!(res);
         Ok(())
diff --git a/src/lib.rs b/src/lib.rs
@@ -665,9 +665,36 @@ impl Connection {
         stmt.query_row(params, f)
     }
 
+    /// Convenience method to execute a query that is expected to return exactly
+    /// one row.
+    ///
+    /// Returns `Err(QueryReturnedMoreThanOneRow)` if the query returns more than one row.
+    ///
+    /// Returns `Err(QueryReturnedNoRows)` if no results are returned. If the
+    /// query truly is optional, you can call
+    /// [`.optional()`](crate::OptionalExtension::optional) on the result of
+    /// this to get a `Result<Option<T>>` (requires that the trait
+    /// `rusqlite::OptionalExtension` is imported).
+    ///
+    /// # Failure
+    ///
+    /// Will return `Err` if the underlying SQLite call fails.
+    pub fn query_one<T, P, F>(&self, sql: &str, params: P, f: F) -> Result<T>
+    where
+        P: Params,
+        F: FnOnce(&Row<'_>) -> Result<T>,
+    {
+        let mut stmt = self.prepare(sql)?;
+        stmt.query_one(params, f)
+    }
+
     // https://sqlite.org/tclsqlite.html#onecolumn
     #[cfg(test)]
     pub(crate) fn one_column<T: types::FromSql>(&self, sql: &str) -> Result<T> {
+        self.query_one(sql, [], |r| r.get(0))
+    }
+    #[cfg(test)]
+    pub(crate) fn first_column<T: types::FromSql>(&self, sql: &str) -> Result<T> {
         self.query_row(sql, [], |r| r.get(0))
     }
 
diff --git a/src/statement.rs b/src/statement.rs
@@ -381,6 +381,33 @@ impl Statement<'_> {
         rows.get_expected_row().and_then(f)
     }
 
+    /// Convenience method to execute a query that is expected to return exactly
+    /// one row.
+    ///
+    /// Returns `Err(QueryReturnedMoreThanOneRow)` if the query returns more than one row.
+    ///
+    /// Returns `Err(QueryReturnedNoRows)` if no results are returned. If the
+    /// query truly is optional, you can call
+    /// [`.optional()`](crate::OptionalExtension::optional) on the result of
+    /// this to get a `Result<Option<T>>` (requires that the trait
+    /// `rusqlite::OptionalExtension` is imported).
+    ///
+    /// # Failure
+    ///
+    /// Will return `Err` if the underlying SQLite call fails.
+    pub fn query_one<T, P, F>(&mut self, params: P, f: F) -> Result<T>
+    where
+        P: Params,
+        F: FnOnce(&Row<'_>) -> Result<T>,
+    {
+        let mut rows = self.query(params)?;
+        let row = rows.get_expected_row().and_then(f)?;
+        if rows.next()?.is_some() {
+            return Err(Error::QueryReturnedMoreThanOneRow);
+        }
+        Ok(row)
+    }
+
     /// Consumes the statement.
     ///
     /// Functionally equivalent to the `Drop` implementation, but allows
@@ -1170,6 +1197,22 @@ mod test {
         Ok(())
     }
 
+    #[test]
+    fn query_one() -> Result<()> {
+        let db = Connection::open_in_memory()?;
+        db.execute_batch("CREATE TABLE foo(x INTEGER, y INTEGER);")?;
+        let mut stmt = db.prepare("SELECT y FROM foo WHERE x = ?1")?;
+        let y: Result<i64> = stmt.query_one([1i32], |r| r.get(0));
+        assert_eq!(Error::QueryReturnedNoRows, y.unwrap_err());
+        db.execute_batch("INSERT INTO foo VALUES(1, 3);")?;
+        let y: Result<i64> = stmt.query_one([1i32], |r| r.get(0));
+        assert_eq!(3i64, y?);
+        db.execute_batch("INSERT INTO foo VALUES(1, 3);")?;
+        let y: Result<i64> = stmt.query_one([1i32], |r| r.get(0));
+        assert_eq!(Error::QueryReturnedMoreThanOneRow, y.unwrap_err());
+        Ok(())
+    }
+
     #[test]
     fn test_query_by_column_name() -> Result<()> {
         let db = Connection::open_in_memory()?;
