diff --git a/backend/v3/storage/database/query.go b/backend/v3/storage/database/query.go index f25927ade..3b594cdd7 100644 --- a/backend/v3/storage/database/query.go +++ b/backend/v3/storage/database/query.go @@ -73,6 +73,13 @@ func WithPermissionCheck(permission string) QueryOption { } } +// WithResultLock locks the results of the query during the transaction. +func WithResultLock() QueryOption { + return func(opts *QueryOpts) { + opts.ShouldLock = true + } +} + type joinType string const ( @@ -119,6 +126,8 @@ type QueryOpts struct { // Permission required to read or write the resource. // When unset, no permission check is made. Permission string + // ShouldLock the results during the transaction. + ShouldLock bool } // Matches implements [gomock.Matcher]. @@ -184,6 +193,7 @@ func (opts *QueryOpts) Write(builder *StatementBuilder) { opts.WriteOrderBy(builder) opts.WriteLimit(builder) opts.WriteOffset(builder) + opts.WriteLock(builder) } func (opts *QueryOpts) WriteCondition(builder *StatementBuilder) { @@ -247,3 +257,10 @@ func (opts *QueryOpts) WriteLeftJoins(builder *StatementBuilder) { join.columns.Write(builder) } } + +func (opts *QueryOpts) WriteLock(builder *StatementBuilder) { + if !opts.ShouldLock { + return + } + builder.WriteString(" FOR UPDATE") +} diff --git a/backend/v3/storage/database/query_test.go b/backend/v3/storage/database/query_test.go index 5e8e1a583..005b64c90 100644 --- a/backend/v3/storage/database/query_test.go +++ b/backend/v3/storage/database/query_test.go @@ -114,6 +114,16 @@ func TestQueryOptions(t *testing.T) { args: nil, }, }, + { + name: "with lock", + options: []QueryOption{ + WithResultLock(), + }, + want: want{ + stmt: " FOR UPDATE", + args: nil, + }, + }, { name: "multiple options", options: []QueryOption{ @@ -122,9 +132,10 @@ func TestQueryOptions(t *testing.T) { WithOrderByDescending(NewColumn("table", "column")), WithLimit(10), WithOffset(5), + WithResultLock(), }, want: want{ - stmt: " LEFT JOIN other_table ON table.id = other_table.table_id WHERE table.column = $1 ORDER BY table.column DESC LIMIT $2 OFFSET $3", + stmt: " LEFT JOIN other_table ON table.id = other_table.table_id WHERE table.column = $1 ORDER BY table.column DESC LIMIT $2 OFFSET $3 FOR UPDATE", args: []any{123, uint32(10), uint32(5)}, }, },