Add query_one convenience method (#559)

mlafeldt · web-flow · commit 4faa961e1b33 · 2025-07-31T16:39:19.000+02:00
Which works exactly like https://docs.rs/rusqlite/latest/rusqlite/struct.Statement.html#method.query_one
diff --git a/crates/duckdb/src/error.rs b/crates/duckdb/src/error.rs
@@ -47,6 +47,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),
@@ -96,6 +100,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)) => {
@@ -170,6 +175,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) => {
@@ -201,6 +207,7 @@ impl error::Error for Error {
             | Self::InvalidParameterName(_)
             | Self::ExecuteReturnedResults
             | Self::QueryReturnedNoRows
+            | Self::QueryReturnedMoreThanOneRow
             | Self::InvalidColumnIndex(_)
             | Self::InvalidColumnName(_)
             | Self::InvalidColumnType(..)
diff --git a/crates/duckdb/src/statement.rs b/crates/duckdb/src/statement.rs
@@ -349,6 +349,33 @@ impl Statement<'_> {
         self.query(params)?.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::OptionalExt::optional) on the result of
+    /// this to get a `Result<Option<T>>` (requires that the trait
+    /// `duckdb::OptionalExt` is imported).
+    ///
+    /// # Failure
+    ///
+    /// Will return `Err` if the underlying DuckDB 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)
+    }
+
     /// Return the row count
     #[inline]
     pub fn row_count(&self) -> usize {
@@ -635,9 +662,8 @@ mod test {
 
         let mut stmt = db.prepare("SELECT id FROM test where name = ?")?;
         {
-            let mut rows = stmt.query([&"one"])?;
-            let id: Result<i32> = rows.next()?.unwrap().get(0);
-            assert_eq!(Ok(1), id);
+            let id: i32 = stmt.query_one([&"one"], |r| r.get(0))?;
+            assert_eq!(id, 1);
         }
         Ok(())
     }
@@ -825,6 +851,71 @@ mod test {
         Ok(())
     }
 
+    #[test]
+    fn test_query_one() -> Result<()> {
+        let db = Connection::open_in_memory()?;
+        let sql = "BEGIN;
+                   CREATE TABLE foo(x INTEGER, y INTEGER);
+                   INSERT INTO foo VALUES(1, 3);
+                   INSERT INTO foo VALUES(2, 4);
+                   END;";
+        db.execute_batch(sql)?;
+
+        // Exactly one row
+        let y: i32 = db
+            .prepare("SELECT y FROM foo WHERE x = ?")?
+            .query_one([1], |r| r.get(0))?;
+        assert_eq!(y, 3);
+
+        // No rows
+        let res: Result<i32> = db
+            .prepare("SELECT y FROM foo WHERE x = ?")?
+            .query_one([99], |r| r.get(0));
+        assert_eq!(res.unwrap_err(), Error::QueryReturnedNoRows);
+
+        // Multiple rows
+        let res: Result<i32> = db.prepare("SELECT y FROM foo")?.query_one([], |r| r.get(0));
+        assert_eq!(res.unwrap_err(), Error::QueryReturnedMoreThanOneRow);
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_query_one_optional() -> Result<()> {
+        use crate::OptionalExt;
+
+        let db = Connection::open_in_memory()?;
+        let sql = "BEGIN;
+                   CREATE TABLE foo(x INTEGER, y INTEGER);
+                   INSERT INTO foo VALUES(1, 3);
+                   INSERT INTO foo VALUES(2, 4);
+                   END;";
+        db.execute_batch(sql)?;
+
+        // Exactly one row
+        let y: Option<i32> = db
+            .prepare("SELECT y FROM foo WHERE x = ?")?
+            .query_one([1], |r| r.get(0))
+            .optional()?;
+        assert_eq!(y, Some(3));
+
+        // No rows
+        let y: Option<i32> = db
+            .prepare("SELECT y FROM foo WHERE x = ?")?
+            .query_one([99], |r| r.get(0))
+            .optional()?;
+        assert_eq!(y, None);
+
+        // Multiple rows - should still return error (not converted by optional)
+        let res = db
+            .prepare("SELECT y FROM foo")?
+            .query_one([], |r| r.get::<_, i32>(0))
+            .optional();
+        assert_eq!(res.unwrap_err(), Error::QueryReturnedMoreThanOneRow);
+
+        Ok(())
+    }
+
     #[test]
     fn test_query_by_column_name() -> Result<()> {
         let db = Connection::open_in_memory()?;
diff --git a/crates/duckdb/src/types/mod.rs b/crates/duckdb/src/types/mod.rs
@@ -312,7 +312,6 @@ mod test {
 
         let mut stmt = db.prepare("SELECT b, t, i, f, n FROM foo")?;
         let mut rows = stmt.query([])?;
-
         let row = rows.next()?.unwrap();
 
         // check the correct types come back as expected
@@ -371,7 +370,6 @@ mod test {
 
         let mut stmt = db.prepare("SELECT b, t, i, f, n FROM foo")?;
         let mut rows = stmt.query([])?;
-
         let row = rows.next()?.unwrap();
         // NOTE: this is different from SQLite
         // assert_eq!(Value::Blob(vec![1, 2]), row.get::<_, Value>(0)?);
diff --git a/crates/duckdb/src/types/to_sql.rs b/crates/duckdb/src/types/to_sql.rs
@@ -316,10 +316,7 @@ mod test {
 
         db.execute("INSERT INTO foo (id) VALUES (gen_random_uuid())", [])?;
 
-        let mut stmt = db.prepare("SELECT id FROM foo")?;
-        let mut rows = stmt.query([])?;
-        let row = rows.next()?.unwrap();
-        let found_id: String = row.get_unwrap(0);
+        let found_id: String = db.prepare("SELECT id FROM foo")?.query_one([], |r| r.get(0))?;
         assert_eq!(found_id.len(), 36);
         Ok(())
     }
@@ -337,11 +334,9 @@ mod test {
         let id_vec = id.as_bytes().to_vec();
         db.execute("INSERT INTO foo (id, label) VALUES (?, ?)", params![id_vec, "target"])?;
 
-        let mut stmt = db.prepare("SELECT id, label FROM foo WHERE id = ?")?;
-        let mut rows = stmt.query(params![id_vec])?;
-        let row = rows.next()?.unwrap();
-        let found_id: Uuid = row.get_unwrap(0);
-        let found_label: String = row.get_unwrap(1);
+        let (found_id, found_label): (Uuid, String) = db
+            .prepare("SELECT id, label FROM foo WHERE id = ?")?
+            .query_one(params![id_vec], |r| Ok((r.get_unwrap(0), r.get_unwrap(1))))?;
         assert_eq!(found_id, id);
         assert_eq!(found_label, "target");
         Ok(())
@@ -359,11 +354,9 @@ mod test {
         let id = Uuid::new_v4();
         db.execute("INSERT INTO foo (id, label) VALUES (?, ?)", params![id, "target"])?;
 
-        let mut stmt = db.prepare("SELECT id, label FROM foo WHERE id = ?")?;
-        let mut rows = stmt.query(params![id])?;
-        let row = rows.next()?.unwrap();
-        let found_id: Uuid = row.get_unwrap(0);
-        let found_label: String = row.get_unwrap(1);
+        let (found_id, found_label): (Uuid, String) = db
+            .prepare("SELECT id, label FROM foo WHERE id = ?")?
+            .query_one(params![id], |r| Ok((r.get_unwrap(0), r.get_unwrap(1))))?;
         assert_eq!(found_id, id);
         assert_eq!(found_label, "target");
         Ok(())
