zcash_client_sqlite::query::{get_balance, get_verified_balance}

zcash_client_sqlite/src/error.rs

@@ -3,6 +3,8 @@ use std::fmt;
 
 #[derive(Debug)]
 pub enum ErrorKind {
+    CorruptedData(&'static str),
+    ScanRequired,
     TableNotEmpty,
     Database(rusqlite::Error),
 }
@@ -13,6 +15,8 @@ pub struct Error(pub(crate) ErrorKind);
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match &self.0 {
+            ErrorKind::CorruptedData(reason) => write!(f, "Data DB is corrupted: {}", reason),
+            ErrorKind::ScanRequired => write!(f, "Must scan blocks first"),
             ErrorKind::TableNotEmpty => write!(f, "Table is not empty"),
             ErrorKind::Database(e) => write!(f, "{}", e),
         }

zcash_client_sqlite/src/lib.rs

@@ -19,6 +19,8 @@
 //! [`CompactBlock`]: zcash_client_backend::proto::compact_formats::CompactBlock
 //! [`init_cache_database`]: crate::init::init_cache_database
 
+use rusqlite::{Connection, NO_PARAMS};
+use std::cmp;
 use zcash_client_backend::{
     constants::testnet::HRP_SAPLING_PAYMENT_ADDRESS, encoding::encode_payment_address,
 };
@@ -28,7 +30,36 @@ pub mod error;
 pub mod init;
 pub mod query;
 
+const ANCHOR_OFFSET: u32 = 10;
+
 fn address_from_extfvk(extfvk: &ExtendedFullViewingKey) -> String {
     let addr = extfvk.default_address().unwrap().1;
     encode_payment_address(HRP_SAPLING_PAYMENT_ADDRESS, &addr)
 }
+
+/// Determines the target height for a transaction, and the height from which to
+/// select anchors, based on the current synchronised block chain.
+fn get_target_and_anchor_heights(data: &Connection) -> Result<(u32, u32), error::Error> {
+    data.query_row_and_then(
+        "SELECT MIN(height), MAX(height) FROM blocks",
+        NO_PARAMS,
+        |row| match (row.get::<_, u32>(0), row.get::<_, u32>(1)) {
+            // If there are no blocks, the query returns NULL.
+            (Err(rusqlite::Error::InvalidColumnType(_, _, _)), _)
+            | (_, Err(rusqlite::Error::InvalidColumnType(_, _, _))) => {
+                Err(error::Error(error::ErrorKind::ScanRequired))
+            }
+            (Err(e), _) | (_, Err(e)) => Err(e.into()),
+            (Ok(min_height), Ok(max_height)) => {
+                let target_height = max_height + 1;
+
+                // Select an anchor ANCHOR_OFFSET back from the target block,
+                // unless that would be before the earliest block we have.
+                let anchor_height =
+                    cmp::max(target_height.saturating_sub(ANCHOR_OFFSET), min_height);
+
+                Ok((target_height, anchor_height))
+            }
+        },
+    )
+}

zcash_client_sqlite/src/query.rs

@@ -2,8 +2,12 @@
 
 use rusqlite::Connection;
 use std::path::Path;
+use zcash_primitives::transaction::components::Amount;
 
-use crate::error::Error;
+use crate::{
+    error::{Error, ErrorKind},
+    get_target_and_anchor_heights,
+};
 
 /// Returns the address for the account.
 ///
@@ -26,3 +30,104 @@ pub fn get_address<P: AsRef<Path>>(db_data: P, account: u32) -> Result<String, E
 
     Ok(addr)
 }
+
+/// Returns the balance for the account, including all mined unspent notes that we know
+/// about.
+///
+/// # Examples
+///
+/// ```
+/// use zcash_client_sqlite::query::get_balance;
+///
+/// let addr = get_balance("/path/to/data.db", 0);
+/// ```
+pub fn get_balance<P: AsRef<Path>>(db_data: P, account: u32) -> Result<Amount, Error> {
+    let data = Connection::open(db_data)?;
+
+    let balance = data.query_row(
+        "SELECT SUM(value) FROM received_notes
+        INNER JOIN transactions ON transactions.id_tx = received_notes.tx
+        WHERE account = ? AND spent IS NULL AND transactions.block IS NOT NULL",
+        &[account],
+        |row| row.get(0).or(Ok(0)),
+    )?;
+
+    match Amount::from_i64(balance) {
+        Ok(amount) if !amount.is_negative() => Ok(amount),
+        _ => Err(Error(ErrorKind::CorruptedData(
+            "Sum of values in received_notes is out of range",
+        ))),
+    }
+}
+
+/// Returns the verified balance for the account, which ignores notes that have been
+/// received too recently and are not yet deemed spendable.
+///
+/// # Examples
+///
+/// ```
+/// use zcash_client_sqlite::query::get_verified_balance;
+///
+/// let addr = get_verified_balance("/path/to/data.db", 0);
+/// ```
+pub fn get_verified_balance<P: AsRef<Path>>(db_data: P, account: u32) -> Result<Amount, Error> {
+    let data = Connection::open(db_data)?;
+
+    let (_, anchor_height) = get_target_and_anchor_heights(&data)?;
+
+    let balance = data.query_row(
+        "SELECT SUM(value) FROM received_notes
+        INNER JOIN transactions ON transactions.id_tx = received_notes.tx
+        WHERE account = ? AND spent IS NULL AND transactions.block <= ?",
+        &[account, anchor_height],
+        |row| row.get(0).or(Ok(0)),
+    )?;
+
+    match Amount::from_i64(balance) {
+        Ok(amount) if !amount.is_negative() => Ok(amount),
+        _ => Err(Error(ErrorKind::CorruptedData(
+            "Sum of values in received_notes is out of range",
+        ))),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use tempfile::NamedTempFile;
+    use zcash_primitives::{
+        transaction::components::Amount,
+        zip32::{ExtendedFullViewingKey, ExtendedSpendingKey},
+    };
+
+    use super::{get_address, get_balance, get_verified_balance};
+    use crate::{
+        error::ErrorKind,
+        init::{init_accounts_table, init_data_database},
+    };
+
+    #[test]
+    fn empty_database_has_no_balance() {
+        let data_file = NamedTempFile::new().unwrap();
+        let db_data = data_file.path();
+        init_data_database(&db_data).unwrap();
+
+        // Add an account to the wallet
+        let extsk = ExtendedSpendingKey::master(&[]);
+        let extfvks = [ExtendedFullViewingKey::from(&extsk)];
+        init_accounts_table(&db_data, &extfvks).unwrap();
+
+        // The account should be empty
+        assert_eq!(get_balance(db_data, 0).unwrap(), Amount::zero());
+
+        // The account should have no verified balance, as we haven't scanned any blocks
+        let e = get_verified_balance(db_data, 0).unwrap_err();
+        match e.kind() {
+            ErrorKind::ScanRequired => (),
+            _ => panic!("Unexpected error: {:?}", e),
+        }
+
+        // An invalid account has zero balance
+        assert!(get_address(db_data, 1).is_err());
+        assert_eq!(get_balance(db_data, 1).unwrap(), Amount::zero());
+    }
+}