Adjust transaction handling via db.Context (#20031)

4 files changed, 29 insertions, 42 deletions

diff --git a/models/db/context.go b/models/db/context.go
index c41d761802..783033c0e0 100644
--- a/models/db/context.go
+++ b/models/db/context.go
@@ -23,23 +23,30 @@ type contextKey struct {
 	name string
 }
 
-// EnginedContextKey is a context key. It is used with context.Value() to get the current Engined for the context
-var EnginedContextKey = &contextKey{"engined"}
+// enginedContextKey is a context key. It is used with context.Value() to get the current Engined for the context
+var enginedContextKey = &contextKey{"engined"}
+var _ Engined = &Context{}
 
 // Context represents a db context
 type Context struct {
 	context.Context
-	e Engine
+	e           Engine
+	transaction bool
 }
 
-// WithEngine returns a Context from a context.Context and Engine
-func WithEngine(ctx context.Context, e Engine) *Context {
+func newContext(ctx context.Context, e Engine, transaction bool) *Context {
 	return &Context{
-		Context: ctx,
-		e:       e.Context(ctx),
+		Context:     ctx,
+		e:           e,
+		transaction: transaction,
 	}
 }
 
+// InTransaction if context is in a transaction
+func (ctx *Context) InTransaction() bool {
+	return ctx.transaction
+}
+
 // Engine returns db engine
 func (ctx *Context) Engine() Engine {
 	return ctx.e
@@ -47,7 +54,7 @@ func (ctx *Context) Engine() Engine {
 
 // Value shadows Value for context.Context but allows us to get ourselves and an Engined object
 func (ctx *Context) Value(key interface{}) interface{} {
-	if key == EnginedContextKey {
+	if key == enginedContextKey {
 		return ctx
 	}
 	return ctx.Context.Value(key)
@@ -55,7 +62,7 @@ func (ctx *Context) Value(key interface{}) interface{} {
 
 // WithContext returns this engine tied to this context
 func (ctx *Context) WithContext(other context.Context) *Context {
-	return WithEngine(other, ctx.e)
+	return newContext(ctx, ctx.e.Context(other), ctx.transaction)
 }
 
 // Engined structs provide an Engine
@@ -68,7 +75,7 @@ func GetEngine(ctx context.Context) Engine {
 	if engined, ok := ctx.(Engined); ok {
 		return engined.Engine()
 	}
-	enginedInterface := ctx.Value(EnginedContextKey)
+	enginedInterface := ctx.Value(enginedContextKey)
 	if enginedInterface != nil {
 		return enginedInterface.(Engined).Engine()
 	}
@@ -89,18 +96,7 @@ func TxContext() (*Context, Committer, error) {
 		return nil, nil, err
 	}
 
-	return &Context{
-		Context: DefaultContext,
-		e:       sess,
-	}, sess, nil
-}
-
-// WithContext represents executing database operations
-func WithContext(f func(ctx *Context) error) error {
-	return f(&Context{
-		Context: DefaultContext,
-		e:       x,
-	})
+	return newContext(DefaultContext, sess, true), sess, nil
 }
 
 // WithTx represents executing database operations on a transaction
@@ -118,10 +114,7 @@ func WithTx(f func(ctx context.Context) error, stdCtx ...context.Context) error
 		return err
 	}
 
-	if err := f(&Context{
-		Context: parentCtx,
-		e:       sess,
-	}); err != nil {
+	if err := f(newContext(parentCtx, sess, true)); err != nil {
 		return err
 	}
 
diff --git a/models/organization/org.go b/models/organization/org.go
index 0d4a5e337b..044ea06563 100644
--- a/models/organization/org.go
+++ b/models/organization/org.go
@@ -680,7 +680,7 @@ type accessibleReposEnv struct {
 	user    *user_model.User
 	team    *Team
 	teamIDs []int64
-	e       db.Engine
+	ctx     context.Context
 	keyword string
 	orderBy db.SearchOrderBy
 }
@@ -706,7 +706,7 @@ func AccessibleReposEnv(ctx context.Context, org *Organization, userID int64) (A
 		org:     org,
 		user:    user,
 		teamIDs: teamIDs,
-		e:       db.GetEngine(ctx),
+		ctx:     ctx,
 		orderBy: db.SearchOrderByRecentUpdated,
 	}, nil
 }
@@ -717,7 +717,7 @@ func (org *Organization) AccessibleTeamReposEnv(team *Team) AccessibleReposEnvir
 	return &accessibleReposEnv{
 		org:     org,
 		team:    team,
-		e:       db.GetEngine(db.DefaultContext),
+		ctx:     db.DefaultContext,
 		orderBy: db.SearchOrderByRecentUpdated,
 	}
 }
@@ -744,7 +744,7 @@ func (env *accessibleReposEnv) cond() builder.Cond {
 }
 
 func (env *accessibleReposEnv) CountRepos() (int64, error) {
-	repoCount, err := env.e.
+	repoCount, err := db.GetEngine(env.ctx).
 		Join("INNER", "team_repo", "`team_repo`.repo_id=`repository`.id").
 		Where(env.cond()).
 		Distinct("`repository`.id").
@@ -761,7 +761,7 @@ func (env *accessibleReposEnv) RepoIDs(page, pageSize int) ([]int64, error) {
 	}
 
 	repoIDs := make([]int64, 0, pageSize)
-	return repoIDs, env.e.
+	return repoIDs, db.GetEngine(env.ctx).
 		Table("repository").
 		Join("INNER", "team_repo", "`team_repo`.repo_id=`repository`.id").
 		Where(env.cond()).
@@ -783,7 +783,7 @@ func (env *accessibleReposEnv) Repos(page, pageSize int) ([]*repo_model.Reposito
 		return repos, nil
 	}
 
-	return repos, env.e.
+	return repos, db.GetEngine(env.ctx).
 		In("`repository`.id", repoIDs).
 		OrderBy(string(env.orderBy)).
 		Find(&repos)
@@ -791,7 +791,7 @@ func (env *accessibleReposEnv) Repos(page, pageSize int) ([]*repo_model.Reposito
 
 func (env *accessibleReposEnv) MirrorRepoIDs() ([]int64, error) {
 	repoIDs := make([]int64, 0, 10)
-	return repoIDs, env.e.
+	return repoIDs, db.GetEngine(env.ctx).
 		Table("repository").
 		Join("INNER", "team_repo", "`team_repo`.repo_id=`repository`.id AND `repository`.is_mirror=?", true).
 		Where(env.cond()).
@@ -812,7 +812,7 @@ func (env *accessibleReposEnv) MirrorRepos() ([]*repo_model.Repository, error) {
 		return repos, nil
 	}
 
-	return repos, env.e.
+	return repos, db.GetEngine(env.ctx).
 		In("`repository`.id", repoIDs).
 		Find(&repos)
 }
diff --git a/models/repo/repo.go b/models/repo/repo.go
index 8ef405a665..bb2a7468ff 100644
--- a/models/repo/repo.go
+++ b/models/repo/repo.go
@@ -319,13 +319,7 @@ func (repo *Repository) LoadUnits(ctx context.Context) (err error) {
 
 // UnitEnabled if this repository has the given unit enabled
 func (repo *Repository) UnitEnabled(tp unit.Type) (result bool) {
-	if err := db.WithContext(func(ctx *db.Context) error {
-		result = repo.UnitEnabledCtx(ctx, tp)
-		return nil
-	}); err != nil {
-		log.Error("repo.UnitEnabled: %v", err)
-	}
-	return result
+	return repo.UnitEnabledCtx(db.DefaultContext, tp)
 }
 
 // UnitEnabled if this repository has the given unit enabled
diff --git a/modules/repository/create.go b/modules/repository/create.go
index 3d04dba70f..9204d7e422 100644
--- a/modules/repository/create.go
+++ b/modules/repository/create.go
@@ -229,7 +229,7 @@ func UpdateRepository(ctx context.Context, repo *repo_model.Repository, visibili
 		}
 
 		// Create/Remove git-daemon-export-ok for git-daemon...
-		if err := CheckDaemonExportOK(db.WithEngine(ctx, e), repo); err != nil {
+		if err := CheckDaemonExportOK(ctx, repo); err != nil {
 			return err
 		}