core: add extractors for trace queries (#1613)

Add some model extractors for getting trace filters and time series queries with expected implicit conversions.
Netflix · Feb 27, 2024 · 47859a4 · 47859a4
1 parent 0d8a345
commit 47859a4
Show file tree

Hide file tree

Showing 2 changed files with 111 additions and 0 deletions.
diff --git a/atlas-core/src/main/scala/com/netflix/atlas/core/model/ModelExtractors.scala b/atlas-core/src/main/scala/com/netflix/atlas/core/model/ModelExtractors.scala
@@ -103,4 +103,32 @@ object ModelExtractors {
       case _             => None
     }
   }
+
+  case object TraceFilterType {
+
+    def unapply(value: Any): Option[TraceQuery.SpanFilter] = value match {
+      case q: TraceQuery            => Some(toFilter(q))
+      case q: Query                 => Some(toFilter(TraceQuery.Simple(q)))
+      case f: TraceQuery.SpanFilter => Some(f)
+      case _                        => None
+    }
+
+    private def toFilter(q: TraceQuery): TraceQuery.SpanFilter = {
+      TraceQuery.SpanFilter(q, Query.True)
+    }
+  }
+
+  case object TraceTimeSeriesType {
+
+    def unapply(value: Any): Option[TraceQuery.SpanTimeSeries] = value match {
+      case q: TraceQuery                => Some(toTimeSeries(q))
+      case q: Query                     => Some(toTimeSeries(TraceQuery.Simple(q)))
+      case t: TraceQuery.SpanTimeSeries => Some(t)
+      case _                            => None
+    }
+
+    private def toTimeSeries(q: TraceQuery): TraceQuery.SpanTimeSeries = {
+      TraceQuery.SpanTimeSeries(q, StyleExpr(DataExpr.Sum(Query.True), Map.empty))
+    }
+  }
 }
diff --git a/atlas-core/src/test/scala/com/netflix/atlas/core/model/ModelExtractorsSuite.scala b/atlas-core/src/test/scala/com/netflix/atlas/core/model/ModelExtractorsSuite.scala
@@ -37,4 +37,87 @@ class ModelExtractorsSuite extends FunSuite {
   completionTest("name,sps,:eq", 21)
   completionTest("name,sps,:eq,app,foo,:eq", 42)
   completionTest("name,sps,:eq,app,foo,:eq,:and,(,asg,)", 12)
+
+  test("TraceQueryType: implicit from Query") {
+    val value = Query.Equal("name", "sps")
+    value match {
+      case ModelExtractors.TraceQueryType(t) =>
+        assertEquals(t, TraceQuery.Simple(value))
+      case _ =>
+        fail(s"did not match value: $value")
+    }
+  }
+
+  test("TraceQueryType: explicit TraceQuery") {
+    val value = TraceQuery.Simple(Query.Equal("name", "sps"))
+    value match {
+      case ModelExtractors.TraceQueryType(t) =>
+        assertEquals(t, value)
+      case _ =>
+        fail(s"did not match value: $value")
+    }
+  }
+
+  test("TraceFilterType: implicit from Query") {
+    val value = Query.Equal("name", "sps")
+    value match {
+      case ModelExtractors.TraceFilterType(t) =>
+        assertEquals(t, TraceQuery.SpanFilter(TraceQuery.Simple(value), Query.True))
+      case _ =>
+        fail(s"did not match value: $value")
+    }
+  }
+
+  test("TraceFilterType: implicit from TraceQuery") {
+    val value = TraceQuery.Simple(Query.Equal("name", "sps"))
+    value match {
+      case ModelExtractors.TraceFilterType(t) =>
+        assertEquals(t, TraceQuery.SpanFilter(value, Query.True))
+      case _ =>
+        fail(s"did not match value: $value")
+    }
+  }
+
+  test("TraceFilterType: explicit TraceQuery.SpanFilter") {
+    val value = TraceQuery.SpanFilter(TraceQuery.Simple(Query.Equal("name", "sps")), Query.True)
+    value match {
+      case ModelExtractors.TraceFilterType(t) =>
+        assertEquals(t, value)
+      case _ =>
+        fail(s"did not match value: $value")
+    }
+  }
+
+  test("TraceTimeSeriesType: implicit from Query") {
+    val expr = StyleExpr(DataExpr.Sum(Query.True), Map.empty)
+    val value = Query.Equal("name", "sps")
+    value match {
+      case ModelExtractors.TraceTimeSeriesType(t) =>
+        assertEquals(t, TraceQuery.SpanTimeSeries(TraceQuery.Simple(value), expr))
+      case _ =>
+        fail(s"did not match value: $value")
+    }
+  }
+
+  test("TraceTimeSeriesType: implicit from TraceQuery") {
+    val expr = StyleExpr(DataExpr.Sum(Query.True), Map.empty)
+    val value = TraceQuery.Simple(Query.Equal("name", "sps"))
+    value match {
+      case ModelExtractors.TraceTimeSeriesType(t) =>
+        assertEquals(t, TraceQuery.SpanTimeSeries(value, expr))
+      case _ =>
+        fail(s"did not match value: $value")
+    }
+  }
+
+  test("TraceTimeSeriesType: explicit TraceQuery.SpanFilter") {
+    val expr = StyleExpr(DataExpr.Sum(Query.True), Map.empty)
+    val value = TraceQuery.SpanTimeSeries(TraceQuery.Simple(Query.Equal("name", "sps")), expr)
+    value match {
+      case ModelExtractors.TraceTimeSeriesType(t) =>
+        assertEquals(t, value)
+      case _ =>
+        fail(s"did not match value: $value")
+    }
+  }
 }