From 7da84678fd0858c9cecd8f4188aece979e900487 Mon Sep 17 00:00:00 2001
From: Manish Dangi <manish.dangi@gojek.com>
Date: Thu, 21 Mar 2024 01:38:57 +0530
Subject: [PATCH] feat: add support for custome labels

---
 query.go | 40 ++++++++++++++++++++++++++++++++++++++--
 1 file changed, 38 insertions(+), 2 deletions(-)

diff --git a/query.go b/query.go
index 3e162ba..616f6e4 100644
--- a/query.go
+++ b/query.go
@@ -3,6 +3,8 @@ package otelsql
 import (
 	"context"
 	"database/sql/driver"
+
+	"go.opentelemetry.io/otel/attribute"
 )
 
 const (
@@ -24,11 +26,41 @@ func skippedQueryContext(_ context.Context, _ string, _ []driver.NamedValue) (dr
 	return nil, driver.ErrSkip
 }
 
+// add custome labels to metrics and traces
+type ctxKey string
+
+const (
+	meterCtxKey  ctxKey = "meterCtxKey"
+	tracerCtxKey ctxKey = "traceCtxKey"
+)
+
+func getLabels(ctx context.Context, key ctxKey) []attribute.KeyValue {
+	labels, ok := ctx.Value(key).([]attribute.KeyValue)
+	if !ok || labels == nil {
+		labels = make([]attribute.KeyValue, 0)
+	}
+	return labels
+}
+
+func addLabels(ctx context.Context, key ctxKey, labels ...attribute.KeyValue) context.Context {
+	new := append(getLabels(ctx, key), labels...)
+	return context.WithValue(ctx, key, new)
+}
+
+func AddMeterLabels(ctx context.Context, labels ...attribute.KeyValue) context.Context {
+	return addLabels(ctx, meterCtxKey, labels...)
+}
+
+func AddTracerLabels(ctx context.Context, labels ...attribute.KeyValue) context.Context {
+	return addLabels(ctx, tracerCtxKey, labels...)
+}
+
 // queryStats records metrics for query.
 func queryStats(r methodRecorder, method string) queryContextFuncMiddleware {
 	return func(next queryContextFunc) queryContextFunc {
 		return func(ctx context.Context, query string, args []driver.NamedValue) (result driver.Rows, err error) {
-			end := r.Record(ctx, method)
+			labels := getLabels(ctx, meterCtxKey)
+			end := r.Record(ctx, method, labels...)
 
 			defer func() {
 				end(err)
@@ -49,7 +81,11 @@ func queryTrace(t methodTracer, traceQuery queryTracer, method string) queryCont
 			ctx, end := t.Trace(ctx, method)
 
 			defer func() {
-				end(err, traceQuery(ctx, query, args)...)
+				labels := append(
+					getLabels(ctx, tracerCtxKey),
+					traceQuery(ctx, query, args)...,
+				)
+				end(err, labels...)
 			}()
 
 			result, err = next(ctx, query, args)