Add initial Address balance endpoint support to the API Server

api-server/api-server-common/src/lib.rs

@@ -16,6 +16,26 @@
 pub mod storage;
 
 use clap::Parser;
+use common::chain::config::ChainType;
+
+#[derive(clap::ValueEnum, Clone, Debug)]
+pub enum Network {
+    Mainnet,
+    Testnet,
+    Regtest,
+    Signet,
+}
+
+impl From<Network> for ChainType {
+    fn from(value: Network) -> Self {
+        match value {
+            Network::Mainnet => ChainType::Mainnet,
+            Network::Testnet => ChainType::Testnet,
+            Network::Regtest => ChainType::Regtest,
+            Network::Signet => ChainType::Signet,
+        }
+    }
+}
 
 #[derive(Parser, Debug)]
 pub struct PostgresConfig {

api-server/api-server-common/src/storage/impls/in_memory/mod.rs

@@ -15,21 +15,23 @@
 
 pub mod transactional;
 
-use std::collections::BTreeMap;
-
+use crate::storage::storage_api::{block_aux_data::BlockAuxData, ApiServerStorageError};
 use common::{
     chain::{Block, ChainConfig, GenBlock, SignedTransaction, Transaction},
-    primitives::{BlockHeight, Id},
+    primitives::{Amount, BlockHeight, Id},
+};
+use std::{
+    collections::BTreeMap,
+    ops::Bound::{Excluded, Unbounded},
 };
-
-use crate::storage::storage_api::{block_aux_data::BlockAuxData, ApiServerStorageError};
 
 use super::CURRENT_STORAGE_VERSION;
 
 #[derive(Debug, Clone)]
 struct ApiServerInMemoryStorage {
     block_table: BTreeMap<Id<Block>, Block>,
     block_aux_data_table: BTreeMap<Id<Block>, BlockAuxData>,
+    address_balance_table: BTreeMap<String, BTreeMap<BlockHeight, Amount>>,
     main_chain_blocks_table: BTreeMap<BlockHeight, Id<Block>>,
     transaction_table: BTreeMap<Id<Transaction>, (Option<Id<Block>>, SignedTransaction)>,
     best_block: (BlockHeight, Id<GenBlock>),
@@ -41,6 +43,7 @@ impl ApiServerInMemoryStorage {
         let mut result = Self {
             block_table: BTreeMap::new(),
             block_aux_data_table: BTreeMap::new(),
+            address_balance_table: BTreeMap::new(),
             main_chain_blocks_table: BTreeMap::new(),
             transaction_table: BTreeMap::new(),
             best_block: (0.into(), chain_config.genesis_block_id()),
@@ -56,6 +59,13 @@ impl ApiServerInMemoryStorage {
         Ok(true)
     }
 
+    fn get_address_balance(&self, address: &str) -> Result<Option<Amount>, ApiServerStorageError> {
+        self.address_balance_table.get(address).map_or_else(
+            || Ok(None),
+            |balance| Ok(balance.last_key_value().map(|(_, &v)| v)),
+        )
+    }
+
     fn get_block(&self, block_id: Id<Block>) -> Result<Option<Block>, ApiServerStorageError> {
         let block_result = self.block_table.get(&block_id);
         let block = match block_result {
@@ -124,6 +134,40 @@ impl ApiServerInMemoryStorage {
         Ok(())
     }
 
+    fn del_address_balance_above_height(
+        &mut self,
+        block_height: BlockHeight,
+    ) -> Result<(), ApiServerStorageError> {
+        // Inefficient, but acceptable for testing with InMemoryStorage
+
+        self.address_balance_table.iter_mut().for_each(|(_, balance)| {
+            balance
+                .range((Excluded(block_height), Unbounded))
+                .map(|b| b.0.into_int())
+                .collect::<Vec<_>>()
+                .iter()
+                .for_each(|&b| {
+                    balance.remove(&BlockHeight::new(b));
+                })
+        });
+
+        Ok(())
+    }
+
+    fn set_address_balance_at_height(
+        &mut self,
+        address: &str,
+        amount: Amount,
+        block_height: BlockHeight,
+    ) -> Result<(), ApiServerStorageError> {
+        self.address_balance_table
+            .entry(address.to_string())
+            .or_default()
+            .insert(block_height, amount);
+
+        Ok(())
+    }
+
     fn set_block(
         &mut self,
         block_id: Id<Block>,

api-server/api-server-common/src/storage/impls/in_memory/transactional/read.rs

@@ -15,7 +15,7 @@
 
 use common::{
     chain::{Block, GenBlock, SignedTransaction, Transaction},
-    primitives::{BlockHeight, Id},
+    primitives::{Amount, BlockHeight, Id},
 };
 
 use crate::storage::storage_api::{
@@ -30,6 +30,13 @@ impl<'t> ApiServerStorageRead for ApiServerInMemoryStorageTransactionalRo<'t> {
         self.transaction.is_initialized()
     }
 
+    async fn get_address_balance(
+        &self,
+        address: &str,
+    ) -> Result<Option<Amount>, ApiServerStorageError> {
+        self.transaction.get_address_balance(address)
+    }
+
     async fn get_block(&self, block_id: Id<Block>) -> Result<Option<Block>, ApiServerStorageError> {
         self.transaction.get_block(block_id)
     }

api-server/api-server-common/src/storage/impls/in_memory/transactional/write.rs

@@ -15,7 +15,7 @@
 
 use common::{
     chain::{Block, ChainConfig, GenBlock, SignedTransaction, Transaction},
-    primitives::{BlockHeight, Id},
+    primitives::{Amount, BlockHeight, Id},
 };
 
 use crate::storage::storage_api::{
@@ -34,6 +34,22 @@ impl<'t> ApiServerStorageWrite for ApiServerInMemoryStorageTransactionalRw<'t> {
         self.transaction.initialize_storage(chain_config)
     }
 
+    async fn del_address_balance_above_height(
+        &mut self,
+        block_height: BlockHeight,
+    ) -> Result<(), ApiServerStorageError> {
+        self.transaction.del_address_balance_above_height(block_height)
+    }
+
+    async fn set_address_balance_at_height(
+        &mut self,
+        address: &str,
+        amount: Amount,
+        block_height: BlockHeight,
+    ) -> Result<(), ApiServerStorageError> {
+        self.transaction.set_address_balance_at_height(address, amount, block_height)
+    }
+
     async fn set_block(
         &mut self,
         block_id: Id<Block>,
@@ -93,6 +109,13 @@ impl<'t> ApiServerStorageRead for ApiServerInMemoryStorageTransactionalRw<'t> {
         Ok(Some(self.transaction.get_storage_version()?))
     }
 
+    async fn get_address_balance(
+        &self,
+        address: &str,
+    ) -> Result<Option<Amount>, ApiServerStorageError> {
+        self.transaction.get_address_balance(address)
+    }
+
     async fn get_best_block(&self) -> Result<(BlockHeight, Id<GenBlock>), ApiServerStorageError> {
         self.transaction.get_best_block()
     }

api-server/api-server-common/src/storage/impls/postgres/queries.rs

@@ -18,7 +18,7 @@ use serialization::{DecodeAll, Encode};
 
 use common::{
     chain::{Block, ChainConfig, GenBlock, SignedTransaction, Transaction},
-    primitives::{BlockHeight, Id},
+    primitives::{Amount, BlockHeight, Id},
 };
 use tokio_postgres::NoTls;
 
@@ -113,6 +113,80 @@ impl<'a, 'b> QueryFromConnection<'a, 'b> {
         Ok(Some(version))
     }
 
+    pub async fn get_address_balance(
+        &self,
+        address: &str,
+    ) -> Result<Option<Amount>, ApiServerStorageError> {
+        self.tx
+            .query_opt(
+                r#"
+                    SELECT amount
+                    FROM ml_address_balance
+                    WHERE address = $1
+                    ORDER BY block_height DESC
+                    LIMIT 1;
+                "#,
+                &[&address],
+            )
+            .await
+            .map_err(|e| ApiServerStorageError::LowLevelStorageError(e.to_string()))?
+            .map_or_else(
+                || Ok(None),
+                |row| {
+                    let amount: Vec<u8> = row.get(0);
+                    let amount = Amount::decode_all(&mut amount.as_slice()).map_err(|e| {
+                        ApiServerStorageError::DeserializationError(format!(
+                            "Amount deserialization failed: {}",
+                            e
+                        ))
+                    })?;
+
+                    Ok(Some(amount))
+                },
+            )
+    }
+
+    pub async fn del_address_balance_above_height(
+        &mut self,
+        block_height: BlockHeight,
+    ) -> Result<(), ApiServerStorageError> {
+        let height = Self::block_height_to_postgres_friendly(block_height);
+
+        self.tx
+            .execute(
+                "DELETE FROM ml_address_balance WHERE block_height > $1;",
+                &[&height],
+            )
+            .await
+            .map_err(|e| ApiServerStorageError::LowLevelStorageError(e.to_string()))?;
+
+        Ok(())
+    }
+
+    pub async fn set_address_balance_at_height(
+        &mut self,
+        address: &str,
+        amount: Amount,
+        block_height: BlockHeight,
+    ) -> Result<(), ApiServerStorageError> {
+        let height = Self::block_height_to_postgres_friendly(block_height);
+
+        self.tx
+            .execute(
+                r#"
+                    INSERT INTO ml_address_balance (address, block_height, amount)
+                    VALUES ($1, $2, $3)
+                    ON CONFLICT (address, block_height)
+                    DO UPDATE SET amount = $3;
+                "#,
+                &[&address.to_string(), &height, &amount.encode()],
+            )
+            .await
+            .map_err(|e| ApiServerStorageError::LowLevelStorageError(e.to_string()))?;
+
+        Ok(())
+    }
+
     pub async fn get_best_block(
         &mut self,
     ) -> Result<(BlockHeight, Id<GenBlock>), ApiServerStorageError> {
@@ -203,6 +277,16 @@ impl<'a, 'b> QueryFromConnection<'a, 'b> {
         )
         .await?;
 
+        self.just_execute(
+            "CREATE TABLE ml_address_balance (
+                    address TEXT NOT NULL,
+                    block_height bigint NOT NULL,
+                    amount bytea NOT NULL,
+                    PRIMARY KEY (address, block_height)
+                );",
+        )
+        .await?;
+
         self.just_execute(
             "CREATE TABLE ml_block_aux_data (
                     block_id bytea PRIMARY KEY,

api-server/api-server-common/src/storage/impls/postgres/transactional/read.rs

@@ -39,6 +39,19 @@ impl<'a> ApiServerStorageRead for ApiServerPostgresTransactionalRo<'a> {
         Ok(res)
     }
 
+    async fn get_address_balance(
+        &self,
+        address: &str,
+    ) -> Result<
+        Option<common::primitives::Amount>,
+        crate::storage::storage_api::ApiServerStorageError,
+    > {
+        let conn = QueryFromConnection::new(self.connection.as_ref().expect(CONN_ERR));
+        let res = conn.get_address_balance(address).await?;
+
+        Ok(res)
+    }
+
     async fn get_best_block(
         &self,
     ) -> Result<

api-server/api-server-common/src/storage/impls/postgres/transactional/write.rs

@@ -15,7 +15,7 @@
 
 use common::{
     chain::{Block, ChainConfig, GenBlock, SignedTransaction, Transaction},
-    primitives::{BlockHeight, Id},
+    primitives::{Amount, BlockHeight, Id},
 };
 
 use crate::storage::{
@@ -40,6 +40,28 @@ impl<'a> ApiServerStorageWrite for ApiServerPostgresTransactionalRw<'a> {
         Ok(())
     }
 
+    async fn del_address_balance_above_height(
+        &mut self,
+        block_height: BlockHeight,
+    ) -> Result<(), ApiServerStorageError> {
+        let mut conn = QueryFromConnection::new(self.connection.as_ref().expect(CONN_ERR));
+        conn.del_address_balance_above_height(block_height).await?;
+
+        Ok(())
+    }
+
+    async fn set_address_balance_at_height(
+        &mut self,
+        address: &str,
+        amount: Amount,
+        block_height: BlockHeight,
+    ) -> Result<(), ApiServerStorageError> {
+        let mut conn = QueryFromConnection::new(self.connection.as_ref().expect(CONN_ERR));
+        conn.set_address_balance_at_height(address, amount, block_height).await?;
+
+        Ok(())
+    }
+
     async fn set_best_block(
         &mut self,
         block_height: BlockHeight,
@@ -123,6 +145,16 @@ impl<'a> ApiServerStorageRead for ApiServerPostgresTransactionalRw<'a> {
         Ok(res)
     }
 
+    async fn get_address_balance(
+        &self,
+        address: &str,
+    ) -> Result<Option<Amount>, ApiServerStorageError> {
+        let conn = QueryFromConnection::new(self.connection.as_ref().expect(CONN_ERR));
+        let res = conn.get_address_balance(address).await?;
+
+        Ok(res)
+    }
+
     async fn get_best_block(&self) -> Result<(BlockHeight, Id<GenBlock>), ApiServerStorageError> {
         let mut conn = QueryFromConnection::new(self.connection.as_ref().expect(CONN_ERR));
         let res = conn.get_best_block().await?;