tx.go 9.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290
  1. // Code generated by ent, DO NOT EDIT.
  2. package ent
  3. import (
  4. "context"
  5. stdsql "database/sql"
  6. "fmt"
  7. "sync"
  8. "entgo.io/ent/dialect"
  9. )
  10. // Tx is a transactional client that is created by calling Client.Tx().
  11. type Tx struct {
  12. config
  13. // Agent is the client for interacting with the Agent builders.
  14. Agent *AgentClient
  15. // BatchMsg is the client for interacting with the BatchMsg builders.
  16. BatchMsg *BatchMsgClient
  17. // Category is the client for interacting with the Category builders.
  18. Category *CategoryClient
  19. // Contact is the client for interacting with the Contact builders.
  20. Contact *ContactClient
  21. // Employee is the client for interacting with the Employee builders.
  22. Employee *EmployeeClient
  23. // EmployeeConfig is the client for interacting with the EmployeeConfig builders.
  24. EmployeeConfig *EmployeeConfigClient
  25. // Label is the client for interacting with the Label builders.
  26. Label *LabelClient
  27. // LabelRelationship is the client for interacting with the LabelRelationship builders.
  28. LabelRelationship *LabelRelationshipClient
  29. // Message is the client for interacting with the Message builders.
  30. Message *MessageClient
  31. // MessageRecords is the client for interacting with the MessageRecords builders.
  32. MessageRecords *MessageRecordsClient
  33. // Msg is the client for interacting with the Msg builders.
  34. Msg *MsgClient
  35. // Server is the client for interacting with the Server builders.
  36. Server *ServerClient
  37. // SopNode is the client for interacting with the SopNode builders.
  38. SopNode *SopNodeClient
  39. // SopStage is the client for interacting with the SopStage builders.
  40. SopStage *SopStageClient
  41. // SopTask is the client for interacting with the SopTask builders.
  42. SopTask *SopTaskClient
  43. // Token is the client for interacting with the Token builders.
  44. Token *TokenClient
  45. // Tutorial is the client for interacting with the Tutorial builders.
  46. Tutorial *TutorialClient
  47. // WorkExperience is the client for interacting with the WorkExperience builders.
  48. WorkExperience *WorkExperienceClient
  49. // Wx is the client for interacting with the Wx builders.
  50. Wx *WxClient
  51. // lazily loaded.
  52. client *Client
  53. clientOnce sync.Once
  54. // ctx lives for the life of the transaction. It is
  55. // the same context used by the underlying connection.
  56. ctx context.Context
  57. }
  58. type (
  59. // Committer is the interface that wraps the Commit method.
  60. Committer interface {
  61. Commit(context.Context, *Tx) error
  62. }
  63. // The CommitFunc type is an adapter to allow the use of ordinary
  64. // function as a Committer. If f is a function with the appropriate
  65. // signature, CommitFunc(f) is a Committer that calls f.
  66. CommitFunc func(context.Context, *Tx) error
  67. // CommitHook defines the "commit middleware". A function that gets a Committer
  68. // and returns a Committer. For example:
  69. //
  70. // hook := func(next ent.Committer) ent.Committer {
  71. // return ent.CommitFunc(func(ctx context.Context, tx *ent.Tx) error {
  72. // // Do some stuff before.
  73. // if err := next.Commit(ctx, tx); err != nil {
  74. // return err
  75. // }
  76. // // Do some stuff after.
  77. // return nil
  78. // })
  79. // }
  80. //
  81. CommitHook func(Committer) Committer
  82. )
  83. // Commit calls f(ctx, m).
  84. func (f CommitFunc) Commit(ctx context.Context, tx *Tx) error {
  85. return f(ctx, tx)
  86. }
  87. // Commit commits the transaction.
  88. func (tx *Tx) Commit() error {
  89. txDriver := tx.config.driver.(*txDriver)
  90. var fn Committer = CommitFunc(func(context.Context, *Tx) error {
  91. return txDriver.tx.Commit()
  92. })
  93. txDriver.mu.Lock()
  94. hooks := append([]CommitHook(nil), txDriver.onCommit...)
  95. txDriver.mu.Unlock()
  96. for i := len(hooks) - 1; i >= 0; i-- {
  97. fn = hooks[i](fn)
  98. }
  99. return fn.Commit(tx.ctx, tx)
  100. }
  101. // OnCommit adds a hook to call on commit.
  102. func (tx *Tx) OnCommit(f CommitHook) {
  103. txDriver := tx.config.driver.(*txDriver)
  104. txDriver.mu.Lock()
  105. txDriver.onCommit = append(txDriver.onCommit, f)
  106. txDriver.mu.Unlock()
  107. }
  108. type (
  109. // Rollbacker is the interface that wraps the Rollback method.
  110. Rollbacker interface {
  111. Rollback(context.Context, *Tx) error
  112. }
  113. // The RollbackFunc type is an adapter to allow the use of ordinary
  114. // function as a Rollbacker. If f is a function with the appropriate
  115. // signature, RollbackFunc(f) is a Rollbacker that calls f.
  116. RollbackFunc func(context.Context, *Tx) error
  117. // RollbackHook defines the "rollback middleware". A function that gets a Rollbacker
  118. // and returns a Rollbacker. For example:
  119. //
  120. // hook := func(next ent.Rollbacker) ent.Rollbacker {
  121. // return ent.RollbackFunc(func(ctx context.Context, tx *ent.Tx) error {
  122. // // Do some stuff before.
  123. // if err := next.Rollback(ctx, tx); err != nil {
  124. // return err
  125. // }
  126. // // Do some stuff after.
  127. // return nil
  128. // })
  129. // }
  130. //
  131. RollbackHook func(Rollbacker) Rollbacker
  132. )
  133. // Rollback calls f(ctx, m).
  134. func (f RollbackFunc) Rollback(ctx context.Context, tx *Tx) error {
  135. return f(ctx, tx)
  136. }
  137. // Rollback rollbacks the transaction.
  138. func (tx *Tx) Rollback() error {
  139. txDriver := tx.config.driver.(*txDriver)
  140. var fn Rollbacker = RollbackFunc(func(context.Context, *Tx) error {
  141. return txDriver.tx.Rollback()
  142. })
  143. txDriver.mu.Lock()
  144. hooks := append([]RollbackHook(nil), txDriver.onRollback...)
  145. txDriver.mu.Unlock()
  146. for i := len(hooks) - 1; i >= 0; i-- {
  147. fn = hooks[i](fn)
  148. }
  149. return fn.Rollback(tx.ctx, tx)
  150. }
  151. // OnRollback adds a hook to call on rollback.
  152. func (tx *Tx) OnRollback(f RollbackHook) {
  153. txDriver := tx.config.driver.(*txDriver)
  154. txDriver.mu.Lock()
  155. txDriver.onRollback = append(txDriver.onRollback, f)
  156. txDriver.mu.Unlock()
  157. }
  158. // Client returns a Client that binds to current transaction.
  159. func (tx *Tx) Client() *Client {
  160. tx.clientOnce.Do(func() {
  161. tx.client = &Client{config: tx.config}
  162. tx.client.init()
  163. })
  164. return tx.client
  165. }
  166. func (tx *Tx) init() {
  167. tx.Agent = NewAgentClient(tx.config)
  168. tx.BatchMsg = NewBatchMsgClient(tx.config)
  169. tx.Category = NewCategoryClient(tx.config)
  170. tx.Contact = NewContactClient(tx.config)
  171. tx.Employee = NewEmployeeClient(tx.config)
  172. tx.EmployeeConfig = NewEmployeeConfigClient(tx.config)
  173. tx.Label = NewLabelClient(tx.config)
  174. tx.LabelRelationship = NewLabelRelationshipClient(tx.config)
  175. tx.Message = NewMessageClient(tx.config)
  176. tx.MessageRecords = NewMessageRecordsClient(tx.config)
  177. tx.Msg = NewMsgClient(tx.config)
  178. tx.Server = NewServerClient(tx.config)
  179. tx.SopNode = NewSopNodeClient(tx.config)
  180. tx.SopStage = NewSopStageClient(tx.config)
  181. tx.SopTask = NewSopTaskClient(tx.config)
  182. tx.Token = NewTokenClient(tx.config)
  183. tx.Tutorial = NewTutorialClient(tx.config)
  184. tx.WorkExperience = NewWorkExperienceClient(tx.config)
  185. tx.Wx = NewWxClient(tx.config)
  186. }
  187. // txDriver wraps the given dialect.Tx with a nop dialect.Driver implementation.
  188. // The idea is to support transactions without adding any extra code to the builders.
  189. // When a builder calls to driver.Tx(), it gets the same dialect.Tx instance.
  190. // Commit and Rollback are nop for the internal builders and the user must call one
  191. // of them in order to commit or rollback the transaction.
  192. //
  193. // If a closed transaction is embedded in one of the generated entities, and the entity
  194. // applies a query, for example: Agent.QueryXXX(), the query will be executed
  195. // through the driver which created this transaction.
  196. //
  197. // Note that txDriver is not goroutine safe.
  198. type txDriver struct {
  199. // the driver we started the transaction from.
  200. drv dialect.Driver
  201. // tx is the underlying transaction.
  202. tx dialect.Tx
  203. // completion hooks.
  204. mu sync.Mutex
  205. onCommit []CommitHook
  206. onRollback []RollbackHook
  207. }
  208. // newTx creates a new transactional driver.
  209. func newTx(ctx context.Context, drv dialect.Driver) (*txDriver, error) {
  210. tx, err := drv.Tx(ctx)
  211. if err != nil {
  212. return nil, err
  213. }
  214. return &txDriver{tx: tx, drv: drv}, nil
  215. }
  216. // Tx returns the transaction wrapper (txDriver) to avoid Commit or Rollback calls
  217. // from the internal builders. Should be called only by the internal builders.
  218. func (tx *txDriver) Tx(context.Context) (dialect.Tx, error) { return tx, nil }
  219. // Dialect returns the dialect of the driver we started the transaction from.
  220. func (tx *txDriver) Dialect() string { return tx.drv.Dialect() }
  221. // Close is a nop close.
  222. func (*txDriver) Close() error { return nil }
  223. // Commit is a nop commit for the internal builders.
  224. // User must call `Tx.Commit` in order to commit the transaction.
  225. func (*txDriver) Commit() error { return nil }
  226. // Rollback is a nop rollback for the internal builders.
  227. // User must call `Tx.Rollback` in order to rollback the transaction.
  228. func (*txDriver) Rollback() error { return nil }
  229. // Exec calls tx.Exec.
  230. func (tx *txDriver) Exec(ctx context.Context, query string, args, v any) error {
  231. return tx.tx.Exec(ctx, query, args, v)
  232. }
  233. // Query calls tx.Query.
  234. func (tx *txDriver) Query(ctx context.Context, query string, args, v any) error {
  235. return tx.tx.Query(ctx, query, args, v)
  236. }
  237. var _ dialect.Driver = (*txDriver)(nil)
  238. // ExecContext allows calling the underlying ExecContext method of the transaction if it is supported by it.
  239. // See, database/sql#Tx.ExecContext for more information.
  240. func (tx *txDriver) ExecContext(ctx context.Context, query string, args ...any) (stdsql.Result, error) {
  241. ex, ok := tx.tx.(interface {
  242. ExecContext(context.Context, string, ...any) (stdsql.Result, error)
  243. })
  244. if !ok {
  245. return nil, fmt.Errorf("Tx.ExecContext is not supported")
  246. }
  247. return ex.ExecContext(ctx, query, args...)
  248. }
  249. // QueryContext allows calling the underlying QueryContext method of the transaction if it is supported by it.
  250. // See, database/sql#Tx.QueryContext for more information.
  251. func (tx *txDriver) QueryContext(ctx context.Context, query string, args ...any) (*stdsql.Rows, error) {
  252. q, ok := tx.tx.(interface {
  253. QueryContext(context.Context, string, ...any) (*stdsql.Rows, error)
  254. })
  255. if !ok {
  256. return nil, fmt.Errorf("Tx.QueryContext is not supported")
  257. }
  258. return q.QueryContext(ctx, query, args...)
  259. }