apollo-language-server 0.7.0

pub mod link;

use apollo_compiler::ast::{Definition, Document, OperationType};
use std::collections::HashMap;

use itertools::Itertools;
use link::{find_link_for_spec, ParsedLink};

use crate::specs::{
    connect::CONNECT_SPEC_NAME, federation::FEDERATION_SPEC_NAME, link::LINK_DIRECTIVE,
    tag::TAG_SPEC_NAME, Spec, KNOWN_SPECS,
};

pub type SpecBuiltins = Option<(String, Option<HashMap<String, String>>, Option<ParsedLink>)>;

// For a given v2 subgraph document, add the @link directive and then parse out the federation `@link` directive if
// present and generate an appropriate document of "federation builtins", to be
// used in conjunction with the original document for schema validation
// purposes.
pub fn get_federation_builtins_for_v2_document(document: &Document) -> SpecBuiltins {
    let implicit_link_builtins = collect_directives_and_dependencies_for_implicit_link(document);

    let (federation_builtins, federation_link) =
        collect_directives_and_dependencies_for_spec(document, FEDERATION_SPEC_NAME);

    let merged_spec_to_aliases_map = implicit_link_builtins
        .spec_to_aliases_map
        .into_iter()
        .chain(federation_builtins.spec_to_aliases_map)
        .collect::<HashMap<_, _>>();

    let builtin_definitions = implicit_link_builtins
        .definitions
        .into_iter()
        .chain(federation_builtins.definitions)
        .sorted()
        .join("\n\n");

    (!builtin_definitions.is_empty()).then_some((
        builtin_definitions,
        Some(merged_spec_to_aliases_map),
        federation_link,
    ))
}

pub fn get_connect_builtins(document: &Document) -> SpecBuiltins {
    let (builtins_metadata, parsed_link) =
        collect_directives_and_dependencies_for_spec(document, CONNECT_SPEC_NAME);

    (!builtins_metadata.definitions.is_empty()).then_some((
        builtins_metadata
            .definitions
            .into_iter()
            .sorted()
            .join("\n\n"),
        Some(builtins_metadata.spec_to_aliases_map),
        parsed_link,
    ))
}

pub fn get_tag_builtins(document: &Document) -> SpecBuiltins {
    let (builtins_metadata, parsed_link) =
        collect_directives_and_dependencies_for_spec(document, TAG_SPEC_NAME);

    (!builtins_metadata.definitions.is_empty()).then_some((
        builtins_metadata
            .definitions
            .into_iter()
            .sorted()
            .join("\n\n"),
        Some(builtins_metadata.spec_to_aliases_map),
        parsed_link,
    ))
}

pub fn get_missing_query_type_for_document(document: &Document) -> Option<String> {
    let maybe_schema_def = document
        .definitions
        .iter()
        .find_map(|def| def.as_schema_definition());
    let schema_extensions = document
        .definitions
        .iter()
        .filter_map(|def| def.as_schema_extension())
        .collect::<Vec<_>>();

    // search the schema definition and extensions for a query root operation
    let has_query_root_operation = maybe_schema_def
        .map(|schema_def| {
            schema_def
                .root_operations
                .iter()
                .any(|root_op| root_op.0 == OperationType::Query)
        })
        .is_some()
        || schema_extensions.iter().any(|schema_ext| {
            schema_ext
                .root_operations
                .iter()
                .any(|root_op| root_op.0 == OperationType::Query)
        });

    if has_query_root_operation {
        return None;
    }

    let query_type = document.definitions.iter().find(|def| match def {
        Definition::ObjectTypeDefinition(obj) => obj.name == "Query",
        _ => false,
    });

    if query_type.is_some() {
        return None;
    }

    let mut to_inject = vec![
        "extend schema { query: _ApolloInjectedQuery }",
        // Ensure we don't cause a naming collision with the name `Query`
        "type _ApolloInjectedQuery { _entities(representations: [_Any!]!): [_Entity]! }",
    ];

    if !document.definitions.iter().any(|def| {
        def.as_union_type_definition()
            .is_some_and(|union_def| union_def.name == "_Entity")
    }) {
        // Using a scalar here is a hack, but it doesn't introduce validation
        // issues due to it being an enum with no members.
        to_inject.push("scalar _Entity");
    }

    if !document.definitions.iter().any(|def| {
        def.as_scalar_type_definition()
            .is_some_and(|scalar_def| scalar_def.name == "_Any")
    }) {
        to_inject.push("scalar _Any");
    }

    Some(to_inject.join("\n\n"))
}

#[derive(Debug, Default)]
pub struct BuiltinsMetadata {
    definitions: Vec<String>,
    spec_to_aliases_map: HashMap<String, String>,
}

fn collect_directives_and_dependencies_for_implicit_link(document: &Document) -> BuiltinsMetadata {
    let implicit_link_directive = KNOWN_SPECS
        .get(LINK_DIRECTIVE)
        .expect("Programming error: `link` spec definition missing")
        .get("1.0")
        .expect("Programming error: `link` spec version 1.0 missing");

    process_document_for_spec(document, implicit_link_directive.clone(), None)
}

fn process_document_for_spec(
    document: &Document,
    spec: Spec,
    parsed_link: Option<&ParsedLink>,
) -> BuiltinsMetadata {
    let mut spec_to_aliases_map = HashMap::new();
    let (mut scalars, mut input_types, mut directives, mut enums) = match parsed_link {
        None => {
            spec.scalars.values().for_each(|scalar| {
                spec_to_aliases_map.insert(
                    scalar.node.name.as_str().to_string(),
                    scalar.node.name.as_str().to_string(),
                );
            });
            spec.input_types.values().for_each(|ty| {
                spec_to_aliases_map.insert(
                    ty.node.name.as_str().to_string(),
                    ty.node.name.as_str().to_string(),
                );
            });
            spec.directives.values().for_each(|directive| {
                spec_to_aliases_map.insert(
                    directive.node.name.as_str().to_string(),
                    directive.node.name.as_str().to_string(),
                );
            });
            spec.enums.values().for_each(|en| {
                spec_to_aliases_map.insert(
                    en.node.name.as_str().to_string(),
                    en.node.name.as_str().to_string(),
                );
            });

            (spec.scalars, spec.input_types, spec.directives, spec.enums)
        }
        Some(parsed_link) => (
            spec.scalars
                .into_values()
                .filter_map(|scalar| {
                    let updated_scalar = scalar.update_with_link(parsed_link).ok()?;
                    spec_to_aliases_map.insert(
                        scalar.node.name.as_str().to_string(),
                        updated_scalar.node.name.as_str().to_string(),
                    );
                    Some((
                        updated_scalar.node.name.as_str().to_string(),
                        updated_scalar,
                    ))
                })
                .collect::<HashMap<_, _>>(),
            spec.input_types
                .into_values()
                .filter_map(|ty| {
                    let updated_type = ty.update_with_link(parsed_link).ok()?;
                    spec_to_aliases_map.insert(
                        ty.node.name.as_str().to_string(),
                        updated_type.node.name.as_str().to_string(),
                    );
                    Some((updated_type.node.name.as_str().to_string(), updated_type))
                })
                .collect::<HashMap<_, _>>(),
            spec.directives
                .into_values()
                .filter_map(|directive| {
                    let updated_directive = directive.update_with_link(parsed_link).ok()?;
                    spec_to_aliases_map.insert(
                        directive.node.name.as_str().to_string(),
                        updated_directive.node.name.as_str().to_string(),
                    );
                    Some((
                        updated_directive.node.name.as_str().to_string(),
                        updated_directive,
                    ))
                })
                .collect::<HashMap<_, _>>(),
            spec.enums
                .into_values()
                .filter_map(|en| {
                    let updated_enum = en.update_with_link(parsed_link).ok()?;
                    spec_to_aliases_map.insert(
                        en.node.name.as_str().to_string(),
                        updated_enum.node.name.as_str().to_string(),
                    );
                    Some((updated_enum.node.name.as_str().to_string(), updated_enum))
                })
                .collect::<HashMap<_, _>>(),
        ),
    };
    // Remove definitions that collide with existing definitions in the document
    for def in &document.definitions {
        match def {
            Definition::DirectiveDefinition(directive_def) => {
                directives.remove(&directive_def.name.to_string());
            }
            Definition::ScalarTypeDefinition(scalar_def) => {
                scalars.remove(&scalar_def.name.to_string());
            }
            Definition::InputObjectTypeDefinition(scalar_def) => {
                input_types.remove(&scalar_def.name.to_string());
            }
            Definition::EnumTypeDefinition(enum_def) => {
                enums.remove(&enum_def.name.to_string());
            }
            _ => {}
        }
    }

    let scalars = scalars.values().map(ToString::to_string).sorted();
    let input_types = input_types.values().map(ToString::to_string).sorted();
    let directives = directives.values().map(ToString::to_string).sorted();
    let enums = enums.values().map(ToString::to_string).sorted();

    BuiltinsMetadata {
        definitions: scalars
            .chain(input_types)
            .chain(directives)
            .chain(enums)
            .collect(),
        spec_to_aliases_map,
    }
}

fn collect_directives_and_dependencies_for_spec(
    document: &Document,
    spec_name: &str,
) -> (BuiltinsMetadata, Option<ParsedLink>) {
    let Some(parsed_link) = find_link_for_spec(spec_name, document) else {
        return (BuiltinsMetadata::default(), None);
    };
    let version_string = format!(
        "{}.{}",
        parsed_link.version.major, parsed_link.version.minor
    );
    let version_string = version_string.as_str();

    let Some(directives_for_spec_version) = KNOWN_SPECS
        .get(spec_name)
        .and_then(|specs_by_version| specs_by_version.get(version_string))
    else {
        return (BuiltinsMetadata::default(), None);
    };

    (
        process_document_for_spec(
            document,
            directives_for_spec_version.clone(),
            Some(&parsed_link),
        ),
        Some(parsed_link),
    )
}

// For a given v1 subgraph document, add the fed 1 directives that aren't
// already defined in the document, to be used in conjunction with the original
// document for schema validation purposes.
pub fn get_federation_builtins_for_v1_document(document: &Document) -> SpecBuiltins {
    let implicit_link_builtins = collect_directives_and_dependencies_for_implicit_link(document);

    let fed1_spec = KNOWN_SPECS
        .get(FEDERATION_SPEC_NAME)
        .expect("Federation spec not found")
        .get("1.0")
        .expect("Federation spec 1.0 not found");

    let federation_builtins = process_document_for_spec(document, fed1_spec.clone(), None);

    let merged_spec_to_aliases_map = implicit_link_builtins
        .spec_to_aliases_map
        .into_iter()
        .chain(federation_builtins.spec_to_aliases_map)
        .collect::<HashMap<_, _>>();

    let builtin_definitions = implicit_link_builtins
        .definitions
        .into_iter()
        .chain(federation_builtins.definitions)
        .sorted()
        .join("\n\n");

    (!builtin_definitions.is_empty()).then_some((
        builtin_definitions,
        Some(merged_spec_to_aliases_map),
        None,
    ))
}

#[cfg(test)]
mod tests {
    use apollo_compiler::parser::Parser;
    use insta::assert_snapshot;

    use crate::testing::{collect_diagnostic_comparisons, pretty_print_spec_builtins};

    use super::*;

    fn parse(source_text: &str) -> Document {
        let mut parser = Parser::new();
        parser.parse_ast(source_text, "test.graphql").unwrap()
    }

    fn get_federation_builtins_for_version(version: &str, imports: &[&str]) -> String {
        let source_text = format!(
            r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v{}", import: [{}])

            type Query {{
                hello: String
            }}
        "#,
            version,
            imports
                .iter()
                .map(|import| format!(r#""{}""#, import))
                .join(", ")
        );
        let document = &parse(&source_text);
        [
            get_federation_builtins_for_v2_document(document).unwrap().0,
            get_connect_builtins(document).unwrap_or_default().0,
        ]
        .iter()
        .join("\n\n")
    }

    fn get_test_connect_builtins() -> String {
        let source_text = r#"
            extend schema
              @link(url: "https://specs.apollo.dev/connect/v0.3", import: ["@connect", "@source", "HTTPHeaderMapping", "SourceHTTP", "JSONSelection", "URLPathTemplate", "ConnectHTTP"])

            type Query {
                hello: String
            }
        "#.to_string();
        let document = parse(&source_text);
        [
            get_federation_builtins_for_v2_document(&document)
                .unwrap()
                .0,
            get_connect_builtins(&document).unwrap().0,
        ]
        .join("\n\n")
    }

    fn get_imports(additional_imports: Vec<&str>) -> Vec<&str> {
        let mut v2_0_imports = vec![
            "@key",
            "@requires",
            "@provides",
            "@extends",
            "@external",
            "@tag",
            "@override",
            "@shareable",
            "@inaccessible",
            "FieldSet",
        ];
        v2_0_imports.extend(additional_imports);
        v2_0_imports
    }

    #[test]
    fn test_builtins_for_1_0() {
        let source_text = r#"
            type Query {
                hello: Hello
            }

            type Hello @key(fields: "id") {
                id: String
            }
        "#;
        let document = &parse(source_text);

        assert_snapshot!(get_federation_builtins_for_v1_document(document).unwrap().0);
    }

    #[test]
    fn test_builtins_for_2_0() {
        insta::assert_snapshot!(&get_federation_builtins_for_version(
            "2.0",
            &get_imports(vec![])
        ));
    }

    #[test]
    fn test_builtins_for_2_1_and_2_2() {
        let v2_2_imports = get_imports(vec!["@composeDirective"]);
        let result = get_federation_builtins_for_version("2.1", &v2_2_imports);
        insta::assert_snapshot!(result);
        assert!(result.contains("@composeDirective"));
        assert_eq!(
            get_federation_builtins_for_version("2.1", &v2_2_imports),
            get_federation_builtins_for_version("2.2", &v2_2_imports)
        );
    }

    #[test]
    fn test_builtins_for_2_3_and_2_4() {
        let v2_4_imports = get_imports(vec!["@composeDirective", "@interfaceObject"]);
        let result = get_federation_builtins_for_version("2.3", &v2_4_imports);
        insta::assert_snapshot!(result);
        assert!(result.contains("@interfaceObject"));
        assert_eq!(
            get_federation_builtins_for_version("2.3", &v2_4_imports),
            get_federation_builtins_for_version("2.4", &v2_4_imports)
        );
    }

    #[test]
    fn test_builtins_for_2_5() {
        let v2_5_imports = get_imports(vec![
            "@composeDirective",
            "@interfaceObject",
            "@authenticated",
            "@requiresScopes",
            "Scope",
        ]);
        let result = get_federation_builtins_for_version("2.5", &v2_5_imports);
        insta::assert_snapshot!(result);
        assert!(result.contains("@authenticated"));
        assert!(result.contains("@requiresScopes"));
        assert!(result.contains("scalar Scope"));
    }

    #[test]
    fn test_builtins_for_2_6() {
        let v2_6_imports = get_imports(vec![
            "@composeDirective",
            "@interfaceObject",
            "@authenticated",
            "@requiresScopes",
            "@policy",
            "Scope",
            "Policy",
        ]);
        let result = get_federation_builtins_for_version("2.6", &v2_6_imports);
        insta::assert_snapshot!(result);
        assert!(result.contains("@policy"));
        assert!(result.contains("scalar Policy"));
    }

    #[test]
    fn test_builtins_for_2_7() {
        let v2_7_imports = get_imports(vec![
            "@composeDirective",
            "@interfaceObject",
            "@authenticated",
            "@requiresScopes",
            "@policy",
            "Scope",
            "Policy",
        ]);
        let result = get_federation_builtins_for_version("2.7", &v2_7_imports);
        insta::assert_snapshot!(result);
        assert!(result.contains("@override(from: String!, label: String)"));
    }

    #[test]
    fn test_builtins_for_2_8() {
        let v2_8_imports = get_imports(vec![
            "@composeDirective",
            "@interfaceObject",
            "@authenticated",
            "@requiresScopes",
            "@policy",
            "Scope",
            "Policy",
            "@context",
            "@fromContext",
            "ContextFieldValue",
        ]);
        let result = get_federation_builtins_for_version("2.8", &v2_8_imports);
        insta::assert_snapshot!(result);
    }

    #[test]
    fn test_builtins_for_2_8_with_connect_link() {
        let result = get_test_connect_builtins();
        insta::assert_snapshot!(result);
    }

    #[test]
    // using `as` argument to alias all federation imports
    fn test_builtins_for_import_spec_as_fed() {
        let source_text = r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v2.7", as: "fed", import: ["@key"])

            type Query {
                hello: String
            }
        "#;

        assert_snapshot!(pretty_print_spec_builtins(
            get_federation_builtins_for_v2_document(&parse(source_text))
        ));
    }

    #[test]
    // using a specific import's `as` argument to alias that import
    fn test_builtins_for_import_key_as_apollokey() {
        let source_text = r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v2.7", import: [{ name: "@key", as: "@apolloKey" }])

            type Query {
                hello: String
            }
        "#;

        let builtins = get_federation_builtins_for_v2_document(&parse(source_text))
            .unwrap()
            .0;
        assert!(builtins.contains("directive @apolloKey"));
        assert!(!builtins.contains("directive @key"));
    }

    #[test]
    // `interfaceObject` shouldn't be available to a 2.0 schema
    fn test_fed_parse_validate_err_directive_not_imported_by_fed_version() {
        let expected_errors = ["cannot find directive `@interfaceObject` in this document"];
        let source_texts = vec![
            r#"
extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key", "@interfaceObject"])

type Query {
    hello: User
}

type User @key(fields: "id") @interfaceObject {
    id: ID!
}"#,
        ];

        assert_snapshot!(collect_diagnostic_comparisons(
            &expected_errors,
            &source_texts,
            None,
        ));
    }

    #[test]
    fn test_fed_parse_validate_err_improper_as_usage() {
        let expected_errors = ["cannot find directive `@fed__interfaceObject` in this document"];
        let source_texts = vec![
            r#"
extend schema @link(url: "https://specs.apollo.dev/federation/v2.3", as: "fed", import: ["@key", "@interfaceObject"])

type Query {
    hello: User
}

type User @key(fields: "id") @fed__interfaceObject {
    id: ID!
}"#,
        ];

        assert_snapshot!(collect_diagnostic_comparisons(
            &expected_errors,
            &source_texts,
            None,
        ));
    }

    #[test]
    fn test_simple_non_federation_usage() {
        let expected_errors = ["`TypeName` has no fields"];
        let source_texts = vec!["type Query { hello: String }\ntype TypeName"];

        assert_snapshot!(collect_diagnostic_comparisons(
            &expected_errors,
            &source_texts,
            None,
        ));
    }

    #[test]
    fn test_handles_invalid_link_usage() {
        let source_text = r#"
            extend schema @link
        "#;
        let expected_errors = ["the required argument `@link(url:)` is not provided"];
        let source_texts = vec![source_text];

        assert_snapshot!(collect_diagnostic_comparisons(
            &expected_errors,
            &source_texts,
            None
        ));
    }

    #[test]
    fn injects_correct_names_when_no_import_as() {
        let source_text = r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v2.6", import: [{ name: "@key", as: "@apolloKey" }, "@interfaceObject", { name: "Policy", as: "apolloPolicy" }, "Scope"])
        "#;

        let result = get_federation_builtins_for_v2_document(&parse(source_text))
            .unwrap()
            .0;

        assert!(result.contains("directive @interfaceObject"));
        assert!(result.contains("directive @apolloKey"));
        assert!(result.contains("directive @federation__requires"));
        assert!(result.contains("scalar federation__FieldSet"));
        assert!(result.contains("scalar Scope"));
        assert!(result.contains("scalar apolloPolicy"));
    }

    #[test]
    fn injects_correct_names_when_using_import_as() {
        let source_text = r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v2.6", as: "fed", import: [{ name: "@key", as: "@apolloKey" }, "@interfaceObject", { name: "Policy", as: "apolloPolicy" }, "Scope"])
        "#;

        let result = get_federation_builtins_for_v2_document(&parse(source_text))
            .unwrap()
            .0;

        assert!(result.contains("directive @interfaceObject"));
        assert!(result.contains("directive @apolloKey"));
        assert!(result.contains("directive @fed__requires"));

        assert!(result.contains("scalar Scope"));
        assert!(result.contains("scalar apolloPolicy"));
        assert!(result.contains("scalar fed__FieldSet"));
    }

    #[test]
    fn allows_providing_link_definition() {
        let source_text = r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: [{ name: "@key", as: "@apolloKey" }])

            directive @link(url: String!, as: String, import: [Import!]) repeatable on SCHEMA
        "#;

        assert_snapshot!(pretty_print_spec_builtins(
            get_federation_builtins_for_v2_document(&parse(source_text))
        ));
    }

    #[test]
    fn supports_no_import_argument() {
        let source_text = r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v2.0")
        "#;

        assert_snapshot!(pretty_print_spec_builtins(
            get_federation_builtins_for_v2_document(&parse(source_text))
        ));
    }

    #[test]
    fn supports_empty_import_argument() {
        let source_text = r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: [])
        "#;

        assert_snapshot!(pretty_print_spec_builtins(
            get_federation_builtins_for_v2_document(&parse(source_text))
        ));
    }

    #[test]
    fn connectors_imports() {
        let source_text = r#"
            extend schema
                @link(url: "https://specs.apollo.dev/connect/v0.1", import: ["@source"])
                @source(name: "abc", http: { baseURL: "abc" })

                type Query {
                    b: String!
                }
            "#;

        assert_snapshot!(pretty_print_spec_builtins(get_connect_builtins(&parse(
            source_text
        ))));
    }

    #[test]
    fn connectors_v0_2_imports() {
        let source_text = r#"
            extend schema
                @link(url: "https://specs.apollo.dev/connect/v0.2", import: ["@source"])
                @source(name: "abc", http: { baseURL: "abc" })

                type Query {
                    b: String!
                }
            "#;

        assert_snapshot!(pretty_print_spec_builtins(get_connect_builtins(&parse(
            source_text
        ))));
    }

    #[test]
    fn connectors_v0_3_imports() {
        let source_text = r#"
            extend schema
                @link(url: "https://specs.apollo.dev/connect/v0.3", import: ["@source"])
                @source(name: "abc", http: { baseURL: "abc" })

                type Query {
                    b: String!
                }
            "#;

        assert_snapshot!(pretty_print_spec_builtins(get_connect_builtins(&parse(
            source_text
        ))));
    }

    #[test]
    fn connectors_v0_4_imports() {
        let source_text = r#"
            extend schema
                @link(url: "https://specs.apollo.dev/connect/v0.4", import: ["@source"])
                @source(name: "abc", http: { baseURL: "abc" })

                type Query {
                    b: String!
                }
            "#;

        assert_snapshot!(pretty_print_spec_builtins(get_connect_builtins(&parse(
            source_text
        ))));
    }

    #[test]
    fn does_not_duplicate_namespaced_imports() {
        let source_text = r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v2.0", import: ["@key"])

            directive @key(fields: federation__FieldSet!, resolvable: Boolean = true) repeatable on OBJECT | INTERFACE

            directive @federation__provides(fields: federation__FieldSet!) on FIELD_DEFINITION

            type Query {
                a: String!
            }
        "#;

        let result = get_federation_builtins_for_v2_document(&parse(source_text))
            .unwrap()
            .0;

        assert!(!result.contains("@federation__provides"));
    }
}