Add derivation of Schema for union types (closes ghostdogpr#1926)

Author: oyvindberg

build.sbt

@@ -39,7 +39,7 @@ Global / onChangedBuildSource := ReloadOnSourceChanges
 
 inThisBuild(
   List(
-    scalaVersion             := scala213,
+    scalaVersion             := scala3,
     crossScalaVersions       := allScala,
     organization             := "com.github.ghostdogpr",
     homepage                 := Some(url("https://github.com/ghostdogpr/caliban")),

core/src/main/scala-3/caliban/schema/SchemaDerivation.scala

@@ -123,6 +123,8 @@ trait SchemaDerivation[R] extends CommonSchemaDerivation {
 
   inline def genDebug[R, A]: Schema[R, A] = PrintDerived(derived[R, A])
 
+  inline def unionType[T]: Schema[R, T] = ${ TypeUnionDerivation.typeUnionSchema[R, T] }
+
   final lazy val auto = new AutoSchemaDerivation[Any] {}
 
   final class SemiAuto[A](impl: Schema[R, A]) extends Schema[R, A] {

core/src/main/scala-3/caliban/schema/TypeUnionDerivation.scala

@@ -0,0 +1,71 @@
+package caliban.schema
+
+import caliban.introspection.adt.__Type
+
+import scala.quoted.*
+
+object TypeUnionDerivation {
+  inline def derived[R, T]: Schema[R, T] = ${ typeUnionSchema[R, T] }
+
+  def typeUnionSchema[R: Type, T: Type](using quotes: Quotes): Expr[Schema[R, T]] = {
+    import quotes.reflect.*
+
+    class TypeAndSchema[A](val typeRef: String, val schema: Expr[Schema[R, A]], val tpe: Type[A])
+
+    def rec[A](using tpe: Type[A]): List[TypeAndSchema[?]] =
+      TypeRepr.of(using tpe).dealias match {
+        case OrType(l, r) =>
+          rec(using l.asType.asInstanceOf[Type[Any]]) ++ rec(using r.asType.asInstanceOf[Type[Any]])
+        case otherRepr =>
+          val otherString: String = otherRepr.show
+          val expr: TypeAndSchema[A] =
+            Expr.summon[Schema[R, A]] match {
+              case Some(foundSchema) =>
+                TypeAndSchema[A](otherString, foundSchema, otherRepr.asType.asInstanceOf[Type[A]])
+              case None =>
+                quotes.reflect.report.errorAndAbort(s"Couldn't resolve Schema[Any, $otherString]")
+            }
+
+          List(expr)
+      }
+
+    val typeAndSchemas: List[TypeAndSchema[?]] = rec[T]
+
+    val schemaByTypeNameList: Expr[List[(String, Schema[R, Any])]] = Expr.ofList(
+      typeAndSchemas.map { case (tas: TypeAndSchema[a]) =>
+        given Type[a] = tas.tpe
+      '{ (${ Expr(tas.typeRef) }, ${ tas.schema }.asInstanceOf[Schema[R, Any]]) }
+      }
+    )
+    val name = TypeRepr.of[T].show
+
+    if (name.contains("|")) {
+      report.error(s"You must explicitly add type parameter to derive Schema for a union type in order to capture the name of the type alias")
+    }
+
+    val ret = '{
+      val schemaByName: Map[String, Schema[R, Any]] = ${ schemaByTypeNameList }.toMap
+      new Schema[R, T] {
+
+        def resolve(value: T): Step[R] = {
+          var ret: Step[R] = null
+          ${
+            Expr.block(
+              typeAndSchemas.map { case (tas: TypeAndSchema[a]) =>
+                given Type[a] = tas.tpe
+                '{ if value.isInstanceOf[a] then ret = schemaByName(${ Expr(tas.typeRef) }).resolve(value) }
+              },
+              '{ require(ret != null, s"no schema for ${value}") }
+            )
+          }
+          ret
+        }
+
+        def toType(isInput: Boolean, isSubscription: Boolean): __Type =
+          Types.makeUnion(Some(${ Expr(name) }), None, schemaByName.values.map(_.toType_(isInput, isSubscription)).toList)
+      }
+    }
+//    quotes.reflect.report.warning(ret.show)
+    ret
+  }
+}

core/src/test/scala-3/caliban/schema/Scala3DerivesSpec.scala

@@ -273,6 +273,51 @@ object Scala3DerivesSpec extends ZIOSpecDefault {
           data1 == """{"enum2String":"ENUM1"}""",
           data2 == """{"enum2String":"ENUM2"}"""
         )
+      },
+      test("union type") {
+        final case class Foo(value: String) derives Schema.SemiAuto
+        final case class Bar(foo: Int) derives Schema.SemiAuto
+        type Payload = Foo | Bar
+
+        given Schema[Any, Payload] = Schema.unionType[Payload]
+
+        final case class QueryInput(isFoo: Boolean) derives ArgBuilder, Schema.SemiAuto
+        final case class Query(testQuery: QueryInput => zio.UIO[Payload]) derives Schema.SemiAuto
+
+        val gql = graphQL(RootResolver(Query(i => ZIO.succeed(if (i.isFoo) Foo("foo") else Bar(1)))))
+
+        assertTrue(
+          gql.render ==
+            """schema {
+              |  query: Query
+              |}
+              |
+              |union Payload = Foo | Bar
+              |
+              |type Bar {
+              |  foo: Int!
+              |}
+              |
+              |type Foo {
+              |  value: String!
+              |}
+              |
+              |type Query {
+              |  testQuery(isFoo: Boolean!): Payload!
+              |}""".stripMargin
+        )
+
+        val interpreter = gql.interpreterUnsafe
+
+        for {
+          res1 <- interpreter.execute("{ testQuery(isFoo: true){ ... on Foo { value } } }")
+          res2 <- interpreter.execute("{ testQuery(isFoo: false){ ... on Bar { foo } } }")
+          data1 = res1.data.toString
+          data2 = res2.data.toString
+        } yield assertTrue(
+          data1 == """{"testQuery":{"value":"foo"}}""",
+          data2 == """{"testQuery":{"foo":1}}"""
+        )
       }
     )
   }