Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[cherry-pick] fix: improve the performance of list artifacts #18639

Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions make/migrations/postgresql/0101_2.7.3_schema.up.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
CREATE INDEX IF NOT EXISTS idx_task_extra_attrs_report_uuids ON task USING gin ((extra_attrs::jsonb->'report_uuids'));
1 change: 1 addition & 0 deletions make/migrations/postgresql/0111_2.8.1_schema.up.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
CREATE INDEX IF NOT EXISTS idx_task_extra_attrs_report_uuids ON task USING gin ((extra_attrs::jsonb->'report_uuids'));
14 changes: 3 additions & 11 deletions src/controller/scan/base_controller.go
Original file line number Diff line number Diff line change
Expand Up @@ -964,15 +964,6 @@ func (bc *basicController) launchScanJob(ctx context.Context, param *launchScanJ
reportUUIDsKey: reportUUIDs,
}

// NOTE: due to the limitation of the beego's orm, the List method of the task manager not support ?! operator for the jsonb field,
// we cann't list the tasks for scan reports of uuid1, uuid2 by SQL `SELECT * FROM task WHERE (extra_attrs->'report_uuids')::jsonb ?| array['uuid1', 'uuid2']`
// or by `SELECT * FROM task WHERE id IN (SELECT id FROM task WHERE (extra_attrs->'report_uuids')::jsonb ?| array['uuid1', 'uuid2'])`
// so save {"report:uuid1": "1", "report:uuid2": "2"} in the extra_attrs of the task, and then list it with
// SQL `SELECT * FROM task WHERE extra_attrs->>'report:uuid1' = '1'` in loop
for _, reportUUID := range reportUUIDs {
extraAttrs["report:"+reportUUID] = "1"
}

_, err = bc.taskMgr.Create(ctx, param.ExecutionID, j, extraAttrs)
return err
}
Expand Down Expand Up @@ -1022,11 +1013,12 @@ func (bc *basicController) listScanTasks(ctx context.Context, reportUUIDs []stri
}

func (bc *basicController) getScanTask(ctx context.Context, reportUUID string) (*task.Task, error) {
query := q.New(q.KeyWords{"extra_attrs." + "report:" + reportUUID: "1"})
tasks, err := bc.taskMgr.List(bc.cloneCtx(ctx), query)
// NOTE: the method uses the postgres' unique operations and should consider here if support other database in the future.
tasks, err := bc.taskMgr.ListScanTasksByReportUUID(ctx, reportUUID)
if err != nil {
return nil, err
}

if len(tasks) == 0 {
return nil, errors.NotFoundError(nil).WithMessage("task for report %s not found", reportUUID)
}
Expand Down
42 changes: 22 additions & 20 deletions src/controller/scan/base_controller_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -321,7 +321,7 @@ func (suite *ControllerTestSuite) TestScanControllerScan() {
walkFn(suite.artifact)
}).Once()

mock.OnAnything(suite.taskMgr, "List").Return([]*task.Task{
mock.OnAnything(suite.taskMgr, "ListScanTasksByReportUUID").Return([]*task.Task{
{ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-001"), Status: "Success"},
}, nil).Once()

Expand All @@ -343,7 +343,7 @@ func (suite *ControllerTestSuite) TestScanControllerScan() {
walkFn(suite.artifact)
}).Once()

mock.OnAnything(suite.taskMgr, "List").Return([]*task.Task{
mock.OnAnything(suite.taskMgr, "ListScanTasksByReportUUID").Return([]*task.Task{
{ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-001"), Status: "Success"},
}, nil).Once()

Expand All @@ -360,7 +360,7 @@ func (suite *ControllerTestSuite) TestScanControllerScan() {
walkFn(suite.artifact)
}).Once()

mock.OnAnything(suite.taskMgr, "List").Return([]*task.Task{
mock.OnAnything(suite.taskMgr, "ListScanTasksByReportUUID").Return([]*task.Task{
{ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-001"), Status: "Running"},
}, nil).Once()

Expand Down Expand Up @@ -409,37 +409,40 @@ func (suite *ControllerTestSuite) TestScanControllerStop() {

// TestScanControllerGetReport ...
func (suite *ControllerTestSuite) TestScanControllerGetReport() {
ctx := orm.NewContext(nil, &ormtesting.FakeOrmer{})
mock.OnAnything(suite.ar, "Walk").Return(nil).Run(func(args mock.Arguments) {
walkFn := args.Get(2).(func(*artifact.Artifact) error)
walkFn(suite.artifact)
}).Once()

mock.OnAnything(suite.taskMgr, "List").Return([]*task.Task{
mock.OnAnything(suite.taskMgr, "ListScanTasksByReportUUID").Return([]*task.Task{
{ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-001")},
}, nil).Once()
mock.OnAnything(suite.accessoryMgr, "List").Return(nil, nil)
rep, err := suite.c.GetReport(context.TODO(), suite.artifact, []string{v1.MimeTypeNativeReport})
rep, err := suite.c.GetReport(ctx, suite.artifact, []string{v1.MimeTypeNativeReport})
require.NoError(suite.T(), err)
assert.Equal(suite.T(), 1, len(rep))
}

// TestScanControllerGetSummary ...
func (suite *ControllerTestSuite) TestScanControllerGetSummary() {
ctx := orm.NewContext(nil, &ormtesting.FakeOrmer{})
mock.OnAnything(suite.accessoryMgr, "List").Return([]accessoryModel.Accessory{}, nil).Once()
mock.OnAnything(suite.ar, "Walk").Return(nil).Run(func(args mock.Arguments) {
walkFn := args.Get(2).(func(*artifact.Artifact) error)
walkFn(suite.artifact)
}).Once()
mock.OnAnything(suite.taskMgr, "List").Return(nil, nil).Once()
mock.OnAnything(suite.taskMgr, "ListScanTasksByReportUUID").Return(nil, nil).Once()

sum, err := suite.c.GetSummary(context.TODO(), suite.artifact, []string{v1.MimeTypeNativeReport})
sum, err := suite.c.GetSummary(ctx, suite.artifact, []string{v1.MimeTypeNativeReport})
require.NoError(suite.T(), err)
assert.Equal(suite.T(), 1, len(sum))
}

// TestScanControllerGetScanLog ...
func (suite *ControllerTestSuite) TestScanControllerGetScanLog() {
mock.OnAnything(suite.taskMgr, "List").Return([]*task.Task{
ctx := orm.NewContext(nil, &ormtesting.FakeOrmer{})
mock.OnAnything(suite.taskMgr, "ListScanTasksByReportUUID").Return([]*task.Task{
{
ID: 1,
ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-001"),
Expand All @@ -448,7 +451,7 @@ func (suite *ControllerTestSuite) TestScanControllerGetScanLog() {

mock.OnAnything(suite.taskMgr, "GetLog").Return([]byte("log"), nil).Once()

bytes, err := suite.c.GetScanLog(context.TODO(), &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, "rp-uuid-001")
bytes, err := suite.c.GetScanLog(ctx, &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, "rp-uuid-001")
require.NoError(suite.T(), err)
assert.Condition(suite.T(), func() (success bool) {
success = len(bytes) > 0
Expand All @@ -457,8 +460,8 @@ func (suite *ControllerTestSuite) TestScanControllerGetScanLog() {
}

func (suite *ControllerTestSuite) TestScanControllerGetMultiScanLog() {
kw1 := q.KeyWords{"extra_attrs.report:rp-uuid-001": "1"}
suite.taskMgr.On("List", context.TODO(), q.New(kw1)).Return([]*task.Task{
ctx := orm.NewContext(nil, &ormtesting.FakeOrmer{})
suite.taskMgr.On("ListScanTasksByReportUUID", ctx, "rp-uuid-001").Return([]*task.Task{
{
ID: 1,
ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-001"),
Expand All @@ -469,8 +472,7 @@ func (suite *ControllerTestSuite) TestScanControllerGetMultiScanLog() {
walkFn(suite.artifact)
})
mock.OnAnything(suite.accessoryMgr, "List").Return(nil, nil)
kw2 := q.KeyWords{"extra_attrs.report:rp-uuid-002": "1"}
suite.taskMgr.On("List", context.TODO(), q.New(kw2)).Return([]*task.Task{
suite.taskMgr.On("ListScanTasksByReportUUID", ctx, "rp-uuid-002").Return([]*task.Task{
{
ID: 2,
ExtraAttrs: suite.makeExtraAttrs(int64(1), "rp-uuid-002"),
Expand All @@ -480,7 +482,7 @@ func (suite *ControllerTestSuite) TestScanControllerGetMultiScanLog() {
// Both success
mock.OnAnything(suite.taskMgr, "GetLog").Return([]byte("log"), nil).Twice()

bytes, err := suite.c.GetScanLog(context.TODO(), &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
bytes, err := suite.c.GetScanLog(ctx, &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
suite.Nil(err)
suite.NotEmpty(bytes)
suite.Contains(string(bytes), "Logs of report rp-uuid-001")
Expand All @@ -489,10 +491,10 @@ func (suite *ControllerTestSuite) TestScanControllerGetMultiScanLog() {

{
// One successfully, one failed
suite.taskMgr.On("GetLog", context.TODO(), int64(1)).Return([]byte("log"), nil).Once()
suite.taskMgr.On("GetLog", context.TODO(), int64(2)).Return(nil, fmt.Errorf("failed")).Once()
suite.taskMgr.On("GetLog", ctx, int64(1)).Return([]byte("log"), nil).Once()
suite.taskMgr.On("GetLog", ctx, int64(2)).Return(nil, fmt.Errorf("failed")).Once()

bytes, err := suite.c.GetScanLog(context.TODO(), &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
bytes, err := suite.c.GetScanLog(ctx, &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
suite.Nil(err)
suite.NotEmpty(bytes)
suite.NotContains(string(bytes), "Logs of report rp-uuid-001")
Expand All @@ -502,7 +504,7 @@ func (suite *ControllerTestSuite) TestScanControllerGetMultiScanLog() {
// Both failed
mock.OnAnything(suite.taskMgr, "GetLog").Return(nil, fmt.Errorf("failed")).Twice()

bytes, err := suite.c.GetScanLog(context.TODO(), &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
bytes, err := suite.c.GetScanLog(ctx, &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
suite.Error(err)
suite.Empty(bytes)
}
Expand All @@ -511,7 +513,7 @@ func (suite *ControllerTestSuite) TestScanControllerGetMultiScanLog() {
// Both empty
mock.OnAnything(suite.taskMgr, "GetLog").Return(nil, nil).Twice()

bytes, err := suite.c.GetScanLog(context.TODO(), &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
bytes, err := suite.c.GetScanLog(ctx, &artifact.Artifact{Artifact: art.Artifact{ID: 1, ProjectID: 1}}, base64.StdEncoding.EncodeToString([]byte("rp-uuid-001|rp-uuid-002")))
suite.Nil(err)
suite.Empty(bytes)
}
Expand Down Expand Up @@ -560,7 +562,7 @@ func (suite *ControllerTestSuite) TestScanAll() {
walkFn(suite.artifact)
}).Once()

mock.OnAnything(suite.taskMgr, "List").Return(nil, nil).Once()
mock.OnAnything(suite.taskMgr, "ListScanTasksByReportUUID").Return(nil, nil).Once()

mock.OnAnything(suite.reportMgr, "Delete").Return(nil).Once()
mock.OnAnything(suite.reportMgr, "Create").Return("uuid", nil).Once()
Expand Down
22 changes: 22 additions & 0 deletions src/pkg/task/dao/task.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,9 @@ type TaskDAO interface {
UpdateStatusInBatch(ctx context.Context, jobIDs []string, status string, batchSize int) (err error)
// ExecutionIDsByVendorAndStatus retrieve the execution id by vendor status
ExecutionIDsByVendorAndStatus(ctx context.Context, vendorType, status string) ([]int64, error)
// ListScanTasksByReportUUID lists scan tasks by report uuid, although it's a specific case but it will be
// more suitable to support multi database in the future.
ListScanTasksByReportUUID(ctx context.Context, uuid string) (tasks []*Task, err error)
}

// NewTaskDAO returns an instance of TaskDAO
Expand Down Expand Up @@ -88,6 +91,25 @@ func (t *taskDAO) List(ctx context.Context, query *q.Query) ([]*Task, error) {
return tasks, nil
}

func (t *taskDAO) ListScanTasksByReportUUID(ctx context.Context, uuid string) ([]*Task, error) {
ormer, err := orm.FromContext(ctx)
if err != nil {
return nil, err
}

tasks := []*Task{}
// Due to the limitation of the beego's orm, the SQL cannot be converted by orm framework,
// so we can only execute the query by raw SQL, the SQL filters the task contains the report uuid in the column extra_attrs,
// consider from performance side which can using indexes to speed up queries.
sql := fmt.Sprintf(`SELECT * FROM task WHERE extra_attrs::jsonb->'report_uuids' @> '["%s"]'`, uuid)
_, err = ormer.Raw(sql).QueryRows(&tasks)
if err != nil {
return nil, err
}

return tasks, nil
}

func (t *taskDAO) Get(ctx context.Context, id int64) (*Task, error) {
task := &Task{
ID: id,
Expand Down
21 changes: 21 additions & 0 deletions src/pkg/task/dao/task_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,27 @@ func (t *taskDAOTestSuite) TestList() {
t.Require().Len(tasks, 0)
}

func (t *taskDAOTestSuite) TestListScanTasksByReportUUID() {
// should not exist if non set
tasks, err := t.taskDAO.ListScanTasksByReportUUID(t.ctx, "fake-report-uuid")
t.Require().Nil(err)
t.Require().Len(tasks, 0)
// create one with report uuid
taskID, err := t.taskDAO.Create(t.ctx, &Task{
ExecutionID: t.executionID,
Status: "success",
StatusCode: 1,
ExtraAttrs: `{"report_uuids": ["fake-report-uuid"]}`,
})
t.Require().Nil(err)
defer t.taskDAO.Delete(t.ctx, taskID)
// should exist as created
tasks, err = t.taskDAO.ListScanTasksByReportUUID(t.ctx, "fake-report-uuid")
t.Require().Nil(err)
t.Require().Len(tasks, 1)
t.Equal(taskID, tasks[0].ID)
}

func (t *taskDAOTestSuite) TestGet() {
// not exist
_, err := t.taskDAO.Get(t.ctx, 10000)
Expand Down
26 changes: 26 additions & 0 deletions src/pkg/task/mock_task_dao_test.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

26 changes: 26 additions & 0 deletions src/pkg/task/mock_task_manager_test.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

17 changes: 17 additions & 0 deletions src/pkg/task/task.go
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,9 @@ type Manager interface {
UpdateStatusInBatch(ctx context.Context, jobIDs []string, status string, batchSize int) error
// ExecutionIDsByVendorAndStatus retrieve execution id by vendor type and status
ExecutionIDsByVendorAndStatus(ctx context.Context, vendorType, status string) ([]int64, error)
// ListScanTasksByReportUUID lists scan tasks by report uuid, although it's a specific case but it will be
// more suitable to support multi database in the future.
ListScanTasksByReportUUID(ctx context.Context, uuid string) (tasks []*Task, err error)
}

// NewManager creates an instance of the default task manager
Expand Down Expand Up @@ -234,6 +237,20 @@ func (m *manager) List(ctx context.Context, query *q.Query) ([]*Task, error) {
return ts, nil
}

func (m *manager) ListScanTasksByReportUUID(ctx context.Context, uuid string) ([]*Task, error) {
tasks, err := m.dao.ListScanTasksByReportUUID(ctx, uuid)
if err != nil {
return nil, err
}
var ts []*Task
for _, task := range tasks {
t := &Task{}
t.From(task)
ts = append(ts, t)
}
return ts, nil
}

func (m *manager) UpdateExtraAttrs(ctx context.Context, id int64, extraAttrs map[string]interface{}) error {
data, err := json.Marshal(extraAttrs)
if err != nil {
Expand Down
13 changes: 13 additions & 0 deletions src/pkg/task/task_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,19 @@ func (t *taskManagerTestSuite) TestList() {
t.dao.AssertExpectations(t.T())
}

func (t *taskManagerTestSuite) TestListScanTasksByReportUUID() {
t.dao.On("ListScanTasksByReportUUID", mock.Anything, mock.Anything).Return([]*dao.Task{
{
ID: 1,
},
}, nil)
tasks, err := t.mgr.ListScanTasksByReportUUID(nil, "uuid")
t.Require().Nil(err)
t.Require().Len(tasks, 1)
t.Equal(int64(1), tasks[0].ID)
t.dao.AssertExpectations(t.T())
}

func TestTaskManagerTestSuite(t *testing.T) {
suite.Run(t, &taskManagerTestSuite{})
}
Loading