From 3b66f2d457c18b9a1f8d28d2e2c7843f0e0912f7 Mon Sep 17 00:00:00 2001
From: Paul Brittain <paul@helio.exchange>
Date: Thu, 7 Mar 2024 14:10:10 +0100
Subject: [PATCH] Add option to set operation name override.

---
 client_interceptors.go | 14 ++++++++++++--
 config.go              | 12 ++++++++++++
 options.go             | 18 ++++++++++++------
 server_interceptors.go | 14 ++++++++++++--
 4 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/client_interceptors.go b/client_interceptors.go
index d773a5f..623bc07 100644
--- a/client_interceptors.go
+++ b/client_interceptors.go
@@ -24,7 +24,12 @@ func UnaryClientInterceptor(opts ...Option) grpc.UnaryClientInterceptor {
 			ctx = sentry.SetHubOnContext(ctx, hub)
 		}
 
-		span := sentry.StartSpan(ctx, "grpc.client")
+		operationName := defaultClientOperationName
+		if o.OperationNameOverride != "" {
+			operationName = o.OperationNameOverride
+		}
+
+		span := sentry.StartSpan(ctx, operationName)
 		ctx = span.Context()
 		md, ok := metadata.FromOutgoingContext(ctx)
 		if ok {
@@ -60,7 +65,12 @@ func StreamClientInterceptor(opts ...Option) grpc.StreamClientInterceptor {
 			ctx = sentry.SetHubOnContext(ctx, hub)
 		}
 
-		span := sentry.StartSpan(ctx, "grpc.client")
+		operationName := defaultClientOperationName
+		if o.OperationNameOverride != "" {
+			operationName = o.OperationNameOverride
+		}
+
+		span := sentry.StartSpan(ctx, operationName)
 		ctx = span.Context()
 		md, ok := metadata.FromOutgoingContext(ctx)
 		if ok {
diff --git a/config.go b/config.go
index 5154101..8a35618 100644
--- a/config.go
+++ b/config.go
@@ -65,3 +65,15 @@ func (r *reportOnOption) Apply(o *options) {
 func WithReportOn(r reporter) Option {
 	return &reportOnOption{ReportOn: r}
 }
+
+type operationNameOverride struct {
+	OperationNameOverride string
+}
+
+func (r *operationNameOverride) Apply(o *options) {
+	o.OperationNameOverride = r.OperationNameOverride
+}
+
+func WithOperationNameOverride(s string) Option {
+	return &operationNameOverride{OperationNameOverride: s}
+}
diff --git a/options.go b/options.go
index 25655a0..33ce8a9 100644
--- a/options.go
+++ b/options.go
@@ -7,14 +7,19 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+const (
+	defaultServerOperationName = "grpc.server"
+	defaultClientOperationName = "grpc.client"
+)
+
 var defaultOptions = &options{
-		Repanic: false,
-		WaitForDelivery: false,
-		ReportOn: ReportAlways,
-		Timeout: 1 * time.Second,
+	Repanic:               false,
+	WaitForDelivery:       false,
+	ReportOn:              ReportAlways,
+	Timeout:               1 * time.Second,
+	OperationNameOverride: "",
 }
 
-
 type options struct {
 	// Repanic configures whether Sentry should repanic after recovery.
 	Repanic bool
@@ -26,8 +31,9 @@ type options struct {
 	Timeout time.Duration
 
 	ReportOn func(error) bool
-}
 
+	OperationNameOverride string
+}
 
 func ReportAlways(error) bool {
 	return true
diff --git a/server_interceptors.go b/server_interceptors.go
index 5bce3ea..fb21e7b 100644
--- a/server_interceptors.go
+++ b/server_interceptors.go
@@ -41,8 +41,13 @@ func UnaryServerInterceptor(opts ...Option) grpc.UnaryServerInterceptor {
 			ctx = sentry.SetHubOnContext(ctx, hub)
 		}
 
+		operationName := defaultServerOperationName
+		if o.OperationNameOverride != "" {
+			operationName = o.OperationNameOverride
+		}
+
 		md, _ := metadata.FromIncomingContext(ctx) // nil check in ContinueFromGrpcMetadata
-		span := sentry.StartSpan(ctx, "grpc.server", ContinueFromGrpcMetadata(md))
+		span := sentry.StartSpan(ctx, operationName, ContinueFromGrpcMetadata(md))
 		ctx = span.Context()
 		defer span.Finish()
 
@@ -79,8 +84,13 @@ func StreamServerInterceptor(opts ...Option) grpc.StreamServerInterceptor {
 			ctx = sentry.SetHubOnContext(ctx, hub)
 		}
 
+		operationName := defaultServerOperationName
+		if o.OperationNameOverride != "" {
+			operationName = o.OperationNameOverride
+		}
+
 		md, _ := metadata.FromIncomingContext(ctx) // nil check in ContinueFromGrpcMetadata
-		span := sentry.StartSpan(ctx, "grpc.server", ContinueFromGrpcMetadata(md))
+		span := sentry.StartSpan(ctx, operationName, ContinueFromGrpcMetadata(md))
 		ctx = span.Context()
 		defer span.Finish()