Add tracing support for BadgerDB.

pkg/gofr/datasource/README.md

@@ -74,15 +74,15 @@ Therefore, GoFr utilizes a pluggable approach for new datasources by separating
 | MySQL          | ✅         | ✅    | ✅     | ✅      |           |
 | REDIS          | ✅         | ✅    | ✅     | ✅      |           |
 | PostgreSQL     | ✅         | ✅    | ✅     | ✅      |           |
-| MongoDB        | ✅         | ✅    | ✅     |        | ✅         |
+| MongoDB        | ✅         | ✅    | ✅     |  ✅      | ✅         |
 | SQLite         | ✅         | ✅    | ✅     | ✅      |           |
-| BadgerDB       | ✅         | ✅    |       |        | ✅         |
-| Cassandra      | ✅         | ✅    | ✅     |        | ✅         |
-| Clickhouse     |           | ✅    | ✅     |        | ✅         |
+| BadgerDB       | ✅         | ✅    | ✅      |    ✅    | ✅         |
+| Cassandra      | ✅         | ✅    | ✅     |  ✅      | ✅         |
+| Clickhouse     |           | ✅    | ✅     |   ✅     | ✅         |
 | FTP            |           | ✅    |       |        | ✅         |
 | SFTP           |           | ✅    |       |        | ✅         |
-| Solr           |           |   ✅   |   ✅      |        |      ✅     |
-| DGraph         | ✅  | ✅ |✅ |        ||
+| Solr           |           |   ✅   |   ✅      |   ✅     |      ✅     |
+| DGraph         | ✅  | ✅ |✅ |   ✅     ||
 | Azure Eventhub |   | ✅ |✅ |        |✅|
 
 

pkg/gofr/datasource/kv-store/badger/badger.go

@@ -3,14 +3,18 @@ package badger
 import (
 	"context"
 	"errors"
+	"fmt"
 	"strings"
 	"time"
 
+	"go.opentelemetry.io/otel/attribute"
 	"go.opentelemetry.io/otel/trace"
 
 	"github.com/dgraph-io/badger/v4"
 )
 
+var errStatusDown = errors.New("status down")
+
 type Configs struct {
 	DirPath string
 }
@@ -63,8 +67,10 @@ func (c *client) Connect() {
 	c.db = db
 }
 
-func (c *client) Get(_ context.Context, key string) (string, error) {
-	defer c.sendOperationStats(time.Now(), "GET", key, "")
+func (c *client) Get(ctx context.Context, key string) (string, error) {
+	span := c.addTrace(ctx, "get", key)
+
+	defer c.sendOperationStats(time.Now(), "GET", "get", span, key)
 
 	var value []byte
 
@@ -88,24 +94,28 @@ func (c *client) Get(_ context.Context, key string) (string, error) {
 
 	err = txn.Commit()
 	if err != nil {
-		c.logger.Debugf("error while commiting transaction: %v", err)
+		c.logger.Debugf("error while committing transaction: %v", err)
 
 		return "", err
 	}
 
 	return string(value), nil
 }
 
-func (c *client) Set(_ context.Context, key, value string) error {
-	defer c.sendOperationStats(time.Now(), "SET", key, value)
+func (c *client) Set(ctx context.Context, key, value string) error {
+	span := c.addTrace(ctx, "set", key)
+
+	defer c.sendOperationStats(time.Now(), "SET", "set", span, key, value)
 
 	return c.useTransaction(func(txn *badger.Txn) error {
 		return txn.Set([]byte(key), []byte(value))
 	})
 }
 
-func (c *client) Delete(_ context.Context, key string) error {
-	defer c.sendOperationStats(time.Now(), "DELETE", key, "")
+func (c *client) Delete(ctx context.Context, key string) error {
+	span := c.addTrace(ctx, "delete", key)
+
+	defer c.sendOperationStats(time.Now(), "DELETE", "delete", span, key, "")
 
 	return c.useTransaction(func(txn *badger.Txn) error {
 		return txn.Delete([]byte(key))
@@ -125,15 +135,16 @@ func (c *client) useTransaction(f func(txn *badger.Txn) error) error {
 
 	err = txn.Commit()
 	if err != nil {
-		c.logger.Debugf("error while commiting transaction: %v", err)
+		c.logger.Debugf("error while committing transaction: %v", err)
 
 		return err
 	}
 
 	return nil
 }
 
-func (c *client) sendOperationStats(start time.Time, methodType string, kv ...string) {
+func (c *client) sendOperationStats(start time.Time, methodType string, method string,
+	span trace.Span, kv ...string) {
 	duration := time.Since(start).Milliseconds()
 
 	c.logger.Debug(&Log{
@@ -142,6 +153,11 @@ func (c *client) sendOperationStats(start time.Time, methodType string, kv ...st
 		Key:      strings.Join(kv, " "),
 	})
 
+	if span != nil {
+		defer span.End()
+		span.SetAttributes(attribute.Int64(fmt.Sprintf("badger.%v.duration(μs)", method), time.Since(start).Microseconds()))
+	}
+
 	c.metrics.RecordHistogram(context.Background(), "app_badger_stats", float64(duration), "database", c.configs.DirPath,
 		"type", methodType)
 }
@@ -162,10 +178,24 @@ func (c *client) HealthCheck(context.Context) (any, error) {
 	if closed {
 		h.Status = "DOWN"
 
-		return &h, errors.New("status down")
+		return &h, errStatusDown
 	}
 
 	h.Status = "UP"
 
 	return &h, nil
 }
+
+func (c *client) addTrace(ctx context.Context, method, key string) trace.Span {
+	if c.tracer != nil {
+		_, span := c.tracer.Start(ctx, fmt.Sprintf("badger-%v", method))
+
+		span.SetAttributes(
+			attribute.String("badger.key", key),
+		)
+
+		return span
+	}
+
+	return nil
+}

pkg/gofr/datasource/kv-store/badger/badger_test.go

@@ -46,6 +46,7 @@ func Test_ClientGet(t *testing.T) {
 	cl := setupDB(t)
 
 	err := cl.Set(context.Background(), "lkey", "lvalue")
+	require.NoError(t, err)
 
 	val, err := cl.Get(context.Background(), "lkey")
 
@@ -58,7 +59,7 @@ func Test_ClientGetError(t *testing.T) {
 
 	val, err := cl.Get(context.Background(), "lkey")
 
-	assert.EqualError(t, err, "Key not found")
+	require.EqualError(t, err, "Key not found")
 	assert.Empty(t, val)
 }
 

pkg/gofr/datasource/kv-store/badger/go.mod

@@ -5,6 +5,7 @@ go 1.22
 require (
 	github.com/dgraph-io/badger/v4 v4.2.0
 	github.com/stretchr/testify v1.9.0
+	go.opentelemetry.io/otel v1.30.0
 	go.opentelemetry.io/otel/trace v1.30.0
 	go.uber.org/mock v0.4.0
 )
@@ -25,7 +26,6 @@ require (
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	go.opencensus.io v0.22.5 // indirect
-	go.opentelemetry.io/otel v1.30.0 // indirect
 	golang.org/x/net v0.23.0 // indirect
 	golang.org/x/sys v0.18.0 // indirect
 	google.golang.org/protobuf v1.33.0 // indirect

pkg/gofr/datasource/kv-store/badger/logger.go

@@ -3,7 +3,6 @@ package badger
 import (
 	"fmt"
 	"io"
-	"strings"
 )
 
 type Logger interface {
@@ -24,5 +23,5 @@ type Log struct {
 
 func (l *Log) PrettyPrint(writer io.Writer) {
 	fmt.Fprintf(writer, "\u001B[38;5;8m%-32s \u001B[38;5;162m%-6s\u001B[0m %8d\u001B[38;5;8mµs\u001B[0m %s \n",
-		l.Type, "BADGR", l.Duration, strings.Join([]string{l.Key, l.Value}, " "))
+		l.Type, "BADGR", l.Duration, l.Key+" "+l.Value)
 }

pkg/gofr/external_db.go

@@ -92,6 +92,10 @@ func (a *App) AddKVStore(db container.KVStoreProvider) {
 	db.UseLogger(a.Logger())
 	db.UseMetrics(a.Metrics())
 
+	tracer := otel.GetTracerProvider().Tracer("gofr-badger")
+
+	db.UseTracer(tracer)
+
 	db.Connect()
 
 	a.container.KVStore = db

pkg/gofr/external_db_test.go

@@ -22,6 +22,7 @@ func TestApp_AddKVStore(t *testing.T) {
 
 		mock.EXPECT().UseLogger(app.Logger())
 		mock.EXPECT().UseMetrics(app.Metrics())
+		mock.EXPECT().UseTracer(otel.GetTracerProvider().Tracer("gofr-badger"))
 		mock.EXPECT().Connect()
 
 		app.AddKVStore(mock)