[IJ Plugin] Inspection: missing directive import

intellij-plugin/build.gradle.kts

@@ -217,6 +217,7 @@ publishing {
 
 dependencies {
   implementation(project(":apollo-gradle-plugin-external"))
+  implementation(project(":apollo-ast"))
   implementation(project(":apollo-tooling"))
   implementation(project(":apollo-normalized-cache-sqlite"))
   implementation(libs.sqlite.jdbc)

intellij-plugin/src/main/kotlin/com/apollographql/ijplugin/gradle/GradleToolingModelService.kt

@@ -117,6 +117,11 @@ class GradleToolingModelService(
     })
   }
 
+  fun triggerFetchToolingModels() {
+    logd()
+    startOrAbortFetchToolingModels()
+  }
+
   private fun startOrAbortFetchToolingModels() {
     logd()
     abortFetchToolingModels()
@@ -181,7 +186,7 @@ class GradleToolingModelService(
           logd("Fetch tooling model for :${gradleProject.name}")
           return@execute try {
             val id = ExternalSystemTaskId.create(GRADLE_SYSTEM_ID, ExternalSystemTaskType.RESOLVE_PROJECT, project)
-            gradleExecutionHelper.getModelBuilder(ApolloGradleToolingModel::class.java, connection,id, executionSettings,  ExternalSystemTaskNotificationListenerAdapter.NULL_OBJECT)
+            gradleExecutionHelper.getModelBuilder(ApolloGradleToolingModel::class.java, connection, id, executionSettings, ExternalSystemTaskNotificationListenerAdapter.NULL_OBJECT)
                 .withCancellationToken(gradleCancellation!!.token())
                 .get()
                 .takeIf {
@@ -296,3 +301,5 @@ class GradleToolingModelService(
     }
   }
 }
+
+val Project.gradleToolingModelService get() = service<GradleToolingModelService>()

intellij-plugin/src/main/kotlin/com/apollographql/ijplugin/inspection/ApolloMissingGraphQLDefinitionImportInspection.kt

@@ -0,0 +1,178 @@
+package com.apollographql.ijplugin.inspection
+
+import com.apollographql.apollo3.ast.GQLDefinition
+import com.apollographql.apollo3.ast.GQLDirectiveDefinition
+import com.apollographql.apollo3.ast.GQLEnumTypeDefinition
+import com.apollographql.apollo3.ast.GQLInputObjectTypeDefinition
+import com.apollographql.apollo3.ast.GQLNamed
+import com.apollographql.apollo3.ast.GQLScalarTypeDefinition
+import com.apollographql.apollo3.ast.rawType
+import com.apollographql.ijplugin.ApolloBundle
+import com.apollographql.ijplugin.gradle.gradleToolingModelService
+import com.apollographql.ijplugin.project.apolloProjectService
+import com.apollographql.ijplugin.telemetry.TelemetryEvent
+import com.apollographql.ijplugin.telemetry.telemetryService
+import com.apollographql.ijplugin.util.cast
+import com.apollographql.ijplugin.util.findChildrenOfType
+import com.apollographql.ijplugin.util.quoted
+import com.apollographql.ijplugin.util.unquoted
+import com.intellij.codeInsight.intention.preview.IntentionPreviewInfo
+import com.intellij.codeInsight.intention.preview.IntentionPreviewUtils
+import com.intellij.codeInspection.LocalInspectionTool
+import com.intellij.codeInspection.LocalQuickFix
+import com.intellij.codeInspection.ProblemDescriptor
+import com.intellij.codeInspection.ProblemHighlightType
+import com.intellij.codeInspection.ProblemsHolder
+import com.intellij.lang.jsgraphql.psi.GraphQLArrayValue
+import com.intellij.lang.jsgraphql.psi.GraphQLDirective
+import com.intellij.lang.jsgraphql.psi.GraphQLElementFactory
+import com.intellij.lang.jsgraphql.psi.GraphQLSchemaExtension
+import com.intellij.lang.jsgraphql.psi.GraphQLVisitor
+import com.intellij.openapi.project.Project
+import com.intellij.psi.PsiElementVisitor
+import com.intellij.psi.util.parentOfType
+
+class ApolloMissingGraphQLDefinitionImportInspection : LocalInspectionTool() {
+  override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean): PsiElementVisitor {
+    return object : GraphQLVisitor() {
+      override fun visitDirective(o: GraphQLDirective) {
+        super.visitDirective(o)
+        if (!o.project.apolloProjectService.apolloVersion.isAtLeastV4) return
+        visitDirective(o, holder, NULLABILITY_DEFINITIONS, NULLABILITY_URL, ProblemHighlightType.LIKE_UNKNOWN_SYMBOL)
+        visitDirective(o, holder, KOTLIN_LABS_DEFINITIONS, KOTLIN_LABS_URL, ProblemHighlightType.WEAK_WARNING)
+      }
+    }
+  }
+
+  private fun visitDirective(
+      directiveElement: GraphQLDirective,
+      holder: ProblemsHolder,
+      definitions: List<GQLDefinition>,
+      definitionsUrl: String,
+      highlightType: ProblemHighlightType,
+  ) {
+    if (directiveElement.name !in definitions.directives().map { it.name }) return
+    val message = if (highlightType == ProblemHighlightType.WEAK_WARNING) "inspection.missingGraphQLDefinitionImport.reportText.warning" else "inspection.missingGraphQLDefinitionImport.reportText.error"
+    if (!directiveElement.isImported(definitionsUrl)) {
+      val typeKind = ApolloBundle.message("inspection.missingGraphQLDefinitionImport.reportText.directive")
+      holder.registerProblem(
+          directiveElement,
+          ApolloBundle.message(message, typeKind, directiveElement.name!!),
+          highlightType,
+          ImportDefinitionQuickFix(typeKind = typeKind, elementName = directiveElement.name!!, definitions = definitions, definitionsUrl = definitionsUrl),
+      )
+    } else {
+      val directiveDefinition = definitions.directives().firstOrNull { it.name == directiveElement.name } ?: return
+      val knownDefinitionNames = definitions.filterIsInstance<GQLNamed>().map { it.name }
+      val arguments = directiveElement.arguments?.argumentList.orEmpty()
+      for (argument in arguments) {
+        val argumentDefinition = directiveDefinition.arguments.firstOrNull { it.name == argument.name } ?: continue
+        val argumentTypeToImport = argumentDefinition.type.rawType().name.takeIf { it in knownDefinitionNames } ?: continue
+        if (!isImported(directiveElement, argumentTypeToImport, definitionsUrl)) {
+          val typeKind = getTypeKind(argumentTypeToImport)
+          holder.registerProblem(
+              argument,
+              ApolloBundle.message(message, typeKind, argumentTypeToImport),
+              highlightType,
+              ImportDefinitionQuickFix(typeKind = typeKind, elementName = argumentTypeToImport, definitions = definitions, definitionsUrl = definitionsUrl),
+          )
+        }
+      }
+    }
+  }
+}
+
+private fun getTypeKind(typeName: String): String {
+  val typeDefinition = NULLABILITY_DEFINITIONS.firstOrNull { it is GQLNamed && it.name == typeName } ?: return "unknown"
+  return ApolloBundle.message(
+      when (typeDefinition) {
+        is GQLDirectiveDefinition -> "inspection.missingGraphQLDefinitionImport.reportText.directive"
+        is GQLEnumTypeDefinition -> "inspection.missingGraphQLDefinitionImport.reportText.enum"
+        is GQLInputObjectTypeDefinition -> "inspection.missingGraphQLDefinitionImport.reportText.input"
+        is GQLScalarTypeDefinition -> "inspection.missingGraphQLDefinitionImport.reportText.scalar"
+        else -> return "unknown"
+      }
+  )
+}
+
+private class ImportDefinitionQuickFix(
+    val typeKind: String,
+    val elementName: String,
+    private val definitions: List<GQLDefinition>,
+    private val definitionsUrl: String,
+) : LocalQuickFix {
+  override fun getName() = ApolloBundle.message("inspection.missingGraphQLDefinitionImport.quickFix", typeKind, "'$elementName'")
+  override fun getFamilyName() = name
+
+  override fun availableInBatchMode() = false
+  override fun generatePreview(project: Project, previewDescriptor: ProblemDescriptor): IntentionPreviewInfo = IntentionPreviewInfo.EMPTY
+
+  override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
+    if (!IntentionPreviewUtils.isIntentionPreviewActive()) project.telemetryService.logEvent(TelemetryEvent.ApolloIjMissingGraphQLDefinitionImportQuickFix())
+
+    val element = descriptor.psiElement.parentOfType<GraphQLDirective>(withSelf = true)!!
+    val schemaFiles = element.schemaFiles()
+    val linkDirective = schemaFiles.flatMap { it.linkDirectives(definitionsUrl) }.firstOrNull()
+
+    if (linkDirective == null) {
+      val linkDirectiveSchemaExtension = createLinkDirectiveSchemaExtension(project, setOf(element.nameForImport), definitions, definitionsUrl)
+      val extraSchemaFile = schemaFiles.firstOrNull { it.name == "extra.graphqls" }
+      if (extraSchemaFile == null) {
+        GraphQLElementFactory.createFile(project, linkDirectiveSchemaExtension.text).also {
+          // Save the file to the project
+          it.name = "extra.graphqls"
+          schemaFiles.first().containingDirectory!!.add(it)
+
+          // There's a new schema file, reload the configuration
+          project.gradleToolingModelService.triggerFetchToolingModels()
+        }
+      } else {
+        val addedElement = extraSchemaFile.addBefore(linkDirectiveSchemaExtension, extraSchemaFile.firstChild)
+        extraSchemaFile.addAfter(GraphQLElementFactory.createWhiteSpace(project, "\n\n"), addedElement)
+      }
+    } else {
+      val importedNames = buildSet {
+        addAll(linkDirective.arguments!!.argumentList.firstOrNull { it.name == "import" }?.value?.cast<GraphQLArrayValue>()?.valueList.orEmpty().map { it.text.unquoted() })
+        add(element.nameForImport)
+      }
+      linkDirective.replace(createLinkDirective(project, importedNames, definitions, definitionsUrl))
+    }
+  }
+}
+
+private fun createLinkDirectiveSchemaExtension(
+    project: Project,
+    importedNames: Set<String>,
+    definitions: List<GQLDefinition>,
+    definitionsUrl: String,
+): GraphQLSchemaExtension {
+  // If any of the imported name is a directive, add its argument types to the import list
+  val knownDefinitionNames = definitions.filterIsInstance<GQLNamed>().map { it.name }
+  val additionalNames = importedNames.flatMap { importedName ->
+    definitions.directives().firstOrNull { "@${it.name}" == importedName }
+        ?.arguments
+        ?.map { it.type.rawType().name }
+        ?.filter { it in knownDefinitionNames }.orEmpty()
+  }.toSet()
+
+  return GraphQLElementFactory.createFile(
+      project,
+      """
+        extend schema
+        @link(
+          url: "$definitionsUrl",
+          import: [${(importedNames + additionalNames).joinToString { it.quoted() }}]
+        )
+      """.trimIndent()
+  )
+      .findChildrenOfType<GraphQLSchemaExtension>().single()
+}
+
+private fun createLinkDirective(
+    project: Project,
+    importedNames: Set<String>,
+    definitions: List<GQLDefinition>,
+    definitionsUrl: String,
+): GraphQLDirective {
+  return createLinkDirectiveSchemaExtension(project, importedNames, definitions, definitionsUrl).directives.single()
+}

intellij-plugin/src/main/kotlin/com/apollographql/ijplugin/inspection/GraphQL.kt

@@ -1,10 +1,16 @@
 package com.apollographql.ijplugin.inspection
 
+import com.apollographql.ijplugin.util.findPsiFileByUrl
+import com.intellij.lang.jsgraphql.ide.config.GraphQLConfigProvider
+import com.intellij.lang.jsgraphql.psi.GraphQLDirective
+import com.intellij.lang.jsgraphql.psi.GraphQLElement
 import com.intellij.lang.jsgraphql.psi.GraphQLFieldDefinition
+import com.intellij.lang.jsgraphql.psi.GraphQLFile
 import com.intellij.lang.jsgraphql.psi.GraphQLIdentifier
 import com.intellij.lang.jsgraphql.psi.GraphQLInterfaceTypeDefinition
 import com.intellij.lang.jsgraphql.psi.GraphQLNamedTypeDefinition
 import com.intellij.lang.jsgraphql.psi.GraphQLObjectTypeDefinition
+import com.intellij.lang.jsgraphql.psi.GraphQLValue
 import com.intellij.psi.util.parentOfType
 
 /**
@@ -66,3 +72,17 @@ private fun matchingFieldCoordinates(
   if (implementedInterfaceTypeDefinitions.isEmpty()) return fieldCoordinates
   return fieldCoordinates + implementedInterfaceTypeDefinitions.flatMap { matchingFieldCoordinates(fieldDefinition, it) }
 }
+
+/**
+ * Return the schema files associated with the given element.
+ */
+fun GraphQLElement.schemaFiles(): List<GraphQLFile> {
+  val containingFile = containingFile ?: return emptyList()
+  val projectConfig = GraphQLConfigProvider.getInstance(project).resolveProjectConfig(containingFile) ?: return emptyList()
+  return projectConfig.schema.mapNotNull { schema ->
+    schema.filePath?.let { path -> project.findPsiFileByUrl(schema.dir.url + "/" + path) } as? GraphQLFile
+  }
+}
+
+fun GraphQLDirective.argumentValue(argumentName: String): GraphQLValue? =
+    arguments?.argumentList.orEmpty().firstOrNull { it.name == argumentName }?.value

intellij-plugin/src/main/kotlin/com/apollographql/ijplugin/inspection/GraphQLUnresolvedReferenceInspectionSuppressor.kt

@@ -0,0 +1,47 @@
+package com.apollographql.ijplugin.inspection
+
+import com.apollographql.apollo3.ast.GQLDirectiveDefinition
+import com.apollographql.apollo3.ast.linkDefinitions
+import com.intellij.codeInspection.InspectionSuppressor
+import com.intellij.codeInspection.SuppressQuickFix
+import com.intellij.lang.jsgraphql.psi.GraphQLArgument
+import com.intellij.lang.jsgraphql.psi.GraphQLDirective
+import com.intellij.lang.jsgraphql.psi.GraphQLDirectivesAware
+import com.intellij.psi.PsiElement
+
+private val KNOWN_DIRECTIVES: List<GQLDirectiveDefinition> by lazy {
+  linkDefinitions().directives() + NULLABILITY_DEFINITIONS.directives()
+}
+
+/**
+ * Do not highlight certain known directives as unresolved references.
+ *
+ * TODO: remove this once https://github.com/JetBrains/js-graphql-intellij-plugin/pull/698 is merged.
+ */
+class GraphQLUnresolvedReferenceInspectionSuppressor : InspectionSuppressor {
+  override fun isSuppressedFor(element: PsiElement, toolId: String): Boolean {
+    val parent = element.parent
+    return when (toolId) {
+      "GraphQLUnresolvedReference" -> parent.isKnownDirective() || parent.isKnownDirectiveArgument()
+
+      "GraphQLMissingType" -> element is GraphQLDirectivesAware && element.directives.all { it.isKnownDirective() }
+
+      // We need to suppress this one too because the plugin doesn't know that @link is repeatable
+      "GraphQLDuplicateDirective" -> element is GraphQLDirective && element.name == "link"
+
+      else -> false
+    }
+  }
+
+  override fun getSuppressActions(psiElement: PsiElement?, s: String): Array<SuppressQuickFix> = SuppressQuickFix.EMPTY_ARRAY
+}
+
+private fun PsiElement.isKnownDirective(): Boolean {
+  return this is GraphQLDirective && (name in KNOWN_DIRECTIVES.map { it.name } || this.isImported(NULLABILITY_URL))
+}
+
+private fun PsiElement.isKnownDirectiveArgument(): Boolean {
+  return this is GraphQLArgument &&
+      parent?.parent?.isKnownDirective() == true &&
+      name in KNOWN_DIRECTIVES.firstOrNull { it.name == (parent.parent as GraphQLDirective).nameWithoutPrefix }?.arguments?.map { it.name }.orEmpty()
+}

intellij-plugin/src/main/kotlin/com/apollographql/ijplugin/inspection/Link.kt

@@ -0,0 +1,88 @@
+@file:OptIn(ApolloInternal::class)
+
+package com.apollographql.ijplugin.inspection
+
+import com.apollographql.apollo3.annotations.ApolloInternal
+import com.apollographql.apollo3.ast.GQLDefinition
+import com.apollographql.apollo3.ast.GQLDirectiveDefinition
+import com.apollographql.apollo3.ast.KOTLIN_LABS_VERSION
+import com.apollographql.apollo3.ast.NULLABILITY_VERSION
+import com.apollographql.apollo3.ast.kotlinLabsDefinitions
+import com.apollographql.apollo3.ast.nullabilityDefinitions
+import com.apollographql.ijplugin.util.quoted
+import com.apollographql.ijplugin.util.unquoted
+import com.intellij.lang.jsgraphql.psi.GraphQLArrayValue
+import com.intellij.lang.jsgraphql.psi.GraphQLDirective
+import com.intellij.lang.jsgraphql.psi.GraphQLElement
+import com.intellij.lang.jsgraphql.psi.GraphQLFile
+import com.intellij.lang.jsgraphql.psi.GraphQLNamedElement
+import com.intellij.lang.jsgraphql.psi.GraphQLSchemaDefinition
+import com.intellij.lang.jsgraphql.psi.GraphQLSchemaExtension
+import com.intellij.lang.jsgraphql.psi.GraphQLStringValue
+
+const val NULLABILITY_URL = "https://specs.apollo.dev/nullability/$NULLABILITY_VERSION"
+
+val NULLABILITY_DEFINITIONS: List<GQLDefinition> by lazy {
+  nullabilityDefinitions(NULLABILITY_VERSION)
+}
+
+const val KOTLIN_LABS_URL = "https://specs.apollo.dev/kotlin_labs/$KOTLIN_LABS_VERSION"
+
+val KOTLIN_LABS_DEFINITIONS: List<GQLDefinition> by lazy {
+  kotlinLabsDefinitions(KOTLIN_LABS_VERSION)
+}
+
+const val CATCH = "catch"
+
+fun List<GQLDefinition>.directives(): List<GQLDirectiveDefinition> {
+  return filterIsInstance<GQLDirectiveDefinition>()
+}
+
+fun GraphQLNamedElement.isImported(definitionsUrl: String): Boolean {
+  for (schemaFile in schemaFiles()) {
+    if (schemaFile.hasImportFor(this.name!!, this is GraphQLDirective, definitionsUrl)) return true
+  }
+  return false
+}
+
+fun isImported(element: GraphQLElement, enumName: String, definitionsUrl: String): Boolean {
+  for (schemaFile in element.schemaFiles()) {
+    if (schemaFile.hasImportFor(enumName, false, definitionsUrl)) return true
+  }
+  return false
+}
+
+fun GraphQLFile.linkDirectives(definitionsUrl: String): List<GraphQLDirective> {
+  val schemaDirectives = typeDefinitions.filterIsInstance<GraphQLSchemaExtension>().flatMap { it.directives } +
+      typeDefinitions.filterIsInstance<GraphQLSchemaDefinition>().flatMap { it.directives }
+  return schemaDirectives.filter { directive ->
+    directive.name == "link" &&
+        directive.arguments?.argumentList.orEmpty().any { arg -> arg.name == "url" && arg.value?.text == definitionsUrl.quoted() }
+  }
+}
+
+private fun GraphQLFile.hasImportFor(name: String, isDirective: Boolean, definitionsUrl: String): Boolean {
+  for (directive in linkDirectives(definitionsUrl)) {
+    val importArgValue = directive.argumentValue("import") as? GraphQLArrayValue
+    if (importArgValue == null) {
+      // Default import is everything - see https://specs.apollo.dev/link/v1.0/#@link.url
+      val asArgValue = directive.argumentValue("as") as? GraphQLStringValue
+      // Default prefix is the name part of the url
+      val prefix = (asArgValue?.text?.unquoted() ?: "nullability") + "__"
+      if (name.startsWith(prefix)) return true
+    } else {
+      if (importArgValue.valueList.any { it.text == name.nameForImport(isDirective).quoted() }) {
+        return true
+      }
+    }
+  }
+  return false
+}
+
+val String.nameWithoutPrefix get() = substringAfter("__")
+
+val GraphQLNamedElement.nameWithoutPrefix get() = name!!.nameWithoutPrefix
+
+fun String.nameForImport(isDirective: Boolean) = "${if (isDirective) "@" else ""}${this.nameWithoutPrefix}"
+
+val GraphQLNamedElement.nameForImport get() = if (this is GraphQLDirective) "@$nameWithoutPrefix" else nameWithoutPrefix

intellij-plugin/src/main/kotlin/com/apollographql/ijplugin/telemetry/TelemetrySession.kt

@@ -384,6 +384,10 @@ sealed class TelemetryEvent(
    */
   class ApolloIjInputConstructorChangeToBuilderIntentionApply : TelemetryEvent("akij_input_constructor_change_to_builder_intention_apply", null)
 
+  /**
+   * User applied the 'Import directive' quickfix for the 'Missing GraphQL definition import' inspection of the Apollo Kotlin IntelliJ plugin.
+   */
+  class ApolloIjMissingGraphQLDefinitionImportQuickFix : TelemetryEvent("akij_missing_graphql_definition_import_quickfix", null)
 }
 
 class TelemetryEventList {