Skip to content

Commit

Permalink
[enhancement] balances -> balance_changes, record balance history by …
Browse files Browse the repository at this point in the history
…height
  • Loading branch information
joel-u410 committed Jul 19, 2024
1 parent 50c042b commit 0bb5d21
Show file tree
Hide file tree
Showing 12 changed files with 124 additions and 37 deletions.
4 changes: 2 additions & 2 deletions chain/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -151,7 +151,7 @@ async fn crawling_fn(

let addresses = block.addresses_with_balance_change(native_token);
let balances =
namada_service::query_balance(&client, &addresses, Some(block_height))
namada_service::query_balance(&client, &addresses, block_height)
.await
.into_rpc_error()?;
tracing::info!("Updating balance for {} addresses...", addresses.len());
Expand Down Expand Up @@ -292,7 +292,7 @@ async fn initial_query(
sleep(Duration::from_secs(initial_query_retry_time)).await;
}

let balances = query_all_balances(client, Some(block_height))
let balances = query_all_balances(client, block_height)
.await
.into_rpc_error()?;

Expand Down
25 changes: 12 additions & 13 deletions chain/src/repository/balance.rs
Original file line number Diff line number Diff line change
@@ -1,27 +1,26 @@
use anyhow::Context;
use diesel::upsert::excluded;
use diesel::{ExpressionMethods, PgConnection, RunQueryDsl};
use orm::balances::BalancesInsertDb;
use orm::schema::balances;
use diesel::{PgConnection, RunQueryDsl};
use orm::balances::BalanceChangesInsertDb;
use orm::schema::balance_changes;
use shared::balance::Balances;

pub fn insert_balance(
transaction_conn: &mut PgConnection,
balances: Balances,
) -> anyhow::Result<()> {
diesel::insert_into(balances::table)
.values::<&Vec<BalancesInsertDb>>(
diesel::insert_into(balance_changes::table)
.values::<&Vec<BalanceChangesInsertDb>>(
&balances
.into_iter()
.map(BalancesInsertDb::from_balance)
.map(BalanceChangesInsertDb::from_balance)
.collect::<Vec<_>>(),
)
.on_conflict((balances::columns::owner, balances::columns::token))
.do_update()
.set(
balances::columns::raw_amount
.eq(excluded(balances::columns::raw_amount)),
)
.on_conflict((
balance_changes::columns::owner,
balance_changes::columns::token,
balance_changes::columns::height,
))
.do_nothing()
.execute(transaction_conn)
.context("Failed to update balances in db")?;

Expand Down
19 changes: 12 additions & 7 deletions chain/src/services/namada.rs
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ pub async fn get_epoch_at_block_height(
pub async fn query_balance(
client: &HttpClient,
balance_changes: &HashSet<BalanceChange>,
block_height: Option<BlockHeight>,
block_height: BlockHeight,
) -> anyhow::Result<Balances> {
Ok(futures::stream::iter(balance_changes)
.filter_map(|balance_change| async move {
Expand All @@ -79,7 +79,7 @@ pub async fn query_balance(
client,
&token,
&owner,
block_height.map(to_block_height),
Some(to_block_height(block_height)),
)
.await
.unwrap_or_default();
Expand All @@ -88,6 +88,7 @@ pub async fn query_balance(
owner: Id::from(owner),
token: Id::from(token),
amount: Amount::from(amount),
height: block_height,
})
})
.map(futures::future::ready)
Expand All @@ -98,7 +99,7 @@ pub async fn query_balance(

pub async fn query_all_balances(
client: &HttpClient,
height: Option<BlockHeight>,
height: BlockHeight,
) -> anyhow::Result<Balances> {
let token_addr = RPC
.shell()
Expand All @@ -108,10 +109,13 @@ pub async fn query_all_balances(

let balance_prefix = namada_token::storage_key::balance_prefix(&token_addr);

let balances =
query_storage_prefix::<token::Amount>(client, &balance_prefix, height)
.await
.context("Failed to query all balances")?;
let balances = query_storage_prefix::<token::Amount>(
client,
&balance_prefix,
Some(height),
)
.await
.context("Failed to query all balances")?;

let mut all_balances: Balances = vec![];

Expand All @@ -128,6 +132,7 @@ pub async fn query_all_balances(
owner: Id::from(o),
token: Id::from(t),
amount: Amount::from(b),
height,
});
}
}
Expand Down
18 changes: 18 additions & 0 deletions orm/migrations/2024-07-19-203433_add_height_to_balances/down.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
-- This file should undo anything in `up.sql`
DROP VIEW balances;

CREATE INDEX index_balances_owner ON balance_changes (OWNER, token);

DROP INDEX index_balance_changes_owner_token_height;

ALTER TABLE balance_changes
DROP CONSTRAINT balance_changes_owner_token_height_key;

ALTER TABLE balance_changes RENAME TO balances;

ALTER TABLE balances
ADD CONSTRAINT balances_owner_token_key UNIQUE (OWNER, token);

ALTER TABLE balances
DROP COLUMN height;

40 changes: 40 additions & 0 deletions orm/migrations/2024-07-19-203433_add_height_to_balances/up.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
-- Your SQL goes here
ALTER TABLE balances
ADD COLUMN height integer NOT NULL DEFAULT 0;

ALTER TABLE balances
ALTER COLUMN height DROP DEFAULT;

ALTER TABLE balances
DROP CONSTRAINT balances_owner_token_key;

ALTER TABLE balances RENAME TO balance_changes;

ALTER TABLE balance_changes
ADD CONSTRAINT balance_changes_owner_token_height_key UNIQUE (OWNER, token, height);

CREATE INDEX index_balance_changes_owner_token_height ON balance_changes (OWNER, token, height);

DROP INDEX index_balances_owner;

CREATE VIEW balances AS
SELECT
bc.id,
bc.owner,
bc.token,
bc.raw_amount
FROM
balance_changes bc
JOIN (
SELECT
OWNER,
token,
MAX(height) AS max_height
FROM
balance_changes
GROUP BY
OWNER,
token) max_heights ON bc.owner = max_heights.owner
AND bc.token = max_heights.token
AND bc.height = max_heights.max_height;

22 changes: 17 additions & 5 deletions orm/src/balances.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,26 +4,38 @@ use bigdecimal::BigDecimal;
use diesel::{Insertable, Queryable, Selectable};
use shared::balance::Balance;

use crate::schema::balances;
use crate::schema::balance_changes;
use crate::views::balances;

#[derive(Insertable, Clone, Queryable, Selectable)]
#[diesel(table_name = balances)]
#[diesel(table_name = balance_changes)]
#[diesel(check_for_backend(diesel::pg::Pg))]
pub struct BalancesInsertDb {
pub struct BalanceChangesInsertDb {
pub owner: String,
pub token: String,
pub raw_amount: BigDecimal,
pub height: i32,
}

pub type BalanceDb = BalancesInsertDb;
pub type BalanceChangeDb = BalanceChangesInsertDb;

#[derive(Clone, Queryable, Selectable)]
#[diesel(table_name = balances)]
#[diesel(check_for_backend(diesel::pg::Pg))]
pub struct BalanceDb {
pub owner: String,
pub token: String,
pub raw_amount: BigDecimal,
}

impl BalancesInsertDb {
impl BalanceChangesInsertDb {
pub fn from_balance(balance: Balance) -> Self {
Self {
owner: balance.owner.to_string(),
token: balance.token.to_string(),
raw_amount: BigDecimal::from_str(&balance.amount.to_string())
.expect("Invalid amount"),
height: balance.height as i32,
}
}
}
1 change: 1 addition & 0 deletions orm/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,3 +13,4 @@ pub mod schema;
pub mod transactions;
pub mod unbond;
pub mod validators;
pub mod views;
5 changes: 3 additions & 2 deletions orm/src/schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -67,11 +67,12 @@ pub mod sql_types {
}

diesel::table! {
balances (id) {
balance_changes (id) {
id -> Int4,
owner -> Varchar,
token -> Varchar,
raw_amount -> Numeric,
height -> Int4,
}
}

Expand Down Expand Up @@ -255,7 +256,7 @@ diesel::joinable!(pos_rewards -> validators (validator_id));
diesel::joinable!(unbonds -> validators (validator_id));

diesel::allow_tables_to_appear_in_same_query!(
balances,
balance_changes,
bonds,
chain_parameters,
crawler_state,
Expand Down
9 changes: 9 additions & 0 deletions orm/src/views.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
// Manually create schema for views - see also https://github.com/diesel-rs/diesel/issues/1482
diesel::table! {
balances (id) {
id -> Int4,
owner -> Varchar,
token -> Varchar,
raw_amount -> Numeric,
}
}
14 changes: 7 additions & 7 deletions seeder/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,16 @@ use anyhow::Context;
use clap::Parser;
use clap_verbosity_flag::LevelFilter;
use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl, SelectableHelper};
use orm::balances::BalancesInsertDb;
use orm::balances::BalanceChangesInsertDb;
use orm::bond::BondInsertDb;
use orm::governance_proposal::{
GovernanceProposalInsertDb, GovernanceProposalUpdateStatusDb,
};
use orm::governance_votes::GovernanceProposalVoteInsertDb;
use orm::pos_rewards::PosRewardInsertDb;
use orm::schema::{
balances, bonds, governance_proposals, governance_votes, pos_rewards,
unbonds, validators,
balance_changes, bonds, governance_proposals, governance_votes,
pos_rewards, unbonds, validators,
};
use orm::unbond::UnbondInsertDb;
use orm::validators::{ValidatorDb, ValidatorInsertDb};
Expand Down Expand Up @@ -138,7 +138,7 @@ async fn main() -> anyhow::Result<(), MainError> {
.execute(transaction_conn)
.context("Failed to remove all validators")?;

diesel::delete(balances::table)
diesel::delete(balance_changes::table)
.execute(transaction_conn)
.context("Failed to remove all validators")?;

Expand Down Expand Up @@ -201,12 +201,12 @@ async fn main() -> anyhow::Result<(), MainError> {
.execute(transaction_conn)
.context("Failed to insert pos rewards in db")?;

diesel::insert_into(balances::table)
.values::<&Vec<BalancesInsertDb>>(
diesel::insert_into(balance_changes::table)
.values::<&Vec<BalanceChangesInsertDb>>(
&balances
.into_iter()
.map(|balance| {
BalancesInsertDb::from_balance(balance)
BalanceChangesInsertDb::from_balance(balance)
})
.collect::<Vec<_>>(),
)
Expand Down
2 changes: 2 additions & 0 deletions shared/src/balance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ pub struct Balance {
pub owner: Id,
pub token: Id,
pub amount: Amount,
pub height: u32,
}

pub type Balances = Vec<Balance>;
Expand All @@ -92,6 +93,7 @@ impl Balance {
owner: Id::Account(address.to_string()),
token: Id::Account(token_address.to_string()),
amount: Amount::fake(),
height: (0..10000).fake::<u32>(),
}
}
}
2 changes: 1 addition & 1 deletion webserver/src/repository/balance.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use axum::async_trait;
use diesel::{ExpressionMethods, QueryDsl, RunQueryDsl, SelectableHelper};
use orm::balances::BalanceDb;
use orm::schema::balances;
use orm::views::balances;

use crate::appstate::AppState;

Expand Down

0 comments on commit 0bb5d21

Please sign in to comment.