From 43ff4abffdda962431aac28abbc8d3cea46763ca Mon Sep 17 00:00:00 2001
From: Gavin Cabbage <dev@gavincabbage.com>
Date: Fri, 6 Dec 2019 13:24:43 -0500
Subject: [PATCH] feat(query): add trace response headers

---
 CHANGELOG.md               |  1 +
 http/query_handler.go      |  7 ++++++-
 http/query_handler_test.go | 16 ++++++++++++++++
 http/swagger.yml           |  5 +++++
 kit/tracing/tracing.go     | 17 +++++++++++++++++
 5 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e47773f1d7..fab233e1b5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@
 1. [15836](https://github.com/influxdata/influxdb/pull/16077): Add stacked line layer option to graphs
 1. [16094](https://github.com/influxdata/influxdb/pull/16094): Annotate log messages with trace ID, if available
 1. [16187](https://github.com/influxdata/influxdb/pull/16187): Bucket create to accept an org name flag
+1. [16158](https://github.com/influxdata/influxdb/pull/16158): Add trace ID response header to query endpoint
 
 ### Bug Fixes
 
diff --git a/http/query_handler.go b/http/query_handler.go
index d8e1c0da0b..318e15154f 100644
--- a/http/query_handler.go
+++ b/http/query_handler.go
@@ -24,6 +24,7 @@ import (
 	"github.com/influxdata/influxdb/http/metric"
 	"github.com/influxdata/influxdb/kit/check"
 	"github.com/influxdata/influxdb/kit/tracing"
+	"github.com/influxdata/influxdb/logger"
 	influxlogger "github.com/influxdata/influxdb/logger"
 	"github.com/influxdata/influxdb/query"
 	"github.com/pkg/errors"
@@ -32,7 +33,8 @@ import (
 )
 
 const (
-	fluxPath = "/api/v2/query"
+	fluxPath      = "/api/v2/query"
+	traceIDHeader = "Trace-Id"
 )
 
 // FluxBackend is all services and associated parameters required to construct
@@ -106,6 +108,9 @@ func (h *FluxHandler) handleQuery(w http.ResponseWriter, r *http.Request) {
 
 	ctx := r.Context()
 	log := h.log.With(influxlogger.TraceFields(ctx)...)
+	if id, _, found := logger.TraceInfo(ctx); found {
+		w.Header().Set(traceIDHeader, id)
+	}
 
 	// TODO(desa): I really don't like how we're recording the usage metrics here
 	// Ideally this will be moved when we solve https://github.com/influxdata/influxdb/issues/13403
diff --git a/http/query_handler_test.go b/http/query_handler_test.go
index b083d4b7c6..9d03f83d95 100644
--- a/http/query_handler_test.go
+++ b/http/query_handler_test.go
@@ -15,6 +15,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/influxdata/influxdb/kit/tracing"
+
 	"github.com/google/go-cmp/cmp"
 	"github.com/influxdata/flux"
 	"github.com/influxdata/flux/csv"
@@ -321,6 +323,8 @@ var _ metric.EventRecorder = noopEventRecorder{}
 
 // Certain error cases must be encoded as influxdb.Error so they can be properly decoded clientside.
 func TestFluxHandler_PostQuery_Errors(t *testing.T) {
+	defer tracing.JaegerTestSetupAndTeardown(t.Name())()
+
 	i := inmem.NewService()
 	b := &FluxBackend{
 		HTTPErrorHandler:    ErrorHandler(0),
@@ -349,6 +353,10 @@ func TestFluxHandler_PostQuery_Errors(t *testing.T) {
 
 		defer resp.Body.Close()
 
+		if actual := resp.Header.Get("Trace-Id"); actual == "" {
+			t.Error("expected trace ID header")
+		}
+
 		if resp.StatusCode != http.StatusUnauthorized {
 			t.Errorf("expected unauthorized status, got %d", resp.StatusCode)
 		}
@@ -380,6 +388,10 @@ func TestFluxHandler_PostQuery_Errors(t *testing.T) {
 
 		h.handleQuery(w, req)
 
+		if actual := w.Header().Get("Trace-Id"); actual == "" {
+			t.Error("expected trace ID header")
+		}
+
 		if w.Code != http.StatusBadRequest {
 			t.Errorf("expected bad request status, got %d", w.Code)
 		}
@@ -413,6 +425,10 @@ func TestFluxHandler_PostQuery_Errors(t *testing.T) {
 		w := httptest.NewRecorder()
 		h.handleQuery(w, req)
 
+		if actual := w.Header().Get("Trace-Id"); actual == "" {
+			t.Error("expected trace ID header")
+		}
+
 		if w.Code != http.StatusBadRequest {
 			t.Errorf("expected bad request status, got %d", w.Code)
 		}
diff --git a/http/swagger.yml b/http/swagger.yml
index b201b12135..8a4ae15380 100644
--- a/http/swagger.yml
+++ b/http/swagger.yml
@@ -3208,6 +3208,11 @@ paths:
                     enum:
                     - gzip
                     - identity
+              Trace-Id:
+                description: The Trace-Id header reports the request's trace ID, if one was generated.
+                schema:
+                  type: string
+                  description: Specifies the request's trace ID.
             content:
               text/csv:
                 schema:
diff --git a/kit/tracing/tracing.go b/kit/tracing/tracing.go
index 1887d8751b..f94db2a68a 100644
--- a/kit/tracing/tracing.go
+++ b/kit/tracing/tracing.go
@@ -7,6 +7,8 @@ import (
 	"runtime"
 	"strings"
 
+	"github.com/uber/jaeger-client-go"
+
 	"github.com/influxdata/httprouter"
 	"github.com/opentracing/opentracing-go"
 	"github.com/opentracing/opentracing-go/ext"
@@ -165,3 +167,18 @@ func StartSpanFromContextWithOperationName(ctx context.Context, operationName st
 
 	return span, ctx
 }
+
+// JaegerTestSetupAndTeardown sets the global tracer to an in memory Jaeger instance for testing.
+// The returned function should be deferred by the caller to tear down this setup after testing is complete.
+func JaegerTestSetupAndTeardown(name string) func() {
+	old := opentracing.GlobalTracer()
+	tracer, closer := jaeger.NewTracer(name,
+		jaeger.NewConstSampler(true),
+		jaeger.NewInMemoryReporter(),
+	)
+	opentracing.SetGlobalTracer(tracer)
+	return func() {
+		_ = closer.Close()
+		opentracing.SetGlobalTracer(old)
+	}
+}