/internal/entity/db.go

https://github.com/photoprism/photoprism · Go · 113 lines · 79 code · 24 blank · 10 comment · 22 complexity · f24419568bcb298653caa766ce4faadb MD5 · raw file

  1. package entity
  2. import (
  3. "fmt"
  4. "sync"
  5. "time"
  6. "github.com/jinzhu/gorm"
  7. _ "github.com/jinzhu/gorm/dialects/mysql"
  8. _ "github.com/jinzhu/gorm/dialects/sqlite"
  9. )
  10. // Database drivers (sql dialects).
  11. const (
  12. MySQL = "mysql"
  13. SQLite = "sqlite3"
  14. )
  15. var dbProvider DbProvider
  16. type DbProvider interface {
  17. Db() *gorm.DB
  18. }
  19. // IsDialect returns true if the given sql dialect is used.
  20. func IsDialect(name string) bool {
  21. return name == Db().Dialect().GetName()
  22. }
  23. // DbDialect returns the sql dialect name.
  24. func DbDialect() string {
  25. return Db().Dialect().GetName()
  26. }
  27. // SetDbProvider sets the provider to get a gorm db connection.
  28. func SetDbProvider(provider DbProvider) {
  29. dbProvider = provider
  30. }
  31. // HasDbProvider returns true if a db provider exists.
  32. func HasDbProvider() bool {
  33. return dbProvider != nil
  34. }
  35. // Db returns a database connection.
  36. func Db() *gorm.DB {
  37. return dbProvider.Db()
  38. }
  39. // UnscopedDb returns an unscoped database connection.
  40. func UnscopedDb() *gorm.DB {
  41. return Db().Unscoped()
  42. }
  43. type Gorm struct {
  44. Driver string
  45. Dsn string
  46. once sync.Once
  47. db *gorm.DB
  48. }
  49. // Db returns the gorm db connection.
  50. func (g *Gorm) Db() *gorm.DB {
  51. g.once.Do(g.Connect)
  52. if g.db == nil {
  53. log.Fatal("entity: database not connected")
  54. }
  55. return g.db
  56. }
  57. // Connect creates a new gorm db connection.
  58. func (g *Gorm) Connect() {
  59. db, err := gorm.Open(g.Driver, g.Dsn)
  60. if err != nil || db == nil {
  61. for i := 1; i <= 12; i++ {
  62. fmt.Printf("gorm.Open(%s, %s) %d\n", g.Driver, g.Dsn, i)
  63. db, err = gorm.Open(g.Driver, g.Dsn)
  64. if db != nil && err == nil {
  65. break
  66. } else {
  67. time.Sleep(5 * time.Second)
  68. }
  69. }
  70. if err != nil || db == nil {
  71. fmt.Println(err)
  72. log.Fatal(err)
  73. }
  74. }
  75. db.LogMode(false)
  76. db.SetLogger(log)
  77. db.DB().SetMaxIdleConns(4)
  78. db.DB().SetMaxOpenConns(256)
  79. g.db = db
  80. }
  81. // Close closes the gorm db connection.
  82. func (g *Gorm) Close() {
  83. if g.db != nil {
  84. if err := g.db.Close(); err != nil {
  85. log.Fatal(err)
  86. }
  87. g.db = nil
  88. }
  89. }