/episode21/vendor/github.com/gobuffalo/pop/connection.go

https://github.com/arschles/go-in-5-minutes · Go · 189 lines · 149 code · 18 blank · 22 comment · 34 complexity · 19040e182b45a084b58f63501c1ea2b2 MD5 · raw file

  1. package pop
  2. import (
  3. "sync/atomic"
  4. "time"
  5. "github.com/jmoiron/sqlx"
  6. "github.com/markbates/going/defaults"
  7. "github.com/markbates/going/randx"
  8. "github.com/pkg/errors"
  9. )
  10. // Connections contains all of the available connections
  11. var Connections = map[string]*Connection{}
  12. // Connection represents all of the necessary details for
  13. // talking with a datastore
  14. type Connection struct {
  15. ID string
  16. Store store
  17. Dialect dialect
  18. Elapsed int64
  19. TX *Tx
  20. }
  21. func (c *Connection) String() string {
  22. return c.URL()
  23. }
  24. // URL returns the datasource connection string
  25. func (c *Connection) URL() string {
  26. return c.Dialect.URL()
  27. }
  28. // MigrationURL returns the datasource connection string used for running the migrations
  29. func (c *Connection) MigrationURL() string {
  30. return c.Dialect.MigrationURL()
  31. }
  32. // NewConnection creates a new connection, and sets it's `Dialect`
  33. // appropriately based on the `ConnectionDetails` passed into it.
  34. func NewConnection(deets *ConnectionDetails) (*Connection, error) {
  35. err := deets.Finalize()
  36. if err != nil {
  37. return nil, errors.WithStack(err)
  38. }
  39. c := &Connection{
  40. ID: randx.String(30),
  41. }
  42. switch deets.Dialect {
  43. case "postgres":
  44. c.Dialect = newPostgreSQL(deets)
  45. case "cockroach":
  46. c.Dialect = newCockroach(deets)
  47. case "mysql":
  48. c.Dialect = newMySQL(deets)
  49. case "sqlite3":
  50. c.Dialect, err = newSQLite(deets)
  51. if err != nil {
  52. return c, errors.WithStack(err)
  53. }
  54. }
  55. return c, nil
  56. }
  57. // Connect takes the name of a connection, default is "development", and will
  58. // return that connection from the available `Connections`. If a connection with
  59. // that name can not be found an error will be returned. If a connection is
  60. // found, and it has yet to open a connection with its underlying datastore,
  61. // a connection to that store will be opened.
  62. func Connect(e string) (*Connection, error) {
  63. if len(Connections) == 0 {
  64. err := LoadConfigFile()
  65. if err != nil {
  66. return nil, errors.WithStack(err)
  67. }
  68. }
  69. e = defaults.String(e, "development")
  70. c := Connections[e]
  71. if c == nil {
  72. return c, errors.Errorf("Could not find connection named %s!", e)
  73. }
  74. err := c.Open()
  75. return c, errors.Wrapf(err, "couldn't open connection for %s", e)
  76. }
  77. // Open creates a new datasource connection
  78. func (c *Connection) Open() error {
  79. if c.Store != nil {
  80. return nil
  81. }
  82. db, err := sqlx.Open(c.Dialect.Details().Dialect, c.Dialect.URL())
  83. db.SetMaxOpenConns(c.Dialect.Details().Pool)
  84. if err == nil {
  85. c.Store = &dB{db}
  86. }
  87. return errors.Wrap(err, "coudn't connection to database")
  88. }
  89. // Close destroys an active datasource connection
  90. func (c *Connection) Close() error {
  91. return errors.Wrap(c.Store.Close(), "couldn't close connection")
  92. }
  93. // Transaction will start a new transaction on the connection. If the inner function
  94. // returns an error then the transaction will be rolled back, otherwise the transaction
  95. // will automatically commit at the end.
  96. func (c *Connection) Transaction(fn func(tx *Connection) error) error {
  97. return c.Dialect.Lock(func() error {
  98. var dberr error
  99. cn, err := c.NewTransaction()
  100. if err != nil {
  101. return err
  102. }
  103. err = fn(cn)
  104. if err != nil {
  105. dberr = cn.TX.Rollback()
  106. } else {
  107. dberr = cn.TX.Commit()
  108. }
  109. if err != nil {
  110. return errors.WithStack(err)
  111. }
  112. return errors.Wrap(dberr, "error committing or rolling back transaction")
  113. })
  114. }
  115. // NewTransaction starts a new transaction on the connection
  116. func (c *Connection) NewTransaction() (*Connection, error) {
  117. var cn *Connection
  118. if c.TX == nil {
  119. tx, err := c.Store.Transaction()
  120. if err != nil {
  121. return cn, errors.Wrap(err, "couldn't start a new transaction")
  122. }
  123. cn = &Connection{
  124. ID: randx.String(30),
  125. Store: tx,
  126. Dialect: c.Dialect,
  127. TX: tx,
  128. }
  129. } else {
  130. cn = c
  131. }
  132. return cn, nil
  133. }
  134. // Rollback will open a new transaction and automatically rollback that transaction
  135. // when the inner function returns, regardless. This can be useful for tests, etc...
  136. func (c *Connection) Rollback(fn func(tx *Connection)) error {
  137. var cn *Connection
  138. if c.TX == nil {
  139. tx, err := c.Store.Transaction()
  140. if err != nil {
  141. return errors.Wrap(err, "couldn't start a new transaction")
  142. }
  143. cn = &Connection{
  144. ID: randx.String(30),
  145. Store: tx,
  146. Dialect: c.Dialect,
  147. TX: tx,
  148. }
  149. } else {
  150. cn = c
  151. }
  152. fn(cn)
  153. return cn.TX.Rollback()
  154. }
  155. // Q creates a new "empty" query for the current connection.
  156. func (c *Connection) Q() *Query {
  157. return Q(c)
  158. }
  159. // TruncateAll truncates all data from the datasource
  160. func (c *Connection) TruncateAll() error {
  161. return c.Dialect.TruncateAll(c)
  162. }
  163. func (c *Connection) timeFunc(name string, fn func() error) error {
  164. now := time.Now()
  165. err := fn()
  166. atomic.AddInt64(&c.Elapsed, int64(time.Now().Sub(now)))
  167. if err != nil {
  168. return errors.WithStack(err)
  169. }
  170. return nil
  171. }