diff --git a/models/db/context.go b/models/db/context.go index 455f3d1c5de..911dbd1c6f7 100644 --- a/models/db/context.go +++ b/models/db/context.go @@ -98,19 +98,31 @@ type Committer interface { // halfCommitter is a wrapper of Committer. // It can be closed early, but can't be committed early, it is useful for reusing a transaction. type halfCommitter struct { - Committer + committer Committer + committed bool } -func (*halfCommitter) Commit() error { - // do nothing +func (c *halfCommitter) Commit() error { + c.committed = true + // should do nothing, and the parent committer will commit later return nil } +func (c *halfCommitter) Close() error { + if c.committed { + // it's "commit and close", should do nothing, and the parent committer will commit later + return nil + } + + // it's "rollback and close", let the parent committer rollback right now + return c.committer.Close() +} + // TxContext represents a transaction Context, // it will reuse the existing transaction in the parent context or create a new one. func TxContext(parentCtx context.Context) (*Context, Committer, error) { if sess, ok := inTransaction(parentCtx); ok { - return newContext(parentCtx, sess, true), &halfCommitter{Committer: sess}, nil + return newContext(parentCtx, sess, true), &halfCommitter{committer: sess}, nil } sess := x.NewSession() @@ -126,7 +138,12 @@ func TxContext(parentCtx context.Context) (*Context, Committer, error) { // this function will reuse it otherwise will create a new one and close it when finished. func WithTx(parentCtx context.Context, f func(ctx context.Context) error) error { if sess, ok := inTransaction(parentCtx); ok { - return f(newContext(parentCtx, sess, true)) + err := f(newContext(parentCtx, sess, true)) + if err != nil { + // rollback immediately, in case the caller ignores returned error and tries to commit the transaction. + _ = sess.Close() + } + return err } return txWithNoCheck(parentCtx, f) } diff --git a/models/db/context_committer_test.go b/models/db/context_committer_test.go new file mode 100644 index 00000000000..38e91f22edb --- /dev/null +++ b/models/db/context_committer_test.go @@ -0,0 +1,102 @@ +// Copyright 2022 The Gitea Authors. All rights reserved. +// SPDX-License-Identifier: MIT + +package db // it's not db_test, because this file is for testing the private type halfCommitter + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/assert" +) + +type MockCommitter struct { + wants []string + gots []string +} + +func NewMockCommitter(wants ...string) *MockCommitter { + return &MockCommitter{ + wants: wants, + } +} + +func (c *MockCommitter) Commit() error { + c.gots = append(c.gots, "commit") + return nil +} + +func (c *MockCommitter) Close() error { + c.gots = append(c.gots, "close") + return nil +} + +func (c *MockCommitter) Assert(t *testing.T) { + assert.Equal(t, c.wants, c.gots, "want operations %v, but got %v", c.wants, c.gots) +} + +func Test_halfCommitter(t *testing.T) { + /* + Do something like: + + ctx, committer, err := db.TxContext(db.DefaultContext) + if err != nil { + return nil + } + defer committer.Close() + + // ... + + if err != nil { + return nil + } + + // ... + + return committer.Commit() + */ + + testWithCommitter := func(committer Committer, f func(committer Committer) error) { + if err := f(&halfCommitter{committer: committer}); err == nil { + committer.Commit() + } + committer.Close() + } + + t.Run("commit and close", func(t *testing.T) { + mockCommitter := NewMockCommitter("commit", "close") + + testWithCommitter(mockCommitter, func(committer Committer) error { + defer committer.Close() + return committer.Commit() + }) + + mockCommitter.Assert(t) + }) + + t.Run("rollback and close", func(t *testing.T) { + mockCommitter := NewMockCommitter("close", "close") + + testWithCommitter(mockCommitter, func(committer Committer) error { + defer committer.Close() + if true { + return fmt.Errorf("error") + } + return committer.Commit() + }) + + mockCommitter.Assert(t) + }) + + t.Run("close and commit", func(t *testing.T) { + mockCommitter := NewMockCommitter("close", "close") + + testWithCommitter(mockCommitter, func(committer Committer) error { + committer.Close() + committer.Commit() + return fmt.Errorf("error") + }) + + mockCommitter.Assert(t) + }) +}