Merge pull request #274 from TencentBlueKing/develop

1.12.10

VERSION

@@ -1 +1 @@
-1.12.9
+1.12.10

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,
 		})

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"`
 }

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的资源进行鉴权

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
+}

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))
+}

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,

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)
+	})
+}

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,13 +62,24 @@ 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,
 	}
 	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,