/core/module/database.go

https://github.com/Clivern/Beetle · Go · 216 lines · 137 code · 53 blank · 26 comment · 14 complexity · 4ba26518c96edf3f14fe79da1e74dff9 MD5 · raw file

  1. // Copyright 2020 Clivern. All rights reserved.
  2. // Use of this source code is governed by the MIT
  3. // license that can be found in the LICENSE file.
  4. package module
  5. import (
  6. "fmt"
  7. "time"
  8. "github.com/clivern/beetle/core/migration"
  9. "github.com/clivern/beetle/core/model"
  10. "github.com/jinzhu/gorm"
  11. _ "github.com/jinzhu/gorm/dialects/mysql"
  12. _ "github.com/jinzhu/gorm/dialects/sqlite"
  13. log "github.com/sirupsen/logrus"
  14. "github.com/spf13/viper"
  15. )
  16. // Database struct
  17. type Database struct {
  18. Connection *gorm.DB
  19. }
  20. // Connect connects to a MySQL database
  21. func (db *Database) Connect(dsn model.DSN) error {
  22. var err error
  23. // Reuse db connections http://go-database-sql.org/surprises.html
  24. if db.Ping() == nil {
  25. return nil
  26. }
  27. db.Connection, err = gorm.Open(dsn.Driver, dsn.ToString())
  28. if err != nil {
  29. return err
  30. }
  31. return nil
  32. }
  33. // Ping check the db connection
  34. func (db *Database) Ping() error {
  35. if db.Connection == nil {
  36. return fmt.Errorf("No DB Connections Found")
  37. }
  38. err := db.Connection.DB().Ping()
  39. if err != nil {
  40. return err
  41. }
  42. // Cleanup stale connections http://go-database-sql.org/surprises.html
  43. db.Connection.DB().SetMaxOpenConns(5)
  44. db.Connection.DB().SetConnMaxLifetime(time.Duration(10) * time.Second)
  45. dbStats := db.Connection.DB().Stats()
  46. log.WithFields(log.Fields{
  47. "dbStats.maxOpenConnections": int(dbStats.MaxOpenConnections),
  48. "dbStats.openConnections": int(dbStats.OpenConnections),
  49. "dbStats.inUse": int(dbStats.InUse),
  50. "dbStats.idle": int(dbStats.Idle),
  51. }).Debug(`Open DB Connection`)
  52. return nil
  53. }
  54. // AutoConnect connects to a MySQL database using loaded configs
  55. func (db *Database) AutoConnect() error {
  56. var err error
  57. // Reuse db connections http://go-database-sql.org/surprises.html
  58. if db.Ping() == nil {
  59. return nil
  60. }
  61. dsn := model.DSN{
  62. Driver: viper.GetString("app.database.driver"),
  63. Username: viper.GetString("app.database.username"),
  64. Password: viper.GetString("app.database.password"),
  65. Hostname: viper.GetString("app.database.host"),
  66. Port: viper.GetInt("app.database.port"),
  67. Name: viper.GetString("app.database.name"),
  68. }
  69. db.Connection, err = gorm.Open(dsn.Driver, dsn.ToString())
  70. if err != nil {
  71. return err
  72. }
  73. return nil
  74. }
  75. // Migrate migrates the database
  76. func (db *Database) Migrate() bool {
  77. status := true
  78. db.Connection.AutoMigrate(&migration.Job{})
  79. status = status && db.Connection.HasTable(&migration.Job{})
  80. return status
  81. }
  82. // Rollback drop tables
  83. func (db *Database) Rollback() bool {
  84. status := true
  85. db.Connection.DropTableIfExists(&migration.Job{})
  86. status = status && !db.Connection.HasTable(&migration.Job{})
  87. return status
  88. }
  89. // HasTable checks if table exists
  90. func (db *Database) HasTable(table string) bool {
  91. return db.Connection.HasTable(table)
  92. }
  93. // CreateJob creates a new job
  94. func (db *Database) CreateJob(job *model.Job) *model.Job {
  95. db.Connection.Create(job)
  96. return job
  97. }
  98. // JobExistByID check if job exists
  99. func (db *Database) JobExistByID(id int) bool {
  100. job := model.Job{}
  101. db.Connection.Where("id = ?", id).First(&job)
  102. return job.ID > 0
  103. }
  104. // GetJobByID gets a job by id
  105. func (db *Database) GetJobByID(id int) model.Job {
  106. job := model.Job{}
  107. db.Connection.Where("id = ?", id).First(&job)
  108. return job
  109. }
  110. // GetJobs gets jobs
  111. func (db *Database) GetJobs() []model.Job {
  112. jobs := []model.Job{}
  113. db.Connection.Select("*").Find(&jobs)
  114. return jobs
  115. }
  116. // JobExistByUUID check if job exists
  117. func (db *Database) JobExistByUUID(uuid string) bool {
  118. job := model.Job{}
  119. db.Connection.Where("uuid = ?", uuid).First(&job)
  120. return job.ID > 0
  121. }
  122. // GetJobByUUID gets a job by uuid
  123. func (db *Database) GetJobByUUID(uuid string) model.Job {
  124. job := model.Job{}
  125. db.Connection.Where("uuid = ?", uuid).First(&job)
  126. return job
  127. }
  128. // GetPendingJobByType gets a job by uuid
  129. func (db *Database) GetPendingJobByType(jobType string) model.Job {
  130. job := model.Job{}
  131. db.Connection.Where("status = ? AND type = ?", model.JobPending, jobType).First(&job)
  132. return job
  133. }
  134. // CountJobs count jobs by status
  135. func (db *Database) CountJobs(status string) int {
  136. count := 0
  137. db.Connection.Model(&model.Job{}).Where("status = ?", status).Count(&count)
  138. return count
  139. }
  140. // DeleteJobByID deletes a job by id
  141. func (db *Database) DeleteJobByID(id int) {
  142. db.Connection.Unscoped().Where("id=?", id).Delete(&migration.Job{})
  143. }
  144. // DeleteJobByUUID deletes a job by uuid
  145. func (db *Database) DeleteJobByUUID(uuid string) {
  146. db.Connection.Unscoped().Where("uuid=?", uuid).Delete(&migration.Job{})
  147. }
  148. // UpdateJobByID updates a job by ID
  149. func (db *Database) UpdateJobByID(job *model.Job) {
  150. db.Connection.Save(&job)
  151. }
  152. // Close closes MySQL database connection
  153. func (db *Database) Close() error {
  154. return db.Connection.Close()
  155. }
  156. // ReleaseChildJobs count jobs by status
  157. func (db *Database) ReleaseChildJobs(parentID int) {
  158. db.Connection.Model(&model.Job{}).Where(
  159. "parent = ? AND status = ?",
  160. parentID,
  161. model.JobOnHold,
  162. ).Update("status", model.JobPending)
  163. }