dbmethods.go 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189
  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.Query)
  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) Query(ctx context.Context, query string, args ...any) (*Rows, error) {
  99. start := time.Now()
  100. rows, err := d.DB.QueryContext(ctx, d.fixQuery(query), args...)
  101. d.log("Query", start, err, false, d.fixQuery(query), args...)
  102. return &Rows{Rows: rows}, err
  103. }
  104. func (d *DBMethods) QueryPrepared(ctx context.Context, prep *Prepared) (*Rows, error) {
  105. return d.Query(ctx, prep.Query, prep.Args...)
  106. }
  107. func (d *DBMethods) QueryRow(ctx context.Context, query string, args ...any) *Row {
  108. start := time.Now()
  109. row := d.DB.QueryRowContext(ctx, d.fixQuery(query), args...)
  110. d.log("QueryRow", start, nil, false, d.fixQuery(query), args...)
  111. return &Row{Row: row}
  112. }
  113. func (d *DBMethods) QueryRowByID(ctx context.Context, id int64, row any) error {
  114. query := queryRowByIDString(row)
  115. return d.QueryRow(ctx, query, id).Scans(row)
  116. }
  117. func (d *DBMethods) QueryRowPrepared(ctx context.Context, prep *Prepared) *Row {
  118. return d.QueryRow(ctx, prep.Query, prep.Args...)
  119. }
  120. func (d *DBMethods) RowExists(ctx context.Context, id int64, row any) bool {
  121. var exists int
  122. query := rowExistsString(row)
  123. if err := d.QueryRow(ctx, query, id).Scan(&exists); err == nil && exists == 1 {
  124. return true
  125. }
  126. return false
  127. }
  128. func (d *DBMethods) SetConnMaxLifetime(t time.Duration) {
  129. start := time.Now()
  130. d.DB.SetConnMaxLifetime(t)
  131. d.log("SetConnMaxLifetime", start, nil, false, "")
  132. }
  133. func (d *DBMethods) SetMaxIdleConns(n int) {
  134. start := time.Now()
  135. d.DB.SetMaxIdleConns(n)
  136. d.log("SetMaxIdleConns", start, nil, false, "")
  137. }
  138. func (d *DBMethods) SetMaxOpenConns(n int) {
  139. start := time.Now()
  140. d.DB.SetMaxOpenConns(n)
  141. d.log("SetMaxOpenConns", start, nil, false, "")
  142. }
  143. func (d *DBMethods) Transaction(ctx context.Context, callback func(ctx context.Context, tx *Tx) error) error {
  144. if callback == nil {
  145. return fmt.Errorf("callback is not set")
  146. }
  147. tx, err := d.Begin(ctx, nil)
  148. if err != nil {
  149. return err
  150. }
  151. if err := callback(ctx, tx); err != nil {
  152. rerr := tx.Rollback()
  153. if rerr != nil {
  154. return fmt.Errorf(
  155. "%v: %v",
  156. rerr,
  157. err,
  158. )
  159. }
  160. return err
  161. }
  162. return tx.Commit()
  163. }