Quick and dirty graphql to Scala case class generator.

GraphQL2Scala.scala

/**
Copyright 2017 Kubit Ltd

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
import java.util.concurrent.atomic.AtomicReference

import com.fortytwodata.huluvu.server.ask.NativeCodeGenerator._
import com.fortytwodata.huluvu.server.ask.TraversalUtils._
import sangria.schema
import sangria.ast._
import sangria.schema.{BuiltinScalars, ObjectLikeType}

import scala.collection.concurrent.TrieMap
import scala.collection.mutable.ListBuffer

case class DocCtx(
  document: AtomicReference[Document],
  types: Map[String, TypeDefinition],
  directives: Map[String, DirectiveDefinition],
  scalars: Map[String, ScalarTypeDefinition],
  native: ScalaDefs = TrieMap()
) {

  def typeOrFail(listType: String): TypeDefinition = types.get(listType) match {
    case Some(tpe) => tpe
    case _ => throw new IllegalArgumentException(s"unknown type $listType")
  }

  def doc: Document = document.get
}

import scala.collection.concurrent.TrieMap

object TraversalUtils {
  val builtinScalars: Map[String, ScalarTypeDefinition] = BuiltinScalars.map(s => s.name -> ScalarTypeDefinition(s.name, s.astDirectives)).toMap
}

trait TraversalUtils {
  def createCtx(doc: Document): DocCtx = {
    val all = doc.definitions.collect { case f: TypeDefinition => f.name -> f }.toMap
    val dirs = doc.definitions.collect { case o: DirectiveDefinition => o.name -> o }.toMap
    val scalars = doc.definitions.collect { case o: ScalarTypeDefinition => o.name -> o }.toMap
    DocCtx(new AtomicReference(doc), all ++ builtinScalars ++ scalars, dirs, scalars)
  }

  def isNumeric(t: Type): Boolean =
    "Int".contentEquals(t.namedType.name) ||
      "Long".contentEquals(t.namedType.name) ||
      "Float".contentEquals(t.namedType.name) ||
      "Double".contentEquals(t.namedType.name) ||
      "BigDecimal".contentEquals(t.namedType.name)

  def isBool(t: Type): Boolean = "Boolean".contentEquals(t.namedType.name)

  def isString(t: Type): Boolean =
    "String".contentEquals(t.namedType.name) ||
      "ID".contentEquals(t.namedType.name)

  def isDateTime(t: Type): Boolean =
  //"Date".contentEquals(f.fieldType.namedType.name) ||
    "DateTime".contentEquals(t.namedType.name)

  def isList(t: Type): Boolean = t match {
    case _: ListType => true
    case NotNullType(ListType(_, _), _) => true //todo this ignores deeply nested list types
    case _ => false
  }

  def isJson(t: Type): Boolean = "map".contentEquals(t.namedType.name.toLowerCase)

  def isEnum(t: Type, ctx: DocCtx): Boolean = typeOf(t.namedType.name, ctx) match {
    case _: EnumTypeDefinition => true
    case _ => false
  }

  def isPrimitive(t: Type, ctx: Option[DocCtx] = None): Boolean =
    isNumeric(t) ||
      isBool(t) ||
      isString(t) ||
      isDateTime(t) ||
      isJson(t) ||
      ctx.exists(isEnum(t, _))

  /*
  def isTable(tpe: WithDirectives, ctx: DocCtx): Boolean = {
    val interfaceInline = tpe match {
      case o: ObjectTypeDefinition => o.interfaces.map(i => typeOf(i.namedType.name, ctx)).exists {
        case i: WithDirectives => isTable(i, ctx)
        case _ => false
      }
      case _ => false
    }
    interfaceInline || tpe.directives.exists(_.name.contentEquals("table"))
  }
  */
  //https://apacheignite-sql.readme.io/docs/data-types
  def primitiveTypeOf(typeName: String, ctx: DocCtx): Option[String] = {
    val name = typeOf(typeName, ctx) match {
      case _: EnumTypeDefinition => "enum"
      case s: ScalarTypeDefinition if "json".contentEquals(s.name.toLowerCase) => "json"
      case _ => typeName.toLowerCase
    }
    Option(name match {
      case "int" => "BIGINT"
      case "long" => "BIGINT"
      case "float" => "REAL"
      case "double" => "DOUBLE"
      case "boolean" => "BOOLEAN"
      case "datetime" => "TIMESTAMP"
      case "string" | "enum" | "id" | "datetime" | "json" => "VARCHAR"
      case "geo" => "GEOMETRY"
      case _ =>
        null //throw UnsupportedPrimitive(s"$name is an unknown/unsupported field type")
    })
  }

  private def hasPK(field: FieldDefinition): Boolean = field.directives.exists(_.name.contentEquals("pk"))

  def isPK(field: FieldDefinition, fields: Vector[FieldDefinition]): Boolean = {
    if (hasPK(field)) return true
    if (fields.exists(hasPK)) return false //if any other field is a PK then ID fields (below condition) are not automatically primary keys
    "ID".contentEquals(field.fieldType.namedType.name)
  }

  def emulateField(name: String, tpe: String, args: Vector[InputValueDefinition] = Vector.empty): FieldDefinition = FieldDefinition(name, NamedType(tpe), args)

  def emulateField(typeName: String, f: Vector[FieldDefinition], ctx: DocCtx): Vector[FieldDefinition] = {
    val (recursiveFields, others) = f.partition(_.fieldType.namedType.name.contentEquals(typeName))
    val pks =
      if (recursiveFields.nonEmpty) {
        val pks = pksFrom(f, ctx)
        pks.flatMap { case (fieldName, (_, _, _, tpeFields)) =>
          tpeFields.map { tpeField =>
            //val name = colName(typeName.toLowerCase, fieldName.titleCase, tpeField.name)
            val name = typeName.lowerCaseFirst + fieldName.titleCase + tpeField.name.titleCase
            FieldDefinition(name, NamedType(tpeField.fieldType.namedType.name), Vector.empty)
          }
        }
      } else {
        Vector.empty
      }
    others ++ pks
  }

  def typeOf(field: FieldDefinition, ctx: DocCtx): (FieldDefinition, AstNode) = (field, typeOf(field.fieldType.namedType.name, ctx))

  def typeOf(name: String, ctx: DocCtx): AstNode = {
    if (ctx.types.contains(name)) ctx.types(name)
    else if (ctx.directives.contains(name)) ctx.directives(name)
    else if (ctx.scalars.contains(name)) ctx.scalars(name)
    else NamedType(name)
  }

  def pksFrom(fields: Vector[FieldDefinition], ctx: DocCtx): Map[String, (FieldDefinition, TypeDefinition, Boolean, Vector[FieldDefinition])] = {
    fields.filter(isPK(_, fields)).map { f =>
      val tpe = ctx.typeOrFail(f.fieldType.namedType.name)
      val tpeFields: (Boolean, Vector[FieldDefinition]) = tpe match {
        case f: InterfaceTypeDefinition => (false, f.fields)
        case o: ObjectTypeDefinition => (false, o.fields ++ o.interfaces.map(iName => ctx.typeOrFail(iName.name)).flatMap {
          case i: InterfaceTypeDefinition => i.fields
          case x => throw new IllegalStateException("The impossible happened. Object appears to implement a non-interface type in :\n" + x.renderCompact)
        })
        case _: ScalarTypeDefinition => (true, Vector(f))
        case _ => (false, Vector.empty)
      }
      f.name -> (f, tpe, tpeFields._1, tpeFields._2)
    }.toMap
  }

  def nFields(o: NamedType, ctx: DocCtx): Vector[FieldDefinition] = ctx.types.get(o.name) match {
    case Some(o: ObjectTypeDefinition) => oFields(o, ctx)
    case Some(o: InterfaceTypeDefinition) => iFields(o, ctx)
    case _ => Vector.empty
  }

  def oFields(o: ObjectLikeType[_, _]): Vector[schema.Field[_, _]] =
    o.ownFields ++ o.interfaces.flatMap(_.fields).map(f => f.name -> f).filterNot(p => o.ownFields.exists(_.name == p._1)).map(_._2)

  def oFields(o: ObjectTypeDefinition, ctx: DocCtx, breakRecursion: Boolean = true): Vector[FieldDefinition] = {
    val fields = TrieMap[String, FieldDefinition]()
    o.interfaces.map(i => typeOf(i.namedType.name, ctx)).collect {
      case i: InterfaceTypeDefinition => i.fields
      case x => throw new UnsupportedOperationException(x.renderPretty)
    }
      .flatten
      .foreach(f => fields += f.name -> f)
    o.fields.foreach(f => fields += f.name -> f) //object fields added last so they replace inherited fields
    deDupeFields(if (breakRecursion) emulateField(o.name, fields.values.toVector, ctx) else fields.values.toVector, ctx)
  }

  /**
    * fields defined in interfaces and on types show up twice.
    * we favour the field that has the @pk definition
    * if no pk definition then either field is taken but not both
    */
  def deDupeFields(fields: Vector[FieldDefinition], ctx: DocCtx): Vector[FieldDefinition] = {
    return fields.distinct
    fields
      .groupBy(_.name)
      .map { case (_, all) =>
        all.find(isPK(_, fields)).getOrElse {
          all.filter {
            case f if !ctx.typeOrFail(f.fieldType.namedType.name).isInstanceOf[InterfaceTypeDefinition] => true
            case _ => false
          }.head //lol you can tell I can't be bothered anymore!
        }
      }.toVector
  }

  def iFields(tpe: InterfaceTypeDefinition, ctx: DocCtx, breakRecursion: Boolean = true): Vector[FieldDefinition] =
    deDupeFields(if (breakRecursion) emulateField(tpe.name, tpe.fields, ctx) else tpe.fields, ctx) //if GQL adds support for interfaces extending other interfaces

}

object NativeCodeGenerator extends TraversalUtils {
  type PropsDef = AtomicReference[Map[String, TrieMap[String, Any]]]
  type TableDefs = TrieMap[String, PropsDef]
  type DefsTpe = TrieMap[String, AtomicReference[String]]
  type ScalaDefs = TrieMap[String, String]
  type TypeDef = Either[TypeDefinition, ObjectLikeType[Any, _]]

  implicit class TitleCase(s: String) {
    def titleCase: String = s.charAt(0).toUpper + s.substring(1)

    def lowerCaseFirst: String = s.charAt(0).toLower + s.substring(1)
  }

  private val gql2Primitives = Map(
    "Int" -> "int",
    "Long" -> "long",
    "Boolean" -> "boolean",
    "ID" -> "UUID",
    "Json" -> "JsonAST.JValue"
  )

  def getType(tpe: Type, ctx: DocCtx, interface: Boolean): String = {
    if (gql2Primitives.contains(tpe.namedType.name)) return gql2Primitives(tpe.namedType.name)
    tpe match {
      case nnt: NotNullType => getType(nnt.ofType, ctx, interface)
      case arr: ListType =>
        val b = getType(arr.ofType, ctx, interface)
        s"List<$b>"
      case nt: NamedType => nt.name
      case o => o.namedType.name
    }
  }

  def generateNative(o: InterfaceTypeDefinition, ctx: DocCtx): Unit = {
    val fields = genAttrs(iFields(o, ctx), ctx, interface = true).foldLeft("") { case (prev, (_, get, set)) => prev + "\n" + get + "\n" + set }.trim
    val code =
      s"""
         |public interface ${o.name} {
         |  $fields
         |}
    """.stripMargin
    //generateClass(ctx, iFields(o, ctx), s"${o.name}", Vector(o.name))
    ctx.native += o.name -> code
  }

  private def genAttrs(fields: Vector[FieldDefinition], ctx: DocCtx, interface: Boolean): Vector[(String, String, String)] = {
    fields.flatMap { field =>
      val tpe = getType(field.fieldType, ctx, interface)
      val name = field.name
      ctx.types(field.fieldType.namedType.name) match {
        case _: ObjectTypeDefinition | _: InterfaceTypeDefinition
          //if we're generating an interface, don't output the list fields because list is invariant on T sub types won't work
          if interface && isList(field.fieldType) =>
          None
        case _ =>
          Option((
            s"public $tpe _$name;",
            //todo - add support for getters and setters
            "", //s"public $tpe get${name.titleCase}()${if (!interface) s"{ return _$name;}" else ";"}",
            "" //s"public void set${name.titleCase}($tpe value)${if (!interface) s"{ _$name = value;}" else ";"}"
          ))
      }
    }
  }

  def generateNative(o: ObjectTypeDefinition, ctx: DocCtx): Unit =
    generateClass(ctx, oFields(o, ctx), o.name, o.interfaces.map(_.name))


  private def generateClass(ctx: DocCtx, all: Vector[FieldDefinition], name: String, interfaces: Vector[String]) = {
    val (fields, methods) = genAttrs(all, ctx, interface = false)
      .foldLeft(("", "")) { case ((prevFields, prevMethods), (field, get, set)) => (prevFields + "\n" + field, prevMethods + "\n" + get + "\n" + set) }
    val body = (fields.trim + "\n" + methods.trim).trim

    val superTypes = interfaces.zipWithIndex.map {
      case (interface, idx) if idx == 0 => s" implements $interface"
      case (interface, _) => s", $interface"
    }.mkString
    val code =
      s"""
         |public class $name $superTypes {
         |  $body
         |}
    """.stripMargin
    ctx.native += name -> code
  }

  def generateNative(o: UnionTypeDefinition, ctx: DocCtx): Unit = {
    val code =
      s"""
         |class ${o.name}{
         |  ${o.types.zipWithIndex.map { case (t, i) => s"public ${t.name} _$i;" }.mkString("\n")}
         |  public Object get(){
         |    ${o.types.zipWithIndex.map { case (_, i) => s"if(_$i != null)return _$i;" }.mkString("\n")}
         |    return null;
         |  }
         |  public void set(Object value){
         |    ${o.types.zipWithIndex.map { case (t, i) => s"if(value instanceof ${t.name}){_$i = (${t.name})value; return;}" }.mkString("\n")}
         |    throw new IllegalArgumentException(String.format("${o.name} can have values of types ${o.types.map(_.name).mkString(",")} but was given %s", value == null ?"null":value.getClass().getSimpleName()));
         |  }
         |}
    """.stripMargin
    ctx.native += o.name -> code
  }

  def generateNative(o: EnumTypeDefinition, ctx: DocCtx): Unit = {
    val types = o.values.map(_.name).mkString(",")
    val code =
      s"""
         |enum ${o.name} {$types}
    """.stripMargin
    ctx.native += o.name -> code
  }

  def generateNative(ctx: DocCtx): String = {
    val enums = ListBuffer.empty[String]
    val interfaces = ListBuffer.empty[String]
    ctx.doc.definitions.foreach {
      case o: InterfaceTypeDefinition =>
        interfaces += o.name
        generateNative(o, ctx)
      case o: ObjectTypeDefinition => generateNative(o, ctx)
      case o: UnionTypeDefinition => generateNative(o, ctx)
      case o: EnumTypeDefinition =>
        enums += o.name
        generateNative(o, ctx)
      case _: DirectiveDefinition => //todo is there anything to be done in these 3 cases
      case _: ScalarTypeDefinition =>
      case _: InputObjectTypeDefinition =>
      case _ => //ignore everything else???
    }
    s"""
       |import java.util.UUID
       |import org.joda.time.DateTime
       |import org.json4s.JsonAST;
       |import java.util.List
       |
       |${ctx.native.values.map(_.trim).mkString("\n")}
    """.stripMargin
  }

  def genNative(doc: Document): (DocCtx, String) = {
    val ctx = createCtx(doc)
    (ctx, generateNative(ctx))
  }
}