Improve introspection (#2290)

ghostdogpr · Jun 18, 2024 · bedb71d · bedb71d
1 parent e728210
commit bedb71d
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 26 deletions.
diff --git a/core/src/main/scala-2/caliban/introspection/IntrospectionDerivation.scala b/core/src/main/scala-2/caliban/introspection/IntrospectionDerivation.scala
@@ -6,7 +6,7 @@ import caliban.schema.Schema
 import caliban.schema.Schema.auto._
 
 trait IntrospectionDerivation {
-  implicit lazy val typeSchema: Schema[Any, __Type] = genAll
+  private implicit lazy val typeSchema: Schema[Any, __Type] = genAll
 
   val introspectionSchema: Schema[Any, __Introspection] = genAll
 }
diff --git a/core/src/main/scala-3/caliban/introspection/IntrospectionDerivation.scala b/core/src/main/scala-3/caliban/introspection/IntrospectionDerivation.scala
@@ -5,14 +5,24 @@ import caliban.Value.StringValue
 import caliban.introspection.adt._
 import caliban.parsing.adt.Directive
 import caliban.schema.Schema
-import caliban.schema.Schema.auto._
-import caliban.schema.ArgBuilder.auto._
+import caliban.schema.ArgBuilder
 
 trait IntrospectionDerivation {
-  import Schema._
+  private given Schema[Any, __InputValue]     = Schema.derived
+  private given Schema[Any, __EnumValue]      = Schema.derived
+  private given Schema[Any, __Field]          = Schema.derived
+  private given Schema[Any, __Type]           = Schema.derived
+  private given Schema[Any, __TypeArgs]       = Schema.derived
+  private given Schema[Any, __Schema]         = Schema.derived
+  private given Schema[Any, __Directive]      = Schema.derived
+  private given Schema[Any, __DeprecatedArgs] = Schema.derived
 
-  implicit lazy val inputValueSchema: Schema[Any, InputValue]   = Schema.gen
-  implicit lazy val typeSchema: Schema[Any, __Type]             = Schema.gen
-  implicit lazy val __directiveSchema: Schema[Any, __Directive] = Schema.gen
-  val introspectionSchema: Schema[Any, __Introspection]         = Schema.gen
+  // Unions, so we can auto-derive them cheaply
+  private given Schema[Any, __TypeKind]          = Schema.Auto.derived
+  private given Schema[Any, __DirectiveLocation] = Schema.Auto.derived
+
+  private given ArgBuilder[__TypeArgs]       = ArgBuilder.derived
+  private given ArgBuilder[__DeprecatedArgs] = ArgBuilder.derived
+
+  val introspectionSchema: Schema[Any, __Introspection] = Schema.derived
 }
diff --git a/core/src/main/scala/caliban/introspection/Introspector.scala b/core/src/main/scala/caliban/introspection/Introspector.scala
@@ -8,10 +8,11 @@ import caliban.parsing.adt.Selection.Field
 import caliban.schema.Step.QueryStep
 import caliban.schema._
 import caliban.wrappers.Wrapper.IntrospectionWrapper
-import zio.ZIO
+import zio.{ Exit, ZIO }
 import zio.query.ZQuery
 
 import scala.annotation.tailrec
+import scala.collection.mutable
 
 object Introspector extends IntrospectionDerivation {
   private[caliban] val directives = List(
@@ -75,13 +76,15 @@ object Introspector extends IntrospectionDerivation {
         case wrapper :: tail => wrap(wrapper.wrap(query))(tail)
       }
 
-    val types = (rootType.types ++ introspectionRootType.types - "__Introspection")
-      .updated("Boolean", Types.boolean) // because of skip and include
-      .updated("String", Types.string)   // because of specifiedBy
-      .values
-      .toList
-      .sortBy(_.name.getOrElse(""))
+    val typesMap = new mutable.HashMap[String, __Type]()
+    typesMap.sizeHint(rootType.types.size + introspectionRootType.types.size + 2)
+    typesMap ++= rootType.types
+    typesMap ++= introspectionRootType.types
+    typesMap.remove("__Introspection")
+    typesMap.update("Boolean", Types.boolean) // because of skip and include
+    typesMap.update("String", Types.string)   // because of specifiedBy
 
+    val types    = typesMap.values.toList.sortBy(_.name)
     val hasOneOf = types.exists(_._isOneOfInput)
 
     val resolver = __Introspection(
@@ -93,25 +96,23 @@ object Introspector extends IntrospectionDerivation {
         types,
         directives ++ (if (hasOneOf) List(oneOfDirective) else Nil) ++ rootType.additionalDirectives
       ),
-      args => types.find(_.name.contains(args.name))
+      args => typesMap.get(args.name)
     )
 
-    RootSchema(
-      Operation(
-        introspectionType,
-        QueryStep(ZQuery.fromZIO(wrap(ZIO.succeed(resolver))(introWrappers)).map(introspectionSchema.resolve))
-      ),
-      None,
-      None
-    )
+    val step = introWrappers match {
+      case Nil => introspectionSchema.resolve(resolver)
+      case ws  => QueryStep(ZQuery.fromZIO(wrap(ZIO.succeed(resolver))(ws).map(introspectionSchema.resolve)))
+    }
+
+    RootSchema(Operation(introspectionType, step), None, None)
   }
 
   private[caliban] def isIntrospection(document: Document): Boolean =
     document.definitions.forall {
       case OperationDefinition(_, _, _, _, selectionSet) =>
         selectionSet.nonEmpty && selectionSet.forall {
-          case Field(_, name, _, _, _, _) => name == "__schema" || name == "__type"
-          case _                          => false
+          case Field(_, "__schema" | "__type", _, _, _, _) => true
+          case _                                           => false
         }
       case _                                             => true
     }