dbmethods.go 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205
  1. package common
  2. import (
  3. "context"
  4. "database/sql"
  5. "fmt"
  6. "os"
  7. "time"
  8. )
  9. type DBMethods struct {
  10. DB *sql.DB
  11. Debug bool
  12. Driver string
  13. }
  14. func (d *DBMethods) fixQuery(query string) string {
  15. if d.Driver == "mysql" {
  16. return fixQuery(query)
  17. }
  18. return query
  19. }
  20. func (d *DBMethods) log(fname string, start time.Time, err error, tx bool, query string, args ...any) {
  21. if d.Debug {
  22. log(os.Stdout, fname, start, err, tx, query, args...)
  23. }
  24. }
  25. func (d *DBMethods) Begin(ctx context.Context, opts *sql.TxOptions) (*Tx, error) {
  26. start := time.Now()
  27. tx, err := d.DB.BeginTx(ctx, opts)
  28. d.log("Begin", start, err, true, "")
  29. return &Tx{tx, d.Debug, d.Driver, start}, err
  30. }
  31. func (d *DBMethods) Close() error {
  32. start := time.Now()
  33. err := d.DB.Close()
  34. d.log("Close", start, err, false, "")
  35. return err
  36. }
  37. func (d *DBMethods) CurrentUnixTimestamp() int64 {
  38. return currentUnixTimestamp()
  39. }
  40. func (d *DBMethods) DeleteRowByID(ctx context.Context, id int64, row any) error {
  41. query := deleteRowByIDString(row)
  42. _, err := d.Exec(ctx, query, id)
  43. return err
  44. }
  45. func (d *DBMethods) Each(ctx context.Context, query string, callback func(ctx context.Context, rows *Rows) error, args ...any) error {
  46. if callback == nil {
  47. return fmt.Errorf("callback is not set")
  48. }
  49. rows, err := d.Query(ctx, query, args...)
  50. if err != nil {
  51. return err
  52. }
  53. defer rows.Close()
  54. for rows.Next() {
  55. select {
  56. case <-ctx.Done():
  57. return ctx.Err()
  58. default:
  59. if err := callback(ctx, rows); err != nil {
  60. return err
  61. }
  62. }
  63. }
  64. if err := rows.Err(); err != nil {
  65. return err
  66. }
  67. return nil
  68. }
  69. func (d *DBMethods) EachPrepared(ctx context.Context, prep *Prepared, callback func(ctx context.Context, rows *Rows) error) error {
  70. return d.Each(ctx, prep.Query, callback, prep.Args...)
  71. }
  72. func (d *DBMethods) Exec(ctx context.Context, query string, args ...any) (sql.Result, error) {
  73. start := time.Now()
  74. res, err := d.DB.ExecContext(ctx, d.fixQuery(query), args...)
  75. d.log("Exec", start, err, false, d.fixQuery(query), args...)
  76. return res, err
  77. }
  78. func (d *DBMethods) ExecPrepared(ctx context.Context, prep *Prepared) (sql.Result, error) {
  79. return d.Exec(ctx, prep.Query, prep.Args...)
  80. }
  81. func (d *DBMethods) InsertRow(ctx context.Context, row any) error {
  82. query, args := insertRowString(row)
  83. _, err := d.Exec(ctx, query, args...)
  84. return err
  85. }
  86. func (d *DBMethods) Ping(ctx context.Context) error {
  87. start := time.Now()
  88. err := d.DB.PingContext(ctx)
  89. d.log("Ping", start, err, false, "")
  90. return err
  91. }
  92. func (d *DBMethods) Prepare(ctx context.Context, query string) (*sql.Stmt, error) {
  93. start := time.Now()
  94. stm, err := d.DB.PrepareContext(ctx, d.fixQuery(query))
  95. d.log("Prepare", start, err, false, d.fixQuery(query))
  96. return stm, err
  97. }
  98. func (d *DBMethods) PrepareSQL(query string, args ...any) *Prepared {
  99. return prepareSQL(query, args...)
  100. }
  101. func (d *DBMethods) Query(ctx context.Context, query string, args ...any) (*Rows, error) {
  102. start := time.Now()
  103. rows, err := d.DB.QueryContext(ctx, d.fixQuery(query), args...)
  104. d.log("Query", start, err, false, d.fixQuery(query), args...)
  105. return &Rows{Rows: rows}, err
  106. }
  107. func (d *DBMethods) QueryPrepared(ctx context.Context, prep *Prepared) (*Rows, error) {
  108. return d.Query(ctx, prep.Query, prep.Args...)
  109. }
  110. func (d *DBMethods) QueryRow(ctx context.Context, query string, args ...any) *Row {
  111. start := time.Now()
  112. row := d.DB.QueryRowContext(ctx, d.fixQuery(query), args...)
  113. d.log("QueryRow", start, nil, false, d.fixQuery(query), args...)
  114. return &Row{Row: row}
  115. }
  116. func (d *DBMethods) QueryRowByID(ctx context.Context, id int64, row any) error {
  117. query := queryRowByIDString(row)
  118. return d.QueryRow(ctx, query, id).Scans(row)
  119. }
  120. func (d *DBMethods) QueryRowPrepared(ctx context.Context, prep *Prepared) *Row {
  121. return d.QueryRow(ctx, prep.Query, prep.Args...)
  122. }
  123. func (d *DBMethods) RowExists(ctx context.Context, id int64, row any) bool {
  124. var exists int
  125. query := rowExistsString(row)
  126. if err := d.QueryRow(ctx, query, id).Scan(&exists); err == nil && exists == 1 {
  127. return true
  128. }
  129. return false
  130. }
  131. func (d *DBMethods) SetConnMaxLifetime(t time.Duration) {
  132. start := time.Now()
  133. d.DB.SetConnMaxLifetime(t)
  134. d.log("SetConnMaxLifetime", start, nil, false, "")
  135. }
  136. func (d *DBMethods) SetMaxIdleConns(n int) {
  137. start := time.Now()
  138. d.DB.SetMaxIdleConns(n)
  139. d.log("SetMaxIdleConns", start, nil, false, "")
  140. }
  141. func (d *DBMethods) SetMaxOpenConns(n int) {
  142. start := time.Now()
  143. d.DB.SetMaxOpenConns(n)
  144. d.log("SetMaxOpenConns", start, nil, false, "")
  145. }
  146. func (d *DBMethods) Transaction(ctx context.Context, callback func(ctx context.Context, tx *Tx) error) error {
  147. if callback == nil {
  148. return fmt.Errorf("callback is not set")
  149. }
  150. tx, err := d.Begin(ctx, nil)
  151. if err != nil {
  152. return err
  153. }
  154. if err := callback(ctx, tx); err != nil {
  155. rerr := tx.Rollback()
  156. if rerr != nil {
  157. return fmt.Errorf(
  158. "%v: %v",
  159. rerr,
  160. err,
  161. )
  162. }
  163. return err
  164. }
  165. return tx.Commit()
  166. }
  167. func (d *DBMethods) UpdateRow(ctx context.Context, row any) error {
  168. query, args := updateRowString(row)
  169. _, err := d.Exec(ctx, query, args...)
  170. return err
  171. }
  172. func (d *DBMethods) UpdateRowOnly(ctx context.Context, row any, fields ...string) error {
  173. query, args := updateRowString(row, fields...)
  174. _, err := d.Exec(ctx, query, args...)
  175. return err
  176. }