diff --git a/infoschema/BUILD.bazel b/infoschema/BUILD.bazel index dba1ce4de6..0c25033d77 100644 --- a/infoschema/BUILD.bazel +++ b/infoschema/BUILD.bazel @@ -24,6 +24,7 @@ go_library( "//kv", "//meta", "//meta/autoid", + "//parser/auth", "//parser/charset", "//parser/model", "//parser/mysql", diff --git a/infoschema/perfschema/const.go b/infoschema/perfschema/const.go index 7ac56ad0f3..65fb0552d1 100644 --- a/infoschema/perfschema/const.go +++ b/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);" diff --git a/infoschema/perfschema/tables.go b/infoschema/perfschema/tables.go index 1774f0d9bc..045307b846 100644 --- a/infoschema/perfschema/tables.go +++ b/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 diff --git a/infoschema/tables.go b/infoschema/tables.go index ba9881f53b..a66cddf142 100644 --- a/infoschema/tables.go +++ b/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 diff --git a/server/BUILD.bazel b/server/BUILD.bazel index 932c2f6648..db313470a6 100644 --- a/server/BUILD.bazel +++ b/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", diff --git a/server/server.go b/server/server.go index e40db3bf24..2d85a3717f 100644 --- a/server/server.go +++ b/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 } diff --git a/server/server_test.go b/server/server_test.go index 4275e81dbf..07fd9dcb5e 100644 --- a/server/server_test.go +++ b/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) +} diff --git a/testkit/BUILD.bazel b/testkit/BUILD.bazel index b6ec22fe93..594fe2a329 100644 --- a/testkit/BUILD.bazel +++ b/testkit/BUILD.bazel @@ -20,6 +20,7 @@ go_library( "//expression", "//kv", "//parser/ast", + "//parser/auth", "//parser/terror", "//planner/core", "//resourcemanager", diff --git a/testkit/mocksessionmanager.go b/testkit/mocksessionmanager.go index 8184bdcf65..7a96588b2a 100644 --- a/testkit/mocksessionmanager.go +++ b/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 } diff --git a/util/BUILD.bazel b/util/BUILD.bazel index 532ff12d6d..972c76bc42 100644 --- a/util/BUILD.bazel +++ b/util/BUILD.bazel @@ -27,6 +27,7 @@ go_library( "//kv", "//metrics", "//parser", + "//parser/auth", "//parser/model", "//parser/mysql", "//parser/terror", diff --git a/util/processinfo.go b/util/processinfo.go index 0425e5dfbf..35aafe5d2d 100644 --- a/util/processinfo.go +++ b/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 }