DBContext is just a Context (#17100)

* DBContext is just a Context

This PR removes some of the specialness from the DBContext and makes it context
This allows us to simplify the GetEngine code to wrap around any context in future
and means that we can change our loadRepo(e Engine) functions to simply take contexts.

Signed-off-by: Andrew Thornton <art27@cantab.net>

* fix unit tests

Signed-off-by: Andrew Thornton <art27@cantab.net>

* another place that needs to set the initial context

Signed-off-by: Andrew Thornton <art27@cantab.net>

* avoid race

Signed-off-by: Andrew Thornton <art27@cantab.net>

* change attachment error

Signed-off-by: Andrew Thornton <art27@cantab.net>

models/db/context.go

@@ -5,14 +5,29 @@
 package db
 
 import (
+	"context"
+
 	"code.gitea.io/gitea/modules/setting"
 
 	"xorm.io/builder"
 	"xorm.io/xorm"
 )
 
+// DefaultContext is the default context to run xorm queries in
+// will be overwritten by Init with HammerContext
+var DefaultContext context.Context
+
+// contextKey is a value for use with context.WithValue.
+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"}
+
 // Context represents a db context
 type Context struct {
+	context.Context
 	e Engine
 }
 
@@ -30,9 +45,48 @@ func (ctx *Context) NewSession() *xorm.Session {
 	return nil
 }
 
-// DefaultContext represents a Context with default Engine
-func DefaultContext() *Context {
-	return &Context{x}
+// 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 {
+		return ctx
+	}
+	return ctx.Context.Value(key)
+}
+
+// Engined structs provide an Engine
+type Engined interface {
+	Engine() Engine
+	NewSession() *xorm.Session
+}
+
+// GetEngine will get a db Engine from this context or return an Engine restricted to this context
+func GetEngine(ctx context.Context) Engine {
+	if engined, ok := ctx.(Engined); ok {
+		return engined.Engine()
+	}
+	enginedInterface := ctx.Value(EnginedContextKey)
+	if enginedInterface != nil {
+		return enginedInterface.(Engined).Engine()
+	}
+	return x.Context(ctx)
+}
+
+// NewSession will get a db Session from this context or return a session restricted to this context
+func NewSession(ctx context.Context) *xorm.Session {
+	if engined, ok := ctx.(Engined); ok {
+		return engined.NewSession()
+	}
+
+	enginedInterface := ctx.Value(EnginedContextKey)
+	if enginedInterface != nil {
+		sess := enginedInterface.(Engined).NewSession()
+		if sess != nil {
+			return sess.Context(ctx)
+		}
+		return nil
+	}
+
+	return x.NewSession().Context(ctx)
 }
 
 // Committer represents an interface to Commit or Close the Context
@@ -49,23 +103,32 @@ func TxContext() (*Context, Committer, error) {
 		return nil, nil, err
 	}
 
-	return &Context{sess}, sess, nil
+	return &Context{
+		Context: DefaultContext,
+		e:       sess,
+	}, sess, nil
 }
 
 // WithContext represents executing database operations
 func WithContext(f func(ctx *Context) error) error {
-	return f(&Context{x})
+	return f(&Context{
+		Context: DefaultContext,
+		e:       x,
+	})
 }
 
 // WithTx represents executing database operations on a transaction
-func WithTx(f func(ctx *Context) error) error {
+func WithTx(f func(ctx context.Context) error) error {
 	sess := x.NewSession()
 	defer sess.Close()
 	if err := sess.Begin(); err != nil {
 		return err
 	}
 
-	if err := f(&Context{sess}); err != nil {
+	if err := f(&Context{
+		Context: DefaultContext,
+		e:       sess,
+	}); err != nil {
 		return err
 	}
 
@@ -73,14 +136,14 @@ func WithTx(f func(ctx *Context) error) error {
 }
 
 // Iterate iterates the databases and doing something
-func Iterate(ctx *Context, tableBean interface{}, cond builder.Cond, fun func(idx int, bean interface{}) error) error {
-	return ctx.e.Where(cond).
+func Iterate(ctx context.Context, tableBean interface{}, cond builder.Cond, fun func(idx int, bean interface{}) error) error {
+	return GetEngine(ctx).Where(cond).
 		BufferSize(setting.Database.IterateBufferSize).
 		Iterate(tableBean, fun)
 }
 
 // Insert inserts records into database
-func Insert(ctx *Context, beans ...interface{}) error {
-	_, err := ctx.e.Insert(beans...)
+func Insert(ctx context.Context, beans ...interface{}) error {
+	_, err := GetEngine(ctx).Insert(beans...)
 	return err
 }

models/db/engine.go

@@ -169,6 +169,11 @@ func NewEngine(ctx context.Context, migrateFunc func(*xorm.Engine) error) (err e
 		return err
 	}
 
+	DefaultContext = &Context{
+		Context: ctx,
+		e:       x,
+	}
+
 	x.SetDefaultContext(ctx)
 
 	if err = x.Ping(); err != nil {

models/db/unit_tests.go

@@ -5,6 +5,7 @@
 package db
 
 import (
+	"context"
 	"fmt"
 	"math"
 	"net/url"
@@ -117,6 +118,11 @@ func CreateTestEngine(fixturesDir string) error {
 		x.ShowSQL(true)
 	}
 
+	DefaultContext = &Context{
+		Context: context.Background(),
+		e:       x,
+	}
+
 	return InitFixtures(fixturesDir)
 }