From 0309b23e70cda79940535f7fbbc8019f1cdbb649 Mon Sep 17 00:00:00 2001
From: Brian Harrington <brharrington@gmail.com>
Date: Tue, 27 Feb 2024 12:25:06 -0600
Subject: [PATCH] core: add extractors for trace queries

Add some model extractors for getting trace filters and
time series queries with expected implicit conversions.
---
 .../atlas/core/model/ModelExtractors.scala    | 28 +++++++
 .../core/model/ModelExtractorsSuite.scala     | 83 +++++++++++++++++++
 2 files changed, 111 insertions(+)

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
index 29bfd4717..33af9a4df 100644
--- 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
index 2405546a1..1b8df4efc 100644
--- 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")
+    }
+  }
 }