@@ -4,74 +4,96 @@ import (
44 "fmt"
55 "time"
66
7- "github.com/fastbill/go-service-toolkit/v3 /observance"
7+ "github.com/fastbill/go-service-toolkit/v4 /observance"
88
9- // import migrate mysql, postgres driver
10- _ "github.com/golang-migrate/migrate/v4/database/mysql"
11- _ "github.com/golang-migrate/migrate/v4/database/postgres"
12- "github.com/jinzhu/gorm"
9+ "gorm.io/gorm"
10+ gormlogger "gorm.io/gorm/logger"
1311)
1412
15- // GormLogrus is a logrus logger that implements the gorm interface for logging.
16- type GormLogrus struct {
17- observance.Logger
18- }
19-
20- // Print implements the gorm.LogWriter interface, courtesy of https://gist.github.com/bnadland/2e4287b801a47dcfcc94.
21- func (g GormLogrus ) Print (v ... interface {}) {
22- if v [0 ] == "sql" {
23- g .WithFields (observance.Fields {"source" : "go-service-toolkit/database" }).Debug (fmt .Sprintf ("%v - %v" , v [3 ], v [4 ]))
24- }
25- if v [0 ] == "log" {
26- g .WithFields (observance.Fields {"source" : "go-service-toolkit/database" }).Debug (fmt .Sprintf ("%v" , v [2 ]))
27- }
28- }
29-
3013// SetupGORM loads the ORM with the given configuration
3114// The setup includes sending a ping and creating the database if it didn't exist.
3215// A logger will be activated if logLevel is 'debug'.
3316func SetupGORM (config Config , logger observance.Logger ) (* gorm.DB , error ) {
34- // We have two connection strings:
35- // 1) For connecting to the server (and maybe creating the database)
36- // 2) For connecting to the database directly.
3717 dbName := config .Name
18+
19+ // First we connect without the database name so we can create the database if it does not exist.
3820 config .Name = ""
39- connectionStringWithoutDatabase := config .ConnectionString ()
40- config .Name = dbName
41- connectionString := config .ConnectionString ()
21+ driverWithoutDatabaseSet := config .Driver ()
22+
23+ gormConfig := & gorm.Config {
24+ Logger : createLogger (logger ),
25+ }
4226
4327 // Open includes sending a ping.
44- db , err := gorm .Open (config . Dialect , connectionStringWithoutDatabase )
28+ db , err := gorm .Open (driverWithoutDatabaseSet , gormConfig )
4529 if err != nil {
46- return nil , err
30+ return nil , fmt . Errorf ( "failed to open DB connection: %w" , err )
4731 }
4832
49- if config .Name != "" {
33+ if dbName != "" {
34+ config .Name = dbName
5035 // Ensure the DB exists.
5136 db .Exec (fmt .Sprintf (config .createDatabaseQuery (), config .Name ))
52- err := db . Close ()
37+ err = Close (db )
5338 if err != nil {
54- return nil , err
39+ return nil , fmt . Errorf ( "failed to close DB connection: %w" , err )
5540 }
5641
5742 // Connect again with DB name.
58- db , err = gorm .Open (config .Dialect , connectionString )
43+ driver := config .Driver ()
44+ db , err = gorm .Open (driver , gormConfig )
5945 if err != nil {
60- return nil , err
46+ return nil , fmt . Errorf ( "failed to open DB connection: %w" , err )
6147 }
6248 }
6349
64- if logger .Level () == "debug" || logger .Level () == "trace" {
65- db .LogMode (true )
66- gormLogger := GormLogrus {logger }
67- db .SetLogger (gormLogger )
68- } else {
69- db .LogMode (false )
50+ dbConn , err := db .DB ()
51+ if err != nil {
52+ return nil , fmt .Errorf ("failed to retrieve DB connection: %w" , err )
7053 }
7154
7255 // This setting addresses "invalid connection" errors in case of connections being closed by the DB server after the wait_timeout (8h).
7356 // See https://github.com/go-sql-driver/mysql/issues/657.
74- db . DB (). SetConnMaxLifetime (3500 * time .Second )
57+ dbConn . SetConnMaxLifetime (time .Hour )
7558
7659 return db , nil
7760}
61+
62+ // Close closes the database connection(s) used by GORM.
63+ func Close (db * gorm.DB ) error {
64+ dbConn , err := db .DB ()
65+ if err != nil {
66+ return fmt .Errorf ("failed to retrieve DB connection: %w" , err )
67+ }
68+
69+ return dbConn .Close ()
70+ }
71+
72+ // GormWriter implements the Writer interface for setting up the GORM logger.
73+ type GormWriter struct {
74+ observance.Logger
75+ }
76+
77+ // Printf writes a log entry.
78+ func (g GormWriter ) Printf (msg string , data ... interface {}) {
79+ g .Logger .Debug (fmt .Sprintf (msg , data ... ))
80+ }
81+
82+ func createLogger (logger observance.Logger ) gormlogger.Interface {
83+ var logLevel gormlogger.LogLevel
84+ if logger .Level () == "debug" || logger .Level () == "trace" {
85+ logLevel = gormlogger .Info
86+ } else {
87+ logLevel = gormlogger .Silent
88+ }
89+
90+ newLogger := gormlogger .New (
91+ GormWriter {Logger : logger },
92+ gormlogger.Config {
93+ LogLevel : logLevel ,
94+ Colorful : false ,
95+ },
96+ )
97+
98+ return newLogger
99+ }
0 commit comments