o11y/pkg/telemetrystore/telemetrystorehook/instrumentation.go

package telemetrystorehook

import (
	"context"

	"github.com/hanzoai/o11y/pkg/factory"
	"github.com/hanzoai/o11y/pkg/telemetrystore"
	"github.com/hanzoai/o11y/pkg/types/ctxtypes"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/metric"
	semconv "go.opentelemetry.io/otel/semconv/v1.12.0"
	"go.opentelemetry.io/otel/trace"
)

type instrumentation struct {
	datastoreCluster string
	tracer           trace.Tracer
	meter            metric.Meter
}

func NewInstrumentationFactory() factory.ProviderFactory[telemetrystore.TelemetryStoreHook, telemetrystore.Config] {
	return factory.NewProviderFactory(factory.MustNewName("instrumentation"), func(ctx context.Context, ps factory.ProviderSettings, c telemetrystore.Config) (telemetrystore.TelemetryStoreHook, error) {
		return NewInstrumentation(ctx, ps, c)
	})
}

func NewInstrumentation(ctx context.Context, providerSettings factory.ProviderSettings, config telemetrystore.Config) (telemetrystore.TelemetryStoreHook, error) {
	meter := providerSettings.MeterProvider.Meter("github.com/hanzoai/o11y/pkg/telemetrystore")

	return &instrumentation{
		datastoreCluster: config.Datastore.Cluster,
		tracer:           providerSettings.TracerProvider.Tracer("github.com/hanzoai/o11y/pkg/telemetrystore"),
		meter:            meter,
	}, nil
}

func (hook *instrumentation) BeforeQuery(ctx context.Context, event *telemetrystore.QueryEvent) context.Context {
	ctx, span := hook.tracer.Start(ctx, "", trace.WithSpanKind(trace.SpanKindClient))

	// add trace_id and span_id to the log_comment
	comment := ctxtypes.CommentFromContext(ctx)
	comment.Set("trace_id", span.SpanContext().TraceID().String())
	comment.Set("span_id", span.SpanContext().SpanID().String())
	ctx = ctxtypes.NewContextWithComment(ctx, comment)

	return ctx
}

func (hook *instrumentation) AfterQuery(ctx context.Context, event *telemetrystore.QueryEvent) {
	span := trace.SpanFromContext(ctx)
	if !span.IsRecording() {
		return
	}

	span.SetName(event.Operation)
	defer span.End()

	var attrs []attribute.KeyValue
	attrs = append(
		attrs,
		semconv.DBStatementKey.String(event.Query),
		semconv.DBSystemKey.String("datastore"),
		semconv.DBOperationKey.String(event.Operation),
		attribute.String("datastore.cluster", hook.datastoreCluster),
	)

	if event.Err != nil {
		span.RecordError(event.Err)
		span.SetStatus(codes.Error, event.Err.Error())
	}

	span.SetAttributes(attrs...)
}