/models/db.go

https://github.com/iwannay/jiacrontab · Go · 93 lines · 76 code · 16 blank · 1 comment · 23 complexity · 66c2ea6294c9ffad5493d147c1fa81e7 MD5 · raw file

  1. package models
  2. import (
  3. "errors"
  4. "fmt"
  5. "os"
  6. "path/filepath"
  7. "github.com/jinzhu/gorm"
  8. )
  9. // D alias DB
  10. type D = gorm.DB
  11. var (
  12. db *D
  13. debugMode bool
  14. )
  15. func CreateDB(dialect string, args ...interface{}) error {
  16. switch dialect {
  17. case "sqlite3":
  18. return createSqlite3(dialect, args...)
  19. case "postgres", "mysql":
  20. var err error
  21. db, err = gorm.Open(dialect, args...)
  22. return err
  23. }
  24. return fmt.Errorf("unknow database type %s", dialect)
  25. }
  26. func createSqlite3(dialect string, args ...interface{}) error {
  27. var err error
  28. if args[0] == nil {
  29. return errors.New("sqlite3:db file cannot empty")
  30. }
  31. dbDir := filepath.Dir(filepath.Clean(fmt.Sprint(args[0])))
  32. err = os.MkdirAll(dbDir, 0755)
  33. if err != nil {
  34. return fmt.Errorf("sqlite3:%s", err)
  35. }
  36. db, err = gorm.Open(dialect, args...)
  37. if err == nil {
  38. db.DB().SetMaxOpenConns(1)
  39. }
  40. return err
  41. }
  42. func DB() *D {
  43. if db == nil {
  44. panic("you must call CreateDb first")
  45. }
  46. if debugMode {
  47. return db.Debug()
  48. }
  49. return db
  50. }
  51. func Transactions(fn func(tx *gorm.DB) error) error {
  52. if fn == nil {
  53. return errors.New("fn is nil")
  54. }
  55. tx := DB().Begin()
  56. defer func() {
  57. if err := recover(); err != nil {
  58. DB().Rollback()
  59. }
  60. }()
  61. if fn(tx) != nil {
  62. tx.Rollback()
  63. }
  64. return tx.Commit().Error
  65. }
  66. func InitModel(driverName string, dsn string, debug bool) error {
  67. if driverName == "" || dsn == "" {
  68. return errors.New("driverName and dsn cannot empty")
  69. }
  70. if err := CreateDB(driverName, dsn); err != nil {
  71. return err
  72. }
  73. debugMode = debug
  74. DB().AutoMigrate(&Node{}, &Group{}, &User{}, &Event{}, &JobHistory{})
  75. DB().Create(&SuperGroup)
  76. return nil
  77. }