Implement appendCallerID feature in vtgate executor

go/flags/endtoend/vtcombo.txt

@@ -16,6 +16,7 @@ Flags:
       --alsologtostderr                                                  log to standard error as well as files
       --app-idle-timeout duration                                        Idle timeout for app connections (default 1m0s)
       --app-pool-size int                                                Size of the connection pool for app connections (default 40)
+      --append-caller-id-to-queries                                      Appends immediate CallerId username into trailing comment of every query through to MySQL
       --backup-engine-implementation string                              Specifies which implementation to use for creating new backups (builtin or xtrabackup). Restores will always be done with whichever engine created a given backup. (default "builtin")
       --backup-storage-block-size int                                    if backup-storage-compress is true, backup-storage-block-size sets the byte size for each block while compressing (default is 250000). (default 250000)
       --backup-storage-compress                                          if set, the backup files will be compressed. (default true)

go/flags/endtoend/vtgate.txt

@@ -28,6 +28,7 @@ Flags:
       --allow-kill-statement                                             Allows the execution of kill statement
       --allowed-tablet-types strings                                     Specifies the tablet types this vtgate is allowed to route queries to. Should be provided as a comma-separated set of tablet types.
       --alsologtostderr                                                  log to standard error as well as files
+      --append-caller-id-to-queries                                      Appends immediate CallerId username into trailing comment of every query through to MySQL
       --balancer-keyspaces strings                                       When in balanced mode, a comma-separated list of keyspaces for which to use the balancer (optional)
       --balancer-vtgate-cells strings                                    When in balanced mode, a comma-separated list of cells that contain vtgates (required)
       --bind-address string                                              Bind address for the server. If empty, the server will listen on all available unicast and anycast IP addresses of the local system.

go/vt/vtgate/caller_id_audit.go

@@ -0,0 +1,16 @@
+package vtgate
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"vitess.io/vitess/go/vt/callerid"
+)
+
+// appends the immediateCallerID to a SQL query as a trailing comment.
+// the comment persists through MySQL to the general logs.
+func addCallerIDUserToQuery(ctx context.Context, sql string) string {
+	safeCallerID := strings.Split(callerid.ImmediateCallerIDFromContext(ctx).GetUsername(), " ")[0]
+	return fmt.Sprintf("%s/* user:%s */", sql, safeCallerID)
+}

go/vt/vtgate/caller_id_audit_test.go

@@ -0,0 +1,172 @@
+/*
+Copyright 2019 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package vtgate
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"vitess.io/vitess/go/vt/callerid"
+)
+
+func TestAddCallerIDUserToQuery(t *testing.T) {
+	tests := []struct {
+		name     string
+		ctx      context.Context
+		sql      string
+		expected string
+	}{
+		{
+			name:     "valid caller ID with simple username",
+			ctx:      callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("testuser")),
+			sql:      "SELECT * FROM users",
+			expected: "SELECT * FROM users/* user:testuser */",
+		},
+		{
+			name:     "valid caller ID with username containing spaces",
+			ctx:      callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("test user with spaces")),
+			sql:      "SELECT * FROM users",
+			expected: "SELECT * FROM users/* user:test */",
+		},
+		{
+			name:     "valid caller ID with username containing multiple spaces",
+			ctx:      callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("first second third")),
+			sql:      "INSERT INTO table VALUES (1)",
+			expected: "INSERT INTO table VALUES (1)/* user:first */",
+		},
+		{
+			name:     "valid caller ID with empty username",
+			ctx:      callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("")),
+			sql:      "DELETE FROM table",
+			expected: "DELETE FROM table/* user: */",
+		},
+		{
+			name:     "nil caller ID",
+			ctx:      callerid.NewContext(context.Background(), nil, nil),
+			sql:      "UPDATE table SET col=1",
+			expected: "UPDATE table SET col=1/* user: */",
+		},
+		{
+			name:     "context without caller ID",
+			ctx:      context.Background(),
+			sql:      "SHOW TABLES",
+			expected: "SHOW TABLES/* user: */",
+		},
+		{
+			name:     "empty SQL query",
+			ctx:      callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("user")),
+			sql:      "",
+			expected: "/* user:user */",
+		},
+		{
+			name:     "complex SQL query with caller ID",
+			ctx:      callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("complex_user")),
+			sql:      "SELECT u.id, u.name FROM users u JOIN orders o ON u.id = o.user_id WHERE u.status = 'active'",
+			expected: "SELECT u.id, u.name FROM users u JOIN orders o ON u.id = o.user_id WHERE u.status = 'active'/* user:complex_user */",
+		},
+		{
+			name:     "SQL query with existing comment",
+			ctx:      callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("user")),
+			sql:      "SELECT * FROM table /* existing comment */",
+			expected: "SELECT * FROM table /* existing comment *//* user:user */",
+		},
+		{
+			name:     "username with leading/trailing spaces",
+			ctx:      callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("  trimmed  ")),
+			sql:      "SELECT 1",
+			expected: "SELECT 1/* user: */",
+		},
+		{
+			name:     "username that is only spaces",
+			ctx:      callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("   ")),
+			sql:      "SELECT 1",
+			expected: "SELECT 1/* user: */",
+		},
+		{
+			name:     "single character username",
+			ctx:      callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("x")),
+			sql:      "SELECT 1",
+			expected: "SELECT 1/* user:x */",
+		},
+		{
+			name:     "username with special characters",
+			ctx:      callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("[email protected]")),
+			sql:      "SELECT 1",
+			expected: "SELECT 1/* user:[email protected] */",
+		},
+		{
+			name:     "username with special characters and spaces",
+			ctx:      callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("[email protected] admin role")),
+			sql:      "SELECT 1",
+			expected: "SELECT 1/* user:[email protected] */",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := addCallerIDUserToQuery(tt.ctx, tt.sql)
+			assert.Equal(t, tt.expected, result)
+		})
+	}
+}
+
+// TestAddCallerIDUserToQuery_EdgeCases tests additional edge cases to ensure 100% coverage
+func TestAddCallerIDUserToQuery_EdgeCases(t *testing.T) {
+	t.Run("multiline SQL query", func(t *testing.T) {
+		ctx := callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("multiline_user"))
+		sql := `SELECT *
+FROM users
+WHERE id = 1`
+		expected := `SELECT *
+FROM users
+WHERE id = 1/* user:multiline_user */`
+		result := addCallerIDUserToQuery(ctx, sql)
+		assert.Equal(t, expected, result)
+	})
+
+	t.Run("SQL with newlines and tabs", func(t *testing.T) {
+		ctx := callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("tab_user"))
+		sql := "SELECT\t*\nFROM\ttable"
+		expected := "SELECT\t*\nFROM\ttable/* user:tab_user */"
+		result := addCallerIDUserToQuery(ctx, sql)
+		assert.Equal(t, expected, result)
+	})
+
+	t.Run("very long username with spaces", func(t *testing.T) {
+		longUsername := "very long username with many spaces and words that should be truncated at first space"
+		ctx := callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID(longUsername))
+		sql := "SELECT 1"
+		expected := "SELECT 1/* user:very */"
+		result := addCallerIDUserToQuery(ctx, sql)
+		assert.Equal(t, expected, result)
+	})
+}
+
+// TestAddCallerIDUserToQuery_Consistency tests that the function consistently behaves the same way
+func TestAddCallerIDUserToQuery_Consistency(t *testing.T) {
+	ctx := callerid.NewContext(context.Background(), nil, callerid.NewImmediateCallerID("consistent_user"))
+	sql := "SELECT * FROM table"
+	expected := "SELECT * FROM table/* user:consistent_user */"
+
+	// Call the function multiple times to ensure consistency
+	for i := 0; i < 5; i++ {
+		result := addCallerIDUserToQuery(ctx, sql)
+		assert.Equal(t, expected, result, "Function should return consistent results on call %d", i+1)
+	}
+}

go/vt/vtgate/executor.go

@@ -145,6 +145,8 @@ type (
 
 		vConfig   econtext.VCursorConfig
 		ddlConfig dynamicconfig.DDL
+
+		appendCallerID bool
 	}
 
 	Metrics struct {
@@ -195,6 +197,7 @@ func NewExecutor(
 		plans:               plans,
 		warmingReadsChannel: make(chan bool, warmingReadsConcurrency),
 		ddlConfig:           ddlConfig,
+		appendCallerID:      appendCallerID,
 	}
 	// setting the vcursor config.
 	e.initVConfig(warnOnShardedOnly, pv)

go/vt/vtgate/plan_execute.go

@@ -70,6 +70,11 @@ func (e *Executor) newExecute(
 	execPlan planExec, // used when there is a plan to execute
 	recResult txResult, // used when it's something simple like begin/commit/rollback/savepoint
 ) (err error) {
+	// Append caller ID if enabled
+	if e.appendCallerID {
+		sql = addCallerIDUserToQuery(ctx, sql)
+	}
+
 	// Start an implicit transaction if necessary.
 	err = e.startTxIfNecessary(ctx, safeSession)
 	if err != nil {

go/vt/vtgate/vtgate.go

@@ -97,6 +97,8 @@ var (
 	lockHeartbeatTime = 5 * time.Second
 	warnShardedOnly   bool
 
+	appendCallerID bool
+
 	// ddl related flags
 	foreignKeyMode     = "allow"
 	dbDDLPlugin        = "fail"
@@ -205,6 +207,7 @@ func registerFlags(fs *pflag.FlagSet) {
 	fs.IntVar(&warmingReadsPercent, "warming-reads-percent", 0, "Percentage of reads on the primary to forward to replicas. Useful for keeping buffer pools warm")
 	fs.IntVar(&warmingReadsConcurrency, "warming-reads-concurrency", 500, "Number of concurrent warming reads allowed")
 	fs.DurationVar(&warmingReadsQueryTimeout, "warming-reads-query-timeout", 5*time.Second, "Timeout of warming read queries")
+	fs.BoolVar(&appendCallerID, "append-caller-id-to-queries", appendCallerID, "Appends immediate CallerId username into trailing comment of every query through to MySQL")
 
 	viperutil.BindFlags(fs,
 		enableOnlineDDL,