dbmethods.go 4.5 KB

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