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 {