/service/db.go

https://github.com/zou2699/mypipe · Go · 83 lines · 52 code · 12 blank · 19 comment · 17 complexity · b068c29eccdabc114f7c126aff9114d5 MD5 · raw file

  1. // Pipe - A small and beautiful blogging platform written in golang.
  2. // Copyright (C) 2017-2019, b3log.org & hacpai.com
  3. //
  4. // This program is free software: you can redistribute it and/or modify
  5. // it under the terms of the GNU General Public License as published by
  6. // the Free Software Foundation, either version 3 of the License, or
  7. // (at your option) any later version.
  8. //
  9. // This program is distributed in the hope that it will be useful,
  10. // but WITHOUT ANY WARRANTY; without even the implied warranty of
  11. // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  12. // GNU General Public License for more details.
  13. //
  14. // You should have received a copy of the GNU General Public License
  15. // along with this program. If not, see <http://www.gnu.org/licenses/>.
  16. package service
  17. import (
  18. "os"
  19. "github.com/jinzhu/gorm"
  20. _ "github.com/jinzhu/gorm/dialects/mysql" // mysql
  21. _ "github.com/jinzhu/gorm/dialects/sqlite" // sqlite
  22. "github.com/zou2699/mypipe/log"
  23. "github.com/zou2699/mypipe/model"
  24. )
  25. // Logger
  26. var logger = log.NewLogger(os.Stdout)
  27. var db *gorm.DB
  28. var useSQLite bool
  29. // ConnectDB connects to the database.
  30. func ConnectDB() {
  31. var err error
  32. useSQLite = false
  33. if "" != model.Conf.SQLite {
  34. db, err = gorm.Open("sqlite3", model.Conf.SQLite)
  35. useSQLite = true
  36. } else if "" != model.Conf.MySQL {
  37. db, err = gorm.Open("mysql", model.Conf.MySQL)
  38. } else {
  39. logger.Fatal("please specify database")
  40. }
  41. if nil != err {
  42. logger.Fatalf("opens database failed: " + err.Error())
  43. }
  44. if useSQLite {
  45. logger.Debug("used [SQLite] as underlying database")
  46. } else {
  47. logger.Debug("used [MySQL] as underlying database")
  48. }
  49. if err = db.AutoMigrate(model.Models...).Error; nil != err {
  50. logger.Fatal("auto migrate tables failed: " + err.Error())
  51. }
  52. if err = db.Model(&model.Article{}).AddIndex("idx_b3_pipe_articles_created_at", "created_at").Error; nil != err {
  53. logger.Fatal("adds index failed: " + err.Error())
  54. }
  55. db.DB().SetMaxIdleConns(10)
  56. db.DB().SetMaxOpenConns(50)
  57. db.LogMode(model.Conf.ShowSQL)
  58. }
  59. // DisconnectDB disconnects from the database.
  60. func DisconnectDB() {
  61. if err := db.Close(); nil != err {
  62. logger.Errorf("Disconnect from database failed: " + err.Error())
  63. }
  64. }
  65. // Database returns the underlying database name.
  66. func Database() string {
  67. if useSQLite {
  68. return "SQLite"
  69. }
  70. return "MySQL"
  71. }