feat: OpenTelemetry tracing support

.github/workflows/build.yml

@@ -19,7 +19,7 @@ jobs:
       - name: Setup Go
         uses: actions/setup-go@v3
         with:
-          go-version: '1.23.2'
+          go-version: '1.25.5'
           cache: true
 
       - name: Run build
@@ -31,7 +31,7 @@ jobs:
           mongodb-version: ${{ matrix.mongodb-version }}
 
       - name: Run tests
-        run: go test `go list ./... | grep -v ./cmd` -v -covermode=count -coverprofile=coverage.out
+        run: GOTOOLCHAIN=go1.25.5+auto go test `go list ./... | grep -v ./cmd` -v -covermode=count -coverprofile=coverage.out
 
       - name: Convert coverage.out to coverage.lcov
         uses: jandelgado/gcov2lcov-action@v1

cmd/tracing/example.go

@@ -0,0 +1,105 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"log"
+
+	"github.com/jensteichert/colt"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/exporters/stdout/stdouttrace"
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/sdk/resource"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
+	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
+)
+
+func initTracer() *sdktrace.TracerProvider {
+	exporter, err := stdouttrace.New(
+		stdouttrace.WithPrettyPrint())
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	tp := sdktrace.NewTracerProvider(
+		sdktrace.WithSampler(sdktrace.AlwaysSample()),
+		sdktrace.WithSyncer(exporter),
+		sdktrace.WithResource(
+			resource.NewWithAttributes(
+				semconv.SchemaURL,
+				semconv.ServiceNameKey.String("colt-example"),
+				semconv.ServiceVersionKey.String("0.0.42"),
+				semconv.DeploymentEnvironmentKey.String("local"),
+			)),
+	)
+	otel.SetTracerProvider(tp)
+	otel.SetTextMapPropagator(propagation.NewCompositeTextMapPropagator(propagation.TraceContext{}))
+	return tp
+}
+
+var tracer = initTracer().Tracer("colt-example")
+
+type Database struct {
+	Todos *colt.Collection[*Todo]
+}
+
+type Todo struct {
+	colt.DocWithTimestamps `bson:",inline"`
+	Title                  string `bson:"title" json:"title"`
+}
+
+func (t *Todo) BeforeInsert() error {
+	t.DocWithTimestamps.BeforeInsert()
+	fmt.Println("BeforeInsert executed")
+	return nil
+}
+
+func initAndConnectToDbWithTracing() colt.Database {
+	connectTraceCtx, span := tracer.Start(context.Background(), "coltInitDB")
+	defer span.End()
+
+	db := colt.NewDatabase().WithContext(connectTraceCtx)
+	db.Connect("mongodb://localhost:27017/colt?readPreference=primary&directConnection=true&ssl=false", "colt")
+	return db
+}
+
+func handleTodos(ctx context.Context) []*Todo {
+	traceCtx, span := tracer.Start(ctx, "handleTodos")
+	defer span.End()
+
+	newTodo := Todo{Title: "Hello"}
+
+	contextAwareTodoCollection := database.Todos.WithContext(traceCtx)
+
+	todo, _ := contextAwareTodoCollection.Insert(&newTodo) // Will return a Todo
+	insertedTodo, _ := contextAwareTodoCollection.FindById(todo.ID)
+
+	fmt.Println(todo)
+
+	contextAwareTodoCollection.UpdateById(todo.ID, todo)
+
+	if insertedTodo != nil {
+		fmt.Println(insertedTodo.Title)
+	}
+
+	allTodos, _ := contextAwareTodoCollection.Find(bson.M{"title": "Hello"})
+	return allTodos
+}
+
+var database Database
+
+func main() {
+	db := initAndConnectToDbWithTracing()
+
+	database = Database{
+		Todos: colt.GetCollection[*Todo](&db, "todos"),
+	}
+
+	// This would usually be a request handler with its own request context
+	allTodos := handleTodos(context.Background())
+
+	for _, todo := range allTodos {
+		fmt.Println(todo.ID)
+	}
+}

collection.go

@@ -1,7 +1,9 @@
 package colt
 
 import (
+	"context"
 	"errors"
+
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo"
@@ -10,6 +12,23 @@ import (
 
 type Collection[T Document] struct {
 	collection *mongo.Collection
+	ctx        context.Context // optional context
+}
+
+func (repo Collection[T]) WithContext(ctx context.Context) *Collection[T] {
+	repo.ctx = ctx
+	return &repo
+}
+
+func (repo *Collection[T]) ctxOrDefault() (context.Context, context.CancelFunc) {
+	if ctx := repo.ctx; ctx == nil {
+		return defaultContext()
+	}
+	if _, hasDeadline := repo.ctx.Deadline(); !hasDeadline {
+		ctxWithTimeout, cf := context.WithTimeout(repo.ctx, defaultTimeout)
+		return ctxWithTimeout, cf
+	}
+	return repo.ctx, nil
 }
 
 func (repo *Collection[T]) Insert(model T) (T, error) {
@@ -23,7 +42,11 @@ func (repo *Collection[T]) Insert(model T) (T, error) {
 		}
 	}
 
-	res, err := repo.collection.InsertOne(DefaultContext(), model)
+	ctx, cancel := repo.ctxOrDefault()
+	if cancel != nil {
+		defer cancel()
+	}
+	res, err := repo.collection.InsertOne(ctx, model)
 	if err != nil && res != nil {
 		model.SetID(res.InsertedID.(string))
 	}
@@ -42,17 +65,29 @@ func (repo *Collection[T]) UpdateOne(filter interface{}, model T) error {
 		}
 	}
 
-	_, err := repo.collection.UpdateOne(DefaultContext(), filter, bson.M{"$set": model})
+	ctx, cancel := repo.ctxOrDefault()
+	if cancel != nil {
+		defer cancel()
+	}
+	_, err := repo.collection.UpdateOne(ctx, filter, bson.M{"$set": model})
 	return err
 }
 
 func (repo *Collection[T]) UpdateMany(filter interface{}, doc bson.M) error {
-	_, err := repo.collection.UpdateMany(DefaultContext(), filter, doc)
+	ctx, cancel := repo.ctxOrDefault()
+	if cancel != nil {
+		defer cancel()
+	}
+	_, err := repo.collection.UpdateMany(ctx, filter, doc)
 	return err
 }
 
 func (repo *Collection[T]) DeleteById(id string) error {
-	res, err := repo.collection.DeleteOne(DefaultContext(), bson.M{"_id": id})
+	ctx, cancel := repo.ctxOrDefault()
+	if cancel != nil {
+		defer cancel()
+	}
+	res, err := repo.collection.DeleteOne(ctx, bson.M{"_id": id})
 
 	if err != nil {
 		return err
@@ -70,43 +105,63 @@ func (repo *Collection[T]) FindById(id interface{}) (T, error) {
 }
 
 func (repo *Collection[T]) FindOne(filter interface{}) (T, error) {
+	ctx, cancel := repo.ctxOrDefault()
+	if cancel != nil {
+		defer cancel()
+	}
 	var target T
-	err := repo.collection.FindOne(DefaultContext(), filter).Decode(&target)
+	err := repo.collection.FindOne(ctx, filter).Decode(&target)
 
 	return target, err
 }
 
 func (repo *Collection[T]) Find(filter interface{}, opts ...*options.FindOptions) ([]T, error) {
-	csr, err := repo.collection.Find(DefaultContext(), filter, opts...)
+	ctx, cancel := repo.ctxOrDefault()
+	if cancel != nil {
+		defer cancel()
+	}
+	csr, err := repo.collection.Find(ctx, filter, opts...)
 	if err != nil {
 		return nil, err
 	}
 	var result = []T{}
-	if err = csr.All(DefaultContext(), &result); err != nil {
+	if err = csr.All(ctx, &result); err != nil {
 		return nil, err
 	}
 
 	return result, nil
 }
 
 func (repo *Collection[T]) CountDocuments(filter interface{}) (int64, error) {
-	count, err := repo.collection.CountDocuments(DefaultContext(), filter)
+	ctx, cancel := repo.ctxOrDefault()
+	if cancel != nil {
+		defer cancel()
+	}
+	count, err := repo.collection.CountDocuments(ctx, filter)
 	return count, err
 }
 
 func (repo *Collection[T]) Aggregate(pipeline mongo.Pipeline, opts ...*options.AggregateOptions) ([]bson.M, error) {
-	csr, err := repo.collection.Aggregate(DefaultContext(), pipeline, opts...)
+	ctx, cancel := repo.ctxOrDefault()
+	if cancel != nil {
+		defer cancel()
+	}
+	csr, err := repo.collection.Aggregate(ctx, pipeline, opts...)
 
 	var result = []bson.M{}
-	if err = csr.All(DefaultContext(), &result); err != nil {
+	if err = csr.All(ctx, &result); err != nil {
 		return nil, err
 	}
 
 	return result, nil
 }
 
 func (repo *Collection[T]) Drop() error {
-	err := repo.collection.Drop(DefaultContext())
+	ctx, cancel := repo.ctxOrDefault()
+	if cancel != nil {
+		defer cancel()
+	}
+	err := repo.collection.Drop(ctx)
 	return err
 }
 

collection_test.go

@@ -1,15 +1,17 @@
 package colt
 
 import (
+	"context"
 	"fmt"
-	"github.com/stretchr/testify/assert"
-	"go.mongodb.org/mongo-driver/bson"
-	"go.mongodb.org/mongo-driver/bson/primitive"
-	"go.mongodb.org/mongo-driver/mongo"
 	"math/rand"
 	"strings"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/assert"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/bson/primitive"
+	"go.mongodb.org/mongo-driver/mongo"
 )
 
 type testdoc struct {
@@ -306,3 +308,71 @@ func TestCollection_Aggregate(t *testing.T) {
 	collection.Drop()
 	mockDb.Disconnect()
 }
+
+func TestCollection_WithContext(t *testing.T) {
+	mockDb.Connect("mongodb://localhost:27017/colt?readPreference=primary&directConnection=true&ssl=false", "colt")
+
+	collection := GetCollection[*testdoc](&mockDb, "testdocs")
+	// Assert that the original collection has nil ctx (default)
+	assert.Nil(t, collection.ctx)
+
+	// Uses with explicit ctx
+	ctxCollection := collection.WithContext(testCtx)
+
+	// Assert that the underlying mongo collection pointer is identical
+	assert.Equal(t, collection.collection, ctxCollection.collection)
+
+	// Assert that the ctx is different
+	assert.NotEqual(t, collection.ctx, ctxCollection.ctx)
+
+	// Assert that the new collection has the context set
+	assert.NotNil(t, ctxCollection.ctx)
+	assert.Equal(t, testCtx, ctxCollection.ctx)
+
+	secondCtx := context.WithValue(context.Background(), "test", "secondCtx")
+	derivedCollection := ctxCollection.WithContext(secondCtx)
+
+	// Assert that the underlying mongo collection pointer is identical
+	assert.Equal(t, collection.collection, derivedCollection.collection)
+
+	// Assert that the ctx is different
+	assert.NotEqual(t, ctxCollection.ctx, derivedCollection.ctx)
+
+	mockDb.Disconnect()
+}
+
+func TestCollection_ctxOrDefault(t *testing.T) {
+
+	mockDb.Connect("mongodb://localhost:27017/colt?readPreference=primary&directConnection=true&ssl=false", "colt")
+
+	// Assert that the original collection has nil ctx (default)
+	staticCollection := GetCollection[*testdoc](&mockDb, "testdocs")
+	defaultCtx, defaultCancelFunc := staticCollection.ctxOrDefault()
+	assert.NotNil(t, defaultCtx)
+	assert.NotNil(t, defaultCancelFunc)
+
+	// Uses with explicit ctx
+	ctxCollection := staticCollection.WithContext(testCtx)
+
+	// Always has a deadline set
+	defaultCtx, defaultCancelFunc = ctxCollection.ctxOrDefault()
+	_, ok := defaultCtx.Deadline()
+	assert.True(t, ok)
+	// Enriches ctx with deadline and returns cancelFunc
+	assert.NotNil(t, defaultCancelFunc)
+
+	// Uses given ctx deadline if set
+	deadline := time.Now().Add(42 * time.Second)
+	deadlineCtx, cancel := context.WithDeadline(testCtx, deadline)
+
+	dbWithManualCtxDeadline := staticCollection.WithContext(deadlineCtx)
+	defaultCtx, defaultCancelFunc = dbWithManualCtxDeadline.ctxOrDefault()
+
+	ctxDeadline, ok := deadlineCtx.Deadline()
+	assert.True(t, ok)
+	assert.Equal(t, ctxDeadline, deadline)
+	// Doesnt return cancelFunc, as its managed by caller
+	assert.Nil(t, defaultCancelFunc)
+
+	cancel()
+}