Simple and Powerful ORM for Go, support mysql,postgres,tidb,sqlite3,mssql,oracle https://xorm.io
您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符
 
 
 
 

231 行
5.3 KiB

  1. // Copyright 2017 The Xorm Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package xorm
  5. import (
  6. "context"
  7. "time"
  8. "xorm.io/xorm/caches"
  9. "xorm.io/xorm/contexts"
  10. "xorm.io/xorm/dialects"
  11. "xorm.io/xorm/log"
  12. "xorm.io/xorm/names"
  13. )
  14. // EngineGroup defines an engine group
  15. type EngineGroup struct {
  16. *Engine
  17. slaves []*Engine
  18. policy GroupPolicy
  19. }
  20. // NewEngineGroup creates a new engine group
  21. func NewEngineGroup(args1 interface{}, args2 interface{}, policies ...GroupPolicy) (*EngineGroup, error) {
  22. var eg EngineGroup
  23. if len(policies) > 0 {
  24. eg.policy = policies[0]
  25. } else {
  26. eg.policy = RoundRobinPolicy()
  27. }
  28. driverName, ok1 := args1.(string)
  29. conns, ok2 := args2.([]string)
  30. if ok1 && ok2 {
  31. engines := make([]*Engine, len(conns))
  32. for i, conn := range conns {
  33. engine, err := NewEngine(driverName, conn)
  34. if err != nil {
  35. return nil, err
  36. }
  37. engine.engineGroup = &eg
  38. engines[i] = engine
  39. }
  40. eg.Engine = engines[0]
  41. eg.slaves = engines[1:]
  42. return &eg, nil
  43. }
  44. master, ok3 := args1.(*Engine)
  45. slaves, ok4 := args2.([]*Engine)
  46. if ok3 && ok4 {
  47. master.engineGroup = &eg
  48. for i := 0; i < len(slaves); i++ {
  49. slaves[i].engineGroup = &eg
  50. }
  51. eg.Engine = master
  52. eg.slaves = slaves
  53. return &eg, nil
  54. }
  55. return nil, ErrParamsType
  56. }
  57. // Close the engine
  58. func (eg *EngineGroup) Close() error {
  59. err := eg.Engine.Close()
  60. if err != nil {
  61. return err
  62. }
  63. for i := 0; i < len(eg.slaves); i++ {
  64. err := eg.slaves[i].Close()
  65. if err != nil {
  66. return err
  67. }
  68. }
  69. return nil
  70. }
  71. // ContextHook returned a group session
  72. func (eg *EngineGroup) Context(ctx context.Context) *Session {
  73. sess := eg.NewSession()
  74. sess.isAutoClose = true
  75. return sess.Context(ctx)
  76. }
  77. // NewSession returned a group session
  78. func (eg *EngineGroup) NewSession() *Session {
  79. sess := eg.Engine.NewSession()
  80. sess.sessionType = groupSession
  81. return sess
  82. }
  83. // Master returns the master engine
  84. func (eg *EngineGroup) Master() *Engine {
  85. return eg.Engine
  86. }
  87. // Ping tests if database is alive
  88. func (eg *EngineGroup) Ping() error {
  89. if err := eg.Engine.Ping(); err != nil {
  90. return err
  91. }
  92. for _, slave := range eg.slaves {
  93. if err := slave.Ping(); err != nil {
  94. return err
  95. }
  96. }
  97. return nil
  98. }
  99. // SetColumnMapper set the column name mapping rule
  100. func (eg *EngineGroup) SetColumnMapper(mapper names.Mapper) {
  101. eg.Engine.SetColumnMapper(mapper)
  102. for i := 0; i < len(eg.slaves); i++ {
  103. eg.slaves[i].SetColumnMapper(mapper)
  104. }
  105. }
  106. // SetConnMaxLifetime sets the maximum amount of time a connection may be reused.
  107. func (eg *EngineGroup) SetConnMaxLifetime(d time.Duration) {
  108. eg.Engine.SetConnMaxLifetime(d)
  109. for i := 0; i < len(eg.slaves); i++ {
  110. eg.slaves[i].SetConnMaxLifetime(d)
  111. }
  112. }
  113. // SetDefaultCacher set the default cacher
  114. func (eg *EngineGroup) SetDefaultCacher(cacher caches.Cacher) {
  115. eg.Engine.SetDefaultCacher(cacher)
  116. for i := 0; i < len(eg.slaves); i++ {
  117. eg.slaves[i].SetDefaultCacher(cacher)
  118. }
  119. }
  120. // SetLogger set the new logger
  121. func (eg *EngineGroup) SetLogger(logger interface{}) {
  122. eg.Engine.SetLogger(logger)
  123. for i := 0; i < len(eg.slaves); i++ {
  124. eg.slaves[i].SetLogger(logger)
  125. }
  126. }
  127. func (eg *EngineGroup) AddHook(hook contexts.Hook) {
  128. eg.Engine.AddHook(hook)
  129. for i := 0; i < len(eg.slaves); i++ {
  130. eg.slaves[i].AddHook(hook)
  131. }
  132. }
  133. // SetLogLevel sets the logger level
  134. func (eg *EngineGroup) SetLogLevel(level log.LogLevel) {
  135. eg.Engine.SetLogLevel(level)
  136. for i := 0; i < len(eg.slaves); i++ {
  137. eg.slaves[i].SetLogLevel(level)
  138. }
  139. }
  140. // SetMapper set the name mapping rules
  141. func (eg *EngineGroup) SetMapper(mapper names.Mapper) {
  142. eg.Engine.SetMapper(mapper)
  143. for i := 0; i < len(eg.slaves); i++ {
  144. eg.slaves[i].SetMapper(mapper)
  145. }
  146. }
  147. // SetMaxIdleConns set the max idle connections on pool, default is 2
  148. func (eg *EngineGroup) SetMaxIdleConns(conns int) {
  149. eg.Engine.DB().SetMaxIdleConns(conns)
  150. for i := 0; i < len(eg.slaves); i++ {
  151. eg.slaves[i].DB().SetMaxIdleConns(conns)
  152. }
  153. }
  154. // SetMaxOpenConns is only available for go 1.2+
  155. func (eg *EngineGroup) SetMaxOpenConns(conns int) {
  156. eg.Engine.DB().SetMaxOpenConns(conns)
  157. for i := 0; i < len(eg.slaves); i++ {
  158. eg.slaves[i].DB().SetMaxOpenConns(conns)
  159. }
  160. }
  161. // SetPolicy set the group policy
  162. func (eg *EngineGroup) SetPolicy(policy GroupPolicy) *EngineGroup {
  163. eg.policy = policy
  164. return eg
  165. }
  166. // SetQuotePolicy sets the special quote policy
  167. func (eg *EngineGroup) SetQuotePolicy(quotePolicy dialects.QuotePolicy) {
  168. eg.Engine.SetQuotePolicy(quotePolicy)
  169. for i := 0; i < len(eg.slaves); i++ {
  170. eg.slaves[i].SetQuotePolicy(quotePolicy)
  171. }
  172. }
  173. // SetTableMapper set the table name mapping rule
  174. func (eg *EngineGroup) SetTableMapper(mapper names.Mapper) {
  175. eg.Engine.SetTableMapper(mapper)
  176. for i := 0; i < len(eg.slaves); i++ {
  177. eg.slaves[i].SetTableMapper(mapper)
  178. }
  179. }
  180. // ShowSQL show SQL statement or not on logger if log level is great than INFO
  181. func (eg *EngineGroup) ShowSQL(show ...bool) {
  182. eg.Engine.ShowSQL(show...)
  183. for i := 0; i < len(eg.slaves); i++ {
  184. eg.slaves[i].ShowSQL(show...)
  185. }
  186. }
  187. // Slave returns one of the physical databases which is a slave according the policy
  188. func (eg *EngineGroup) Slave() *Engine {
  189. switch len(eg.slaves) {
  190. case 0:
  191. return eg.Engine
  192. case 1:
  193. return eg.slaves[0]
  194. }
  195. return eg.policy.Slave(eg)
  196. }
  197. // Slaves returns all the slaves
  198. func (eg *EngineGroup) Slaves() []*Engine {
  199. return eg.slaves
  200. }