|
|
|
@ -6,6 +6,12 @@ package db |
|
|
|
|
import ( |
|
|
|
|
"context" |
|
|
|
|
"database/sql" |
|
|
|
|
"errors" |
|
|
|
|
"runtime" |
|
|
|
|
"slices" |
|
|
|
|
"sync" |
|
|
|
|
|
|
|
|
|
"code.gitea.io/gitea/modules/setting" |
|
|
|
|
|
|
|
|
|
"xorm.io/builder" |
|
|
|
|
"xorm.io/xorm" |
|
|
|
@ -15,45 +21,23 @@ import ( |
|
|
|
|
// 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 |
|
|
|
|
} |
|
|
|
|
type engineContextKeyType struct{} |
|
|
|
|
|
|
|
|
|
// enginedContextKey is a context key. It is used with context.Value() to get the current Engined for the context
|
|
|
|
|
var ( |
|
|
|
|
enginedContextKey = &contextKey{"engined"} |
|
|
|
|
_ Engined = &Context{} |
|
|
|
|
) |
|
|
|
|
var engineContextKey = engineContextKeyType{} |
|
|
|
|
|
|
|
|
|
// Context represents a db context
|
|
|
|
|
type Context struct { |
|
|
|
|
context.Context |
|
|
|
|
e Engine |
|
|
|
|
transaction bool |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func newContext(ctx context.Context, e Engine, transaction bool) *Context { |
|
|
|
|
return &Context{ |
|
|
|
|
Context: ctx, |
|
|
|
|
e: e, |
|
|
|
|
transaction: transaction, |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
// InTransaction if context is in a transaction
|
|
|
|
|
func (ctx *Context) InTransaction() bool { |
|
|
|
|
return ctx.transaction |
|
|
|
|
engine Engine |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
// Engine returns db engine
|
|
|
|
|
func (ctx *Context) Engine() Engine { |
|
|
|
|
return ctx.e |
|
|
|
|
func newContext(ctx context.Context, e Engine) *Context { |
|
|
|
|
return &Context{Context: ctx, engine: e} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
// Value shadows Value for context.Context but allows us to get ourselves and an Engined object
|
|
|
|
|
func (ctx *Context) Value(key any) any { |
|
|
|
|
if key == enginedContextKey { |
|
|
|
|
if key == engineContextKey { |
|
|
|
|
return ctx |
|
|
|
|
} |
|
|
|
|
return ctx.Context.Value(key) |
|
|
|
@ -61,30 +45,66 @@ func (ctx *Context) Value(key any) any { |
|
|
|
|
|
|
|
|
|
// WithContext returns this engine tied to this context
|
|
|
|
|
func (ctx *Context) WithContext(other context.Context) *Context { |
|
|
|
|
return newContext(ctx, ctx.e.Context(other), ctx.transaction) |
|
|
|
|
return newContext(ctx, ctx.engine.Context(other)) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
// Engined structs provide an Engine
|
|
|
|
|
type Engined interface { |
|
|
|
|
Engine() Engine |
|
|
|
|
var ( |
|
|
|
|
contextSafetyOnce sync.Once |
|
|
|
|
contextSafetyDeniedFuncPCs []uintptr |
|
|
|
|
) |
|
|
|
|
|
|
|
|
|
func contextSafetyCheck(e Engine) { |
|
|
|
|
if setting.IsProd && !setting.IsInTesting { |
|
|
|
|
return |
|
|
|
|
} |
|
|
|
|
if e == nil { |
|
|
|
|
return |
|
|
|
|
} |
|
|
|
|
// Only do this check for non-end-users. If the problem could be fixed in the future, this code could be removed.
|
|
|
|
|
contextSafetyOnce.Do(func() { |
|
|
|
|
// try to figure out the bad functions to deny
|
|
|
|
|
type m struct{} |
|
|
|
|
_ = e.SQL("SELECT 1").Iterate(&m{}, func(int, any) error { |
|
|
|
|
callers := make([]uintptr, 32) |
|
|
|
|
callerNum := runtime.Callers(1, callers) |
|
|
|
|
for i := 0; i < callerNum; i++ { |
|
|
|
|
if funcName := runtime.FuncForPC(callers[i]).Name(); funcName == "xorm.io/xorm.(*Session).Iterate" { |
|
|
|
|
contextSafetyDeniedFuncPCs = append(contextSafetyDeniedFuncPCs, callers[i]) |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
return nil |
|
|
|
|
}) |
|
|
|
|
if len(contextSafetyDeniedFuncPCs) != 1 { |
|
|
|
|
panic(errors.New("unable to determine the functions to deny")) |
|
|
|
|
} |
|
|
|
|
}) |
|
|
|
|
|
|
|
|
|
// it should be very fast: xxxx ns/op
|
|
|
|
|
callers := make([]uintptr, 32) |
|
|
|
|
callerNum := runtime.Callers(3, callers) // skip 3: runtime.Callers, contextSafetyCheck, GetEngine
|
|
|
|
|
for i := 0; i < callerNum; i++ { |
|
|
|
|
if slices.Contains(contextSafetyDeniedFuncPCs, callers[i]) { |
|
|
|
|
panic(errors.New("using database context in an iterator would cause corrupted results")) |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
// GetEngine will get a db Engine from this context or return an Engine restricted to this context
|
|
|
|
|
// GetEngine gets an existing db Engine/Statement or creates a new Session
|
|
|
|
|
func GetEngine(ctx context.Context) Engine { |
|
|
|
|
if e := getEngine(ctx); e != nil { |
|
|
|
|
if e := getExistingEngine(ctx); e != nil { |
|
|
|
|
return e |
|
|
|
|
} |
|
|
|
|
return x.Context(ctx) |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
// getEngine will get a db Engine from this context or return nil
|
|
|
|
|
func getEngine(ctx context.Context) Engine { |
|
|
|
|
if engined, ok := ctx.(Engined); ok { |
|
|
|
|
return engined.Engine() |
|
|
|
|
// getExistingEngine gets an existing db Engine/Statement from this context or returns nil
|
|
|
|
|
func getExistingEngine(ctx context.Context) (e Engine) { |
|
|
|
|
defer func() { contextSafetyCheck(e) }() |
|
|
|
|
if engined, ok := ctx.(*Context); ok { |
|
|
|
|
return engined.engine |
|
|
|
|
} |
|
|
|
|
enginedInterface := ctx.Value(enginedContextKey) |
|
|
|
|
if enginedInterface != nil { |
|
|
|
|
return enginedInterface.(Engined).Engine() |
|
|
|
|
if engined, ok := ctx.Value(engineContextKey).(*Context); ok { |
|
|
|
|
return engined.engine |
|
|
|
|
} |
|
|
|
|
return nil |
|
|
|
|
} |
|
|
|
@ -132,23 +152,23 @@ func (c *halfCommitter) Close() error { |
|
|
|
|
// d. It doesn't mean rollback is forbidden, but always do it only when there is an error, and you do want to rollback.
|
|
|
|
|
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), &halfCommitter{committer: sess}, nil |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
sess := x.NewSession() |
|
|
|
|
if err := sess.Begin(); err != nil { |
|
|
|
|
sess.Close() |
|
|
|
|
_ = sess.Close() |
|
|
|
|
return nil, nil, err |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
return newContext(DefaultContext, sess, true), sess, nil |
|
|
|
|
return newContext(DefaultContext, sess), sess, nil |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
// WithTx represents executing database operations on a transaction, if the transaction exist,
|
|
|
|
|
// 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 { |
|
|
|
|
err := f(newContext(parentCtx, sess, true)) |
|
|
|
|
err := f(newContext(parentCtx, sess)) |
|
|
|
|
if err != nil { |
|
|
|
|
// rollback immediately, in case the caller ignores returned error and tries to commit the transaction.
|
|
|
|
|
_ = sess.Close() |
|
|
|
@ -165,7 +185,7 @@ func txWithNoCheck(parentCtx context.Context, f func(ctx context.Context) error) |
|
|
|
|
return err |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
if err := f(newContext(parentCtx, sess, true)); err != nil { |
|
|
|
|
if err := f(newContext(parentCtx, sess)); err != nil { |
|
|
|
|
return err |
|
|
|
|
} |
|
|
|
|
|
|
|
|
@ -312,7 +332,7 @@ func InTransaction(ctx context.Context) bool { |
|
|
|
|
} |
|
|
|
|
|
|
|
|
|
func inTransaction(ctx context.Context) (*xorm.Session, bool) { |
|
|
|
|
e := getEngine(ctx) |
|
|
|
|
e := getExistingEngine(ctx) |
|
|
|
|
if e == nil { |
|
|
|
|
return nil, false |
|
|
|
|
} |
|
|
|
|