From 7432aa86305a4836c9665677bc2a2b88dcaf4139 Mon Sep 17 00:00:00 2001
From: budougumi0617 <budougumi0617@gmail.com>
Date: Tue, 5 Jan 2021 02:54:10 +0900
Subject: [PATCH] :sparkles: add buildDeferStmtWithHttpRequest

---
 nrseg.go | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/nrseg.go b/nrseg.go
index 9374091..41df547 100644
--- a/nrseg.go
+++ b/nrseg.go
@@ -99,3 +99,36 @@ func buildDeferStmt(pkgName, ctxName, segName string) *ast.DeferStmt {
 		},
 	}
 }
+
+// buildDeferStmt builds the defer statement with *http.Request.
+// ex:
+//    defer newrelic.FromContext(req.Context()).StartSegment("slow").End()
+func buildDeferStmtWithHttpRequest(pkgName, reqName, segName string) *ast.DeferStmt {
+	return &ast.DeferStmt{
+		Call: &ast.CallExpr{
+			Fun: &ast.SelectorExpr{
+				X: &ast.CallExpr{
+					Fun: &ast.SelectorExpr{
+						X: &ast.CallExpr{
+							Fun: &ast.SelectorExpr{
+								X:   &ast.Ident{Name: pkgName},
+								Sel: &ast.Ident{Name: "FromContext"},
+							},
+							Args: []ast.Expr{
+								&ast.CallExpr{
+									Fun: &ast.SelectorExpr{
+										X:   &ast.Ident{Name: reqName},
+										Sel: &ast.Ident{Name: "Context"},
+									},
+								},
+							},
+						},
+						Sel: &ast.Ident{Name: "StartSegment"},
+					},
+					Args: []ast.Expr{&ast.BasicLit{Kind: token.STRING, Value: strconv.Quote(segName)}},
+				},
+				Sel: &ast.Ident{Name: "End"},
+			},
+		},
+	}
+}