From 9f0d378aa6c5433627a916ae53538397dedac419 Mon Sep 17 00:00:00 2001
From: Ajeet D'Souza <98ajeet@gmail.com>
Date: Tue, 25 May 2021 15:36:09 +0530
Subject: [PATCH] chore(tracing): Add ns attribute to doQuery (#7850)

(cherry picked from commit b80efb69e0b61132aec9dbebbb00177324faf774)
---
 edgraph/server.go | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/edgraph/server.go b/edgraph/server.go
index 2a46b772813..47679eda4d0 100644
--- a/edgraph/server.go
+++ b/edgraph/server.go
@@ -527,8 +527,12 @@ func (s *Server) Alter(ctx context.Context, op *api.Operation) (*api.Payload, er
 	return empty, nil
 }
 
+func annotateNamespace(span *otrace.Span, ns uint64) {
+	span.AddAttributes(otrace.Int64Attribute("ns", int64(ns)))
+}
+
 func annotateStartTs(span *otrace.Span, ts uint64) {
-	span.Annotate([]otrace.Attribute{otrace.Int64Attribute("startTs", int64(ts))}, "")
+	span.AddAttributes(otrace.Int64Attribute("startTs", int64(ts)))
 }
 
 func (s *Server) doMutate(ctx context.Context, qc *queryContext, resp *api.Response) error {
@@ -1175,6 +1179,10 @@ func (s *Server) doQuery(ctx context.Context, req *Request) (resp *api.Response,
 
 	var measurements []ostats.Measurement
 	ctx, span := otrace.StartSpan(ctx, methodRequest)
+	if ns, err := x.ExtractNamespace(ctx); err == nil {
+		annotateNamespace(span, ns)
+	}
+
 	ctx = x.WithMethod(ctx, methodRequest)
 	defer func() {
 		span.End()
@@ -1540,6 +1548,9 @@ func (s *Server) CommitOrAbort(ctx context.Context, tc *api.TxnContext) (*api.Tx
 		return &api.TxnContext{}, errors.Errorf(
 			"StartTs cannot be zero while committing a transaction")
 	}
+	if ns, err := x.ExtractJWTNamespace(ctx); err == nil {
+		annotateNamespace(span, ns)
+	}
 	annotateStartTs(span, tc.StartTs)
 
 	if err := validateNamespace(ctx, tc); err != nil {