Partial union type query #446

client/src/main/scala/caliban/client/FieldBuilder.scala

@@ -20,19 +20,19 @@ trait FieldBuilder[+A] {
 }
 
 object FieldBuilder {
-  case class Scalar[A]()(implicit decoder: ScalarDecoder[A])       extends FieldBuilder[A]         {
+  case class Scalar[A]()(implicit decoder: ScalarDecoder[A])       extends FieldBuilder[A]               {
     override def fromGraphQL(value: __Value): Either[DecodingError, A] = decoder.decode(value)
     override def toSelectionSet: List[Selection]                       = Nil
   }
-  case class Obj[Origin, A](builder: SelectionBuilder[Origin, A])  extends FieldBuilder[A]         {
+  case class Obj[Origin, A](builder: SelectionBuilder[Origin, A])  extends FieldBuilder[A]               {
     override def fromGraphQL(value: __Value): Either[DecodingError, A] =
       value match {
         case o: __ObjectValue => builder.fromGraphQL(o)
         case _                => Left(DecodingError(s"Field $value is not an object"))
       }
     override def toSelectionSet: List[Selection]                       = builder.toSelectionSet
   }
-  case class ListOf[A](builder: FieldBuilder[A])                   extends FieldBuilder[List[A]]   {
+  case class ListOf[A](builder: FieldBuilder[A])                   extends FieldBuilder[List[A]]         {
     override def fromGraphQL(value: __Value): Either[DecodingError, List[A]] =
       value match {
         case __ListValue(items) =>
@@ -43,15 +43,15 @@ object FieldBuilder {
       }
     override def toSelectionSet: List[Selection]                             = builder.toSelectionSet
   }
-  case class OptionOf[A](builder: FieldBuilder[A])                 extends FieldBuilder[Option[A]] {
+  case class OptionOf[A](builder: FieldBuilder[A])                 extends FieldBuilder[Option[A]]       {
     override def fromGraphQL(value: __Value): Either[DecodingError, Option[A]] =
       value match {
         case `__NullValue` => Right(None)
         case other         => builder.fromGraphQL(other).map(Some(_))
       }
     override def toSelectionSet: List[Selection]                               = builder.toSelectionSet
   }
-  case class ChoiceOf[A](builderMap: Map[String, FieldBuilder[A]]) extends FieldBuilder[A]         {
+  case class ChoiceOf[A](builderMap: Map[String, FieldBuilder[A]]) extends FieldBuilder[A]               {
     override def fromGraphQL(value: __Value): Either[DecodingError, A] =
       value match {
         case __ObjectValue(fields) =>
@@ -67,8 +67,21 @@ object FieldBuilder {
         case _                     => Left(DecodingError(s"Field $value is not an object"))
       }
 
-    override def toSelectionSet: List[Selection] =
+    override def toSelectionSet: List[Selection] = {
+      val filteredBuilderMap = builderMap.filter((f) => !isNullField(f._2));
       Selection.Field(None, "__typename", Nil, Nil, Nil, 0) ::
-        builderMap.map { case (k, v) => Selection.InlineFragment(k, v.toSelectionSet) }.toList
+        filteredBuilderMap.map { case (k, v) =>
+          Selection.InlineFragment(k, v.toSelectionSet)
+        }.toList
+    }
+  }
+  case object NullField                                            extends FieldBuilder[Option[Nothing]] {
+    override def fromGraphQL(value: __Value): Either[DecodingError, Option[Nothing]] = Right(None)
+    override def toSelectionSet: List[Selection]                                     = Nil
+  }
+
+  def isNullField(p: Any): Boolean = p match {
+    case NullField => true
+    case _         => false
   }
 }

client/src/test/scala/caliban/client/SelectionBuilderSpec.scala

@@ -35,7 +35,12 @@ object SelectionBuilderSpec extends DefaultRunnableSpec {
               Character.name ~
                 Character.nicknames ~
                 Character
-                  .role(Role.Captain.shipName, Role.Pilot.shipName, Role.Mechanic.shipName, Role.Engineer.shipName)
+                  .role(
+                    onCaptain = Role.Captain.shipName,
+                    onPilot = Role.Pilot.shipName,
+                    onMechanic = Role.Mechanic.shipName,
+                    onEngineer = Role.Engineer.shipName
+                  )
             }
           val (s, _) = SelectionBuilder.toGraphQL(query.toSelectionSet, useVariables = false)
           assert(s)(
@@ -44,6 +49,21 @@ object SelectionBuilderSpec extends DefaultRunnableSpec {
             )
           )
         },
+        test("union type with optional parameters") {
+          val query  =
+            Queries.characters() {
+              Character.name ~
+                Character.nicknames ~
+                Character
+                  .roleOption(onEngineer = Some(Role.Engineer.shipName))
+            }
+          val (s, _) = SelectionBuilder.toGraphQL(query.toSelectionSet, useVariables = false)
+          assert(s)(
+            equalTo(
+              "characters{name nicknames role{__typename ... on Engineer{shipName}}}"
+            )
+          )
+        },
         test("argument") {
           val query  =
             Queries.characters(Some(Origin.MARS)) {
@@ -154,7 +174,12 @@ object SelectionBuilderSpec extends DefaultRunnableSpec {
               (Character.name ~
                 Character.nicknames ~
                 Character
-                  .role(Role.Captain.shipName, Role.Pilot.shipName, Role.Mechanic.shipName, Role.Engineer.shipName))
+                  .role(
+                    onCaptain = Role.Captain.shipName,
+                    onPilot = Role.Pilot.shipName,
+                    onMechanic = Role.Mechanic.shipName,
+                    onEngineer = Role.Engineer.shipName
+                  ))
                 .mapN(CharacterView)
             }
 
@@ -183,6 +208,44 @@ object SelectionBuilderSpec extends DefaultRunnableSpec {
             isRight(equalTo(List(CharacterView("Amos Burton", List("Amos"), Some("Rocinante")))))
           )
         },
+        test("union type with optional parameters") {
+          case class CharacterView(name: String, nicknames: List[String], role: Option[Option[String]])
+          val query =
+            Queries.characters() {
+              (Character.name ~
+                Character.nicknames ~
+                Character
+                  .roleOption(
+                    onMechanic = Some(Role.Mechanic.shipName)
+                  ))
+                .mapN(CharacterView)
+            }
+
+          val response =
+            __ObjectValue(
+              List(
+                "characters" -> __ListValue(
+                  List(
+                    __ObjectValue(
+                      List(
+                        "name"      -> __StringValue("Amos Burton"),
+                        "nicknames" -> __ListValue(List(__StringValue("Amos"))),
+                        "role"      -> __ObjectValue(
+                          List(
+                            "__typename" -> __StringValue("Mechanic"),
+                            "shipName"   -> __StringValue("Rocinante")
+                          )
+                        )
+                      )
+                    )
+                  )
+                )
+              )
+            )
+          assert(query.fromGraphQL(response))(
+            isRight(equalTo(List(CharacterView("Amos Burton", List("Amos"), Some(Some("Rocinante"))))))
+          )
+        },
         test("aliases") {
           val query    =
             Queries

client/src/test/scala/caliban/client/TestData.scala

@@ -74,6 +74,25 @@ object TestData {
           )
         )
       )
+    def roleOption[A](
+      onCaptain: Option[SelectionBuilder[Captain, A]] = None,
+      onEngineer: Option[SelectionBuilder[Engineer, A]] = None,
+      onMechanic: Option[SelectionBuilder[Mechanic, A]] = None,
+      onPilot: Option[SelectionBuilder[Pilot, A]] = None
+    ): SelectionBuilder[Character, Option[Option[A]]]        =
+      Field(
+        "role",
+        OptionOf(
+          ChoiceOf(
+            Map(
+              "Captain"  -> onCaptain.fold[FieldBuilder[Option[A]]](NullField)(a => OptionOf(Obj(a))),
+              "Engineer" -> onEngineer.fold[FieldBuilder[Option[A]]](NullField)(a => OptionOf(Obj(a))),
+              "Mechanic" -> onMechanic.fold[FieldBuilder[Option[A]]](NullField)(a => OptionOf(Obj(a))),
+              "Pilot"    -> onPilot.fold[FieldBuilder[Option[A]]](NullField)(a => OptionOf(Obj(a)))
+            )
+          )
+        )
+      )
   }
 
   // Auto-generated query

examples/src/main/scala/caliban/client/Client.scala

@@ -59,6 +59,25 @@ object Client {
           )
         )
       )
+    def roleOption[A](
+      onCaptain: Option[SelectionBuilder[Captain, A]] = None,
+      onEngineer: Option[SelectionBuilder[Engineer, A]] = None,
+      onMechanic: Option[SelectionBuilder[Mechanic, A]] = None,
+      onPilot: Option[SelectionBuilder[Pilot, A]] = None
+    ): SelectionBuilder[Character, Option[Option[A]]] =
+      Field(
+        "role",
+        OptionOf(
+          ChoiceOf(
+            Map(
+              "Captain"  -> onCaptain.fold[FieldBuilder[Option[A]]](NullField)(a => OptionOf(Obj(a))),
+              "Engineer" -> onEngineer.fold[FieldBuilder[Option[A]]](NullField)(a => OptionOf(Obj(a))),
+              "Mechanic" -> onMechanic.fold[FieldBuilder[Option[A]]](NullField)(a => OptionOf(Obj(a))),
+              "Pilot"    -> onPilot.fold[FieldBuilder[Option[A]]](NullField)(a => OptionOf(Obj(a)))
+            )
+          )
+        )
+      )
   }
 
   type Pilot

tools/src/main/scala/caliban/tools/ClientWriter.scala

@@ -178,9 +178,19 @@ object ClientWriter {
     mappingClashedTypeNames: Map[String, String],
     genView: Boolean
   ): String = {
-    val objectName: String = safeTypeName(typedef.name, mappingClashedTypeNames)
-    val fields             = typedef.fields.map(collectFieldInfo(_, objectName, typesMap, mappingClashedTypeNames))
-    val view               =
+    val objectName: String      = safeTypeName(typedef.name, mappingClashedTypeNames)
+    val unionTypes              = typesMap.filter(_._2.isInstanceOf[UnionTypeDefinition]).keys
+    val optionalUnionTypeFields = typedef.fields.map { field =>
+      val isOptionalUnionType = unionTypes.exists(_.compareToIgnoreCase(field.ofType.toString) == 0)
+      if (isOptionalUnionType) {
+        Some(collectFieldInfo(field, objectName, typesMap, mappingClashedTypeNames, true))
+      } else {
+        None
+      }
+    }.flatten
+    val fields                  =
+      typedef.fields.map(collectFieldInfo(_, objectName, typesMap, mappingClashedTypeNames)) ::: optionalUnionTypeFields
+    val view                    =
       if (genView && typedef.fields.length <= MaxTupleLength)
         "\n  " + writeView(typedef.name, fields.map(_.typeInfo), mappingClashedTypeNames)
       else ""
@@ -400,9 +410,10 @@ object ClientWriter {
     field: FieldDefinition,
     typeName: String,
     typesMap: Map[String, TypeDefinition],
-    mappingClashedTypeNames: Map[String, String]
+    mappingClashedTypeNames: Map[String, String],
+    optionalType: Boolean = false
   ): String =
-    writeFieldInfo(collectFieldInfo(field, typeName, typesMap, mappingClashedTypeNames))
+    writeFieldInfo(collectFieldInfo(field, typeName, typesMap, mappingClashedTypeNames, optionalType))
 
   def writeFieldInfo(fieldInfo: FieldInfo): String = {
     val FieldInfo(
@@ -424,13 +435,33 @@ object ClientWriter {
     s"""$description${deprecated}def $safeName$typeParam$args$innerSelection: SelectionBuilder[$typeName, $outputType] = Field("$name", $builder$argBuilder)"""
   }
 
+  def isOptionalUnionType(
+    field: FieldDefinition,
+    typeName: String,
+    typesMap: Map[String, TypeDefinition],
+    mappingClashedTypeNames: Map[String, String],
+    optionalType: Boolean = false
+  ): Boolean = {
+    val fieldType  = safeTypeName(getTypeName(field.ofType), mappingClashedTypeNames)
+    val unionTypes = typesMap
+      .get(fieldType)
+      .collect { case UnionTypeDefinition(_, _, _, memberTypes) =>
+        memberTypes.flatMap(name => typesMap.get(safeTypeName(name, mappingClashedTypeNames)))
+      }
+      .getOrElse(Nil)
+      .collect { case o: ObjectTypeDefinition =>
+        o
+      }
+    unionTypes.nonEmpty
+  }
+
   def collectFieldInfo(
     field: FieldDefinition,
     typeName: String,
     typesMap: Map[String, TypeDefinition],
-    mappingClashedTypeNames: Map[String, String]
+    mappingClashedTypeNames: Map[String, String],
+    optionalType: Boolean = false
   ): FieldInfo = {
-    val name                                             = safeName(field.name)
     val description                                      = field.description match {
       case Some(d) if d.trim.nonEmpty => s"/**\n * ${d.trim}\n */\n"
       case _                          => ""
@@ -488,15 +519,27 @@ object ClientWriter {
           writeTypeBuilder(field.ofType, "Scalar()")
         )
       } else if (unionTypes.nonEmpty) {
-        (
-          s"[$typeLetter]",
-          s"(${unionTypes.map(t => s"""on${t.name}: SelectionBuilder[${safeTypeName(t.name, mappingClashedTypeNames)}, $typeLetter]""").mkString(", ")})",
-          writeType(field.ofType, mappingClashedTypeNames).replace(fieldType, typeLetter),
-          writeTypeBuilder(
-            field.ofType,
-            s"ChoiceOf(Map(${unionTypes.map(t => s""""${t.name}" -> Obj(on${t.name})""").mkString(", ")}))"
+        if (optionalType) {
+          (
+            s"[$typeLetter]",
+            s"(${unionTypes.map(t => s"""on${t.name}: Option[SelectionBuilder[${safeTypeName(t.name, mappingClashedTypeNames)}, $typeLetter]] = None""").mkString(", ")})",
+            s"Option[${writeType(field.ofType, mappingClashedTypeNames).replace(fieldType, typeLetter)}]",
+            writeTypeBuilder(
+              field.ofType,
+              s"ChoiceOf(Map(${unionTypes.map(t => s""""${t.name}" -> on${t.name}.fold[FieldBuilder[Option[A]]](NullField)(a => OptionOf(Obj(a)))""").mkString(", ")}))"
+            )
           )
-        )
+        } else {
+          (
+            s"[$typeLetter]",
+            s"(${unionTypes.map(t => s"""on${t.name}: SelectionBuilder[${safeTypeName(t.name, mappingClashedTypeNames)}, $typeLetter]""").mkString(", ")})",
+            writeType(field.ofType, mappingClashedTypeNames).replace(fieldType, typeLetter),
+            writeTypeBuilder(
+              field.ofType,
+              s"ChoiceOf(Map(${unionTypes.map(t => s""""${t.name}" -> Obj(on${t.name})""").mkString(", ")}))"
+            )
+          )
+        }
       } else if (interfaceTypes.nonEmpty) {
         (
           s"[$typeLetter]",
@@ -526,6 +569,7 @@ object ClientWriter {
     }
 
     val owner         = if (typeParam.nonEmpty) Some(fieldType) else None
+    val name          = if (optionalType && unionTypes.nonEmpty) safeName(field.name + "Option") else safeName(field.name)
     val fieldTypeInfo = FieldTypeInfo(
       field.name,
       name,

tools/src/test/scala/caliban/tools/ClientWriterSpec.scala

@@ -338,8 +338,22 @@ object Client {
     def role[A](
       onCaptain: SelectionBuilder[Captain, A],
       onPilot: SelectionBuilder[Pilot, A]
-    ): SelectionBuilder[Character, Option[A]] =
+    ): SelectionBuilder[Character, Option[A]]         =
       Field("role", OptionOf(ChoiceOf(Map("Captain" -> Obj(onCaptain), "Pilot" -> Obj(onPilot)))))
+    def roleOption[A](
+      onCaptain: Option[SelectionBuilder[Captain, A]] = None,
+      onPilot: Option[SelectionBuilder[Pilot, A]] = None
+    ): SelectionBuilder[Character, Option[Option[A]]] = Field(
+      "role",
+      OptionOf(
+        ChoiceOf(
+          Map(
+            "Captain" -> onCaptain.fold[FieldBuilder[Option[A]]](NullField)(a => OptionOf(Obj(a))),
+            "Pilot"   -> onPilot.fold[FieldBuilder[Option[A]]](NullField)(a => OptionOf(Obj(a)))
+          )
+        )
+      )
+    )
   }
 
 }