Address PR comments

crates/sui-graphql-rpc/src/data.rs

@@ -1,6 +1,7 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
+pub(crate) mod move_registry_data_loader;
 pub(crate) mod package_resolver;
 pub(crate) mod pg;
 

crates/sui-graphql-rpc/src/types/dot_move/data_loader.rs → crates/sui-graphql-rpc/src/data/move_registry_data_loader.rs

@@ -6,11 +6,16 @@ use std::{collections::HashMap, str::FromStr, sync::Arc};
 use async_graphql::dataloader::{DataLoader, Loader};
 use serde::{Deserialize, Serialize};
 
-use crate::{config::MoveRegistryConfig, error::Error, types::base64::Base64};
-
-use super::{
-    error::MoveRegistryError,
-    on_chain::{AppRecord, Name},
+use crate::{
+    config::MoveRegistryConfig,
+    error::Error,
+    types::{
+        base64::Base64,
+        dot_move::{
+            error::MoveRegistryError,
+            on_chain::{AppRecord, Name},
+        },
+    },
 };
 
 /// GraphQL fragment to query the values of the dynamic fields.
@@ -21,24 +26,23 @@ pub(crate) struct ExternalNamesLoader {
     client: reqwest::Client,
     config: MoveRegistryConfig,
 }
+
 /// Helper types for accessing a shared `DataLoader` instance.
 #[derive(Clone)]
-pub(crate) struct DotMoveDataLoader(pub Arc<DataLoader<ExternalNamesLoader>>);
+pub(crate) struct MoveRegistryDataLoader(pub Arc<DataLoader<ExternalNamesLoader>>);
 
 impl ExternalNamesLoader {
-    pub(crate) fn new(config: &MoveRegistryConfig) -> Self {
+    pub(crate) fn new(config: MoveRegistryConfig) -> Self {
         Self {
             client: reqwest::Client::new(),
-            config: config.clone(),
+            config,
         }
     }
 
-    /// Constructs the GraphQL Query to query the names on a mainnet graphql endpoint.
-    pub(crate) fn construct_names_graphql_query(
-        &self,
-        names: &[Name],
-        mapping: &mut HashMap<Name, usize>,
-    ) -> String {
+    /// Constructs the GraphQL Query to query the names on an external graphql endpoint.
+    fn construct_names_graphql_query(&self, names: &[Name]) -> (String, HashMap<Name, usize>) {
+        let mut mapping: HashMap<Name, usize> = HashMap::new();
+
         let mut result = format!(r#"{{ owner(address: "{}") {{"#, self.config.registry_id);
 
         // we create the GraphQL query keys with a `fetch_{id}` prefix, which is accepted on graphql fields.
@@ -48,27 +52,26 @@ impl ExternalNamesLoader {
             // retain the mapping here (id to bcs representation, so we can pick the right response later on)
             mapping.insert(name.clone(), index);
 
-            let field_str = format!(
+            result.push_str(&format!(
                 r#"{}: dynamicField(name: {{ type: "{}::name::Name", bcs: {} }}) {{ ...RECORD_VALUES }}"#,
                 fetch_key(&index),
                 self.config.package_address,
                 bcs_base64
-            );
-
-            result.push_str(&field_str);
+            ));
         }
 
         result.push_str("}} ");
         result.push_str(QUERY_FRAGMENT);
 
-        result
+        (result, mapping)
     }
 }
 
-impl DotMoveDataLoader {
-    pub(crate) fn new(config: &MoveRegistryConfig) -> Self {
+impl MoveRegistryDataLoader {
+    pub(crate) fn new(config: MoveRegistryConfig) -> Self {
+        let batch_size = config.page_limit as usize;
         let data_loader = DataLoader::new(ExternalNamesLoader::new(config), tokio::spawn)
-            .max_batch_size(config.page_limit as usize);
+            .max_batch_size(batch_size);
         Self(Arc::new(data_loader))
     }
 }
@@ -78,28 +81,29 @@ impl Loader<Name> for ExternalNamesLoader {
     type Value = AppRecord;
     type Error = Error;
 
-    /// This function queries the mainnet API to fetch the app records for the requested names.
+    /// This function queries the external API to fetch the app records for the requested names.
     /// This is part of the data loader, so all queries are bulked-up to the maximum of {config.page_limit}.
     /// We handle the cases where individual queries fail, to ensure that a failed query cannot affect
     /// a successful one.
     async fn load(&self, keys: &[Name]) -> Result<HashMap<Name, AppRecord>, Error> {
-        let Some(mainnet_api_url) = self.config.external_api_url.as_ref() else {
+        let Some(api_url) = self.config.external_api_url.as_ref() else {
             return Err(Error::MoveNameRegistry(
                 MoveRegistryError::ExternalApiUrlUnavailable,
             ));
         };
 
         let mut results: HashMap<Name, AppRecord> = HashMap::new();
-        let mut mapping: HashMap<Name, usize> = HashMap::new();
+
+        let (query, mapping) = self.construct_names_graphql_query(keys);
 
         let request_body = GraphQLRequest {
-            query: self.construct_names_graphql_query(keys, &mut mapping),
+            query,
             variables: serde_json::Value::Null,
         };
 
         let res = self
             .client
-            .post(mainnet_api_url)
+            .post(api_url)
             .json(&request_body)
             .send()
             .await
@@ -125,7 +129,7 @@ impl Loader<Name> for ExternalNamesLoader {
         let names = response_json.data.owner.names;
 
         for k in mapping.keys() {
-            // Safe unwrap: we inserted the keys in the mapping before.
+            // SAFETY: we inserted the keys in the mapping before.
             let idx = mapping.get(k).unwrap();
 
             let Some(Some(bcs)) = names.get(&fetch_key(idx)) else {
@@ -150,10 +154,7 @@ impl Loader<Name> for ExternalNamesLoader {
 
 impl Default for ExternalNamesLoader {
     fn default() -> Self {
-        Self {
-            client: reqwest::Client::new(),
-            config: MoveRegistryConfig::default(),
-        }
+        Self::new(MoveRegistryConfig::default())
     }
 }
 

crates/sui-graphql-rpc/src/server/builder.rs

@@ -7,12 +7,12 @@ use crate::config::{
     ConnectionConfig, ServiceConfig, Version, MAX_CONCURRENT_REQUESTS,
     RPC_TIMEOUT_ERR_SLEEP_RETRY_PERIOD,
 };
+use crate::data::move_registry_data_loader::MoveRegistryDataLoader;
 use crate::data::package_resolver::{DbPackageStore, PackageResolver};
 use crate::data::{DataLoader, Db};
 use crate::metrics::Metrics;
 use crate::mutation::Mutation;
 use crate::types::chain_identifier::ChainIdentifier;
-use crate::types::dot_move::data_loader::DotMoveDataLoader;
 use crate::types::move_object::IMoveObject;
 use crate::types::object::IObject;
 use crate::types::owner::IOwner;
@@ -429,7 +429,7 @@ impl ServerBuilder {
             .context_data(metrics.clone())
             .context_data(config.clone())
             .context_data(move_registry_config.clone())
-            .context_data(DotMoveDataLoader::new(&move_registry_config));
+            .context_data(MoveRegistryDataLoader::new(move_registry_config));
 
         if config.internal_features.feature_gate {
             builder = builder.extension(FeatureGate);

crates/sui-graphql-rpc/src/types/dot_move/mod.rs

@@ -1,7 +1,6 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-pub(crate) mod data_loader;
 pub(crate) mod error;
 pub(crate) mod named_move_package;
 pub(crate) mod named_type;

crates/sui-graphql-rpc/src/types/dot_move/named_move_package.rs

@@ -7,6 +7,7 @@ use async_graphql::Context;
 
 use crate::{
     config::{MoveRegistryConfig, ResolutionType},
+    data::move_registry_data_loader::MoveRegistryDataLoader,
     error::Error,
     types::{
         chain_identifier::ChainIdentifier, move_object::MoveObject, move_package::MovePackage,
@@ -15,7 +16,6 @@ use crate::{
 };
 
 use super::{
-    data_loader::DotMoveDataLoader,
     error::MoveRegistryError,
     on_chain::{AppInfo, AppRecord, VersionedName},
 };
@@ -35,94 +35,94 @@ impl NamedMovePackage {
 
         match config.resolution_type {
             ResolutionType::Internal => {
-                Self::query_internal(ctx, config, versioned, checkpoint_viewed_at).await
+                query_internal(ctx, config, versioned, checkpoint_viewed_at).await
             }
             ResolutionType::External => {
-                Self::query_external(ctx, config, versioned, checkpoint_viewed_at).await
+                query_external(ctx, config, versioned, checkpoint_viewed_at).await
             }
         }
     }
+}
 
-    async fn query_internal(
-        ctx: &Context<'_>,
-        config: &MoveRegistryConfig,
-        versioned: VersionedName,
-        checkpoint_viewed_at: u64,
-    ) -> Result<Option<MovePackage>, Error> {
-        let df_id = versioned.name.to_dynamic_field_id(config).map_err(|e| {
-            Error::Internal(format!("Failed to convert name to dynamic field id: {}", e))
-        })?;
-
-        let Some(df) =
-            MoveObject::query(ctx, df_id.into(), Object::latest_at(checkpoint_viewed_at)).await?
-        else {
-            return Ok(None);
-        };
-
-        let app_record = AppRecord::try_from(df.native)?;
-
-        let Some(app_info) = app_record.app_info else {
-            return Ok(None);
-        };
+async fn query_internal(
+    ctx: &Context<'_>,
+    config: &MoveRegistryConfig,
+    versioned: VersionedName,
+    checkpoint_viewed_at: u64,
+) -> Result<Option<MovePackage>, Error> {
+    let df_id = versioned.name.to_dynamic_field_id(config).map_err(|e| {
+        Error::Internal(format!("Failed to convert name to dynamic field id: {}", e))
+    })?;
+
+    let Some(df) =
+        MoveObject::query(ctx, df_id.into(), Object::latest_at(checkpoint_viewed_at)).await?
+    else {
+        return Ok(None);
+    };
+
+    let app_record = AppRecord::try_from(df.native)?;
+
+    let Some(app_info) = app_record.app_info else {
+        return Ok(None);
+    };
+
+    package_from_app_info(ctx, app_info, versioned.version, checkpoint_viewed_at).await
+}
 
-        Self::package_from_app_info(ctx, app_info, versioned.version, checkpoint_viewed_at).await
+async fn query_external(
+    ctx: &Context<'_>,
+    config: &MoveRegistryConfig,
+    versioned: VersionedName,
+    checkpoint_viewed_at: u64,
+) -> Result<Option<MovePackage>, Error> {
+    if config.external_api_url.is_none() {
+        return Err(MoveRegistryError::ExternalApiUrlUnavailable.into());
     }
 
-    async fn query_external(
-        ctx: &Context<'_>,
-        config: &MoveRegistryConfig,
-        versioned: VersionedName,
-        checkpoint_viewed_at: u64,
-    ) -> Result<Option<MovePackage>, Error> {
-        if config.external_api_url.is_none() {
-            return Err(MoveRegistryError::ExternalApiUrlUnavailable.into());
-        }
-
-        let identifier: ChainIdentifier = *ctx.data_unchecked();
+    let identifier: ChainIdentifier = *ctx.data_unchecked();
 
-        let Some(chain_id) = identifier.0 else {
-            return Err(MoveRegistryError::ChainIdentifierUnavailable.into());
-        };
+    let Some(chain_id) = identifier.0 else {
+        return Err(MoveRegistryError::ChainIdentifierUnavailable.into());
+    };
 
-        let DotMoveDataLoader(loader) = &ctx.data_unchecked();
+    let MoveRegistryDataLoader(loader) = &ctx.data_unchecked();
 
-        let Some(result) = loader.load_one(versioned.name).await? else {
-            return Ok(None);
-        };
+    let Some(result) = loader.load_one(versioned.name).await? else {
+        return Ok(None);
+    };
 
-        let Some(app_info) = result.networks.get(&chain_id.to_string()) else {
-            return Ok(None);
-        };
+    let Some(app_info) = result.networks.get(&chain_id.to_string()) else {
+        return Ok(None);
+    };
 
-        Self::package_from_app_info(
-            ctx,
-            app_info.clone(),
-            versioned.version,
-            checkpoint_viewed_at,
-        )
-        .await
-    }
+    package_from_app_info(
+        ctx,
+        app_info.clone(),
+        versioned.version,
+        checkpoint_viewed_at,
+    )
+    .await
+}
 
-    async fn package_from_app_info(
-        ctx: &Context<'_>,
-        app_info: AppInfo,
-        version: Option<u64>,
-        checkpoint_viewed_at: u64,
-    ) -> Result<Option<MovePackage>, Error> {
-        let Some(package_address) = app_info.package_address else {
-            return Ok(None);
-        };
-
-        // let's now find the package at a specified version (or latest)
-        MovePackage::query(
-            ctx,
-            package_address.into(),
-            if let Some(v) = version {
-                MovePackage::by_version(v, checkpoint_viewed_at)
-            } else {
-                MovePackage::latest_at(checkpoint_viewed_at)
-            },
-        )
-        .await
-    }
+async fn package_from_app_info(
+    ctx: &Context<'_>,
+    app_info: AppInfo,
+    version: Option<u64>,
+    checkpoint_viewed_at: u64,
+) -> Result<Option<MovePackage>, Error> {
+    let Some(package_address) = app_info.package_address else {
+        return Ok(None);
+    };
+
+    // let's now find the package at a specified version (or latest)
+    MovePackage::query(
+        ctx,
+        package_address.into(),
+        if let Some(v) = version {
+            MovePackage::by_version(v, checkpoint_viewed_at)
+        } else {
+            MovePackage::latest_at(checkpoint_viewed_at)
+        },
+    )
+    .await
 }

crates/sui-graphql-rpc/tests/dot_move/dot_move/sources/dotmove.move

@@ -18,8 +18,8 @@ module dotmove::dotmove {
     public struct AppRecord has store {
         // The Capability object used for managing the `AppRecord`.
         app_cap_id: ID,
-        // The mainnet `AppInfo` object.
-        // This is optional until a `mainnet` package is mapped to a record, making 
+        // The core `AppInfo` object.
+        // This is optional until a `mainnet` (base network) package is mapped to a record, making 
         // the record immutable.
         app_info: Option<AppInfo>,
         // This is what being resolved across networks.

crates/sui-graphql-rpc/tests/dot_move_e2e.rs

@@ -82,7 +82,7 @@ mod tests {
 
         // Initialize the internal and external clients of GraphQL.
 
-        // The first cluster uses internal resolution (mimics mainnet, does not rely on external chain).
+        // The first cluster uses internal resolution (mimics our base network, does not rely on external chain).
         let internal_client = init_dot_move_gql(
             8000,
             9184,