test

Author: wxiaoguang

models/db/context.go

@@ -15,58 +15,31 @@ 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,
-	}
+	engine Engine
 }
 
-// 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
+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)
 }
 
 // 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)
-}
-
-// Engined structs provide an Engine
-type Engined interface {
-	Engine() Engine
+	return newContext(ctx, ctx.engine.Context(other))
 }
 
 // GetEngine will get a db Engine from this context or return an Engine restricted to this context
@@ -79,12 +52,11 @@ func GetEngine(ctx context.Context) Engine {
 
 // 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()
+	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 +104,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 +137,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
 	}
 

models/db/engine.go

@@ -161,10 +161,7 @@ func InitEngine(ctx context.Context) error {
 // SetDefaultEngine sets the default engine for db
 func SetDefaultEngine(ctx context.Context, eng *xorm.Engine) {
 	x = eng
-	DefaultContext = &Context{
-		Context: ctx,
-		e:       x,
-	}
+	DefaultContext = &Context{Context: ctx, engine: x}
 }
 
 // UnsetDefaultEngine closes and unsets the default engine

models/db/install/db.go

@@ -11,7 +11,7 @@ import (
 )
 
 func getXORMEngine() *xorm.Engine {
-	return db.DefaultContext.(*db.Context).Engine().(*xorm.Engine)
+	return db.GetEngine(db.DefaultContext).(*xorm.Engine)
 }
 
 // CheckDatabaseConnection checks the database connection

models/unittest/fixtures.go

@@ -25,7 +25,7 @@ func GetXORMEngine(engine ...*xorm.Engine) (x *xorm.Engine) {
 	if len(engine) == 1 {
 		return engine[0]
 	}
-	return db.DefaultContext.(*db.Context).Engine().(*xorm.Engine)
+	return db.GetEngine(db.DefaultContext).(*xorm.Engine)
 }
 
 // InitFixtures initialize test fixtures for a test database