Merge pull request #4 from amarrella/master

Added opentracing middleware

build.sbt

@@ -22,7 +22,8 @@ libraryDependencies ++= Seq(
   "org.sangria-graphql" %% "sangria" % "1.4.1",
   "io.dropwizard.metrics" % "metrics-core" % "4.0.2",
   "org.slf4j" % "slf4j-api" % "1.7.25",
-
+  "io.opentracing.contrib" %% "opentracing-scala-concurrent" % "0.0.4",
+  "io.opentracing" % "opentracing-mock" % "0.31.0",
   "org.scalatest" %% "scalatest" % "3.0.5" % Test,
   "org.sangria-graphql" %% "sangria-json4s-native" % "1.0.0" % Test,
   "org.slf4j" % "slf4j-simple" % "1.7.25" % Test

src/main/scala/sangria/slowlog/OpenTracing.scala

@@ -0,0 +1,76 @@
+package sangria.slowlog
+
+
+import io.opentracing.{Span, Tracer}
+import sangria.execution._
+import sangria.schema.Context
+
+import scala.collection.concurrent.TrieMap
+
+final case class SpanAttachment(span: Span) extends MiddlewareAttachment
+
+class OpenTracing(implicit private val tracer: Tracer, 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 = tracer
+      .buildSpan(context.operationName.getOrElse(defaultOperationName))
+      .withTag("type", "graphql-query")
+      .start()
+    TrieMap(Vector() -> span)
+  }
+
+  def afterQuery(queryVal: QueryVal, context: MiddlewareQueryContext[Any, _, _]) =
+    queryVal.get(Vector()).foreach(_.finish())
+
+  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 span =
+      queryVal
+        .get(parentPath)
+        .map { parentSpan =>
+          tracer
+            .buildSpan(ctx.field.name)
+            .withTag("type", "graphql-field")
+            .asChildOf(parentSpan)
+        }
+        .getOrElse {
+          tracer
+            .buildSpan(ctx.field.name)
+            .withTag("type", "graphql-field")
+        }
+        .start()
+
+    BeforeFieldResult(queryVal.update(ctx.path.path, span), attachment = Some(SpanAttachment(span)))
+  }
+
+  def afterField(
+                  queryVal: QueryVal,
+                  fieldVal: FieldVal,
+                  value: Any,
+                  mctx: MiddlewareQueryContext[Any, _, _],
+                  ctx: Context[Any, _]) = {
+    queryVal.get(ctx.path.path).foreach(_.finish())
+    None
+  }
+
+  def fieldError(
+                  queryVal: QueryVal,
+                  fieldVal: FieldVal,
+                  error: Throwable,
+                  mctx: MiddlewareQueryContext[Any, _, _],
+                  ctx: Context[Any, _]) =
+    queryVal.get(ctx.path.path).foreach(_.finish())
+
+}

src/main/scala/sangria/slowlog/SlowLog.scala

@@ -1,7 +1,8 @@
 package sangria.slowlog
 
-import language.postfixOps
+import io.opentracing.Tracer
 
+import language.postfixOps
 import org.slf4j.Logger
 import sangria.ast.Document
 import sangria.execution._
@@ -95,4 +96,7 @@ object SlowLog {
     result.middlewareVals.collectFirst {case (v: QueryMetrics, _: SlowLog) ⇒ v}
 
   lazy val apolloTracing: Middleware[Any] = ApolloTracingExtension
+
+  def openTracing(implicit tracer: Tracer): Middleware[Any] =
+    new OpenTracing
 }

src/test/scala/sangria/slowlog/OpenTracingSpec.scala

@@ -0,0 +1,121 @@
+package sangria.slowlog
+
+import org.json4s.JsonAST._
+import org.scalatest.{BeforeAndAfter, Matchers, OptionValues, WordSpec}
+import sangria.execution.Executor
+import sangria.macros._
+import sangria.marshalling.ScalaInput
+import sangria.marshalling.json4s.native._
+import sangria.slowlog.util.{FutureResultSupport, StringMatchers}
+import io.opentracing.mock.{MockSpan, MockTracer}
+import io.opentracing.mock.MockTracer.Propagator
+import io.opentracing.util.ThreadLocalScopeManager
+import io.opentracing.contrib.concurrent.TracedExecutionContext
+import scala.concurrent.ExecutionContext.global
+import scala.language.postfixOps
+import scala.collection.JavaConverters._
+
+final case class SimpleMockSpan(traceId: Long, spanId: Long, parentId: Long, operationName: String)
+object SimpleMockSpan {
+  def apply(s: MockSpan): SimpleMockSpan =
+    SimpleMockSpan(s.context().traceId(), s.context().spanId(), s.parentId(), s.operationName())
+}
+
+class OpenTracingSpec extends WordSpec with Matchers with FutureResultSupport with StringMatchers with OptionValues with BeforeAndAfter  {
+  import TestSchema._
+
+
+  implicit val mockTracer = new MockTracer(new ThreadLocalScopeManager, Propagator.TEXT_MAP)
+  implicit val ec = new TracedExecutionContext(global, mockTracer, false)
+
+  before {
+    mockTracer.reset()
+  }
+
+  val mainQuery =
+    gql"""
+      query Foo {
+        friends {
+          ...Name
+          ...Name2
+        }
+      }
+
+      query Test($$limit: Int!) {
+        __typename
+        name
+             ...Name1
+             pets(limit: $$limit) {
+          ... on Cat {
+            name
+            meows
+            ...Name
+          }
+          ... on Dog {
+            ...Name1
+            ...Name1
+            foo: name
+            barks
+          }
+        }
+      }
+
+      fragment Name on Named {
+        name
+        ...Name1
+      }
+
+      fragment Name1 on Named {
+        ... on Person {
+          name
+        }
+      }
+
+      fragment Name2 on Named {
+        name
+      }
+    """
+
+  "OpenTracing middleware" should {
+    "Nest the spans correctly" in {
+
+      val vars = ScalaInput.scalaInput(Map("limit" → 4))
+
+      val scope = mockTracer.buildSpan("root").startActive(false)
+
+      val result =
+        Executor.execute(schema, mainQuery,
+          root = bob,
+          operationName = Some("Test"),
+          variables = vars,
+          middleware = SlowLog.openTracing :: Nil).await
+
+      scope.span.finish()
+
+      val finishedSpans = mockTracer.finishedSpans.asScala.map(SimpleMockSpan.apply).toSet
+      println(finishedSpans)
+      finishedSpans.forall(_.traceId == 1) shouldBe true
+      val querySpan = finishedSpans.find(_.operationName == "Test").get
+      querySpan.parentId shouldEqual 2
+      val typeNameSpan = finishedSpans.find(_.operationName == "__typename").get
+      typeNameSpan.parentId shouldEqual querySpan.spanId
+      val bobSpan = finishedSpans.filter(s => s.operationName == "name" && s.parentId == querySpan.spanId)
+      bobSpan.size shouldBe 1
+      val petsSpan = finishedSpans.filter(s => s.operationName == "pets" && s.parentId == querySpan.spanId)
+      petsSpan.size shouldBe 1
+      val petsNameSpan = finishedSpans.filter(s => s.operationName == "name" && s.parentId == petsSpan.head.spanId)
+      petsNameSpan.size shouldBe 4
+      val petsMeowsSpan = finishedSpans.filter(s => s.operationName == "meows" && s.parentId == petsSpan.head.spanId)
+      petsMeowsSpan.size shouldBe 4
+    }
+  }
+
+  def removeTime(res: JValue) =
+    res.transformField {
+      case (name @ "startOffset", _) ⇒ name → JInt(0)
+      case (name @ "duration", _) ⇒ name → JInt(0)
+      case (name @ "startTime", _) ⇒ name → JString("DATE")
+      case (name @ "endTime", _) ⇒ name → JString("DATE")
+      case (name @ "resolvers", JArray(elems)) ⇒ name → JArray(elems.sortBy(e ⇒ (e \ "path").toString))
+    }
+}