From 0cab25d8d710241cedcf8fd79ed72174893c410d Mon Sep 17 00:00:00 2001
From: Yurii Koba <kobayurii88@gmail.com>
Date: Wed, 2 Oct 2024 17:41:06 +0300
Subject: [PATCH] imrovement

---
 configuration/src/configs/lake.rs | 18 ++++++++++++------
 rpc-server/src/config.rs          |  3 ++-
 state-indexer/src/main.rs         |  5 ++++-
 tx-indexer/src/main.rs            |  2 +-
 4 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/configuration/src/configs/lake.rs b/configuration/src/configs/lake.rs
index caf5de5c..ea6c00aa 100644
--- a/configuration/src/configs/lake.rs
+++ b/configuration/src/configs/lake.rs
@@ -1,8 +1,7 @@
+use crate::configs::deserialize_optional_data_or_env;
 use near_lake_framework::near_indexer_primitives::near_primitives;
 use serde_derive::Deserialize;
 
-use crate::configs::deserialize_optional_data_or_env;
-
 #[derive(Debug, Clone)]
 pub struct LakeConfig {
     pub num_threads: Option<u64>,
@@ -12,13 +11,20 @@ impl LakeConfig {
     pub async fn lake_config(
         &self,
         start_block_height: near_primitives::types::BlockHeight,
+        chain_id: crate::ChainId,
     ) -> anyhow::Result<near_lake_framework::FastNearConfig> {
-        let config_builder = near_lake_framework::FastNearConfigBuilder::default();
+        let mut config_builder = near_lake_framework::FastNearConfigBuilder::default();
+        match chain_id {
+            crate::ChainId::Mainnet => config_builder = config_builder.mainnet(),
+            // Testnet is the default chain for other chain_id
+            _ => config_builder = config_builder.testnet(),
+        };
+        if let Some(num_threads) = self.num_threads {
+            config_builder = config_builder.num_threads(num_threads);
+        };
         Ok(config_builder
-            .mainnet()
             .start_block_height(start_block_height)
-            .build()
-            .expect("Failed to build LakeConfig"))
+            .build()?)
     }
 }
 
diff --git a/rpc-server/src/config.rs b/rpc-server/src/config.rs
index 32c537b8..5999ee51 100644
--- a/rpc-server/src/config.rs
+++ b/rpc-server/src/config.rs
@@ -44,7 +44,7 @@ impl GenesisInfo {
 
 #[derive(Clone)]
 pub struct ServerContext {
-    /// Lake s3 client
+    /// Fastnear client
     pub fastnear_client: near_lake_framework::fastnear_client::FastNearClient,
     /// Database manager
     pub db_manager: std::sync::Arc<Box<dyn database::ReaderDbManager + Sync + Send + 'static>>,
@@ -109,6 +109,7 @@ impl ServerContext {
                     .optimistic_cache_block()
                     .await
                     .block_height,
+                rpc_server_config.general.chain_id.clone(),
             )
             .await?
             .client();
diff --git a/state-indexer/src/main.rs b/state-indexer/src/main.rs
index df2a3f8d..534e9b2f 100644
--- a/state-indexer/src/main.rs
+++ b/state-indexer/src/main.rs
@@ -36,7 +36,10 @@ async fn main() -> anyhow::Result<()> {
     )
     .await?;
 
-    let lake_config = indexer_config.lake_config.lake_config(start_block_height).await?;
+    let lake_config = indexer_config
+        .lake_config
+        .lake_config(start_block_height, indexer_config.general.chain_id.clone())
+        .await?;
     let (sender, stream) = near_lake_framework::streamer(lake_config);
 
     // Initiate metrics http server
diff --git a/tx-indexer/src/main.rs b/tx-indexer/src/main.rs
index 97d5ffda..c8f8d45b 100644
--- a/tx-indexer/src/main.rs
+++ b/tx-indexer/src/main.rs
@@ -63,7 +63,7 @@ async fn main() -> anyhow::Result<()> {
     tracing::info!(target: INDEXER, "Generating LakeConfig...");
     let lake_config = indexer_config
         .lake_config
-        .lake_config(start_block_height)
+        .lake_config(start_block_height, indexer_config.general.chain_id.clone())
         .await?;
 
     tracing::info!(target: INDEXER, "Creating cache storage...");