From ed2cfa5e34b2e43997c8eda76d89aaab5568f4c2 Mon Sep 17 00:00:00 2001
From: lWarne <laurencewarne@gmail.com>
Date: Mon, 5 Jul 2021 18:20:57 +0100
Subject: [PATCH] Allow for abstracted effect type in schema gen

Allow for the use of an abstracted effect type in schema generation
and add tests for usage.
---
 .../scala/caliban/tools/SchemaWriter.scala    | 17 +++---
 .../caliban/tools/SchemaWriterSpec.scala      | 59 ++++++++++++++++++-
 2 files changed, 67 insertions(+), 9 deletions(-)

diff --git a/tools/src/main/scala/caliban/tools/SchemaWriter.scala b/tools/src/main/scala/caliban/tools/SchemaWriter.scala
index 5965bcd2cf..55459a43f4 100644
--- a/tools/src/main/scala/caliban/tools/SchemaWriter.scala
+++ b/tools/src/main/scala/caliban/tools/SchemaWriter.scala
@@ -12,7 +12,8 @@ object SchemaWriter {
     schema: Document,
     packageName: Option[String] = None,
     effect: String = "zio.UIO",
-    imports: Option[List[String]] = None
+    imports: Option[List[String]] = None,
+    isEffectTypeAbstract: Boolean = false
   )(implicit scalarMappings: ScalarMappings): String = {
     val schemaDef = schema.schemaDefinition
 
@@ -43,12 +44,12 @@ object SchemaWriter {
 
     val queries = schema
       .objectTypeDefinition(schemaDef.flatMap(_.query).getOrElse("Query"))
-      .map(t => writeRootQueryOrMutationDef(t, effect))
+      .map(t => writeRootQueryOrMutationDef(t, effect, isEffectTypeAbstract))
       .getOrElse("")
 
     val mutations = schema
       .objectTypeDefinition(schemaDef.flatMap(_.mutation).getOrElse("Mutation"))
-      .map(t => writeRootQueryOrMutationDef(t, effect))
+      .map(t => writeRootQueryOrMutationDef(t, effect, isEffectTypeAbstract))
       .getOrElse("")
 
     val subscriptions = schema
@@ -103,17 +104,19 @@ object SchemaWriter {
     s"${safeName(field.name)} :$argsTypeName $effect[${writeType(field.ofType)}]"
   }
 
-  def writeRootQueryOrMutationDef(op: ObjectTypeDefinition, effect: String)(implicit
+  def writeRootQueryOrMutationDef(op: ObjectTypeDefinition, effect: String, isEffectTypeAbstract: Boolean)(implicit
     scalarMappings: ScalarMappings
-  ): String =
+  ): String = {
+    val typeParamOrEmpty = if (isEffectTypeAbstract) s"[$effect[_]]" else ""
     s"""
-       |${writeDescription(op.description)}case class ${op.name}(
+       |${writeDescription(op.description)}case class ${op.name}$typeParamOrEmpty(
        |${op.fields.map(c => writeRootField(c, op, effect)).mkString(",\n")}
        |)""".stripMargin
 
+  }
   def writeSubscriptionField(field: FieldDefinition, od: ObjectTypeDefinition)(implicit
     scalarMappings: ScalarMappings
-  ): String =
+  ): String                                                       =
     "%s:%s ZStream[Any, Nothing, %s]".format(
       safeName(field.name),
       if (field.args.nonEmpty) s" ${argsName(field, od)} =>" else "",
diff --git a/tools/src/test/scala/caliban/tools/SchemaWriterSpec.scala b/tools/src/test/scala/caliban/tools/SchemaWriterSpec.scala
index 72e7c6a488..26abe2212c 100644
--- a/tools/src/test/scala/caliban/tools/SchemaWriterSpec.scala
+++ b/tools/src/test/scala/caliban/tools/SchemaWriterSpec.scala
@@ -76,7 +76,9 @@ object SchemaWriterSpec extends DefaultRunnableSpec {
         val result = Parser
           .parseQuery(schema)
           .map(
-            _.objectTypeDefinition("Query").map(SchemaWriter.writeRootQueryOrMutationDef(_, "zio.UIO")).mkString("\n")
+            _.objectTypeDefinition("Query")
+              .map(SchemaWriter.writeRootQueryOrMutationDef(_, "zio.UIO", false))
+              .mkString("\n")
           )
           .flatMap(Formatter.format(_, None).map(_.trim))
 
@@ -100,7 +102,7 @@ object SchemaWriterSpec extends DefaultRunnableSpec {
           .parseQuery(schema)
           .map(
             _.objectTypeDefinition("Mutation")
-              .map(SchemaWriter.writeRootQueryOrMutationDef(_, "zio.UIO"))
+              .map(SchemaWriter.writeRootQueryOrMutationDef(_, "zio.UIO", false))
               .mkString("\n")
           )
           .flatMap(Formatter.format(_, None).map(_.trim))
@@ -134,6 +136,59 @@ object SchemaWriterSpec extends DefaultRunnableSpec {
           )
         )
       },
+      testM("simple queries with abstracted effect type") {
+        val schema =
+          """
+         type Query {
+           user(id: Int): User
+           userList: [User]!
+         }
+         type User {
+           id: Int
+           name: String
+           profilePic: String
+         }"""
+
+        val result = Parser
+          .parseQuery(schema)
+          .map(
+            _.objectTypeDefinition("Query").map(SchemaWriter.writeRootQueryOrMutationDef(_, "F", true)).mkString("\n")
+          )
+          .flatMap(Formatter.format(_, None).map(_.trim))
+
+        assertM(result)(
+          equalTo(
+            """case class Query[F[_]](
+  user: QueryUserArgs => F[Option[User]],
+  userList: F[List[Option[User]]]
+)""".stripMargin
+          )
+        )
+      },
+      testM("simple mutation with abstracted effect type") {
+        val schema =
+          """
+         type Mutation {
+           setMessage(message: String): String
+         }
+         """
+        val result = Parser
+          .parseQuery(schema)
+          .map(
+            _.objectTypeDefinition("Mutation")
+              .map(SchemaWriter.writeRootQueryOrMutationDef(_, "F", true))
+              .mkString("\n")
+          )
+          .flatMap(Formatter.format(_, None).map(_.trim))
+
+        assertM(result)(
+          equalTo(
+            """case class Mutation[F[_]](
+              |  setMessage: MutationSetMessageArgs => F[Option[String]]
+              |)""".stripMargin
+          )
+        )
+      },
       testM("schema test") {
         val schema =
           """