datadog_wrapper.go

package traces

import (
	"context"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
)

type DataDogTracerProvider struct {
	base              trace.TracerProvider
	spanNameFormatter DataDogTracerProviderSpanNameFormatter
	operationName     string
}

func NewDataDogTraceProvider(operationName string, opts ...DataDogTracerProviderOption) trace.TracerProvider {
	ret := &DataDogTracerProvider{
		operationName: operationName,
	}
	for _, opt := range opts {
		opt(ret)
	}
	if ret.base == nil {
		ret.base = otel.GetTracerProvider()
	}
	if ret.spanNameFormatter == nil {
		ret.spanNameFormatter = defaultDataDogTracerProviderSpanNameFormatter
	}
	return ret
}

func (d *DataDogTracerProvider) Tracer(name string, options ...trace.TracerOption) trace.Tracer {
	return &dataDogTracer{
		base:              d.base.Tracer(name, options...),
		spanNameFormatter: d.spanNameFormatter,
		operationName:     d.operationName,
	}
}

type DataDogTracerProviderOption func(*DataDogTracerProvider)

func WithDataDogTracerProviderBase(base trace.TracerProvider) DataDogTracerProviderOption {
	return func(p *DataDogTracerProvider) {
		p.base = base
	}
}

func WithDataDogTracerProviderSpanNameFormatter(spanNameFormatter DataDogTracerProviderSpanNameFormatter) DataDogTracerProviderOption {
	return func(p *DataDogTracerProvider) {
		p.spanNameFormatter = spanNameFormatter
	}
}

type DataDogTracerProviderSpanNameFormatter func(ctx context.Context, operationName string, spanName string) string

type dataDogTracer struct {
	base              trace.Tracer
	spanNameFormatter DataDogTracerProviderSpanNameFormatter
	operationName     string
}

func (d *dataDogTracer) Start(ctx context.Context, spanName string,
	opts ...trace.SpanStartOption) (context.Context, trace.Span) {
	spanName = d.spanNameFormatter(ctx, d.operationName, spanName)
	rctx, span := d.base.Start(ctx, spanName, opts...)
	span.SetAttributes(
		attribute.String("operation.name", d.operationName),
		attribute.String("resource.name", spanName),
	)
	return rctx, span
}

func defaultDataDogTracerProviderSpanNameFormatter(_ context.Context, _ string, spanName string) string {
	return spanName
}