11package database
22
33import (
4+ "errors"
5+ "log"
46 "os"
57 "path/filepath"
68 "time"
79
810 "github.com/gotify/server/v2/auth/password"
11+ "github.com/gotify/server/v2/mode"
912 "github.com/gotify/server/v2/model"
10- "github.com/jinzhu/gorm"
11- _ "github.com/jinzhu/gorm/dialects/mysql" // enable the mysql dialect.
12- _ "github.com/jinzhu/gorm/dialects/postgres" // enable the postgres dialect.
13- _ "github.com/jinzhu/gorm/dialects/sqlite" // enable the sqlite3 dialect.
13+ "github.com/mattn/go-isatty"
14+ "gorm.io/driver/mysql"
15+ "gorm.io/driver/postgres"
16+ "gorm.io/driver/sqlite"
17+ "gorm.io/gorm"
18+ "gorm.io/gorm/logger"
1419)
1520
1621var mkdirAll = os .MkdirAll
@@ -19,40 +24,68 @@ var mkdirAll = os.MkdirAll
1924func New (dialect , connection , defaultUser , defaultPass string , strength int , createDefaultUserIfNotExist bool ) (* GormDatabase , error ) {
2025 createDirectoryIfSqlite (dialect , connection )
2126
22- db , err := gorm .Open (dialect , connection )
27+ logLevel := logger .Info
28+ if mode .Get () == mode .Prod {
29+ logLevel = logger .Warn
30+ }
31+
32+ dbLogger := logger .New (log .New (os .Stderr , "\r \n " , log .LstdFlags ), logger.Config {
33+ SlowThreshold : 200 * time .Millisecond ,
34+ LogLevel : logLevel ,
35+ IgnoreRecordNotFoundError : true ,
36+ Colorful : isatty .IsTerminal (os .Stderr .Fd ()),
37+ })
38+ gormConfig := & gorm.Config {
39+ Logger : dbLogger ,
40+ DisableForeignKeyConstraintWhenMigrating : true ,
41+ }
42+
43+ var db * gorm.DB
44+ err := errors .New ("unsupported dialect: " + dialect )
45+
46+ switch dialect {
47+ case "mysql" :
48+ db , err = gorm .Open (mysql .Open (connection ), gormConfig )
49+ case "postgres" :
50+ db , err = gorm .Open (postgres .Open (connection ), gormConfig )
51+ case "sqlite3" :
52+ db , err = gorm .Open (sqlite .Open (connection ), gormConfig )
53+ }
54+
55+ if err != nil {
56+ return nil , err
57+ }
58+
59+ sqldb , err := db .DB ()
2360 if err != nil {
2461 return nil , err
2562 }
2663
2764 // We normally don't need that much connections, so we limit them. F.ex. mysql complains about
2865 // "too many connections", while load testing Gotify.
29- db . DB () .SetMaxOpenConns (10 )
66+ sqldb .SetMaxOpenConns (10 )
3067
3168 if dialect == "sqlite3" {
3269 // We use the database connection inside the handlers from the http
3370 // framework, therefore concurrent access occurs. Sqlite cannot handle
3471 // concurrent writes, so we limit sqlite to one connection.
3572 // see https://github.com/mattn/go-sqlite3/issues/274
36- db . DB () .SetMaxOpenConns (1 )
73+ sqldb .SetMaxOpenConns (1 )
3774 }
3875
3976 if dialect == "mysql" {
4077 // Mysql has a setting called wait_timeout, which defines the duration
4178 // after which a connection may not be used anymore.
4279 // The default for this setting on mariadb is 10 minutes.
4380 // See https://github.com/docker-library/mariadb/issues/113
44- db . DB () .SetConnMaxLifetime (9 * time .Minute )
81+ sqldb .SetConnMaxLifetime (9 * time .Minute )
4582 }
4683
47- if err := db .AutoMigrate (new (model.User ), new (model.Application ), new (model.Message ), new (model.Client ), new (model.PluginConf )). Error ; err != nil {
84+ if err := db .AutoMigrate (new (model.User ), new (model.Application ), new (model.Message ), new (model.Client ), new (model.PluginConf )); err != nil {
4885 return nil , err
4986 }
5087
51- if err := prepareBlobColumn (dialect , db ); err != nil {
52- return nil , err
53- }
54-
55- userCount := 0
88+ userCount := int64 (0 )
5689 db .Find (new (model.User )).Count (& userCount )
5790 if createDefaultUserIfNotExist && userCount == 0 {
5891 db .Create (& model.User {Name : defaultUser , Pass : password .CreatePassword (defaultPass , strength ), Admin : true })
@@ -61,31 +94,6 @@ func New(dialect, connection, defaultUser, defaultPass string, strength int, cre
6194 return & GormDatabase {DB : db }, nil
6295}
6396
64- func prepareBlobColumn (dialect string , db * gorm.DB ) error {
65- blobType := ""
66- switch dialect {
67- case "mysql" :
68- blobType = "longblob"
69- case "postgres" :
70- blobType = "bytea"
71- }
72- if blobType != "" {
73- for _ , target := range []struct {
74- Table interface {}
75- Column string
76- }{
77- {model.Message {}, "extras" },
78- {model.PluginConf {}, "config" },
79- {model.PluginConf {}, "storage" },
80- } {
81- if err := db .Model (target .Table ).ModifyColumn (target .Column , blobType ).Error ; err != nil {
82- return err
83- }
84- }
85- }
86- return nil
87- }
88-
8997func createDirectoryIfSqlite (dialect , connection string ) {
9098 if dialect == "sqlite3" {
9199 if _ , err := os .Stat (filepath .Dir (connection )); os .IsNotExist (err ) {
@@ -103,5 +111,9 @@ type GormDatabase struct {
103111
104112// Close closes the gorm database connection.
105113func (d * GormDatabase ) Close () {
106- d .DB .Close ()
114+ sqldb , err := d .DB .DB ()
115+ if err != nil {
116+ return
117+ }
118+ sqldb .Close ()
107119}
0 commit comments