Skip to content

Commit ee89261

Browse files
ZackarySantanaZackarySantana
authored
DEVPROD-53 Add context to db queries (Remove, Find, FindAll, UpdateAll, and UpdateId) (#8821)
1 parent d11ac41 commit ee89261

File tree

116 files changed

+479
-485
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

116 files changed

+479
-485
lines changed

‎db/db_utils.go

+6-25
Original file line numberDiff line numberDiff line change
@@ -199,17 +199,7 @@ func EnsureIndex(collection string, index mongo.IndexModel) error {
199199
}
200200

201201
// Remove removes one item matching the query from the specified collection.
202-
func Remove(collection string, query any) error {
203-
session, db, err := GetGlobalSessionFactory().GetSession()
204-
if err != nil {
205-
return err
206-
}
207-
defer session.Close()
208-
209-
return db.C(collection).Remove(query)
210-
}
211-
212-
func RemoveContext(ctx context.Context, collection string, query any) error {
202+
func Remove(ctx context.Context, collection string, query any) error {
213203
session, db, err := GetGlobalSessionFactory().GetContextSession(ctx)
214204
if err != nil {
215205
return err
@@ -220,7 +210,7 @@ func RemoveContext(ctx context.Context, collection string, query any) error {
220210
}
221211

222212
// RemoveAll removes all items matching the query from the specified collection.
223-
func RemoveAll(collection string, query any) error {
213+
func RemoveAll(ctx context.Context, collection string, query any) error {
224214
session, db, err := GetGlobalSessionFactory().GetSession()
225215
if err != nil {
226216
return err
@@ -232,6 +222,8 @@ func RemoveAll(collection string, query any) error {
232222
}
233223

234224
// Update updates one matching document in the collection.
225+
// DEPRECATED (DEVPROD-15398): This is only here to support a cache
226+
// with Gimlet, use UpdateContext instead.
235227
func Update(collection string, query any, update any) error {
236228
session, db, err := GetGlobalSessionFactory().GetSession()
237229
if err != nil {
@@ -326,19 +318,6 @@ func UpdateAllContext(ctx context.Context, collection string, query any, update
326318
return &db.ChangeInfo{Updated: int(res.ModifiedCount)}, nil
327319
}
328320

329-
// UpdateId updates one _id-matching document in the collection.
330-
func UpdateId(collection string, id, update any) error {
331-
session, db, err := GetGlobalSessionFactory().GetSession()
332-
if err != nil {
333-
grip.Errorf("error establishing db connection: %+v", err)
334-
335-
return err
336-
}
337-
defer session.Close()
338-
339-
return db.C(collection).UpdateId(id, update)
340-
}
341-
342321
// UpdateIdContext updates one _id-matching document in the collection.
343322
func UpdateIdContext(ctx context.Context, collection string, id, update any) error {
344323
res, err := evergreen.GetEnvironment().DB().Collection(collection).UpdateOne(ctx,
@@ -356,6 +335,8 @@ func UpdateIdContext(ctx context.Context, collection string, id, update any) err
356335
}
357336

358337
// UpdateAll updates all matching documents in the collection.
338+
// DEPRECATED (DEVPROD-15398): This is only here to support a cache
339+
// with Gimlet, use UpdateAllContext instead.
359340
func UpdateAll(collection string, query any, update any) (*db.ChangeInfo, error) {
360341
switch query.(type) {
361342
case *Q, Q:

‎db/db_utils_test.go

+6-5
Original file line numberDiff line numberDiff line change
@@ -111,7 +111,7 @@ func TestDBUtils(t *testing.T) {
111111
So(count, ShouldEqual, 2)
112112

113113
// remove just the first
114-
So(Remove(collection, bson.M{"field_one": "1"}),
114+
So(Remove(t.Context(), collection, bson.M{"field_one": "1"}),
115115
ShouldBeNil)
116116
count, err = Count(collection, bson.M{})
117117
So(err, ShouldBeNil)
@@ -151,7 +151,7 @@ func TestDBUtils(t *testing.T) {
151151
So(count, ShouldEqual, 3)
152152

153153
// remove just the first
154-
So(RemoveAll(collection, bson.M{"field_one": "1"}),
154+
So(RemoveAll(t.Context(), collection, bson.M{"field_one": "1"}),
155155
ShouldBeNil)
156156
count, err = Count(collection, bson.M{})
157157
So(err, ShouldBeNil)
@@ -205,7 +205,7 @@ func TestDBUtils(t *testing.T) {
205205
// one and limit to one (meaning only the second struct should be
206206
// returned)
207207
out := []insertableStruct{}
208-
err = FindAllQ(collection, Query(bson.M{"field_two": 1}).
208+
err = FindAllQContext(t.Context(), collection, Query(bson.M{"field_two": 1}).
209209
Project(bson.M{"field_three": 0}).
210210
Sort([]string{"-field_one"}).
211211
Limit(1).
@@ -288,7 +288,8 @@ func TestDBUtils(t *testing.T) {
288288
So(count, ShouldEqual, 3)
289289

290290
// update the first and third
291-
_, err = UpdateAll(
291+
_, err = UpdateAllContext(
292+
t.Context(),
292293
collection,
293294
bson.M{
294295
"field_one": "1",
@@ -302,7 +303,7 @@ func TestDBUtils(t *testing.T) {
302303
So(err, ShouldBeNil)
303304

304305
out := []insertableStruct{}
305-
err = FindAllQ(collection, Query(bson.M{"field_two": 3}), &out)
306+
err = FindAllQContext(t.Context(), collection, Query(bson.M{"field_two": 3}), &out)
306307
So(err, ShouldBeNil)
307308
So(len(out), ShouldEqual, 2)
308309

‎db/query.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -150,8 +150,8 @@ func CountQContext(ctx context.Context, collection string, q Q) (int, error) {
150150
}
151151

152152
// RemoveAllQ removes all docs that satisfy the query
153-
func RemoveAllQ(collection string, q Q) error {
154-
return Remove(collection, q.filter)
153+
func RemoveAllQ(ctx context.Context, collection string, q Q) error {
154+
return Remove(ctx, collection, q.filter)
155155
}
156156

157157
// implement custom marshaller interfaces to prevent the bug where we

‎db/query_test.go

+4-4
Original file line numberDiff line numberDiff line change
@@ -56,14 +56,14 @@ func TestQueryExecution(t *testing.T) {
5656

5757
Convey("BelowFive should return 4 documents", func() {
5858
out := []insertableStruct{}
59-
err := FindAllQ(collection, BelowFive, &out)
59+
err := FindAllQContext(t.Context(), collection, BelowFive, &out)
6060
So(err, ShouldBeNil)
6161
So(len(out), ShouldEqual, 4)
6262
})
6363

6464
Convey("BelowFiveSorted should return 4 documents, sorted in reverse", func() {
6565
out := []insertableStruct{}
66-
err := FindAllQ(collection, BelowFiveSorted, &out)
66+
err := FindAllQContext(t.Context(), collection, BelowFiveSorted, &out)
6767
So(err, ShouldBeNil)
6868
So(len(out), ShouldEqual, 4)
6969
So(out[0].FieldTwo, ShouldEqual, 4)
@@ -74,14 +74,14 @@ func TestQueryExecution(t *testing.T) {
7474

7575
Convey("BelowFiveLimit should return 2 documents", func() {
7676
out := []insertableStruct{}
77-
err := FindAllQ(collection, BelowFiveLimit, &out)
77+
err := FindAllQContext(t.Context(), collection, BelowFiveLimit, &out)
7878
So(err, ShouldBeNil)
7979
So(len(out), ShouldEqual, 2)
8080
})
8181

8282
Convey("JustOneField should return 1 document", func() {
8383
out := []bson.M{}
84-
err := FindAllQ(collection, JustOneField, &out)
84+
err := FindAllQContext(t.Context(), collection, JustOneField, &out)
8585
So(err, ShouldBeNil)
8686
So(out[0]["three"], ShouldEqual, "COOL")
8787
})

‎graphql/annotation_resolver.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ func (r *annotationResolver) WebhookConfigured(ctx context.Context, obj *restMod
2020
if t == nil {
2121
return false, ResourceNotFound.Send(ctx, fmt.Sprintf("task '%s' not found", taskID))
2222
}
23-
_, ok, _ := model.IsWebhookConfigured(t.Project, t.Version)
23+
_, ok, _ := model.IsWebhookConfigured(ctx, t.Project, t.Version)
2424
return ok, nil
2525
}
2626

‎graphql/mutation_resolver.go

+5-5
Original file line numberDiff line numberDiff line change
@@ -346,7 +346,7 @@ func (r *mutationResolver) SchedulePatch(ctx context.Context, patchID string, co
346346
// AttachProjectToNewRepo is the resolver for the attachProjectToNewRepo field.
347347
func (r *mutationResolver) AttachProjectToNewRepo(ctx context.Context, project MoveProjectInput) (*restModel.APIProjectRef, error) {
348348
usr := mustHaveUser(ctx)
349-
pRef, err := data.FindProjectById(project.ProjectID, false, false)
349+
pRef, err := data.FindProjectById(ctx, project.ProjectID, false, false)
350350
if err != nil {
351351
return nil, InternalServerError.Send(ctx, fmt.Sprintf("fetching project '%s': %s", project.ProjectID, err.Error()))
352352
}
@@ -370,7 +370,7 @@ func (r *mutationResolver) AttachProjectToNewRepo(ctx context.Context, project M
370370
// AttachProjectToRepo is the resolver for the attachProjectToRepo field.
371371
func (r *mutationResolver) AttachProjectToRepo(ctx context.Context, projectID string) (*restModel.APIProjectRef, error) {
372372
usr := mustHaveUser(ctx)
373-
pRef, err := data.FindProjectById(projectID, false, false)
373+
pRef, err := data.FindProjectById(ctx, projectID, false, false)
374374
if err != nil {
375375
return nil, InternalServerError.Send(ctx, fmt.Sprintf("fetching project '%s': %s", projectID, err.Error()))
376376
}
@@ -519,7 +519,7 @@ func (r *mutationResolver) DeleteProject(ctx context.Context, projectID string)
519519
// DetachProjectFromRepo is the resolver for the detachProjectFromRepo field.
520520
func (r *mutationResolver) DetachProjectFromRepo(ctx context.Context, projectID string) (*restModel.APIProjectRef, error) {
521521
usr := mustHaveUser(ctx)
522-
pRef, err := data.FindProjectById(projectID, false, false)
522+
pRef, err := data.FindProjectById(ctx, projectID, false, false)
523523
if err != nil {
524524
return nil, InternalServerError.Send(ctx, fmt.Sprintf("fetching project '%s': %s", projectID, err.Error()))
525525
}
@@ -597,7 +597,7 @@ func (r *mutationResolver) SetLastRevision(ctx context.Context, opts SetLastRevi
597597
return nil, InternalServerError.Send(ctx, fmt.Sprintf("updating last revision for '%s': %s", opts.ProjectIdentifier, err.Error()))
598598
}
599599

600-
if err = project.SetRepotrackerError(&model.RepositoryErrorDetails{}); err != nil {
600+
if err = project.SetRepotrackerError(ctx, &model.RepositoryErrorDetails{}); err != nil {
601601
return nil, InternalServerError.Send(ctx, fmt.Sprintf("clearing repotracker error for '%s': %s", opts.ProjectIdentifier, err.Error()))
602602
}
603603

@@ -1188,7 +1188,7 @@ func (r *mutationResolver) SaveSubscription(ctx context.Context, subscription re
11881188
return false, ResourceNotFound.Send(ctx, fmt.Sprintf("version '%s' not found", id))
11891189
}
11901190
case "project":
1191-
p, projectErr := data.FindProjectById(id, false, false)
1191+
p, projectErr := data.FindProjectById(ctx, id, false, false)
11921192
if projectErr != nil {
11931193
return false, InternalServerError.Send(ctx, fmt.Sprintf("fetching project '%s': %s", id, projectErr.Error()))
11941194
}

‎graphql/patch_resolver.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,7 @@ func (r *patchResolver) GeneratedTaskCounts(ctx context.Context, obj *restModel.
129129
// PatchTriggerAliases is the resolver for the patchTriggerAliases field.
130130
func (r *patchResolver) PatchTriggerAliases(ctx context.Context, obj *restModel.APIPatch) ([]*restModel.APIPatchTriggerDefinition, error) {
131131
projectID := utility.FromStringPtr(obj.ProjectId)
132-
projectRef, err := data.FindProjectById(projectID, true, true)
132+
projectRef, err := data.FindProjectById(ctx, projectID, true, true)
133133
if err != nil {
134134
return nil, InternalServerError.Send(ctx, fmt.Sprintf("fetching project '%s': %s", projectID, err.Error()))
135135
}
@@ -153,7 +153,7 @@ func (r *patchResolver) PatchTriggerAliases(ctx context.Context, obj *restModel.
153153
projectCache[alias.ChildProject] = project
154154
}
155155

156-
matchingTasks, err := project.VariantTasksForSelectors([]patch.PatchTriggerDefinition{alias}, utility.FromStringPtr(obj.Requester))
156+
matchingTasks, err := project.VariantTasksForSelectors(ctx, []patch.PatchTriggerDefinition{alias}, utility.FromStringPtr(obj.Requester))
157157
if err != nil {
158158
return nil, InternalServerError.Send(ctx, fmt.Sprintf("matching tasks to definitions for alias '%s': %s", alias.Alias, err.Error()))
159159
}

‎graphql/query_resolver.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -405,7 +405,7 @@ func (r *queryResolver) Patch(ctx context.Context, patchID string) (*restModel.A
405405

406406
// GithubProjectConflicts is the resolver for the githubProjectConflicts field.
407407
func (r *queryResolver) GithubProjectConflicts(ctx context.Context, projectID string) (*model.GithubProjectConflicts, error) {
408-
pRef, err := model.FindMergedProjectRef(projectID, "", false)
408+
pRef, err := model.FindMergedProjectRef(ctx, projectID, "", false)
409409
if err != nil {
410410
return nil, InternalServerError.Send(ctx, fmt.Sprintf("fetching project '%s': %s", projectID, err.Error()))
411411
}
@@ -422,7 +422,7 @@ func (r *queryResolver) GithubProjectConflicts(ctx context.Context, projectID st
422422

423423
// Project is the resolver for the project field.
424424
func (r *queryResolver) Project(ctx context.Context, projectIdentifier string) (*restModel.APIProjectRef, error) {
425-
project, err := data.FindProjectById(projectIdentifier, true, false)
425+
project, err := data.FindProjectById(ctx, projectIdentifier, true, false)
426426
if err != nil {
427427
return nil, InternalServerError.Send(ctx, fmt.Sprintf("fetching project '%s': %s", projectIdentifier, err.Error()))
428428
}

‎graphql/task_resolver.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -453,7 +453,7 @@ func (r *taskResolver) IsPerfPluginEnabled(ctx context.Context, obj *restModel.A
453453
if !evergreen.IsFinishedTaskStatus(utility.FromStringPtr(obj.Status)) {
454454
return false, nil
455455
}
456-
if !model.IsPerfEnabledForProject(utility.FromStringPtr(obj.ProjectId)) {
456+
if !model.IsPerfEnabledForProject(ctx, utility.FromStringPtr(obj.ProjectId)) {
457457
return false, nil
458458
}
459459
opts := apimodels.GetPerfCountOptions{
@@ -527,7 +527,7 @@ func (r *taskResolver) Pod(ctx context.Context, obj *restModel.APITask) (*restMo
527527
// Project is the resolver for the project field.
528528
func (r *taskResolver) Project(ctx context.Context, obj *restModel.APITask) (*restModel.APIProjectRef, error) {
529529
projectID := utility.FromStringPtr(obj.ProjectId)
530-
pRef, err := data.FindProjectById(projectID, true, false)
530+
pRef, err := data.FindProjectById(ctx, projectID, true, false)
531531
if err != nil {
532532
return nil, InternalServerError.Send(ctx, fmt.Sprintf("fetching project '%s': %s", projectID, err.Error()))
533533
}

‎graphql/util.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ const (
4848

4949
// getGroupedFiles returns the files of a Task inside a GroupedFile struct
5050
func getGroupedFiles(ctx context.Context, name string, taskID string, execution int) (*GroupedFiles, error) {
51-
taskFiles, err := artifact.GetAllArtifacts([]artifact.TaskIDAndExecution{{TaskID: taskID, Execution: execution}})
51+
taskFiles, err := artifact.GetAllArtifacts(ctx, []artifact.TaskIDAndExecution{{TaskID: taskID, Execution: execution}})
5252
if err != nil {
5353
return nil, ResourceNotFound.Send(ctx, err.Error())
5454
}
@@ -882,7 +882,7 @@ func getHostRequestOptions(ctx context.Context, usr *user.DBUser, spawnHostInput
882882
}
883883

884884
func getProjectMetadata(ctx context.Context, projectId *string, patchId *string) (*restModel.APIProjectRef, error) {
885-
projectRef, err := model.FindMergedProjectRef(*projectId, *patchId, false)
885+
projectRef, err := model.FindMergedProjectRef(ctx, *projectId, *patchId, false)
886886
if err != nil {
887887
return nil, InternalServerError.Send(ctx, fmt.Sprintf("finding merged project ref for project '%s': %s", utility.FromStringPtr(projectId), err.Error()))
888888
}

‎graphql/version_resolver.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ func (r *versionResolver) ChildVersions(ctx context.Context, obj *restModel.APIV
140140
// ExternalLinksForMetadata is the resolver for the externalLinksForMetadata field.
141141
func (r *versionResolver) ExternalLinksForMetadata(ctx context.Context, obj *restModel.APIVersion) ([]*ExternalLinkForMetadata, error) {
142142
projectID := utility.FromStringPtr(obj.Project)
143-
pRef, err := data.FindProjectById(projectID, false, false)
143+
pRef, err := data.FindProjectById(ctx, projectID, false, false)
144144
if err != nil {
145145
return nil, InternalServerError.Send(ctx, fmt.Sprintf("fetching project '%s': %s", projectID, err.Error()))
146146
}

‎model/alertrecord/alert_bookkeeping_test.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -175,7 +175,7 @@ func (s *alertRecordSuite) TestFindOneWithUnsetIDQuery() {
175175
s.Equal(2, rec.RevisionOrderNumber)
176176

177177
records := []AlertRecord{}
178-
err = db.FindAllQ(Collection, ByLastFailureTransition(legacyAlertsSubscription, "task", "variant", "project").Limit(999), &records)
178+
err = db.FindAllQContext(s.T().Context(), Collection, ByLastFailureTransition(legacyAlertsSubscription, "task", "variant", "project").Limit(999), &records)
179179
s.NoError(err)
180180
s.Len(records, 3)
181181
}

‎model/annotations/db.go

+6-6
Original file line numberDiff line numberDiff line change
@@ -44,9 +44,9 @@ func FindOne(ctx context.Context, query db.Q) (*TaskAnnotation, error) {
4444
}
4545

4646
// Find gets every TaskAnnotation matching the given query.
47-
func Find(query db.Q) ([]TaskAnnotation, error) {
47+
func Find(ctx context.Context, query db.Q) ([]TaskAnnotation, error) {
4848
annotations := []TaskAnnotation{}
49-
err := db.FindAllQ(Collection, query, &annotations)
49+
err := db.FindAllQContext(ctx, Collection, query, &annotations)
5050
if err != nil {
5151
return nil, errors.Wrap(err, "finding task annotations")
5252
}
@@ -58,12 +58,12 @@ func FindOneByTaskIdAndExecution(ctx context.Context, id string, execution int)
5858
return FindOne(ctx, db.Query(ByTaskIdAndExecution(id, execution)))
5959
}
6060

61-
func FindByTaskIds(ids []string) ([]TaskAnnotation, error) {
62-
return Find(ByTaskIds(ids))
61+
func FindByTaskIds(ctx context.Context, ids []string) ([]TaskAnnotation, error) {
62+
return Find(ctx, ByTaskIds(ids))
6363
}
6464

65-
func FindByTaskId(id string) ([]TaskAnnotation, error) {
66-
return Find(db.Query(ByTaskId(id)))
65+
func FindByTaskId(ctx context.Context, id string) ([]TaskAnnotation, error) {
66+
return Find(ctx, db.Query(ByTaskId(id)))
6767
}
6868

6969
// Upsert writes the task_annotation to the database.

‎model/annotations/task_annotations_test.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ func TestGetLatestExecutions(t *testing.T) {
3535
assert.NoError(t, a.Upsert())
3636
}
3737

38-
taskAnnotations, err := FindByTaskIds([]string{"t1", "t2"})
38+
taskAnnotations, err := FindByTaskIds(t.Context(), []string{"t1", "t2"})
3939
assert.NoError(t, err)
4040
assert.Len(t, taskAnnotations, 3)
4141
assert.Len(t, GetLatestExecutions(taskAnnotations), 2)

‎model/artifact/artifact_file.go

+3-3
Original file line numberDiff line numberDiff line change
@@ -129,8 +129,8 @@ func presignFile(ctx context.Context, file File) (string, error) {
129129
return pail.PreSign(ctx, requestParams)
130130
}
131131

132-
func GetAllArtifacts(tasks []TaskIDAndExecution) ([]File, error) {
133-
artifacts, err := FindAll(ByTaskIdsAndExecutions(tasks))
132+
func GetAllArtifacts(ctx context.Context, tasks []TaskIDAndExecution) ([]File, error) {
133+
artifacts, err := FindAll(ctx, ByTaskIdsAndExecutions(tasks))
134134
if err != nil {
135135
return nil, errors.Wrap(err, "finding artifact files for task")
136136
}
@@ -139,7 +139,7 @@ func GetAllArtifacts(tasks []TaskIDAndExecution) ([]File, error) {
139139
for _, t := range tasks {
140140
taskIds = append(taskIds, t.TaskID)
141141
}
142-
artifacts, err = FindAll(ByTaskIds(taskIds))
142+
artifacts, err = FindAll(ctx, ByTaskIds(taskIds))
143143
if err != nil {
144144
return nil, errors.Wrap(err, "finding artifact files for task without execution number")
145145
}

0 commit comments

Comments
 (0)