init.go 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102
  1. package db
  2. import (
  3. "fmt"
  4. _ "github.com/go-sql-driver/mysql"
  5. "github.com/jmoiron/sqlx"
  6. log "github.com/sirupsen/logrus"
  7. _ "modernc.org/sqlite"
  8. "pmail/config"
  9. "pmail/dto"
  10. )
  11. var Instance *sqlx.DB
  12. func Init() {
  13. dsn := config.Instance.DbDSN
  14. var err error
  15. switch config.Instance.DbType {
  16. case "mysql":
  17. Instance, err = sqlx.Open("mysql", dsn)
  18. case "sqlite":
  19. Instance, err = sqlx.Open("sqlite", dsn)
  20. default:
  21. return
  22. }
  23. if err != nil {
  24. panic(err)
  25. }
  26. Instance.SetMaxOpenConns(100)
  27. Instance.SetMaxIdleConns(10)
  28. //showMySQLCharacterSet()
  29. checkTable()
  30. }
  31. func WithContext(ctx *dto.Context, sql string) string {
  32. if ctx != nil {
  33. logId := ctx.GetValue(dto.LogID)
  34. return fmt.Sprintf("/* %s */ %s", logId, sql)
  35. }
  36. return sql
  37. }
  38. type tables struct {
  39. TablesInPmail string `db:"Tables_in_pmail"`
  40. }
  41. func checkTable() {
  42. var res []*tables
  43. var err error
  44. if config.Instance.DbType == "sqlite" {
  45. err = Instance.Select(&res, "select name as `Tables_in_pmail` from sqlite_master where type='table'")
  46. } else {
  47. err = Instance.Select(&res, "show tables")
  48. }
  49. if err != nil {
  50. panic(err)
  51. }
  52. existTable := map[string]struct{}{}
  53. for _, tableName := range res {
  54. existTable[tableName.TablesInPmail] = struct{}{}
  55. }
  56. for tableName, createSQL := range config.Instance.Tables {
  57. if _, ok := existTable[tableName]; !ok {
  58. _, err = Instance.Exec(createSQL)
  59. log.Infof("Create Table: %s", createSQL)
  60. if err != nil {
  61. panic(err)
  62. }
  63. if initData, ok := config.Instance.TablesInitData[tableName]; ok {
  64. _, err = Instance.Exec(initData)
  65. log.Infof("Init Table: %s", initData)
  66. if err != nil {
  67. panic(err)
  68. }
  69. }
  70. }
  71. }
  72. }
  73. func showMySQLCharacterSet() {
  74. var res []struct {
  75. Variable_name string `db:"Variable_name"`
  76. Value string `db:"Value"`
  77. }
  78. err := Instance.Select(&res, "show variables like '%character%';")
  79. log.Debugf("%+v %+v", res, err)
  80. }
  81. func testSlowLog() {
  82. var res []struct {
  83. Value string `db:"Value"`
  84. }
  85. err := Instance.Select(&res, "/* asddddasad */select /* this is test */ sleep(4) as Value")
  86. log.Debugf("%+v %+v", res, err)
  87. }