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.

Sign up for GitHub

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

Jump to bottom

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

Merged
merged 1 commit into from
May 5, 2023
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
1 change: 1 addition & 0 deletions make/migrations/postgresql/0101_2.7.3_schema.up.sql
View file
Open in desktop
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
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -970,15 +970,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 @@ -1028,11 +1019,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
View file
Open in desktop
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
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,9 @@
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 @@
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
}

Check warning on line 98 in src/pkg/task/dao/task.go

View check run for this annotation

Codecov / codecov/patch

src/pkg/task/dao/task.go#L97-L98 

Added lines #L97 - L98 were not covered by tests

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
}

Check warning on line 108 in src/pkg/task/dao/task.go

View check run for this annotation

Codecov / codecov/patch

src/pkg/task/dao/task.go#L107-L108 

Added lines #L107 - L108 were not covered by tests

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
View file
Open in desktop
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
23 changes: 23 additions & 0 deletions src/pkg/task/mock_task_dao_test.go
View file
Open in desktop

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

23 changes: 23 additions & 0 deletions src/pkg/task/mock_task_manager_test.go
View file
Open in desktop

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
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,9 @@
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 @@ -232,6 +235,20 @@
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
}

Check warning on line 242 in src/pkg/task/task.go

View check run for this annotation

Codecov / codecov/patch

src/pkg/task/task.go#L241-L242 

Added lines #L241 - L242 were not covered by tests
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
View file
Open in desktop
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{})
}
23 changes: 23 additions & 0 deletions src/testing/pkg/task/manager.go
View file
Open in desktop

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