Merge pull request #4950 from Dokiys/feature/orm_otel_filter

Feat: add orm opentelemetry filter

client/orm/filter/opentelemetry/filter.go

@@ -0,0 +1,90 @@
+// Copyright 2020 beego
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package opentelemetry
+
+import (
+	"context"
+	"strings"
+
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	otelTrace "go.opentelemetry.io/otel/trace"
+
+	"github.com/beego/beego/v2/client/orm"
+)
+
+type (
+	CustomSpanFunc    func(ctx context.Context, span otelTrace.Span, inv *orm.Invocation)
+	FilterChainOption func(fcv *FilterChainBuilder)
+)
+
+// FilterChainBuilder provides an opentelemtry Filter
+type FilterChainBuilder struct {
+	// customSpanFunc users are able to custom their span
+	customSpanFunc CustomSpanFunc
+}
+
+func NewFilterChainBuilder(options ...FilterChainOption) *FilterChainBuilder {
+	fcb := &FilterChainBuilder{}
+	for _, o := range options {
+		o(fcb)
+	}
+	return fcb
+}
+
+// WithCustomSpanFunc add function to custom span
+func WithCustomSpanFunc(customSpanFunc CustomSpanFunc) FilterChainOption {
+	return func(fcv *FilterChainBuilder) {
+		fcv.customSpanFunc = customSpanFunc
+	}
+}
+
+// FilterChain traces invocation with opentelemetry
+// Unless invocation.Method is Begin*, Commit or Rollback
+func (builder *FilterChainBuilder) FilterChain(next orm.Filter) orm.Filter {
+	return func(ctx context.Context, inv *orm.Invocation) []interface{} {
+		if strings.HasPrefix(inv.Method, "Begin") || inv.Method == "Commit" || inv.Method == "Rollback" {
+			return next(ctx, inv)
+		}
+		spanCtx, span := otel.Tracer("beego_orm").Start(ctx, invOperationName(ctx, inv))
+		defer span.End()
+
+		res := next(spanCtx, inv)
+		builder.buildSpan(spanCtx, span, inv)
+		return res
+	}
+}
+
+// buildSpan add default span attributes and custom attributes with customSpanFunc
+func (builder *FilterChainBuilder) buildSpan(ctx context.Context, span otelTrace.Span, inv *orm.Invocation) {
+	span.SetAttributes(attribute.String("orm.method", inv.Method))
+	span.SetAttributes(attribute.String("orm.table", inv.GetTableName()))
+	span.SetAttributes(attribute.Bool("orm.insideTx", inv.InsideTx))
+	v, _ := ctx.Value(orm.TxNameKey).(string)
+	span.SetAttributes(attribute.String("orm.txName", v))
+	span.SetAttributes(attribute.String("span.kind", "client"))
+	span.SetAttributes(attribute.String("component", "beego"))
+
+	if builder.customSpanFunc != nil {
+		builder.customSpanFunc(ctx,span, inv)
+	}
+}
+
+func invOperationName(ctx context.Context, inv *orm.Invocation) string {
+	if n, ok := ctx.Value(orm.TxNameKey).(string); ok {
+		return inv.Method + "#tx(" + n + ")"
+	}
+	return inv.Method + "#" + inv.GetTableName()
+}

client/orm/filter/opentelemetry/filter_test.go

@@ -0,0 +1,61 @@
+// Copyright 2020 beego
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package opentelemetry
+
+import (
+	"bytes"
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/exporters/stdout/stdouttrace"
+	"go.opentelemetry.io/otel/sdk/trace"
+	otelTrace "go.opentelemetry.io/otel/trace"
+
+	"github.com/beego/beego/v2/client/orm"
+)
+
+func TestFilterChainBuilderFilterChain(t *testing.T) {
+	// Init Trace
+	buf := bytes.NewBuffer([]byte{})
+	exp, err := stdouttrace.New(stdouttrace.WithWriter(buf))
+	if err != nil {
+		t.Error(err)
+	}
+	tp := trace.NewTracerProvider(trace.WithBatcher(exp))
+	otel.SetTracerProvider(tp)
+
+	// Build FilterChain
+	csf := func(ctx context.Context, span otelTrace.Span, inv *orm.Invocation) {
+		span.SetAttributes(attribute.String("hello", "work"))
+	}
+	builder := NewFilterChainBuilder(WithCustomSpanFunc(csf))
+
+	inv := &orm.Invocation{Method: "Hello"}
+	next := func(ctx context.Context, inv *orm.Invocation) []interface{} { return nil }
+
+	builder.FilterChain(next)(context.Background(), inv)
+
+	// Close tp
+	err = tp.Shutdown(context.Background())
+	if err != nil {
+		t.Error(err)
+	}
+
+	// Assert opentelemetry span name
+	assert.Equal(t, "Hello#", string(buf.Bytes()[9:15]))
+}