44package db
55
66import (
7+ "bytes"
78 "context"
89 "database/sql"
10+ "errors"
11+ "runtime"
12+
13+ "code.gitea.io/gitea/modules/setting"
914
1015 "xorm.io/builder"
1116 "xorm.io/xorm"
@@ -15,62 +20,47 @@ import (
1520// will be overwritten by Init with HammerContext
1621var DefaultContext context.Context
1722
18- // contextKey is a value for use with context.WithValue.
19- type contextKey struct {
20- name string
21- }
23+ type engineContextKeyType struct {}
2224
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- )
25+ var engineContextKey = engineContextKeyType {}
2826
2927// Context represents a db context
3028type Context struct {
3129 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- }
42- }
43-
44- // InTransaction if context is in a transaction
45- func (ctx * Context ) InTransaction () bool {
46- return ctx .transaction
30+ engine Engine
4731}
4832
49- // Engine returns db engine
50- func (ctx * Context ) Engine () Engine {
51- return ctx .e
33+ func newContext (ctx context.Context , e Engine ) * Context {
34+ return & Context {Context : ctx , engine : e }
5235}
5336
5437// Value shadows Value for context.Context but allows us to get ourselves and an Engined object
5538func (ctx * Context ) Value (key any ) any {
56- if key == enginedContextKey {
39+ if key == engineContextKey {
5740 return ctx
5841 }
5942 return ctx .Context .Value (key )
6043}
6144
6245// WithContext returns this engine tied to this context
6346func (ctx * Context ) WithContext (other context.Context ) * Context {
64- return newContext (ctx , ctx .e .Context (other ), ctx . transaction )
47+ return newContext (ctx , ctx .engine .Context (other ))
6548}
6649
67- // Engined structs provide an Engine
68- type Engined interface {
69- Engine () Engine
50+ func contextSafetyCheck () {
51+ if setting .IsProd && ! setting .IsInTesting {
52+ return
53+ }
54+ buf := make ([]byte , 1024 * 100 )
55+ runtime .Stack (buf , false )
56+ if bytes .Contains (buf , []byte ("xorm.io/xorm.(*Session).Iterate(" )) {
57+ panic (errors .New ("using database context in an iterator would cause corrupted results" ))
58+ }
7059}
7160
7261// GetEngine will get a db Engine from this context or return an Engine restricted to this context
7362func GetEngine (ctx context.Context ) Engine {
63+ contextSafetyCheck ()
7464 if e := getEngine (ctx ); e != nil {
7565 return e
7666 }
@@ -79,12 +69,11 @@ func GetEngine(ctx context.Context) Engine {
7969
8070// getEngine will get a db Engine from this context or return nil
8171func getEngine (ctx context.Context ) Engine {
82- if engined , ok := ctx .(Engined ); ok {
83- return engined .Engine ()
72+ if engined , ok := ctx .(* Context ); ok {
73+ return engined .engine
8474 }
85- enginedInterface := ctx .Value (enginedContextKey )
86- if enginedInterface != nil {
87- return enginedInterface .(Engined ).Engine ()
75+ if engined , ok := ctx .Value (engineContextKey ).(* Context ); ok {
76+ return engined .engine
8877 }
8978 return nil
9079}
@@ -132,23 +121,23 @@ func (c *halfCommitter) Close() error {
132121// d. It doesn't mean rollback is forbidden, but always do it only when there is an error, and you do want to rollback.
133122func TxContext (parentCtx context.Context ) (* Context , Committer , error ) {
134123 if sess , ok := inTransaction (parentCtx ); ok {
135- return newContext (parentCtx , sess , true ), & halfCommitter {committer : sess }, nil
124+ return newContext (parentCtx , sess ), & halfCommitter {committer : sess }, nil
136125 }
137126
138127 sess := x .NewSession ()
139128 if err := sess .Begin (); err != nil {
140- sess .Close ()
129+ _ = sess .Close ()
141130 return nil , nil , err
142131 }
143132
144- return newContext (DefaultContext , sess , true ), sess , nil
133+ return newContext (DefaultContext , sess ), sess , nil
145134}
146135
147136// WithTx represents executing database operations on a transaction, if the transaction exist,
148137// this function will reuse it otherwise will create a new one and close it when finished.
149138func WithTx (parentCtx context.Context , f func (ctx context.Context ) error ) error {
150139 if sess , ok := inTransaction (parentCtx ); ok {
151- err := f (newContext (parentCtx , sess , true ))
140+ err := f (newContext (parentCtx , sess ))
152141 if err != nil {
153142 // rollback immediately, in case the caller ignores returned error and tries to commit the transaction.
154143 _ = sess .Close ()
@@ -165,7 +154,7 @@ func txWithNoCheck(parentCtx context.Context, f func(ctx context.Context) error)
165154 return err
166155 }
167156
168- if err := f (newContext (parentCtx , sess , true )); err != nil {
157+ if err := f (newContext (parentCtx , sess )); err != nil {
169158 return err
170159 }
171160
0 commit comments