Add new relic segments for auth middleware.

Comfy-Org · Jan 19, 2025 · ff03fcf · ff03fcf
1 parent 2e310d7
commit ff03fcf
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 0 deletions.
diff --git a/server/middleware/authentication/firebase_auth.go b/server/middleware/authentication/firebase_auth.go
@@ -9,6 +9,7 @@ import (
 	"registry-backend/ent"
 	"strings"
 
+	"github.com/newrelic/go-agent/v3/newrelic"
 	"github.com/rs/zerolog/log"
 
 	firebase "firebase.google.com/go"
@@ -50,6 +51,10 @@ func FirebaseAuthMiddleware(entClient *ent.Client) echo.MiddlewareFunc {
 
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
 		return func(ctx echo.Context) error {
+			if txn := newrelic.FromContext(ctx.Request().Context()); txn != nil {
+				segment := txn.StartSegment("FirebaseAuthMiddleware")
+				defer segment.End()
+			}
 			// Check if the request is in the allow list.
 			reqPath := ctx.Request().URL.Path
 			reqMethod := ctx.Request().Method

diff --git a/server/middleware/authentication/jwt_admin_auth.go b/server/middleware/authentication/jwt_admin_auth.go
@@ -10,6 +10,7 @@ import (
 
 	"github.com/golang-jwt/jwt/v5"
 	"github.com/labstack/echo/v4"
+	"github.com/newrelic/go-agent/v3/newrelic"
 )
 
 // JWTAdminAuthMiddleware checks for a JWT token in the Authorization header,
@@ -33,6 +34,10 @@ func JWTAdminAuthMiddleware(entClient *ent.Client, secret string) echo.Middlewar
 
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
 		return func(c echo.Context) error {
+			if txn := newrelic.FromContext(c.Request().Context()); txn != nil {
+				segment := txn.StartSegment("JWTAdminAuthMiddleware")
+				defer segment.End()
+			}
 			reqPath := c.Request().URL.Path
 
 			// Check if the request path matches any of the protected endpoints

diff --git a/server/middleware/authentication/service_account_auth.go b/server/middleware/authentication/service_account_auth.go
@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/labstack/echo/v4"
+	"github.com/newrelic/go-agent/v3/newrelic"
 	"github.com/rs/zerolog/log"
 	"google.golang.org/api/idtoken"
 )
@@ -22,6 +23,10 @@ func ServiceAccountAuthMiddleware() echo.MiddlewareFunc {
 
 	return func(next echo.HandlerFunc) echo.HandlerFunc {
 		return func(ctx echo.Context) error {
+			if txn := newrelic.FromContext(ctx.Request().Context()); txn != nil {
+				segment := txn.StartSegment("ServiceAccountAuthMiddleware")
+				defer segment.End()
+			}
 			// Check if the request reqPath and method are in the checklist
 			reqPath := ctx.Request().URL.Path
 			reqMethod := ctx.Request().Method