github.com/newrelic/go-agent@v3.26.0+incompatible/_integrations/nrawssdk/internal/internal.go (about)

     1  // Copyright 2020 New Relic Corporation. All rights reserved.
     2  // SPDX-License-Identifier: Apache-2.0
     3  
     4  package internal
     5  
     6  import (
     7  	"context"
     8  	"net/http"
     9  	"reflect"
    10  
    11  	newrelic "github.com/newrelic/go-agent"
    12  	"github.com/newrelic/go-agent/internal/integrationsupport"
    13  )
    14  
    15  type contextKeyType struct{}
    16  
    17  var segmentContextKey = contextKeyType(struct{}{})
    18  
    19  type endable interface{ End() error }
    20  
    21  func getTableName(params interface{}) string {
    22  	var tableName string
    23  
    24  	v := reflect.ValueOf(params)
    25  	if v.IsValid() && v.Kind() == reflect.Ptr {
    26  		e := v.Elem()
    27  		if e.Kind() == reflect.Struct {
    28  			n := e.FieldByName("TableName")
    29  			if n.IsValid() {
    30  				if name, ok := n.Interface().(*string); ok {
    31  					if nil != name {
    32  						tableName = *name
    33  					}
    34  				}
    35  			}
    36  		}
    37  	}
    38  
    39  	return tableName
    40  }
    41  
    42  func getRequestID(hdr http.Header) string {
    43  	id := hdr.Get("X-Amzn-Requestid")
    44  	if id == "" {
    45  		// Alternative version of request id in the header
    46  		id = hdr.Get("X-Amz-Request-Id")
    47  	}
    48  	return id
    49  }
    50  
    51  // StartSegmentInputs is used as the input to StartSegment.
    52  type StartSegmentInputs struct {
    53  	HTTPRequest *http.Request
    54  	ServiceName string
    55  	Operation   string
    56  	Region      string
    57  	Params      interface{}
    58  }
    59  
    60  // StartSegment starts a segment of either type DatastoreSegment or
    61  // ExternalSegment given the serviceName provided. The segment is then added to
    62  // the request context.
    63  func StartSegment(input StartSegmentInputs) *http.Request {
    64  
    65  	httpCtx := input.HTTPRequest.Context()
    66  	txn := newrelic.FromContext(httpCtx)
    67  
    68  	var segment endable
    69  	// Service name capitalization is different for v1 and v2.
    70  	if input.ServiceName == "dynamodb" || input.ServiceName == "DynamoDB" {
    71  		segment = &newrelic.DatastoreSegment{
    72  			Product:            newrelic.DatastoreDynamoDB,
    73  			Collection:         getTableName(input.Params),
    74  			Operation:          input.Operation,
    75  			ParameterizedQuery: "",
    76  			QueryParameters:    nil,
    77  			Host:               input.HTTPRequest.URL.Host,
    78  			PortPathOrID:       input.HTTPRequest.URL.Port(),
    79  			DatabaseName:       "",
    80  			StartTime:          newrelic.StartSegmentNow(txn),
    81  		}
    82  	} else {
    83  		segment = newrelic.StartExternalSegment(txn, input.HTTPRequest)
    84  	}
    85  
    86  	integrationsupport.AddAgentSpanAttribute(txn, newrelic.SpanAttributeAWSOperation, input.Operation)
    87  	integrationsupport.AddAgentSpanAttribute(txn, newrelic.SpanAttributeAWSRegion, input.Region)
    88  
    89  	ctx := context.WithValue(httpCtx, segmentContextKey, segment)
    90  	return input.HTTPRequest.WithContext(ctx)
    91  }
    92  
    93  // EndSegment will end any segment found in the given context.
    94  func EndSegment(ctx context.Context, hdr http.Header) {
    95  	if segment, ok := ctx.Value(segmentContextKey).(endable); ok {
    96  		if id := getRequestID(hdr); "" != id {
    97  			txn := newrelic.FromContext(ctx)
    98  			integrationsupport.AddAgentSpanAttribute(txn, newrelic.SpanAttributeAWSRequestID, id)
    99  		}
   100  		segment.End()
   101  	}
   102  }