This is an automated email from the ASF dual-hosted git repository.

hanahmily pushed a commit to branch trace
in repository 
https://gitbox.apache.org/repos/asf/skywalking-banyandb-java-client.git

commit dd1b83a47f5f8e7343d664641014229b8eaf12ee
Author: Gao Hongtao <[email protected]>
AuthorDate: Wed Jun 26 08:32:20 2024 +0800

    Bump up the API of BanyanDB Server to support the query trace.
    
    Signed-off-by: Gao Hongtao <[email protected]>
---
 CHANGES.md                                         |  3 +-
 .../banyandb/v1/client/AbstractQuery.java          | 13 ++++++++
 .../banyandb/v1/client/MeasureQuery.java           |  1 +
 .../skywalking/banyandb/v1/client/StreamQuery.java |  1 +
 src/main/proto/banyandb/v1/banyandb-common.proto   | 37 ++++++++++++++++++++++
 src/main/proto/banyandb/v1/banyandb-measure.proto  |  5 ++-
 src/main/proto/banyandb/v1/banyandb-stream.proto   |  4 +++
 .../v1/client/BanyanDBClientMeasureQueryTest.java  | 23 ++++++++++++++
 .../v1/client/BanyanDBClientStreamQueryTest.java   | 24 ++++++++++++++
 9 files changed, 109 insertions(+), 2 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index cb63790..26de637 100644
--- a/CHANGES.md
+++ b/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
 
diff --git 
a/src/main/java/org/apache/skywalking/banyandb/v1/client/AbstractQuery.java 
b/src/main/java/org/apache/skywalking/banyandb/v1/client/AbstractQuery.java
index ed05aa1..1de95e2 100644
--- a/src/main/java/org/apache/skywalking/banyandb/v1/client/AbstractQuery.java
+++ b/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 abstract class AbstractQuery<T> {
         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.
diff --git 
a/src/main/java/org/apache/skywalking/banyandb/v1/client/MeasureQuery.java 
b/src/main/java/org/apache/skywalking/banyandb/v1/client/MeasureQuery.java
index 716949c..f3e0874 100644
--- a/src/main/java/org/apache/skywalking/banyandb/v1/client/MeasureQuery.java
+++ b/src/main/java/org/apache/skywalking/banyandb/v1/client/MeasureQuery.java
@@ -200,6 +200,7 @@ public class MeasureQuery extends 
AbstractQuery<BanyandbMeasure.QueryRequest> {
         }
         // add all criteria
         buildCriteria().ifPresent(builder::setCriteria);
+        builder.setTrace(this.trace);
         return builder.build();
     }
 
diff --git 
a/src/main/java/org/apache/skywalking/banyandb/v1/client/StreamQuery.java 
b/src/main/java/org/apache/skywalking/banyandb/v1/client/StreamQuery.java
index c553356..6ea9742 100644
--- a/src/main/java/org/apache/skywalking/banyandb/v1/client/StreamQuery.java
+++ b/src/main/java/org/apache/skywalking/banyandb/v1/client/StreamQuery.java
@@ -82,6 +82,7 @@ public class StreamQuery extends 
AbstractQuery<BanyandbStream.QueryRequest> {
         if (orderBy != null) {
             builder.setOrderBy(orderBy.build());
         }
+        builder.setTrace(this.trace);
         return builder.build();
     }
 }
diff --git a/src/main/proto/banyandb/v1/banyandb-common.proto 
b/src/main/proto/banyandb/v1/banyandb-common.proto
index f8d0f3c..3e0ac35 100644
--- a/src/main/proto/banyandb/v1/banyandb-common.proto
+++ b/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;
+}
diff --git a/src/main/proto/banyandb/v1/banyandb-measure.proto 
b/src/main/proto/banyandb/v1/banyandb-measure.proto
index fc80f88..8294081 100644
--- a/src/main/proto/banyandb/v1/banyandb-measure.proto
+++ b/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
diff --git a/src/main/proto/banyandb/v1/banyandb-stream.proto 
b/src/main/proto/banyandb/v1/banyandb-stream.proto
index cd8e0a2..0f29944 100644
--- a/src/main/proto/banyandb/v1/banyandb-stream.proto
+++ b/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 {
diff --git 
a/src/test/java/org/apache/skywalking/banyandb/v1/client/BanyanDBClientMeasureQueryTest.java
 
b/src/test/java/org/apache/skywalking/banyandb/v1/client/BanyanDBClientMeasureQueryTest.java
index c71f1a5..624c356 100644
--- 
a/src/test/java/org/apache/skywalking/banyandb/v1/client/BanyanDBClientMeasureQueryTest.java
+++ 
b/src/test/java/org/apache/skywalking/banyandb/v1/client/BanyanDBClientMeasureQueryTest.java
@@ -116,6 +116,7 @@ public class BanyanDBClientMeasureQueryTest extends 
AbstractBanyanDBClientTest {
                 parseProjectionList(request.getTagProjection()));
         assertCollectionEqual(Lists.newArrayList("total"),
                 request.getFieldProjection().getNamesList());
+        Assert.assertFalse(request.getTrace());
     }
 
     @Test
@@ -152,6 +153,28 @@ public class BanyanDBClientMeasureQueryTest extends 
AbstractBanyanDBClientTest {
                 (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));
     }
diff --git 
a/src/test/java/org/apache/skywalking/banyandb/v1/client/BanyanDBClientStreamQueryTest.java
 
b/src/test/java/org/apache/skywalking/banyandb/v1/client/BanyanDBClientStreamQueryTest.java
index 84177da..d98a618 100644
--- 
a/src/test/java/org/apache/skywalking/banyandb/v1/client/BanyanDBClientStreamQueryTest.java
+++ 
b/src/test/java/org/apache/skywalking/banyandb/v1/client/BanyanDBClientStreamQueryTest.java
@@ -126,6 +126,7 @@ public class BanyanDBClientStreamQueryTest extends 
AbstractBanyanDBClientTest {
         // 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 class BanyanDBClientStreamQueryTest extends 
AbstractBanyanDBClientTest {
         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 class BanyanDBClientStreamQueryTest extends 
AbstractBanyanDBClientTest {
                 "    }\n" +
                 "  }\n" +
                 "}\n", request.getCriteria().toString());
+        Assert.assertFalse(request.getTrace());
     }
 
     @Test
@@ -361,6 +364,27 @@ public class BanyanDBClientStreamQueryTest extends 
AbstractBanyanDBClientTest {
                 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));
     }

Reply via email to