Bump up the API of BanyanDB Server to support the query trace. (#59)

Author: hanahmily

CHANGES.md

@@ -2,11 +2,12 @@ Changes by Version
 ==================
 Release Notes.
 
-0.7.0
+0.7.0-rc0
 ------------------
 
 ### Features
 
+* Bump up the API of BanyanDB Server to support the query trace.
 
 ### Bugs
 

src/main/java/org/apache/skywalking/banyandb/v1/client/AbstractQuery.java

@@ -62,6 +62,11 @@ public abstract class AbstractQuery<T> {
      */
     protected AbstractCriteria criteria;
 
+    /**
+     * Enable or disable trace.
+     */
+    protected boolean trace;
+
     public AbstractQuery(List<String> groups, String name, TimestampRange timestampRange, Set<String> tagProjections) {
         this.groups = groups;
         this.name = name;
@@ -100,6 +105,14 @@ public AbstractQuery<T> criteria(AbstractCriteria criteria) {
         return this;
     }
 
+    /**
+     * Enable trace for the query.
+     */
+    public AbstractQuery<T> enableTrace() {
+        this.trace = true;
+        return this;
+    }
+
     /**
      * @return QueryRequest for gRPC level query.
      * @throws BanyanDBException thrown from entity build, e.g. invalid reference to non-exist fields or tags.

src/main/java/org/apache/skywalking/banyandb/v1/client/MeasureQuery.java

@@ -200,6 +200,7 @@ BanyandbMeasure.QueryRequest build(MetadataCache.EntityMetadata entityMetadata)
         }
         // add all criteria
         buildCriteria().ifPresent(builder::setCriteria);
+        builder.setTrace(this.trace);
         return builder.build();
     }
 

src/main/java/org/apache/skywalking/banyandb/v1/client/StreamQuery.java

@@ -82,6 +82,7 @@ BanyandbStream.QueryRequest build(MetadataCache.EntityMetadata entityMetadata) t
         if (orderBy != null) {
             builder.setOrderBy(orderBy.build());
         }
+        builder.setTrace(this.trace);
         return builder.build();
     }
 }

src/main/proto/banyandb/v1/banyandb-common.proto

@@ -77,3 +77,40 @@ message Group {
   // updated_at indicates when resources of the group are updated
   google.protobuf.Timestamp updated_at = 4;
 }
+
+
+// Trace is the top level message of a trace.
+message Trace {
+  // trace_id is the unique identifier of the trace.
+  string trace_id = 1;
+  // spans is a list of spans in the trace.
+  repeated Span spans = 2;
+  // error indicates whether the trace is an error trace.
+  bool error = 3;
+}
+
+// Span is the basic unit of a trace.
+message Span {
+  // start_time is the start time of the span.
+  google.protobuf.Timestamp start_time = 1;
+  // end_time is the end time of the span.
+  google.protobuf.Timestamp end_time = 2;
+  // error indicates whether the span is an error span.
+  bool error = 3;
+  // tags is a list of tags of the span.
+  repeated Tag tags = 4;
+  // message is the message generated by the span.
+  string message = 5;
+  // children is a list of child spans of the span.
+  repeated Span children = 6;
+  // duration is the duration of the span.
+  int64 duration = 7;
+}
+
+// Tag is the key-value pair of a span.
+message Tag {
+  // key is the key of the tag.
+  string key = 1;
+  // value is the value of the tag.
+  string value = 2;
+}

src/main/proto/banyandb/v1/banyandb-measure.proto

@@ -44,9 +44,10 @@ message DataPoint {
 message QueryResponse {
   // data_points are the actual data returned
   repeated DataPoint data_points = 1;
+  // trace contains the trace information of the query when trace is enabled
+  common.v1.Trace trace = 2;
 }
 
-// QueryRequest is the request contract for query.
 // QueryRequest is the request contract for query.
 message QueryRequest {
   // groups indicate where the data points are stored.
@@ -101,6 +102,8 @@ message QueryRequest {
   uint32 limit = 11;
   // order_by is given to specify the sort for a tag.
   model.v1.QueryOrder order_by = 12;
+  // trace is used to enable trace for the query
+  bool trace = 13;
 }
 
 // TopNList contains a series of topN items

src/main/proto/banyandb/v1/banyandb-stream.proto

@@ -49,6 +49,8 @@ message Element {
 message QueryResponse {
   // elements are the actual data returned
   repeated Element elements = 1;
+  // trace contains the trace information of the query when trace is enabled
+  common.v1.Trace trace = 2;
 }
 
 // QueryRequest is the request contract for query.
@@ -72,6 +74,8 @@ message QueryRequest {
   model.v1.Criteria criteria = 7;
   // projection can be used to select the key names of the element in the response
   model.v1.TagProjection projection = 8 [(validate.rules).message.required = true];
+  // trace is used to enable trace for the query
+  bool trace = 9;
 }
 
 message ElementValue {

src/test/java/org/apache/skywalking/banyandb/v1/client/BanyanDBClientMeasureQueryTest.java

@@ -116,6 +116,7 @@ public void testQuery_tableScan() throws BanyanDBException {
                 parseProjectionList(request.getTagProjection()));
         assertCollectionEqual(Lists.newArrayList("total"),
                 request.getFieldProjection().getNamesList());
+        Assert.assertFalse(request.getTrace());
     }
 
     @Test
@@ -152,6 +153,28 @@ public void testQuery_responseConversion() {
                 (Number) resp.getDataPoints().get(0).getFieldValue("total"));
     }
 
+    @Test
+    public void testQuery_enableTrace() throws BanyanDBException {
+        ArgumentCaptor<BanyandbMeasure.QueryRequest> requestCaptor = ArgumentCaptor.forClass(BanyandbMeasure.QueryRequest.class);
+
+        Instant end = Instant.now();
+        Instant begin = end.minus(15, ChronoUnit.MINUTES);
+        MeasureQuery query = new MeasureQuery(Lists.newArrayList("sw_metric"), "service_cpm_minute",
+                new TimestampRange(begin.toEpochMilli(), end.toEpochMilli()),
+                ImmutableSet.of("entity_id"),
+                ImmutableSet.of("total"));
+        query.maxBy("total", ImmutableSet.of("entity_id"));
+        // search with conditions
+        query.and(PairQueryCondition.StringQueryCondition.eq("entity_id", "abc"));
+        query.enableTrace();
+        client.query(query);
+
+        verify(measureQueryService).query(requestCaptor.capture(), ArgumentMatchers.any());
+
+        final BanyandbMeasure.QueryRequest request = requestCaptor.getValue();
+        Assert.assertTrue(request.getTrace());
+    }
+
     static <T> void assertCollectionEqual(Collection<T> c1, Collection<T> c2) {
         Assert.assertTrue(c1.size() == c2.size() && c1.containsAll(c2) && c2.containsAll(c1));
     }

src/test/java/org/apache/skywalking/banyandb/v1/client/BanyanDBClientStreamQueryTest.java

@@ -126,6 +126,7 @@ public void testQuery_tableScan() throws BanyanDBException {
         // assert projections
         assertCollectionEqual(Lists.newArrayList("searchable:duration", "searchable:state", "searchable:start_time", "searchable:trace_id"),
                 parseProjectionList(request.getProjection()));
+        Assert.assertFalse(request.getTrace());
     }
 
     @Test
@@ -279,6 +280,7 @@ public void testQuery_indexScan() throws BanyanDBException {
         Assert.assertEquals("start_time", request.getOrderBy().getIndexRuleName());
         // assert projections
         assertCollectionEqual(Lists.newArrayList("searchable:duration", "searchable:state", "searchable:start_time", "searchable:trace_id"), parseProjectionList(request.getProjection()));
+        Assert.assertFalse(request.getTrace());
     }
 
     @Test
@@ -305,6 +307,7 @@ public void testQuery_TraceIDFetch() throws BanyanDBException {
                 "    }\n" +
                 "  }\n" +
                 "}\n", request.getCriteria().toString());
+        Assert.assertFalse(request.getTrace());
     }
 
     @Test
@@ -361,6 +364,27 @@ public void testQuery_responseConversion() throws BanyanDBException {
                 resp.getElements().get(0).getTagValue("data_binary"));
     }
 
+    @Test
+    public void testQuery_enableTrace() throws BanyanDBException {
+        ArgumentCaptor<BanyandbStream.QueryRequest> requestCaptor = ArgumentCaptor.forClass(BanyandbStream.QueryRequest.class);
+
+        Instant end = Instant.now();
+        Instant begin = end.minus(15, ChronoUnit.MINUTES);
+        StreamQuery query = new StreamQuery(Lists.newArrayList("default"), "sw",
+                new TimestampRange(begin.toEpochMilli(), end.toEpochMilli()),
+                ImmutableSet.of("state", "start_time", "duration", "trace_id"));
+        // search for all states
+        query.and(PairQueryCondition.LongQueryCondition.eq("state", 0L));
+        query.setOrderBy(new StreamQuery.OrderBy("duration", AbstractQuery.Sort.DESC));
+        query.enableTrace();
+        client.query(query);
+
+        verify(streamQueryServiceImpl).query(requestCaptor.capture(), ArgumentMatchers.any());
+
+        final BanyandbStream.QueryRequest request = requestCaptor.getValue();
+        Assert.assertTrue(request.getTrace());
+    }
+
     static <T> void assertCollectionEqual(Collection<T> c1, Collection<T> c2) {
         Assert.assertTrue(c1.size() == c2.size() && c1.containsAll(c2) && c2.containsAll(c1));
     }