Add trace to response (#61)

Author: hanahmily

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

@@ -31,12 +31,16 @@ public class MeasureQueryResponse {
     @Getter
     private final List<DataPoint> dataPoints;
 
+    @Getter
+    private final Trace trace;
+
     MeasureQueryResponse(BanyandbMeasure.QueryResponse response) {
         final List<BanyandbMeasure.DataPoint> dataPointList = response.getDataPointsList();
         this.dataPoints = new ArrayList<>(dataPointList.size());
         for (final BanyandbMeasure.DataPoint dp : dataPointList) {
             dataPoints.add(DataPoint.create(dp));
         }
+        this.trace = Trace.convertFromProto(response.getTrace());
     }
 
     /**

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

@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.apache.skywalking.banyandb.v1.client;
+
+import com.google.protobuf.Timestamp;
+import lombok.AccessLevel;
+import lombok.Getter;
+import lombok.Setter;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * Span represents the span of a {@link Trace}.
+ */
+@Getter
+@Setter(value = AccessLevel.PRIVATE)
+public class Span {
+    private Timestamp startTime;
+    private Timestamp endTime;
+    private boolean error;
+    private List<Tag> tags;
+    private String message;
+    private List<Span> children;
+    private long duration;
+
+    static Span convertSpanFromProto(org.apache.skywalking.banyandb.common.v1.BanyandbCommon.Span protoSpan) {
+        Span spanBean = new Span();
+        spanBean.setStartTime(protoSpan.getStartTime());
+        spanBean.setEndTime(protoSpan.getEndTime());
+        spanBean.setError(protoSpan.getError());
+        spanBean.setMessage(protoSpan.getMessage());
+        spanBean.setDuration(protoSpan.getDuration());
+        spanBean.setTags(protoSpan.getTagsList().stream()
+                .map(Tag::convertTagFromProto)
+                .collect(Collectors.toList()));
+        spanBean.setChildren(protoSpan.getChildrenList().stream()
+                .map(Span::convertSpanFromProto)
+                .collect(Collectors.toList()));
+        return spanBean;
+    }
+}

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

@@ -18,23 +18,27 @@
 
 package org.apache.skywalking.banyandb.v1.client;
 
-import java.util.ArrayList;
-import java.util.List;
-
 import lombok.Getter;
 import org.apache.skywalking.banyandb.stream.v1.BanyandbStream;
 
+import java.util.ArrayList;
+import java.util.List;
+
 /**
  * StreamQueryResponse represents the stream query result.
  */
 public class StreamQueryResponse {
     @Getter
     private final List<Element> elements;
 
+    @Getter
+    private final Trace trace;
+
     StreamQueryResponse(BanyandbStream.QueryResponse response) {
         final List<BanyandbStream.Element> elementsList = response.getElementsList();
         elements = new ArrayList<>(elementsList.size());
         elementsList.forEach(element -> elements.add(Element.create(element)));
+        this.trace = Trace.convertFromProto(response.getTrace());
     }
 
     /**

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

@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.apache.skywalking.banyandb.v1.client;
+
+import lombok.AccessLevel;
+import lombok.Getter;
+import lombok.Setter;
+
+
+/**
+ * Tag represents the key-value pair of a tag in a {@link Span}.
+ */
+@Getter
+@Setter(value = AccessLevel.PRIVATE)
+public class Tag {
+    private String key;
+    private String value;
+
+    static Tag convertTagFromProto(org.apache.skywalking.banyandb.common.v1.BanyandbCommon.Tag protoTag) {
+        Tag tagBean = new Tag();
+        tagBean.setKey(protoTag.getKey());
+        tagBean.setValue(protoTag.getValue());
+        return tagBean;
+    }
+}
+

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

@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package org.apache.skywalking.banyandb.v1.client;
+
+import lombok.AccessLevel;
+import lombok.Getter;
+import lombok.Setter;
+
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * Trace represents the trace of a request.
+ */
+@Getter
+@Setter(value = AccessLevel.PRIVATE)
+public class Trace {
+    private String traceId;
+    private List<Span> spans;
+    private boolean error;
+
+    static Trace convertFromProto(org.apache.skywalking.banyandb.common.v1.BanyandbCommon.Trace protoTrace) {
+        Trace traceBean = new Trace();
+        traceBean.setTraceId(protoTrace.getTraceId());
+        traceBean.setError(protoTrace.getError());
+        traceBean.setSpans(protoTrace.getSpansList().stream()
+                .map(Span::convertSpanFromProto)
+                .collect(Collectors.toList()));
+        return traceBean;
+    }
+}

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

@@ -22,6 +22,7 @@
 import com.google.common.collect.Lists;
 import com.google.protobuf.Timestamp;
 import io.grpc.stub.StreamObserver;
+import org.apache.skywalking.banyandb.common.v1.BanyandbCommon;
 import org.apache.skywalking.banyandb.measure.v1.BanyandbMeasure;
 import org.apache.skywalking.banyandb.measure.v1.MeasureServiceGrpc;
 import org.apache.skywalking.banyandb.model.v1.BanyandbModel;
@@ -124,6 +125,10 @@ public void testQuery_responseConversion() {
         final String entityIDValue = "entity_id_a";
         final Instant now = Instant.now();
         final BanyandbMeasure.QueryResponse responseObj = BanyandbMeasure.QueryResponse.newBuilder()
+                .setTrace(BanyandbCommon.Trace.newBuilder().addSpans(BanyandbCommon.Span.newBuilder()
+                        .setMessage("test")
+                                .addTags(BanyandbCommon.Tag.newBuilder().setKey("b").setValue("a").build())
+                        .build()).build())
                 .addDataPoints(BanyandbMeasure.DataPoint.newBuilder()
                         .setTimestamp(Timestamp.newBuilder()
                                 .setSeconds(now.toEpochMilli() / 1000)
@@ -151,6 +156,11 @@ public void testQuery_responseConversion() {
         Assert.assertEquals(entityIDValue, resp.getDataPoints().get(0).getTagValue("entity_id"));
         Assert.assertEquals(10L,
                 (Number) resp.getDataPoints().get(0).getFieldValue("total"));
+        Assert.assertEquals(1, resp.getTrace().getSpans().size());
+        Assert.assertEquals("test", resp.getTrace().getSpans().get(0).getMessage());
+        Assert.assertEquals(1, resp.getTrace().getSpans().get(0).getTags().size());
+        Assert.assertEquals("b", resp.getTrace().getSpans().get(0).getTags().get(0).getKey());
+        Assert.assertEquals("a", resp.getTrace().getSpans().get(0).getTags().get(0).getValue());
     }
 
     @Test

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

@@ -24,6 +24,7 @@
 import com.google.protobuf.NullValue;
 import com.google.protobuf.Timestamp;
 import io.grpc.stub.StreamObserver;
+import org.apache.skywalking.banyandb.common.v1.BanyandbCommon;
 import org.apache.skywalking.banyandb.model.v1.BanyandbModel;
 import org.apache.skywalking.banyandb.stream.v1.BanyandbStream;
 import org.apache.skywalking.banyandb.stream.v1.StreamServiceGrpc;
@@ -318,6 +319,10 @@ public void testQuery_responseConversion() throws BanyanDBException {
         final long duration = 200L;
         final Instant now = Instant.now();
         final BanyandbStream.QueryResponse responseObj = BanyandbStream.QueryResponse.newBuilder()
+                .setTrace(BanyandbCommon.Trace.newBuilder().addSpans(BanyandbCommon.Span.newBuilder()
+                        .setMessage("test")
+                        .addTags(BanyandbCommon.Tag.newBuilder().setKey("b").setValue("a").build())
+                        .build()).build())
                 .addElements(BanyandbStream.Element.newBuilder()
                         .setElementId(elementId)
                         .setTimestamp(Timestamp.newBuilder()
@@ -362,6 +367,11 @@ public void testQuery_responseConversion() throws BanyanDBException {
         Assert.assertNull(resp.getElements().get(0).getTagValue("mq.broker"));
         Assert.assertArrayEquals(binaryData,
                 resp.getElements().get(0).getTagValue("data_binary"));
+        Assert.assertEquals(1, resp.getTrace().getSpans().size());
+        Assert.assertEquals("test", resp.getTrace().getSpans().get(0).getMessage());
+        Assert.assertEquals(1, resp.getTrace().getSpans().get(0).getTags().size());
+        Assert.assertEquals("b", resp.getTrace().getSpans().get(0).getTags().get(0).getKey());
+        Assert.assertEquals("a", resp.getTrace().getSpans().get(0).getTags().get(0).getValue());
     }
 
     @Test