diff --git a/bin/drk/src/walletdb.rs b/bin/drk/src/walletdb.rs index fa03647fe..fbbdbaf57 100644 --- a/bin/drk/src/walletdb.rs +++ b/bin/drk/src/walletdb.rs @@ -245,7 +245,7 @@ impl WalletDb { ) -> WalletDbResult>> { // Generate `SELECT` query let query = self.generate_select_query(table, col_names, params); - debug!(target: "walletdb::multiple", "[WalletDb] Executing SQL query:\n{query}"); + debug!(target: "walletdb::query_multiple", "[WalletDb] Executing SQL query:\n{query}"); // First we prepare the query let Ok(conn) = self.conn.lock() else { return Err(WalletDbError::FailedToAquireLock) }; @@ -296,6 +296,54 @@ impl WalletDb { Ok(result) } + /// Query provided table using provided query for multiple rows. + pub fn query_custom( + &self, + query: &str, + params: &[&dyn ToSql], + ) -> WalletDbResult>> { + debug!(target: "walletdb::query_custom", "[WalletDb] Executing SQL query:\n{query}"); + + // First we prepare the query + let Ok(conn) = self.conn.lock() else { return Err(WalletDbError::FailedToAquireLock) }; + let Ok(mut stmt) = conn.prepare(query) else { + return Err(WalletDbError::QueryPreparationFailed) + }; + + // Execute the query using provided converted params + let Ok(mut rows) = stmt.query(params) else { + return Err(WalletDbError::QueryExecutionFailed) + }; + + // Loop over returned rows and parse them + let mut result = vec![]; + loop { + // Check if an error occured + let row = match rows.next() { + Ok(r) => r, + Err(_) => return Err(WalletDbError::QueryExecutionFailed), + }; + + // Check if no row was returned + let row = match row { + Some(r) => r, + None => break, + }; + + // Grab row returned values + let mut row_values = vec![]; + let mut idx = 0; + loop { + let Ok(value) = row.get(idx) else { break }; + row_values.push(value); + idx += 1; + } + result.push(row_values); + } + + Ok(result) + } + /// Auxiliary function to store provided inverse query into our cache. pub fn cache_inverse(&self, query: String) -> WalletDbResult<()> { debug!(target: "walletdb::cache_inverse", "[WalletDb] Storing query:\n{query}"); @@ -558,6 +606,12 @@ mod tests { assert_eq!(ret.len(), 1); let numba: i64 = if let Value::Integer(numba) = ret[0] { numba } else { -1 }; assert_eq!(numba, 42); + + let ret = wallet.query_custom("SELECT numba FROM mista;", &[]).unwrap(); + assert_eq!(ret.len(), 1); + assert_eq!(ret[0].len(), 1); + let numba: i64 = if let Value::Integer(numba) = ret[0][0] { numba } else { -1 }; + assert_eq!(numba, 42); } #[test] @@ -585,6 +639,13 @@ mod tests { assert_eq!(ret[1], Value::Text(are.clone())); assert_eq!(ret[2], Value::Integer(you)); assert_eq!(ret[3], Value::Blob(gae.clone())); + let ret = wallet.query_custom("SELECT why, are, you, gae FROM mista;", &[]).unwrap(); + assert_eq!(ret.len(), 1); + assert_eq!(ret[0].len(), 4); + assert_eq!(ret[0][0], Value::Integer(why)); + assert_eq!(ret[0][1], Value::Text(are.clone())); + assert_eq!(ret[0][2], Value::Integer(you)); + assert_eq!(ret[0][3], Value::Blob(gae.clone())); let ret = wallet .query_single( @@ -594,7 +655,16 @@ mod tests { ) .unwrap(); assert_eq!(ret.len(), 1); - assert_eq!(ret[0], Value::Blob(gae)); + assert_eq!(ret[0], Value::Blob(gae.clone())); + let ret = wallet + .query_custom( + "SELECT gae FROM mista WHERE why = ?1 AND are = ?2 AND you = ?3;", + rusqlite::params![why, are, you], + ) + .unwrap(); + assert_eq!(ret.len(), 1); + assert_eq!(ret[0].len(), 1); + assert_eq!(ret[0][0], Value::Blob(gae)); } #[test] @@ -631,6 +701,15 @@ mod tests { assert_eq!(row[2], Value::Integer(you)); assert_eq!(row[3], Value::Blob(gae.clone())); } + let ret = wallet.query_custom("SELECT * FROM mista;", &[]).unwrap(); + assert_eq!(ret.len(), 2); + for row in ret { + assert_eq!(row.len(), 4); + assert_eq!(row[0], Value::Integer(why)); + assert_eq!(row[1], Value::Text(are.clone())); + assert_eq!(row[2], Value::Integer(you)); + assert_eq!(row[3], Value::Blob(gae.clone())); + } let ret = wallet .query_multiple( @@ -644,6 +723,17 @@ mod tests { assert_eq!(row.len(), 1); assert_eq!(row[0], Value::Blob(gae.clone())); } + let ret = wallet + .query_custom( + "SELECT gae FROM mista WHERE why = ?1 AND are = ?2 AND you = ?3;", + rusqlite::params![why, are, you], + ) + .unwrap(); + assert_eq!(ret.len(), 2); + for row in ret { + assert_eq!(row.len(), 1); + assert_eq!(row[0], Value::Blob(gae.clone())); + } } #[test]