Merge pull request #204 from neocybereth/feat/add-total-reward-for-epoch

feat: add total reward for epoch for each user

Author: sug0

chain/src/repository/pos.rs

@@ -248,10 +248,8 @@ pub fn delete_claimed_rewards(
         return Ok(());
     }
 
-    let mut query = diesel::delete(pos_rewards::table).into_boxed();
-
     for (owner, validator_id) in reward_claimers {
-        query = query.or_filter(
+        let target = pos_rewards::table.filter(
             pos_rewards::owner.eq(owner.to_string()).and(
                 pos_rewards::validator_id.eq_any(
                     validators::table.select(validators::columns::id).filter(
@@ -261,11 +259,12 @@ pub fn delete_claimed_rewards(
                 ),
             ),
         );
-    }
 
-    query
-        .execute(transaction_conn)
-        .context("Failed to remove pos rewards from db")?;
+        diesel::update(target)
+            .set(pos_rewards::claimed.eq(true))
+            .execute(transaction_conn)
+            .context("Failed to update pos rewards in db")?;
+    }
 
     anyhow::Ok(())
 }

orm/migrations/2024-12-16-231256_add_epoch_pos_rewards_table/down.sql

@@ -0,0 +1,4 @@
+ALTER table pos_rewards ADD CONSTRAINT pos_rewards_owner_validator_id_key UNIQUE (owner, validator_id);
+ALTER table pos_rewards DROP CONSTRAINT pos_rewards_owner_validator_id_epoch_key;
+ALTER TABLE pos_rewards DROP COLUMN epoch;
+ALTER TABLE pos_rewards DROP COLUMN claimed;

orm/migrations/2024-12-16-231256_add_epoch_pos_rewards_table/up.sql

@@ -0,0 +1,9 @@
+-- Populate existing records with claimed = false
+ALTER TABLE pos_rewards ADD COLUMN claimed BOOLEAN DEFAULT FALSE;
+-- Populate existing records with epoch = 0
+ALTER TABLE pos_rewards ADD COLUMN epoch INTEGER NOT NULL DEFAULT 0;
+-- Now we can safely drop the default
+ALTER TABLE pos_rewards ALTER COLUMN epoch DROP DEFAULT;
+-- Also update the UNIQUE constraint to include the epoch column
+ALTER table pos_rewards ADD CONSTRAINT pos_rewards_owner_validator_id_epoch_key unique (owner, validator_id, epoch);
+ALTER table pos_rewards DROP CONSTRAINT pos_rewards_owner_validator_id_key;

orm/src/pos_rewards.rs

@@ -6,24 +6,27 @@ use shared::rewards::Reward;
 
 use crate::schema::pos_rewards;
 
-#[derive(Insertable, Clone, Queryable, Selectable)]
+#[derive(Insertable, Queryable, Selectable, Clone)]
 #[diesel(table_name = pos_rewards)]
-#[diesel(check_for_backend(diesel::pg::Pg))]
 pub struct PosRewardInsertDb {
     pub owner: String,
     pub validator_id: i32,
     pub raw_amount: BigDecimal,
+    pub epoch: i32,
+    pub claimed: bool,
 }
 
 pub type PoSRewardDb = PosRewardInsertDb;
 
 impl PosRewardInsertDb {
-    pub fn from_reward(reward: Reward, validator_id: i32) -> Self {
-        Self {
+    pub fn from_reward(reward: Reward, validator_id: i32, epoch: i32) -> Self {
+        PosRewardInsertDb {
             owner: reward.delegation_pair.delegator_address.to_string(),
-            raw_amount: BigDecimal::from_str(&reward.amount.to_string())
-                .expect("Invalid amount"),
             validator_id,
+            raw_amount: BigDecimal::from_str(&reward.amount.to_string())
+                .unwrap(),
+            epoch,
+            claimed: false,
         }
     }
 }

orm/src/schema.rs

@@ -359,6 +359,8 @@ diesel::table! {
         owner -> Varchar,
         validator_id -> Int4,
         raw_amount -> Numeric,
+        epoch -> Int4,
+        claimed -> Bool,
     }
 }
 

rewards/src/main.rs

@@ -117,20 +117,25 @@ async fn crawling_fn(
         "Queried rewards successfully",
     );
 
+    let current_epoch = namada_service::get_current_epoch(&client)
+        .await
+        .into_rpc_error()?;
+
     conn.interact(move |conn| {
         conn.build_transaction().read_write().run(
-            |transaction_conn: &mut diesel::prelude::PgConnection| {
+            |transaction_conn: &mut diesel::pg::PgConnection| {
                 repository::pos_rewards::upsert_rewards(
                     transaction_conn,
                     non_zero_rewards,
+                    current_epoch as i32,
                 )?;
 
                 repository::crawler_state::upsert_crawler_state(
                     transaction_conn,
                     (CrawlerName::Rewards, crawler_state).into(),
                 )?;
 
-                anyhow::Ok(())
+                Ok(())
             },
         )
     })

rewards/src/repository/pos_rewards.rs

@@ -11,6 +11,7 @@ use super::utils::MAX_PARAM_SIZE;
 pub fn upsert_rewards(
     transaction_conn: &mut PgConnection,
     rewards: Vec<Reward>,
+    epoch: i32,
 ) -> anyhow::Result<()> {
     let rewards_col_count = pos_rewards::all_columns.len() as i64;
 
@@ -19,7 +20,7 @@ pub fn upsert_rewards(
         .collect::<Vec<_>>()
         .chunks((MAX_PARAM_SIZE as i64 / rewards_col_count) as usize)
     {
-        upsert_rewards_chunk(transaction_conn, chunk.to_vec())?;
+        upsert_rewards_chunk(transaction_conn, chunk.to_vec(), epoch)?;
     }
 
     anyhow::Ok(())
@@ -28,6 +29,7 @@ pub fn upsert_rewards(
 fn upsert_rewards_chunk(
     transaction_conn: &mut PgConnection,
     rewards: Vec<Reward>,
+    epoch: i32,
 ) -> anyhow::Result<()> {
     diesel::insert_into(pos_rewards::table)
         .values::<Vec<PosRewardInsertDb>>(
@@ -45,13 +47,14 @@ fn upsert_rewards_chunk(
                         .first(transaction_conn)
                         .expect("Failed to get validator");
 
-                    PosRewardInsertDb::from_reward(reward, validator_id)
+                    PosRewardInsertDb::from_reward(reward, validator_id, epoch)
                 })
                 .collect::<Vec<_>>(),
         )
         .on_conflict((
             pos_rewards::columns::owner,
             pos_rewards::columns::validator_id,
+            pos_rewards::columns::epoch, // Add epoch to conflict target
         ))
         .do_update()
         .set(

rewards/src/services/namada.rs

@@ -125,6 +125,8 @@ async fn process_batch(
     client: &HttpClient,
     batch: Vec<DelegationPair>,
 ) -> anyhow::Result<Vec<Reward>> {
+    let epoch = get_current_epoch(client).await?;
+
     Ok(futures::stream::iter(batch)
         .filter_map(|delegation| async move {
             tracing::debug!(
@@ -154,6 +156,7 @@ async fn process_batch(
             Some(Reward {
                 delegation_pair: delegation.clone(),
                 amount: Amount::from(reward),
+                epoch: epoch as i32,
             })
         })
         .map(futures::future::ready)

seeder/src/main.rs

@@ -177,7 +177,8 @@ async fn main() -> anyhow::Result<(), MainError> {
                             .into_iter()
                             .map(|reward| {
                                 let validator_id = reward.delegation_pair.validator_address.to_string().parse::<i32>().unwrap();
-                                PosRewardInsertDb::from_reward(reward, validator_id)
+                                let epoch = reward.epoch;
+                                PosRewardInsertDb::from_reward(reward, validator_id, epoch)
                             })
                             .collect::<Vec<_>>(),
                     )

shared/src/rewards.rs

@@ -6,6 +6,7 @@ use crate::utils::DelegationPair;
 pub struct Reward {
     pub delegation_pair: DelegationPair,
     pub amount: Amount,
+    pub epoch: i32,
 }
 
 impl Reward {
@@ -19,6 +20,7 @@ impl Reward {
                 delegator_address: Id::Account(delegator_address.to_string()),
             },
             amount: Amount::fake(),
+            epoch: 0,
         }
     }
 }

webserver/src/app.rs

@@ -75,6 +75,10 @@ impl ApplicationServer {
                     get(pos_handlers::get_withdraws),
                 )
                 .route("/pos/reward/:address", get(pos_handlers::get_rewards))
+                .route(
+                    "/pos/reward/:delegator/:validator/:epoch",
+                    get(pos_handlers::get_rewards_by_delegator_and_validator_and_epoch),
+                )
                 .route(
                     "/pos/voting-power",
                     get(pos_handlers::get_total_voting_power),

webserver/src/dto/pos.rs

@@ -88,3 +88,9 @@ pub struct WithdrawsDto {
     #[validate(range(min = 1, max = 10000))]
     pub epoch: Option<u64>,
 }
+
+#[derive(Clone, Serialize, Deserialize, Validate)]
+pub struct RewardsDto {
+    #[validate(range(min = 1, max = 10000))]
+    pub epoch: Option<u64>,
+}

webserver/src/handler/pos.rs

@@ -5,8 +5,8 @@ use axum_extra::extract::Query;
 use axum_macros::debug_handler;
 
 use crate::dto::pos::{
-    AllValidatorsQueryParams, BondsDto, UnbondsDto, ValidatorQueryParams,
-    ValidatorStateDto, WithdrawsDto,
+    AllValidatorsQueryParams, BondsDto, RewardsDto, UnbondsDto,
+    ValidatorQueryParams, ValidatorStateDto, WithdrawsDto,
 };
 use crate::error::api::ApiError;
 use crate::response::pos::{
@@ -149,10 +149,29 @@ pub async fn get_withdraws(
 #[debug_handler]
 pub async fn get_rewards(
     _headers: HeaderMap,
+    query: Query<RewardsDto>,
     Path(address): Path<String>,
     State(state): State<CommonState>,
 ) -> Result<Json<Vec<Reward>>, ApiError> {
-    let rewards = state.pos_service.get_rewards_by_address(address).await?;
+    let rewards = state
+        .pos_service
+        .get_rewards_by_address(address, query.epoch)
+        .await?;
+    Ok(Json(rewards))
+}
+
+#[debug_handler]
+pub async fn get_rewards_by_delegator_and_validator_and_epoch(
+    _headers: HeaderMap,
+    Path((delegator, validator, epoch)): Path<(String, String, u64)>,
+    State(state): State<CommonState>,
+) -> Result<Json<Vec<Reward>>, ApiError> {
+    let rewards = state
+        .pos_service
+        .get_rewards_by_delegator_and_validator_and_epoch(
+            delegator, validator, epoch,
+        )
+        .await?;
     Ok(Json(rewards))
 }
 

webserver/src/repository/pos.rs

@@ -49,6 +49,11 @@ pub trait PosRepositoryTrait {
         id: i32,
     ) -> Result<Option<ValidatorDb>, String>;
 
+    async fn find_validator_by_address(
+        &self,
+        address: String,
+    ) -> Result<Option<ValidatorDb>, String>;
+
     async fn find_merged_bonds_by_address(
         &self,
         address: String,
@@ -97,6 +102,14 @@ pub trait PosRepositoryTrait {
     async fn find_rewards_by_address(
         &self,
         address: String,
+        epoch: Option<u64>,
+    ) -> Result<Vec<PoSRewardDb>, String>;
+
+    async fn find_rewards_by_delegator_and_validator_and_epoch(
+        &self,
+        delegator: String,
+        validator_id: i32,
+        epoch: u64,
     ) -> Result<Vec<PoSRewardDb>, String>;
 
     async fn get_total_voting_power(&self) -> Result<Option<i64>, String>;
@@ -188,6 +201,23 @@ impl PosRepositoryTrait for PosRepository {
         .map_err(|e| e.to_string())
     }
 
+    async fn find_validator_by_address(
+        &self,
+        address: String,
+    ) -> Result<Option<ValidatorDb>, String> {
+        let conn = self.app_state.get_db_connection().await;
+
+        conn.interact(move |conn| {
+            validators::table
+                .filter(validators::dsl::namada_address.eq(address))
+                .select(ValidatorDb::as_select())
+                .first(conn)
+                .ok()
+        })
+        .await
+        .map_err(|e| e.to_string())
+    }
+
     async fn find_bonds_by_address(
         &self,
         address: String,
@@ -344,12 +374,50 @@ impl PosRepositoryTrait for PosRepository {
     async fn find_rewards_by_address(
         &self,
         address: String,
+        epoch: Option<u64>,
+    ) -> Result<Vec<PoSRewardDb>, String> {
+        let conn = self.app_state.get_db_connection().await;
+
+        conn.interact(
+            move |conn| -> Result<Vec<PoSRewardDb>, diesel::result::Error> {
+                let epoch = match epoch {
+                    Some(e) => e as i32,
+                    None => {
+                        // Properly propagate database errors with ?
+                        let max_epoch = pos_rewards::table
+                            .select(diesel::dsl::max(pos_rewards::epoch))
+                            .first::<Option<i32>>(conn)?;
+
+                        max_epoch.unwrap_or(0) // This is fine - just unwrapping the Option, not a Result
+                    }
+                };
+
+                // Propagate errors from this query too
+                pos_rewards::table
+                    .filter(pos_rewards::epoch.eq(&epoch))
+                    .filter(pos_rewards::dsl::owner.eq(&address))
+                    .select(PoSRewardDb::as_select())
+                    .get_results(conn)
+            },
+        )
+        .await
+        .map_err(|e| e.to_string())?
+        .map_err(|e| e.to_string())
+    }
+
+    async fn find_rewards_by_delegator_and_validator_and_epoch(
+        &self,
+        delegator: String,
+        validator_id: i32,
+        epoch: u64,
     ) -> Result<Vec<PoSRewardDb>, String> {
         let conn = self.app_state.get_db_connection().await;
 
         conn.interact(move |conn| {
             pos_rewards::table
-                .filter(pos_rewards::dsl::owner.eq(address))
+                .filter(pos_rewards::dsl::owner.eq(delegator))
+                .filter(pos_rewards::dsl::validator_id.eq(validator_id))
+                .filter(pos_rewards::dsl::epoch.eq(epoch as i32))
                 .select(PoSRewardDb::as_select())
                 .get_results(conn)
         })