sqldelight · aperfilyev · Feb 23, 2022 · Feb 20, 2022
diff --git a/...lugin/src/main/kotlin/app/cash/sqldelight/intellij/SqlDelightClassNameElementAnnotator.kt b/...lugin/src/main/kotlin/app/cash/sqldelight/intellij/SqlDelightClassNameElementAnnotator.kt
@@ -16,9 +16,11 @@
 package app.cash.sqldelight.intellij
 
 import app.cash.sqldelight.core.lang.SqlDelightFile
-import app.cash.sqldelight.core.lang.psi.ImportStmtMixin
 import app.cash.sqldelight.core.lang.psi.JavaTypeMixin
+import app.cash.sqldelight.core.lang.util.findChildOfType
 import app.cash.sqldelight.core.lang.util.findChildrenOfType
+import app.cash.sqldelight.core.psi.SqlDelightImportStmt
+import app.cash.sqldelight.core.psi.SqlDelightImportStmtList
 import app.cash.sqldelight.intellij.intentions.AddImportIntention
 import app.cash.sqldelight.intellij.util.PsiClassSearchHelper
 import com.intellij.codeInspection.ProblemHighlightType
@@ -30,47 +32,55 @@ import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiElement
 import com.intellij.psi.search.GlobalSearchScope
 
-class SqlDelightClassNameElementAnnotator : Annotator {
+internal data class AnnotationData(
+  val element: PsiElement,
+  val intentionAvailable: Boolean,
+  val classes: List<PsiClass> = emptyList(),
+)
+
+internal class SqlDelightClassNameElementAnnotator : Annotator {
   override fun annotate(element: PsiElement, holder: AnnotationHolder) {
     if (element !is JavaTypeMixin || element.reference.resolve() != null) {
       return
     }
-
-    val project = element.project
-    val sqlDelightFile = element.containingFile as SqlDelightFile
-    val module = ModuleUtil.findModuleForFile(sqlDelightFile) ?: return
-    val scope = GlobalSearchScope.moduleWithDependenciesAndLibrariesScope(module, false)
-    val outerClassElement = element.firstChild
-    val outerClassName = outerClassElement.text
-    val classes = PsiClassSearchHelper.getClassesByShortName(outerClassName, project, scope)
-
-    val hasImport = sqlDelightFile.hasImport(outerClassName)
-    val psiElement = if (hasImport) {
-      missingNestedClass(classes, element)
+    val data = if (element.context is SqlDelightImportStmt) {
+      AnnotationData(element, false)
     } else {
-      outerClassElement
+      val project = element.project
+      val sqlDelightFile = element.containingFile as SqlDelightFile
+      val module = ModuleUtil.findModuleForFile(sqlDelightFile) ?: return
+      val scope = GlobalSearchScope.moduleWithDependenciesAndLibrariesScope(module, false)
+      val outerClassElement = element.firstChild
+      val outerClassName = outerClassElement.text
+      val classes = PsiClassSearchHelper.getClassesByShortName(outerClassName, project, scope)
+      val hasImport = sqlDelightFile.hasImport(outerClassName)
+      val enabled = classes.isNotEmpty() && !hasImport
+      val psiElement = if (hasImport) {
+        missingNestedClass(classes, element)
+      } else {
+        outerClassElement
+      }
+      AnnotationData(psiElement, enabled, classes)
     }
 
-    val needsQuickFix = classes.isNotEmpty() && !hasImport
-
-    holder.newAnnotation(HighlightSeverity.ERROR, "Unresolved reference: ${psiElement.text}")
-      .range(psiElement)
+    holder.newAnnotation(HighlightSeverity.ERROR, "Unresolved reference: ${data.element.text}")
+      .range(data.element)
       .highlightType(ProblemHighlightType.LIKE_UNKNOWN_SYMBOL)
-      .apply {
-        if (needsQuickFix) {
-          withFix(AddImportIntention(outerClassName))
-        }
-      }
+      .withFix(AddImportIntention(data.classes, data.intentionAvailable))
       .create()
   }
 
   private fun SqlDelightFile.hasImport(outerClassName: String): Boolean {
-    return sqlStmtList?.findChildrenOfType<ImportStmtMixin>()
+    return sqlStmtList?.findChildOfType<SqlDelightImportStmtList>()
+      ?.importStmtList
       .orEmpty()
       .any { it.javaType.text.endsWith(outerClassName) }
   }
 
-  private fun missingNestedClass(classes: List<PsiClass>, javaTypeMixin: JavaTypeMixin): PsiElement {
+  private fun missingNestedClass(
+    classes: List<PsiClass>,
+    javaTypeMixin: JavaTypeMixin
+  ): PsiElement {
     val elementText = javaTypeMixin.text
     val className = classes.map { clazz -> findMissingNestedClassName(clazz, elementText) }
       .maxByOrNull { it.length }
@@ -83,10 +93,10 @@ class SqlDelightClassNameElementAnnotator : Annotator {
     if (psiClass.name != nestedClassName) {
       return nestedClassName
     }
-    val nextName = className.removePrefix(nestedClassName).removePrefix(".")
+    val nextName = className.substringAfter(".")
     val lookupString = nextName.substringBefore(".")
     val nextClass = psiClass.innerClasses.firstOrNull { clazz ->
-      clazz.textMatches(lookupString)
+      clazz.name == lookupString
     }
     return nextClass?.let { findMissingNestedClassName(it, nextName) } ?: nextName
   }

diff --git a/...idea-plugin/src/main/kotlin/app/cash/sqldelight/intellij/intentions/AddImportIntention.kt b/...idea-plugin/src/main/kotlin/app/cash/sqldelight/intellij/intentions/AddImportIntention.kt
@@ -1,15 +1,14 @@
 package app.cash.sqldelight.intellij.intentions
 
+import app.cash.sqldelight.core.lang.psi.JavaTypeMixin
 import app.cash.sqldelight.core.lang.util.findChildrenOfType
 import app.cash.sqldelight.core.psi.SqlDelightImportStmt
-import app.cash.sqldelight.intellij.util.PsiClassSearchHelper
 import com.intellij.codeInsight.daemon.QuickFixBundle
 import com.intellij.codeInsight.intention.BaseElementAtCaretIntentionAction
 import com.intellij.codeInsight.navigation.NavigationUtil
 import com.intellij.openapi.command.WriteCommandAction
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.editor.Editor
-import com.intellij.openapi.module.ModuleUtil
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.popup.PopupStep
 import com.intellij.openapi.ui.popup.util.BaseListPopupStep
@@ -18,23 +17,23 @@ import com.intellij.psi.PsiClass
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiElement
 import com.intellij.psi.PsiFile
-import com.intellij.psi.search.GlobalSearchScope
 import com.intellij.ui.popup.list.ListPopupImpl
 import javax.swing.Icon
 
-class AddImportIntention(private val key: String) : BaseElementAtCaretIntentionAction() {
-  override fun getFamilyName(): String = INTENTIONS_FAMILY_NAME_IMPORTS
+internal class AddImportIntention(
+  private val classes: List<PsiClass>,
+  private val isAvailable: Boolean,
+) : BaseElementAtCaretIntentionAction() {
 
-  override fun getText(): String = "Add import for $key"
+  override fun getFamilyName(): String = INTENTIONS_FAMILY_NAME_IMPORTS
 
   override fun isAvailable(project: Project, editor: Editor, element: PsiElement): Boolean {
-    return true
+    if (element !is JavaTypeMixin || element.context is SqlDelightImportStmt) return false
+    text = "Add import for ${element.text}"
+    return isAvailable
   }
 
   override fun invoke(project: Project, editor: Editor, element: PsiElement) {
-    val module = ModuleUtil.findModuleForPsiElement(element) ?: return
-    val scope = GlobalSearchScope.moduleWithDependenciesAndLibrariesScope(module)
-    val classes = PsiClassSearchHelper.getClassesByShortName(key, project, scope)
     val document = editor.document
     val file = element.containingFile
     if (classes.size == 1) {

diff --git a/...plugin/src/test/kotlin/app/cash/sqldelight/intellij/annotations/ClassNameAnnotatorTest.kt b/...plugin/src/test/kotlin/app/cash/sqldelight/intellij/annotations/ClassNameAnnotatorTest.kt
@@ -1,6 +1,8 @@
 package app.cash.sqldelight.intellij.annotations
 
+import app.cash.sqldelight.core.lang.SqlDelightFileType
 import app.cash.sqldelight.intellij.SqlDelightProjectTestCase
+import org.jetbrains.kotlin.idea.KotlinFileType
 
 class ClassNameAnnotatorTest : SqlDelightProjectTestCase() {
   fun testResolveOnSamePackageImport() {
@@ -10,4 +12,73 @@ class ClassNameAnnotatorTest : SqlDelightProjectTestCase() {
 
     myFixture.checkHighlighting()
   }
+
+  fun testUnresolvedClassName() {
+    myFixture.configureByText(
+      SqlDelightFileType,
+      """
+      |CREATE TABLE test (
+      |  value TEXT AS <error descr="Unresolved reference: KoolKidz">KoolKidz</error>
+      |);
+    """.trimMargin()
+    )
+
+    myFixture.checkHighlighting()
+  }
+
+  fun testInnerClassWorksFine() {
+    myFixture.configureByText(
+      SqlDelightFileType,
+      """
+      |import com.example.KotlinClass;
+      |
+      |CREATE TABLE test (
+      |  value TEXT AS KotlinClass.InnerClass
+      |);
+    """.trimMargin()
+    )
+
+    myFixture.checkHighlighting()
+  }
+
+  fun testUnresolvedImport() {
+    myFixture.configureByText(
+      SqlDelightFileType,
+      """
+      |import <error descr="Unresolved reference: com.somepackage.SomeClass">com.somepackage.SomeClass</error>;
+      |
+      |CREATE TABLE new_table (
+      |  col TEXT AS <error descr="Unresolved reference: SomeClass">SomeClass</error> NOT NULL
+      |);
+    """.trimMargin()
+    )
+
+    myFixture.checkHighlighting()
+  }
+
+  fun testUnresolvedNestedClass() {
+    myFixture.configureByText(
+      KotlinFileType.INSTANCE,
+      """
+        |package com.somepackage
+        |
+        |class SomeClass {
+        |  class FirstLevel
+        |}
+      """.trimMargin()
+    )
+
+    myFixture.configureByText(
+      SqlDelightFileType,
+      """
+        |import com.somepackage.SomeClass;
+        |
+        |CREATE TABLE new_table (
+        |  col TEXT AS SomeClass.FirstLevel.<error descr="Unresolved reference: SecondLevel">SecondLevel</error> NOT NULL
+        |);
+        """.trimMargin()
+    )
+
+    myFixture.checkHighlighting()
+  }
 }
diff --git a/...lugin/src/test/kotlin/app/cash/sqldelight/intellij/autocomplete/ClassAutocompleteTests.kt b/...lugin/src/test/kotlin/app/cash/sqldelight/intellij/autocomplete/ClassAutocompleteTests.kt