feat(query): add trace response headers

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
 

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

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)
 		}

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:

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)
+	}
+}