From f1727b1d5e7b2d10e76502dce5a93e3bb87fadba Mon Sep 17 00:00:00 2001
From: kezhenxu94 <kezhenxu94@apache.org>
Date: Fri, 16 May 2025 12:16:33 +0800
Subject: [PATCH] feat: add support to query trace from cold-stage storage

---
 assets/graphqls/trace/ColdTrace.graphql | 66 +++++++++++++++++++++++++
 internal/commands/trace/list.go         | 22 +++++++++
 pkg/graphql/trace/trace.go              | 11 +++++
 3 files changed, 99 insertions(+)
 create mode 100644 assets/graphqls/trace/ColdTrace.graphql

diff --git a/assets/graphqls/trace/ColdTrace.graphql b/assets/graphqls/trace/ColdTrace.graphql
new file mode 100644
index 00000000..eb53ef7c
--- /dev/null
+++ b/assets/graphqls/trace/ColdTrace.graphql
@@ -0,0 +1,66 @@
+# Licensed to 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. Apache Software Foundation (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.
+
+query ($traceId: ID!, $duration: Duration!) {
+    result: queryTraceFromColdStage(traceId: $traceId, duration: $duration) {
+        spans {
+            traceId
+            segmentId
+            spanId
+            parentSpanId
+            refs {
+                traceId
+                parentSegmentId
+                parentSpanId
+                type
+            }
+            serviceCode
+            serviceInstanceName
+            startTime
+            endTime
+            endpointName
+            type
+            peer
+            component
+            isError
+            layer
+            tags {
+                key value
+            }
+            logs {
+                time data {
+                    key value
+                }
+            }
+            attachedEvents {
+                startTime {
+                    seconds nanos
+                }
+                event
+                endTime {
+                    seconds nanos
+                }
+                tags {
+                    key value
+                }
+                summary {
+                    key value
+                }
+            }
+        }
+    }
+}
diff --git a/internal/commands/trace/list.go b/internal/commands/trace/list.go
index 6160bb2a..ec4a4ece 100644
--- a/internal/commands/trace/list.go
+++ b/internal/commands/trace/list.go
@@ -53,6 +53,9 @@ $ swctl trace ls --service-name "business-zone::projectB" --endpoint-name "/proj
 
 3. Query the monitored trace of id "321661b1-9a31-4e12-ad64-c8f6711f108d":
 $ swctl trace ls --trace-id "321661b1-9a31-4e12-ad64-c8f6711f108d"
+
+4. Query the monitored trace of id "321661b1-9a31-4e12-ad64-c8f6711f108d" from cold-stage storage:
+$ swctl trace ls --trace-id "321661b1-9a31-4e12-ad64-c8f6711f108d" --cold
 `,
 	Flags: flags.Flags(
 		flags.DurationFlags,
@@ -74,6 +77,11 @@ $ swctl trace ls --trace-id "321661b1-9a31-4e12-ad64-c8f6711f108d"
 				Usage: "`order` of the returned traces, can be `duration` or `startTime`",
 				Value: "duration",
 			},
+			&cli.BoolFlag{
+				Name:  "cold",
+				Usage: "query trace from cold-stage storage, must be used with trace-id",
+				Value: false,
+			},
 		},
 	),
 	Before: interceptor.BeforeChain(
@@ -96,6 +104,20 @@ $ swctl trace ls --trace-id "321661b1-9a31-4e12-ad64-c8f6711f108d"
 		serviceInstanceID := ctx.String("instance-id")
 		traceID := ctx.String("trace-id")
 		tagStr := ctx.String("tags")
+		cold := ctx.Bool("cold")
+
+		if cold && traceID == "" {
+			return fmt.Errorf("cold storage must be queried with trace-id")
+		}
+
+		if cold {
+			trace, err := trace.ColdTrace(ctx.Context, duration, traceID)
+			if err != nil {
+				return err
+			}
+			return display.Display(ctx.Context, &displayable.Displayable{Data: trace})
+		}
+
 		var tags []*api.SpanTag = nil
 		if tagStr != "" {
 			tagArr := strings.SplitSeq(tagStr, ",")
diff --git a/pkg/graphql/trace/trace.go b/pkg/graphql/trace/trace.go
index 1a440603..fa6669e4 100644
--- a/pkg/graphql/trace/trace.go
+++ b/pkg/graphql/trace/trace.go
@@ -39,6 +39,17 @@ func Trace(ctx context.Context, traceID string) (api.Trace, error) {
 	return response["result"], err
 }
 
+func ColdTrace(ctx context.Context, duration api.Duration, traceID string) (api.Trace, error) {
+	var response map[string]api.Trace
+
+	request := graphql.NewRequest(assets.Read("graphqls/trace/ColdTrace.graphql"))
+	request.Var("traceId", traceID)
+	request.Var("duration", duration)
+	err := client.ExecuteQuery(ctx, request, &response)
+
+	return response["result"], err
+}
+
 func Traces(ctx context.Context, condition *api.TraceQueryCondition) (api.TraceBrief, error) {
 	var response map[string]api.TraceBrief