From f271d2cd2c02feb48307208ecdbb2fc2a34af028 Mon Sep 17 00:00:00 2001 From: Timmy Date: Mon, 23 Oct 2023 10:15:35 +0800 Subject: [PATCH 1/2] feat: add subject cache batch get (#271) --- pkg/abac/pap/group.go | 25 +++++++--- pkg/abac/pap/types.go | 1 + pkg/cacheimpls/init.go | 4 +- pkg/cacheimpls/local_subject.go | 43 +++++++++++++++++ pkg/cacheimpls/local_subject_test.go | 71 ++++++++++++++++++++++++++++ pkg/database/dao/mock/subject.go | 15 ++++++ pkg/database/dao/subject.go | 24 ++++++++++ pkg/database/dao/subject_test.go | 29 ++++++++++++ pkg/service/mock/subject.go | 15 ++++++ pkg/service/subject.go | 13 +++++ pkg/service/subject_test.go | 46 ++++++++++++++++++ pkg/service/types/subject.go | 1 + 12 files changed, 278 insertions(+), 9 deletions(-) diff --git a/pkg/abac/pap/group.go b/pkg/abac/pap/group.go index 067f8923..d915e2dd 100644 --- a/pkg/abac/pap/group.go +++ b/pkg/abac/pap/group.go @@ -797,21 +797,32 @@ func convertToSubjectGroups(svcSubjectGroups []types.SubjectGroup) ([]SubjectGro } func convertToGroupMembers(svcGroupMembers []types.GroupMember) ([]GroupMember, error) { - members := make([]GroupMember, 0, len(svcGroupMembers)) + subjectPKs := make([]int64, 0, len(svcGroupMembers)) for _, m := range svcGroupMembers { - subject, err := cacheimpls.GetSubjectByPK(m.SubjectPK) - if err != nil { - if errors.Is(err, sql.ErrNoRows) { - continue - } + subjectPKs = append(subjectPKs, m.SubjectPK) + } + subjects, err := cacheimpls.BatchGetSubjectByPKs(subjectPKs) + if err != nil { + return nil, err + } - return nil, err + subjectMap := make(map[int64]types.Subject, len(subjects)) + for _, subject := range subjects { + subjectMap[subject.PK] = subject + } + + members := make([]GroupMember, 0, len(svcGroupMembers)) + for _, m := range svcGroupMembers { + subject, ok := subjectMap[m.SubjectPK] + if !ok { + continue } members = append(members, GroupMember{ PK: m.PK, Type: subject.Type, ID: subject.ID, + Name: subject.Name, ExpiredAt: m.ExpiredAt, CreatedAt: m.CreatedAt, }) diff --git a/pkg/abac/pap/types.go b/pkg/abac/pap/types.go index 05fba8b7..7339124b 100644 --- a/pkg/abac/pap/types.go +++ b/pkg/abac/pap/types.go @@ -30,6 +30,7 @@ type GroupMember struct { PK int64 `json:"pk"` Type string `json:"type"` ID string `json:"id"` + Name string `json:"name"` ExpiredAt int64 `json:"expired_at"` CreatedAt time.Time `json:"created_at"` } diff --git a/pkg/cacheimpls/init.go b/pkg/cacheimpls/init.go index 67d8c678..6fc6bb65 100644 --- a/pkg/cacheimpls/init.go +++ b/pkg/cacheimpls/init.go @@ -103,8 +103,8 @@ func InitCaches(disabled bool) { "local_subject", disabled, retrieveSubject, - 1*time.Minute, - newRandomDuration(30), + 30*time.Minute, + newRandomDuration(600), ) // 影响: job查询cmdb的资源进行鉴权 diff --git a/pkg/cacheimpls/local_subject.go b/pkg/cacheimpls/local_subject.go index ef903def..58ed1823 100644 --- a/pkg/cacheimpls/local_subject.go +++ b/pkg/cacheimpls/local_subject.go @@ -12,6 +12,7 @@ package cacheimpls import ( "errors" + "fmt" "github.com/TencentBlueKing/gopkg/cache" @@ -46,3 +47,45 @@ func GetSubjectByPK(pk int64) (subject types.Subject, err error) { return } + +// BatchGetSubjectByPKs ... +func BatchGetSubjectByPKs(pks []int64) (subjects []types.Subject, err error) { + subjects = make([]types.Subject, 0, len(pks)) + missPKs := make([]int64, 0, len(pks)) + for _, pk := range pks { + key := SubjectPKCacheKey{ + PK: pk, + } + value, ok := LocalSubjectCache.DirectGet(key) + if !ok { + missPKs = append(missPKs, pk) + continue + } + + subject, ok := value.(types.Subject) + if !ok { + err = fmt.Errorf("not types.Subject in cache for pk=%d", pk) + return + } + subjects = append(subjects, subject) + } + + if len(missPKs) > 0 { + svc := service.NewSubjectService() + missSubjects, err := svc.ListByPKs(missPKs) + if err != nil { + return nil, err + } + + for _, subject := range missSubjects { + key := SubjectPKCacheKey{ + PK: subject.PK, + } + LocalSubjectCache.Set(key, subject) + } + + subjects = append(subjects, missSubjects...) + } + + return subjects, nil +} diff --git a/pkg/cacheimpls/local_subject_test.go b/pkg/cacheimpls/local_subject_test.go index 0e1c18f4..2eb19735 100644 --- a/pkg/cacheimpls/local_subject_test.go +++ b/pkg/cacheimpls/local_subject_test.go @@ -17,8 +17,12 @@ import ( "github.com/TencentBlueKing/gopkg/cache" "github.com/TencentBlueKing/gopkg/cache/memory" + "github.com/agiledragon/gomonkey/v2" + "github.com/golang/mock/gomock" "github.com/stretchr/testify/assert" + "iam/pkg/service" + "iam/pkg/service/mock" svctypes "iam/pkg/service/types" ) @@ -47,3 +51,70 @@ func TestGetSubjectByPK(t *testing.T) { _, err = GetSubjectByPK(1) assert.Error(t, err) } + +func TestBatchGetSubjectByPKsFail(t *testing.T) { + ctl := gomock.NewController(t) + mockSvc := mock.NewMockSubjectService(ctl) + mockSvc.EXPECT().ListByPKs([]int64{1}).Return(nil, errors.New("error here")) + + patches := gomonkey.ApplyFunc(service.NewSubjectService, + func() service.SubjectService { + return mockSvc + }) + defer patches.Reset() + + expiration := 5 * time.Minute + + // valid + retrieveFunc := func(key cache.Key) (interface{}, error) { + return svctypes.Subject{}, nil + } + mockCache := memory.NewCache( + "mockCache", false, retrieveFunc, expiration, nil) + LocalSubjectCache = mockCache + + _, err := BatchGetSubjectByPKs([]int64{1}) + assert.Error(t, err) +} + +func TestBatchGetSubjectByPKsOK(t *testing.T) { + ctl := gomock.NewController(t) + mockSvc := mock.NewMockSubjectService(ctl) + mockSvc.EXPECT().ListByPKs([]int64{2, 3}).Return([]svctypes.Subject{ + { + PK: 3, + Type: "department", + ID: "department", + Name: "department", + }, + }, nil) + + patches := gomonkey.ApplyFunc(service.NewSubjectService, + func() service.SubjectService { + return mockSvc + }) + defer patches.Reset() + + expiration := 5 * time.Minute + + // valid + retrieveFunc := func(key cache.Key) (interface{}, error) { + return svctypes.Subject{}, nil + } + mockCache := memory.NewCache( + "mockCache", false, retrieveFunc, expiration, nil) + LocalSubjectCache = mockCache + + LocalSubjectCache.Set(SubjectPKCacheKey{ + PK: 1, + }, svctypes.Subject{ + PK: 1, + Type: "user", + ID: "admin", + Name: "admin", + }) + + subjects, err := BatchGetSubjectByPKs([]int64{1, 2, 3}) + assert.Nil(t, err) + assert.Equal(t, 2, len(subjects)) +} diff --git a/pkg/database/dao/mock/subject.go b/pkg/database/dao/mock/subject.go index d5c164d7..10c6bbe0 100644 --- a/pkg/database/dao/mock/subject.go +++ b/pkg/database/dao/mock/subject.go @@ -137,6 +137,21 @@ func (mr *MockSubjectManagerMockRecorder) ListByIDs(_type, ids interface{}) *gom return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListByIDs", reflect.TypeOf((*MockSubjectManager)(nil).ListByIDs), _type, ids) } +// ListByPKs mocks base method. +func (m *MockSubjectManager) ListByPKs(pks []int64) ([]dao.Subject, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListByPKs", pks) + ret0, _ := ret[0].([]dao.Subject) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListByPKs indicates an expected call of ListByPKs. +func (mr *MockSubjectManagerMockRecorder) ListByPKs(pks interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListByPKs", reflect.TypeOf((*MockSubjectManager)(nil).ListByPKs), pks) +} + // ListPaging mocks base method. func (m *MockSubjectManager) ListPaging(_type string, limit, offset int64) ([]dao.Subject, error) { m.ctrl.T.Helper() diff --git a/pkg/database/dao/subject.go b/pkg/database/dao/subject.go index 563618ab..1c508ef8 100644 --- a/pkg/database/dao/subject.go +++ b/pkg/database/dao/subject.go @@ -34,6 +34,7 @@ type Subject struct { type SubjectManager interface { Get(pk int64) (Subject, error) GetPK(_type, id string) (int64, error) + ListByPKs(pks []int64) (subjects []Subject, err error) ListByIDs(_type string, ids []string) ([]Subject, error) ListPaging(_type string, limit, offset int64) ([]Subject, error) GetCount(_type string) (int64, error) @@ -60,6 +61,18 @@ func (m *subjectManager) Get(pk int64) (subject Subject, err error) { return } +// ListByPKs ... +func (m *subjectManager) ListByPKs(pks []int64) (subjects []Subject, err error) { + if len(pks) == 0 { + return + } + err = m.selectByPKs(&subjects, pks) + if errors.Is(err, sql.ErrNoRows) { + return subjects, nil + } + return +} + // GetPK ... func (m *subjectManager) GetPK(_type, id string) (int64, error) { var pk int64 @@ -141,6 +154,17 @@ func (m *subjectManager) selectPK(pk *int64, _type string, id string) error { return database.SqlxGet(m.DB, pk, query, _type, id) } +func (m *subjectManager) selectByPKs(subjects *[]Subject, pks []int64) error { + query := `SELECT + pk, + type, + id, + name + FROM subject + WHERE pk IN (?)` + return database.SqlxSelect(m.DB, subjects, query, pks) +} + func (m *subjectManager) selectSubjectsByIDs(subjects *[]Subject, _type string, ids []string) error { query := `SELECT pk, diff --git a/pkg/database/dao/subject_test.go b/pkg/database/dao/subject_test.go index 28f05d0b..88f748ae 100644 --- a/pkg/database/dao/subject_test.go +++ b/pkg/database/dao/subject_test.go @@ -89,3 +89,32 @@ func Test_subjectManager_BulkUpdate(t *testing.T) { assert.NoError(t, err, "query from db fail.") }) } + +func TestListByPKs(t *testing.T) { + database.RunWithMock(t, func(db *sqlx.DB, mock sqlmock.Sqlmock, t *testing.T) { + mockData := []interface{}{ + Subject{ + PK: 1, + Type: "group", + ID: "1", + Name: "group1", + }, + Subject{ + PK: 3, + Type: "group", + ID: "3", + Name: "group3", + }, + } + + mockQuery := `^SELECT pk, type, id, name FROM subject WHERE pk IN` + mockRows := database.NewMockRows(mock, mockData...) + mock.ExpectQuery(mockQuery).WithArgs(1, 3).WillReturnRows(mockRows) + + manager := &subjectManager{DB: db} + subject, err := manager.ListByPKs([]int64{1, 3}) + + assert.NoError(t, err, "query from db fail.") + assert.Equal(t, len(subject), 2) + }) +} diff --git a/pkg/service/mock/subject.go b/pkg/service/mock/subject.go index fed01997..f20ab56c 100644 --- a/pkg/service/mock/subject.go +++ b/pkg/service/mock/subject.go @@ -122,6 +122,21 @@ func (mr *MockSubjectServiceMockRecorder) GetPK(_type, id interface{}) *gomock.C return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPK", reflect.TypeOf((*MockSubjectService)(nil).GetPK), _type, id) } +// ListByPKs mocks base method. +func (m *MockSubjectService) ListByPKs(pks []int64) ([]types.Subject, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListByPKs", pks) + ret0, _ := ret[0].([]types.Subject) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListByPKs indicates an expected call of ListByPKs. +func (mr *MockSubjectServiceMockRecorder) ListByPKs(pks interface{}) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListByPKs", reflect.TypeOf((*MockSubjectService)(nil).ListByPKs), pks) +} + // ListPKsBySubjects mocks base method. func (m *MockSubjectService) ListPKsBySubjects(subjects []types.Subject) ([]int64, error) { m.ctrl.T.Helper() diff --git a/pkg/service/subject.go b/pkg/service/subject.go index 70008c32..51aef46a 100644 --- a/pkg/service/subject.go +++ b/pkg/service/subject.go @@ -31,6 +31,7 @@ type SubjectService interface { // web api GetCount(_type string) (int64, error) + ListByPKs(pks []int64) ([]types.Subject, error) ListPaging(_type string, limit, offset int64) ([]types.Subject, error) ListPKsBySubjects(subjects []types.Subject) ([]int64, error) BulkCreate(subjects []types.Subject) error @@ -61,6 +62,7 @@ func (l *subjectService) Get(pk int64) (subject types.Subject, err error) { } subject = types.Subject{ + PK: s.PK, Type: s.Type, ID: s.ID, Name: s.Name, @@ -68,6 +70,16 @@ func (l *subjectService) Get(pk int64) (subject types.Subject, err error) { return } +// ListByPKs ... +func (l *subjectService) ListByPKs(pks []int64) ([]types.Subject, error) { + daoSubjects, err := l.manager.ListByPKs(pks) + if err != nil { + return nil, errorx.Wrapf(err, SubjectSVC, + "ListByPKs", "manager.ListByPKs pks=`%v`", pks) + } + return convertToSubjects(daoSubjects), nil +} + // GetPK ... func (l *subjectService) GetPK(_type, id string) (pk int64, err error) { return l.manager.GetPK(_type, id) @@ -82,6 +94,7 @@ func convertToSubjects(daoSubjects []dao.Subject) []types.Subject { subjects := make([]types.Subject, 0, len(daoSubjects)) for _, s := range daoSubjects { subjects = append(subjects, types.Subject{ + PK: s.PK, Type: s.Type, ID: s.ID, Name: s.Name, diff --git a/pkg/service/subject_test.go b/pkg/service/subject_test.go index fdce46ab..b3eb9ddc 100644 --- a/pkg/service/subject_test.go +++ b/pkg/service/subject_test.go @@ -153,6 +153,52 @@ var _ = Describe("SubjectService", func() { }) }) + Describe("ListByPKs", func() { + var ctl *gomock.Controller + BeforeEach(func() { + ctl = gomock.NewController(GinkgoT()) + }) + AfterEach(func() { + ctl.Finish() + }) + + It("ListByPKs fail", func() { + mockSubjectService := mock.NewMockSubjectManager(ctl) + mockSubjectService.EXPECT().ListByPKs([]int64{1, 3}).Return( + []dao.Subject{}, errors.New("list by pk fail"), + ).AnyTimes() + + manager := &subjectService{ + manager: mockSubjectService, + } + + _, err := manager.ListByPKs([]int64{1, 3}) + assert.Error(GinkgoT(), err) + assert.Contains(GinkgoT(), err.Error(), "ListByPKs") + }) + + It("ok", func() { + mockSubjectService := mock.NewMockSubjectManager(ctl) + mockSubjectService.EXPECT().ListByPKs([]int64{1, 3}).Return( + []dao.Subject{ + {PK: 1, Type: "user", ID: "test", Name: "test"}, + {PK: 3, Type: "department", ID: "test", Name: "test"}, + }, nil, + ).AnyTimes() + + manager := &subjectService{ + manager: mockSubjectService, + } + + subjects, err := manager.ListByPKs([]int64{1, 3}) + assert.NoError(GinkgoT(), err) + assert.Equal(GinkgoT(), subjects, []types.Subject{ + {PK: 1, Type: "user", ID: "test", Name: "test"}, + {PK: 3, Type: "department", ID: "test", Name: "test"}, + }) + }) + }) + Describe("BulkCreate", func() { var ctl *gomock.Controller BeforeEach(func() { diff --git a/pkg/service/types/subject.go b/pkg/service/types/subject.go index 624e8557..624fab22 100644 --- a/pkg/service/types/subject.go +++ b/pkg/service/types/subject.go @@ -14,6 +14,7 @@ import "time" // Subject ... type Subject struct { + PK int64 Type string `json:"type"` ID string `json:"id"` Name string `json:"name"` From 43afc46d6a6f8f6eb560cf817c35797a002d2865 Mon Sep 17 00:00:00 2001 From: Timmy Date: Wed, 1 Nov 2023 11:01:47 +0800 Subject: [PATCH 2/2] Update VERSION --- VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/VERSION b/VERSION index 166a50ff..393ccdb5 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -1.12.9 +1.12.10