From 59fdf1acb00d807de3ad6d11ad0683bab643b451 Mon Sep 17 00:00:00 2001
From: duffn <3457341+duffn@users.noreply.github.com>
Date: Sat, 20 Jul 2024 14:07:35 -0600
Subject: [PATCH] fix: Add Go controller comment

---
 go/database/sql/connection.go      | 4 ++++
 go/database/sql/connection_test.go | 2 +-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/go/database/sql/connection.go b/go/database/sql/connection.go
index 2dbb5f1d..d01d09cb 100644
--- a/go/database/sql/connection.go
+++ b/go/database/sql/connection.go
@@ -101,6 +101,10 @@ func (conn *sqlCommenterConn) withComment(ctx context.Context, query string) str
 		commentsMap[core.Action] = ctx.Value(core.Action).(string)
 	}
 
+	if config.EnableController && (ctx.Value(core.Controller) != nil) {
+		commentsMap[core.Controller] = ctx.Value(core.Controller).(string)
+	}
+
 	// `driver` information should not be coming from framework.
 	// So, explicitly adding that here.
 	if config.EnableDBDriver {
diff --git a/go/database/sql/connection_test.go b/go/database/sql/connection_test.go
index 1fcdc9f5..d977c829 100644
--- a/go/database/sql/connection_test.go
+++ b/go/database/sql/connection_test.go
@@ -147,7 +147,7 @@ func TestWithComment_WithContext(t *testing.T) {
 				},
 			),
 			query:     "SELECT 1;",
-			wantQuery: "SELECT 1/*action='any+action',application='app-1',db_driver='database%2Fsql%3Apostgres',framework='custom-golang',route='listData'*/;",
+			wantQuery: "SELECT 1/*action='any+action',application='app-1',controller='custom-controller',db_driver='database%2Fsql%3Apostgres',framework='custom-golang',route='listData'*/;",
 		},
 	}
 	for _, tc := range testCases {