init.go 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273
  1. // Copyright 2019 getensh.com. All rights reserved.
  2. // Use of this source code is governed by getensh.com.
  3. package database
  4. import (
  5. "fmt"
  6. "git.getensh.com/common/gopkgsv2/logger"
  7. "go.uber.org/zap"
  8. "gorm.io/driver/mysql"
  9. "gorm.io/gorm"
  10. "log"
  11. "os"
  12. "time"
  13. )
  14. var (
  15. db *gorm.DB
  16. )
  17. // Setup 建立连接
  18. func Setup(user, passwd, addr, dbname, charset string, maxIdle, maxConn int, logMode bool) *gorm.DB {
  19. conf := Config{
  20. SlowThreshold: time.Second,
  21. Colorful: true,
  22. IgnoreRecordNotFoundError: false,
  23. LogLevel: zap.ErrorLevel,
  24. TraceWithLevel: zap.ErrorLevel,
  25. Zap: logger.Logger,
  26. }
  27. // 是否开启调试模式
  28. if logMode {
  29. conf.LogLevel = zap.DebugLevel
  30. conf.TraceWithLevel = zap.DebugLevel
  31. conf.SlowThreshold = time.Millisecond * 200
  32. }
  33. config := &gorm.Config{
  34. Logger: NewGormLogger(log.New(os.Stdout, "\r\n", log.LstdFlags), conf),
  35. }
  36. // 组装参数
  37. dsn := fmt.Sprintf("%s:%s@tcp(%s)/%s?charset=%s&parseTime=True&loc=Local",
  38. user, passwd, addr, dbname, charset)
  39. // 打开新连接
  40. var err error
  41. db, err = gorm.Open(mysql.Open(dsn), config)
  42. if err != nil {
  43. log.Fatal("open mysql connection failed. err: ", err)
  44. }
  45. // 其他设置
  46. sqlDB, err := db.DB()
  47. if err != nil {
  48. log.Fatal("open mysql connection failed. err: ", err)
  49. }
  50. sqlDB.SetMaxIdleConns(maxIdle)
  51. sqlDB.SetMaxOpenConns(maxConn)
  52. RegisterCallback(db)
  53. return db
  54. }
  55. // DB 获取连接
  56. func DB() *gorm.DB {
  57. return db
  58. }
  59. // Close 关闭连接
  60. func Close() {
  61. if db != nil {
  62. sqlDB, _ := db.DB()
  63. _ = sqlDB.Close()
  64. }
  65. }