@@ -15,58 +15,31 @@ import (
1515// will be overwritten by Init with HammerContext
1616var DefaultContext context.Context
1717
18- // contextKey is a value for use with context.WithValue.
19- type contextKey struct {
20- name string
21- }
18+ type engineContextKeyType struct {}
2219
23- // enginedContextKey is a context key. It is used with context.Value() to get the current Engined for the context
24- var (
25- enginedContextKey = & contextKey {"engined" }
26- _ Engined = & Context {}
27- )
20+ var engineContextKey = engineContextKeyType {}
2821
2922// Context represents a db context
3023type Context struct {
3124 context.Context
32- e Engine
33- transaction bool
34- }
35-
36- func newContext (ctx context.Context , e Engine , transaction bool ) * Context {
37- return & Context {
38- Context : ctx ,
39- e : e ,
40- transaction : transaction ,
41- }
25+ engine Engine
4226}
4327
44- // InTransaction if context is in a transaction
45- func (ctx * Context ) InTransaction () bool {
46- return ctx .transaction
47- }
48-
49- // Engine returns db engine
50- func (ctx * Context ) Engine () Engine {
51- return ctx .e
28+ func newContext (ctx context.Context , e Engine ) * Context {
29+ return & Context {Context : ctx , engine : e }
5230}
5331
5432// Value shadows Value for context.Context but allows us to get ourselves and an Engined object
5533func (ctx * Context ) Value (key any ) any {
56- if key == enginedContextKey {
34+ if key == engineContextKey {
5735 return ctx
5836 }
5937 return ctx .Context .Value (key )
6038}
6139
6240// WithContext returns this engine tied to this context
6341func (ctx * Context ) WithContext (other context.Context ) * Context {
64- return newContext (ctx , ctx .e .Context (other ), ctx .transaction )
65- }
66-
67- // Engined structs provide an Engine
68- type Engined interface {
69- Engine () Engine
42+ return newContext (ctx , ctx .engine .Context (other ))
7043}
7144
7245// 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 {
7952
8053// getEngine will get a db Engine from this context or return nil
8154func getEngine (ctx context.Context ) Engine {
82- if engined , ok := ctx .(Engined ); ok {
83- return engined .Engine ()
55+ if engined , ok := ctx .(* Context ); ok {
56+ return engined .engine
8457 }
85- enginedInterface := ctx .Value (enginedContextKey )
86- if enginedInterface != nil {
87- return enginedInterface .(Engined ).Engine ()
58+ if engined , ok := ctx .Value (engineContextKey ).(* Context ); ok {
59+ return engined .engine
8860 }
8961 return nil
9062}
@@ -132,23 +104,23 @@ func (c *halfCommitter) Close() error {
132104// d. It doesn't mean rollback is forbidden, but always do it only when there is an error, and you do want to rollback.
133105func TxContext (parentCtx context.Context ) (* Context , Committer , error ) {
134106 if sess , ok := inTransaction (parentCtx ); ok {
135- return newContext (parentCtx , sess , true ), & halfCommitter {committer : sess }, nil
107+ return newContext (parentCtx , sess ), & halfCommitter {committer : sess }, nil
136108 }
137109
138110 sess := x .NewSession ()
139111 if err := sess .Begin (); err != nil {
140- sess .Close ()
112+ _ = sess .Close ()
141113 return nil , nil , err
142114 }
143115
144- return newContext (DefaultContext , sess , true ), sess , nil
116+ return newContext (DefaultContext , sess ), sess , nil
145117}
146118
147119// WithTx represents executing database operations on a transaction, if the transaction exist,
148120// this function will reuse it otherwise will create a new one and close it when finished.
149121func WithTx (parentCtx context.Context , f func (ctx context.Context ) error ) error {
150122 if sess , ok := inTransaction (parentCtx ); ok {
151- err := f (newContext (parentCtx , sess , true ))
123+ err := f (newContext (parentCtx , sess ))
152124 if err != nil {
153125 // rollback immediately, in case the caller ignores returned error and tries to commit the transaction.
154126 _ = sess .Close ()
@@ -165,7 +137,7 @@ func txWithNoCheck(parentCtx context.Context, f func(ctx context.Context) error)
165137 return err
166138 }
167139
168- if err := f (newContext (parentCtx , sess , true )); err != nil {
140+ if err := f (newContext (parentCtx , sess )); err != nil {
169141 return err
170142 }
171143
0 commit comments