From 87760c5892c21acc1a48eb2372c572febd26db24 Mon Sep 17 00:00:00 2001
From: Joey Kraut <joey@renegade.fi>
Date: Thu, 25 Jul 2024 13:15:55 -0700
Subject: [PATCH] funds-manager: fee-indexer: Fix misc bugs with fee indexing +
 redemption

---
 .../src/fee_indexer/index_fees.rs             |  2 +-
 .../src/fee_indexer/queries.rs                | 39 ++++++++++++++++---
 .../src/fee_indexer/redeem_fees.rs            | 24 +++++++-----
 3 files changed, 49 insertions(+), 16 deletions(-)

diff --git a/funds-manager/funds-manager-server/src/fee_indexer/index_fees.rs b/funds-manager/funds-manager-server/src/fee_indexer/index_fees.rs
index dcac217..2dc586a 100644
--- a/funds-manager/funds-manager-server/src/fee_indexer/index_fees.rs
+++ b/funds-manager/funds-manager-server/src/fee_indexer/index_fees.rs
@@ -66,7 +66,7 @@ impl Indexer {
         let note = match maybe_note {
             Some(note) => note,
             None => {
-                info!("note not found, skipping");
+                info!("not the note receiver, skipping...");
                 return Ok(());
             },
         };
diff --git a/funds-manager/funds-manager-server/src/fee_indexer/queries.rs b/funds-manager/funds-manager-server/src/fee_indexer/queries.rs
index 8eccd82..9295bcb 100644
--- a/funds-manager/funds-manager-server/src/fee_indexer/queries.rs
+++ b/funds-manager/funds-manager-server/src/fee_indexer/queries.rs
@@ -5,6 +5,7 @@ use std::collections::HashMap;
 use bigdecimal::BigDecimal;
 use diesel::deserialize::Queryable;
 use diesel::deserialize::QueryableByName;
+use diesel::result::Error as DieselError;
 use diesel::sql_function;
 use diesel::sql_query;
 use diesel::sql_types::SingleValue;
@@ -12,7 +13,9 @@ use diesel::sql_types::{Array, Integer, Nullable, Numeric, Text};
 use diesel::PgArrayExpressionMethods;
 use diesel::{ExpressionMethods, QueryDsl};
 use diesel_async::RunQueryDsl;
+use renegade_common::types::wallet::WalletIdentifier;
 use renegade_constants::MAX_BALANCES;
+use tracing::warn;
 
 use crate::db::models::WalletMetadata;
 use crate::db::models::{Metadata, NewFee};
@@ -44,6 +47,11 @@ sql_function! {
     fn coalesce<T: SingleValue>(x: Nullable<T>, y: T) -> T;
 }
 
+sql_function! {
+    /// Append an element to an array
+    fn array_append<T: SingleValue>(arr: Array<T>, elem: T) -> Array<T>;
+}
+
 // ---------------
 // | Query Types |
 // ---------------
@@ -111,12 +119,17 @@ impl Indexer {
 
     /// Insert a fee into the fees table
     pub(crate) async fn insert_fee(&mut self, fee: NewFee) -> Result<(), FundsManagerError> {
-        diesel::insert_into(fees_table)
-            .values(vec![fee])
-            .execute(&mut self.db_conn)
-            .await
-            .map_err(|e| FundsManagerError::db(format!("failed to insert fee: {e}")))
-            .map(|_| ())
+        match diesel::insert_into(fees_table).values(vec![fee]).execute(&mut self.db_conn).await {
+            Ok(_) => Ok(()),
+            Err(DieselError::DatabaseError(
+                diesel::result::DatabaseErrorKind::UniqueViolation,
+                _,
+            )) => {
+                warn!("Fee already exists in the database, skipping insertion...",);
+                Ok(())
+            },
+            Err(e) => Err(FundsManagerError::db(format!("failed to insert fee: {e}"))),
+        }
     }
 
     /// Get all mints that have unredeemed fees
@@ -237,4 +250,18 @@ impl Indexer {
             .map_err(|_| FundsManagerError::db("failed to insert wallet"))
             .map(|_| ())
     }
+
+    /// Add a new mint to a wallet's managed mints
+    pub(crate) async fn add_mint_to_wallet(
+        &mut self,
+        wallet_id: &WalletIdentifier,
+        mint: &str,
+    ) -> Result<(), FundsManagerError> {
+        diesel::update(wallet_table.find(wallet_id))
+            .set(managed_mints_col.eq(array_append(managed_mints_col, mint)))
+            .execute(&mut self.db_conn)
+            .await
+            .map_err(|_| FundsManagerError::db("failed to add mint to wallet"))
+            .map(|_| ())
+    }
 }
diff --git a/funds-manager/funds-manager-server/src/fee_indexer/redeem_fees.rs b/funds-manager/funds-manager-server/src/fee_indexer/redeem_fees.rs
index f4c0ac4..4c3862e 100644
--- a/funds-manager/funds-manager-server/src/fee_indexer/redeem_fees.rs
+++ b/funds-manager/funds-manager-server/src/fee_indexer/redeem_fees.rs
@@ -4,6 +4,7 @@ use std::collections::HashMap;
 use std::str::FromStr;
 
 use aws_sdk_secretsmanager::Client as SecretsManagerClient;
+use diesel::IntoSql;
 use ethers::core::rand::thread_rng;
 use ethers::signers::LocalWallet;
 use ethers::types::TxHash;
@@ -66,20 +67,25 @@ impl Indexer {
         &mut self,
         mint: &str,
     ) -> Result<WalletMetadata, FundsManagerError> {
+        // Find a wallet with an existing balance
         let maybe_wallet = self.get_wallet_for_mint(mint).await?;
-        let maybe_wallet = if maybe_wallet.is_none() {
-            self.find_wallet_with_empty_balance().await?
-        } else {
-            maybe_wallet
-        };
+        if let Some(wallet) = maybe_wallet {
+            return Ok(wallet);
+        }
 
-        match maybe_wallet {
-            Some(wallet) => Ok(wallet),
+        // Otherwise find a wallet with an empty balance slot, create a new one if no
+        // such wallet exists
+        let maybe_wallet = self.find_wallet_with_empty_balance().await?;
+        let wallet = match maybe_wallet {
+            Some(wallet) => wallet,
             None => {
                 info!("creating new wallet for {mint}");
-                self.create_new_wallet().await
+                self.create_new_wallet().await?
             },
-        }
+        };
+
+        self.add_mint_to_wallet(&wallet.id, mint).await?;
+        Ok(wallet)
     }
 
     /// Create a new wallet for managing a given mint