dbmethods.go 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132
  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 (db *DBMethods) fixQuery(query string) string {
  15. if db.Driver == "mysql" {
  16. return fixQuery(query)
  17. }
  18. return query
  19. }
  20. func (db *DBMethods) Begin(ctx context.Context, opts *sql.TxOptions) (*Tx, error) {
  21. if db.Debug {
  22. t := time.Now()
  23. tx, err := db.DB.BeginTx(ctx, opts)
  24. log(os.Stdout, "[func Begin]", t, err, true, "")
  25. return &Tx{tx, db.Debug, db.Driver, t}, err
  26. }
  27. tx, err := db.DB.BeginTx(ctx, opts)
  28. if err != nil {
  29. return nil, err
  30. }
  31. return &Tx{tx, db.Debug, db.Driver, time.Now()}, err
  32. }
  33. func (db *DBMethods) Close() error {
  34. if db.Debug {
  35. t := time.Now()
  36. err := db.DB.Close()
  37. log(os.Stdout, "[func Close]", t, err, false, "")
  38. return err
  39. }
  40. return db.DB.Close()
  41. }
  42. func (db *DBMethods) Exec(ctx context.Context, query string, args ...any) (sql.Result, error) {
  43. if db.Debug {
  44. t := time.Now()
  45. res, err := db.DB.ExecContext(ctx, db.fixQuery(query), args...)
  46. log(os.Stdout, "[func Exec]", t, err, false, db.fixQuery(query), args...)
  47. return res, err
  48. }
  49. return db.DB.ExecContext(ctx, db.fixQuery(query), args...)
  50. }
  51. func (db *DBMethods) Ping(ctx context.Context) error {
  52. if db.Debug {
  53. t := time.Now()
  54. err := db.DB.PingContext(ctx)
  55. log(os.Stdout, "[func Ping]", t, err, false, "")
  56. return err
  57. }
  58. return db.DB.PingContext(ctx)
  59. }
  60. func (db *DBMethods) Prepare(ctx context.Context, query string) (*sql.Stmt, error) {
  61. if db.Debug {
  62. t := time.Now()
  63. stm, err := db.DB.PrepareContext(ctx, db.fixQuery(query))
  64. log(os.Stdout, "[func Prepare]", t, err, false, db.fixQuery(query))
  65. return stm, err
  66. }
  67. return db.DB.PrepareContext(ctx, db.fixQuery(query))
  68. }
  69. func (db *DBMethods) Query(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
  70. if db.Debug {
  71. t := time.Now()
  72. rows, err := db.DB.QueryContext(ctx, db.fixQuery(query), args...)
  73. log(os.Stdout, "[func Query]", t, err, false, db.fixQuery(query), args...)
  74. return rows, err
  75. }
  76. return db.DB.QueryContext(ctx, db.fixQuery(query), args...)
  77. }
  78. func (db *DBMethods) QueryRow(ctx context.Context, query string, args ...any) *sql.Row {
  79. if db.Debug {
  80. t := time.Now()
  81. row := db.DB.QueryRowContext(ctx, db.fixQuery(query), args...)
  82. log(os.Stdout, "[func QueryRow]", t, nil, false, db.fixQuery(query), args...)
  83. return row
  84. }
  85. return db.DB.QueryRowContext(ctx, db.fixQuery(query), args...)
  86. }
  87. func (db *DBMethods) SetConnMaxLifetime(d time.Duration) {
  88. db.DB.SetConnMaxLifetime(d)
  89. }
  90. func (db *DBMethods) SetMaxIdleConns(n int) {
  91. db.DB.SetMaxIdleConns(n)
  92. }
  93. func (db *DBMethods) SetMaxOpenConns(n int) {
  94. db.DB.SetMaxOpenConns(n)
  95. }
  96. func (db *DBMethods) Transaction(ctx context.Context, queries func(ctx context.Context, tx *Tx) error) error {
  97. if queries == nil {
  98. return fmt.Errorf("queries is not set for transaction")
  99. }
  100. tx, err := db.Begin(ctx, nil)
  101. if err != nil {
  102. return err
  103. }
  104. if err := queries(ctx, tx); err != nil {
  105. rerr := tx.Rollback()
  106. if rerr != nil {
  107. return fmt.Errorf(
  108. "%v: %v",
  109. rerr,
  110. err,
  111. )
  112. }
  113. return err
  114. }
  115. return tx.Commit()
  116. }