get_balance supports custom constraints

This commit is contained in:
Lex Berezhny 2018-07-09 09:55:59 -04:00
parent ae56b81a08
commit 2006ad68ee
2 changed files with 21 additions and 12 deletions

View file

@ -173,5 +173,5 @@ class BaseAccount(object):
assert not self.encrypted, "Cannot get private key on encrypted wallet account." assert not self.encrypted, "Cannot get private key on encrypted wallet account."
return self.private_key.child(chain).child(index) return self.private_key.child(chain).child(index)
def get_balance(self): def get_balance(self, **constraints):
return self.ledger.db.get_balance_for_account(self) return self.ledger.db.get_balance_for_account(self, **constraints)

View file

@ -156,6 +156,15 @@ class BaseDatabase(SQLiteMixin):
CREATE_TXI_TABLE CREATE_TXI_TABLE
) )
def txo_to_row(self, tx, address, txo):
return {
'txhash': sqlite3.Binary(tx.hash),
'address': sqlite3.Binary(address),
'position': txo.index,
'amount': txo.amount,
'script': sqlite3.Binary(txo.script.source)
}
def save_transaction_io(self, save_tx, tx, height, is_verified, address, hash, history): def save_transaction_io(self, save_tx, tx, height, is_verified, address, hash, history):
def _steps(t): def _steps(t):
@ -181,13 +190,7 @@ class BaseDatabase(SQLiteMixin):
if txo.index in existing_txos: if txo.index in existing_txos:
continue continue
if txo.script.is_pay_pubkey_hash and txo.script.values['pubkey_hash'] == hash: if txo.script.is_pay_pubkey_hash and txo.script.values['pubkey_hash'] == hash:
t.execute(*self._insert_sql("txo", { t.execute(*self._insert_sql("txo", self.txo_to_row(tx, address, txo)))
'txhash': sqlite3.Binary(tx.hash),
'address': sqlite3.Binary(address),
'position': txo.index,
'amount': txo.amount,
'script': sqlite3.Binary(txo.script.source)
}))
elif txo.script.is_pay_script_hash: elif txo.script.is_pay_script_hash:
# TODO: implement script hash payments # TODO: implement script hash payments
print('Database.save_transaction_io: pay script hash is not implemented!') print('Database.save_transaction_io: pay script hash is not implemented!')
@ -233,15 +236,21 @@ class BaseDatabase(SQLiteMixin):
defer.returnValue((None, None, False)) defer.returnValue((None, None, False))
@defer.inlineCallbacks @defer.inlineCallbacks
def get_balance_for_account(self, account): def get_balance_for_account(self, account, **constraints):
extra_sql = ""
if constraints:
extra_sql = ' AND ' + ' AND '.join(
'{} = :{}'.format(c, c) for c in constraints.keys()
)
values = {'account': sqlite3.Binary(account.public_key.address)}
values.update(constraints)
result = yield self.db.runQuery( result = yield self.db.runQuery(
""" """
SELECT SUM(amount) FROM txo SELECT SUM(amount) FROM txo
JOIN pubkey_address ON pubkey_address.address=txo.address JOIN pubkey_address ON pubkey_address.address=txo.address
WHERE account=:account AND WHERE account=:account AND
txoid NOT IN (SELECT txoid FROM txi) txoid NOT IN (SELECT txoid FROM txi)
""", """+extra_sql, values
{'account': sqlite3.Binary(account.public_key.address)}
) )
if result: if result:
defer.returnValue(result[0][0] or 0) defer.returnValue(result[0][0] or 0)