server,infoschema,util: Implement session_account_connect_attrs P_S t…

…able (pingcap#46431)

ref pingcap#44341

infoschema/BUILD.bazel

@@ -24,6 +24,7 @@ go_library(
         "//kv",
         "//meta",
         "//meta/autoid",
+        "//parser/auth",
         "//parser/charset",
         "//parser/model",
         "//parser/mysql",

infoschema/perfschema/const.go

@@ -17,6 +17,7 @@ package perfschema
 // perfSchemaTables is a shortcut to involve all table names.
 var perfSchemaTables = []string{
 	tableGlobalStatus,
+	tableSessionAccountConnectAttrs,
 	tableSessionConnectAttrs,
 	tableSessionStatus,
 	tableSetupActors,
@@ -557,3 +558,10 @@ const tableSessionConnectAttrs = "CREATE TABLE IF NOT EXISTS " + tableNameSessio
 	"ATTR_NAME varchar(32) COLLATE utf8mb4_bin NOT NULL," +
 	"ATTR_VALUE varchar(1024) COLLATE utf8mb4_bin DEFAULT NULL," +
 	"ORDINAL_POSITION int DEFAULT NULL);"
+
+// tableSessionAccountConnectAttrs contains the column name definitions for the table session_connect_attrs
+const tableSessionAccountConnectAttrs = "CREATE TABLE IF NOT EXISTS " + tableNameSessionAccountConnectAttrs + " (" +
+	"PROCESSLIST_ID bigint unsigned NOT NULL," +
+	"ATTR_NAME varchar(32) COLLATE utf8mb4_bin NOT NULL," +
+	"ATTR_VALUE varchar(1024) COLLATE utf8mb4_bin DEFAULT NULL," +
+	"ORDINAL_POSITION int DEFAULT NULL);"

infoschema/perfschema/tables.go

@@ -70,6 +70,7 @@ const (
 	tableNamePDProfileAllocs                 = "pd_profile_allocs"
 	tableNamePDProfileBlock                  = "pd_profile_block"
 	tableNamePDProfileGoroutines             = "pd_profile_goroutines"
+	tableNameSessionAccountConnectAttrs      = "session_account_connect_attrs"
 	tableNameSessionConnectAttrs             = "session_connect_attrs"
 	tableNameSessionVariables                = "session_variables"
 )
@@ -107,6 +108,7 @@ var tableIDMap = map[string]int64{
 	tableNamePDProfileGoroutines:             autoid.PerformanceSchemaDBID + 30,
 	tableNameSessionVariables:                autoid.PerformanceSchemaDBID + 31,
 	tableNameSessionConnectAttrs:             autoid.PerformanceSchemaDBID + 32,
+	tableNameSessionAccountConnectAttrs:      autoid.PerformanceSchemaDBID + 33,
 }
 
 // perfSchemaTable stands for the fake table all its data is in the memory.
@@ -254,7 +256,9 @@ func (vt *perfSchemaTable) getRows(ctx context.Context, sctx sessionctx.Context,
 	case tableNameSessionVariables:
 		fullRows, err = infoschema.GetDataFromSessionVariables(ctx, sctx)
 	case tableNameSessionConnectAttrs:
-		fullRows, err = infoschema.GetDataFromSessionConnectAttrs(sctx)
+		fullRows, err = infoschema.GetDataFromSessionConnectAttrs(sctx, false)
+	case tableNameSessionAccountConnectAttrs:
+		fullRows, err = infoschema.GetDataFromSessionConnectAttrs(sctx, true)
 	}
 	if err != nil {
 		return

infoschema/tables.go

@@ -39,6 +39,7 @@ import (
 	"github.com/pingcap/tidb/domain/infosync"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/meta/autoid"
+	"github.com/pingcap/tidb/parser/auth"
 	"github.com/pingcap/tidb/parser/charset"
 	"github.com/pingcap/tidb/parser/model"
 	"github.com/pingcap/tidb/parser/mysql"
@@ -2030,12 +2031,16 @@ func GetDataFromSessionVariables(ctx context.Context, sctx sessionctx.Context) (
 }
 
 // GetDataFromSessionConnectAttrs produces the rows for the session_connect_attrs table.
-func GetDataFromSessionConnectAttrs(sctx sessionctx.Context) ([][]types.Datum, error) {
+func GetDataFromSessionConnectAttrs(sctx sessionctx.Context, sameAccount bool) ([][]types.Datum, error) {
 	sm := sctx.GetSessionManager()
 	if sm == nil {
 		return nil, nil
 	}
-	allAttrs := sm.GetConAttrs()
+	var user *auth.UserIdentity
+	if sameAccount {
+		user = sctx.GetSessionVars().User
+	}
+	allAttrs := sm.GetConAttrs(user)
 	rows := make([][]types.Datum, 0, len(allAttrs)*10) // 10 Attributes per connection
 	for pid, attrs := range allAttrs {                 // Note: PID is not ordered.
 		// Sorts the attributes by key and gives ORDINAL_POSITION based on this. This is needed as we didn't store the

server/BUILD.bazel

@@ -137,7 +137,7 @@ go_test(
     data = glob(["testdata/**"]),
     embed = [":server"],
     flaky = True,
-    shard_count = 47,
+    shard_count = 48,
     deps = [
         "//config",
         "//domain",

server/server.go

@@ -56,6 +56,7 @@ import (
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/metrics"
 	"github.com/pingcap/tidb/parser/ast"
+	"github.com/pingcap/tidb/parser/auth"
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/parser/terror"
 	"github.com/pingcap/tidb/planner/core"
@@ -834,11 +835,19 @@ func (s *Server) GetProcessInfo(id uint64) (*util.ProcessInfo, bool) {
 }
 
 // GetConAttrs returns the connection attributes
-func (s *Server) GetConAttrs() map[uint64]map[string]string {
+func (s *Server) GetConAttrs(user *auth.UserIdentity) map[uint64]map[string]string {
 	s.rwlock.RLock()
 	defer s.rwlock.RUnlock()
 	rs := make(map[uint64]map[string]string)
 	for _, client := range s.clients {
+		if user != nil {
+			if user.Username != client.user {
+				continue
+			}
+			if user.Hostname != client.peerHost {
+				continue
+			}
+		}
 		if pi := client.ctx.ShowProcess(); pi != nil {
 			rs[pi.ID] = client.attrs
 		}

server/server_test.go

@@ -24,6 +24,7 @@ import (
 	"testing"
 
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/tidb/parser/auth"
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/server/internal"
 	"github.com/pingcap/tidb/server/internal/testutil"
@@ -150,3 +151,45 @@ func TestIssue46197(t *testing.T) {
 	path := testdata.ConvertRowsToStrings(tk.MustQuery("select @@tidb_last_plan_replayer_token").Rows())
 	require.NoError(t, os.Remove(filepath.Join(replayer.GetPlanReplayerDirName(), path[0])))
 }
+
+func TestGetConAttrs(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tidbdrv := NewTiDBDriver(store)
+	cfg := util.NewTestConfig()
+	cfg.Port, cfg.Status.StatusPort = 0, 0
+	cfg.Status.ReportStatus = false
+	server, err := NewServer(cfg, tidbdrv)
+	require.NoError(t, err)
+
+	cc := &clientConn{
+		server:     server,
+		alloc:      arena.NewAllocator(1024),
+		chunkAlloc: chunk.NewAllocator(),
+		pkt:        internal.NewPacketIOForTest(bufio.NewWriter(bytes.NewBuffer(nil))),
+		attrs: map[string]string{
+			"_client_name": "tidb_test",
+		},
+		user:     "userA",
+		peerHost: "foo.example.com",
+	}
+	cc.SetCtx(&TiDBContext{Session: tk.Session(), stmts: make(map[int]*TiDBStatement)})
+	server.registerConn(cc)
+
+	// Get attributes for all clients
+	attrs := server.GetConAttrs(nil)
+	require.Equal(t, attrs[1]["_client_name"], "tidb_test")
+
+	// Get attributes for userA@foo.example.com, which should have results for connID 1.
+	userA := &auth.UserIdentity{Username: "userA", Hostname: "foo.example.com"}
+	attrs = server.GetConAttrs(userA)
+	require.Equal(t, attrs[1]["_client_name"], "tidb_test")
+	_, hasClientName := attrs[1]
+	require.True(t, hasClientName)
+
+	// Get attributes for userB@foo.example.com, which should NOT have results for connID 1.
+	userB := &auth.UserIdentity{Username: "userB", Hostname: "foo.example.com"}
+	attrs = server.GetConAttrs(userB)
+	_, hasClientName = attrs[1]
+	require.False(t, hasClientName)
+}

testkit/BUILD.bazel

@@ -20,6 +20,7 @@ go_library(
         "//expression",
         "//kv",
         "//parser/ast",
+        "//parser/auth",
         "//parser/terror",
         "//planner/core",
         "//resourcemanager",

testkit/mocksessionmanager.go

@@ -20,6 +20,7 @@ import (
 
 	"github.com/pingcap/tidb/domain"
 	"github.com/pingcap/tidb/parser/ast"
+	"github.com/pingcap/tidb/parser/auth"
 	"github.com/pingcap/tidb/planner/core"
 	"github.com/pingcap/tidb/session"
 	"github.com/pingcap/tidb/session/txninfo"
@@ -105,7 +106,7 @@ func (msm *MockSessionManager) GetProcessInfo(id uint64) (*util.ProcessInfo, boo
 }
 
 // GetConAttrs returns the connection attributes of all connections
-func (msm *MockSessionManager) GetConAttrs() map[uint64]map[string]string {
+func (msm *MockSessionManager) GetConAttrs(user *auth.UserIdentity) map[uint64]map[string]string {
 	return msm.ConAttrs
 }
 

util/BUILD.bazel

@@ -27,6 +27,7 @@ go_library(
         "//kv",
         "//metrics",
         "//parser",
+        "//parser/auth",
         "//parser/model",
         "//parser/mysql",
         "//parser/terror",

util/processinfo.go

@@ -21,6 +21,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/pingcap/tidb/parser/auth"
 	"github.com/pingcap/tidb/parser/mysql"
 	"github.com/pingcap/tidb/session/txninfo"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
@@ -204,5 +205,5 @@ type SessionManager interface {
 	// KillNonFlashbackClusterConn kill all non flashback cluster connections.
 	KillNonFlashbackClusterConn()
 	// GetConAttrs gets the connection attributes
-	GetConAttrs() map[uint64]map[string]string
+	GetConAttrs(user *auth.UserIdentity) map[uint64]map[string]string
 }