OpenCensus support

[siderakis]

OpenCensus support would be nice. :-)

[siderakis]

Here is a simple middleware using the Open Census Scala wrapper.

"com.github.sebruck" %% "opencensus-scala-core" % "0.6.1"

package tracing

import io.opencensus.trace.{EndSpanOptions, Span, Status}
import sangria.execution._
import sangria.schema.Context

import scala.collection.concurrent.TrieMap

class TracingMiddleware(parentSpan: Option[Span] = None, defaultOperationName: String = "UNNAMED")
  extends Middleware[Any] with MiddlewareAfterField[Any] with MiddlewareErrorField[Any] {
  type QueryVal = TrieMap[Vector[Any], Span]
  type FieldVal = Unit

  def beforeQuery(context: MiddlewareQueryContext[Any, _, _]) = {
    val span =
      parentSpan match {
        case Some(parent) ⇒ Tracing.startSpanWithParent(context.operationName.getOrElse(defaultOperationName), parent)
        case None ⇒ Tracing.startSpan(context.operationName.getOrElse(defaultOperationName))
      }

    TrieMap(Vector.empty → span)
  }

  def afterQuery(queryVal: QueryVal, context: MiddlewareQueryContext[Any, _, _]) =
    queryVal.get(Vector.empty).foreach(resp => {
      resp.end()
    })

  def beforeField(queryVal: QueryVal, mctx: MiddlewareQueryContext[Any, _, _], ctx: Context[Any, _]) = {
    val path = ctx.path.path
    val parentPath = path
      .dropRight(1)
      .reverse
      .dropWhile {
        case _: String ⇒ false
        case _: Int ⇒ true
      }
      .reverse

    val scope =
      queryVal
        .get(parentPath)
        .map { parentScope ⇒
          Tracing.startSpanWithParent(ctx.field.name, parentScope)
        }
        .getOrElse {
          Tracing.startSpan(ctx.field.name)
        }

    BeforeFieldResult(queryVal.update(ctx.path.path, scope), attachment = Some(ScopeAttachment(scope)))
  }

  def afterField(
                  queryVal: QueryVal,
                  fieldVal: FieldVal,
                  value: Any,
                  mctx: MiddlewareQueryContext[Any, _, _],
                  ctx: Context[Any, _]) = {
    queryVal.get(ctx.path.path).foreach(scope => {
      scope.end()
    })
    None
  }

  def fieldError(
                  queryVal: QueryVal,
                  fieldVal: FieldVal,
                  error: Throwable,
                  mctx: MiddlewareQueryContext[Any, _, _],
                  ctx: Context[Any, _]) =
    queryVal.get(ctx.path.path).foreach(scope => {
      scope.end(EndSpanOptions.builder().setStatus(Status.UNKNOWN).build())
    })
}

final case class ScopeAttachment(scope: Span) extends MiddlewareAttachment {}