Adds config support for GRPC and improves test assertions

instrumentation/opencensus/database/hypersql/sql_test.go

@@ -1,6 +1,6 @@
 package hypersql
 
-// highly inspired by https://github.com/openzipkin-contrib/zipkin-go-sql/blob/master/driver_test.go
+// highly inspired in https://github.com/openzipkin-contrib/zipkin-go-sql/blob/master/driver_test.go
 
 import (
 	"context"

instrumentation/opentelemetry/database/hypersql/sql.go

@@ -111,6 +111,7 @@ func (in *interceptor) ConnPrepareContext(ctx context.Context, conn driver.ConnP
 		span.SetAttribute(key, value)
 	}
 
+	span.SetAttribute("db.statement", query)
 	defer span.End()
 
 	tx, err := conn.PrepareContext(ctx, query)

sdk/google.golang.org/grpc/attributes_test.go

@@ -13,14 +13,20 @@ func TestSetScalarAttributeSuccess(t *testing.T) {
 	span := mock.NewSpan()
 	setAttributesFromMetadata("request", md, span)
 
-	assert.Equal(t, "value_1", span.Attributes["rpc.request.metadata.key_1"].(string))
+	assert.Equal(t, "value_1", span.ReadAttribute("rpc.request.metadata.key_1").(string))
+
+	_ = span.ReadAttribute("container_id") // needed in containarized envs
+	assert.Zero(t, span.RemainingAttributes(), "unexpected remaining attribute: %v", span.Attributes)
 }
 
 func TestSetMultivalueAttributeSuccess(t *testing.T) {
 	md := metadata.Pairs("key_1", "value_1", "key_1", "value_2")
 	span := mock.NewSpan()
 	setAttributesFromMetadata("request", md, span)
 
-	assert.Equal(t, "value_1", span.Attributes["rpc.request.metadata.key_1[0]"].(string))
-	assert.Equal(t, "value_2", span.Attributes["rpc.request.metadata.key_1[1]"].(string))
+	assert.Equal(t, "value_1", span.ReadAttribute("rpc.request.metadata.key_1[0]").(string))
+	assert.Equal(t, "value_2", span.ReadAttribute("rpc.request.metadata.key_1[1]").(string))
+
+	_ = span.ReadAttribute("container_id") // needed in containarized envs
+	assert.Zero(t, span.RemainingAttributes(), "unexpected remaining attribute: %v", span.Attributes)
 }

sdk/google.golang.org/grpc/client.go

@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/hypertrace/goagent/sdk"
+	internalconfig "github.com/hypertrace/goagent/sdk/internal/config"
 	"github.com/hypertrace/goagent/sdk/internal/container"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
@@ -20,6 +21,8 @@ func WrapUnaryClientInterceptor(delegateInterceptor grpc.UnaryClientInterceptor,
 		defaultAttributes["container_id"] = containerID
 	}
 
+	dataCaptureConfig := internalconfig.GetConfig().GetDataCapture()
+
 	return func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
 		var header metadata.MD
 		var trailer metadata.MD
@@ -45,22 +48,26 @@ func WrapUnaryClientInterceptor(delegateInterceptor grpc.UnaryClientInterceptor,
 			span.SetAttribute("rpc.method", pieces[1])
 
 			reqBody, err := marshalMessageableJSON(req)
-			if len(reqBody) > 0 && err == nil {
+			if dataCaptureConfig.RpcBody.Request.Value && len(reqBody) > 0 && err == nil {
 				span.SetAttribute("rpc.request.body", string(reqBody))
 			}
 
-			setAttributesFromRequestOutgoingMetadata(ctx, span)
+			if dataCaptureConfig.RpcMetadata.Request.Value {
+				setAttributesFromRequestOutgoingMetadata(ctx, span)
+			}
 
 			err = invoker(ctx, method, req, reply, cc, opts...)
 			if err != nil {
 				return err
 			}
 
-			setAttributesFromMetadata("response", header, span)
-			setAttributesFromMetadata("response", trailer, span)
+			if dataCaptureConfig.RpcMetadata.Response.Value {
+				setAttributesFromMetadata("response", header, span)
+				setAttributesFromMetadata("response", trailer, span)
+			}
 
 			resBody, err := marshalMessageableJSON(reply)
-			if len(resBody) > 0 && err == nil {
+			if dataCaptureConfig.RpcBody.Response.Value && len(resBody) > 0 && err == nil {
 				span.SetAttribute("rpc.response.body", string(resBody))
 			}
 

sdk/google.golang.org/grpc/client_test.go

@@ -68,28 +68,32 @@ func TestUnaryClientHelloWorldSuccess(t *testing.T) {
 
 	span := spans[0]
 
-	assert.Equal(t, "grpc", span.Attributes["rpc.system"].(string))
-	assert.Equal(t, "helloworld.Greeter", span.Attributes["rpc.service"].(string))
-	assert.Equal(t, "SayHello", span.Attributes["rpc.method"].(string))
-	assert.Equal(t, "test_value_1", span.Attributes["rpc.request.metadata.test_key_1"].(string))
-	assert.Equal(t, "test_header_value", span.Attributes["rpc.response.metadata.test_header_key"].(string))
-	assert.Equal(t, "test_trailer_value", span.Attributes["rpc.response.metadata.test_trailer_key"].(string))
+	assert.Equal(t, "grpc", span.ReadAttribute("rpc.system").(string))
+	assert.Equal(t, "helloworld.Greeter", span.ReadAttribute("rpc.service").(string))
+	assert.Equal(t, "SayHello", span.ReadAttribute("rpc.method").(string))
+	assert.Equal(t, "test_value_1", span.ReadAttribute("rpc.request.metadata.test_key_1").(string))
+	assert.Equal(t, "test_header_value", span.ReadAttribute("rpc.response.metadata.test_header_key").(string))
+	assert.Equal(t, "test_trailer_value", span.ReadAttribute("rpc.response.metadata.test_trailer_key").(string))
+	assert.Equal(t, "application/grpc", span.ReadAttribute("rpc.response.metadata.content-type").(string))
 
 	expectedBody := "{\"name\":\"Cuchi\"}"
-	actualBody := span.Attributes["rpc.request.body"].(string)
+	actualBody := span.ReadAttribute("rpc.request.body").(string)
 	if ok, err := jsonEqual(expectedBody, actualBody); err == nil {
 		assert.True(t, ok, "incorrect request body:\nwant %s,\nhave %s", expectedBody, actualBody)
 	} else {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
 	expectedBody = "{\"message\":\"Hello Cuchi\"}"
-	actualBody = span.Attributes["rpc.response.body"].(string)
+	actualBody = span.ReadAttribute("rpc.response.body").(string)
 	if ok, err := jsonEqual(expectedBody, actualBody); err == nil {
 		assert.True(t, ok, "incorrect response body:\nwant %s,\nhave %s", expectedBody, actualBody)
 	} else {
 		t.Fatalf("unexpected error: %v", err)
 	}
+
+	_ = span.ReadAttribute("container_id") // needed in containarized envs
+	assert.Zero(t, span.RemainingAttributes(), "unexpected remaining attribute: %v", span.Attributes)
 }
 
 func TestClientHandlerHelloWorldSuccess(t *testing.T) {
@@ -135,23 +139,23 @@ func TestClientHandlerHelloWorldSuccess(t *testing.T) {
 
 	span := mockHandler.Spans[0]
 
-	assert.Equal(t, "grpc", span.Attributes["rpc.system"].(string))
-	assert.Equal(t, "helloworld.Greeter", span.Attributes["rpc.service"].(string))
-	assert.Equal(t, "SayHello", span.Attributes["rpc.method"].(string))
-	assert.Equal(t, "test_value_1", span.Attributes["rpc.request.metadata.test_key_1"].(string))
-	assert.Equal(t, "test_header_value", span.Attributes["rpc.response.metadata.test_header_key"].(string))
-	assert.Equal(t, "test_trailer_value", span.Attributes["rpc.response.metadata.test_trailer_key"].(string))
+	assert.Equal(t, "grpc", span.ReadAttribute("rpc.system").(string))
+	assert.Equal(t, "helloworld.Greeter", span.ReadAttribute("rpc.service").(string))
+	assert.Equal(t, "SayHello", span.ReadAttribute("rpc.method").(string))
+	assert.Equal(t, "test_value_1", span.ReadAttribute("rpc.request.metadata.test_key_1").(string))
+	assert.Equal(t, "test_header_value", span.ReadAttribute("rpc.response.metadata.test_header_key").(string))
+	assert.Equal(t, "test_trailer_value", span.ReadAttribute("rpc.response.metadata.test_trailer_key").(string))
 
 	expectedBody := "{\"name\":\"Cuchi\"}"
-	actualBody := span.Attributes["rpc.request.body"].(string)
+	actualBody := span.ReadAttribute("rpc.request.body").(string)
 	if ok, err := jsonEqual(expectedBody, actualBody); err == nil {
 		assert.True(t, ok, "incorrect request body:\nwant %s,\nhave %s", expectedBody, actualBody)
 	} else {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
 	expectedBody = "{\"message\":\"Hello Cuchi\"}"
-	actualBody = span.Attributes["rpc.response.body"].(string)
+	actualBody = span.ReadAttribute("rpc.response.body").(string)
 	if ok, err := jsonEqual(expectedBody, actualBody); err == nil {
 		assert.True(t, ok, "incorrect response body:\nwant %s,\nhave %s", expectedBody, actualBody)
 	} else {

sdk/google.golang.org/grpc/server.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"strings"
 
+	"github.com/hypertrace/goagent/config"
 	"github.com/hypertrace/goagent/sdk"
+	internalconfig "github.com/hypertrace/goagent/sdk/internal/config"
 	"github.com/hypertrace/goagent/sdk/internal/container"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/stats"
@@ -36,7 +38,7 @@ func WrapUnaryServerInterceptor(
 			ctx,
 			req,
 			info,
-			wrapHandler(info.FullMethod, handler, spanFromContext, defaultAttributes),
+			wrapHandler(info.FullMethod, handler, spanFromContext, defaultAttributes, internalconfig.GetConfig().GetDataCapture()),
 		)
 	}
 }
@@ -46,6 +48,7 @@ func wrapHandler(
 	delegateHandler grpc.UnaryHandler,
 	spanFromContext sdk.SpanFromContext,
 	defaultAttributes map[string]string,
+	dataCaptureConfig *config.DataCapture,
 ) grpc.UnaryHandler {
 	return func(ctx context.Context, req interface{}) (interface{}, error) {
 		span := spanFromContext(ctx)
@@ -65,19 +68,23 @@ func wrapHandler(
 		span.SetAttribute("rpc.method", pieces[1])
 
 		reqBody, err := marshalMessageableJSON(req)
-		if len(reqBody) > 0 && err == nil {
+		if dataCaptureConfig.RpcBody.Request.Value &&
+			len(reqBody) > 0 && err == nil {
 			span.SetAttribute("rpc.request.body", string(reqBody))
 		}
 
-		setAttributesFromRequestIncomingMetadata(ctx, span)
+		if dataCaptureConfig.RpcMetadata.Request.Value {
+			setAttributesFromRequestIncomingMetadata(ctx, span)
+		}
 
 		res, err := delegateHandler(ctx, req)
 		if err != nil {
 			return res, err
 		}
 
 		resBody, err := marshalMessageableJSON(res)
-		if len(resBody) > 0 && err == nil {
+		if dataCaptureConfig.RpcBody.Response.Value &&
+			len(resBody) > 0 && err == nil {
 			span.SetAttribute("rpc.response.body", string(resBody))
 		}
 
@@ -91,6 +98,7 @@ type handler struct {
 	stats.Handler
 	spanFromContext   sdk.SpanFromContext
 	defaultAttributes map[string]string
+	dataCaptureConfig *config.DataCapture
 }
 
 // HandleRPC implements per-RPC tracing and stats instrumentation.
@@ -102,7 +110,7 @@ func (s *handler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
 		// isNoop means either the span is not sampled or there was no span
 		// in the request context which means this Handler is not used
 		// inside an instrumented Handler, hence we just invoke the delegate
-		// round tripper.
+		// handler.
 		return
 	}
 
@@ -117,21 +125,21 @@ func (s *handler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
 			return
 		}
 
-		if rs.IsClient() {
+		if rs.IsClient() && s.dataCaptureConfig.RpcBody.Response.Value {
 			span.SetAttribute("rpc.response.body", string(body))
-		} else {
+		} else if !rs.IsClient() && s.dataCaptureConfig.RpcBody.Request.Value {
 			span.SetAttribute("rpc.request.body", string(body))
 		}
 	case *stats.InHeader:
-		if rs.IsClient() {
+		if rs.IsClient() && s.dataCaptureConfig.RpcMetadata.Response.Value {
 			setAttributesFromMetadata("response", rs.Header, span)
-		} else {
+		} else if !rs.IsClient() && s.dataCaptureConfig.RpcMetadata.Request.Value {
 			setAttributesFromMetadata("request", rs.Header, span)
 		}
 	case *stats.InTrailer:
-		if rs.IsClient() {
+		if rs.IsClient() && s.dataCaptureConfig.RpcMetadata.Response.Value {
 			setAttributesFromMetadata("response", rs.Trailer, span)
-		} else {
+		} else if !rs.IsClient() && s.dataCaptureConfig.RpcMetadata.Request.Value {
 			setAttributesFromMetadata("request", rs.Trailer, span)
 		}
 	case *stats.OutPayload:
@@ -140,21 +148,21 @@ func (s *handler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
 			return
 		}
 
-		if rs.IsClient() {
+		if rs.IsClient() && s.dataCaptureConfig.RpcBody.Request.Value {
 			span.SetAttribute("rpc.request.body", string(body))
-		} else {
+		} else if !rs.IsClient() && s.dataCaptureConfig.RpcBody.Response.Value {
 			span.SetAttribute("rpc.response.body", string(body))
 		}
 	case *stats.OutHeader:
-		if rs.IsClient() {
+		if rs.IsClient() && s.dataCaptureConfig.RpcMetadata.Request.Value {
 			setAttributesFromMetadata("request", rs.Header, span)
-		} else {
+		} else if !rs.IsClient() && s.dataCaptureConfig.RpcMetadata.Response.Value {
 			setAttributesFromMetadata("response", rs.Header, span)
 		}
 	case *stats.OutTrailer:
-		if rs.IsClient() {
+		if rs.IsClient() && s.dataCaptureConfig.RpcMetadata.Request.Value {
 			setAttributesFromMetadata("request", rs.Trailer, span)
-		} else {
+		} else if !rs.IsClient() && s.dataCaptureConfig.RpcMetadata.Response.Value {
 			setAttributesFromMetadata("response", rs.Trailer, span)
 		}
 	}
@@ -188,5 +196,10 @@ func WrapStatsHandler(delegate stats.Handler, spanFromContext sdk.SpanFromContex
 		defaultAttributes["container_id"] = containerID
 	}
 
-	return &handler{Handler: delegate, spanFromContext: spanFromContext, defaultAttributes: defaultAttributes}
+	return &handler{
+		Handler:           delegate,
+		spanFromContext:   spanFromContext,
+		defaultAttributes: defaultAttributes,
+		dataCaptureConfig: internalconfig.GetConfig().GetDataCapture(),
+	}
 }

sdk/google.golang.org/grpc/server_test.go

@@ -61,21 +61,21 @@ func TestServerInterceptorHelloWorldSuccess(t *testing.T) {
 
 	span := spans[0]
 
-	assert.Equal(t, "grpc", span.Attributes["rpc.system"].(string))
-	assert.Equal(t, "helloworld.Greeter", span.Attributes["rpc.service"].(string))
-	assert.Equal(t, "SayHello", span.Attributes["rpc.method"].(string))
-	assert.Equal(t, "test_value", span.Attributes["rpc.request.metadata.test_key"].(string))
+	assert.Equal(t, "grpc", span.ReadAttribute("rpc.system").(string))
+	assert.Equal(t, "helloworld.Greeter", span.ReadAttribute("rpc.service").(string))
+	assert.Equal(t, "SayHello", span.ReadAttribute("rpc.method").(string))
+	assert.Equal(t, "test_value", span.ReadAttribute("rpc.request.metadata.test_key").(string))
 
 	expectedBody := "{\"name\":\"Pupo\"}"
-	actualBody := span.Attributes["rpc.request.body"].(string)
+	actualBody := span.ReadAttribute("rpc.request.body").(string)
 	if ok, err := jsonEqual(expectedBody, actualBody); err == nil {
 		assert.True(t, ok, "incorrect request body:\nwant %s,\nhave %s", expectedBody, actualBody)
 	} else {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
 	expectedBody = "{\"message\":\"Hello Pupo\"}"
-	actualBody = span.Attributes["rpc.response.body"].(string)
+	actualBody = span.ReadAttribute("rpc.response.body").(string)
 	if ok, err := jsonEqual(expectedBody, actualBody); err == nil {
 		assert.True(t, ok, "incorrect response body:\nwant %s,\nhave %s", expectedBody, actualBody)
 	} else {
@@ -101,6 +101,7 @@ func TestServerHandlerHelloWorldSuccess(t *testing.T) {
 		"bufnet",
 		grpc.WithContextDialer(dialer),
 		grpc.WithInsecure(),
+		grpc.WithUserAgent("test_agent"),
 	)
 	if err != nil {
 		t.Fatalf("failed to dial bufnet: %v", err)
@@ -110,6 +111,7 @@ func TestServerHandlerHelloWorldSuccess(t *testing.T) {
 	client := helloworld.NewGreeterClient(conn)
 
 	ctx = metadata.NewOutgoingContext(ctx, metadata.Pairs("test_key", "test_value"))
+
 	_, err = client.SayHello(ctx, &helloworld.HelloRequest{
 		Name: "Pupo",
 	})
@@ -121,25 +123,31 @@ func TestServerHandlerHelloWorldSuccess(t *testing.T) {
 
 	span := mockHandler.Spans[0]
 
-	assert.Equal(t, "grpc", span.Attributes["rpc.system"].(string))
-	assert.Equal(t, "helloworld.Greeter", span.Attributes["rpc.service"].(string))
-	assert.Equal(t, "SayHello", span.Attributes["rpc.method"].(string))
-	assert.Equal(t, "test_value", span.Attributes["rpc.request.metadata.test_key"].(string))
-	assert.Equal(t, "test_value", span.Attributes["rpc.request.metadata.test_key"].(string))
+	assert.Equal(t, "grpc", span.ReadAttribute("rpc.system").(string))
+	assert.Equal(t, "helloworld.Greeter", span.ReadAttribute("rpc.service").(string))
+	assert.Equal(t, "SayHello", span.ReadAttribute("rpc.method").(string))
+	assert.Equal(t, "test_value", span.ReadAttribute("rpc.request.metadata.test_key").(string))
+
+	assert.Equal(t, "bufnet", span.ReadAttribute("rpc.request.metadata.:authority").(string))
+	assert.Equal(t, "application/grpc", span.ReadAttribute("rpc.request.metadata.content-type").(string))
+	assert.Contains(t, span.ReadAttribute("rpc.request.metadata.user-agent").(string), "test_agent")
 
 	expectedBody := "{\"name\":\"Pupo\"}"
-	actualBody := span.Attributes["rpc.request.body"].(string)
+	actualBody := span.ReadAttribute("rpc.request.body").(string)
 	if ok, err := jsonEqual(expectedBody, actualBody); err == nil {
 		assert.True(t, ok, "incorrect request body:\nwant %s,\nhave %s", expectedBody, actualBody)
 	} else {
 		t.Fatalf("unexpected error: %v", err)
 	}
 
 	expectedBody = "{\"message\":\"Hello Pupo\"}"
-	actualBody = span.Attributes["rpc.response.body"].(string)
+	actualBody = span.ReadAttribute("rpc.response.body").(string)
 	if ok, err := jsonEqual(expectedBody, actualBody); err == nil {
 		assert.True(t, ok, "incorrect response body:\nwant %s,\nhave %s", expectedBody, actualBody)
 	} else {
 		t.Fatalf("unexpected error: %v", err)
 	}
+
+	_ = span.ReadAttribute("container_id") // needed in containarized envs
+	assert.Zero(t, span.RemainingAttributes(), "unexpected remaining attribute: %v", span.Attributes)
 }