core, graph, node: use ChainName instead of String

Author: isum

core/src/amp_subgraph/manager.rs

@@ -9,7 +9,7 @@ use graph::{
     components::{
         link_resolver::{LinkResolver, LinkResolverContext},
         metrics::MetricsRegistry,
-        network_provider::{AmpChainConfig, AmpClients},
+        network_provider::{AmpChainConfig, AmpClients, ChainName},
         store::{DeploymentLocator, SubgraphStore},
         subgraph::SubgraphInstanceManager,
     },
@@ -35,7 +35,7 @@ pub struct Manager<SS, NC> {
     subgraph_store: Arc<SS>,
     link_resolver: Arc<dyn LinkResolver>,
     amp_clients: AmpClients<NC>,
-    amp_chain_configs: HashMap<String, AmpChainConfig>,
+    amp_chain_configs: HashMap<ChainName, AmpChainConfig>,
 }
 
 impl<SS, NC> Manager<SS, NC>
@@ -52,7 +52,7 @@ where
         subgraph_store: Arc<SS>,
         link_resolver: Arc<dyn LinkResolver>,
         amp_clients: AmpClients<NC>,
-        amp_chain_configs: HashMap<String, AmpChainConfig>,
+        amp_chain_configs: HashMap<ChainName, AmpChainConfig>,
     ) -> Self {
         let logger = logger_factory.component_logger("AmpSubgraphManager", None);
         let logger_factory = logger_factory.with_parent(logger);
@@ -145,7 +145,7 @@ where
                         }
                     };
 
-                    let amp_context = network_name.as_deref().and_then(|chain| {
+                    let amp_context = network_name.as_ref().and_then(|chain| {
                         manager
                             .amp_chain_configs
                             .get(chain)

core/src/subgraph/instance_manager.rs

@@ -273,7 +273,7 @@ impl<S: SubgraphStore, AC: amp::Client> SubgraphInstanceManager<S, AC> {
         // Look up the per-chain Amp client based on the network from the
         // raw manifest (before the manifest is moved into parse).
         let amp_client = network_name_from_raw_manifest(&raw_manifest)
-            .as_deref()
+            .as_ref()
             .and_then(|network| self.amp_clients.get(network));
 
         let manifest =

core/src/subgraph/registrar.rs

@@ -3,6 +3,7 @@ use std::collections::{HashMap, HashSet};
 use async_trait::async_trait;
 use graph::amp;
 use graph::blockchain::{Blockchain, BlockchainKind, BlockchainMap};
+use graph::components::network_provider::ChainName;
 use graph::components::{
     link_resolver::LinkResolverContext,
     network_provider::{AmpChainConfig, AmpChainNames, AmpClients},
@@ -26,7 +27,7 @@ pub struct SubgraphRegistrar<P, S, SM, AC> {
     store: Arc<S>,
     subscription_manager: Arc<SM>,
     amp_clients: AmpClients<AC>,
-    amp_chain_configs: HashMap<String, AmpChainConfig>,
+    amp_chain_configs: HashMap<ChainName, AmpChainConfig>,
     chains: Arc<BlockchainMap>,
     node_id: NodeId,
     version_switching_mode: SubgraphVersionSwitchingMode,
@@ -49,7 +50,7 @@ where
         store: Arc<S>,
         subscription_manager: Arc<SM>,
         amp_clients: AmpClients<AC>,
-        amp_chain_configs: HashMap<String, AmpChainConfig>,
+        amp_chain_configs: HashMap<ChainName, AmpChainConfig>,
         chains: Arc<BlockchainMap>,
         node_id: NodeId,
         version_switching_mode: SubgraphVersionSwitchingMode,
@@ -302,15 +303,13 @@ where
 
         // Extract the network name from the raw manifest and resolve the
         // per-chain Amp client (if any).
-        let resolved_amp_chain = network_name_from_raw_manifest(&raw).map(|network| {
-            let resolved = self.amp_chain_names.resolve(&Word::from(network));
-            resolved.to_string()
-        });
+        let resolved_amp_chain = network_name_from_raw_manifest(&raw)
+            .map(|network| self.amp_chain_names.resolve(&network));
         let amp_client = resolved_amp_chain
-            .as_deref()
+            .as_ref()
             .and_then(|chain| self.amp_clients.get(chain));
         let amp_context = resolved_amp_chain
-            .as_deref()
+            .as_ref()
             .and_then(|chain| self.amp_chain_configs.get(chain))
             .map(|cfg| cfg.context());
 

graph/src/components/network_provider/mod.rs

@@ -55,18 +55,18 @@ impl AmpChainConfig {
 /// instead of a single global `Option<Arc<AC>>`. Use `get(chain_name)` to
 /// retrieve the client for a specific chain.
 pub struct AmpClients<AC> {
-    clients: HashMap<String, Arc<AC>>,
+    clients: HashMap<ChainName, Arc<AC>>,
 }
 
 impl<AC> AmpClients<AC> {
     /// Creates a new `AmpClients` from a map of chain names to clients.
-    pub fn new(clients: HashMap<String, Arc<AC>>) -> Self {
+    pub fn new(clients: HashMap<ChainName, Arc<AC>>) -> Self {
         Self { clients }
     }
 
     /// Returns the Amp client for the given chain, or `None` if no client
     /// is configured for that chain.
-    pub fn get(&self, chain_name: &str) -> Option<Arc<AC>> {
+    pub fn get(&self, chain_name: &ChainName) -> Option<Arc<AC>> {
         self.clients.get(chain_name).cloned()
     }
 
@@ -166,26 +166,26 @@ mod tests {
     #[test]
     fn amp_clients_returns_client_for_configured_chain() {
         let mut map = HashMap::new();
-        map.insert("mainnet".to_string(), Arc::new(42u32));
+        map.insert(ChainName::from("mainnet"), Arc::new(42u32));
         let clients = AmpClients::new(map);
-        let client = clients.get("mainnet");
+        let client = clients.get(&ChainName::from("mainnet"));
         assert!(client.is_some());
         assert_eq!(*client.unwrap(), 42);
     }
 
     #[test]
     fn amp_clients_returns_none_for_unconfigured_chain() {
-        let map: HashMap<String, Arc<u32>> = HashMap::new();
+        let map: HashMap<ChainName, Arc<u32>> = HashMap::new();
         let clients = AmpClients::new(map);
-        assert!(clients.get("mainnet").is_none());
+        assert!(clients.get(&ChainName::from("mainnet")).is_none());
     }
 
     /// Verifies the condition that causes Amp manager registration:
     /// `!amp_clients.is_empty()` is true when at least one chain has config.
     #[test]
     fn amp_manager_registered_when_chain_has_config() {
         let mut map = HashMap::new();
-        map.insert("mainnet".to_string(), Arc::new(42u32));
+        map.insert(ChainName::from("mainnet"), Arc::new(42u32));
         let clients = AmpClients::new(map);
         assert!(
             !clients.is_empty(),
@@ -210,12 +210,12 @@ mod tests {
     #[test]
     fn amp_clients_error_for_unconfigured_amp_chain() {
         let mut map = HashMap::new();
-        map.insert("mainnet".to_string(), Arc::new(1u32));
+        map.insert(ChainName::from("mainnet"), Arc::new(1u32));
         let clients = AmpClients::new(map);
 
         // "matic" is not configured.
         let result = clients
-            .get("matic")
+            .get(&ChainName::from("matic"))
             .ok_or_else(|| "Amp is not configured for chain 'matic'".to_string());
         assert!(result.is_err());
         assert_eq!(

graph/src/data/subgraph/mod.rs

@@ -14,7 +14,11 @@ pub mod status;
 
 pub use features::{SubgraphFeature, SubgraphFeatureValidationError};
 
-use crate::{cheap_clone::CheapClone, components::store::BLOCK_NUMBER_MAX, object};
+use crate::{
+    cheap_clone::CheapClone,
+    components::{network_provider::ChainName, store::BLOCK_NUMBER_MAX},
+    object,
+};
 use anyhow::{anyhow, Context, Error};
 use futures03::future::try_join_all;
 use itertools::Itertools;
@@ -1393,15 +1397,15 @@ fn display_vector(input: &[impl std::fmt::Display]) -> impl std::fmt::Display {
 ///
 /// Navigates `dataSources[0].network` and returns the network name as an owned string,
 /// or `None` if any step in the path is missing.
-pub fn network_name_from_raw_manifest(raw: &serde_yaml::Mapping) -> Option<String> {
+pub fn network_name_from_raw_manifest(raw: &serde_yaml::Mapping) -> Option<ChainName> {
     use serde_yaml::Value;
     raw.get(Value::String("dataSources".to_owned()))
         .and_then(|ds| ds.as_sequence())
         .and_then(|ds| ds.first())
         .and_then(|ds| ds.as_mapping())
         .and_then(|ds| ds.get(Value::String("network".to_owned())))
         .and_then(|n| n.as_str())
-        .map(|s| s.to_owned())
+        .map(|s| s.into())
 }
 
 #[test]
@@ -1473,7 +1477,7 @@ mod tests {
 
         assert_eq!(
             network_name_from_raw_manifest(&raw),
-            Some("mainnet".to_string())
+            Some(ChainName::from("mainnet"))
         );
     }
 

node/src/config.rs

@@ -139,7 +139,7 @@ impl Config {
     /// that has an `[amp]` section. The `AmpConfig.address` string is parsed
     /// into a `Uri`; this is expected to always succeed because
     /// `ChainSection::validate` already rejects invalid URIs.
-    pub fn amp_chain_configs(&self) -> Result<HashMap<String, AmpChainConfig>> {
+    pub fn amp_chain_configs(&self) -> Result<HashMap<ChainName, AmpChainConfig>> {
         let mut map = HashMap::new();
         for (chain_name, chain) in &self.chains.chains {
             if let Some(amp) = &chain.amp {
@@ -152,7 +152,7 @@ impl Config {
                     )
                 })?;
                 map.insert(
-                    chain_name.clone(),
+                    chain_name.as_str().into(),
                     AmpChainConfig {
                         address: uri,
                         token: amp.token.clone(),
@@ -1375,8 +1375,8 @@ mod tests {
     use crate::config::{default_polling_interval, ChainSection, Web3Rule};
 
     use super::{
-        AmpConfig, Chain, Config, FirehoseProvider, Provider, ProviderDetails, Shard, Transport,
-        Web3Provider,
+        AmpConfig, Chain, ChainName, Config, FirehoseProvider, Provider, ProviderDetails, Shard,
+        Transport, Web3Provider,
     };
     use graph::blockchain::BlockchainKind;
     use graph::firehose::SubgraphLimit;
@@ -2425,9 +2425,11 @@ fdw_pool_size = [
 
         // Only mainnet (with amp) should be in the map
         assert_eq!(map.len(), 1);
-        assert!(!map.contains_key("sepolia"));
+        assert!(!map.contains_key(&ChainName::from("sepolia")));
 
-        let mainnet = map.get("mainnet").expect("mainnet should be in map");
+        let mainnet = map
+            .get(&ChainName::from("mainnet"))
+            .expect("mainnet should be in map");
         assert_eq!(mainnet.address.to_string(), "http://localhost:50051/");
         assert_eq!(mainnet.token.as_deref(), Some("my-token"));
         assert_eq!(mainnet.context_dataset, "eth");
@@ -2563,7 +2565,9 @@ fdw_pool_size = [
         };
 
         let map = config.amp_chain_configs().unwrap();
-        let mainnet = map.get("mainnet").expect("mainnet should be in map");
+        let mainnet = map
+            .get(&ChainName::from("mainnet"))
+            .expect("mainnet should be in map");
 
         // Identifiers with special characters should be double-quoted
         assert_eq!(mainnet.context_dataset, "\"ns/data@v1\"");
@@ -2603,7 +2607,9 @@ fdw_pool_size = [
         };
 
         let map = config.amp_chain_configs().unwrap();
-        let mainnet = map.get("mainnet").expect("mainnet should be in map");
+        let mainnet = map
+            .get(&ChainName::from("mainnet"))
+            .expect("mainnet should be in map");
 
         // Simple identifiers should be lowercased and unquoted
         assert_eq!(mainnet.context_dataset, "eth");

node/src/launcher.rs

@@ -1,13 +1,17 @@
 use std::{
+    collections::HashMap,
     io::{BufRead, BufReader},
     path::Path,
     time::Duration,
 };
 
 use anyhow::Result;
 use git_testament::{git_testament, render_testament};
-use graph::components::link_resolver::{ArweaveClient, FileSizeLimit};
 use graph::components::subgraph::Settings;
+use graph::components::{
+    link_resolver::{ArweaveClient, FileSizeLimit},
+    network_provider::ChainName,
+};
 use graph::data::graphql::load_manager::LoadManager;
 use graph::endpoint::EndpointMetrics;
 use graph::env::EnvVars;
@@ -276,10 +280,7 @@ fn build_subgraph_registrar<AC>(
     arweave_service: ArweaveService,
     ipfs_service: IpfsService,
     amp_clients: AmpClients<AC>,
-    amp_chain_configs: std::collections::HashMap<
-        String,
-        graph::components::network_provider::AmpChainConfig,
-    >,
+    amp_chain_configs: HashMap<ChainName, graph::components::network_provider::AmpChainConfig>,
     cancel_token: CancellationToken,
     amp_chain_names: Arc<AmpChainNames>,
 ) -> Arc<

server/index-node/src/resolver.rs

@@ -524,7 +524,7 @@ where
         // Extract the network name from the raw yaml to look up the
         // per-chain Amp client.
         let amp_client = network_name_from_raw_manifest(&raw_yaml)
-            .as_deref()
+            .as_ref()
             .and_then(|network| self.amp_clients.get(network));
 
         let max_spec_version = ENV_VARS.max_spec_version.clone();