Sangria Federation

EntityResolver.scala

package sangria.federation

import play.api.libs.json.{JsError, JsObject, JsSuccess, Reads}
import sangria.schema.{LeafAction, ObjectLikeType, Value}

trait EntityResolver[Ctx] {
  def typename: String

  def resolve(obj: JsObject): LeafAction[Ctx, Option[Any]]
}

object EntityResolver {

  def apply[Ctx, Arg: Reads, Val](
    resolver: Arg => LeafAction[Ctx, Option[Val]]
  )(implicit outputType: ObjectLikeType[Ctx, Val]): EntityResolver[Ctx] =
    new EntityResolver[Ctx] {

      def typename: String = outputType.name

      def resolve(obj: JsObject): LeafAction[Ctx, Option[Any]] = obj.validate[Arg] match {
        case JsSuccess(value, _) => resolver(value)
        case JsError(_)          => Value(None)
      }
    }
}

FederationSupport.scala

package sangria.federation

import play.api.libs.json.{JsObject, JsPath, Json, Reads}
import sangria.ast
import sangria.ast.ObjectValue
import sangria.renderer.SchemaFilter
import sangria.schema._
import sangria.validation.{StringCoercionViolation, ValueCoercionViolation}

import scala.reflect.ClassTag

trait FederationSupport {
  import FederationSupport._

  def federate[Ctx, Val: ClassTag](_schema: Schema[Ctx, Val], resolvers: EntityResolver[Ctx]*): Schema[Ctx, Val] = {

    val entityTypes = _schema.allTypes.values.collect {
      case obj: ObjectType[Ctx, _] @unchecked if obj.astDirectives.exists(_.name == "key") => obj
    }.toList

    val EntityType: UnionType[Ctx] =
      UnionType(
        "_Entity",
        Some(
          "Union of all the available entities. Entities are objects which have been annotated with a @key directive"
        ),
        entityTypes
      )

    val _ServiceType: ObjectType[Ctx, Unit] = ObjectType(
      "_Service",
      fields[Ctx, Unit](
        Field("sdl", StringType, resolve = _ => _schema.renderCompact(SchemaFilter.withoutGraphQLBuiltIn))
      )
    )

    val representationArgs = Argument("representations", ListInputType(AnyType))
    val allResolversMap    = resolvers.map(r => r.typename -> r).toMap

    val alwaysIncludeFields = fields[Ctx, Val](
      Field(
        "_service",
        _ServiceType,
        resolve = _ => ()
      )
    )

    val entityFields =
      if (entityTypes.nonEmpty)
        fields[Ctx, Val](
          Field(
            "_entities",
            ListType(OptionType(EntityType)),
            arguments = representationArgs :: Nil,
            resolve = ctx =>
              ctx.withArgs(representationArgs) { args =>
                Action.sequence(
                  args.map(a => allResolversMap.get(a.__typename).map(_.resolve(a.value)).getOrElse(Value(None)))
                )
            }
          )
        )
      else List.empty

    val federationFields = alwaysIncludeFields ++ entityFields

    val newQueryFields: () => List[Field[Ctx, Val]] = () => _schema.query.fieldsFn() ++ federationFields

    _schema.copy(
      query = _schema.query.copy(fieldsFn = newQueryFields),
      directives = _schema.directives ++ federationDirectives,
      additionalTypes = _schema.additionalTypes ++ Seq(_FieldSetType)
    )
  }

}

object FederationSupport {
  case object AnyCoercionViolation extends ValueCoercionViolation("Could not parse type into an _Any")

  case class _Any(__typename: String, value: JsObject)

  object _Any {
    import play.api.libs.functional.syntax._
    implicit val reads: Reads[_Any] =
      (Reads.at[String](JsPath \ "__typename") ~ Reads.of[JsObject])(_Any.apply _)
  }

  private[federation] val AnyType: ScalarType[_Any] = ScalarType[_Any](
    "_Any",
    None,
    coerceUserInput = {
      case obj: JsObject => obj.asOpt[_Any].toRight(AnyCoercionViolation)
      case v             => Left(AnyCoercionViolation)
    },
    coerceOutput = {
      case (a, _) => a.__typename
    },
    coerceInput = {
      case ObjectValue(fields, _, _) if fields.exists(_.name == "__typename") =>
        Right(_Any(fields.find(_.name == "__typename").map(_.name).get, Json.obj()))
      case v => Left(AnyCoercionViolation)

    }
  )

  case class _FieldSet(fields: String)

  private[federation] val _FieldSetType: ScalarType[_FieldSet] = ScalarType[_FieldSet](
    "_FieldSet",
    description = Some(
      """The `_FieldSet` is structurally a string but semantically it represents a selection set of fields
        |which can be used to uniquely identify an entity.""".stripMargin
    ),
    coerceOutput = (v, _) => v,
    coerceUserInput = {
      case s: String ⇒ Right(_FieldSet(s))
      case _         ⇒ Left(StringCoercionViolation)
    },
    coerceInput = {
      case ast.StringValue(s, _, _, _, _) ⇒ Right(_FieldSet(s))
      case _                              ⇒ Left(StringCoercionViolation)
    }
  )

  val Extend = ast.Directive("extends", Vector.empty)

  val External = ast.Directive("external", Vector.empty)

  val federationDirectives = List(
    Directive(
      "key",
      None,
      List(Argument("fields", _FieldSetType, "The set of fields which describe this key")),
      locations = Set(DirectiveLocation.Object, DirectiveLocation.Interface)
    ),
    Directive(
      "requires",
      None,
      List(
        Argument("fields",
                 _FieldSetType,
                 "The set of fields which will be required by the service when resolving an entity")
      ),
      locations = Set(DirectiveLocation.FieldDefinition)
    ),
    Directive(
      "provides",
      None,
      List(
        Argument("fields",
                 _FieldSetType,
                 "The set of fields which be always provided by the service in addition to the minimum key set")
      ),
      locations = Set(DirectiveLocation.FieldDefinition)
    ),
    Directive("extends", None, List.empty, locations = Set(DirectiveLocation.Object, DirectiveLocation.Interface)),
    Directive("external", None, List.empty, locations = Set(DirectiveLocation.FieldDefinition))
  )

}

package.scala

package sangria

import sangria.ast
import sangria.ast.StringValue
import sangria.schema.{Field, InterfaceType, ObjectType}

import scala.reflect.ClassTag

package object federation extends FederationSupport {

  object Key {

    def apply[Ctx, Val: ClassTag](objectType: ObjectType[Ctx, Val],
                                  key: String,
                                  others: String*): ObjectType[Ctx, Val] =
      objectType.copy(astDirectives = objectType.astDirectives ++ (key +: others).map(Key.apply).toVector)

    def apply[Ctx, Val](interfaceType: InterfaceType[Ctx, Val], key: String, others: String*): InterfaceType[Ctx, Val] =
      interfaceType.copy(astDirectives = interfaceType.astDirectives ++ (key +: others).map(Key.apply).toVector)

    def apply(fields: String): ast.Directive =
      ast.Directive("key", arguments = Vector(ast.Argument("fields", StringValue(fields))))
  }

  object Provides {

    def apply(fields: String): ast.Directive =
      ast.Directive("provides", arguments = Vector(ast.Argument("fields", StringValue(fields))))

    def apply[Ctx, Val](field: Field[Ctx, Val], key: String): Field[Ctx, Val] =
      field.copy(astDirectives = field.astDirectives :+ apply(key))

  }

  object Requires {

    def apply(fields: String): ast.Directive =
      ast.Directive("requires", arguments = Vector(ast.Argument("fields", StringValue(fields))))

    def apply[Ctx, Val](field: Field[Ctx, Val], key: String): Field[Ctx, Val] =
      field.copy(astDirectives = field.astDirectives :+ apply(key))
  }

  object External {

    def apply(): ast.Directive =
      ast.Directive("external", arguments = Vector.empty)

    def apply[Ctx, Val](field: Field[Ctx, Val]): Field[Ctx, Val] =
      field.copy(astDirectives = field.astDirectives :+ apply())
  }

  object Extends {

    def apply(): ast.Directive =
      ast.Directive("extends", Vector.empty)

    def apply[Ctx, Val: ClassTag](objectType: ObjectType[Ctx, Val]): ObjectType[Ctx, Val] =
      objectType.copy(astDirectives = objectType.astDirectives :+ apply())

    def apply[Ctx, Val](interfaceType: InterfaceType[Ctx, Val]): InterfaceType[Ctx, Val] =
      interfaceType.copy(astDirectives = interfaceType.astDirectives :+ apply())
  }
}