support logging session connect attrs to slow query log

Author: jiong-nba

pkg/executor/adapter_slow_log.go

@@ -264,6 +264,9 @@ func SetSlowLogItems(a *ExecStmt, txnTS uint64, hasMoreResults bool, items *vari
 	items.StorageKV = stmtCtx.IsTiKV.Load()
 	items.StorageMPP = stmtCtx.IsTiFlash.Load()
 	items.MemArbitration = stmtCtx.MemTracker.MemArbitration().Seconds()
+	if sessVars.ConnectionInfo != nil && len(sessVars.ConnectionInfo.Attributes) > 0 {
+		items.SessionConnectAttrs = sessVars.ConnectionInfo.Attributes
+	}
 
 	if a.retryCount > 0 {
 		items.ExecRetryTime = items.TimeTotal - sessVars.DurationParse - sessVars.DurationCompile - time.Since(a.retryStartTime)

pkg/executor/slow_query.go

@@ -982,6 +982,9 @@ func (e *slowQueryRetriever) parseLog(ctx context.Context, sctx sessionctx.Conte
 				} else if strings.HasPrefix(line, variable.SlowLogWarnings+variable.SlowLogSpaceMarkStr) {
 					line = line[len(variable.SlowLogWarnings+variable.SlowLogSpaceMarkStr):]
 					valid = e.setColumnValue(sctx, row, tz, variable.SlowLogWarnings, line, e.checker, fileLine)
+				} else if strings.HasPrefix(line, variable.SlowLogSessionConnectAttrs+variable.SlowLogSpaceMarkStr) {
+					line = line[len(variable.SlowLogSessionConnectAttrs+variable.SlowLogSpaceMarkStr):]
+					valid = e.setColumnValue(sctx, row, tz, variable.SlowLogSessionConnectAttrs, line, e.checker, fileLine)
 				} else if strings.HasPrefix(line, variable.SlowLogDBStr+variable.SlowLogSpaceMarkStr) {
 					line = line[len(variable.SlowLogDBStr+variable.SlowLogSpaceMarkStr):]
 					valid = e.setColumnValue(sctx, row, tz, variable.SlowLogDBStr, line, e.checker, fileLine)
@@ -1170,6 +1173,18 @@ func getColumnValueFactoryByName(colName string, columnIdx int) (slowQueryColumn
 			row[columnIdx] = types.NewDatum(v)
 			return true, nil
 		}, nil
+	case variable.SlowLogSessionConnectAttrs:
+		return func(row []types.Datum, value string, _ *time.Location, _ *slowLogChecker) (valid bool, err error) {
+			if len(value) == 0 {
+				return true, nil
+			}
+			bj, err := types.ParseBinaryJSONFromString(value)
+			if err != nil {
+				return false, err
+			}
+			row[columnIdx] = types.NewDatum(bj)
+			return true, nil
+		}, nil
 	}
 	return nil, nil
 }

pkg/executor/slow_query_sql_test.go

@@ -544,3 +544,54 @@ func TestStorageEnginesInSlowQuery(t *testing.T) {
 		"where query like 'select%tablesample%;'").
 		Check(testkit.Rows("1 0"))
 }
+
+func TestSessionConnectAttrsInSlowQuery(t *testing.T) {
+	originCfg := config.GetGlobalConfig()
+	newCfg := *originCfg
+	f, err := os.CreateTemp("", "tidb-slow-*.log")
+	require.NoError(t, err)
+	_, err = f.WriteString(`# Time: 2024-01-15T10:00:00.000000+08:00
+# Txn_start_ts: 123456789
+# User@Host: root[root] @ localhost [127.0.0.1]
+# Query_time: 0.5
+# Digest: 42a1c8aae6f133e934d4bf0147491709a8812ea05ff8819ec522780fe657b772
+# Is_internal: false
+# Succ: true
+# Session_connect_attrs: {"_client_name":"Go-MySQL-Driver","_os":"linux","app_name":"test_app"}
+select * from t;
+`)
+	require.NoError(t, err)
+	require.NoError(t, f.Close())
+	newCfg.Log.SlowQueryFile = f.Name()
+	config.StoreGlobalConfig(&newCfg)
+	defer func() {
+		config.StoreGlobalConfig(originCfg)
+		require.NoError(t, os.Remove(newCfg.Log.SlowQueryFile))
+	}()
+	require.NoError(t, logutil.InitLogger(newCfg.Log.ToLogConfig()))
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("set @@time_zone='+08:00'")
+	tk.MustExec(fmt.Sprintf("set @@tidb_slow_query_file='%v'", f.Name()))
+
+	// Verify Session_connect_attrs column is present and returns the correct JSON value.
+	rows := tk.MustQuery("select Session_connect_attrs from information_schema.slow_query " +
+		"where query = 'select * from t;'").Rows()
+	require.Len(t, rows, 1)
+	attrsStr := rows[0][0].(string)
+	require.Contains(t, attrsStr, `"_client_name"`)
+	require.Contains(t, attrsStr, `"Go-MySQL-Driver"`)
+	require.Contains(t, attrsStr, `"_os"`)
+	require.Contains(t, attrsStr, `"linux"`)
+	require.Contains(t, attrsStr, `"app_name"`)
+	require.Contains(t, attrsStr, `"test_app"`)
+
+	// Verify individual keys are accessible via JSON_EXTRACT.
+	tk.MustQuery("select JSON_EXTRACT(Session_connect_attrs, '$._client_name') from information_schema.slow_query " +
+		"where query = 'select * from t;'").
+		Check(testkit.Rows(`"Go-MySQL-Driver"`))
+	tk.MustQuery("select JSON_EXTRACT(Session_connect_attrs, '$.app_name') from information_schema.slow_query " +
+		"where query = 'select * from t;'").
+		Check(testkit.Rows(`"test_app"`))
+}

pkg/executor/slow_query_test.go

@@ -191,7 +191,7 @@ select * from t;`
 		`0.1,0.2,0.03,127.0.0.1:20160,0.05,0.6,0.8,0.0.0.0:20160,70724,23333,65536,0,0,0,30000,3000,10000,1000,500000,500005,300000,300005,0,0,,` +
 		`Cop_backoff_regionMiss_total_times: 200 Cop_backoff_regionMiss_total_time: 0.2 Cop_backoff_regionMiss_max_time: 0.2 Cop_backoff_regionMiss_max_addr: 127.0.0.1 Cop_backoff_regionMiss_avg_time: 0.2 Cop_backoff_regionMiss_p90_time: 0.2 Cop_backoff_rpcPD_total_times: 200 Cop_backoff_rpcPD_total_time: 0.2 Cop_backoff_rpcPD_max_time: 0.2 Cop_backoff_rpcPD_max_addr: 127.0.0.1 Cop_backoff_rpcPD_avg_time: 0.2 Cop_backoff_rpcPD_p90_time: 0.2 Cop_backoff_rpcTiKV_total_times: 200 Cop_backoff_rpcTiKV_total_time: 0.2 Cop_backoff_rpcTiKV_max_time: 0.2 Cop_backoff_rpcTiKV_max_addr: 127.0.0.1 Cop_backoff_rpcTiKV_avg_time: 0.2 Cop_backoff_rpcTiKV_p90_time: 0.2,` +
 		`0,0,1,0,1,1,0,default,2.158,2.123,0.05,0.01,0.021,1,1,,60e9378c746d9a2be1c791047e008967cf252eb6de9167ad3aa6098fa2d523f4,` +
-		`,update t set i = 1;,select * from t;`
+		`,update t set i = 1;,null,select * from t;`
 	require.Equal(t, expectRecordString, recordString)
 
 	// Issue 20928
@@ -214,7 +214,7 @@ select * from t;`
 		`0.1,0.2,0.03,127.0.0.1:20160,0.05,0.6,0.8,0.0.0.0:20160,70724,23333,65536,0,0,0,30000,3000,10000,1000,500000,500005,300000,300005,0,0,,` +
 		`Cop_backoff_regionMiss_total_times: 200 Cop_backoff_regionMiss_total_time: 0.2 Cop_backoff_regionMiss_max_time: 0.2 Cop_backoff_regionMiss_max_addr: 127.0.0.1 Cop_backoff_regionMiss_avg_time: 0.2 Cop_backoff_regionMiss_p90_time: 0.2 Cop_backoff_rpcPD_total_times: 200 Cop_backoff_rpcPD_total_time: 0.2 Cop_backoff_rpcPD_max_time: 0.2 Cop_backoff_rpcPD_max_addr: 127.0.0.1 Cop_backoff_rpcPD_avg_time: 0.2 Cop_backoff_rpcPD_p90_time: 0.2 Cop_backoff_rpcTiKV_total_times: 200 Cop_backoff_rpcTiKV_total_time: 0.2 Cop_backoff_rpcTiKV_max_time: 0.2 Cop_backoff_rpcTiKV_max_addr: 127.0.0.1 Cop_backoff_rpcTiKV_avg_time: 0.2 Cop_backoff_rpcTiKV_p90_time: 0.2,` +
 		`0,0,1,0,1,1,0,default,2.158,2.123,0.05,0.01,0.021,1,1,,60e9378c746d9a2be1c791047e008967cf252eb6de9167ad3aa6098fa2d523f4,` +
-		`,update t set i = 1;,select * from t;`
+		`,update t set i = 1;,null,select * from t;`
 	require.Equal(t, expectRecordString, recordString)
 
 	// fix sql contain '# ' bug
@@ -284,6 +284,62 @@ select * from t;
 	require.Equal(t, value, "[t:i: a]")
 }
 
+func TestParseSlowLogSessionConnectAttrs(t *testing.T) {
+	// Slow log entry that includes Session_connect_attrs JSON.
+	slowLogStr := `# Time: 2019-04-28T15:24:04.309074+08:00
+# Txn_start_ts: 405888132465033227
+# User@Host: root[root] @ localhost [127.0.0.1]
+# Query_time: 0.216905
+# Digest: 42a1c8aae6f133e934d4bf0147491709a8812ea05ff8819ec522780fe657b772
+# Is_internal: false
+# Succ: true
+# Session_connect_attrs: {"_client_name":"Go-MySQL-Driver","_os":"linux","app_name":"test_app"}
+# Prev_stmt: begin;
+select * from t;
+`
+	loc, err := time.LoadLocation("Asia/Shanghai")
+	require.NoError(t, err)
+	ctx := mock.NewContext()
+	ctx.ResetSessionAndStmtTimeZone(loc)
+
+	// Use the retriever directly (without initialize) to avoid reading
+	// from actual slow log files on disk, which can produce extra rows.
+	retriever, err := newSlowQueryRetriever()
+	require.NoError(t, err)
+	retriever.columnValueFactoryMap = make(map[string]slowQueryColumnValueFactory, len(retriever.outputCols))
+	for idx, col := range retriever.outputCols {
+		factory, err := getColumnValueFactoryByName(col.Name.O, idx)
+		require.NoError(t, err)
+		require.NotNil(t, factory, "column %s should have a factory", col.Name.O)
+		retriever.columnValueFactoryMap[col.Name.O] = factory
+	}
+
+	reader := bufio.NewReader(bytes.NewBufferString(slowLogStr))
+	rows, err := parseLog(retriever, ctx, reader)
+	require.NoError(t, err)
+	require.Len(t, rows, 1)
+
+	// Find the Session_connect_attrs column.
+	colIdx := -1
+	for i, col := range retriever.outputCols {
+		if col.Name.L == strings.ToLower(variable.SlowLogSessionConnectAttrs) {
+			colIdx = i
+			break
+		}
+	}
+	require.NotEqual(t, -1, colIdx, "Session_connect_attrs column should exist")
+
+	// Verify the parsed JSON contains the expected keys.
+	bj := rows[0][colIdx].GetMysqlJSON()
+	bjStr := bj.String()
+	require.Contains(t, bjStr, `"_client_name"`)
+	require.Contains(t, bjStr, `"Go-MySQL-Driver"`)
+	require.Contains(t, bjStr, `"_os"`)
+	require.Contains(t, bjStr, `"linux"`)
+	require.Contains(t, bjStr, `"app_name"`)
+	require.Contains(t, bjStr, `"test_app"`)
+}
+
 // It changes variable.MaxOfMaxAllowedPacket, so must be stayed in SerialSuite.
 func TestParseSlowLogFileSerial(t *testing.T) {
 	loc, err := time.LoadLocation("Asia/Shanghai")

pkg/infoschema/tables.go

@@ -983,6 +983,7 @@ var slowQueryCols = []columnInfo{
 	{name: variable.SlowLogPlanDigest, tp: mysql.TypeVarchar, size: 128},
 	{name: variable.SlowLogBinaryPlan, tp: mysql.TypeLongBlob, size: types.UnspecifiedLength},
 	{name: variable.SlowLogPrevStmt, tp: mysql.TypeLongBlob, size: types.UnspecifiedLength},
+	{name: variable.SlowLogSessionConnectAttrs, tp: mysql.TypeJSON, size: types.UnspecifiedLength},
 	{name: variable.SlowLogQuerySQLStr, tp: mysql.TypeLongBlob, size: types.UnspecifiedLength},
 }