snapshot of https://github.com/ent/ent.git (master, 2026-09-15; blobs > 131072B omitted)

dialect/gremlin/ocgremlin/client.go

@@ -0,0 +1,64 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ocgremlin
+
+import (
+	"context"
+
+	"entgo.io/ent/dialect/gremlin"
+
+	"go.opencensus.io/trace"
+)
+
+// Transport is an gremlin.RoundTripper that instruments all outgoing requests with
+// OpenCensus stats and tracing.
+type Transport struct {
+	// Base is a wrapped gremlin.RoundTripper that does the actual requests.
+	Base gremlin.RoundTripper
+
+	// StartOptions are applied to the span started by this Transport around each
+	// request.
+	//
+	// StartOptions.SpanKind will always be set to trace.SpanKindClient
+	// for spans started by this transport.
+	StartOptions trace.StartOptions
+
+	// GetStartOptions allows to set start options per request. If set,
+	// StartOptions is going to be ignored.
+	GetStartOptions func(context.Context, *gremlin.Request) trace.StartOptions
+
+	// NameFromRequest holds the function to use for generating the span name
+	// from the information found in the outgoing Gremlin Request. By default the
+	// name equals the URL Path.
+	FormatSpanName func(context.Context, *gremlin.Request) string
+
+	// WithQuery, if set to true, will enable recording of gremlin queries in spans.
+	// Only allow this if it is safe to have queries recorded with respect to
+	// security.
+	WithQuery bool
+}
+
+// RoundTrip implements gremlin.RoundTripper, delegating to Base and recording stats and traces for the request.
+func (t *Transport) RoundTrip(ctx context.Context, req *gremlin.Request) (*gremlin.Response, error) {
+	spanNameFormatter := t.FormatSpanName
+	if spanNameFormatter == nil {
+		spanNameFormatter = func(context.Context, *gremlin.Request) string {
+			return "gremlin:traversal"
+		}
+	}
+	startOpts := t.StartOptions
+	if t.GetStartOptions != nil {
+		startOpts = t.GetStartOptions(ctx, req)
+	}
+
+	var rt gremlin.RoundTripper = &traceTransport{
+		base:           t.Base,
+		formatSpanName: spanNameFormatter,
+		startOptions:   startOpts,
+		withQuery:      t.WithQuery,
+	}
+	rt = statsTransport{rt}
+	return rt.RoundTrip(ctx, req)
+}

dialect/gremlin/ocgremlin/client_test.go

@@ -0,0 +1,74 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ocgremlin
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin"
+
+	"github.com/stretchr/testify/mock"
+	"go.opencensus.io/trace"
+)
+
+type mockExporter struct {
+	mock.Mock
+}
+
+func (e *mockExporter) ExportSpan(s *trace.SpanData) {
+	e.Called(s)
+}
+
+func TestTransportOptions(t *testing.T) {
+	tests := []struct {
+		name     string
+		spanName string
+		wantName string
+	}{
+		{
+			name:     "Default formatter",
+			wantName: "gremlin:traversal",
+		},
+		{
+			name:     "Custom formatter",
+			spanName: "tester",
+			wantName: "tester",
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			var exporter mockExporter
+			exporter.On(
+				"ExportSpan",
+				mock.MatchedBy(func(s *trace.SpanData) bool { return s.Name == tt.wantName })).
+				Once()
+			defer exporter.AssertExpectations(t)
+			trace.RegisterExporter(&exporter)
+			defer trace.UnregisterExporter(&exporter)
+
+			transport := &mockTransport{}
+			transport.On("RoundTrip", mock.Anything, mock.Anything).
+				Return(nil, errors.New("noop")).
+				Once()
+			defer transport.AssertExpectations(t)
+
+			rt := &Transport{
+				Base: transport,
+				GetStartOptions: func(context.Context, *gremlin.Request) trace.StartOptions {
+					return trace.StartOptions{Sampler: trace.AlwaysSample()}
+				},
+			}
+			if tt.spanName != "" {
+				rt.FormatSpanName = func(context.Context, *gremlin.Request) string {
+					return tt.spanName
+				}
+			}
+			_, _ = rt.RoundTrip(context.Background(), gremlin.NewEvalRequest("g.E()"))
+		})
+	}
+}

dialect/gremlin/ocgremlin/stats.go

@@ -0,0 +1,117 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ocgremlin
+
+import (
+	"context"
+	"strconv"
+	"time"
+
+	"entgo.io/ent/dialect/gremlin"
+
+	"go.opencensus.io/stats"
+	"go.opencensus.io/stats/view"
+	"go.opencensus.io/tag"
+)
+
+// The following measures are supported for use in custom views.
+var (
+	RequestCount = stats.Int64(
+		"gremlin/request_count",
+		"Number of Gremlin requests started",
+		stats.UnitDimensionless,
+	)
+	ResponseBytes = stats.Int64(
+		"gremlin/response_bytes",
+		"Total number of bytes in response data",
+		stats.UnitBytes,
+	)
+	RoundTripLatency = stats.Float64(
+		"gremlin/roundtrip_latency",
+		"End-to-end latency",
+		stats.UnitMilliseconds,
+	)
+)
+
+// The following tags are applied to stats recorded by this package.
+var (
+	// StatusCode is the numeric Gremlin response status code,
+	// or "error" if a transport error occurred and no status code was read.
+	StatusCode, _ = tag.NewKey("gremlin_status_code")
+)
+
+// Default distributions used by views in this package.
+var (
+	DefaultSizeDistribution    = view.Distribution(32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144, 524288, 1048576)
+	DefaultLatencyDistribution = view.Distribution(1, 2, 3, 4, 5, 6, 8, 10, 13, 16, 20, 25, 30, 40, 50, 65, 80, 100, 130, 160, 200, 250, 300, 400, 500, 650, 800, 1000, 2000, 5000, 10000, 20000, 50000, 100000)
+)
+
+// Package ocgremlin provides some convenience views for measures.
+// You still need to register these views for data to actually be collected.
+var (
+	RequestCountView = &view.View{
+		Name:        "gremlin/request_count",
+		Measure:     RequestCount,
+		Aggregation: view.Count(),
+		Description: "Count of Gremlin requests started",
+	}
+
+	ResponseCountView = &view.View{
+		Name:        "gremlin/response_count",
+		Measure:     RoundTripLatency,
+		Aggregation: view.Count(),
+		Description: "Count of responses received, by response status",
+		TagKeys:     []tag.Key{StatusCode},
+	}
+
+	ResponseBytesView = &view.View{
+		Name:        "gremlin/response_bytes",
+		Measure:     ResponseBytes,
+		Aggregation: DefaultSizeDistribution,
+		Description: "Total number of bytes in response data",
+	}
+
+	RoundTripLatencyView = &view.View{
+		Name:        "gremlin/roundtrip_latency",
+		Measure:     RoundTripLatency,
+		Aggregation: DefaultLatencyDistribution,
+		Description: "End-to-end latency, by response code",
+		TagKeys:     []tag.Key{StatusCode},
+	}
+)
+
+// Views are the default views provided by this package.
+func Views() []*view.View {
+	return []*view.View{
+		RequestCountView,
+		ResponseCountView,
+		ResponseBytesView,
+		RoundTripLatencyView,
+	}
+}
+
+// statsTransport is an gremlin.RoundTripper that collects stats for the outgoing requests.
+type statsTransport struct {
+	base gremlin.RoundTripper
+}
+
+func (t statsTransport) RoundTrip(ctx context.Context, req *gremlin.Request) (*gremlin.Response, error) {
+	stats.Record(ctx, RequestCount.M(1))
+	start := time.Now()
+	rsp, err := t.base.RoundTrip(ctx, req)
+	latency := float64(time.Since(start)) / float64(time.Millisecond)
+	var (
+		tags = make([]tag.Mutator, 1)
+		ms   = []stats.Measurement{RoundTripLatency.M(latency)}
+	)
+	if err == nil {
+		tags[0] = tag.Upsert(StatusCode, strconv.Itoa(rsp.Status.Code))
+		ms = append(ms, ResponseBytes.M(int64(len(rsp.Result.Data))))
+	} else {
+		tags[0] = tag.Upsert(StatusCode, "error")
+	}
+	_ = stats.RecordWithTags(ctx, tags, ms...)
+	return rsp, err
+}

dialect/gremlin/ocgremlin/stats_test.go

@@ -0,0 +1,94 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ocgremlin
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin"
+	"entgo.io/ent/dialect/gremlin/encoding/graphson"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"go.opencensus.io/stats/view"
+)
+
+func TestStatsCollection(t *testing.T) {
+	err := view.Register(
+		RequestCountView,
+		ResponseCountView,
+		ResponseBytesView,
+		RoundTripLatencyView,
+	)
+	require.NoError(t, err)
+
+	req := gremlin.NewEvalRequest("g.E()")
+	rsp := &gremlin.Response{RequestID: req.RequestID}
+	rsp.Status.Code = gremlin.StatusSuccess
+	rsp.Result.Data = graphson.RawMessage(
+		`{"@type": "g:List", "@value": [{"@type": "g:Int32", "@value": 42}]}`,
+	)
+
+	transport := &mockTransport{}
+	transport.On("RoundTrip", mock.Anything, mock.Anything).
+		Return(rsp, nil).
+		Once()
+	defer transport.AssertExpectations(t)
+
+	rt := &statsTransport{transport}
+	_, _ = rt.RoundTrip(context.Background(), req)
+
+	tests := []struct {
+		name   string
+		expect func(*testing.T, *view.Row)
+	}{
+		{
+			name: "gremlin/request_count",
+			expect: func(t *testing.T, row *view.Row) {
+				count, ok := row.Data.(*view.CountData)
+				require.True(t, ok)
+				assert.Equal(t, int64(1), count.Value)
+			},
+		},
+		{
+			name: "gremlin/response_count",
+			expect: func(t *testing.T, row *view.Row) {
+				count, ok := row.Data.(*view.CountData)
+				require.True(t, ok)
+				assert.Equal(t, int64(1), count.Value)
+			},
+		},
+		{
+			name: "gremlin/response_bytes",
+			expect: func(t *testing.T, row *view.Row) {
+				data, ok := row.Data.(*view.DistributionData)
+				require.True(t, ok)
+				assert.EqualValues(t, len(rsp.Result.Data), data.Sum())
+			},
+		},
+		{
+			name: "gremlin/roundtrip_latency",
+			expect: func(t *testing.T, row *view.Row) {
+				data, ok := row.Data.(*view.DistributionData)
+				require.True(t, ok)
+				assert.NotZero(t, data.Sum())
+			},
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name[strings.Index(tt.name, "/")+1:], func(t *testing.T) {
+			v := view.Find(tt.name)
+			assert.NotNil(t, v)
+			rows, err := view.RetrieveData(tt.name)
+			require.NoError(t, err)
+			require.Len(t, rows, 1)
+			tt.expect(t, rows[0])
+		})
+	}
+}

dialect/gremlin/ocgremlin/trace.go

@@ -0,0 +1,133 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ocgremlin
+
+import (
+	"context"
+	"fmt"
+
+	"entgo.io/ent/dialect/gremlin"
+
+	"go.opencensus.io/trace"
+)
+
+// Attributes recorded on the span for the requests.
+const (
+	RequestIDAttribute = "gremlin.request_id"
+	OperationAttribute = "gremlin.operation"
+	QueryAttribute     = "gremlin.query"
+	BindingAttribute   = "gremlin.binding"
+	CodeAttribute      = "gremlin.code"
+	MessageAttribute   = "gremlin.message"
+)
+
+type traceTransport struct {
+	base           gremlin.RoundTripper
+	startOptions   trace.StartOptions
+	formatSpanName func(context.Context, *gremlin.Request) string
+	withQuery      bool
+}
+
+func (t *traceTransport) RoundTrip(ctx context.Context, req *gremlin.Request) (*gremlin.Response, error) {
+	ctx, span := trace.StartSpan(ctx,
+		t.formatSpanName(ctx, req),
+		trace.WithSampler(t.startOptions.Sampler),
+		trace.WithSpanKind(trace.SpanKindClient),
+	)
+	defer span.End()
+
+	span.AddAttributes(requestAttrs(req, t.withQuery)...)
+	rsp, err := t.base.RoundTrip(ctx, req)
+	if err != nil {
+		span.SetStatus(trace.Status{Code: trace.StatusCodeUnknown, Message: err.Error()})
+		return rsp, err
+	}
+
+	span.AddAttributes(responseAttrs(rsp)...)
+	span.SetStatus(TraceStatus(rsp.Status.Code))
+	return rsp, err
+}
+
+func requestAttrs(req *gremlin.Request, withQuery bool) []trace.Attribute {
+	attrs := []trace.Attribute{
+		trace.StringAttribute(RequestIDAttribute, req.RequestID),
+		trace.StringAttribute(OperationAttribute, req.Operation),
+	}
+	if withQuery {
+		query, _ := req.Arguments[gremlin.ArgsGremlin].(string)
+		attrs = append(attrs, trace.StringAttribute(QueryAttribute, query))
+		if bindings, ok := req.Arguments[gremlin.ArgsBindings].(map[string]any); ok {
+			attrs = append(attrs, bindingsAttrs(bindings)...)
+		}
+	}
+	return attrs
+}
+
+func bindingsAttrs(bindings map[string]any) []trace.Attribute {
+	attrs := make([]trace.Attribute, 0, len(bindings))
+	for key, val := range bindings {
+		key = BindingAttribute + "." + key
+		attrs = append(attrs, bindingToAttr(key, val))
+	}
+	return attrs
+}
+
+func bindingToAttr(key string, val any) trace.Attribute {
+	switch v := val.(type) {
+	case nil:
+		return trace.StringAttribute(key, "")
+	case int64:
+		return trace.Int64Attribute(key, v)
+	case float64:
+		return trace.Float64Attribute(key, v)
+	case string:
+		return trace.StringAttribute(key, v)
+	case bool:
+		return trace.BoolAttribute(key, v)
+	default:
+		s := fmt.Sprintf("%v", v)
+		if len(s) > 256 {
+			s = s[:256]
+		}
+		return trace.StringAttribute(key, s)
+	}
+}
+
+func responseAttrs(rsp *gremlin.Response) []trace.Attribute {
+	attrs := []trace.Attribute{
+		trace.Int64Attribute(CodeAttribute, int64(rsp.Status.Code)),
+	}
+	if rsp.Status.Message != "" {
+		attrs = append(attrs, trace.StringAttribute(MessageAttribute, rsp.Status.Message))
+	}
+	return attrs
+}
+
+// TraceStatus is a utility to convert the gremlin status code to a trace.Status.
+func TraceStatus(status int) trace.Status {
+	var code int32
+	switch status {
+	case gremlin.StatusSuccess,
+		gremlin.StatusNoContent,
+		gremlin.StatusPartialContent:
+		code = trace.StatusCodeOK
+	case gremlin.StatusUnauthorized:
+		code = trace.StatusCodePermissionDenied
+	case gremlin.StatusAuthenticate:
+		code = trace.StatusCodeUnauthenticated
+	case gremlin.StatusMalformedRequest,
+		gremlin.StatusInvalidRequestArguments,
+		gremlin.StatusScriptEvaluationError:
+		code = trace.StatusCodeInvalidArgument
+	case gremlin.StatusServerError,
+		gremlin.StatusServerSerializationError:
+		code = trace.StatusCodeInternal
+	case gremlin.StatusServerTimeout:
+		code = trace.StatusCodeDeadlineExceeded
+	default:
+		code = trace.StatusCodeUnknown
+	}
+	return trace.Status{Code: code, Message: gremlin.StatusText(status)}
+}

dialect/gremlin/ocgremlin/trace_test.go

@@ -0,0 +1,254 @@
+// Copyright 2019-present Facebook Inc. All rights reserved.
+// This source code is licensed under the Apache 2.0 license found
+// in the LICENSE file in the root directory of this source tree.
+
+package ocgremlin
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"testing"
+
+	"entgo.io/ent/dialect/gremlin"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	"go.opencensus.io/trace"
+)
+
+type mockTransport struct {
+	mock.Mock
+}
+
+func (t *mockTransport) RoundTrip(ctx context.Context, req *gremlin.Request) (*gremlin.Response, error) {
+	args := t.Called(ctx, req)
+	rsp, _ := args.Get(0).(*gremlin.Response)
+	return rsp, args.Error(1)
+}
+
+func TestTraceTransportRoundTrip(t *testing.T) {
+	_, parent := trace.StartSpan(context.Background(), "parent")
+	tests := []struct {
+		name   string
+		parent *trace.Span
+	}{
+		{
+			name:   "no parent",
+			parent: nil,
+		},
+		{
+			name:   "parent",
+			parent: parent,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			transport := &mockTransport{}
+			transport.On("RoundTrip", mock.Anything, mock.Anything).
+				Run(func(args mock.Arguments) {
+					span := trace.FromContext(args.Get(0).(context.Context))
+					require.NotNil(t, span)
+					if tt.parent != nil {
+						assert.Equal(t, tt.parent.SpanContext().TraceID, span.SpanContext().TraceID)
+					}
+				}).
+				Return(nil, errors.New("noop")).
+				Once()
+			defer transport.AssertExpectations(t)
+
+			ctx, req := context.Background(), gremlin.NewEvalRequest("g.V()")
+			if tt.parent != nil {
+				ctx = trace.NewContext(ctx, tt.parent)
+			}
+			rt := &Transport{Base: transport}
+			_, _ = rt.RoundTrip(ctx, req)
+		})
+	}
+}
+
+type testExporter struct {
+	spans []*trace.SpanData
+}
+
+func (t *testExporter) ExportSpan(s *trace.SpanData) {
+	t.spans = append(t.spans, s)
+}
+
+func TestEndToEnd(t *testing.T) {
+	trace.ApplyConfig(trace.Config{DefaultSampler: trace.AlwaysSample()})
+	var exporter testExporter
+	trace.RegisterExporter(&exporter)
+	defer trace.UnregisterExporter(&exporter)
+
+	req := gremlin.NewEvalRequest("g.V()")
+	rsp := &gremlin.Response{
+		RequestID: req.RequestID,
+	}
+	rsp.Status.Code = 200
+	rsp.Status.Message = "OK"
+
+	var transport mockTransport
+	transport.On("RoundTrip", mock.Anything, mock.Anything).
+		Return(rsp, nil).
+		Once()
+	defer transport.AssertExpectations(t)
+
+	rt := &Transport{Base: &transport, WithQuery: true}
+	_, err := rt.RoundTrip(context.Background(), req)
+	require.NoError(t, err)
+
+	require.Len(t, exporter.spans, 1)
+	attrs := exporter.spans[0].Attributes
+	assert.Len(t, attrs, 5)
+	assert.Equal(t, req.RequestID, attrs["gremlin.request_id"])
+	assert.Equal(t, req.Operation, attrs["gremlin.operation"])
+	assert.Equal(t, req.Arguments[gremlin.ArgsGremlin], attrs["gremlin.query"])
+	assert.Equal(t, int64(200), attrs["gremlin.code"])
+	assert.Equal(t, "OK", attrs["gremlin.message"])
+}
+
+func TestRequestAttributes(t *testing.T) {
+	tests := []struct {
+		name      string
+		makeReq   func() *gremlin.Request
+		wantAttrs []trace.Attribute
+	}{
+		{
+			name: "Query without bindings",
+			makeReq: func() *gremlin.Request {
+				req := gremlin.NewEvalRequest("g.E().count()")
+				req.RequestID = "a8b5c664-03ca-4175-a9e7-569b46f3551c"
+				return req
+			},
+			wantAttrs: []trace.Attribute{
+				trace.StringAttribute("gremlin.request_id", "a8b5c664-03ca-4175-a9e7-569b46f3551c"),
+				trace.StringAttribute("gremlin.operation", "eval"),
+				trace.StringAttribute("gremlin.query", "g.E().count()"),
+			},
+		},
+		{
+			name: "Query with bindings",
+			makeReq: func() *gremlin.Request {
+				bindings := map[string]any{
+					"$1": "user", "$2": int64(42),
+					"$3": 3.14, "$4": bytes.Repeat([]byte{0xff}, 257),
+					"$5": true, "$6": nil,
+				}
+				req := gremlin.NewEvalRequest(
+					`g.V().hasLabel($1).has("age",$2).has("v",$3).limit($4).valueMap($5)`,
+					gremlin.WithBindings(bindings),
+				)
+				req.RequestID = "d3d986fa-bd22-41bd-b2f7-ef2f1f639260"
+				return req
+			},
+			wantAttrs: []trace.Attribute{
+				trace.StringAttribute("gremlin.request_id", "d3d986fa-bd22-41bd-b2f7-ef2f1f639260"),
+				trace.StringAttribute("gremlin.operation", "eval"),
+				trace.StringAttribute("gremlin.query", `g.V().hasLabel($1).has("age",$2).has("v",$3).limit($4).valueMap($5)`),
+				trace.StringAttribute("gremlin.binding.$1", "user"),
+				trace.Int64Attribute("gremlin.binding.$2", 42),
+				trace.Float64Attribute("gremlin.binding.$3", 3.14),
+				trace.StringAttribute("gremlin.binding.$4", func() string {
+					str := fmt.Sprintf("%v", bytes.Repeat([]byte{0xff}, 256))
+					return str[:256]
+				}()),
+				trace.BoolAttribute("gremlin.binding.$5", true),
+				trace.StringAttribute("gremlin.binding.$6", ""),
+			},
+		},
+		{
+			name: "Authentication",
+			makeReq: func() *gremlin.Request {
+				return gremlin.NewAuthRequest(
+					"d239d950-59a1-41a7-a103-908f976ebd89",
+					"user", "pass",
+				)
+			},
+			wantAttrs: []trace.Attribute{
+				trace.StringAttribute("gremlin.request_id", "d239d950-59a1-41a7-a103-908f976ebd89"),
+				trace.StringAttribute("gremlin.operation", "authentication"),
+				trace.StringAttribute("gremlin.query", ""),
+			},
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			req := tt.makeReq()
+			attrs := requestAttrs(req, true)
+			for _, attr := range attrs {
+				assert.Contains(t, tt.wantAttrs, attr)
+			}
+			assert.Len(t, attrs, len(tt.wantAttrs))
+		})
+	}
+}
+
+func TestResponseAttributes(t *testing.T) {
+	tests := []struct {
+		name      string
+		makeRsp   func() *gremlin.Response
+		wantAttrs []trace.Attribute
+	}{
+		{
+			name: "Success no message",
+			makeRsp: func() *gremlin.Response {
+				var rsp gremlin.Response
+				rsp.Status.Code = 204
+				return &rsp
+			},
+			wantAttrs: []trace.Attribute{
+				trace.Int64Attribute("gremlin.code", 204),
+			},
+		},
+		{
+			name: "Authenticate with message",
+			makeRsp: func() *gremlin.Response {
+				var rsp gremlin.Response
+				rsp.Status.Code = 407
+				rsp.Status.Message = "login required"
+				return &rsp
+			},
+			wantAttrs: []trace.Attribute{
+				trace.Int64Attribute("gremlin.code", 407),
+				trace.StringAttribute("gremlin.message", "login required"),
+			},
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			rsp := tt.makeRsp()
+			attrs := responseAttrs(rsp)
+			assert.Equal(t, tt.wantAttrs, attrs)
+		})
+	}
+}
+
+func TestTraceStatus(t *testing.T) {
+	tests := []struct {
+		in   int
+		want trace.Status
+	}{
+		{200, trace.Status{Code: trace.StatusCodeOK, Message: "Success"}},
+		{204, trace.Status{Code: trace.StatusCodeOK, Message: "No Content"}},
+		{206, trace.Status{Code: trace.StatusCodeOK, Message: "Partial Content"}},
+		{401, trace.Status{Code: trace.StatusCodePermissionDenied, Message: "Unauthorized"}},
+		{407, trace.Status{Code: trace.StatusCodeUnauthenticated, Message: "Authenticate"}},
+		{498, trace.Status{Code: trace.StatusCodeInvalidArgument, Message: "Malformed Request"}},
+		{499, trace.Status{Code: trace.StatusCodeInvalidArgument, Message: "Invalid Request Arguments"}},
+		{500, trace.Status{Code: trace.StatusCodeInternal, Message: "Server Error"}},
+		{597, trace.Status{Code: trace.StatusCodeInvalidArgument, Message: "Script Evaluation Error"}},
+		{598, trace.Status{Code: trace.StatusCodeDeadlineExceeded, Message: "Server Timeout"}},
+		{599, trace.Status{Code: trace.StatusCodeInternal, Message: "Server Serialization Error"}},
+		{600, trace.Status{Code: trace.StatusCodeUnknown, Message: ""}},
+	}
+	for _, tt := range tests {
+		assert.Equal(t, tt.want, TraceStatus(tt.in))
+	}
+}