bug: insert balances in chunks

chain/src/main.rs

@@ -362,7 +362,7 @@ async fn initial_query(
         conn.build_transaction()
             .read_write()
             .run(|transaction_conn| {
-                repository::balance::insert_balance(
+                repository::balance::insert_balance_in_chunks(
                     transaction_conn,
                     balances,
                 )?;

chain/src/repository/balance.rs

@@ -1,8 +1,18 @@
+use diesel::sql_types::{BigInt, Integer};
 use diesel::upsert::excluded;
-use diesel::{ExpressionMethods, PgConnection, RunQueryDsl};
+use diesel::{
+    sql_query, ExpressionMethods, PgConnection, QueryableByName, RunQueryDsl,
+};
 use orm::balances::BalancesInsertDb;
 use orm::schema::balances;
 use shared::balance::Balances;
+pub const MAX_PARAM_SIZE: u16 = u16::MAX;
+
+#[derive(QueryableByName)]
+struct BalanceColCount {
+    #[diesel(sql_type = BigInt)]
+    count: i64,
+}
 
 pub fn insert_balance(
     transaction_conn: &mut PgConnection,
@@ -27,6 +37,32 @@ pub fn insert_balance(
     anyhow::Ok(())
 }
 
+pub fn insert_balance_in_chunks(
+    transaction_conn: &mut PgConnection,
+    balances: Balances,
+) -> anyhow::Result<()> {
+    let balances_col_count = sql_query(
+        "SELECT COUNT(*)
+            FROM information_schema.columns
+            WHERE table_schema = 'public'
+            AND table_name = 'balances';",
+    )
+    .get_result::<BalanceColCount>(transaction_conn)?;
+
+    for chunk in balances
+        // We have to devide MAX_PARAM_SIZE by the number of columns in the balances table to get
+        // the correct number of rows in the chunk.
+        .chunks((MAX_PARAM_SIZE as i64 / balances_col_count.count) as usize)
+    {
+        transaction_conn
+            .build_transaction()
+            .read_write()
+            .run(|conn| insert_balance(conn, chunk.to_vec()))?;
+    }
+
+    anyhow::Ok(())
+}
+
 #[cfg(test)]
 mod tests {
 
@@ -303,6 +339,48 @@ mod tests {
         .expect("Failed to run test");
     }
 
+    /// Test that we can insert more than u16::MAX balances
+    #[tokio::test]
+    async fn test_insert_balance_in_chunks_with_max_param_size_plus_one() {
+        let config = TestConfig::parse();
+        let db = TestDb::new(&config);
+
+        db.run_test_without_transaction(|conn| {
+            let mps = MAX_PARAM_SIZE as u32;
+
+            let balances =
+                (0..mps + 1).map(|_| Balance::fake()).collect::<Vec<_>>();
+
+            let res = insert_balance_in_chunks(conn, balances)?;
+
+            assert_eq!(res, ());
+
+            anyhow::Ok(())
+        })
+        .await
+        .expect("Failed to run test");
+    }
+
+    /// Test that we can insert less than u16::MAX balances using chunks
+    #[tokio::test]
+    async fn test_insert_balance_in_chunks_with_1000_params() {
+        let config = TestConfig::parse();
+        let db = TestDb::new(&config);
+
+        db.run_test_without_transaction(|conn| {
+            let balances =
+                (0..1000).map(|_| Balance::fake()).collect::<Vec<_>>();
+
+            let res = insert_balance_in_chunks(conn, balances)?;
+
+            assert_eq!(res, ());
+
+            anyhow::Ok(())
+        })
+        .await
+        .expect("Failed to run test");
+    }
+
     fn seed_balance(
         conn: &mut PgConnection,
         balances: Vec<Balance>,

chain/src/test_db.rs

@@ -74,6 +74,25 @@ impl TestDb {
 
         anyhow::Ok(())
     }
+
+    pub async fn run_test_without_transaction(
+        &self,
+        test: impl Fn(&mut PgConnection) -> anyhow::Result<()>
+            + namada_sdk::MaybeSend
+            + 'static,
+    ) -> anyhow::Result<()> {
+        let conn = &mut self.pool.get().await.unwrap();
+
+        run_migrations(conn).await.unwrap();
+
+        conn.interact(move |conn| test(conn))
+            .await
+            .context_db_interact_error()
+            .and_then(identity)
+            .into_db_error()?;
+
+        anyhow::Ok(())
+    }
 }
 
 impl Drop for TestDb {

orm/src/balances.rs

@@ -19,18 +19,11 @@ pub type BalanceDb = BalancesInsertDb;
 
 impl BalancesInsertDb {
     pub fn from_balance(balance: Balance) -> Self {
-        let asd = Self {
+        Self {
             owner: balance.owner.to_string(),
             token: balance.token.to_string(),
             raw_amount: BigDecimal::from_str(&balance.amount.to_string())
                 .expect("Invalid amount"),
-        };
-
-        println!(
-            "Balance_test: {:?}",
-            BigDecimal::from_str(&balance.amount.to_string())
-        );
-
-        asd
+        }
     }
 }