Skip to content

Commit 4faa961

Browse files
mlafeldtmlafeldt
authored
Add query_one convenience method (#559)
Which works exactly like https://docs.rs/rusqlite/latest/rusqlite/struct.Statement.html#method.query_one
2 parents 2b5c932 + 9f2daa4 commit 4faa961

File tree

4 files changed

+108
-19
lines changed

4 files changed

+108
-19
lines changed

‎crates/duckdb/src/error.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,10 @@ pub enum Error {
4747
/// for [`query_row`](crate::Connection::query_row)) did not return any.
4848
QueryReturnedNoRows,
4949

50+
/// Error when a query that was expected to return only one row (e.g.,
51+
/// for [`query_one`](crate::Connection::query_one)) did return more than one.
52+
QueryReturnedMoreThanOneRow,
53+
5054
/// Error when the value of a particular column is requested, but the index
5155
/// is out of range for the statement.
5256
InvalidColumnIndex(usize),
@@ -96,6 +100,7 @@ impl PartialEq for Error {
96100
(Self::InvalidPath(p1), Self::InvalidPath(p2)) => p1 == p2,
97101
(Self::ExecuteReturnedResults, Self::ExecuteReturnedResults) => true,
98102
(Self::QueryReturnedNoRows, Self::QueryReturnedNoRows) => true,
103+
(Self::QueryReturnedMoreThanOneRow, Self::QueryReturnedMoreThanOneRow) => true,
99104
(Self::InvalidColumnIndex(i1), Self::InvalidColumnIndex(i2)) => i1 == i2,
100105
(Self::InvalidColumnName(n1), Self::InvalidColumnName(n2)) => n1 == n2,
101106
(Self::InvalidColumnType(i1, n1, t1), Self::InvalidColumnType(i2, n2, t2)) => {
@@ -170,6 +175,7 @@ impl fmt::Display for Error {
170175
write!(f, "Execute returned results - did you mean to call query?")
171176
}
172177
Self::QueryReturnedNoRows => write!(f, "Query returned no rows"),
178+
Self::QueryReturnedMoreThanOneRow => write!(f, "Query returned more than one row"),
173179
Self::InvalidColumnIndex(i) => write!(f, "Invalid column index: {i}"),
174180
Self::InvalidColumnName(ref name) => write!(f, "Invalid column name: {name}"),
175181
Self::InvalidColumnType(i, ref name, ref t) => {
@@ -201,6 +207,7 @@ impl error::Error for Error {
201207
| Self::InvalidParameterName(_)
202208
| Self::ExecuteReturnedResults
203209
| Self::QueryReturnedNoRows
210+
| Self::QueryReturnedMoreThanOneRow
204211
| Self::InvalidColumnIndex(_)
205212
| Self::InvalidColumnName(_)
206213
| Self::InvalidColumnType(..)

‎crates/duckdb/src/statement.rs

Lines changed: 94 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -349,6 +349,33 @@ impl Statement<'_> {
349349
self.query(params)?.get_expected_row().and_then(f)
350350
}
351351

352+
/// Convenience method to execute a query that is expected to return exactly
353+
/// one row.
354+
///
355+
/// Returns `Err(QueryReturnedMoreThanOneRow)` if the query returns more than one row.
356+
///
357+
/// Returns `Err(QueryReturnedNoRows)` if no results are returned. If the
358+
/// query truly is optional, you can call
359+
/// [`.optional()`](crate::OptionalExt::optional) on the result of
360+
/// this to get a `Result<Option<T>>` (requires that the trait
361+
/// `duckdb::OptionalExt` is imported).
362+
///
363+
/// # Failure
364+
///
365+
/// Will return `Err` if the underlying DuckDB call fails.
366+
pub fn query_one<T, P, F>(&mut self, params: P, f: F) -> Result<T>
367+
where
368+
P: Params,
369+
F: FnOnce(&Row<'_>) -> Result<T>,
370+
{
371+
let mut rows = self.query(params)?;
372+
let row = rows.get_expected_row().and_then(f)?;
373+
if rows.next()?.is_some() {
374+
return Err(Error::QueryReturnedMoreThanOneRow);
375+
}
376+
Ok(row)
377+
}
378+
352379
/// Return the row count
353380
#[inline]
354381
pub fn row_count(&self) -> usize {
@@ -635,9 +662,8 @@ mod test {
635662

636663
let mut stmt = db.prepare("SELECT id FROM test where name = ?")?;
637664
{
638-
let mut rows = stmt.query([&"one"])?;
639-
let id: Result<i32> = rows.next()?.unwrap().get(0);
640-
assert_eq!(Ok(1), id);
665+
let id: i32 = stmt.query_one([&"one"], |r| r.get(0))?;
666+
assert_eq!(id, 1);
641667
}
642668
Ok(())
643669
}
@@ -825,6 +851,71 @@ mod test {
825851
Ok(())
826852
}
827853

854+
#[test]
855+
fn test_query_one() -> Result<()> {
856+
let db = Connection::open_in_memory()?;
857+
let sql = "BEGIN;
858+
CREATE TABLE foo(x INTEGER, y INTEGER);
859+
INSERT INTO foo VALUES(1, 3);
860+
INSERT INTO foo VALUES(2, 4);
861+
END;";
862+
db.execute_batch(sql)?;
863+
864+
// Exactly one row
865+
let y: i32 = db
866+
.prepare("SELECT y FROM foo WHERE x = ?")?
867+
.query_one([1], |r| r.get(0))?;
868+
assert_eq!(y, 3);
869+
870+
// No rows
871+
let res: Result<i32> = db
872+
.prepare("SELECT y FROM foo WHERE x = ?")?
873+
.query_one([99], |r| r.get(0));
874+
assert_eq!(res.unwrap_err(), Error::QueryReturnedNoRows);
875+
876+
// Multiple rows
877+
let res: Result<i32> = db.prepare("SELECT y FROM foo")?.query_one([], |r| r.get(0));
878+
assert_eq!(res.unwrap_err(), Error::QueryReturnedMoreThanOneRow);
879+
880+
Ok(())
881+
}
882+
883+
#[test]
884+
fn test_query_one_optional() -> Result<()> {
885+
use crate::OptionalExt;
886+
887+
let db = Connection::open_in_memory()?;
888+
let sql = "BEGIN;
889+
CREATE TABLE foo(x INTEGER, y INTEGER);
890+
INSERT INTO foo VALUES(1, 3);
891+
INSERT INTO foo VALUES(2, 4);
892+
END;";
893+
db.execute_batch(sql)?;
894+
895+
// Exactly one row
896+
let y: Option<i32> = db
897+
.prepare("SELECT y FROM foo WHERE x = ?")?
898+
.query_one([1], |r| r.get(0))
899+
.optional()?;
900+
assert_eq!(y, Some(3));
901+
902+
// No rows
903+
let y: Option<i32> = db
904+
.prepare("SELECT y FROM foo WHERE x = ?")?
905+
.query_one([99], |r| r.get(0))
906+
.optional()?;
907+
assert_eq!(y, None);
908+
909+
// Multiple rows - should still return error (not converted by optional)
910+
let res = db
911+
.prepare("SELECT y FROM foo")?
912+
.query_one([], |r| r.get::<_, i32>(0))
913+
.optional();
914+
assert_eq!(res.unwrap_err(), Error::QueryReturnedMoreThanOneRow);
915+
916+
Ok(())
917+
}
918+
828919
#[test]
829920
fn test_query_by_column_name() -> Result<()> {
830921
let db = Connection::open_in_memory()?;

‎crates/duckdb/src/types/mod.rs

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -312,7 +312,6 @@ mod test {
312312

313313
let mut stmt = db.prepare("SELECT b, t, i, f, n FROM foo")?;
314314
let mut rows = stmt.query([])?;
315-
316315
let row = rows.next()?.unwrap();
317316

318317
// check the correct types come back as expected
@@ -371,7 +370,6 @@ mod test {
371370

372371
let mut stmt = db.prepare("SELECT b, t, i, f, n FROM foo")?;
373372
let mut rows = stmt.query([])?;
374-
375373
let row = rows.next()?.unwrap();
376374
// NOTE: this is different from SQLite
377375
// assert_eq!(Value::Blob(vec![1, 2]), row.get::<_, Value>(0)?);

‎crates/duckdb/src/types/to_sql.rs

Lines changed: 7 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -316,10 +316,7 @@ mod test {
316316

317317
db.execute("INSERT INTO foo (id) VALUES (gen_random_uuid())", [])?;
318318

319-
let mut stmt = db.prepare("SELECT id FROM foo")?;
320-
let mut rows = stmt.query([])?;
321-
let row = rows.next()?.unwrap();
322-
let found_id: String = row.get_unwrap(0);
319+
let found_id: String = db.prepare("SELECT id FROM foo")?.query_one([], |r| r.get(0))?;
323320
assert_eq!(found_id.len(), 36);
324321
Ok(())
325322
}
@@ -337,11 +334,9 @@ mod test {
337334
let id_vec = id.as_bytes().to_vec();
338335
db.execute("INSERT INTO foo (id, label) VALUES (?, ?)", params![id_vec, "target"])?;
339336

340-
let mut stmt = db.prepare("SELECT id, label FROM foo WHERE id = ?")?;
341-
let mut rows = stmt.query(params![id_vec])?;
342-
let row = rows.next()?.unwrap();
343-
let found_id: Uuid = row.get_unwrap(0);
344-
let found_label: String = row.get_unwrap(1);
337+
let (found_id, found_label): (Uuid, String) = db
338+
.prepare("SELECT id, label FROM foo WHERE id = ?")?
339+
.query_one(params![id_vec], |r| Ok((r.get_unwrap(0), r.get_unwrap(1))))?;
345340
assert_eq!(found_id, id);
346341
assert_eq!(found_label, "target");
347342
Ok(())
@@ -359,11 +354,9 @@ mod test {
359354
let id = Uuid::new_v4();
360355
db.execute("INSERT INTO foo (id, label) VALUES (?, ?)", params![id, "target"])?;
361356

362-
let mut stmt = db.prepare("SELECT id, label FROM foo WHERE id = ?")?;
363-
let mut rows = stmt.query(params![id])?;
364-
let row = rows.next()?.unwrap();
365-
let found_id: Uuid = row.get_unwrap(0);
366-
let found_label: String = row.get_unwrap(1);
357+
let (found_id, found_label): (Uuid, String) = db
358+
.prepare("SELECT id, label FROM foo WHERE id = ?")?
359+
.query_one(params![id], |r| Ok((r.get_unwrap(0), r.get_unwrap(1))))?;
367360
assert_eq!(found_id, id);
368361
assert_eq!(found_label, "target");
369362
Ok(())

0 commit comments

Comments
 (0)