Merge pull request #27 from apollographql/jeffrey/support-passing-in-a-custom-scalar-map-file

Jephuff · web-flow · commit 34b3ed40f145 · 2025-04-29T09:52:18.000-07:00
add support for passing in a custom_scalar_map file
diff --git a/crates/mcp-apollo-server/src/custom_scalar_map.rs b/crates/mcp-apollo-server/src/custom_scalar_map.rs
@@ -0,0 +1,226 @@
+use crate::errors::ServerError;
+use rmcp::{
+    schemars::schema::{Schema, SchemaObject, SingleOrVec},
+    serde_json,
+};
+use std::{collections::HashMap, path::PathBuf, str::FromStr};
+
+impl FromStr for CustomScalarMap {
+    type Err = ServerError;
+
+    fn from_str(string_custom_scalar_file: &str) -> Result<Self, Self::Err> {
+        // Parse the string into an initial map of serde_json::Values
+        let parsed_custom_scalar_file: serde_json::Map<String, serde_json::Value> =
+            serde_json::from_str(string_custom_scalar_file)
+                .map_err(ServerError::CustomScalarConfig)?;
+
+        // Validate each of the values in the map and coerce into schemars::schema::SchemaObject
+        let custom_scalar_map = parsed_custom_scalar_file
+            .into_iter()
+            .map(|(key, value)| {
+                let value_string = value.to_string();
+                // The only way I could find to do this was to reparse it.
+                let schema: SchemaObject = serde_json::from_str(value_string.as_str())
+                    .map_err(ServerError::CustomScalarConfig)?;
+
+                if has_invalid_schema(&Schema::Object(schema.clone())) {
+                    Err(ServerError::CustomScalarJsonSchema(value))
+                } else {
+                    Ok((key, schema))
+                }
+            })
+            .collect::<Result<_, _>>()?;
+
+        // panic!("hello2! {:?}", parsed_custom_scalar_file);
+
+        Ok::<_, ServerError>(CustomScalarMap(custom_scalar_map))
+    }
+}
+
+impl TryFrom<&PathBuf> for CustomScalarMap {
+    type Error = ServerError;
+
+    fn try_from(file_path_buf: &PathBuf) -> Result<Self, Self::Error> {
+        let custom_scalars_config_path = file_path_buf.as_path();
+        tracing::info!(custom_scalars_config=?custom_scalars_config_path, "Loading custom_scalars_config");
+        let string_custom_scalar_file = std::fs::read_to_string(custom_scalars_config_path)?;
+        CustomScalarMap::from_str(string_custom_scalar_file.as_str())
+    }
+}
+
+#[derive(Debug)]
+pub struct CustomScalarMap(HashMap<String, SchemaObject>);
+
+impl CustomScalarMap {
+    pub fn get(&self, key: &str) -> Option<&SchemaObject> {
+        self.0.get(key)
+    }
+}
+
+// Unknown keys will be put into "extensions" in the schema object, check for those and consider those invalid
+fn has_invalid_schema(schema: &Schema) -> bool {
+    match schema {
+        Schema::Object(schema_object) => {
+            !schema_object.extensions.is_empty()
+                || schema_object
+                    .object
+                    .as_ref()
+                    .is_some_and(|object| object.properties.values().any(has_invalid_schema))
+                || schema_object.array.as_ref().is_some_and(|object| {
+                    object.items.as_ref().is_some_and(|items| match items {
+                        SingleOrVec::Single(item) => has_invalid_schema(item),
+                        SingleOrVec::Vec(items) => items.iter().any(has_invalid_schema),
+                    })
+                })
+        }
+        Schema::Bool(_) => false,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::{
+        collections::{BTreeMap, HashMap},
+        str::FromStr,
+    };
+
+    use rmcp::schemars::schema::{
+        InstanceType, ObjectValidation, Schema, SchemaObject, SingleOrVec,
+    };
+
+    use crate::custom_scalar_map::CustomScalarMap;
+
+    #[test]
+    fn empty_file() {
+        let result = CustomScalarMap::from_str("").err().unwrap();
+
+        insta::assert_debug_snapshot!(result, @r###"
+            CustomScalarConfig(
+                Error("EOF while parsing a value", line: 1, column: 0),
+            )
+        "###)
+    }
+
+    #[test]
+    fn only_spaces() {
+        let result = CustomScalarMap::from_str("    ").err().unwrap();
+
+        insta::assert_debug_snapshot!(result, @r###"
+            CustomScalarConfig(
+                Error("EOF while parsing a value", line: 1, column: 4),
+            )
+        "###)
+    }
+
+    #[test]
+    fn invalid_json() {
+        let result = CustomScalarMap::from_str("Hello: }").err().unwrap();
+
+        insta::assert_debug_snapshot!(result, @r###"
+            CustomScalarConfig(
+                Error("expected value", line: 1, column: 1),
+            )
+        "###)
+    }
+
+    #[test]
+    fn invalid_simple_schema() {
+        let result = CustomScalarMap::from_str(
+            r###"{
+                "custom": {
+                    "test": true
+                }
+            }"###,
+        )
+        .err()
+        .unwrap();
+
+        insta::assert_debug_snapshot!(result, @r###"
+            CustomScalarJsonSchema(
+                Object {
+                    "test": Bool(true),
+                },
+            )
+        "###)
+    }
+
+    #[test]
+    fn invalid_complex_schema() {
+        let result = CustomScalarMap::from_str(
+            r###"{
+                "custom": {
+                    "type": "object",
+                    "properties": {
+                        "test": {
+                            "test": true
+                        }
+                    }
+                }
+            }"###,
+        )
+        .err()
+        .unwrap();
+
+        insta::assert_debug_snapshot!(result, @r###"
+            CustomScalarJsonSchema(
+                Object {
+                    "properties": Object {
+                        "test": Object {
+                            "test": Bool(true),
+                        },
+                    },
+                    "type": String("object"),
+                },
+            )
+        "###)
+    }
+
+    #[test]
+    fn valid_schema() {
+        let result = CustomScalarMap::from_str(
+            r###"
+        {
+            "simple": {
+                "type": "string"
+            },
+            "complex": {
+                "type": "object",
+                "properties": { "name": { "type": "string" } }
+            }
+        }
+        "###,
+        )
+        .unwrap()
+        .0;
+
+        let expected_data = HashMap::from_iter([
+            (
+                "simple".to_string(),
+                SchemaObject {
+                    instance_type: Some(SingleOrVec::Single(Box::new(InstanceType::String))),
+                    ..Default::default()
+                },
+            ),
+            (
+                "complex".to_string(),
+                SchemaObject {
+                    instance_type: Some(SingleOrVec::Single(Box::new(InstanceType::Object))),
+                    object: Some(Box::new(ObjectValidation {
+                        properties: BTreeMap::from_iter([(
+                            "name".to_string(),
+                            Schema::Object(SchemaObject {
+                                instance_type: Some(SingleOrVec::Single(Box::new(
+                                    InstanceType::String,
+                                ))),
+                                ..Default::default()
+                            }),
+                        )]),
+                        ..Default::default()
+                    })),
+                    ..Default::default()
+                },
+            ),
+        ]);
+        assert_eq!(result, expected_data);
+    }
+}
diff --git a/crates/mcp-apollo-server/src/errors.rs b/crates/mcp-apollo-server/src/errors.rs
@@ -53,6 +53,12 @@ pub enum ServerError {
 
     #[error("invalid header: {0}")]
     Header(String),
+
+    #[error("invalid custom_scalar_config: {0}")]
+    CustomScalarConfig(serde_json::Error),
+
+    #[error("invalid json schema: {0}")]
+    CustomScalarJsonSchema(serde_json::Value),
 }
 
 /// An MCP tool error
diff --git a/crates/mcp-apollo-server/src/lib.rs b/crates/mcp-apollo-server/src/lib.rs
@@ -1,3 +1,4 @@
+pub mod custom_scalar_map;
 pub mod errors;
 mod graphql;
 mod introspection;
diff --git a/crates/mcp-apollo-server/src/main.rs b/crates/mcp-apollo-server/src/main.rs
@@ -3,6 +3,7 @@ use apollo_compiler::Schema;
 use clap::builder::Styles;
 use clap::builder::styling::{AnsiColor, Effects};
 use clap::{Parser, ValueEnum};
+use mcp_apollo_server::custom_scalar_map::CustomScalarMap;
 use mcp_apollo_server::errors::ServerError;
 use mcp_apollo_server::server::Server;
 use rmcp::ServiceExt;
@@ -35,6 +36,10 @@ struct Args {
     #[clap(long, short = 's')]
     schema: PathBuf,
 
+    /// The path to the GraphQL custom_scalars_config file
+    #[clap(long, short = 'c', required = false)]
+    custom_scalars_config: Option<PathBuf>,
+
     /// The GraphQL endpoint the server will invoke
     #[clap(long, short = 'e', default_value = "http://127.0.0.1:4000")]
     endpoint: String,
@@ -102,6 +107,11 @@ async fn main() -> anyhow::Result<()> {
         .operations(args.operations)
         .headers(args.headers)
         .introspection(args.introspection)
+        .and_custom_scalar_map(
+            args.custom_scalars_config
+                .map(|custom_scalars_config| CustomScalarMap::try_from(&custom_scalars_config))
+                .transpose()?,
+        )
         .and_persisted_query_manifest(
             args.pq_manifest
                 .map(
diff --git a/crates/mcp-apollo-server/src/operations.rs b/crates/mcp-apollo-server/src/operations.rs
@@ -1,5 +1,3 @@
-use std::collections::HashMap;
-
 use apollo_compiler::ast::{FragmentDefinition, Selection};
 use apollo_compiler::{
     Name, Node, Schema as GraphqlSchema,
@@ -17,6 +15,7 @@ use rmcp::{
 use rover_copy::pq_manifest::ApolloPersistedQueryManifest;
 use serde::Serialize;
 
+use crate::custom_scalar_map::CustomScalarMap;
 use crate::errors::{McpError, OperationError};
 use crate::graphql;
 use crate::tree_shake::TreeShaker;
@@ -42,7 +41,7 @@ impl Operation {
     pub fn from_document(
         source_text: &str,
         graphql_schema: &GraphqlSchema,
-        custom_scalar_map: Option<&HashMap<String, SchemaObject>>,
+        custom_scalar_map: Option<&CustomScalarMap>,
     ) -> Result<Self, OperationError> {
         let document = Parser::new()
             .parse_ast(source_text, "operation.graphql")
@@ -110,14 +109,15 @@ impl Operation {
     pub fn from_manifest(
         schema: &GraphqlSchema,
         manifest: ApolloPersistedQueryManifest,
+        custom_scalar_map: Option<&CustomScalarMap>,
     ) -> Result<Vec<Self>, OperationError> {
         manifest
             .operations
             .into_iter()
             .map(|pq| {
                 tracing::info!(pesisted_query = pq.name, "Loading persisted query");
 
-                Self::from_document(&pq.body, schema, None)
+                Self::from_document(&pq.body, schema, custom_scalar_map)
             })
             .collect::<Result<Vec<_>, _>>()
     }
@@ -230,7 +230,7 @@ impl Operation {
 fn get_json_schema(
     operation: &Node<OperationDefinition>,
     graphql_schema: &GraphqlSchema,
-    custom_scalar_map: Option<&HashMap<String, SchemaObject>>,
+    custom_scalar_map: Option<&CustomScalarMap>,
 ) -> RootSchema {
     let mut obj = ObjectValidation::default();
 
@@ -321,7 +321,7 @@ fn type_to_schema(
     description: Option<String>,
     variable_type: &Type,
     graphql_schema: &GraphqlSchema,
-    custom_scalar_map: Option<&HashMap<String, SchemaObject>>,
+    custom_scalar_map: Option<&CustomScalarMap>,
 ) -> Schema {
     match variable_type {
         Type::NonNullNamed(named) | Type::Named(named) => match named.as_str() {
@@ -459,20 +459,13 @@ impl graphql::Executable for Operation {
 
 #[cfg(test)]
 mod tests {
-    use std::{
-        collections::{HashMap, HashSet},
-        sync::LazyLock,
-    };
+    use std::{collections::HashSet, str::FromStr, sync::LazyLock};
 
     use apollo_compiler::{Schema, parser::Parser, validation::Valid};
-    use rmcp::{
-        model::Tool,
-        schemars::schema::{InstanceType, SchemaObject, SingleOrVec},
-        serde_json,
-    };
+    use rmcp::{model::Tool, serde_json};
     use rover_copy::pq_manifest::ApolloPersistedQueryManifest;
 
-    use crate::operations::Operation;
+    use crate::{custom_scalar_map::CustomScalarMap, operations::Operation};
 
     // Example schema for tests
     static SCHEMA: LazyLock<Valid<Schema>> = LazyLock::new(|| {
@@ -1069,7 +1062,7 @@ mod tests {
         let operation = Operation::from_document(
             "query QueryName($id: RealCustomScalar) { id }",
             &SCHEMA,
-            Some(&HashMap::new()),
+            Some(&CustomScalarMap::from_str("{}").unwrap()),
         )
         .unwrap();
         let tool = Tool::from(operation);
@@ -1092,18 +1085,13 @@ mod tests {
 
     #[test]
     fn custom_scalar_with_map() {
-        let custom_scalar_map = HashMap::from([(
-            "RealCustomScalar".to_string(),
-            SchemaObject {
-                instance_type: Some(SingleOrVec::Single(Box::new(InstanceType::String))),
-                ..Default::default()
-            },
-        )]);
+        let custom_scalar_map =
+            CustomScalarMap::from_str("{ \"RealCustomScalar\": { \"type\": \"string\" }}");
 
         let operation = Operation::from_document(
             "query QueryName($id: RealCustomScalar) { id }",
             &SCHEMA,
-            Some(&custom_scalar_map),
+            custom_scalar_map.ok().as_ref(),
         )
         .unwrap();
         let tool = Tool::from(operation);
@@ -1355,7 +1343,7 @@ mod tests {
         }))
         .expect("apollo pq should be valid");
 
-        let operations = Operation::from_manifest(&SCHEMA, apollo_pq.clone())
+        let operations = Operation::from_manifest(&SCHEMA, apollo_pq.clone(), None)
             .expect("operations from manifest should parse");
         assert_eq!(
             operations
diff --git a/crates/mcp-apollo-server/src/server.rs b/crates/mcp-apollo-server/src/server.rs

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+pub mod custom_scalar_map;`
`1`	`2`	`pub mod errors;`
`2`	`3`	`mod graphql;`
`3`	`4`	`mod introspection;`